Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks

AI summary

Overview

Research area: Mechanistic interpretability of looped (weight-tied, depth-recurrent) transformers on algorithmic reasoning tasks. Technical level: Advanced — assumes familiarity with transformer internals, linear probing, and activation patching. Scope: A controlled study of whether a looped transformer, at matched effective depth, learns an iterative algorithm that uses one loop iteration per algorithmic step, compared against an untied transformer of the same effective depth.

What This Paper Is About

Transformers that solve tasks requiring many sequential steps usually need either more layers (more parameters) or longer chain-of-thought (more inference compute). Looped transformers reuse one block of layers across several iterations, scaling effective depth for free — but it is unclear whether each iteration actually executes one step of an algorithm, or whether the loop merely imitates a deeper untied stack. The authors study this on a synthetic task, directed permutation traversal, where every intermediate step has a known correct answer, so they can see exactly where in the network each step becomes decodable.

Key Contributions

1. Looped models are better learners at matched effective depth. Trained identically, looped models solve the task where untied models of the same effective depth do not, even though the untied model has roughly T times more parameters (3.22M for the 4×8 looped model versus 25.72M for the untied 32-layer model).
2. Looped models build a "staircase" of intermediate states. Without any supervision of intermediate hops, looped models make successive reference nodes linearly decodable a few layers apart, at a rate set by the iteration budget.
3. The staircase is causally used. Logit lens and activation patching show the model's own output head reads the current node and that the patched state determines which walk the model finishes.
4. An open benchmark task and code release. The directed permutation traversal (dePo) task and code at https://github.com/magnus-sesodia/insidetheloop.

Main Findings

• Effective depth is necessary but not sufficient. No configuration at or below effective depth 8 solved the task, but at effective depth 32 the 8×4, 4×8 and 2×16 looped models solved it in at least four of five seeds, while the untied 32-layer model and the 16×2 model (looping only twice) never did.
• No untied model solved at the default rate. At η = 10⁻³, no untied model from 2 to 32 layers solved the task; the untied 32-layer model's training loss changed by less than 0.03 between the end of the hop curriculum at 30k steps and the end of training at 100k.
• A lower learning rate rescues the untied model. At η = 3×10⁻⁴, the untied 32-layer model solved in four of five seeds and the untied 16-layer model in one; the looped 4×8 model still solved in all five and reached the solution sooner (45k steps versus 55k). At η = 10⁻⁴ both solved in only one of five seeds.
• Intermediate supervision also rescues it. Adding a next-token loss at layers 8, 16 and 24 of the untied 32-layer model at η = 10⁻³ (total weight 0.3) made it solve in two of five seeds.
• Different solution geometry. At H = 4, the looped 4×8 model revealed successive nodes a few layers apart, reaching the answer node by layer 16. The untied model at the same rate never made more than one hop decodable beyond the start. At η = 3×10⁻⁴ the untied model exposed no node beyond the start for the first 25 layers, then all three intermediate nodes and the answer between layers 26 and 30.
• One iteration ≈ one hop, when budget allows. At H = 8, the 2×16 model advanced about one hop per iteration in two of three seeds and finished by iteration 12 of 16; the 4×8 model roughly followed the identity line in all three seeds (though needing about two iterations for the first hop and making two hops in one later iteration); the 8×4 model packed several hops into single iterations and finished after its fourth and final iteration.
• Hop order is preserved. At H = 12, where four backward steps would reach the answer, the 2×16 model still made every forward hop in turn.
• The output head reads the staircase. Logit lens on the 4×8 model showed the head writing the node of hop t after iteration t, with probability rising from 0.53 at hop 4 to 1.00 at hop 7, and never putting more than 0.10 of its probability on already-passed nodes.
• Causal control. Replacing the receiver's states at its two answer positions with a donor prompt's (same permutation, different start node) made the share of donor answers reach 1.00 after the second iteration in two of three seeds and after the sixth in the third; an untrained model wrote neither answer.
• Longer chains at N = 32. The untied 32-layer model never got past the first hop in any of three seeds, while the looped 4×8 model, with an eighth of the parameters, solved at least the first six hops in one of three seeds.
• Caveats the authors report. The logit lens tracked every hop in only one seed per iteration budget; patching with a donor prompt having a different permutation yielded the composed node in only one of five seeds and in no seed of the 2×16 model; and the 4×8 model is the only looped configuration probed in depth (four configurations probed in total).

Methodology in Plain English

The authors build a task where the answer is a chain of lookups with no ambiguity. A model is given the edges of a cycle of 16 nodes in shuffled order, plus a start node and a number of hops H, and must report the node reached after H hops. Because the cycle is closed, at most min(H, 16−H) ≤ 8 forward-or-backward steps are ever needed, and every intermediate node is known in advance even though it is never supervised.

