ViT-5: Vision Transformers for The Mid-2020s

TL;DR

ViT-5 is a component-wise modernization of the plain Vision Transformer that ports the architectural refinements LLMs accumulated over the last five years — LayerScale, RMSNorm, 2D RoPE (jointly with APE), register tokens, QK-Norm, bias-free QKV — back into the canonical Attention–FFN ViT stack. The non-obvious negative result is that SwiGLU should not be ported: combined with LayerScale it over-gates and hurts ViT performance, so ViT-5 keeps GeLU. At ViT-L/384², ViT-5 reaches 86.0% ImageNet-1k top-1 (vs DeiT-III’s 85.4% and ConvNeXt-L’s 85.5%); as a SiT diffusion backbone at XL it hits 1.84 FID vs SiT-ViT-XL’s 2.06 at matched compute. The paper is best read as a drop-in upgrade recipe for any pretraining or generation pipeline currently using a vanilla ViT.

Key claims

LayerScale and post-RMSNorm act through the same mechanism (both rescale block output relative to residual); empirically they yield matched accuracy across S/B/L scales, so ViT-5 picks LayerScale for its lower compute overhead [§3.1, Table 1].
Replacing all LayerNorm with RMSNorm gives ≈+0.2% top-1 on ViT-B with slightly lower FLOPs — the re-centering step in LayerNorm is unnecessary in ViTs as well as in LLMs [§3.2].
Combining LayerScale with SwiGLU MLP degrades ViT performance because both act as channel-wise filters, causing excessive activation sparsity (“over-gating”); ablation: LayerScale+SwiGLU = 83.70% / FID 16.22 vs LayerScale+GeLU = 84.16% / FID 14.57 at ViT-B scale [§3.3, Table 2].
ViT-5 uses both absolute positional embeddings and 2D RoPE; RoPE-only is rejected because it makes patch-level image flips fully invariant, which is fine for classification but harmful for spatial reasoning [§3.4, Fig. 2].
2D RoPE delivers much stronger dynamic-resolution robustness than DeiT-III’s APE-only baseline — trained at 224², ViT-5 maintains accuracy across 128–512 test resolutions without fine-tuning, while DeiT-III degrades sharply away from 224 [§3.4, Fig. 3].
Register tokens must also be assigned RoPE (with a higher frequency base than patch tokens); leaving registers un-rotated when patches are rotated introduces an undesired positional bias via cosine-similarity asymmetry [§3.5, Table 3].
QK-Normalization yields a modest accuracy bump but a large reduction in sharp loss spikes during optimization [§3.6, Fig. 5]; ViT-5 additionally removes bias terms from QKV projections to let QK-Norm operate cleanly [§3.7].
On ImageNet-1k, ViT-5-B reaches 84.2% top-1 vs DeiT-III-B’s 83.8% at matched ≈87M params / 17.9G FLOPs; ViT-5-L/384² reaches 86.0% vs DeiT-III-L’s 85.4% and ConvNeXt-L’s 85.5% [Table 5].
As a SiT diffusion backbone at ImageNet-256, ViT-5-XL achieves FID 1.84 vs SiT-ViT-XL’s 2.06 at 7M steps under identical training configurations, with consistent improvements across S/B/L/XL and smooth scaling curves [Table 7, Fig. 6].
Direct transfer of LLaMA/Qwen/Gemma3-style configurations to ViTs underperforms ViT-5 by 0.25–0.37% at ViT-L/384² — language-model design choices do not trivially translate to optimal vision performance, motivating a vision-specific component search rather than blind porting [§4.4, Table 9].
On ADE20k semantic segmentation with UperNet, the ViT-5 advantage widens with scale: +2.3 mIoU at S, +1.1 at B, +2.7 at L (52.0% vs DeiT-III-L’s 49.3%) [§4.3, Table 8].

Method

ViT-5 preserves the canonical plain-ViT Attention–FFN topology and changes only the constituent components. The final recipe is: LayerScale on each block output (initialized ≈10⁻⁴), RMSNorm replacing all LayerNorms (bias-free), GeLU MLP (not SwiGLU — explicit negative result), joint APE + 2D RoPE positional encoding, 4 register tokens with their own 2D RoPE at a higher frequency base than patches, QK-Normalization on Q and K before attention, and bias-free QKV projections. Configurations: ViT-5-S (12L/d384/h6, 22M), ViT-5-B (12L/d768/h12, 87M), ViT-5-L (24L/d1024/h16, 304M), ViT-5-XL (28L/d1152/h16, 449M).

The classification training recipe follows DeiT-III with minor modifications; for generation, ViT-5 is plugged in as the backbone of a SiT diffusion transformer under SiT’s exact training configuration. The two-part ablation strategy — (1) compare against existing prevalent configurations drawn from prior ViT/LLM models, (2) leave-one-out ablation of each of the seven ViT-5 components — isolates that no single change dominates, but the set is mutually reinforcing, and that the LayerScale × SwiGLU interaction is the only meaningful negative pair.

Results

ImageNet-1k classification (224²): ViT-5-S 82.2% vs DeiT-III-S 81.4%; ViT-5-B 84.2% vs DeiT-III-B 83.8%; ViT-5-L 84.9% vs DeiT-III-L 84.5%. At 384² resolution, ViT-5-L reaches 86.0% — a new plain-ViT SoTA exceeding both DeiT-III-L (85.4%) and ConvNeXt-L (85.5%) [Table 5]. ImageNet-256 generation as a SiT backbone at 7M steps: ViT-5-XL 1.84 FID / 282.7 IS vs SiT-ViT-XL 2.06 / 277.5 vs DiT-ViT-XL 2.27 / 278.2 [Table 7]; the FID gap holds across all four scales and remains stable as training extends. ADE20k semantic segmentation with UperNet at 512²: ViT-5-S 47.5 mIoU (DeiT-III-S 45.2), ViT-5-B 49.1 (DeiT-III-B 48.0), ViT-5-L 52.0 (DeiT-III-L 49.3) — the gain widens with scale, suggesting the modernization is more valuable for larger models [§4.3, Table 8]. Cross-architecture comparison at ViT-L/384²: ViT-5 (86.00%) outperforms every prevalent vision (DeiT-III, DINO v2/v3, VisionLLaMA, NEPA) or language (LLaMA, Qwen, GPT-oss, Gemma3) configuration by 0.25–1.15% on ImageNet [Table 9].

Why it’s interesting

For Luma, ViT-5 is most relevant as a drop-in upgrade for the ViT backbones used in diffusion / generation stacks — the 1.84 FID at SiT-XL is achieved purely by swapping the backbone, with no change to the SiT objective, data, or schedule, and the scaling curves stay smooth. The negative result that LayerScale + SwiGLU over-gates in vision is the kind of surprise that would otherwise be discovered the expensive way during a from-scratch DiT pretraining run. The paper also adds a vision-specific data point to the Training stability at scale thread: QK-Norm shows up here as the load-bearing stability fix for ViTs (loss spikes disappear, training stabilizes) and not as a substitute for µP-style parametrization, mirroring the How to Set the Learning Rate for Large-Scale Pre-training? finding for LLMs. The 2D-RoPE-on-registers detail (registers need a separate, higher-frequency-base RoPE to avoid cosine-similarity-induced positional bias) is a non-obvious implementation lesson for anyone adopting both Darcet-style registers and RoPE simultaneously.