PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Adrian Bulat, Yassine Ouali, Georgios Tzimiropoulos

arXiv:2610.12448v1cs.CVcs.LG

TL;DR

Vision Transformers use depth-wise stacks, raising the question of whether one recurrent block can retain their depth-dependent computation. reViT programs a shared FFN expert bank with normalized depth; it outperforms tested alternatives at matched one-FFN compute and matches DeiT III with fewer parameters.

  • Problem

    Can a ViT replace its depth-wise stack with one recurrent block while retaining the depth-dependent computation needed for accuracy?

    Source: §1 · p. 1

  • Method

    ReViT recurrently applies one Transformer module and uses normalized depth to mix a shared FFN expert bank into a dense FFN at each step.

    Source: §1 · p. 1, §6 · p. 9

  • Results

    At matched one-FFN compute, reViT outperforms tested token-dispatch and output-mixture alternatives; scratch-trained E=4 reViT-B/16 matches DeiT III at near-matched inference FLOPs with 73% fewer parameters.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The same learned depth trajectory can be sampled at multiple tested depths or materialized as fixed-depth dense FFNs for deployment.

    Source: §1 · p. 2, §6 · p. 10

    In the higher-U comparisons, method-specific adaptations mean the configurations test whether relaxing the one-FFN constraint closes the gap, not a common scaling curve.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

reViT asks whether one recurrent Transformer block can preserve the depth-dependent computation needed for ViT accuracy. It programs a shared FFN expert bank by normalized depth, matching DeiT III with fewer parameters and outperforming tested alternatives at matched compute.

1 Introduction

  • The motivating question is whether one recurrent block can replace a ViT’s depth-wise stack while retaining the depth-dependent computation needed for accuracy.Prior recurrent-sharing results in language leave open which expert formulations preserve accuracy under full-depth sharing in vision.Source: §1 · p. 1
  • A normalized depth coordinate softly mixes a shared FFN expert bank into one dense FFN per recurrent depth, defining a continuous trajectory through parameter space.The same trajectory can be resampled at a different depth; a fixed-depth deployment can materialize its FFNs and remove online routing and merging.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • At matched one-FFN compute, depth-programmed weight merging outperforms tested token-dispatch and output-mixture alternatives.The comparison uses the same recurrent backbone, task loss, and base training recipe.Source: §1 · p. 2
  • 73% fewer parameters lets scratch-trained E=4 reViT-B/16 match DeiT III at near-matched inference FLOPs.An E=8 model nearly matches its DINOv2 teacher’s ImageNet linear-probe accuracy.Source: §1 · p. 2
  • Under final-layer distillation alone, the router learns dominant expert changes and blends across transitions, while larger backbones use fewer effective directions and gain less from added experts.The learned gate-expert assignment matters for teacher alignment, and increasing expert capacity improves coverage of teacher layers.Source: §1 · p. 2

2 Closely Related Work

Prior work studies recurrent sharing, token-routed experts, and conditioned parameter generation, but reViT targets full-depth recurrence with a shared expert bank programmed by depth. Its gate composes FFN parameters by depth rather than choosing computation from token or image features.

2 Closely Related Work

  • Unlike prior vision approaches that retain multiple recurrent blocks or align segments with intermediate teacher features, reViT reuses one block throughout and uses only final-layer teacher features when distilling.Raptor retains several distinct blocks and underperforms the network it approximates.Source: §2 · pp. 2–3
  • Sparse and soft MoE approaches primarily route tokens, whereas reViT shares a depth-derived gate across inputs and tokens and merges expert parameters before execution.The paper compares these mechanisms under full-depth recurrence.Source: §2 · p. 3
  • ReViT uses normalized recurrent depth instead of image-dependent representations to gate mixtures of stored modules, composing complete FFN parameters rather than modulating activations.Its router predicts softmax coefficients over a learned bank of complete FFNs.Source: §2 · p. 3

3 Method

