Representations Have Trajectories: Characterizing Temporal Representation Dynamics in Neural Networks

Shivaprasad B. Gowda; Sparsh Mittal

cs.LGOct 17, 2026

Representations Have Trajectories: Characterizing Temporal Representation Dynamics in Neural Networks

Authors: Shivaprasad B. Gowda, Sparsh Mittal

Organizations: Indian Institute of Information Technology, Nagpur · Indian Institute of Technology, Roorkee

Abstract

Recent work has shown that deep neural networks exhibit an implicit low-rank simplicity bias. In this paper, we investigate the questions that naturally arise from this observation. What are the temporal consequences of low-rank simplicity bias during optimization? How does representation geometry evolve throughout optimization rather than only at convergence? Do different neural architectures exhibit distinct geometric optimization behaviors? To answer these questions, we conduct a large-scale empirical study tracking representational geometry throughout training. We find that neural networks follow three recurring geometric trajectories: persistent high-rank, collapse-and-recovery, and persistent low-rank. We further identify Destructive Representation Collapse (DRC) as a pathological regime in which representations collapse to rank-one and fail to recover, leading to a sustained loss of representational capacity. While collapse-and-recovery is consistently associated with successful optimization in controlled settings, experiments on real-world models show that the relationship between geometric trajectories and optimization outcomes is configuration dependent. Overall, our findings suggest that representation geometry evolves dynamically during training and provide a new perspective on how deep neural networks learn.

Figures & tables

Figure 1: Archetypal representation geometry trajectories: (a) persistent high rank, (b) collapse-and-recovery, and (c) persistent low rank. The x-axis uses a symlog scale to resolve early dynamics.

Archetype	Criterion
A1 (Persistent High-Rank)	$\text{SR}_{t}>1.5$ at every logged epoch t
A2 (Transient Collapse and Recovery)	$(\text{SR}_{\min}\leq 1.5)\;\land\;(\text{SR}_{\mathrm{final}}>1.5)\;\land\;(\text{SR}_{\mathrm{final}}-\text{SR}_{\min}>0.1)$
A3 (Persistent Low-Rank)	$(\text{SR}_{\min}\leq 1.5)\;\land\;\left[\text{SR}_{\mathrm{final}}\leq 1.5\;\lor\;(\text{SR}_{\mathrm{final}}-\text{SR}_{\min}\leq 0.1)\right]\;\land\;$
	$(\text{ER}_{\mathrm{final}}>1.5)$
DRC (Destructive Representation Collapse)	$(\text{SR}_{\min}\leq 1.05)\;\land\;(\text{SR}_{\mathrm{final}}\leq 1.05)\;\land\;$ $(\text{ER}_{\mathrm{final}}\leq 1.5)$

Table 1: Rank dynamics archetypes . Here,

s_{t}

denotes stable rank at epoch

t

s_{\min}=

min over logged epochs

t\in{1,10,50,100,\dots,T}

s_{t}

. The same notation holds for effective rank

e_{t}

as well.

Metric	$\text{Loss}<0.05$	$\text{Loss}<10^{-2}$	$\text{Loss}<10^{-3}$
Contingency $(a,b,c,d)$	$(164,23,1517,2380)$	$(156,31,1132,2765)$	$(155,32,1051,2846)$
Success Rate (A2 vs. No A2)	$87.70\%$ vs. $38.93\%$	$83.42\%$ vs. $29.05\%$	$82.89\%$ vs. $26.97\%$
RR [95% CI]	$2.25$ [ $2.11,2.41$ ]	$2.87$ [ $2.65,3.11$ ]	$3.07$ [ $2.83,3.34$ ]
Unadjusted OR [95% CI]	$10.98$ [ $7.09,17.00$ ]	$12.13$ [ $8.22,17.90$ ]	$12.95$ [ $8.82,19.03$ ]
Fisher $p$ -value	$1.07\times 10^{-41}$	$4.74\times 10^{-51}$	$2.67\times 10^{-54}$
Adjusted OR [95% CI]	6.77 [ $3.92,11.70$ ] $p<0.0001$	4.41 [ $2.40,8.08$ ] $p<0.0001$	4.59 [ $2.39,8.81$ ] $p<0.0001$

