[KongchangAI]
· 3 min read· 1,653 words

New Breakthrough in Recurrent Vision Transformers: A Single Reused Block Approaches Full-Depth Encoder Performance

New Breakthrough in Recurrent Vision Transformers: A Single Reused Block Approaches Full-Depth Encoder Performance

reViT uses a single recurrent block plus a depth-programmed expert bank to reach full-depth ViT accuracy at equal inference FLOPs with ~70% fewer stored parameters.

reViT challenges the assumption that vision transformers must stack different parameters across depth. It recursively applies one Transformer block and expresses each depth's FFN as a convex combination of a shared expert bank. Its key innovation is a continuous normalized depth coordinate that programs expert mixing weights, creating a resamplable parameter-space trajectory and enabling elastic-depth training. Controlled experiments show weight-space merging beats token-dispatch and output-mixture under a single-FFN budget. reViT-B/16 matches DeiT III accuracy with ~70% fewer stored parameters from scratch supervised training; an 8-expert distilled model using only DINOv2 teacher output features retains near-full linear probe accuracy and transfers across tasks. At deployment, recursive blocks can route and merge online or be materialized as dense graphs, trading storage for removing routing overhead.

Title

The Efficient Vision Transformer Path Is Being Rethought

The efficiency path for vision transformers is being reconsidered. A study named reViT (Recurrent Vision Transformer) proposes that recursively applying a single Transformer block can achieve the accuracy of a full-depth vision encoder at the same inference FLOPs, and without intermediate feature distillation—a design that challenges the conventional assumption that "depth must stack different parameters."

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

Core Problem: What Recursive Shared Parameters Lose

The intuitive advantage of a recursive structure is straightforward: one block is reused across multiple depths, greatly reducing stored parameters. But naive approaches face a fundamental obstacle—different depths require different transformations. Shallow layers capture textures and edges, while deeper layers perform semantic aggregation. If each recursive step executes exactly the same feed-forward network (FFN), it becomes difficult for the model to express this depth-specialized hierarchy.

reViT's solution is not to assign an independent set of parameters to each layer, but to introduce a shared expert bank, so that the FFN at each recursive depth becomes a convex combination of this expert bank. In other words, what stays constant across depth is the parameter pool; what changes across depth is the combination weights.

Mixture of Experts (MoE) is an architectural paradigm that expands model capacity by combining multiple sub-networks (experts) while controlling per-inference computation. Each expert is typically an independent FFN, and a gating network dynamically decides the activation or combination weights of each expert based on the input. Traditional MoE often dispatches at the token granularity (token-dispatch), meaning different tokens go to different experts; there are also variants that weight at the output end (output-mixture). The weight-space merging adopted by reViT instead combines multiple experts' weights linearly according to mixing coefficients at the parameter level into an equivalent FFN, and then applies it uniformly to all tokens. The advantage of this approach is that the merged weights are a single dense matrix, and the forward computation is equivalent to an ordinary FFN, preserving the expressive diversity of multiple experts while avoiding the irregular computation and memory overhead caused by token dispatch. When matching a single-FFN computation budget, this "merge first, then compute" path is often more efficient than "compute separately, then merge outputs," because it moves the mixing complexity from activation space to weight space.

"Programming" the FFN with Continuous Depth Coordinates

reViT's most ingenious design lies in how it represents depth. It defines a continuous normalized depth coordinate used to program the mixing weights over the expert bank. This coordinate is not a discrete layer index, but a position on a continuous interval, naturally producing a resamplable trajectory in FFN parameter space.

This continuous representation brings two benefits. The first is elastic-depth training: the same checkpoint can run at multiple test depths by resampling the same normalized coordinate interval, without training a separate model for each depth. The second is that it leaves room for subsequent deployment optimization—when deploying at a fixed depth, the recursive block can be materialized as a conventional dense graph.

The research team conducted controlled comparisons of the mixing mechanism implementation, comparing three types of MoE families: token-dispatch, output-mixture, and weight-space merging. The conclusion is clear—under the condition of matching a single FFN budget, weight-space merging is the strongest approach, ahead of the other two alternatives.

Convex combination refers to a linear weighted sum in which the coefficients are non-negative and sum to 1. Using a convex combination to mix expert weights means that the FFN at each recursive depth always lies within the convex hull spanned by the expert bank weights. This constraint brings two practical benefits: first, the mixing result will not become numerically unstable due to coefficient extrapolation, and the weight scale is confined within the range formed by the expert weights; second, because the weights are bounded, activation explosion or gradient vanishing is less likely even after multiple recursive iterations. By contrast, if arbitrary linear combinations were allowed (coefficients could be positive or negative, or not sum to 1), deep recursion could theoretically amplify the weight norm, leading to unstable training. The continuous normalized depth coordinate combined with convex combination enables reViT to interpolate a reasonable FFN at any depth, which is precisely the mathematical basis for elastic-depth training and resampling.

Validation Results Under Two Training Paradigms

reViT was evaluated under two distinctly different training mechanisms to test the generality of its design.

Supervised Learning: Training from Scratch on ImageNet-1k

