Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic
2026-06-08 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study a special behavior in neural networks learning modular addition, where instead of spreading out representations in many dimensions (as neural collapse predicts), the network arranges them in a simple two-dimensional circle. They explain this by showing that classifier weights first form a flat, rank-2 pattern, which then constrains the embeddings to lie on the same circle. They also describe this circular arrangement mathematically and show it has practical benefits in reducing certain complexity measures more than the usual neural collapse solution. Overall, they clarify why the network prefers this circular structure when learning modular arithmetic tasks.
Neural CollapseEquiangular Tight FrameModular AdditionCross-Entropy GradientRank-2 GeometryBackpropagationEntropy-Regularized TransportPhase AlignmentWeight DecaySchatten Norm
Authors
Hu Tan, Kuo Gai, Shihua Zhang
Abstract
While neural collapse (NC) predicts that a $K$-class-balanced classifier should organize terminal representations as a $(K-1)$-dimensional simplex equiangular tight frame (ETF), modular addition consistently enters a different regime: networks compress to a two-dimensional cyclic geometry in which both classifier weights and token embeddings lie on circles. We refine the explanation of this phenomenon in three directions. First, we formalize a layerwise non-uniform training mechanism: downstream classifier weights are driven by dense cross-entropy gradients into a rank-2 equiangular configuration before upstream embeddings fully reorganize, and once this classifier plane forms, backpropagated feature gradients constrain embedding motion to the same plane while weight decay suppresses orthogonal components. Second, after this subspace locking, the induced in-plane dynamics admit an entropy-regularized transport interpretation on $S^1$; combined with modular-addition labels, this reduces embedding formation to phase alignment, whose minimizers are single-frequency characters of $\mathbb{Z}/P\mathbb{Z}$ and hence equal-angle points on a circle. Third, we quantify why this solution prevails over NC: a simplex ETF gains only an $O(1)$ advantage in cross-entropy, whereas the cyclic rank-2 solution enjoys a $Θ(K)$ advantage under Schatten or weight-decay surrogates, yielding a critical threshold $λ_{\mathrm{crit}} = Θ(1/K)$. Our results explain both why classifier weights move first and why embeddings subsequently align with them, showing that grokking on modular arithmetic is governed not by maximal separation alone but by a task-structured trade-off between separation, symmetry, and complexity.