Source
Dehghani, Mustafa et al. (2023), LIT-657 — ARXIV-2307.06304.
The recipe
- Keep aspect ratio. Resize each image to a target area without squashing or cropping it square.
- Sample the area per image. NaViT samples side length r ~ U(64, 256) and resizes to r² pixels (App. A.1). Sampling side length beats sampling area, and a distribution biased toward small sizes does best (Fig. 7).
- Pack. Put several images’ patches into one fixed-length sequence. Mask attention and pooling per example, and mask padded examples out of the loss. Greedy packing leaves under 2% padding.
- Use factorized x/y position embeddings. They extrapolate to unseen resolutions better than learned 1D or 2D tables (Fig. 10).
What the evidence supports, part by part
| Part | Evidence | Isolated? |
|---|---|---|
| The whole package | Leads ViT at each of 12 compute-matched budgets. Matches the best ViT with 4x less compute (Fig. 1) | No: aspect ratio, resolution, token dropping and images seen all change together |
| Variable resolution | U(64, R_max) matches or beats fixed R_max at equal FLOPs, with aspect ratio preserved in both (Fig. 5) | Yes |
| Aspect ratio | Native beats square resize at equal area in a linear fairness probe, p = 0.02 (Fig. 12 right) | Only at probe time, not in pretraining |
| Throughput | About 5x more images seen at the same compute (Table 2) | The paper names this “the chief contributor” |
The recipe and every row of the table are NaViT’s (LIT-657). On JFT-4B it runs 12 compute-matched ViT budgets against the packed, aspect-preserved, resolution-sampled alternative it proposes, and attributes the gain chiefly to seeing about 5x more images, not to aspect ratio.
So the recommendation holds as a package. Its title names aspect ratio because that is what the package lets you keep, but aspect ratio is not shown to be what makes it win.
Conditions
- ViT encoders trained for classification or contrastive learning. JFT-4B and WebLI, B/32 to L/16, one lab. Sora (LIT-652) and SDXL (LIT-566, crop conditioning) argue for the generative case with example images only. The one measured generative case is FiT (LIT-658, Table 3). A DiT-B trained on uncropped, aspect-preserved images scores FID 43.34 against 44.83 for square crops at a 256² output, with a worse sFID. At non-square outputs it wins by about 40–57 FID, but the square model never trained on those shapes. For generation, preserving aspect ratio is what makes non-square output possible. It is not shown to improve square output.
- Evaluation has to match. Much of the out-of-distribution gain appears when the square baseline is evaluated by square resize (Table 5). With an aspect-preserving crop, ObjectNet is roughly tied.
- Packing costs attention. The overhead from attending over longer packed sequences shrinks as the model gets wider (Fig. 4). Per-example masks also require an attention kernel that supports them.
- Example-level losses need care. Pooled, contrastive losses must handle a variable number of examples per sequence. NaViT uses a chunked contrastive loss for this (§2.3).
Relation to the record
This is an alternative to SOTA-251, which holds resolution low and raises it only in the decay phase. NaViT mixes resolutions throughout training. It argues that a staged increase cannot be reversed, while a mixed-trained model can be run at any resolution afterwards. The record holds no comparison of the two.
SOTA-152 recommends best-fit document packing for language models. This practice applies the same idea to images, with per-example masks that SOTA-152’s setting does not need in the same form.