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

arXiv: https://arxiv.org/abs/2606.08985 · HTML: https://mlsanity.com/paper/2606.08985/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 31 references that are in the map.
- First posted on arXiv 8 Jun 2026.

## Builds on (most cited first, 24 of 24)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks](https://mlsanity.com/paper/2306.17844/) (2023)
- [Extended Unconstrained Features Model for Exploring Deep Neural Collapse](https://mlsanity.com/paper/2202.08087/) (2022)
- [On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features](https://mlsanity.com/paper/2203.01238/) (2022)
- [Grokking modular arithmetic](https://mlsanity.com/paper/2301.02679/) (2023)
- [Neural Collapse in Deep Linear Networks: From Balanced to Imbalanced Data](https://mlsanity.com/paper/2301.00437/) (2023)
- [Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition](https://mlsanity.com/paper/2407.12332/) (2024)
- [Understanding Deep Representation Learning via Layerwise Feature Compression and Discrimination](https://mlsanity.com/paper/2311.02960/) (2023)
- [Implicit bias of SGD in L2-regularized linear DNNs: One-way jumps from high to low rank](https://mlsanity.com/paper/2305.16038/) (2023)
- [Average gradient outer product as a mechanism for deep neural collapse](https://mlsanity.com/paper/2402.13728/) (2024)
- [Bottleneck Structure in Learned Features: Low-Dimension vs Regularity Tradeoff](https://mlsanity.com/paper/2305.19008/) (2023)
- [Progressive Feedforward Collapse of ResNet Training](https://mlsanity.com/paper/2405.00985/) (2024)
- [Provable Emergence of Deep Neural Collapse and Low-Rank Bias in L2-Regularized Nonlinear Networks](https://mlsanity.com/paper/2402.03991/) (2024)
- [Mechanistic Insights into Grokking from the Embedding Layer](https://mlsanity.com/paper/2505.15624/) (2025)
- [Clustering and Alignment: Understanding the Training Dynamics in Modular Addition](https://mlsanity.com/paper/2408.09414/) (2024)
- [OTAD: An Optimal Transport-Induced Robust Model for Agnostic Adversarial Attack](https://mlsanity.com/paper/2408.00329/) (2024)
- [Deciphering Two Training Clocks in Grokking via Deep Linear Network Theory with Conditional ReLU Reduction](https://mlsanity.com/paper/2606.05863/) (2026)

## Built on by (most cited first, 2 of 2)

- [Structure-Specific Representational Priors Causally Control the Grokking Delay](https://mlsanity.com/paper/2607.04333/) (2026)
- [Symmetry without a manifold: intrinsic dimension on orbits](https://mlsanity.com/paper/2609.17926/) (2026)

## Abstract (arXiv)

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 S1; combined with modular-addition labels, this reduces embedding formation to phase alignment, whose minimizers are single-frequency characters of mathbbZ/PmathbbZ 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 Theta(K) advantage under Schatten or weight-decay surrogates, yielding a critical threshold λcrit = Theta(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.