They train two families of decoder-only transformers with identical width (256), 8 heads, SwiGLU MLPs and RoPE: looped models that apply one L-layer block T times with a learnable per-iteration embedding, and untied models with L·T distinct layers. Effective depth D_eff = L·T is matched, so both do the same number of layer evaluations per forward pass; the looped model just has about T times fewer parameters. They sweep L ∈ {2, 4, 8, 16} with T up to 16, plus untied models from 2 to 32 layers, training 100k steps with five seeds each under a hop curriculum that widens from H = 1 to H ≤ 15.

To see inside the model they fit a linear probe per layer and per hop that tries to decode each intermediate node from the residual stream at the position just before the answer, and they report probe selectivity — the gain in held-out accuracy over a probe trained on shuffled labels. A hop is "decodable" once selectivity reaches 0.5. They also apply the model's own final norm and unembedding after each iteration (logit lens) and swap hidden states between prompts (activation patching) to check the states are actually used.

Why This Matters

The paper converts a vague claim — that weight tying gives an inductive bias toward iterative computation — into a measurable one, and it shows the loop budget is a real design parameter: a model with fewer iterations than task steps will pack multiple steps into one iteration rather than fail. It also offers a partial explanation for why earlier mechanistic work on pretrained looped models (Huginn-3.5B) found only limited evidence of a per-iteration chain: those studies used the logit lens without training probes, and in these models the lens tracks every hop in only one seed per budget whereas probes decode every hop in most seeds at H = 4. The authors are careful to note their task has a correct node at every step and trains every hop length, conditions generic pretraining lacks.

Real-world applications that could benefit:

• Latent reasoning models that must scale sequential computation without emitting chain-of-thought tokens.
• Algorithmic and graph-reachability workloads where each step is verifiable and depth, not parameters, is the bottleneck.
• Parameter-constrained deployment, since a looped 4×8 model stores roughly an eighth of the weights of an untied 32-layer model at the same effective depth.
• Model-design tooling for choosing iteration budgets and early-exit policies in recurrent-depth architectures.

Industrial relevance: labs building recurrent-depth or "looped" LLMs (the paper cites Huginn and Ouro-style architectures and middle-cycle sharing between an untied prelude and coda) can use these results to reason about how many iterations to train with, whether to supervise intermediate layers, and why an untied stack at the same effective depth may need a lower learning rate.

Future Directions

• Testing whether the staircase survives on tasks that mix retrieval with computation, and on natural language, rather than a single synthetic permutation task.
• Scaling beyond the ~26M-parameter models studied here, including the regime where the loop is wrapped in an untied prelude and coda, where Huginn and Ouro report their largest gains.
• Understanding mechanistically why tied training reaches the staircase without intermediate supervision while the untied stack needs it.
• Exploring early exit once a walk ends, and whether running a hop-packing model for more iterations than it was trained on extends its walk — which the authors leave untested beyond the trained budget.

Target Audience

Interpretability and architecture researchers working on recurrent-depth and latent-reasoning transformers; practitioners deciding between depth, parameter count and loop iterations; and graduate students looking for a clean, fully supervised testbed for mechanistic claims about iterative computation. Readers should be comfortable with residual streams, probing, logit lens and activation patching.

Authors’ abstract

Solving algorithmic tasks with many sequential steps typically requires scaling up transformer depth or chain-of-thought length, incurring unfavorable parameter and inference-compute costs, respectively. Looped transformers, which reuse one block of layers across many iterations, are poised to address this, since they scale effective depth without adding parameters or tokens. Yet, it remains unclear what each iteration is actually computing, or whether the loop merely simulates a deeper untied stack. Here, we study directed permutation traversal, in which a model follows $H$ hops from a starting node and reports the node it reaches, so every intermediate hop has a known ground truth. We find that a minimum effective depth is necessary to solve the task, beyond which the degree of weight tying becomes the determining factor. For instance, a looped transformer with a four-layer block iterated eight times solves the task in every seed, while an identically-trained untied transformer of the same effective depth never does, succeeding only with a lower learning rate or supervision of its intermediate layers. Probing and intervening on the residual stream, we find that the looped model advances through the intermediate nodes gradually, forming a staircase that adapts to its iteration budget, whereas the untied model, at that lower rate, exposes them only in its final layers. Taken together, our results recast looped transformers as iterative program learners, which allocate their iterations to successive steps of the algorithm and solve tasks that identically-trained untied models of the same effective depth do not. Code is available at https://github.com/magnus-sesodia/insidetheloop.