Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks
Organizations: University of Oxford · BBC R&D
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 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.
Figures & tables
| Task | Looped | Untied -layer | |
|---|---|---|---|
| , untied layers | – | ||
| , auxiliary loss | – | ||
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Model, | Probes | Seed | Hop 1 | Hop 2 | Hop 3 | Answer | Accuracy |
|---|---|---|---|---|---|---|---|
| Looped , | row split | 42 | 8 | 12 | 16 | 16 | 1.000 |
| row split | 43 | 8 | 12 | 16 | 20 | 1.000 | |
| row split | 44 | 8 | 12 | 16 | 20 | 1.000 | |
| row split | 45 | 8 | 16 | 20 | 28 | 1.000 | |
| row split | 46 | 8 | 12 | – | 32 | 0.650 | |
| prompt-disjoint | 42 | 8 | 11 | 15 | 16 | 1.000 |
Recurrent Transformers
Recurrent Transformers reuse layers or state across steps to extend context, reasoning depth, and sequence modeling with less parameter growth.
See all 172 papers