Jun 8, 2026
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
Z/PZ 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
λ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.
Hu Tan, Kuo Gai, Shihua Zhang
Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; and School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China · Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS), Shanghai, China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China; and Key Laboratory of Systems Health Science of Zhejiang Province, School of Life Science, Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Chinese Academy of Sciences, Hangzhou 310024, China