Three things everyone should know about Vision Transformers

TL;DR

Three modular ViT changes that are each simple to drop into a vanilla ViT-B/L. (1) Reorganize residual blocks into parallel MHSA+FFN pairs — same parameters and FLOPs but shallower depth, easier optimization for large models, and up to ~45% latency reduction at batch size 1–4 on V100. (2) When adapting a pretrained ViT to a new resolution or downstream classification task, fine-tune only the MHSA weights — statistically indistinguishable from full fine-tuning at higher resolution, saves ~10% memory and time, and lets the same FFN weights (2/3 of parameters) be shared across tasks. (3) A hierarchical MLP (hMLP) stem — small MLP blocks interleaved with patch-aggregation that go 2×2 → 4×4 → 8×8 → 16×16 with no cross-patch communication — matches the best convolutional stems in supervised training and is compatible with BeiT-style masked self-supervised learning because masking commutes with the stem.

Key claims

Parallel blocks $x = x + \text{mhsa}_1(x) + \text{mhsa}_2(x); x = x + \text{ffn}_1(x) + \text{ffn}_2(x)$ preserve parameter count and FLOPs vs a strictly sequential ViT with the same total MHSA/FFN count, but reduce depth by half [§3.2].
Parallel two-branch ViT-B18×2 matches sequential ViT-B36×1 on ImageNet-val (83.9 vs 83.9 at 400 ep + LayerScale) and helps more as models get deeper and harder to optimize; three or more parallel branches typically hurt accuracy [Fig. 4, Table 4].
On V100 GPUs at batch 1 the parallel ViT-B18×2 is 45% faster than its sequential counterpart (61 vs 42 im/s); the speedup shrinks with batch size and disappears by batch 16, and reaching full potential requires custom CUDA kernels [Table 4].
Fine-tuning only the MHSA weights of a pretrained ViT-B / ViT-L to 384×384 matches full fine-tuning on ImageNet-val and ImageNet-V2 (within one standard deviation) while fine-tuning FFN alone is measurably worse (e.g. ViT-S 82.5 attn vs 82.2 ffn vs 82.7 full) [Table 5].
Attention-only fine-tuning saves ~10% training memory and ~10% wall-clock, and lets a base model + N task-specific MHSA weight-sets share the ~2/3 of parameters that live in FFNs across tasks [§4].
Attention-only fine-tuning is a strong option for small downstream classification datasets (CARS, Flowers — even beats full fine-tuning) but has a measurable gap on large downstream sets (iNat-18) with small backbones (ViT-S); the gap closes for ViT-L [Table 6].
A hierarchical MLP stem that processes each 16×16 patch independently through interleaved 2×2 / 4×4 / 8×8 / 16×16 MLP-plus-patchify blocks matches the best convolutional stems in supervised ImageNet training and is fully compatible with BeiT masking — because there is no cross-patch communication in the stem, masking commutes with pre-processing [§5, Fig. 5].
Existing convolutional stems break BeiT: the convolutions leak information between neighbouring patches, so masked-patch predictions get a shortcut and reconstruction quality drops; hMLP avoids this by construction [§5].

Method

Contribution 1 — Parallel ViT: instead of $2N$ layers each containing one MHSA followed by one FFN residual block, group them into $N$ layers each containing two MHSA blocks in parallel followed by two FFN blocks in parallel. Same parameter count and FLOPs, but width in parallel doubles and effective depth halves. This is combined with LayerScale for stability on ViT-L. Contribution 2 — Attention-only fine-tuning: freeze the FFN and LayerNorm and only update MHSA weights (Q, K, V, and output projection) during transfer or resolution adaptation, evaluated with AdamW and SGD. Contribution 3 — hMLP stem: process each 16×16 patch as four 2×2 sub-patches; apply MLP-in-patch, then patchify (concatenate spatially, project), yielding 4×4 sub-patches; MLP; patchify to 8×8; MLP; patchify to 16×16; MLP. The stem is implemented as a stack of stride-matched convolutions to remain patch-local. Combined with BeiT, masking can be applied either before or after the stem and yield equivalent results because the stem never mixes information across patches.

Results

Parallel vs sequential ImageNet-1k top-1: ViT-B18×2 83.9 vs ViT-B36×1 83.9 (400 epochs + LayerScale), with V100 batch-1 latency 61 vs 42 im/s (45% faster) [Table 4].
Attention-only fine-tuning at 384² (AdamW): ViT-S 82.5 (attn) vs 82.7 (full), ViT-B 84.3 vs 84.3, ViT-L 85.5 vs 85.5. ImageNet-V2: ViT-S 72.4 vs 72.5, ViT-B 74.0 vs 73.7, ViT-L 75.4 vs 75.5 [Table 5].
Transfer to smaller datasets with ViT-B attn-only fine-tuning: CARS 92.8 (attn) vs 92.7 (full), Flowers 98.5 vs 97.8, CIFAR-10 99.2 vs 99.3, CIFAR-100 91.8 vs 92.5; iNat-18 71.1 (attn) vs 74.1 (full) is the clearest gap, closing at ViT-L (75.3 vs 75.9) [Table 6].
hMLP + BeiT self-supervised pretraining: hMLP stem is compatible with BeiT masking and matches the best convolutional stem in the supervised regime; convolutional stems break BeiT because information leaks across patches during masked-patch reconstruction [§5, Fig. 5].

Why it’s interesting

The hMLP stem is the vision-side counterpart to dMel (dMel: Speech Tokenization made Simple) in the encoder-free early-fusion recipe used by Interaction Models: A Scalable Approach to Human-AI Collaboration: images become 40×40-pixel patches encoded by a light hMLP embedding, no ViT encoder, no VAE — the transformer backbone processes them directly alongside dMel audio and text tokens. That’s the same “pretrained-encoder replaced by a train-free / lightweight patch primitive” move that NEO-unify: Building Native Multimodal Unified Models End to End (NEO-unify) and Tuna-2: Pixel Embeddings Beat Vision Encoders for Multimodal Understanding and Generation make in the pure vision UMM setting. Attention-only fine-tuning is the earliest incarnation of the “adapt attention, freeze FFN” pattern that later parameter-efficient methods (LoRA, IA³) systematized — it’s the maximally-simple upper bound of PEFT because it changes real weights but only the ones in attention, sharing FFN as a hard prior. The parallel-ViT trick is a shallow-wide contrast to the loop-transformer direction filed in Looped Transformers and ELT: Elastic Looped Transformers for Visual Generation — both split the “same FLOPs, different depth/width tradeoff” axis in opposite directions, and the ViT-B18×2 result is a useful upper bound on how much can be gained by widening without changing parameter count.