Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic

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.

• The paper shows that for modular addition tasks, trained networks converge to a 2D cyclic geometry instead of the canonical high-dimensional Neural Collapse structure.
• It reveals a hierarchical collapse where classifier weights first compress to a 2D subspace, guiding embeddings to align with cyclic group symmetries.
• The findings imply that task-specific algebraic structures and regularization can drive low-rank representations, offering insights into grokking and algorithmic learning.

Task-Intrinsic Geometric Structure in Neural Representations of Modular Arithmetic

Introduction

The recent study "Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic" (2606.08985) challenges canonical predictions of the Neural Collapse (NC) paradigm, particularly in the modular addition task. While NC asserts that, for $K$-class problems, neural networks' learned representations in the terminal phase of training collapse to a $(K-1)$-dimensional simplex equiangular tight frame (ETF), empirical and theoretical analysis in this work reveals a fundamentally distinct outcome for modular addition: the emergence of a rank-2 cyclic geometry, with both classifier weights and upstream token embeddings arranged as points along circles in a two-dimensional subspace. This work not only delineates the mechanisms and optimality behind this phenomenon but also articulates the circumstances under which structured, low-rank codes predominate over maximal separation in high dimensions.

Modular Addition Task and Empirical Phenomenology

Modular addition tasks, such as computing $(a+b) \bmod P$ with $P$ prime, present a highly structured algebraic label space. Inputs $(a, b)$ are one-hot encoded and processed by a multilayer perceptron (MLP), while the model predicts $c = (a + b) \bmod P$ (Figure 1A). Notably, prior empirical studies have repeatedly observed that, instead of distributing representations across high-dimensional spaces as predicted by ETF, both embeddings and classifier weights converge into two-dimensional circular arrangements:

Figure 1: Task setup and cyclic representations for modular addition modulo 97. (A) Each input pair $(a,b)$ is one-hot encoded, concatenated, and processed by an MLP; the softmax output predicts $c=(a+b)\bmod 97$. (B) PCA of the learned token embeddings after automorphism-based label reordering. (C) PCA of the output weights. The matched circular geometry in panels (B) and (C) shows that the network learns the task-intrinsic cyclic code rather than a generic high-dimensional separator.

This configuration corresponds to learning the characters of the cyclic group $\mathbb{Z}/P\mathbb{Z}$, reflecting the algebraic symmetries of the task instead of maximizing mutual class separation. For $P = 97$, this entails compressing the representational geometry from the $(K-1)$0 dimensions of an ETF to a 2D cyclic structure.

Mechanistic and Theoretical Dissection

Hierarchical Emergence: Classifier-First Collapse

Contrary to a synchronous layerwise collapse, analysis reveals a hierarchical process:

1. Classifier Collapse: Classifier weights are rapidly compressed into a 2D plane, attaining an equiangular configuration optimal for cross-entropy in the angular regime. Dense gradients from cross-entropy efficiently drive this collapse.
2. Subspace Locking: Once the classifier's effective rank is capped at two, upstream embedding updates—delivered via backpropagated gradients—are constrained to this plane, with weight decay and implicit regularization suppressing components orthogonal to it.
3. Cyclic Alignment: The residual in-plane optimization becomes a problem of phase alignment, reducing further to soft assignment of embedding phases to class prototypes under an entropy-regularized angular cost, leading to embeddings distributed as equally-spaced points on the circle—precisely, characters of $(K-1)$1.

Mathematically, after collapse, the in-plane optimization is captured by an angular potential minimized uniquely by the regular $(K-1)$2-gon (Theorem: Equiangular optimality).

Task Symmetry and Regularization: Circumventing Neural Collapse

While a simplex ETF remains optimal where class identity is unstructured and the main objective is maximal separation, modular addition's cyclic group structure, in concert with explicit or implicit low-rank biases (e.g., weight decay, deep factorization), alters the comparison class of viable solutions. Theoretical analysis demonstrates:

• Cross-Entropy Trade-Off: The cyclic code incurs only an $(K-1)$3 disadvantage in cross-entropy loss relative to the ETF.
• Regularization Gain: The same cyclic code yields a $(K-1)$4 reduction in regularization cost (e.g., Schatten norm $(K-1)$5, nuclear norm) compared with high-rank ETFs.
• Phase Transition: There exists a sharp threshold in regularization coefficient, $(K-1)$6, above which the rank-2 cyclic solution becomes globally optimal.

This analysis refines the perceived universality of NC, highlighting that low-rank, task-structured solutions are not pathological but are, in fact, optimal in symmetry-dominated, regularized learning contexts.

Implications for Grokking and Representation Learning

This geometry offers a mechanistic explanation for grokking in modular tasks, where perfect generalization emerges abruptly after a long period of poor validation performance. The classifier's early collapse to a 2D subspace sets the stage for upstream embeddings to realign, and only after both alignments does full generalization occur. The learned cyclical geometry reflects the algebraic operations underlying the task, emphasizing that neural networks, when regularized appropriately and trained on structured tasks, materialize efficient algorithmic representations rather than high-dimensional lookup codes.

Notably, this work implies that in algorithmic or group-structured classification tasks:

• Downstream layers (classifier weights) can spontaneously set intrinsic dimensionality, eliding the necessity for high-dimensional separation.
• Regularization strength and label symmetries play a pivotal role in determining the emergent geometry.
• The adopted geometric code directly mirrors the generative algebraic structure of the task.

Theoretical and Practical Outlook

These results prompt a methodological shift in the study of learned geometries within neural networks. Analyses must consider not only class-cardinality-induced dimensionality expansions but also the task's group-theoretic symmetries and the implicit or explicit spectral biases introduced by the training dynamics. In practice, this work suggests that algorithmic tasks with rich symmetries are likely to induce highly organized, low-rank, task-specific codes in learned representations—assuming sufficient regularization.

Future research directions include:

• Generalization to broader classes of group-structured or algorithmic tasks and their induced representational manifolds.
• Empirical exploration of modular addition at higher $(K-1)$7 and in different neural architectures, including transformers and RNNs.
• Investigation of explicit architectural priors or penalties tuned to latent symmetries, augmenting the effect of stochastic or weight-based regularization.
• Analytical characterization of phase transitions between ETF and structured codes in mixed or partially symmetric tasks.

Conclusion

This study rigorously demonstrates that, for modular arithmetic, neural networks eschew the conventional $(K-1)$8-dimensional NC geometry in favor of a rank-2 cyclic structure driven by task-intrinsic symmetry and optimization's implicit low-rank bias. Hierarchical contraction across layers, regularization-induced preference for simplicity, and label space algebra combine to yield embeddings and classifier weights that encode group structure efficiently on a circle. This work refines theoretical understandings of representation learning, underscoring the primacy of task geometry and regularization in governing learned codes over class-count alone.