In the ImageNet-1k supervised training setting, reViT-B/16 trained from scratch achieves accuracy comparable to DeiT III, while reducing stored parameters by about 70%. This number is noteworthy: the reduction in parameter count does not come at the cost of accuracy, but is achieved at the same inference FLOPs.

Distillation: Transferring from a DINOv2 Teacher

In the distillation paradigm, reViT uses an 8-expert model and distills only from the teacher's output features, preserving nearly the full linear probe accuracy of the DINOv2 teacher. More importantly, this capability transfers to multiple downstream tasks such as classification, segmentation, and depth prediction—indicating that what the model learns is not task-specific features, but generalizable visual representations.

The distillation setting is especially convincing: using only output features, without intermediate-layer features, it still achieves such transfer performance. This echoes the opening claim of "no intermediate feature distillation."

DINOv2 is a self-supervised visual pretraining method proposed by Meta. It trains ViT on large-scale unlabeled images through self-distillation, and the learned features have strong linear separability on many downstream tasks, the so-called linear probe accuracy. A linear probe refers to the standard protocol of freezing the backbone and training only a linear classification head to evaluate feature quality, which can exclude the interference of fine-tuning on task-specific capability. In the distillation experiment, reViT uses DINOv2 as the teacher and aligns only its output features without using intermediate-layer features, which is equivalent to letting the student model directly learn the teacher's final representation manifold. This approach is more lightweight than intermediate-layer distillation, but places higher demands on the student model's capacity and optimization path. That reViT can approach full linear probe accuracy with 8 experts and transfer across classification, segmentation, and depth prediction shows that the recurrent shared-parameter structure does not become a bottleneck for representation learning; instead, it retains sufficient expressive flexibility through depth programming.

Deployment Perspective: Merge or Expand

reViT's design offers two forms at deployment time, each with trade-offs.

At inference time, the recursive block implies online routing and merging. Since the FFN at each depth is a dynamic combination of the expert bank, the inference process needs to compute the mixing weights in real time.

If the simplicity of fixed-depth deployment is desired, the recursive block can be materialized into a conventional dense graph. This removes online routing and merging operations, and the computation per depth remains unchanged for one FFN; the cost is that deployment storage expands—essentially unfolding the compact representation during training into a dense form at inference.

This is a typical trade-off between "training storage vs. inference efficiency." reViT does not avoid this contradiction, but explicitly places it in the design space: stay compact during training and elastic use, and accept expansion cost for fixed deployment.

Inference FLOPs is one of the core metrics for measuring model deployment cost. It refers to the number of floating-point operations required for a single forward pass, directly related to latency and energy consumption; stored parameter count determines the GPU memory or disk size occupied by model weights. The two do not always move in the same direction: MoE structures can trade a larger parameter count for a smaller per-activation computation, while parameter-sharing structures do the opposite, carrying a larger equivalent depth computation with a smaller parameter count. reViT's recursive design sits exactly in the sensitive region of this trade-off—parameters are compact during training but the equivalent depth computation is large; if expanded into a dense graph at deployment, the parameter count expands but routing overhead disappears. When evaluating such architectures, one must clearly distinguish among three metrics: "total parameter count," "activated parameter count," and "inference FLOPs"; otherwise it is easy to reach misleading conclusions. The comparison emphasized in the reViT report—"achieving full-depth accuracy at the same inference FLOPs with about 70% fewer stored parameters"—anchors both dimensions simultaneously.

Significance of This Work

The value of reViT does not lie in proposing an isolated trick, but in integrating recursive depth, mixture of experts, and continuous depth programming into a self-consistent framework. It proves that recurrent vision encoders need not compromise on accuracy, and provides reproducible evidence that weight-space merging outperforms other MoE approaches under a single-FFN budget. For researchers and engineers focused on model compression, parameter-efficient vision backbones, and elastic inference, this route of "one block for multiple uses, programmable depth" offers a direction worth following.

One-line Summary

reViT uses a single recurrent block plus a depth-programmed expert bank to achieve full-depth ViT accuracy at the same inference FLOPs, with about 70% fewer stored parameters.

Paragraph Summary

reViT (Recurrent Vision Transformer) challenges the conventional assumption that vision transformers must stack different parameter depths. It proposes repeatedly applying a single Transformer block recursively and expressing the FFN at each recursive depth through a convex combination of a shared expert bank. Its core innovation is a continuous normalized depth coordinate, which programs expert mixing weights using a continuous interval's depth position, forming a resamplable trajectory in parameter space and thereby supporting elastic-depth training. Controlled experiments show that, under a matched single-FFN budget, weight-space merging outperforms token-dispatch and output-mixture. reViT-B/16 trained from scratch under supervision reaches DeiT III accuracy with about 70% fewer stored parameters; an 8-expert distilled model using only DINOv2 teacher output features preserves nearly full linear probe accuracy and transfers across tasks. At deployment, the recursive block can route and merge in real time, or be materialized as a dense graph to eliminate routing overhead at the cost of storage expansion.

#reViT#Recurrent Vision Transformer#vision transformer efficiency#mixture of experts#weight-space merging#elastic-depth training#DINOv2 distillation#parameter-efficient vision backbone
Share:

Related articles