Rotate the Loop: Blockwise Conformality for Looped Transformers

AI summary

Overview

Research area: Machine learning, specifically looped (recurrent-depth) transformers, linear transition operators, state tracking, and language modelling.

Technical level: Advanced.

Scope: This paper argues that whether a looped transformer can carry information across recurrent steps is a geometric property of the linear map applied between steps, and it proposes a blockwise conformal replacement for the usual diagonal map.

What This Paper Is About

Looped transformers reuse the same transformer block across many recurrent steps, so additional computation can be bought at test time without adding parameters. Those extra steps only help if information written on an early step remains readable on later ones, and the paper shows that this depends on whether the linear map between steps preserves angles rather than merely shrinking magnitudes. The authors replace the standard learned diagonal transition operator with a blockwise conformal map: the residual state is split into coordinate-aligned groups, each group rotates internally at its own rate, and the resulting operator is tested on finite-group state tracking and language modelling.

Key Contributions

1. A theoretical link between readability and conformality. The paper defines a transition operator as “readable” when it preserves inner products up to one shared scale, and proves this holds if and only if the operator is a scaled orthogonal matrix, (M = cQ) with (Q^\top Q = I). It further shows that repeatedly applying a non-conformal operator collapses state directions toward a dominant eigenspace, while a conformal operator keeps directions moving and distinguishable across depth.

2. A blockwise conformal parameterization. The residual is partitioned into (G) contiguous groups of equal width; coordinates inside each group are paired into rotation planes. Each group has its own contraction rate (c_g \in (0,1)), and each plane has its own frequency (\omega_i), with the phase tied to the group’s discretization rate. The operator is never materialized and is implemented through the same rotation kernel used by rotary position embeddings, applied on the depth axis rather than the token axis.

3. Evidence on state tracking. On (A_5) word problems with sequence length (n=128), width (512), and a training recurrence budget (T_{\max}=32), the blockwise conformal operator at (G=2) reaches last-token accuracy (0.894 \pm 0.052) at (T=32), versus (0.692 \pm 0.097) for the published diagonal baseline. On the harder (S_5) task, the corresponding numbers are (0.808 \pm 0.054) versus (0.531 \pm 0.092).

4. Evidence on language modelling and measured geometry. Across (60)M, (140)M, and (370)M parameter models, the conformal operator lowers validation perplexity and raises LAMBADA exact match once more than one recurrent step is used. A “Toeplitz defect” measurement confirms that the trained conformal operators preserve the predicted geometry to numerical precision, while the diagonal operators’ defect grows with recurrent depth.

Main Findings

• Readability is exactly conformality. For a matrix (M), preserving (\langle Mx, My\rangle = c^2\langle x,y\rangle) for all (x,y) is equivalent to (M = cQ) with (Q^\top Q = I). Because transformer blocks compare states through inner products, a conformal map keeps those comparisons intact up to a shared scale that pre-normalization removes.

• Repeated non-conformal maps collapse the state. For a normal operator with a strictly dominant eigenvalue magnitude, the normalized state converges geometrically to the dominant eigenspace. Writes from different recurrent steps end up pointing in the same direction, so the block can no longer tell which step produced them. A conformal operator instead keeps the direction moving on the unit sphere.

• Conformality gives depth invariance. If (A = cQ), the normalized Gram matrix between recurrent steps (t) and (s) depends only on (t-s), not on where in the loop the pair sits. Extra recurrent steps therefore do not rewrite the geometry the block reads.

• Blockwise grouping yields multiple timescales. Splitting the state into (G) coordinate subspaces gives each group its own scale (c_g) and orthogonal map (Q_g). Readability holds inside each group with constant (c_g); across groups, relative component norms evolve as ((c_g/c_h)^t). The construction is unitarily equivalent to a complex diagonal with free phases per plane and magnitudes tied within a group.

• State-tracking gains are substantial at the trained budget. On (A_5), the diagonal baseline scores (0.537 \pm 0.098), (0.692 \pm 0.097), and (0.771 \pm 0.079) at (T=16,32,48). The (G=2) conformal model scores (0.812 \pm 0.058), (0.894 \pm 0.052), and (0.926 \pm 0.040). The gap is present at the trained budget and persists at (T=48), a recurrence budget the checkpoints never saw.

