Sparse Prototype Code Underlies Classification and Prediction Across Modalities

Yehonatan Avidan; Daniel D. Lee; Haim Sompolinsky

cs.LGAug 16, 2026

Sparse Prototype Code Underlies Classification and Prediction Across Modalities

Authors: Yehonatan Avidan, Daniel D. Lee, Haim Sompolinsky

Organizations: The Racah Institute of Physics, Hebrew University · Edmond and Lily Safra Center for Brain Sciences, Hebrew University · Department of Electrical and Computer Engineering, Cornell Tech, Cornell University · Center for Brain Science, Harvard University

Abstract

Neural representations have become a central tool for studying the internal mechanisms of modern AI models, yet their complex high-dimensional structure makes them difficult to interpret. We show that classification tasks give rise to a universal representational geometry, shared across state-of-the-art models in vision, audio, and language processing. The key structure is that within-class variability is not random in representation space. Instead, its classifier-relevant component has strong and structured correlations with the class's own centroid and with the centroids of its competing classes. Building on this observation, we derive an analytical mean-field theory governed mainly by the variability along true-class and rival-class centroid coordinates, together with a global renormalization of the class radius that compensates for the non-Gaussian statistics of real representations. The theory accurately predicts classification accuracy across architectures and modalities. The relevant geometric quantities improve systematically with model scale, mirroring the observed gains in accuracy. A striking feature of the theory is its sparsity: accurate prediction requires only a small set of centroid coordinates associated with the true class and its strongest rivals - connecting our framework to sparse-feature extraction approaches such as sparse autoencoders. Together, these results provide a parsimonious predictive theory of neural representations and suggest that classification in deep networks is governed by a sparse, centroid-aligned structure embedded within the full high-dimensional representation space.

Figures & tables

Figure 1: Centroid-aligned variability, Qwen3-32B, TinyStories representations. The variability of each class is highly correlated with other class centroids. (a) To visualize it, we use Qwen3-32B representations ( Yang et al., 2025 ) on the dataset TinyStories ( Eldan and Li, 2023 ) . We project each example onto the three centroids with the largest correlations among the class variability. For 217/320 classes ( $68\%$ ), the true class centroid is among these top 3 (the nine shown here are of this sort). The rivals carry semantic meaning (e.g., the rivals of "cookie" are "food" and "delicious"; the rivals of "head" are "face" and "hand"). The mean prototype classification accuracy of these classes (see Def. 2.1) is $76.4\%$ . (b) To causally test the role of the variability-centroid correlations, we rotate the class manifold around its mean (denoted as $\delta$ -rotation), breaking these correlations while leaving the class radius and the distance between class centroids unchanged. The manipulation removes all classification errors, and the projections onto the top-3 directions shrink to $11.3\%$ of their original magnitude on average, and those directions no longer carry semantic meaning.

Figure 2: Geometric signatures of multiclass classification (a) Per-class accuracy vs. empirical radius $R$ : sharp decrease across the narrow empirical $R$ range and a wide spread at fixed $R$ . After random rotation of within-class variability $\delta\textbf{x}$ ( $\delta$ -rotation), the accuracy is perfect. (b) We artificially scale the variability $\delta\mathbf{x}\!\to\!(R^{\star}/R)\delta\mathbf{x}$ . The accuracy decreases fast but stabilizes at $~39\%$ accuracy for $R^{\star}\rightarrow\infty$ due to strong true centroid correlations (blue line). After random rotation of the variability $\delta$ -rotation, the accuracy is well described by the isotropic model ( $\mathrm{MSE}=1.26\times 10^{-3}$ , orange line). Additional breaking of centroid-centroid correlations (full rotation) has a very small effect; lines overlap.

Figure 3: Sparsity and heavy tail statistics: (a) K-way classification of rivals chosen by centroid projection variance

\sigma_{\mu}

converges quickly to the full

P

-way classification. (b) We define the strongest rival of each example

\mu^{*}(\mathbf{x})=\arg\max_{\mu\neq k}t_{\mu}(\mathbf{x})

, and count the number of examples captured by each rival. The strongest rival of each class captures more than half of the examples on average, with

90\%