reViT replaces independently parameterized ViT blocks with a recurrent pre-norm module whose shared attention and normalization accompany a depth-programmed FFN mixture. Training supports supervised classification or final-feature distillation, while elastic-depth inference resamples the same program at a selected depth.

3.1 Recurrent block and depth-programmed experts

  • The recurrent module shares attention and normalization across steps and uses a depth-programmed FFN mixture in place of independently parameterized blocks.The image is embedded as tokens with a prepended class token and learned spatial positional embeddings; a final LayerNorm and task head consume the last recurrent representation.Source: §3 · p. 3, §3.1 · pp. 3–4
  • At each depth, a normalized coordinate t/(L−1) maps to expert logits and mixture weights, which merge expert projections and biases into one FFN applied to all tokens.Only the merged FFN executes; the router and expert bank are shared across recurrent steps.Source: §3.1 · p. 3
  • Changing L resamples the same continuous depth program on a new grid, while fixed-depth FFNs can be precomputed for a dense graph without online routing or merging.The compact model stores E expert FFNs; the fixed graph stores L merged FFNs, with one dense FFN evaluated at each depth.Source: §3.1 · p. 4

3.2 Training and elastic-depth inference

  • Supervised training uses the classification objective, while distillation matches final post-LayerNorm teacher tokens without intermediate feature distillation.For E>1, auxiliary losses encourage expert usage balance, control router-logit scale, and promote distinct expert parameters; the tied E=1 baseline omits them.Source: §3.2 · p. 4
  • Elastic-depth inference selects a global L and recomputes the normalized-depth grid for its L recurrent steps.During training, dropped recurrent steps are re-numbered and their coordinates spread from 0 to 1, exposing the router to multiple effective depths.Source: §3.2 · p. 4

4 Experiments

The experiments test reViT under supervised ImageNet-1k training and DINOv2 distillation, comparing recurrent models, expert-composition mechanisms, elastic-depth behavior, and expert-bank sizes. Results assess whether one recurrent block can approach deeper ViTs and how its performance varies across training regimes and inference depths.

4.1 Experimental setup

  • The study tests supervised ImageNet-1k training and frozen-DINOv2 distillation, evaluates frozen-backbone probes, and varies expert-composition mechanisms and stored expert-bank size.For the supervised comparison, reViT, DeiT III, and Raptor use the same data, base recipe, recurrent depth, and near-matched inference FLOPs.Source: §4 · p. 4, §4.1 · p. 5

4.2 Can one recurrent block replace a deep ViT?

  • At near-matched inference FLOPs, reViT-B/16 reaches 83.0% ImageNet-1k top-1 versus 82.8% for DeiT III, while E=4 reViT models score 0.1–0.2 points above four-block Raptor across scales.The E=4 supervised comparison uses the same data, base recipe, recurrent depth, and near-matched inference FLOPs.Source: §4.2 · p. 5, Table 1 · p. 5
  • At E=k ∈{2, 3, 4}, distilled reViT has better reported values than Raptor on all four probes, while the E=8 model is compared only with its DINOv2 teacher.The comparisons are not controlled ablations because the distillation objectives differ: reViT matches final-layer features and Raptor uses intermediate features.Source: §4.2 · p. 5, Table 2 · p. 6

4.3 Which MoE formulation suits a single recurrent block?

  • The MoE comparison adapts mechanisms to a common recurrent ViT testbed, with U=1 matching the nominal computation of one standard dense FFN.Matching U does not match parameter count, total FLOPs, or latency; higher-U configurations can also change expert number, width, or activation pattern.Source: §4.3 · pp. 5–6
  • At U=1, reViT leads the supervised MoE comparison at 78.2% ImageNet-1k top-1 and has the best reported ImageNet, ADE20k, and linear-depth results under distillation.For supervised training, Lory reaches 77.9% and SMEAR 77.7%; under distillation, SMEAR has the lower MLP-head RMSE, 0.485 versus 0.491, and remains close otherwise.Source: §4.3 · pp. 6–7, Table 3 · p. 6, Table 4 · p. 7

