Native-Resolution Image Synthesis

TL;DR

NiT (NeurIPS 2025) is a diffusion transformer that trains directly on images at their original resolutions and aspect ratios — no resize, no crop, no aspect-ratio buckets, no padding. A single 675M NiT-XL achieves dual SOTA in class-guided ImageNet generation (2.08 FID at 256², 1.48 FID at 512²) and, despite being trained only on ImageNet, zero-shots to unseen resolutions and aspect ratios (4.52 FID at 1024², 4.11 FID at 432×768). The three load-bearing pieces are NaViT-style sequence packing for variable-length latents, FlashAttention-2 with cumulative-sequence-length masking for “packed full attention” within each instance, and axial 2D-RoPE to inject 2D structural priors decoupled across height and width.

Key claims

A single NiT-XL achieves FID 2.16 (256²), 1.57 (512²), 4.05 (768²), 4.52 (1024²) — i.e. dual SOTA on both standard ImageNet benchmarks with one model instead of the usual two resolution-specialist models [§4.2, Model Zoo table].
NiT-XL-1500K pushes 256² → 2.03 FID and 512² → 1.45 FID, beating EDM2-XXL at roughly half the parameter count [§4.2, Model Zoo].
The architecture handles native resolutions through three pieces: (1) Dynamic Tokenization — variable-length token sequences with (h, w) tuples, no padding; (2) Packed Full-Attention via FlashAttention-2’s cu_seqlens API, restricting attention within each instance without explicit masks; (3) axial 2D-RoPE that generates independent height/width rotary-frequency components [§3.2].
Packing uses Longest-Pack-First Histogram Packing (LPFHP) to pre-set image indices per pack at a fixed maximum sequence length L (16384 for NiT-XL), so the number of image instances per iteration is dynamic [§3.2, §4.1].
Packed Adaptive LayerNorm — per-instance scale/shift parameters broadcast across that instance’s tokens — replaces standard AdaLN to support the heterogeneous packed sequence [§3.2].
Training compute is reported in token budget (sum of tokens across iterations) rather than batch size, mirroring LLM conventions; standard runs use 2³³ tokens per iteration [§4.1].
Ablation finds that native-resolution data is essential for aspect-ratio generalization: training on only 256+512 (no native) matches native-only on those two benchmarks but collapses on 1024² and on non-square aspects [§4.3, Tables 4–5].
mFID (averaged across multiple resolutions) is reported as the unified metric; NiT-XL beats DiT-XL and SiT-XL on mFID with 50% of the training compute, and SiT-REPA with 75% of the token budget [§4.2].
Baselines fail in characteristic ways at extreme resolutions: EDM2-L produces non-informative textures past 1024², FlowDCN-XL duplicates objects, and both exhibit cropping-induced bias (truncated objects) at extreme aspect ratios — consistent with SDXL’s reported crop-conditioning effect [§4.2, Fig. 4].
NiT-B (130M) trains at 1.28 iter/s vs DiT-B’s 1.08 iter/s at matched token count per iteration, and infers at 0.246 s vs 0.322 s — packing is reportedly a net positive on efficiency over fixed-resolution training, not just a flexibility win [§4.3 Efficiency].
Authors use DC-AE (f32c32) as the image autoencoder (downsample 32, latent dim 32) rather than the standard f8 VAE; this is what makes 1024² and beyond tractable as a token count [§4.1].
For REPA-style auxiliary supervision NiT uses RADIO-v2.5-H as the encoder for >NiT-S models (RADIO-v2.5-B for NiT-S); the multi-resolution-trained RADIO is the natural pairing with multi-resolution diffusion training [Training Scripts section, repo README].

Method