of the examples captured by only 14.4 rivals. (c) The standardized centroid projections

s_{\mu}

(see Eq. 4 ) have heavy tail power law statistics, which inflate the extreme values that control classification accuracy. The magnified panel shows the log density of the far tail.

model	class	top-5 rivals (by $\sigma_{\mu}$ )
Qwen3-32B	happy	excite, proud, glad, surprised, curious
	bad	sad, sorry, scare, angry, tired
	voice	loud, noise, sound, big, funny
	smile	reply, nod, say, think, tell
	cake	delicious, cookie, food, ice, big
	forest	park, garden, grass, world, sky

Table 1: Rival classes selected by the projection’s variance

\sigma_{\mu}

are semantically related words (Qwen3-32B, TinyStories) or visually similar and co-occurring objects (DINOv3 ViT-7B, ImageNet). For more examples see also Fig. 1 a.

Notation	Name	Definition
$\mathbf{c}_{k}$	Centroid	$\mathbb{E}_{k}[\mathbf{x}]$
$\hat{\mathbf{c}}_{\mu}$	Unit centroid	$\mathbf{c}_{\mu}/\lVert\mathbf{c}_{\mu}\rVert$
$\delta\mathbf{x}$	Residual	$\mathbf{x}-\mathbf{c}_{k}$
$R^{2}$	Radius 2	$\mathbb{E}_{k}\!\big[\lVert\delta\mathbf{x}\rVert^{2}\big]\big/\lVert\mathbf{c}_{k}\rVert^{2}$
$\sigma_{\mu}^{2}$	Variance fraction	$\mathbb{E}_{k}\!\big[(\delta\mathbf{x}\cdot\hat{\mathbf{c}}_{\mu})^{2}\big]\big/\mathbb{E}_{k}\!\big[\lVert\delta\mathbf{x}\rVert^{2}\big]$
$s_{\mu}$	Std. projection	$\left(\delta\mathbf{x}\cdot\hat{\mathbf{c}}_{\mu}\right)/\mathrm{std}_{k}\left[\delta\mathbf{x}\cdot\hat{\mathbf{c}}_{\mu}\right]$

Table 2: Definitions.

\mathbb{E}_{k}

is defined as average over all samples

\mathbf{x}

belong to the manifold

k

Figure 4: Centroid-aligned theory matches per-class accuracy across modalities. (a, c) Per-class accuracy vs. empirical radius $R$ for DINOv3 ViT-7B on ImageNet ( $P{=}1000$ , $K{=}20$ ) and Qwen3-32B on TinyStories next-token prediction ( $P{=}320$ , $K{=}20$ ). Per-class scatter is presented with empirical binned mean. The $\lambda$ -rescaled theory ( $\lambda^{\star}{=}1.316$ , $1.119$ resp.) accounts for the empirical accuracy accurately. (b, d) Accurate per-class prediction of the accuracy. MSE ${=}1.83\times 10^{-3}$ (vision) and $4.43\times 10^{-3}$ (language), Pearson $r{=}0.951$ and $0.955$ . (e) Optimal $\lambda^{\star}$ vs. heavy tail measurement (1/ $\nu$ ) on representations for 15 vision backbones, 14 autoregressive language models, 4 masked language models, and 3 audio models. Overall Pearson correlations between $\lambda$ and $1/\nu$ are $r{=}0.90$ , confirming that $\lambda$ is a byproduct of the heavy-tail shape of the projection variables.

Backbone	$P$	$N$	$R$ range	$\langle\mathrm{Acc}\rangle$	Fixed $R$ fluctuations	$\delta$ -rot	$R^{\star}\rightarrow\infty$	$\langle K_{90\%}\rangle$	$\nu$
Vision – ImageNet-1K
DINOv3 ViT-7B/16	1000	4096	[0.4, 4.3]	0.855	0.099	1.000	0.392	14.4	4.5
DINOv3 ConvNeXt-L	1000	1536	[0.2, 3.1]	0.814	0.100	1.000	0.312	14.2	4.4
CLIP ViT-L/14	1000	768	[0.6, 2.7]	0.792	0.124	0.999	0.077	13.9	16.4
ViT-B/16 (IN21K)	1000	768	[0.5, 2.7]	0.853	0.088	1.000	0.265	14.1	4.1
Audio – VGGSound