Table 2: Association between the recovery trajectory (A2) and optimization success at three loss thresholds (

M_{\mathrm{models}}=4{,}084

\text{seed }42

). Contingency counts, RR (Relative Risk) and the unadjusted OR (Odd’s Ratio) compare A2 with all other runs (A1+A3+DRC); the unadjusted OR uses the Haldane–Anscombe correction ((a+0.5)(d+0.5))/((b+0.5)(c+0.5)) with Fisher’s exact p. The adjusted OR is the A2 coefficient of the multivariate logistic regression (adjusting for activation, architecture, depth, task rank and optimizer; not initialization) and therefore compares A2 with A1 only. a: A2+success; b: A2+failure; c: non-A2+success; d: non-A2+failure.

Factor / Variable	Coef ( $\beta$ )	SE	Adjusted OR [95% CI]	$p$ -value
Intercept	$+2.0358$	$0.4662$	$7.66$ [ $3.07,19.10$ ]	$<0.0001$
Archetype A2	$+1.5229$	$0.3334$	4.59 [ $2.39,8.81$ ]	$<0.0001$
Archetype A3	$-1.2377$	$0.1862$	0.29 [ $0.20,0.42$ ]	$<0.0001$
DRC	$-5.2012$	$0.3896$	0.0055 [ $0.0026,0.0118$ ]	$<10^{-20}$
Residual MLP	$+1.6988$	$0.3917$	5.47 [ $2.54,11.78$ ]	$<0.0001$
Deep Linear MLP	$+0.4347$	$0.3881$	1.54 [ $0.72,3.30$ ]	$0.2627$

Table 3: Multivariate logistic regression of factors influencing optimization success (

\mathcal{L}<10^{-3}

M_{\mathrm{models}}=4{,}084

). Reference baselines: Archetype = A1, Activation = GELU, Architecture = Plain MLP, Optimizer = Adam, Seed = 42, Initialization Scheme is not a covariate. Archetype is defined by the trajectory, so the coefficients of the other factors are conditional on the archetype.

Figure 2: Multivariate Forest Plot: Controlled Predictors of Optimization Success.

Figure 3: Empirical separation of DRC from moderate low-rank collapse (A3).

Figure 4: Layer-wise stable rank dynamics: deeper layers undergo DRC without recovery.

Figure 5: Active learning in a successful model (A) versus vanished layerwise gradients (

\|g\|=0.0000

) in an un-converged model (B).

Appendix figures & tables30 assets

Supplementary material from the paper’s appendix.

Appendix

Seed	A1 (%)	A2 (%)	A3 (%)	DRC (%)	A2 Succ. (%)	A3 Succ. (%)	DRC Succ. (%)
42	$31.70\%$	$5.42\%$	$21.08\%$	$41.80\%$	$80.49\%$	$32.43\%$	$1.48\%$
43	$33.55\%$	$4.72\%$	$20.50\%$	$41.23\%$	$88.79\%$	$30.97\%$	$1.28\%$
44	$33.99\%$	$5.34\%$	$18.21\%$	$42.46\%$	$81.82\%$	$34.87\%$	$1.56\%$
Mean $\pm$ SD	$33.08\pm 1.22\%$	$5.16\pm 0.38\%$	$19.93\pm 1.52\%$	$41.83\pm 0.62\%$	$83.70\pm 4.46\%$	$32.75\pm 1.97\%$	$1.44\pm 0.14\%$

Appendix

Table S.1: Trajectory archetype distribution and optimization success across three random initialization seeds. Each seed evaluates the same

N=2{,}268

configurations, yielding

6{,}804

fully trained runs in total. Success is defined as

\mathcal{L}_{\mathrm{train}}<0.01

Type	$M_{\mathrm{models}}$	%	Mean Loss	Median Loss	Mean SR
A1	1,746	42.75	0.7127	0.0887	7.8135
A2	187	4.58	0.0580	0.00001	2.7860
A3	1,045	25.59	0.1846	0.0223	1.1748
DRC	1,106	27.08	$1.05\times 10^{15\dagger}$	0.3676	1.00

Appendix

Table S.2: Summary statistics across trajectory archetypes.

	$<10^{-2}$	$<10^{-3}$
A1	37.80%	35.11%
A2	83.42%	82.89%
A3	43.73%	40.96%
DRC	1.36%	0.90%

Appendix

