One Transformer block runs the whole vision encoder: depth-programmed experts match DeiT III with 70% fewer stored parameters
Synopsis
A vision Transformer's depth-wise stack is compressed into a single recurrent Transformer block, where the FFN at each recurrent depth is a convex combination of a small shared expert bank programmed by a normalized depth coordinate, so shared weights still yield depth-dependent computation. Under supervised ImageNet-1k training, reViT-B/16 reaches DeiT III accuracy with about 70% fewer stored parameters; under distillation that matches only the DINOv2 teacher's final-layer features, an 8-expert model approaches the teacher on frozen classification, segmentation, and depth probes. One checkpoint also runs at multiple tested depths by resampling the depth coordinate.
Interpretation
A vision encoder no longer needs a stack of independently parameterized blocks; one Transformer block reused at every recurrent depth carries the whole depth computation. Vision Transformers were built as fixed-depth stacks of separate blocks, and recurrent sharing schemes either kept several distinct blocks or relied on intermediate teacher features. In supervised ImageNet-1k training, reViT-B/16 reaches 83.0% top-1 against DeiT III's 82.8%, and reViT-L/16 reaches 83.8% against 84.1%, with roughly a quarter of the stored parameters; the four-block Raptor trails reViT at the S, B, and L scales.
The FFN at each recurrent depth is not learned independently but merged from a small expert bank into one dense FFN by the depth coordinate. The standard approach learns a separate set of FFN weights per depth, or dispatches tokens to different experts. A two-layer MLP router maps the normalized depth coordinate to expert weights, which then merge the expert projections and biases, and only the merged FFN executes at each depth; at the same one-FFN compute budget this weight-space merge outperforms the tested token-dispatch and output-mixture alternatives.
One trained checkpoint can run at multiple inference depths, because changing depth only resamples the same normalized coordinate interval. Conventional models fix depth during training, so changing depth normally requires retraining or a separate set of weights. Under distillation, the depth-only gate variant improves ADE20k to 43.6 mIoU, reaches 83.1% ImageNet linear-probe top-1, and gives 0.555 NYUv2 RMSE, while a feature-only gate is nearly flat on ImageNet and ADE20k at greater depth and worsens on NYUv2.
Perspective
The report covers ImageNet-1k supervised training and DINOv2 distillation, with probes on image classification, ADE20k segmentation, and NYUv2 depth prediction, at the S/16, B/16, L/16, and B/14 scales. It suits researchers and engineers working on vision backbone compression and recurrent weight sharing. The reported latency measurements come from a direct PyTorch implementation without custom kernels or cross-input caching, which the authors describe as a property of the current implementation rather than an optimized deployment result.
In the distillation comparison, reViT matches only the teacher's final-layer features while Raptor uses intermediate features, so the authors describe that table as a reference comparison rather than a controlled ablation. Expert deletion analysis shows the alignment drop from removing an expert is usually smaller than a random gate perturbation of matched size, and the earliest-used experts cause the largest losses, so deletion sensitivity alone cannot separate an expert's role from its position in the recurrence. Permuting gate columns while keeping the expert bank fixed lowers teacher alignment. On larger backbones, merged weights concentrate in fewer directions and gains from additional experts shrink.