Recurrent Vision Transformers with depth-programmed experts

In this work we show that a single Transformer block applied recurrently can match the accuracy of a full-depth vision encoder at comparable inference FLOPs without intermediate feature distillation. reViT restores depth-specific transformations by representing the FFN at each recurrent depth as a convex combination of a small shared expert bank. A continuous normalized-depth coordinate programs this mixture, defining a resampleable trajectory through FFN parameter space.

From single block to many depths

The core idea is to replace a deep stacked ViT with a single block that is applied repeatedly. At each recurrence, the FFN (feed-forward network) is not a fixed parameter set but a mixture of a small shared expert bank. A continuous normalized-depth coordinate controls the mixture weights, effectively traversing a trajectory through FFN parameter space that changes with depth. This design restores depth-specific transformations that would normally require distinct FFN parameters at each layer, while keeping the total parameter count low.

Mixture-of-Experts families compared

The authors evaluate several MoE families under a matching one-FFN budget: weight-space merging, token-dispatch, and output-mixture. Across both supervised ImageNet-1k training and distillation from a DINOv2 teacher, weight-space merging emerges as the strongest tested MoE family, outperforming token-dispatch and output-mixture at the same computational budget. This result holds across both training regimes, suggesting that merging FFN weights across recurrent depths is a particularly effective strategy.

Accuracy and parameter efficiency

Trained from scratch, reViT-B/16 attains DeiT III accuracy with about 70% fewer stored parameters. This is a striking efficiency gain: the model achieves the same top-1 accuracy as a established DeiT III model while storing substantially fewer parameters. The reduction comes from sharing expert parameters across depths and using the depth coordinate to select mixtures rather than maintaining separate FFN weights for each layer.

Distillation and transferability

An 8-experts model distilled using only the teacher's output features retains nearly all of its DINOv2 teacher's linear-probe accuracy and transfers across classification, segmentation, and depth prediction. This indicates that the reViT architecture not only matches supervised accuracy but also preserves the rich representation quality of large teacher models, making it suitable for downstream tasks beyond image classification.

Elastic-depth operation

Elastic-depth training allows one checkpoint (trained model) to operate at multiple tested depths by resampling the same normalized coordinate interval. A single trained model can be deployed at different computational budgets simply by varying how many times the recurrent block is applied and how the depth coordinate is sampled. This eliminates the need to train separate models for different FLOP budgets.

Fixed-depth deployment

For fixed-depth deployment, the recurrent block can be materialized as a conventional dense graph, removing online routing and merging without changing the one-FFN-per-depth compute but expanding deployment storage. When online routing is undesirable—for instance, in latency-sensitive deployments—the model can be unrolled into a static computation graph, at the cost of increased storage for the expanded FFN parameters.

Conclusion

reViT demonstrates that a single Transformer block applied recurrently, with FFN represented as a depth-programmed mixture of shared experts, can match the accuracy of a full-depth vision encoder at comparable inference FLOPs. Weight-space merging emerges as the strongest MoE family at a matching one-FFN budget. Trained from scratch, reViT-B/16 attains DeiT III accuracy with about 70% fewer stored parameters. An 8-experts model distilled using only the teacher's output features retains nearly all of its DINOv2 teacher's linear-probe accuracy across classification, segmentation, and depth prediction. Elastic-depth training allows one checkpoint to operate at multiple depths, and for fixed-depth deployment the recurrent block can be materialized as a conventional dense graph. This work opens a pathway to more efficient ViT deployment without distillation or architectural overhead. Read the paper on arXiv