Understanding Dynamic Compute Allocation in Recurrent Transformers

Paper: 2602.08864 Authors: Ibraheem Muhammad Moosa, Suhas Lohit, Ye Wang, Moitreya Chatterjee, Wenpeng Yin Categories: cs.CL, cs.LG

Abstract

This paper studies token-level adaptive computation in recurrent transformers through a rigorous empirical framework. It introduces complexity-controlled evaluation using algorithmic and synthetic tasks with parameterized difficulty, and proposes ANIRA, a unified recurrent Transformer framework supporting per-token variable-depth computation. Key findings include: compute allocation aligned with task complexity can emerge without explicit supervision, but such alignment does not imply algorithmic generalization — models fail to extrapolate to unseen problem sizes. Furthermore, early halting decisions rely on static structural cues, while online halting tracks the algorithmic execution state more faithfully.

Key Contributions

• ANIRA framework: A unified recurrent Transformer architecture supporting per-token variable-depth computation
• Complexity-controlled evaluation: Algorithmic tasks with parameterized difficulty for rigorous analysis
• Emergent compute alignment: Demonstrates that complexity-aware allocation arises without explicit supervision
• Alignment-generalization gap: Shows that aligned compute allocation does not guarantee out-of-distribution generalization
• Halting mechanism analysis: Distinguishes between static structural halting and dynamic execution-state halting

Background: Adaptive Computation in Transformers

The Fixed-Depth Problem

Standard transformers apply a fixed number of layers to every token, regardless of the computational difficulty of processing that token. This is wasteful for easy tokens and potentially insufficient for hard ones.

Recurrent Transformers

Recurrent transformers address this by allowing layers to be applied iteratively. A single block of parameters is reused across multiple “depth” steps:

$h^{(t+1)} = f_\theta(h^{(t)})$

where $f_\theta$ is the shared transformer block and $t$ indexes the recurrence depth. The key question is: how many steps should each token receive?

Adaptive Computation Time (ACT)

ACT and its variants learn a halting probability at each step:

$p_{\text{halt}}^{(t)} = \sigma(w^T h^{(t)})$

The model stops recurring when the cumulative halting probability exceeds a threshold. This allows different tokens to receive different amounts of computation.

The ANIRA Framework

ANIRA (Adaptive N-Iteration Recurrent Architecture) unifies several recurrent transformer variants under a single framework, enabling controlled comparison. It supports:

• Fixed-depth: All tokens receive the same number of iterations (baseline)
• Per-token adaptive depth: Each token independently decides when to halt
• Global adaptive depth: A single halting decision for the entire sequence
• Hybrid: Combination of per-token and global signals

Complexity-Controlled Tasks

To rigorously study compute allocation, the paper designs tasks with parameterized difficulty:

• Sorting: Difficulty controlled by sequence length and value range
• Graph reachability: Difficulty controlled by graph size and density
• Arithmetic: Difficulty controlled by operand size and operation complexity
• String manipulation: Difficulty controlled by pattern complexity

Each task has a known algorithmic solution with well-understood computational complexity, enabling precise measurement of whether the model’s compute allocation matches the theoretical requirements.

Key Findings

Finding 1: Emergent Complexity Alignment

Without any explicit supervision about task difficulty, recurrent transformers learn to allocate more computation to harder instances. For example:

• Longer sequences in sorting tasks receive more recurrence steps
• Larger graphs in reachability tasks trigger deeper computation
• More complex arithmetic expressions use more iterations

This alignment emerges purely from the task loss signal.

Finding 2: Alignment Does Not Imply Generalization

Despite learning to allocate compute proportionally to difficulty, models fail to extrapolate to problem sizes not seen during training. A model trained on sorting sequences of length 10-50 that correctly allocates more compute to length 50 will still fail on length 100, even though it “knows” length 100 should require more compute.

This reveals a fundamental distinction: knowing how much to compute is not the same as knowing what to compute.

Finding 3: Two Types of Halting

The paper identifies two distinct halting strategies:

1. Static structural halting: The model decides how many steps to use based on input features (e.g., sequence length) before processing begins. This is like estimating difficulty from surface features.

2. Dynamic execution-state halting: The model monitors its internal computation state and halts when the computation converges. This tracks the actual algorithmic execution.

Early layers tend to use static halting, while later layers shift toward dynamic halting. Dynamic halting is more faithful to the actual computation but harder to learn.

Results

• ANIRA provides a clean framework for comparing adaptive computation strategies
• Complexity-aligned allocation emerges reliably across tasks
• Generalization to unseen sizes fails consistently, even with correct compute allocation
• Dynamic halting correlates better with algorithmic correctness than static halting
• Per-token adaptive depth outperforms global adaptive depth on heterogeneous-difficulty inputs

Significance

1. Sobering finding: Adaptive compute alignment is necessary but not sufficient for algorithmic generalization
2. Evaluation methodology: Complexity-controlled tasks provide a rigorous framework for studying adaptive computation
3. Mechanistic insight: The static vs. dynamic halting distinction reveals how models internally manage computation
4. Practical implications: Suggests that achieving true algorithmic generalization requires more than just adaptive depth — it may require explicit algorithmic inductive biases
5. Unified framework: ANIRA enables fair comparison across adaptive computation methods, advancing the field’s ability to draw reliable conclusions