4.4 Depth-programmed experts enable elastic-depth inference

  • At L=12, the depth-only gate reaches 44.4 ADE20k mIoU versus approximately 35.6 for the feature-only control, and the proposed variant improves across three probes through L=16 and stays near its best at L=24.ADE20k increases from 40.3 mIoU at L=8 to 44.9 at L=16 and 44.7 at L=24; NYUv2 RMSE changes from 0.522 to 0.485 and 0.487.Source: §4.4 · p. 7, Figure 2 · p. 7

4.5 Number of experts

  • From E=1 to E=8, ImageNet-1k top-1 increases by 9.6, 6.8, and 5.7 points for S/16, B/16, and L/16, respectively, while retaining one dense FFN evaluation per step.The largest single improvement at each scale occurs from E=1 to E=2; beyond E=4, additional capacity primarily benefits S/16.Source: §4.5 · p. 8, Table 5 · p. 8

5 Analysis: what does the recurrent block learn?

The analysis finds structured, resampleable expert routing across recurrent depth, while showing that learned assignments matter and that expert-bank use varies by model scale.

5.1 How the router allocates experts over depth

  • Dominant experts occupy contiguous depth intervals with soft overlap at transitions; the learned schedule is resampleable, and E=8 mixtures are more diffuse.No objective specifies the intervals or their boundaries.Source: §5.1 · p. 8, Figure 3 · p. 8
  • Most expert deletions reduce teacher alignment less than matched random gate perturbations, and removing expert 5 at E=8 has no measurable adverse effect.The matched perturbations have the same merged-weight displacement.Source: §5.1 · p. 8, Figure 4 · p. 8
  • The largest deletion losses occur for earliest-used experts, so deletion sensitivity alone cannot separate an expert’s role from its recurrence position.Source: §5.1 · p. 9
  • All 23 nonidentity E=4 permutations and all 64 tested E=8 permutations reduce teacher CKA, indicating that retaining the bank alone does not preserve alignment.Source: §5.1 · p. 9

5.2 How fully is the expert bank used?

  • Increasing E from four to eight improves S/16 by 2.0 ImageNet points but L/16 by only 0.2.Source: §5.2 · p. 9
  • At E=8, S/16 expert, gate, and merged-weight effective ranks are 7.7, 7.6, and 7.3; L/16 ranks are 6.2, 5.2, and 4.0.The authors interpret S/16 as realizing almost all available directions and L/16 as concentrating merged weights in fewer directions.Source: §5.2 · p. 9, Figure 5 · p. 9

6 Conclusion

reViT replaces the depth-wise stack with one recurrent Transformer module whose normalized-depth coordinate softly merges shared experts into a depth-specific FFN. The resulting trajectory recovers much of the accuracy lost under plain weight tying, while supporting multiple tested depths and fixed-depth dense deployment.

6 Conclusion

  • A continuous normalized-depth coordinate softly merges a shared expert bank into one FFN at each recurrent depth, forming a learned parameter trajectory.Source: §6 · p. 9
  • The trajectory recovers much of the accuracy lost under plain weight tying while retaining substantially fewer stored parameters than a full-depth ViT.Source: §6 · p. 10
  • Elastic-depth training lets the same trajectory be sampled at multiple tested depths; fixed-depth deployment can materialize depth-specific FFNs for dense execution at the cost of compact storage.Source: §6 · p. 10

A Implementation, training, and evaluation details

The implementation routes using normalized recurrent depth, initializes the router for equal expert weights, and applies expert-specific regularization and elastic-depth training.

A.1 Model implementation

  • A two-layer router MLP with 16 hidden units and SiLU maps normalized depth to E logits; its zero-initialized output gives experts equal starting weights, with no learned step embedding and τ=1.Source: §A.1 · p. 10
  • Experts are initialized independently with the dense ViT FFN initialization and no additional variance correction for merged weights.Source: §A.1 · p. 10

