Summary of "Learning to Compute is Not Learning to Retain or Reuse in Looped Transformers"

Summary (Overview)

This paper investigates whether looped transformers that learn to compute answers also learn to retain those answers and reuse their internal states. Using in-context Mealy machine execution with exact symbolic targets, the authors disentangle three distinct properties: computation, retention, and reuse.

• Key finding 1: Training objective dramatically affects retention. Models trained to answer at loop 32 drop to 22.3% accuracy at loop 64, but training with one randomly sampled supervised position per example (SAMPLED-PREFIX) retains 99.8% accuracy at loop 64.
• Key finding 2: Computation is exact—position $k$ is computed at loop $k$, and swapping a single hidden vector at the frontier redirects all later positions to the exact counterfactual execution (237/237 pairs in each of ten models).
• Key finding 3: Computation and reuse are dissociable. Freezing each position's hidden state at the loop its answer first appears (which per-position halting would do) drops endpoint accuracy from 100% to 4.1% under SAMPLED-PREFIX, even though answers are perfectly retained.
• Key finding 4: Randomizing training exits (Jitter-8) helps retain answers much more than a fixed excess budget (Slack-4), while truncated gradients (progressive training) do not help.
• Key finding 5: The released maze network of Bansal et al. (2022) shows the same pattern: it solves all mazes when frozen 4 iterations after answer arrival but none when frozen at the moment of arrival.

---

Introduction and Theoretical Foundation

Background and Motivation

Looped transformers apply one shared block repeatedly, effectively trading recurrence depth for a computation budget. This enables running for more loops than the model was trained on, which can help length generalization (Schwarzschild et al., 2021) but can also degrade performance. Bansal et al. (2022) identified this as "overthinking" and proposed mitigations like progressive training.

The authors note that an overthinking curve conflates three distinct failure modes:

1. The answer was never computed (computation failure)
2. The answer was computed but not retained (retention failure)
3. The state carrying the answer cannot be reused by later computation (reuse failure)

These are impossible to distinguish with accuracy curves alone. The paper's key insight: in-context Mealy machine execution provides exact counterfactual targets, making these distinctions measurable via hidden-state interventions.

Theoretical Framework

A deterministic Mealy machine consists of:

• Finite state set $Q$, input alphabet $\Sigma$, output alphabet $\Gamma$
• Initial state $q_0$
• Transition function $\delta: Q \times \Sigma \to Q$
• Output function $\lambda: Q \times \Sigma \to \Gamma$

Reading actions $a_1, \ldots, a_H$, the machine computes:

$$q_k = \delta(q_{k-1}, a_k), \quad y_k = \lambda(q_{k-1}, a_k), \quad k = 1, \dots, H, \tag{1}$$

where $q_k$ is the state after action $k$ and $y_k$ the output emitted on that transition. The frontier is the diagonal $r = k$ where position $k$ first becomes correct at loop $k$.

---

Methodology

Model Architecture

The model repeats a two-layer transformer block with width 128, four attention heads, and feedforward width 512:

$$X^{0} = 0, \qquad X^{r} = B_{\theta}\big(X^{r-1} + E(x)\big). \tag{2}$$

Here $r$ indexes loops and $X^r$ collects hidden vectors of all tokens after loop $r$. All loops share the same block parameters, with no positional or loop-depth embeddings—only the causal attention mask provides order information.

Training Objectives

The main comparison fixes the training exit $R = H$ while varying supervised positions. With $\ell_{i,k}^{R}$ as the sum of state and output cross-entropies for example $i$, the batch objective is:

$$\mathcal{L}_{S} = \frac{1}{\sum_i |S_i|} \sum_i \sum_{k \in S_i} \ell_{i,k}^{R}, \qquad S_i = \begin{cases} \{H\}, & \text{FINAL-ONLY}, \\ \{1, \ldots, H\}, & \text{ALL-PREFIX}, \\ \{K_i\}, K_i \sim U\{1, \ldots, H\}, & \text{SAMPLED-PREFIX}. \end{cases} \tag{3}$$

• FINAL-ONLY: supervises only the last position's answer
• ALL-PREFIX: supervises all positions (H supervised positions per example)
• SAMPLED-PREFIX: supervises one uniformly chosen position per example (separating position coverage from dense supervision)

Evaluation Metrics

Endpoint accuracy at readout loop $r$ (fraction of examples with both predictions correct at the last action):

$$A(H, r) = \frac{1}{N} \sum_{i=1}^{N} \mathbf{1}\left[\hat{q}_{i,H}^{r} = q_{i,H} \wedge \hat{y}_{i,H}^{r} = y_{i,H}\right]. \tag{4}$$

Intervention Types

