Block-Recurrent Dynamics in Vision Transformers

TL;DR

<3-5 sentence high-level summary> The paper investigates why Vision Transformers (ViTs) work so well by proposing the Block-Recurrent Hypothesis (BRH): after training, ViT depth can be represented by a small set of reusable, parameter-tied blocks applied recurrently. It operationalizes BRH with Raptor, a weight-tied recurrent surrogate that reconstructs the full layerwise activations using a few blocks discovered via a max-cut partition of layer similarities, and it validates BRH on foundation models (e.g., DINOv2) with 2–3 blocks recovering a large majority of the teacher’s performance (roughly 96–98% of linear-probe accuracy). The authors then develop a Dynamical Interpretability framework, showing that depth behaves like a discrete dynamical system: token directions converge to angular attractors, token-specific phase dynamics arise at block boundaries, and late-depth updates are low-rank and coherent. Together, these results argue for a compact, interpretable recurrent core in ViTs and provide principled tools for mechanistic analysis and potential efficiency gains without sacrificing performance.

Abstract

As Vision Transformers (ViTs) become standard vision backbones, a mechanistic account of their computational phenomenology is essential. Despite architectural cues that hint at dynamical structure, there is no settled framework that interprets Transformer depth as a well-characterized flow. In this work, we introduce the Block-Recurrent Hypothesis (BRH), arguing that trained ViTs admit a block-recurrent depth structure such that the computation of the original $L$ blocks can be accurately rewritten using only $k \ll L$ distinct blocks applied recurrently. Across diverse ViTs, between-layer representational similarity matrices suggest few contiguous phases. To determine whether these phases reflect genuinely reusable computation, we train block-recurrent surrogates of pretrained ViTs: Recurrent Approximations to Phase-structured TransfORmers (Raptor). In small-scale, we demonstrate that stochastic depth and training promote recurrent structure and subsequently correlate with our ability to accurately fit Raptor. We then provide an empirical existence proof for BRH by training a Raptor model to recover $96\%$ of DINOv2 ImageNet-1k linear probe accuracy in only 2 blocks at equivalent computational cost. Finally, we leverage our hypothesis to develop a program of Dynamical Interpretability. We find i) directional convergence into class-dependent angular basins with self-correcting trajectories under small perturbations, ii) token-specific dynamics, where cls executes sharp late reorientations while patch tokens exhibit strong late-stage coherence toward their mean direction, and iii) a collapse to low rank updates in late depth, consistent with convergence to low-dimensional attractors. Altogether, we find a compact recurrent program emerges along ViT depth, pointing to a low-complexity normative solution that enables these models to be studied through principled dynamical systems analysis.

Figures (15)

• Figure 1: Layer–layer similarity matrices across diverse Vision Transformers reveal block-structure. Despite differences in scale and training objectives, all models exhibit contiguous block structure along depth, visible as phase-segmented regions of high similarity. Beyond representational similarity, this raises the question of whether a deeper functional recurrence underlies these patterns, hinting at block-wise reusability of computation across layers. In this work, we investigate this hypothesis, showing that these phase segments correspond to blocks with functional similarity, which can be approximated by a single shared block applied recurrently along depth.
• Figure 2: Block discovery via max-cut segmentation of the layer–layer similarity matrix. Our algorithm partitions depth into contiguous segments by maximizing within-block similarity and minimizing cross-block cosine similarity. Shown are two cuts of the same ViT-B: with 3-blocks (left, green) and 2-blocks (right, magenta). These cuts reveal candidate block boundaries where the representation dynamics undergo sharp transitions, providing an operational method for detecting contiguous recurrent phases in trained Vision Transformers.
• Figure 3: Evaluation of Raptor models on CIFAR-100 using our max-cut partitioning algorithm versus random partitions. Reported values are classification accuracy. Results for random partitions are aggregated over 10 different random partitions. Random shuffle refers to non-contiguous (fragmented) random partitions.
• Figure 4: Stochastic depth promotes representational similarity across layers block-recurrence.A) ViT layer-layer cosine similarity matrices for models trained with increasing stochastic depth (SD) dropout probability $p$ (probabilities of 0.0-0.9, uniform over layer depth). Dashed red lines delineate blocks, as defined by the max-cut algorithm. Higher SD $p$ values lead to a more similar representation across layers. B) Layerwise teacher-student representational alignment $R^2$ (Raptor vs. ViT) of the class cls and patch tokens. Increases in SD $p$ correspond to an increase in the ability of Raptor to match the ViT's layerwise representations. ViT models for SD=0.7-0.9 show abberant training dynamics and are excluded from this and further analysis. C) CIFAR classification accuracy for a ViT trained on CIFAR, and a Raptor model with $k=3$ blocks trained to match the hidden state of the ViT. D) Last layer hidden-state similarity $R^2$ error (equivalent to $1 - R^2$) of the ViT and Raptor model as a function of SD $p$. Increases in stochastic depth lead to a greater ability to reconstruct ViT function using Raptor . E) Association between layer-layer representational similarity and Raptor reconstruction $R^2$. Stochastic depth encourages the formation of more similar blocks of layers within the ViT, which facilitates approximation by the recurrent Raptor model.
• Figure 5: Three training paradigms for learning recurrent approximations. Each panel shows three token trajectories through depth. Gray dashed lines with filled circles represent the ground-truth teacher trajectories; black solid lines with filled circles show the student's predictions; colored dotted lines (with $\varepsilon$ labels) indicate the error signal between predicted and ground-truth states. Left (Distillation): The student network directly predicts the final layer from the initial state, with no supervision on intermediate representations. Error is measured only at the terminal state, providing no guidance on the representational trajectory. Middle (Teacher Forcing): At each depth step $\ell$, the student block predicts $\hat{\bm{x}}_{\ell+1} = \bm{B}(\bm{x}_\ell)$ using the ground-truth activation $\bm{x}_\ell$ from the teacher. Vertical arrows indicate where the student "resets" to ground-truth states. This enables efficient parallel training and prevents error accumulation, but creates a train-test mismatch since the model never learns to handle its own prediction errors. Right (Autoregressive): The student autoregressively predicts $\hat{\bm{x}}_{\ell+1} = \bm{B}(\hat{\bm{x}}_\ell)$ using its own previous predictions, matching inference conditions. Errors can compound across depth (shown by increasing deviation between trajectories), requiring the model to learn self-consistent, closed-loop dynamics. Our two-stage training (Sec. \ref{['sec:constructive_validation']}) combines both approaches: Stage 1 uses teacher forcing for stable, parallelizable pretraining; Stage 2 switches to autoregressive training to ensure self-consistency at inference.

Theorems & Definitions (4)

• Definition 1: Block-Recurrent Hypothesis (BRH)
• Claim 1: BRH guarantees low Levin complexity
• Claim 2: BRH guarantees low Levin complexity
• proof