One Block, Multiple Depths: Recurrent Vision Transformers with Depth-Programmed Experts

One Block, Multiple Depths. Recurrent Vision Transformers with Depth-Programmed Experts. By Adrian Bulat and two others, from Samsung AI Cambridge, the Technical University of Iasi, and one other institution. Published on the eighth of October.

A vision encoder normally stacks many distinct Transformer blocks, each with its own weights. Reusing one block over and over would store far fewer parameters, but then every pass through the block applies the same transformation, and the depth-specific behaviour of a real deep network is gone. The question is whether a single recurrent block can still match a full-depth encoder, at comparable inference FLOPs, and without leaning on intermediate feature distillation. The authors show that it can.

reViT restores depth-specific transformations inside the feed-forward network. At each recurrent depth, the FFN is a convex combination of a small, shared bank of experts. A continuous coordinate, normalized depth, programs that mixture. The block is shared, but the weights it applies shift as the coordinate moves. The result is a trajectory through FFN parameter space, and because the coordinate is continuous, the trajectory can be resampled.

The design is evaluated in two regimes: supervised training on ImageNet-1k, and distillation from a DINOv2 teacher. In both, controlled adaptations compare ways of mixing experts at a matching budget of one FFN. Weight-space merging comes out as the strongest family tested, ahead of token dispatch and ahead of mixing outputs.

Two deployment properties follow. With elastic-depth training, one checkpoint can run at several tested depths, simply by resampling the same normalized coordinate interval. And for deployment at a fixed depth, the recurrent block can be materialized as an ordinary dense graph. That removes online routing and merging, keeps the compute at one FFN per depth, and costs more storage.

Trained from scratch, reViT-B/16 reaches DeiT III accuracy with about seventy percent fewer stored parameters.

In the distillation regime, an eight-expert model trained only on the teacher's output features keeps nearly all of its DINOv2 teacher's linear-probe accuracy. It transfers across classification, segmentation, and depth prediction.

Both results are reached without intermediate feature distillation. The supervised model learns from labels alone, and the distilled one sees only what the teacher outputs at the end, not the features from its inner layers. The depth-specific behaviour has to come from the expert mixture itself.

Depth here stops being a count of separate layers and becomes a coordinate. One block plus a small expert bank, steered along that coordinate, stands in for a full stack, and the same weights can be walked in more steps or fewer depending on what deployment needs.