Two-Scale Latent Dynamics for Recurrent-Depth Transformers (2509.23314v1)

Abstract: Recurrent-depth transformers scale test-time compute by iterating latent computations before emitting tokens. We study the geometry of these iterates and argue for a simple, \emph{two-scale} operational picture: (i) within a looped block, updates act as \emph{small-scale refinements}; (ii) across consecutive blocks, states undergo a \emph{larger-scale drift}. Across checkpoints, our measurements show that loop steps become \emph{smaller} and increasingly \emph{orthogonal} to one another, indicating better local modeling of fine structure rather than merely pushing in a single direction. These dynamics motivate an early-exit mechanism based on the model's second-order difference in step-size, which we show is superior in terms of performance, stability and time-efficiency, when compared to the KL-divergence exit strategy of Geiping et al. and its naive first-order counterpart.

• The paper presents a two-scale latent dynamics framework that distinguishes fine-grained loop refinements from coarse block transitions.
• It introduces a second-order early-exit rule based on stabilization metrics that reduces latency while preserving output quality.
• Empirical results on a GPT-2-style model demonstrate that the acceleration exit efficiently balances compute allocation and performance.

Two-Scale Latent Dynamics in Recurrent-Depth Transformers

The paper "Two-Scale Latent Dynamics for Recurrent-Depth Transformers" (2509.23314) presents a geometric analysis of latent-space computation in recurrent-depth transformer architectures. The authors investigate the internal dynamics of looped transformer blocks, revealing a two-scale separation: fine-grained refinements within looped blocks and coarser drift across block boundaries. This insight motivates a second-order, geometry-inspired early-exit criterion that improves the latency-quality trade-off in test-time compute allocation.

Geometric Characterization of Recurrent Depth

The central thesis is that recurrent-depth transformers, which iterate block computations before emitting tokens, exhibit distinct geometric patterns in latent space. Within a looped block, the model performs small, increasingly orthogonal updates, effectively tracing a stable-curvature spiral around a local representation. In contrast, transitions between blocks induce larger, directional changes—constituting a slow drift at a coarser scale.

Empirical analysis is conducted on a GPT-2-style model with recurrence enabled in three regions: layer 4 (self-loop), layers 5–6 (paired loop), and layer 7 (self-loop). The authors track two diagnostics: the norm of the update step ($\|\Delta^{(k)}\|_2$) and the cosine of the angle between consecutive updates ($\cos\angle(\Delta^{(k)}, \Delta^{(k-1)})$). These statistics are visualized using PCA projections of hidden states.

Figure 1: 2D PCA trajectory of hidden states, showing tight loop refinements (inset) and larger cross-block moves (main view).

The PCA trajectories consistently show tight arcs within loops and larger jumps at block hand-offs, supporting the two-scale separation. Step norms decay rapidly within $5$–$10$ loop steps, and step angles stabilize at moderate values ($\approx 0.5$–$0.65$), indicating non-collinear, refinement-like updates.

Figure 2: Loop dynamics for groups 4, 5–6, and 7. (a) Step norms decay rapidly; (b) step angles stabilize, indicating orthogonal refinements.

Early-Exit Criteria: Geometry-Inspired Acceleration

The geometric findings motivate a second-order early-exit rule for looped blocks. The standard approaches—step-norm and KL-divergence—have limitations: step-norm is sensitive to feature scaling and can mis-halt under spiral dynamics, while KL-divergence is computationally expensive and reacts late to latent rotations.

The proposed acceleration criterion triggers when the norm of the difference between consecutive updates ($a^{(k)} = \|\Delta^{(k)} - \Delta^{(k-1)}\|_2$) falls below a threshold for two consecutive steps. This directly detects stabilization of the local spiral, enabling earlier and more reliable exits.

Algorithmically, the two-hit acceleration exit is implemented as follows:

def acceleration_exit(f, x0, tau, k_max): k = 0 delta_prev = None prev_small = False while k < k_max: x1 = f(x0) delta_cur = x1 - x0 if delta_prev is not None: a = np.linalg.norm(delta_cur - delta_prev) small = (a < tau) if small and prev_small: break prev_small = small delta_prev = delta_cur x0 = x1 k += 1 return x0, k

This method is computationally cheap ($\mathcal{O}(d)$ per token-step), decoding-free, and robust to threshold variation.

Empirical Comparison of Exit Policies

The authors benchmark three exit policies—step-norm, KL-divergence, and acceleration—on latency, perplexity, and cross-entropy. Acceleration achieves substantial latency reductions (e.g., $580 \rightarrow 360$ ms/token as $\tau$ increases from $10^{-5}$ to $10^{-2}$) while maintaining quality. KL-divergence preserves quality but is slower due to vocabulary decoding. Step-norm is fast but suffers a sharp quality cliff at aggressive thresholds.

Figure 3: Latency vs. threshold for step-norm, KL, and acceleration exit policies.

Figure 4: Full $\cos\angle(\Delta^{(k)},\Delta^{(k-1)})$ curves across checkpoints, showing stabilization of step angles in recurrent groups.

Acceleration consistently delivers a superior latency-quality Pareto frontier, avoiding the instability of step-norm and the computational overhead of KL-divergence. The method is robust across a wide range of thresholds, reducing the need for hyperparameter tuning.

Theoretical and Practical Implications

The geometric analysis provides a mechanistic understanding of how recurrent-depth transformers allocate compute in latent space. The two-scale separation—local spiral refinements within blocks and slow drift across blocks—suggests that looped blocks are well-suited for fine-grained reasoning, while depth transitions handle coarser semantic changes.

Practically, the acceleration-based early-exit enables dynamic allocation of test-time compute without sacrificing output quality or incurring significant computational overhead. This is particularly relevant for latency-sensitive applications and for scaling models to larger architectures or datasets.

Theoretically, the findings invite further investigation into the role of latent-space geometry in model reasoning and the design of controllers for dynamic-depth or token-level compute allocation. Extensions could include per-block calibration, integration with hierarchical controllers, and scaling to larger models.

Conclusion

This work provides a geometry-first perspective on recurrent-depth transformers, revealing a two-scale separation in latent dynamics and introducing a second-order, decoding-free early-exit criterion. The acceleration-based exit achieves lower latency at matched quality compared to KL-divergence and avoids the instability of step-norm thresholds. The approach is computationally efficient and robust, with actionable implications for dynamic compute allocation in transformer architectures. Future directions include scaling to larger models, learning per-block calibration, and integrating geometry-aware controllers for adaptive test-time compute.