Table S.3: Success probability conditioned on trajectory archetype. The probability of successful optimization varies substantially across A1–A3 and DRC.

Architecture	Total	DRC	DRC Rate	Diverged	Frozen
Deep Linear	252	152	60.32%	15	137
Plain MLP	1,788	780	43.62%	39	741
Residual MLP	2,044	174	8.51%	98	76
Total Synthetic Pool	4,084	1,106	27.08%	152	954

Appendix

Table S.4: Architecture class-wise DRC breakdown. Frozen = DRC runs that did not diverge

(L\leq 100)

, this includes runs on low-rank teachers that reach low absolute loss.

Init	A1	A2	A3	DRC	Total
Default (PyTorch Kaiming-uniform; Table S.14 )	1,649 (48.5%)	215 (6.3%)	542 (15.9%)	996 (29.3%)	3,402
Low-rank ( $r=1$ )	602 (17.7%)	136 (4.0%)	814 (23.9%)	1,850 (54.4%)	3,402
Total	2,251 (33.1%)	351 (5.2%)	1,356 (19.9%)	2,846 (41.8%)	6,804

Appendix

Table S.5: The 2,268-configuration grid is balanced across initialization schemes (1,134 configurations each per seed). Counts are summed over seeds 42, 43, 44.

Target	Runs	Mean SR	Median SR	Mean ER	Median ER
Rank	( $N$ )
$r=1$	378	$2.7229\pm 3.997$	1.0003	$8.8891\pm 11.482$	1.1900
$r=4$	378	$2.7630\pm 4.010$	1.0004	$8.9129\pm 11.451$	1.2003
$r=16$	378	$2.8345\pm 4.110$	1.0003	$8.9942\pm 11.449$	1.2088
$r=24$	378	$2.7768\pm 3.960$	1.0004	$9.0070\pm 11.452$	1.2312
$r=32$	378	$2.7448\pm 3.966$	1.0004	$8.9510\pm 11.414$	1.1969

Appendix

Table S.6: Statistical Analysis of the Influence of Task Rank on Representation Metrics. SR and ER denote Stable Rank and Effective Rank, respectively.

Figure S.1: Mean min-observed representation rank remains uniform across target task ranks $r\in\{1,\ldots,64\}$ , confirming that early-layer rank collapse is governed by architectural depth and activation geometry rather than target task rank.

Hyperparameter	Synthetic Sweeps ( $N=4{,}084$ )	CIFAR-10 Plain CNN sweep ( $N=192$ )
Random Seed	42	42
Learning Rate ( $\eta$ )	0.001	0.001
Weight Decay ( $\lambda$ )	0.0	0.0001
Batch Size ( $B$ )	Full-batch (32)	128
Optimizers	Adam, Momentum SGD, SGD	Adam, SGD
Loss Criterion	Mean Squared Error (MSE)	Cross-Entropy Loss

Appendix

Table S.8: Common hyperparameter specifications across synthetic and vision experiment sweeps.

Model	Mean AUC	SD	95% CI	Acc.
1 Static Final-Rank Endpoint	0.7743	0.0193	[0.7574, 0.7912]	69.20%
2 Hyperparameter Controls Only	0.9459	0.0073	[0.9395, 0.9523]	90.77%
3 Full Trajectory + Controls	0.9660	0.0048	[0.9618, 0.9702]	92.58%

Appendix

Table S.9: 5-fold stratified cross-validation performance (

M_{\mathrm{models}}=4{,}084

). SD=Standard deviation

Figure S.8: 5-Fold Cross-Validation ROC Curve Comparison

