Low-Dimensional Execution Manifolds in Transformer Learning Dynamics: Evidence from Modular Arithmetic Tasks

TL;DR

The paper investigates the geometry of learning dynamics in overparameterized transformer models trained on a controlled marker-based modular arithmetic task. It shows that training trajectories rapidly collapse onto low-dimensional execution manifolds of about $3$--$4$ dimensions within a $d=128$ parameter space, though the manifold orientation depends on the seed. Core contributions include (i) a geometric explanation for attention bubbling via saturation along routing coordinates, (ii) evidence that SGD updates are approximately integrable when projected onto the execution subspace, with non-integrability localized to orthogonal staging directions, (iii) demonstration that sparse autoencoders capture peripheral routing structure but do not isolate execution, and (iv) analysis of curriculum versus mixed training on manifold discovery and compositional generalization. The results suggest a unifying, trajectory-centered view of transformer learning, with implications for interpretability, curriculum design, and understanding the role of overparameterization in enabling stable, low-dimensional computation amid high-dimensional optimization noise.

Abstract

We investigate the geometric structure of learning dynamics in overparameterized transformer models through carefully controlled modular arithmetic tasks. Our primary finding is that despite operating in high-dimensional parameter spaces ($d=128$), transformer training trajectories rapidly collapse onto low-dimensional execution manifolds of dimension $3$--$4$. This dimensional collapse is robust across random seeds and moderate task difficulties, though the orientation of the manifold in parameter space varies between runs. We demonstrate that this geometric structure underlies several empirically observed phenomena: (1) sharp attention concentration emerges as saturation along routing coordinates within the execution manifold, (2) stochastic gradient descent (SGD) exhibits approximately integrable dynamics when projected onto the execution subspace, with non-integrability confined to orthogonal staging directions, and (3) sparse autoencoders capture auxiliary routing structure but fail to isolate execution itself, which remains distributed across the low-dimensional manifold. Our results suggest a unifying geometric framework for understanding transformer learning, where the vast majority of parameters serve to absorb optimization interference while core computation occurs in a dramatically reduced subspace. These findings have implications for interpretability, training curriculum design, and understanding the role of overparameterization in neural network learning.

Figures (6)

• Figure 1: Effective rank (at $90\%$ variance) of attention weight matrices ($W_Q$, $W_K$, $W_V$, $W_O$) across layers. All matrices collapse to dimension $3$--$5$ despite ambient dimension $d^2 = 16{,}384$, with $W_V$ and $W_O$ consistently lower-dimensional than $W_Q$ and $W_K$. Error bars show standard deviation across 5 random seeds.
• Figure 2: Mean attention entropy decreases as training progresses, indicating the emergence of sharp attention "bubbles." The decrease is continuous and smooth, consistent with geometric saturation along the execution manifold.
• Figure 3: Raw commutator norms $\|[\nabla_A, \nabla_B]\|$ remain large throughout training, with frequent spikes indicating strong non-commutativity in the full parameter space.
• Figure 4: The ratio $\rho = \|[\nabla_A, \nabla_B]_{\parallel}\| / \|[\nabla_A, \nabla_B]\|$ rapidly approaches zero, indicating that non-commutativity is confined to directions orthogonal to the execution manifold.
• Figure 5: Decomposition of commutators into components parallel ($\parallel$) and perpendicular ($\perp$) to the execution manifold. The perpendicular component dominates, accounting for $>95\%$ of total magnitude.