1. Swapping (Equation 5): Replace a hidden vector in a base run with the corresponding vector from a source run on a different input prefix: $h_{c}^{r} \leftarrow h_{c,S}^{r}$
2. Resetting: Return a position to its own earlier vector
3. Freezing: Hold each position's hidden vector fixed from a chosen loop onward ($cache$ keeps predictions outside the model; $freeze$ keeps hidden states)

Datasets and Training

• 132,000 AdamW steps, batch size 256, 35-stage curriculum reaching $H = 32$
• Machines have 4, 6, or 8 states, 2 or 4 input symbols, and 2 output symbols
• Main evaluations use 512 traces per model per length (16 machines × 32 traces)

---

Empirical Validation / Results

1. Retention is Determined by Training Requirements (Section 3)

Endpoint retention (loop 64 accuracy after training at loop 32):

Whole-trajectory exactness at loop 64:

• SAMPLED-PREFIX: 96.0%
• ALL-PREFIX: 74.3%

Exit distribution effects (original cohort, matched expected loop budgets):

Gradient truncation: Progressive training (truncated gradients) fails—only 1 of 3 seeds learns length 32—whereas MATCHED RANDOM EXIT (same exits, full gradients) learns in all 3 seeds and keeps 97.9% at loop 64.

2. Computation is Exact, Answers are Transient (Section 4)

Swapping one 128-dimensional hidden vector at position 8, loop 8 (the frontier where $\tau(8) = 8$):

• All 237 eligible pairs in all 10 FINAL-ONLY models follow the source's exact continuation
• One vector steers 24 further steps of exact symbolic execution (positions 9–32)
• Answers fade quickly above the frontier in both swapped and unmodified runs

3. Answers Become Readable Before States Become Reusable (Section 5)

Freezing results (H32, loop 32):

• Errors begin at position 3 (68% accuracy in FINAL-ONLY models), which reads position 1's frozen vector that missed its loop-2 update
• Delaying capture by 2 loops restores SAMPLED-PREFIX to 99.8%; MATCHED RANDOM EXIT needs 4 loops
• Freezing only positions 1–8 in FINAL-ONLY models leaves endpoint correct on at most 11.9% of examples

Maze network (Bansal et al., 2022): Freezing each pixel at its answer's arrival solves none of 32 mazes; freezing 4 iterations later solves all 32.

4. Sample Size and Training Details

• Common cohort: seeds 9–11 (3 seeds per condition)
• Original cohort: seeds 9–14 (6 paired seeds)
• Position-8 transport audit: 10 FINAL-ONLY models (seeds 9–18)

---

Theoretical and Practical Implications

Theoretical Contributions

1. Three-way dissociation: The paper demonstrates that computation, retention, and reuse are distinct properties requiring distinct tests. Correctness at one exit does not test retention; retained predictions do not test whether a hidden state can be held or reused.

2. Supervision structure matters more than supervision density: One uniformly sampled position per example (SAMPLED-PREFIX) outperforms dense supervision of all positions (ALL-PREFIX) and vastly outperforms last-position-only supervision (FINAL-ONLY) for retention.

3. Exit variability vs. gradient truncation: Randomizing training exits (Jitter-8) improves retention more than a fixed excess budget (Slack-4) at matched expected loops. Gradient truncation (progressive training) hinders rather than helps.

4. Halting too early is harmful: Per-position halting at the first readable answer is too early—states need 2–4 additional loops of updates before they can be frozen without disrupting downstream computation.

Practical Implications

• For adaptive computation: Adaptive methods (Universal Transformers, Huginn, progressive training) that halt positions when answers appear may freeze states before they are reusable, causing errors in downstream positions.
• For training design: Randomly sampling supervised positions (even one per example) and randomizing training exits improve robustness to inference-time loop budgets.
• For interpretation: Full-vector interventions with exact counterfactuals can establish causal sufficiency of hidden states, complementing correlation-based analyses.

---

Conclusion

Main takeaways:

1. Learning to compute an answer does not imply learning to retain it: FINAL-ONLY models compute perfectly but retain only 22% of answers 32 loops past the training exit.

2. Retained answers need not reside in reusable states: models that retain 99.8% of answers collapse to 4.1% when positions are frozen at the loop their answers first appear.

3. Retention is trainable via position coverage, exit variability, and full gradient flow—but not via gradient truncation.

4. A single frontier hidden vector is causally sufficient to drive exact subsequent execution, but the states become reusable only 2–4 loops after their answers first appear.

Future directions: The authors note that full-vector interventions do not identify unique encodings or microscopic causes of forgetting; downstream resets on new training controls were not tested. The scope is one synthetic task family plus a limited external maze check. Open questions include the timing mechanism (possibly related to the absence of positional embeddings), whether stochastic-budget models share downstream mechanisms, and whether findings generalize to other architectures and tasks.

Author statement: "Learning to compute is thus not learning to retain, and avoiding overthinking does not make a state reusable when its answer first appears."