NiT is built on the DiT backbone with patch size 1 (operating directly on DC-AE latent tokens). The pipeline: (1) encode each native-resolution image with DC-AE into a variable-shape latent z of size (32, h, w); (2) flatten z to a sequence of (h·w) tokens and concatenate multiple instances’ sequences into a single packed sequence of length ≤ L, with cumulative offsets (s_1, s_2, …, s_B) recording instance boundaries; (3) add per-instance noise and time according to a Logit-Normal time distribution under linear-path Flow Matching; (4) patch-embed to hidden size d; (5) attach axial 2D-RoPE built independently from each token’s (h, w) coordinate; (6) run the DiT blocks, with self-attention computed via FlashAttention-2’s variable-length API (cu_seqlens) so each instance attends only within itself, and with Packed AdaLN broadcasting instance-conditional scale/shift across that instance’s tokens.

The axial 2D-RoPE is the key generalization mechanism: base angular frequencies follow standard RoPE θ_i = 10000^(-2i/d_h) on a hidden-size-d/2 partition; the height angle vector is θ_i · h, the width angle vector is θ_i · w, and the two are concatenated before the standard pairwise rotation. This decouples the height and width rotational signals, so a model trained on (e.g.) 16×16 and 32×32 latent grids can be invoked at 12×18, 24×48, or 64×64 with consistent relative-position behavior.

The repository ships NiT-XL-1000K (256²: 2.16 FID; 512²: 1.57 FID; 768²: 4.05 FID; 1024²: 4.52 FID) and NiT-XL-1500K (256²: 2.03; 512²: 1.45). Sampling uses SDE with 250 NFE for 256/512 and ODE with 50 NFE for the higher resolutions, with CFG interval [0, 0.7] for most settings and [0, 0.9] at 1024+.

Results

Headline FIDs (NiT-XL-1000K, 675M):

256×256: 2.16 (NiT-XL-1500K: 2.03) — beats EDM2-XXL (FID 1.81 at 1.5B params) and SiT-XL-REPA at the same model scale.
512×512: 1.57 (NiT-XL-1500K: 1.45) — the same single model that wins at 256.
768×768: 4.05, 1024×1024: 4.52, 1536×1536: 6.51 — unseen resolutions, no resolution-specific training.
2048×2048: 24.76 FID — degrades, named as a limitation in §5.
Non-square aspects: 480×640 (3.72), 640×480 (3.41), 432×768 (4.11), 768×432 (5.27), 320×960 (16.85), 960×320 (9.90).

mFID comparison reports NiT-XL beats DiT-XL / SiT-XL / FiTv2-XL / SiT-REPA on resolution-averaged FID at 50–75% of their training compute [§4.2]. The ablation in §4.3 / Tables 4–5 is the source of the “native-resolution data is essential for generalization” claim — removing native data from the mix (only 256+512 training) leaves 1024² and non-square aspects badly broken even when 256/512 metrics are preserved.

Why it’s interesting

For a video lab, NiT is the cleanest published existence proof that the SDXL bucket-sampling + crop-conditioning era can be replaced with NaViT-style native-resolution packing in image diffusion at SOTA quality — i.e. that the training-side cost of variable-shape inputs in DiTs is not fundamental, only an artifact of pre-flash-attention infrastructure. The combination “DC-AE f32c32 latents + FlashAttention-2 cu_seqlens + axial 2D-RoPE + LPFHP packing” is portable to video transformers, where the resize/crop tax compounds across the temporal axis. This complements Diffusion training efficiency by adding a data-side efficiency lever (eliminate resize/crop information loss, reuse compute across resolutions in one model) that’s orthogonal to the REPA / RAE / Self-Flow internal-supervision recipes already tracked there — and ablation §4.3 shows the levers are independent (NiT also uses REPA-style RADIO-v2.5 auxiliary supervision). It contrasts with Foveated Diffusion: Efficient Spatially Adaptive Image and Video Generation, which reduces tokens per image via spatially-adaptive resolution in the periphery; NiT instead reduces wasted tokens across the batch by eliminating padding and per-resolution model duplication. There is also an interesting connection to Don't Look Twice: Faster Video Transformers with Run-Length Tokenization (RLT for video ViTs), which uses the same NaViT example-packing + block-diagonal masking primitive but for content-aware temporal token reduction rather than spatial flexibility.