Adrian Bulat, Yassine Ouali, Georgios Tzimiropoulos
9 min
Abstract
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. We evaluate this design in two regimes: supervised ImageNet-1k training and distillation from a DINOv2 teacher. Across both regimes, controlled adaptations identify weight-space merging as the strongest tested MoE family at a matching one-FFN budget, ahead of the token-dispatch and output-mixture alternatives. 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 and transfers across classification, segmentation, and depth prediction. Elastic-depth training allows one checkpoint (trained model) to operate at multiple tested depths by resampling the same normalized coordinate interval. 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.
Sam: That means the decision does not depend on the image, or on which patch is being processed?
Alex: Correct: every image and every patch uses the same mixture at a given depth. The position runs from the start to the end of the sequence. The router turns that position into soft weights over the expert bank.
Sam: There is a distinction here between blending networks and running several networks. Which one actually happens?
Alex: They blend the parameters first, including both projections and their biases. Then they execute only the resulting feed-forward network. They do not run every expert and average its output, so each step still executes one dense feed-forward network.
Sam: Give me a concrete example of what this buys, rather than just the mechanism.
Alex: On ImageNet image classification, the base-sized reViT reaches eighty-three percent accuracy. Its conventional DeiT III baseline reaches eighty-two point eight percent. Those are near-matched computation comparisons under the same training recipe, with the large storage reduction we mentioned.
Sam: That is a close result. How much repetition is behind it, and does it hold across model sizes?
Alex: The main reViT results are means over three independently trained runs. The large model also comes close to its conventional baseline, but the small model has a larger gap with the default expert bank. Adding more experts narrows that gap.
Sam: So I would read this as evidence that depth-specific weights recover much of what plain sharing loses. Not as evidence that sharing always wins.
Alex: That is the supported reading. The largest accuracy gain comes when moving from a single shared feed-forward network to a small expert bank. Beyond that, extra stored capacity helps the smaller backbone more than the larger one.
Sam: Classification is one test. Does the shared model preserve features useful for other tasks?
Alex: They also use distillation: training a student to match a frozen teacher’s representations. The teacher is DINOv2, a pretrained vision encoder. The student matches its final token features, without targets from intermediate layers.
Sam: If only the final features are matched, how do they check whether the representation transfers?
Alex: They freeze the student and train lightweight prediction heads. These test classification, image segmentation, and depth prediction. Segmentation labels image regions; depth prediction estimates their distance.
Sam: Does the strongest student actually match the teacher across those tasks?
Alex: With eight experts, it comes within about six-tenths of a percentage point of the teacher’s classification result. It remains below the teacher on segmentation, while depth results depend on the prediction head. So “nearly matches” describes classification, not every measurement.
Sam: And the Raptor comparison uses different teacher supervision. I would hesitate to attribute that entire difference to the architecture.
Alex: The authors make that distinction too. Those distilled comparisons are references, not controlled ablations. Their cleaner mechanism test puts alternative expert formulations inside the same recurrent backbone, with the same task objective and base recipe.
Sam: What does that cleaner test say about the decision to merge weights?
Alex: At the budget of one feed-forward network per step, weight merging is the strongest tested family overall. ReViT beats the tested token-dispatch approaches, which send patches to selected experts, and output-mixture approaches, which combine expert results. But another merging method is close under distillation and better on one depth measurement.
Sam: Does the advantage survive if the other approaches get more active computation?
Alex: Some become competitive, and an output-mixture method leads the distilled tasks at the higher tested budget. Those configurations also change expert widths or activation patterns. This ranks adaptations under full recurrence, not the original methods in every architecture.
Sam: Let’s return to flexible depth. Why should a model trained for one processing length tolerate another?
Alex: Because its router reads relative position, not a fixed step identity. During elastic-depth training, complete recurrent steps are randomly dropped for each image. The retained steps are renumbered and spread across the same start-to-finish interval.
Sam: So choosing a longer sequence samples the learned weight progression more densely. Did extra steps actually help?
Alex: In the tested depth-conditioned model, all three probes improve through sixteen recurrent steps. Going deeper leaves performance near its best. A feature-conditioned control largely plateaus beyond its reference depth and worsens on depth prediction.
Sam: Is there evidence that the experts learn different stages, rather than an arbitrary mixture?
Alex: Dominant experts occupy contiguous depth intervals, with soft blending around transitions. No training target specifies those boundaries. Reassigning the learned mixtures to different experts reduces teacher alignment in every tested permutation, so the pairing matters.
Sam: The remaining practical question is speed. Does the compact model deliver the same computation more quickly?
Alex: Not in the measured dynamic implementation. On a single-image batch, it has thirteen percent higher latency, despite lower runtime memory. The implementation reconstructs merged weights with generic operations; the authors explicitly distinguish that from optimized deployment.
Sam: Since the mixtures are image-independent, could a deployment avoid reconstructing them?
Alex: Yes. At a fixed depth, the merged weights can be precomputed and exported into a conventional dense graph. That removes online routing and merging, but stores a feed-forward network for each depth. Compact storage and conventional execution are different deployment choices.
Sam: Given those trade-offs, who should read the full document, and where should they start?
Alex: Researchers working on vision compression, recurrent models, or expert mixtures should read it. Start with the method and the common recurrent comparison of expert formulations. Then read elastic-depth inference and the measured inference latency appendix before drawing deployment conclusions.
Sam: And what should everyone else remember?
Alex: Sharing a block does not have to mean using identical weights at every step. ReViT saves storage by learning how shared weights should change across depth—not by eliminating the work of depth.
Sam: Keep the storage saving separate from the speed claim.