Architecture	Depth ( $L$ )	$N$	Train Acc (%)	Test Acc (%)	Deepest Block SR	Deepest Block ER	Penult SR	Catastrophic DRC
Standard ViT (Pre-LN + Residual)	3	8	$87.0\pm 8.7$	$76.2\pm 3.6$	$6.50\pm 1.85$	$72.7\pm 23.8$	$5.19\pm 0.49$	0 / 8 (0.0%)
Standard ViT (Pre-LN + Residual)	6	8	$83.2\pm 17.5$	$72.3\pm 10.6$	$6.27\pm 2.00$	$71.2\pm 22.8$	$5.08\pm 1.09$	0 / 8 (0.0%)
Standard ViT (Pre-LN + Residual)	9	8	$91.9\pm 8.6$	$77.2\pm 3.6$	$7.25\pm 1.35$	$75.1\pm 25.4$	$5.79\pm 0.78$	0 / 8 (0.0%)
No-Residual ViT (Pre-LN Only)	3	8	$74.3\pm 6.8$	$68.0\pm 4.8$	$3.53\pm 0.51$	$29.7\pm 26.5$	$3.58\pm 0.59$	0 / 8 (0.0%)
No-Residual ViT (Pre-LN Only)	6	8	$66.6\pm 8.8$	$62.8\pm 7.8$	$2.94\pm 0.62$	$22.9\pm 19.7$	$3.29\pm 0.27$	0 / 8 (0.0%)
No-Residual ViT (Pre-LN Only)	9	8	$52.8\pm 27.0$	$49.2\pm 24.6$	$2.46\pm 1.04$	$13.5\pm 11.6$	$2.62\pm 1.07$	2 / 8 (25.0%)

Appendix

Table S.10: Depth-Wise Representation & Performance Scaling (

N=48

Models, 200 Epochs)

Layer	Standard ViT SR	Standard ViT ER	No-Residual ViT SR	No-Residual ViT ER
Block 1	$5.95\pm 1.19$	$57.7\pm 13.8$	$3.49\pm 2.39$	$39.7\pm 27.1$
Block 2	$7.08\pm 1.40$	$68.8\pm 20.2$	$3.00\pm 2.35$	$29.4\pm 31.3$
Block 3	$7.40\pm 1.92$	$73.3\pm 22.8$	$2.66\pm 1.90$	$21.6\pm 26.8$
Block 4	$8.10\pm 1.68$	$75.4\pm 23.4$	$2.55\pm 1.65$	$15.7\pm 18.8$
Block 5	$8.34\pm 1.94$	$76.6\pm 23.8$	$2.68\pm 1.75$	$12.4\pm 13.9$
Block 6	$8.82\pm 2.43$	$77.2\pm 23.9$	$2.78\pm 1.81$	$11.5\pm 12.2$

Appendix

Table S.11: Layer-by-Layer Metrics across Depth at

L=9

(

N=16

Runs, Epoch 200)

Training Milestone	Arch 1: Never Collapsed (Success %)	Arch 2: Early Recovery (Success %)	Arch 3: Early Stuck (Success %)	Odds Ratio (Arch 2 vs 3)	Mean Final Accuracy (A2 vs A3)
10% of Training (Epoch 20)	31 ( $90.3\%$ )	14 ( $78.6\%$ )	3 ( $0.0\%$ )	$23.00\times$	$64.4\%$ vs $25.9\%$
25% of Training (Epoch 50)	31 ( $90.3\%$ )	15 ( $73.3\%$ )	2 ( $0.0\%$ )	$12.78\times$	$64.0\%$ vs $10.0\%$
50% of Training (Epoch 100)	31 ( $90.3\%$ )	15 ( $73.3\%$ )	2 ( $0.0\%$ )	$12.78\times$	$64.0\%$ vs $10.0\%$

Appendix

Table S.12: Early Trajectory Archetype Forecasting in Vision Transformers

Optimizer	Description	Default Hyperparameters
SGD	Plain gradient descent	—
SGD + Momentum	Momentum-based GD	$\beta=0.9$
Adam	Adaptive moment estimation	$\beta_{1}=0.9$ , $\beta_{2}=0.999$ , $\epsilon=10^{-8}$

Appendix

Table S.13: Optimization algorithms used in the experimental study. All optimizers use learning rate

\eta=10^{-3}

and weight decay

\lambda=0

(synthetic),

10^{-4}

(CIFAR-10 CNN)

Scheme	Initialization	Key Parameter
Default	PyTorch Kaiming Uniform	–
Low-Rank	$W=\dfrac{AB}{\sqrt{rd_{\mathrm{in}}}},\;A,B\sim\mathcal{N}(0,1)$	Rank $r$ ( $r=1$ )

Appendix

Table S.14: Initialization schemes evaluated.

Activation	Definition
Linear	$\Phi(x)=x$
ReLU	$\Phi(x)=\max(0,x)$
Leaky ReLU	$\Phi(x)=\max(0.01x,x)$
Tanh	$\Phi(x)=\tanh(x)$
GELU	Standard GELU approximation
SELU	Standard SELU activation