A.2 Objectives and regularization

  • Training uses balance, router z-loss, and expert-parameter similarity regularization, with λbal=10^-2, λz=10^-3, and λdiv=10^-3 when E>1.All three expert-specific terms are omitted for E=1.Source: §A.2 · p. 10
  • Elastic-depth training uses per-image depth dropout at rate 0.3; dropped recurrent steps act as identity and retained steps receive recomputed normalized-depth coordinates.Source: §A.2 · p. 10

A.3 Training and evaluation protocols

The evaluation uses a shared supervised ImageNet-1k recipe and a separate fixed-depth DINOv2 distillation setup, then tests frozen-backbone probes across classification, segmentation, and depth prediction.

A.3 Training and evaluation protocols

  • Supervised ImageNet-1k training follows the 300-epoch DeiT III recipe with shared data, augmentation, optimizer, and schedule; recurrent depth is 12 for S/16 and B/16 and 24 for L/16.Source: §A.3 · p. 10, Table 6 · p. 11
  • Fixed-depth DINOv2 distillation trains a reViT-B/14 student at L=12 to match a frozen teacher’s final post-LayerNorm representation; elastic-depth variants are trained afterward.Source: §A.3 · p. 10
  • Frozen-backbone probes use an ImageNet-1k linear classifier, an ADE20k linear segmentation head, and linear and two-layer MLP depth heads on NYUv2.The linear depth head measures information available to a linear readout; the MLP measures what a shallow nonlinear decoder can recover.Source: §A.3 · p. 11

A.4 Computational accounting

  • All reViT results in Tables 1 and 2 average three independently trained runs; downstream probes use fixed seeds, and training takes roughly 192–384 H100 GPU-hours per run.Source: §A.3 · p. 11, §A.4 · p. 11

A.5 Measured inference latency

The appendix reports compiled inference costs and specifies how recurrent adaptations of several MoE families were configured. The comparisons match nominal expert-FFN computation, not necessarily parameter count, total FLOPs, or latency.

A.5 Measured inference latency

  • At batch 1, dynamic reViT uses a 3.6× smaller deployment graph, 48% less total runtime memory, and 13% more latency; at batch 64, latency overhead is 39% and memory is similar.Measurements use compiled PyTorch inference for ViT-S at 224×224 in bfloat16 on one H100.Source: §A.5 · pp. 11–12, Table 7 · p. 12

B MoE baseline adaptations

  • U=1 matches one dense FFN’s arithmetic, but matching U does not match parameter count, total FLOPs, or latency.Higher-budget variants can change expert number, width, or activation pattern.Source: B MoE baseline adaptations · p. 12
  • Lory routes by the preceding recurrent step to merge the next FFN, whereas SMEAR uses a content-conditioned gate recomputed at every depth.Both produce one merged FFN; adding experts changes stored capacity rather than U in these adaptations.Source: B MoE baseline adaptations · p. 13
  • The DeepSeek-V3 adaptation combines routed experts with always-active shared experts, omits its sequence-level balance loss, and applies the same MoE block at all recurrent steps.Its U=1 version uses four routed experts, one active per token, plus one shared expert, all at expert hidden ratio 2.Source: B MoE baseline adaptations · pp. 13–14
  • The U=1 Qwen3 adaptation activates one of four experts and failed to converge under distillation, while MoEUT activates one of four standard FFN experts instead of its usual narrow-expert regime.Only MoEUT’s U=4 configuration reaches its typical K=16 activation regime.Source: B MoE baseline adaptations · p. 14
  • The ReMoE adaptation evaluates all experts as a dense output mixture, while Soft-MoE uses full-width slot experts whose slot counts rise with its computation budget.Because ReMoE evaluates every expert, its nominal computation grows directly with expert count.Source: B MoE baseline adaptations · p. 14

B.1 Layer-wise LoRA baseline