Table 3: Universality of the centroid-aligned geometry. For each backbone, we report: the number of classes

P

; feature dimension

N

; per-class radius range; per-class mean accuracy

\langle\mathrm{Acc}\rangle

and its fluctuation around the mean accuracy at a fixed

R

; mean accuracy after

\delta

-rotation; saturation accuracy at

R^{\star}\rightarrow\infty

(controlled only by the residual

\delta{\bf x}

);

K_{90\%}

(mean rivals covering

90\%

of examples); right-tail Student-

t

tail index

\nu

(

K{=}20

Model	$\langle R\rangle$	$\sigma_{k}/\sigma^{*}_{\mu}$	$\sigma^{*}_{\mu}$	$g^{*}_{\mu}$	$\rho_{\mu}^{*}$	$\nu$	Acc
DINOv3 ViT-S	1.68	1.06	0.172	0.574	0.558	19.3	0.706
DINOv3 ViT-S $+$	1.60	1.09	0.172	0.555	0.535	14.4	0.738
DINOv3 ViT-B	1.52	1.19	0.163	0.542	0.529	9.0	0.794
DINOv3 ViT-L	1.39	1.36	0.147	0.447	0.322	7.0	0.845
DINOv3 ViT-H $+$	1.29	1.59	0.164	0.236	-0.011	5.3	0.855
DINOv3 ViT-7B	1.50	2.01	0.120	0.289	0.038	4.5	0.855

Table 4: Geometric measures across model scales: DINOv3 ViT family on ImageNet-1K (top block) and the Qwen3 language-model family on TinyStories (bottom block). The geometric quantities are defined in Sec. 5 and Table 2 ; the superscript ∗ flags the hardest competitor of a class, identified by

\sigma_{\mu}

. Best measure within each family marked in bold.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Centroid vs trained linear-readout accuracy, per backbone. The gap between methods is small (see also Table 5 ).

Vision (ImageNet-1K)
Backbone	$r$	$\Delta$ Acc
ViT-B/16 (IN21K)	$0.937$	$-0.005$
ConvNeXt-B (IN1K)	$0.893$	$-0.001$
DINOv3 ViT-L/16	$0.885$	$+0.028$
DINOv3 ViT-H + /16	$0.873$	$+0.028$
DINOv3 ConvNeXt-S	$0.872$	$+0.058$

Appendix

Table 5: Per-class Pearson correlation

r

between trained linear-readout and prototype-classifier accuracy, together with the mean accuracy gap

\Delta\mathrm{Acc}

(readout minus prototype) on the held-out split. The two classifiers align closely, with the exception of masked language models (see Sec. A.1 ).

Family	Model	$N$	Data / task	$P$	Objective
Vision — ImageNet-1K
Vision	DINOv3 ViT-S/16	$384$	ImageNet-1K	$1000$	Self-supervised
Vision	DINOv3 ViT-S+/16	$384$	ImageNet-1K	$1000$	Self-supervised
Vision	DINOv3 ViT-B/16	$768$	ImageNet-1K	$1000$	Self-supervised
Vision	DINOv3 ViT-L/16	$1024$	ImageNet-1K	$1000$	Self-supervised
Vision	DINOv3 ViT-H+/16	$1280$	ImageNet-1K	$1000$	Self-supervised

Appendix

Table 6: Model roster (

36

backbones).

N

is the feature dimension of the extracted representation;

P

is the number of classes in the task.

Figure 6: Validation of the isotropic model’s theory on synthetic data. (a) Class-averaged accuracy versus

R/\sqrt{N}

: simulation (markers, mean

\pm

between-class std) against Theorem 3.3 (line),

N=512

P=500

. (b) Self-averaging: quenched per-class accuracy (one shared centroid bank) versus the fresh-rival accuracy (rivals redrawn per example), the line is

y=x

r=0.985

, confirming self-averaging - the thermal average is equivalent to the quenched average. (c) Across-realization variance of the accuracy versus

P

at fixed radius

R=9.05

, fitted power law

\propto P^{-1.00}

, predicted by Proposition C.1.5. (d) Variance versus radius

R

(swept via the dimension at fixed

R/\sqrt{N}

), fitted power law

\propto R^{-1.96}

, predicted by Proposition C.1.5.

Figure 7: Synthetic validation of the centroid-aligned model. Examples are generated according to the centroid-aligned variability model (Sec. 5 and Sec. C.2 ) with

K=20

rivals,

N=200

P=500

. The geometry is calibrated to DINOv3 ViT-7B: own-class projected std

\sigma_{k}=0.24

, rival stds

\sigma_{\mu}=0.12\,\mu^{-0.4}

. (a) Class-averaged accuracy versus class radius

R

: simulation (markers) against Theorem 5.5 (line). The dotted line is the

R\to\infty

plateau

0.316

, which is set by

\sigma_{k}/\sigma_{\mu}

alone. (b) Quenched centroid versus fresh per-example sampling of the centroids,

y=x

line,

r=0.993

, validating self-averaging. (c) Across-realization variance versus

P

at fixed radius, fitted

\propto P^{-0.97}

, similar to the isotropic model (Proposition C.1.5).

Selector	within $5\%$	within $3\%$	within $2\%$	within $1\%$
Centroid correlations ( $g_{\mu}$ )	14	24	36	87
Projection variance ( $\sigma_{\mu}$ )	14	20	30	54
Theory margin	11	18	25	49
Gaussian margin	10	19	28	73

Appendix

Table 7: Rivals needed to reach the full accuracy. Median number of rivals

K

(over

36

backbones) at which the restricted

K

-way accuracy comes within

5/3/2/1\%

of the full

P

-way accuracy, per selector. Smaller is better.

Modality	$n$	Centroids corr.	Proj. variance	Theory margin	Gaussian margin
Vision	15	76	47	45	56
Autoregressive LM	14	102	89	82	108
Masked LM	4	48	42	35	43
Audio	3	60	39	37	49

Appendix

Table 8: Rivals needed within

1\%

, by modality. Median

K

to reach within

1\%

of the full accuracy, per selector and modality (

n

backbones).

Modality	$n$	RMSE $K{=}10$	RMSE $K{=}20$	floor	within $0.003$ by $K{=}20$
Vision	15	0.049	0.047	0.045	13/15
Audio	3	0.081	0.076	0.064	0/3
Masked LM	4	0.064	0.062	0.061	4/4
Autoregressive LM	14	0.080	0.075	0.075	14/14
All	36	0.065	0.062	0.060	31/36

Appendix

Table 9: The theory’s error saturates by

K\approx 10

–

20

. Mean per-class accuracy RMSE of the theory at

K=10

and

K=20

, its large-

K

floor, and the number of backbones already within

0.003

of that floor (large

K

) at

K=20

, by modality.

Figure 8: Standardized centroid projections have a heavy right tail that varies across models. Density of the standardized projection

s

(filled) with its fitted half-Student-t right tail

p(s)\propto(a^{2}+s^{2})^{-\nu/2}

(solid),

K=20

. Panels run light

\to

heavy tail (left to right, top to bottom), with the fitted exponent

\nu

and model in each title. Color encodes modality (vision blue, autoregressive LM orange, masked LM green). As

\nu

decreases, the peak sharpens, and the right tail fattens.

Modality	$n$	$\langle\mathrm{Acc}^{\mathrm{}}\rangle$	$\langle\mathrm{Acc}^{\mathrm{th}}\rangle$	RMSE	$\langle\mathrm{Acc}^{\mathrm{th}}\rangle$	RMSE	$\langle\lambda\rangle$
			$\lambda=1$		fitted $\lambda$
Vision	15	0.812	0.876	0.080	0.819	0.047	1.228
Audio	3	0.613	0.657	0.087	0.617	0.076	1.084
Masked LM	4	0.715	0.785	0.090	0.724	0.062	1.131
Autoregressive LM	14	0.522	0.583	0.096	0.524	0.075	1.111
All	36	0.672	0.734	0.088	0.677	0.062	1.160

Appendix

Table 10: Without global radius rescaling, the theory is over-optimistic. Per-class accuracy predicted by the centroid-aligned theory of Sec. 5.2 , evaluated at

\lambda=1

(no rescaling of the variability) and at the fitted

\lambda

, compared with the empirical accuracy. The theoretical accuracy with

\lambda=1

is higher than the empirical one, while the theory with

\lambda

is within

1\%

of the empirical one, with significantly better RMSE.

Vision (ImageNet-1K)
Backbone	$\nu$	$\lambda$
DINOv2 ViT-g	$3.7$	$1.272$
ConvNeXt-B	$4.0$	$1.254$
ViT-B	$4.1$	$1.344$
DINOv3 CNX-L	$4.4$	$1.305$
DINOv3 ViT-7B	$4.5$	$1.316$

Appendix

Table 11: Right-tail exponent

\nu

across all backbones.

\nu

is the power-law exponent of the fitted right tail

p(s)\propto(a^{2}+s^{2})^{-\nu/2}

; a smaller

\nu

means a heavier tail.

\lambda

is the global radius rescaling of the theory. Grouped by modality, sorted from heaviest to lightest tail. Pearson correlations between

\lambda

and

1/\nu

are high in every family (

r=0.89,0.95,0.81,0.89

in vision, audio, masked language models and autoregressive models, respectively) and globally over all models (

r=0.90

Vision (ImageNet-1K)
Backbone	$P$	$N$	Acc	Acc th	$r$	$\lambda$	RMSE
ConvNeXt-B	1000	1024	$0.943$	$0.952$	$0.954$	$1.254$	$0.0256$
ResNet50	1000	2048	$0.861$	$0.871$	$0.936$	$1.228$	$0.0454$
DINOv3 ViT-H +	1000	1280	$0.856$	$0.862$	$0.960$	$1.297$	$0.0390$
DINOv3 ViT-7B	1000	4096	$0.855$	$0.861$	$0.951$	$1.316$	$0.0427$
ViT-B	1000	768	$0.853$	$0.861$	$0.959$	$1.344$	$0.0378$

Appendix

Table 12: Centroid-aligned theory : The theory was tested on all 36 models. For each backbone, we report: number of classes

P

, feature dimension

N

, empirical prototype accuracy, per-class Pearson

r

between predicted and empirical accuracy, the fitted global rescaling

\lambda

, and the per-class prediction error RMSE. The theory consistently predicts the accuracy well across all models (

r>0.93

Figure 9: The best prototype layer is a few blocks before the final one for language, but is the final block for vision. Per-layer prototype classification accuracy versus fractional depth. Color runs from light (small) to dark (large) within each modality, and the peak performance is marked. (a) The six Qwen3 autoregressive language models; the optimal accuracy is a few blocks before the final layer, with a big drop at the final layer in most models. (b) the six DINOv3 vision transformers. The final layer is the best, with a plateau on DINOv3 7B.

Layer	$\langle R\rangle$	$\sigma_{k}/\sigma^{*}_{\mu}$	$\sigma^{*}_{\mu}$	$g^{*}_{\mu}$	$\rho^{*}_{\mu}$	$\nu$	Acc
L58 (0.91)	2.40	0.993	0.172	0.514	0.496	10.9	0.521
L59 (0.92)	2.32	0.996	0.175	0.511	0.502	11.0	0.542
L60 (0.94)	2.24	0.999	0.175	0.510	0.474	10.9	0.557
L61 (0.95)	2.20	1.007	0.175	0.503	0.451	11.4	0.566
L62 (0.97)	2.17	0.957	0.181	0.517	0.463	13.7	0.553
L63 (0.98)	2.16	0.894	0.191	0.521	0.448	18.1	0.528

Appendix

Table 13: Qwen3-32B across depth. Class-averaged geometric measures for the seven last blocks of Qwen3-32B (

64

blocks). ∗ marks each class’s hardest competitor

\mu^{\star}

, measured by

\sigma_{\mu}

, the projection standard deviation.

\langle R\rangle

is the mean class radius;

\sigma_{k}/\sigma^{*}_{\mu}

is the ratio between the standard deviation of the projection onto the true-class and hardest-rival directions.

\sigma^{*}_{\mu}

is the hardest rival projected std,

g^{*}_{\mu}=\hat{\mathbf{c}}_{k}\!\cdot\!\hat{\mathbf{c}}_{\mu^{\star}}

the centroid overlap and

\rho^{*}_{\mu}

the correlation between projections.

\nu

is the right-tail exponent. Best measure highlighted in bold.

Figure 10: 15 vision models. For each backbone: (left) per-class accuracy versus empirical radius $R$ , empirical binned mean and theory binned mean nearly overlap in most models; (right) per-class predicted versus empirical accuracy with $r$ , RMSE and $\lambda$ . The dashed line is $y=x$ . Mean Pearson $r=0.951$ , mean RMSE $=0.047$ .

Figure 11: Masked language models (4 models, TinyStories content-word prediction, best layer). Panels as in Fig. 10 . Mean $r=0.943$ , mean RMSE $=0.062$ .

Figure 12: Autoregressive language models (14 models, TinyStories content-word prediction, best layer). Panels as in Fig. 10 ; equally spaced $R$ bins fixed for this modality. Mean $r=0.946$ , mean RMSE $=0.075$ .

Figure 13: Audio (3 encoders, VGGSound training split, final layer). Panels as in Fig. 10 . Mean $r=0.948$ , mean RMSE $=0.076$ .

Vision (ImageNet-1K)
	Uncorrelated		Main theory		Cluster expansion
Backbone	$\lambda$	RMSE	$\lambda$	RMSE	$\lambda$	RMSE
ConvNeXt-B	$0.979$	$0.0493$	$1.254$	$0.0256$	$1.281$	$0.0234$
ResNet50	$1.114$	$0.0589$	$1.228$	$0.0454$	$1.320$	$0.0434$
DINOv3 ViT-H +	$1.212$	$0.0578$	$1.297$	$0.0390$	$1.317$	$0.0377$
DINOv3 ViT-7B	$1.248$	$0.0574$	$1.316$	$0.0427$	$1.337$	$0.0402$

Appendix

Table 14: Fitted global rescaling

\lambda

and per-class prediction error RMSE for three predictors at

K{=}20

, on all 36 backbones: the uncorrelated-projections baseline (

\rho_{\mu}{=}0

), the main theory (rank-one true–rival correlation), and the first-order cluster expansion.

\lambda

is refitted separately for each predictor.

Neural Representation Geometry

Neural representation geometry studies how learned features are arranged and how that structure relates to model behavior and performance.

See all 256 papers

Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic
Jun 8, 2026Hu Tan, Kuo Gai, Shihua ZhangNeural Representation Geometry Classification
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
Anchoring the Eigengap: Cross-Modal Spectral Stabilization for Sample-Efficient Representation Learning
May 9, 2026Nikhil J. Dhinagar, Vidhi Chhatbar, Chirag Jagad +6Representation Learning Spectral Methods
Imaging Genetics Center, Mark & Mary Stevens Neuroimaging & Informatics Institute, Keck School of Medicine, University of Southern California, Los Angeles, CA, USA · Neuroscience Graduate Program, Mark & Mary Stevens Neuroimaging & Informatics Institute, Chan Division of Occupational Science & Occupational Therapy, Biomedical Engineering, University of Southern California, Los Angeles, CA, USA
Directional Linear Separability of Neural Representations: Geometry and Transformations
Jun 7, 2026Yi Wei, Xuan Qi, Suorong Yang +1Neural Representation Geometry Linear Classification
DITEN University of Genoa, Genoa, Italy · State Key Laboratory of Novel Software Technology School of Artificial Intelligence Nanjing University, Jiangsu, China

Jun 8, 2026J/K move · Enter open · S save

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

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

Neural Representation Geometry Classification Modular Addition Neural Collapse Neural Network Training Dynamics Modular Arithmetic+1

Open paper AskPDF 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

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.

Abstract

Neural Representation Geometry

Explore similar work

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

Anchoring the Eigengap: Cross-Modal Spectral Stabilization for Sample-Efficient Representation Learning

Directional Linear Separability of Neural Representations: Geometry and Transformations