Appendix

Table S.15: Activation functions used in the experiments.

Variant	Residual Mapping
Post-activation	$h^{(l)}=h^{(l-1)}+\Phi(W^{(l)}h^{(l-1)})$
Pre-activation	$h^{(l)}=h^{(l-1)}+W^{(l)}\Phi(h^{(l-1)})$
Depth-scaled	$h^{(l)}=h^{(l-1)}+\frac{1}{L-1}\Phi(\cdot)$
Fixed-scale	$h^{(l)}=h^{(l-1)}+\alpha\Phi(\cdot)$
Learnable gate (ReZero-style)	$h^{(l)}=h^{(l-1)}+\alpha_{l}\Phi(\cdot)$

Appendix

Table S.16: Residual connection variants. Depth scaled variant defined for

L\geq 2

Plain MLP / DeepLinear
Input: Flattened Vector
Linear ( $d_{\mathrm{in}}\rightarrow H$ )
$(L-1)\times$ Linear ( $H\rightarrow H$ )
Linear ( $H\rightarrow d_{\mathrm{out}}$ )
DeepLinear uses identity activation throughout.
Note: $L$ counts hidden layers, the network has $L+1$ weight matrices.

Appendix

Table S.17: Description of Plain MLP and DeepLinear MLP architectures used

ResidualMLP
Input: Flattened Vector
Linear ( $d_{\mathrm{in}}\rightarrow H$ )
$(L-1)\times$ Residual Block ( $H\rightarrow H$ )
Linear ( $H\rightarrow d_{\mathrm{out}}$ )
Residual connections begin from Block 2.

Appendix

Table S.18: Description of Residual MLP used.

PlainCNN

Input: RGB Image (

32\times 32

)

Conv

3\times 3

, BN, Activation

(L-1)\times

Conv

3\times 3

, BN, Activation

MaxPool at regular intervals (see released code)

Global Average Pooling

Linear Classifier

Appendix

Table S.19: Description of Plain CNN used.

ResNet-18 / ResNet-50
Input: RGB Image ( $32\times 32$ )
$3\times 3$ Conv Stem, BN, ReLU
Stage 1: Residual Blocks (64)
Stage 2: Residual Blocks (128, stride 2)
Stage 3: Residual Blocks (256, stride 2)
Stage 4: Residual Blocks (512, stride 2)

Appendix

Table S.20: Description of ResNet-18 and ResNet-50 used

Vision Transformer (ViT)
Input: RGB Image ( $32\times 32$ )
Patch Embedding: Conv $4\times 4$ , Stride 4, Dim $d$ ( $64$ Patches)
Prepend $[$ CLS $]$ Token ( $1\times d$ ) + Positional Embedding ( $65\times d$ )
$L\times$ Pre-LN Transformer Blocks ( $L\in\{3,6,9\}$ ):
Multi-Head Self-Attention ( $h\in\{2,4\}$ , $d_{k}=32$ ) $\pm$ Residual
Transformer MLP ( $d\to 2d\to d$ , Activation) $\pm$ Residual

Appendix

Table S.21: Description of Vision Transformer (ViT) used.

Common Settings
Depth	$\{1,2,4,8,16,32,64\}$
Task Rank	$\{1,4,16,24,32,64\}$
Optimizer	SGD, Momentum SGD, Adam
Initialization	Default, Low-Rank
Runs	2268 (total)
Architecture	Configuration

Appendix

Table S.22: Synthetic Full Grid Sweep

Setting	Value
Architecture	Plain MLP
Runs	24
Activations	SIREN, Leaky ReLU
Depth	$\{2,4,8,16,32,64\}$
Task Rank	24
Optimizers	SGD, Adam

Appendix

Table S.23: Synthetic Recovery grid sweep.

Setting	Value
Architecture	Residual MLP
Runs	448
Depth	$\{1,2,4,8,16,32,64\}$
Activations	SIREN, Leaky ReLU
Residual Types	Post-activation, Depth-scaled,
	Fixed- $0.5$ , Learnable Gate

Appendix

Table S.24: Synthetic Residual Runs Sweep

Setting	Value
Architecture	Residual MLP
Runs	1344
Depth	$\{2,4,8,16,32,64\}$
Task Rank	24
Activations	Linear, ReLU, Tanh, Leaky ReLU,
	SELU, GELU, SIREN