• The group count has a narrow sweet spot. (G=1) already beats the diagonal: (0.628), (0.746), and (0.795) at (T=16,32,48) on (A_5). (G=2) is the peak. (G=4) and (G=8) fall below the diagonal baseline, at (0.573) and (0.614) at (T=32) versus (0.692). Larger counts recover partially: (G=256), one group per rotation plane, reaches (0.741) at (T=32), close to (G=1)’s (0.746). The authors interpret this as tied magnitudes helping at small (G), while phases help throughout; giving many groups their own rates reintroduces the spectral spread the method is meant to remove.

• Tying rates is the larger of the two effects. Removing rotation but tying rates inside (G=2) groups raises (T=32) accuracy from (0.692) to (0.810 \pm 0.036), recovering (0.118) of the (0.202) gap to (0.894). Rotating without tying magnitudes, the (G=256) case, recovers (0.049). Doing both recovers the full (0.202). Neither ingredient alone reproduces the full gain.

• The parameterization chart matters. At (T_{\max}=16), a diagonal baseline scores (0.408 \pm 0.061); the pairwise rotation chart scores (0.721 \pm 0.042); Householder QR over (O(d)) scores (0.552 \pm 0.054); and Cayley over (SO(d)) scores (0.274 \pm 0.059), below the diagonal. The theory treats all orthogonal bases as equivalent, but the optimizer does not, consistent with the privileged-basis account the paper cites.

• Language modelling improves with depth only for the conformal operator. Validation perplexity is measured on (780)k FineWeb tokens and LAMBADA exact match on (5153) examples. At (T=1) the methods are close: at (60)M, perplexity is (28.40) versus (29.10), and LAMBADA is (0.168) versus (0.174); LAMBADA is slightly behind at two of three sizes. From (T=8), the conformal operator is ahead by about (15%) on perplexity and about (20%) on LAMBADA. At (T=48), it is ahead by (17.1%) on perplexity and (21.1%) on LAMBADA at (60)M, (16.9%) and (19.9%) at (140)M, and (17.0%) and (21.8%) at (370)M.

• The diagonal model stops benefiting from extra depth. Its validation perplexity is flat from (T=8) onward: (22.45) at (60)M for all of (T \in {8,16,32,48}), and the same to within a hundredth at (140)M and (370)M. The conformal model keeps improving through (T=48) at all three sizes, ending at (18.62), (16.18), and (12.37) for the three model sizes.

• The trained geometry matches the theory. Averaged over (256) probe directions, the conformal Toeplitz defect is between (3.6\cdot10^{-16}) and (5.4\cdot10^{-16}) across sizes and depths, within precision error. The diagonal defect grows with depth and width: (0.059 \pm 0.011) at (60)M and (T=8), rising toward (0.283 \pm 0.023) by (T=32), and already (0.220 \pm 0.047) at (370)M and (T=8).

Methodology in Plain English

A looped transformer applies one shared block repeatedly to a residual state and re-injects the same prelude embedding at every pass. Between passes, a linear transition operator is applied. The published Parcae baseline uses a learned diagonal operator with positive entries in ((0,1)), which guarantees contraction but gives every coordinate its own independent decay rate.

The authors replace that diagonal with a structured rotation. They split the residual dimensions into a small number of contiguous groups, pair coordinates inside each group into two-dimensional planes, and apply a planar rotation to each pair, scaled by a per-group decay factor. Each group therefore has its own lifetime, and each plane has its own angle. The operator is implemented with the same rotation machinery used in rotary position embeddings, so it is applied efficiently without constructing a full (d \times d) matrix.

For state tracking, the model reads generators of a finite group and

Authors’ abstract

Looped transformers reuse one shared block across many recurrent steps, so depth becomes a quantity that can be spent at runtime rather than a stack of distinct layers. Those extra steps help only if a write made on one recurrent step is still readable on later ones, and that depends on the linear map applied to the residual between recurrent steps. We show that whether the loop learns to track state is a geometric question about the form of this map, and we propose blockwise conformal injections: the residual splits into coordinate-aligned groups and rotates inside each one, so every group has its own timescale while the geometry stays exact within it. Our experiments show a performance increase on a demanding form of state tracking - last-token accuracy when the rest of the prefix is already correct - and on language modelling, where a signal written earlier has to be retrieved at the end of the sequence.