This comparison evaluates layer-wise LoRA updates on a recurrent ViT-S/16 against reViT under the supervised ImageNet-1k recipe. LoRA adds independently trained depth-specific updates, while reViT uses a resampleable depth program.

B.1 Layer-wise LoRA baseline

  • Layer-wise LoRA reaches 77.6% top-1 versus 78.2% for reViT, while storing 7.3M rather than 6.1M compact parameters.Table 10 reports compact checkpoint parameters and distinguishes online from folded GFLOPs.Source: Table 10 · p. 14, §B.1 · p. 15
  • The LoRA baseline reuses one block for 12 depths and adds an independent rank-64 update to the QKV, attention-output, and both FFN projections at each depth.The base weights, biases, and normalization parameters remain shared; training uses the 300-epoch supervised S/16 recipe from random initialization.Source: §B.1 · pp. 14–15
  • Changing recurrent depth requires selecting or interpolating LoRA adapters, whereas reViT evaluates its continuous normalized-depth program on the new grid.Each LoRA adapter corresponds to one trained depth.Source: §B.1 · p. 15

C Regularization ablations

The regularization ablation varies router z-loss and usage-balance coefficients individually while holding the other at its default. Across tested settings, accuracy stays close to the default, which gives the highest observed value in both sweeps.

C Regularization ablations

  • Across tested z-loss and usage-balance coefficients, ImageNet-1k top-1 stays within 0.6 points of the default, which gives the highest observed accuracy in both sweeps.Each sweep varies one coefficient while the other remains at its default; Table 11 reports changes from (λz, λbal)=(10−3, 10−2).Source: Table 11 · p. 15, C Regularization ablations · p. 15

D Additional representation and expert analyses

The additional analyses examine recurrent representation trajectories, their alignment with teacher depth, and expert-bank structure. They also test expert deletion and reassignment while documenting the CKA evaluation protocol.

D Additional representation and expert analyses

  • The CKA analyses compare mean-pooled patch-token residual states after the embedding and each block or recurrent depth, before final LayerNorm, using linear CKA.Self- and cross-CKA use 1,000 validation images; ablation and permutation analyses use 256.Source: D Additional representation and expert… · p. 15
  • Centering the gate and merged-weight trajectories across depth gives each rank at most E−1 because every gate row sums to one.The centered merged projection is GcW1,bank, and Figure 6 marks the resulting rank ceiling.Source: D Additional representation and expert… · p. 15, Figure 6 · p. 16
  • Depthwise self-CKA is banded for DINOv2, Raptor, and both distilled reViT models, including the plain tied E=1 block, indicating that recurrence alone yields a noncollapsed trajectory.The analysis uses mean-pooled patch-token residual states from the embedding through depth 12 on 1,000 ImageNet validation images.Source: D Additional representation and expert… · p. 15, Figure 7 · p. 16
  • Supervised reViT also shows banded self-CKA at S/16 and L/16 scales; the L/16 band narrows visibly over the final quarter of recurrence.The panels compare reViT with same-size DeiT III baselines and describe trajectories rather than establish the expert mechanism.Source: Figure 8 · p. 16, D Additional representation and expert… · p. 17
  • As the expert bank grows from E=1 to E=8, the maximum-CKA path approaches proportional teacher depth, with ρ/MAD improving from 0.969/1.54 to 0.999/0.23.Raptor follows the proportional diagonal exactly under direct intermediate-layer teacher supervision; reViT is distilled only from the teacher’s final output features.Source: D Additional representation and expert… · p. 17, Figure 9 · p. 18
  • No tested gate-column permutation matches the learned assignment: teacher CKA averages fall from 0.793 to 0.208 for E=4 and from 0.849 to 0.202 for E=8.Most expert deletions are less damaging than matched-size random gate perturbations, and the earliest-routed experts cause the largest ablation losses.Source: D Additional representation and expert… · pp. 17–18, Figure 10 · p. 18

Related papers

Loading 2610.12448v1…