Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks
Magnus Sesodia ⋅ Philip Torr ⋅ Christian Schroeder de Witt
Abstract
Solving algorithmic tasks that involve many sequential steps typically requires scaling up transformer depth or chain-of-thought length, incurring unfavorable parameter and inference-compute costs, respectively, as tasks deepen. Looped transformers promise to recover that depth at a fraction of the cost by reusing a single block across iterations, but it remains unclear how the loop actually performs its computation, or whether it merely simulates a deeper untied stack. We test this on \emph{directed permutation traversal} (\dePo), a synthetic task in the physics-of-LMs lineage in which solving an $H$-hop query requires following $H$ pointers along a fixed permutation of $N{=}16$ entities, with $H$ up to $15$. We demonstrate that, above a minimum effective depth, a $3.2$\,M-parameter looped transformer trained with hop length sampled uniformly per batch traverses all $15$ hops across all seeds, while a $25.8$\,M-parameter fixed-depth baseline at the same effective depth fails to exceed three hops; weight tying, not effective depth, predicts capability. In the same models, linear probes on the residual stream uncover a clean \emph{staircase}: at iteration $k$ the stream linearly decodes hop $k$, with slope exactly one hop per iteration, and each hop stays decodable for roughly two further iterations before being overwritten. The model's own logit lens, applied through the trained final layers, traces the same one-hop-per-iteration ordering at lower amplitude. Taken together, our results recast looped transformers as iterative program learners: weight tying forces the optimizer to discover a single per-iteration update rule, and that rule corresponds to one algorithmic step of the underlying task. Code is available at \url{https://anonymous.4open.science/r/insidetheloop-3ED0}.
Chat is not available.
Successful Page Load