Appendix

Table S.25: Synthetic Residual sweep extended.

Setting	Value
Architectures	MLP, Residual MLP, CNN
Expansion Factor	$\{1,2,4,8\}$
Samples/Class	Full, 20, 50, 100, 200

Appendix

Table S.26: CIFAR10 CNN, Plain MLPs, Residual MLPs sweep

Setting	Value
Architecture	Plain CNN
Depth	$\{8,32\}$
Activations	Leaky ReLU, SIREN
Expansion Factor	$\{1,2,4\}$
Optimizers	SGD, Adam
Initialization	Default, Low-Rank

Appendix

Table S.27: CNN only on CIFAR10 sweep

Setting	Value
Architecture	ResNet
Models	ResNet-18, ResNet-50
Residual Skip Modes	On: $x+F(x)$ Off: $F(x)$ Scaled: $x+\frac{1}{N_{blocks}}F(x)$

Appendix

Table S.28: ResNet CIFAR10 sweeps

Hyperparameter	Values / Levels	Count
Residual Mode	Standard (with residuals), No-Residual (without residuals)	2
Depth ( $L$ )	3, 6, 9 blocks	3
Embedding Dim ( $d$ )	64 ( $h=2,d_{\text{mlp}}=128$ ), 128 ( $h=4,d_{\text{mlp}}=256$ )	2
Activation ( $\Phi$ )	GELU, ReLU	2
Optimizer	Adam ( $\eta=10^{-3}$ ), Momentum SGD ( $\eta=10^{-2},\beta=0.9$ )	2
Total Configurations	$2\times 3\times 2\times 2\times 2$	48

Appendix

Table S.29: Factorial Grid Sweep Configurations for Vision Transformers (

N=48

Models).

Low-Rank Representation Learning

See all 13 papers

How to Tame Grokking: Representation Geometry as a Control Signal
Jul 13, 2026Maksim A KazanskiiNeural Network Generalization Spectral Regularization
Independent Researcher
The Implicit Bias of Depth: From Neural Collapse to Softmax Codes
May 21, 2026Connall Garrod, Jonathan P. Keating, Christos ThrampoulidisDeep Linear Networks Implicit Bias
Mathematical Institute, University of Oxford · Department of Electrical and Computer Engineering, University of British Columbia
Scalar Representations of Neural Network Training Dynamics
Jun 29, 2026Pedro Jiménez-González, Miguel C. Soriano, Lucas LacasaNeural Network Optimization Chaotic Dynamical Systems
Institute for Cross-Disciplinary Physics and Complex Systems (IFISC, CSIC-UIB), Campus UIB, 07122 Palma de Mallorca, Spain

Jul 13, 2026J/K move · Enter open · S save

How to Tame Grokking: Representation Geometry as a Control Signal

Maksim A Kazanskii

Independent Researcher

Neural Network Generalization Spectral Regularization Representation Collapse Neural Representation Geometry Neural Network Training Dynamics Grokking

Open paper AskPDF arXiv

Grokking is a phenomenon in which neural networks initially memorize training data and only later exhibit strong generalization after prolonged optimization. Despite extensive recent study, the factors influencing the emergence and timing of grokking remain incompletely understood. We investigate the relationship between representation geometry and delayed generalization. We find that dimensionality collapse consistently precedes the onset of grokking in all evaluated settings. Motivated by these observations, we introduce Geometric Dimensionality Regularization (GeomDR), a simple spectral regularizer that modifies the effective dimensionality of hidden representations during training. Across modular addition, modular division, and permutation composition tasks, GeomDR consistently alters grokking dynamics and can substantially accelerate the onset of generalization depending on the intervention schedule and target dimensionality. In several settings, grokking is accelerated by up to 52 times relative to standard AdamW training. Similar qualitative effects are observed in both multilayer perceptrons and transformers. Together, these results suggest that representation geometry can serve as an effective control signal for grokking and provide evidence that geometric interventions offer a practical approach for studying and influencing delayed generalization in neural networks.

Abstract

Low-Rank Representation Learning

Explore similar work

How to Tame Grokking: Representation Geometry as a Control Signal

The Implicit Bias of Depth: From Neural Collapse to Softmax Codes

Scalar Representations of Neural Network Training Dynamics