Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Authors: Chi-Ning Chou
Hang Le
Yichen Wang
SueYeon Chung
Venue: ICML 2025 spotlightposter (2025)
Keywords: Neural manifold
feature learning
manifold capacity
representational geometry
Primary area: applications->neuroscience_cognitive_science

Abstract: Integrating task-relevant information into neural representations is a fundamental ability of both biological and artificial intelligence systems. 
Recent theories have categorized learning into two regimes: the rich regime, where neural networks actively learn task-relevant features, and the lazy regime, where networks behave like random feature models. Yet this simple lazy–rich dichotomy overlooks a diverse underlying taxonomy of feature learning, shaped by differences in learning algorithms, network architectures, and data properties. To address this gap, we introduce an analysis framework to study feature learning via the geometry of neural representations. Rather than inspecting individual learned features, we characterize how task-relevant representational manifolds evolve throughout the learning process. We show, in both theoretical and empirical settings, that as networks learn features, task-relevant manifolds untangle, with changes in manifold geometry revealing distinct learning stages and strategies beyond the lazy–rich dichotomy. This framework provides novel insights into feature learning across neuroscience and machine learning, shedding light on structural inductive biases in neural circuits and the mechanisms underlying out-of-distribution generalization.

Learning induces changes in brain activity, whether it in- volves navigating a new city, adapting novel motor skills, or Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Chi-Ning Chou
Hang Le
Yichen Wang? Yeo Chung 13
Abstract
a
Task relevant manifolds
Integrating task-relevant information into neural representations is fundamental ability of both biological and artificial intelligence systems. Re- cent theories have categorized learning into two regimes: the rich regime, where neural networks actively learn rask-relevant features, and the lazy regime, where networks behave like random fea- ture models. Yet this simple lazy-rich dichotomy overlooks diverse underlying taxonomy of fea- ture learning, shaped by differences in learning algorithms, network architectures, and data prop- erties. To address this gap, we introduce an anal- ysis framework to study feature learning via the geometry of neural representations. Rather than inspecting individual learned features, char- acterize how task-relevant representational mani- folds evolve throughout the learning process. We show, in both theoretical and empirical settings, that as networks learn features, task-relevant man- ifolds untangle, with changes in manifold geom- etry revealing distinct learning stages and strate- gies beyond the lazy-rich dichotomy. This frame- work provides novel insights into feature learning across neuroscience and machine learning, shed- ding light on structural inductive biases in neural circuits and the mechanisms underlying out-of- distribution generalization.
Cat? Dog?
cat manifold
dog manifold
Neural activity space (a.k.a, festure space)
b
Key takeaway:
When a network learns useful features from data, task-relevant manifolds become more untangled
Lazy regime (random festure)
After training
Rich regime (feature leaming)
leg, smaller margin
(egg larger margin sokaion volume capacity)
solution volume capacity)
more tangled
more untangled
Figure 1. Schematic illustration. a. We propose investigate fea- ture learning using representational geometry and task-relevant manifolds. b, Feature learning (i.c., rich regime) can be viewed as process of untangling task-relevant strroudddssssruuuiinnnnnnn the neural activity space improve separation among manifolds.
acquiring cognitive tasks. These changes reflected in the incorporation of task-relevant features into neural rep- resentations (Olshausen & Field, 1996; Poort et al., 2015; Niv, 2019; Reinert et al., 2021; Gurnan & Gajic, 2023). Similarly, the remarkable success of deep learning is of- ten attributed to the ability of neural networks learn problem-speciffc features For example, in deep neural networks (DNNSs) (LeCun al., 1998; Krizhevsky al., 2012), the ability to learn rich feature hierarchies enables superior image classification performance (Girshick et al., 2014). Meanwhile, the seminal work (Chizat et al., 2019) demonstrated that neural networks can perform well even when there are negligible changes in the weights of the net- works. These observations raise important questions: Do neural networks always need to learn task-relevant features? How can we evaluate the quality of the features they learn?
1. Introduction
"qqual contribution 'Center for Computational Neuroscience, Flatiron Institute, New York, NY, USA 'University of Cal- ifornia, UCLA, Los Angeles, CA, USA *Center for Neural Science, New York University, New York, NY. USA. Corre- spondence Chi-Ning Chou <echou Ilatironnstitute.org Hang <hle@ fattroninstiuuccorg... SucYeon Chung <schung@ ffatironinstiuute.rgg>>
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
In this paper, features broadly refer to measurable properties or characteristics patterns data/input.

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
To answer these questions, researchers in representation learning have developed several methods to determine whether a neural network operates in the lazy regime (leam- ing without changing internal features) or the rich regime (learning task-relevant features)². These methods include measuring changes in the weights of the network, tracking activated neurons, and assessing differences in the linearized model (also known as the neural tangent kernel, NTK (Jacot et al., 2018)). Factors such as initial weight norm, learning rate, and readout weight have been found to play a role in whether network is lazy or rich (Chizat et al., 2019). Moreover, recent theoretical evidence has suggested that networks could perform better in the rich regime compared to the lazy regime (Yang & Hu, 2021; Shi et al., 2022; Karp et al., 2021; Damian al., 2022; Ba et al., 2022).
From this perspective, feature learning can be viewed as process of untangling tangled manifolds at initialization (i.e., random features) to facilitate easy separation.
To make this intuition concrete and quantitative, we use manifold capacity² (Chung et al., 2018; Chou et al., 2025) (Definition 2.1 and Definition B.3), to quantify the degree of richness in feature learning (Figure 2b). Additionally, we use geometric measures that are analytically connected to manifold capacity as mechanistic descriptors to explain how task-relevant manifolds untangle. To demonstrate our proposed method, we systematically study a wide range of settings, from theoretical models to machine learning and neuroscience problems. Our contributions include:
(Section 3) We use manifold capacity as representation-basee method to quantify the degree of feature learning and demonstrate that it is better than conventional measures across a wide range of settings.
However, feature learning is much richer than the lazy ver- sus rich dichotomy. For example, changes in representations are not always beneficial as they can lead to issues such as catastrophic forgetting (Kirkpatrick et al., 2017). More- over, different network architectures, training procedures, objective functions, and initializations can result in different inductive biases for feature learning (Chizat et al., 2019; Bordelon & Pehlevan, 2022; Ba al., 2022; Damian al., 2022), yet all of these scenarios could fall under the broad category of rich learning. Lastly, limitations in neuroscience technology for tracking synaptic weight changes in neural circuits necessitate a framework based on neural activities rather than network weights or neural tangent kernel.
(Section 4) Manifold geometry reveals previously un- reported subtypes of feature learning. We find that the training of neural networks undergoes various learning stages as shown by the dynamies of manifold geometry, and there are emergent learning strategies as networks having different degree of richness in learning.
(Section 5) We find new geometric insights that have not been reported in problems from neuroscience (e.g.. structural inductive biases in neural circuits) and ma- chine learning (e.g.. out-of-distribution generalization).
1.1. Contributions
1.2. Related work
We go beyond the lazy-veree sizy-versus-rccl dichotomy and address the above-mentioned gaps by investigating feature learning though the geometric properties of task-relevant manifolds. Here, task-relevant manifolds refer to the point clouds of neural activity patterns that are related to the tasks. In clas- sification, for instance, a manifold could be the point cloud of neural activations corresponding to stimulus category (e.g., the cat and dog manifolds in Figure 1 a, left). In other domains, a manifold could correspond to a context (c.g., environmental cues) in a neuroscience experiment or a concept (e.g., semantic categories) in a language model.
Feature learning has been a fundamental research problem in various domains, including neuroscience and machine learning. In neuroscience, understanding the relationship between neural representations and task performance is a central focus (Gao & Ganguli, 2015). Representational ge- ometry (Chung & Abbott, 2021) has emerged as promising approach to investigate how different organizations of fea- tures can lead to better task performance (Bernardi et al., 2020; Flesch et al., 2022; Gurnani & Gajic, 2023). Several studies have explored methods for inferring the underlying learning rules of neural network using representational geometry (Cao et al., 2020; Sorscher al., 2022) and low- order statistics (Nayebi et al., 2020). In machine learning, visualization techniques (Zeiler & Fergus, 2014) have been widely used to gain intuitive insights into learned represen-
Instead of explicitly identifying which features a network learns, we ask: what changes occur in task-relevant mani- folds as a network undergoes feature learning?
When a neural network learns useful features from data for solving some task(s), task-relevant mani- folds become better organized and more separable.
*Manifold untangling originates from neuroscience (DiCarlo & Cox, 2007) and refers to the intuition that task-relevant manifolds become inereasingly separable a nigh-dimensional state space. While there have been methods (Yamins & DiCarlo, 2016; Hong et al., 2016) using worst-case decoding accuracy to quantify the degree of manifold untangling, manifold capacity is average- case method that known be better capturing complex shapes of manifolds. See Section 2.2 and Appendix B for more details.
'These two regimes are also known as kernel regime and feature learning regime.
2

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Interpolating between lazy and rich learning
b Capacity quantifies the amount of
Effective geometric measures explain the changes in capacity
task-relevant features in neural al manifolds
low learning high leerning rate
Fewer task-relevant features
Lower capacity
Geometric measure
Lower capacity
Higher capacity
Lazy learning
Rich leaminn
Radius
packable
Higher capacity
Dimension
task-relovant features
packable
Center alignment
Changes in capacity quantifies task-relevant richness/laziness
richer
Fiefeer
alignment
lazlet
Relative T
Cspachy
Center axis
alignment
lozker
(e.g.. VGG-11 trained and evalusted on CIFAR-10
opach
Contribution I (Section 3): Quantify feature learning via representation geometry
Contribution II (Section 4):
Contribution IlI (Section 5); Reveal insights neuroscience ML
Delineate geometric differences in feature learning
Lazy lsamiing
Manifold
Rich lcamiing
Learning stages
Learning strategies
In-O
OCD
capacity
righ
Packabilty drtangledhess
(BGL
(69) dee making)
Figure 2. Our methods. a, We adopt the method from (Chizat et al., 2019) which interpolates lazy (in blue) and rich (in red) learning via adjusting scale factor of learning rate (see Section 2.1). b, We propose use changes of capacity across training to study task task-relevant richness/laziness in feature learning (Section 3). c, Effective geometric measures are analytically connected to the capacity value (Chou al., 2025), providing mechanistic descriptors to study representational changes in feature learning. d. Our three main contributions.
tations. On the theoretical front, the kernel method (Jacot et al., 2018; Lee et al., 2019) has been a leading approach to analytically characterize the behavior of neural networks, particularly in terms their deviation from the correspond- ing kernel. This line of research includes studies on the distinction between lazy and rich regimes (Chizat et al., 2019; Geiger al., 2020) and identifying problem settings where neural networks with feature learning outperform kernel methods (Ba et al., 2022; Dandi et al., 2023; Yang & Hu, 2021). See Appendix A for more on related work.
higher learning rate of intermediate layers compared to that of the readout weights, resulting in richer learning process (Figure 2a). See Appendix D and Appendix E for details.
Neural networks. For 2-layer networks, we considered synthetic data models with random point clouds as input manifolds. This setting serves as testbed for testing the proposed method and building intuitions. See Appendix D for details. For feedforward DNNs, as the goal of this work is to understand neural representations rather than push- ing benchmarks, we focused on models and settings that are large enough to see interesting phenomena. Specif- ically. we considered VGG-11 (Simonyan & Zisserman, 2015) and ResNet-18 (He al., 2016) and datasets CIFAR- 10 (Krizhevsky & Hinton, 2009), CIFAR- 100 (Krizhevsky & Hinton, 2009), CIFAR-1OC (Hendrycks & Dietterich, 2018), See Appendix E for details. In all reported experimental results, we trained 5 models initialized with different seed.
2. Method and Setup
2.1. Rich and lazy training in neural networks
We studied rich versus lazy learning in two standard set- tings: 2-layer non-linear neural networks synthetic data and feedforward deep neural networks on real image clas- sification datasets (Chizat et al.. 2019). All analyses were performed on the test data representations in the last layer.
Tesk-relevant manifolds. Let be the number of classes and N be the number of neurons (in the layer of interest, e.g., the last layer). The i-th class manifold is modeled as the convex set (as in linear classification, it is mathemat- ically equivalent to study the convex hull of a manifold) M; = conv(f : E X,)) where X; is the collection of inputs in the i-th class, ((x) is the representation for x, and conv(-) denotes the convex hull.
Interpolating between rich and lazy regime. In all ex- periments, we use the inverse scale factor π as tunable ground truth for the degree of feature learning. In particu- lar, o controls the magnitude of the output of the network, as (Chizat al., 2019). Intuitively, a larger îj leads to
3

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
2.2. Manifold capacity theory
and convex geometry (Definition B.3):
Manifold capacity theory (Chung et al., 2018; Chou al., 2025) extends the classic notion of storage capacity of points (Cover, 1965; Gardner & Derrida, 1988; Gardner, 1988) to object manifolds for study the untangling hypothe- sis" of invariant object recognition in vision neuroscience (DiCarlo & Cox, 2007). Manifold capacity quantifies mani- fold packability in terms of linear classification, making it a useful metric for assessing the degree of untangling.
O P
1
max B.EMM.
{Iproolcone((ns))) ((nsa)) til (1)
y- ((.
where (+:-) denotes the multivariate Gaussian distribution and cone(-) is the convex cone spanned by the vectors, i.e., cone (w:8,)) }) =
: X 2 0}. Specifically, O'sim - ow = ((1/N) and au can be computed by solving quadratic programs (Algorithm 2).
Definition 2.1 (Simulated manifold capacity (Chou et al., 2025)). Let P. N E N and M, S TR- be convex sets for each E [P] {1...., P). For each n E [N], define
2.3. Effective geometric measures
Equation 1 connects manifold capacity to the structure of the manifolds (M.)- Concretely, for each y; t, de t)} Yi arg (s. [pro)cone([ t) as the anchor points with respect to y and t. Intuitively, these anchor points are the support vectors with respect to some random projection and dichotomy y. Namely, the random- ness induces a distribution of anchor points supported on the manifolds {M&) J manifold capacity then emerges naturally summary statistic derived from this distribution. This connection between manifold capacity and anchor points motivated previous work (Chung et al., 2018; Chou et al., 2025) to define the following effective manifold geometric measures that capture the structure of manifolds while being analytically connected to capacity (see Figure 2c and Ap- pendix B).
Pn: = Pr 30 NVER" R" (0 s) 2 O, [P) SE S.)) y.
where y is dichotomy vector randomly sampled from Y' C {=1}" and II, is a random projection from R.V to IR.. Suppose px = (which means that the input manifolds are linearly separable in R,), the simulated capacity of (Miheefp is defined as
P
Orsim =
Em(1 (I-
1.0
Probability of lineatly separable
0.5
N
Definition 2.2 (Effective manifold geometric mea- sures (Chung et al., 2018; Chou et al., 2025), informal), For each i E [P]. define s := Ey, (oyyy,))) t)] as the center of the i-th manifold and define s (y, = (y t) -s to be the axis part of sy t) for each pair of (y,t).
0.0
N
(projected cimennsoonalllt)
Manifold dimension captures the degree of freedom of the noises/variations within the manifolds. It is approx-
Figure The area above the curve of linearly separable probability is the critical dimension N : (1-p.) z min. p. (n), which appears in the denominator the definition for ckaim-
imately DM Ri Ey,t P Σ; (el(0.))))))
/) which is
Intuitively, the simulated manifold capacity measures the packability (Chung et al., 2018) of manifolds by determin- ing the smallest dimensional subspace needed ensure that the manifolds can be separated. Namely, manifolds that are more packable (i.e., separable when projected to smaller dimensional subspaces) exhibit higher manifold capacity. While Definition 2.1 can be numerically simulated (Algo- rithm 1), a recent theory, Geometry Linked to Untangling Efficiency (GLUE) (Chou et al., 2025) further derived an an- alytical formula through techniques from statistical physics
analogous to the Gaussian width of the manifolds (Ver- shynin, 2018, Section 7.7).
Manifold radius captures the noise-to-signal ra- tio of the manifolds. It is approximately Rar P [
Ey. [ E
Center alignment captures the correlation between the center of manifolds and is defined as P& := P(P-1) Eizs (s{,s)))-
Axis alignment captures the correlation between the axis of manifolds and is defined as P == My-n Eiy, Ey. (y, t), sy, t)]]]-
'The "untangling hypothesis" posits that the brain transforms complex, entangled sensory inputs into more linearly separable representations, facilitating efficient object recognition. 'Typically we pick y to be the collection of I-versus-rest dichotomies.
Center-axis alignment captures the correlation be- tween the center and axis of manifolds and is defined as V'M P(P-1) Σisy Ey. Ey.el(ss) t))
'See Figure 4 for how NTK-label alignment and representation- label alignment could fail at quantifying task-relevant features.

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Table I. Comparison conventional measures used in lazy versus rich learning-
Our approach
Weight changes
NTK-label alignment
Representation-label
Accuracy
(manifold geometry)
alignment
Detect the changes in features
x
V
V
Quantify the amount of lask-relevant features
X
X
X
Representation-bassed
X
X
X
V
Delineate subtypes of feature learning
X
X X
x
Three important remarks on effective manifold geometric measures to be made: First, the changes in manifold ca- pacity can be explained by the changes of these geomet- ric measures. For example, the decrease of manifold ra- dius and dimension makes the capacity higher (see Fig- ure 2c, Section B.4). Second, these effective geometric measures faithfully track the corresponding underlying ge- ometric properties in well-studied mathematical settings (see Section B.5). Moreover, there a simple formula con- necting manifold capacity with effective geometric measure: ay R (1+ //Dr (see Appendix B for details). Finally, combining the above two points, these effective geometric measures serve as intermediate-level descriptors to investi- gate how different structural properties of neural manifolds contribute to the changes of task-level performance.
manifold untangling to quantify feature learning.
Theoretical justification on 2-layer non-linear neural net- works. To strengthen the connection between capacity and feature learning, we next consider well-studied theoretical model (Ba et al., 2022; Montanari et al., 2025) and analyti- cally characterize the relationship between capacity, predic- tion error, and the effective degree of richness. Concretely, we consider the training of a fully-connected 2-layer net- work of the form fxx = / (WT x), where x E Rd is an input, W E Nxd xd is the hidden layer matrix, a E R is the readout weight, and : R. + R is the (non-linear) activation function. To study feature learning in this setting, it is common to consider W to be randomly initialized (i.c., random feature model (Rahimi & Recht, 2007)) and update via gradient descent with squared loss. Meanwhile, the readout weight a is randomly initialized and fixed to avoid lazy learning (where the network minimally adjusts the hid- den layer and focuses on learning a good readout weight) as well enable mathematical analysis (Ba et al., 2022). Input data and label (x1, vI),.... ...., PX MPPan) were ran- domly generated by teacher-student setting. where there is a hidden signal direction 3 that correlates with the label (see Setting C.2 for the full setting). As previously proved in (Ba et al., 2022) (see Theorem C.5), in the proportional asymptotic limit (i.e., Pusin. d, N - oo at the same rate), the first-step gradient update can be approximated by a rank- matrix that contains label information, resulting in the updated weight to be more aligned with the hidden signal 8r. Hence, in this setting, the learning rate can be used as the ground-truth to measure the amount of task-relevant information (i.e., richness in learning) in the model repre- sentation after gradient updates.
3. Manifold capacity quantifies the degree of feature learning
In this section, we provide both empirical and theoretical justifications for using the increase in capacity during train- ing as a measure to quantify the degree of richness (or the amount of task-relevant features) in feature learning. Furthermore, we compare our method with conventional approaches in the study of lazy versus rich learning, high- lighting the new insights uncovered by our approach.
3.1. Justificatiins of capacity for quantifying the lazy versus rich dichotomy
Empirical justification in standard settings. We start with empirically justifying the use of capacity to quantify the degree of feature learning, A classic result in the literature of lazy versus rich training is to train a lazy network where the test accuracy improves, but the weight matrices (or kernels) do not change much before and after training. We consider two settings in (Chizat et al., 2019), one is feedforward DNNs (VGG-11 11 and ResNet- 18) trained on CIFAR- 10 (Fig- ure 2a-b, Figure 15, and Figure 14), and the other is 2-layer non-linear NNs trained on random point clouds (Figure 4a). In both cases, we observe that the manifolds are more un- tangled when training is richer and capacity correctly tracks the degree of feature learning (the ground truth is i7). This provides empirical justification for the use of capacity and
We extend the previous results in (Ba et al., 2022) from regression setting to classification setting. Specifically, We prove that capacity correctly tracks the effective degree of richness after one gradient step". Moreover, we derive a monotone connection between capacity and prediction accu-
'Here we follow the convention in (Ba et al., 2022) and study only the first gradient step as the key Gaussian equivalence step might not hold for more steps remarked in footnote 2 of (Ba et al., 2022).
5

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Interpolating between rich and lazy regime by scale factor
Inverse scale
factor
Capecity
accurcyy
Weight chargess
1.20
NTKilbbl aiigrmmet
0.24
Reg. aliggment
Inteer
0.21
21
0.44
0.80
0.60
0.60
trained gaussian clouds)
(2-layer non-linear NNs
iihhor
0.60
0.00
0.40
10"100"04
10"10F301
10"10"10"
10"10"10"
10"10"10P
epoch
epoch
epoch
spoch
epoch
b
Interpolating between wealthy and poor regime by varying input data dimension
Input data dimension
1.20
Test sccuracy
1.00
Weight changes
0.20
Kltall aiigmment
0.90
Rap. labbl alignment
0.45
1.0 wealther
Capacitty
0.90
0.50
0.10
0.60
0.40
0.60
0.30
0.35
12-layer non-linear NNs trained geussian clouds)
4.0 pooner
0.50
ondering
100110100 epoch
10"10"10' epoch
10"10'10'
10-10"100
10'100100
epoch
epoch
epoch
Figure 4. Capacity as a measure for the degree of feature learning. See Section D.I for details. a, We interpolated between lazy and rich regime in 2-layer NNs trained to classify Gaussian clouds. We found that capacity could tell the difference hetween the underlying scale parameter better than the other conventional methods. b. We fixed a scale parameter and initialized the input Gaussian clouds with different dimensions (higher dimension leads to poorer initial representations). We found that capacity could tell the difference in the amount of tasks-relevant features at initialization. Other conventional methods, e.g., the representation-lan alignment, could characterize the wrong ordering of wealthiness in initial features.
racy. Here, we provide an informal statement of our results and leave the formal version and proof in Appendix C.
changes, and alignment methods (Table 1). Concretely, weight changes at the t-th epoch is defined as [[W Wolle/||Wa, where W. is the weight matrix at the 1-th epoch. NTK-label alignment and representation-label align- ment at the t-th epoch are defined as CKA(KATK, .yy') and CKA(X;X, yyy") respectively, where y is the label vec- tor, CKA(: the center kernel alignment measure (Ko- rbblith al., 2019), KNIK is the neural tangent kernel and X₃ is the representational matrix at the t-th epoch. To test these measures in a wide variety of settings, we consider 2- 2-layer NNs with synthetic data where we can vary a wide range of parameters. See Appendix A for details about these methods and Appendix D for details about experiments.
Theorem 3.1. Given Theorem C.I and Setting C.2. Let 0 < n < oo be the learning rate of a one-step gradient descent NN
with squared loss and 2/1 = d 1/2 Pam where Pirain is the number of training points, d is the impul dimension, and N is the number of hidden neurons. Let d. (1) be the capacity and let ACCp d.N (()) be the prediction accuracy after a gradient step with learning rate n. We have
d
1. (Capacity tracks the degree of richness) d.
a(1, d1, 122 where a(s;';) is defined in Theorem C.4. Specifically, a(1, 122 < cu(y', Hi, 122 for every 0 < <n= < vf.
Poinnd (n)
Capacity can detect task-relevant features in the pres- ence of complex structures in data (Figure 4a). We consider 2-layer NNs trained on random Gaussian clouds with gradient descent. We vary the scale parameter of the network to interpolate between lazy and rich regimes as done in (Chizat et al., 2019). We find that capacity is better at telling the difference of effective richness (i.e., the scale parameter) of the training than other conventional measures (Figure 4a). In particular, when the training is richer, we expect the representations to exhibit more complex struc- tures. Manifold capacity excels at extracting task-relevant structures in representations because it is data-driven and free from additional statistical assumptions on the data.
2. (Capacity links to prediction
accuracy) AcCp N. (n) No Acd(n, 11, 1/2) where Acc(n, 11, 122) is formally defined in Theorem C.4. In particular, there exists an increasing and in- vertible function 16g1,02 R + [0, such that Acc(1) 1/1, 4/2) he (a(n, (/2))-
The above theorem justifies the usage of capacity as a mea- sure for the degree of richness in feature learning within well-studied theoretical setting. We remark that our proof requires substantial technical improvements from (Ba et al., 2022) due to the difference between regression and classifi- cation (e.g., analyzing the margin of the Gaussian equivalent model after one-step gradient using tools from (Montanari et al., 2025), Theorem C.7).
Capacity can quantify the differences in task-relevant features at initialization (Figure 4b). When comparing two networks with different initializations, focusing solely on network changes can overlook differences in features present at initialization. Here, we use the capacity value at initialization to determine whether a network a wealthy
3.2. Comparison to conventional methods
Here we compare our method with several common measures for feature learning: accuracy curves, weight
6

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
An example of learning strategies (2-layer network on synthetic data, Fig. 3a)
b
example of learning strategies (2-layer network synthetic data, Fig. 3b)
richer
poorer
last epoch
first epoch
Inyerse scale factor: 2 Inverse scale factor: 2
Input dimensionaalty
Dimension
5.75
240
0.75
0.90
Inverse scale factor: 2 Inyerse scale factor:
0.54
Radius
0.22
0.50
lazier
9.6
0.90
0.36 Radius
wealthier
An example of learning stages (VGG-11 trained with CiFAR-10)
Intuitive Picture
Stage Clustering
Stage Structuring
Stage Separating
Stage 4 Stabilizing
1.0
Capacity
0.8
Dimension
0.6
Center A.
0.4
Axes
Center-Axes
Training steps
Figure Manifold geometry characterizes learning strategies and learning stages. Capacity contour plot of example from Figure 4a. The x-axis is the average manifold radius Ra. the y-axis is the average manifold dimension Dar, and the contour is the geometric approximation capacity, i.c., Ow & (1 + Ry /DD Normalized manifold geometry dynamics plot of VGG-11 trained with CIFAR-10. The values in each row rescaled so that the min and max values are 0 and
Dy (see Appendix B). b, Capacity contour plot of the example from Figure 4b. c,
regime (i.e., possessing more task-relevant features) or poor regime (i.e., possessing less task-relevant features), shown in (Figure 4b). For example, network is in wealthy (resp. poor) regime when the manifolds are lower (resp. higher) dimensional because there are more (resp. less) features that are useful for telling the manifolds apart. The wealthy versus poor distinction provides insight into the network's initial state, allowing for more comprehensive comparison of different settings (see Section 5.1).
We demonstrate in 2-layer NNs the emergence of distinct learning strategies driven by different factors of training. In Figure 5a, consider 2-layer networks trained on Gaus- sian clouds (as in Figure 4a). As training moves from the lazy to a richer regime (blue to gray), the network com- presses both the radius and dimension to increase capacity. Interestingly, an even richer regime (gray to red), the network sacrifices radius to further reduce dimension.
In Figure 5b, we consider the setting in Figure 4b where interpolate the wealth of initialization by varying input data dimension. For the wealthiest initialization (purple), the network primarily compresses radius. For poorer initial- ization (green), both radius and dimension are compressed in lazier training, while in the richer regime (e.g.. inverse scale factor 24), the network sacrifices radius for further dimension compression. In summary, varying degrees of richness in feature learning can exhibit different learning mechanisms, captured by manifold geometry.
4. Manifold Geometry Reveals Subtypes of Feature Learning
In this section, demonstrate that feature learning is much richer than the lazy versus rich dichotomy. In particular, use manifold geometric measures (Figure 2c, and Ap- pendix B for details) to delineate the differences in the learned features (learning strategies) and representational changes throughout training (learning stages). The key take- away from this section is the ability of our method reveal task-relevant changes in neural representations.
4.2. Geometric changes during training: Learning stages
4.1. Geometric differences in learned features: Learning strategies
Neural networks learn in a highly non-monotonic manner throughout the training period. Examples include double descent (Belkin et al., 2019; Nakkiran et al., 2021; Mei & Montanari, 2022) and grokking (Power et al., 2022; Liu
To increase capacity, network can shrink the radius and/or compress the dimension of neural manifolds (Figure 2c).

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
et al., 2022; Nanda et al., 2023; Kumar et al., 2024). Pre- vious works have analytically or empirically described the different stages/phases such comprehension, grokking, memorization, and confusion (Liu et al., 2022) through the trajectory of accuracy curves.
a RNNs trained on neuroscience tasks with
b
Effective laziness [Liu et al. ICLR 2024]
different initial weight rank
Activation Aazura
riche
inter
(e.g. alternative
22
choice)
Inibal weight rank
Initill weight rank
From Figure 5a.b we observe distinct stages of manifold geometry evolution during training in 2-layer networks. In the very rich regime, the network initially compresses both radius and dimension, then increases radius further reduce dimension. In Figure 5c, we examine standard setting where VGG-11 is trained on CIFAR-10. Despite the rapid saturation of training and test accuracy, at least four stages of geometric changes are evident (see Figure 2c for analytical connections between geometric measures and capacity): a clustering stage (initial manifold compression), followed by structuring stage (increasing alignment), separating stage (decreasing alignment to push manifolds apart), and final stabilizing stage (further reducing center alignment).
Regresentation signment
Tost CUITE
Beflne trareeg
richer
ladis
ap" weight
ad
weight
weight rank
d Structural blases reflected by geometry
Before training After training
inttiulicatoo
Radius
Ares Algriment
2 R 8
Capacity coe
12t 12t sat
121 Initial weight rank
Init al weight rank
Dimersson
Alignment
%
606
2
Poorer inbillzatioon
5. Applications to Neuroscience and Machine Learning Problems
sa Initial weight rank
Inital weight rank
Intial weight rank
Figure 6. Structural inductive biases in neural circuits. a, We con- sider RNNs trained on standard neuroscience tasks. b, Previous work (Liu et al., 2024) found that the initial weight rank of the recurrent connectivity matrix leads to inductive bias toward effectively richer or lazier training. c, We find that RNNs trained with different initial weight rank reach the same capacity value at final epoch. It is the difference in capacity at initialization that makes RNNs with small initial weight rank richer in training. Despite having the same capacity at final epoch, RNNs with differ- ent initial weight rank exhibit different manifold geometry.
In previous sections, we used capacity to quantify the degree of feature learning and delineate the learning stages and strategies through effective geometry. In this section, we apply our framework to find geometric insights in problems from neuroscience and machine learning.
5.1. Structural inductive biases in neural circuits
Structural variations (e.g., neural connectivity pattern) can play crucial role in learning dynamics (Campagnola al., 2022; Goudar et al., 2023; Xie et al., 2023; Raman & O'Leary, 2021). Theoretical studies have shown that certain patterns of neural connections can enhance the speed and/or performance at which specific tasks are learned (Canatar et al., 2021; Flesch et al., 2021; Braun et al., 2022). How- ever, there is a gap between theory and experiments as currently there is no experimental technology that can easily measure the precise changes in synaptic weight of multiple neurons simultancously. Mcanwhile, studying the changes in neural activity can provide an alternative lens to probe the functional changes in neural circuits. We study recurrent neural networks (RNNs) that are trained on standard neuro- science tasks such as perceptual decision making (Britten et al., 1992) (Figure 6a). We adopt the setting from previous work (Liu et al., 2024) on investigating how differences in connectivity initialization affect the learning process.
works (RNN) weights with varying ranks via Singular Value Decomposition. The RNN have 300 hidden units, 1 layer, with ReI.U activations, and are trained for 10000 iterations using SGD optimizer. Manifold capacity and effective geo- metric measures are computed using representations from the hidden states. See Appendix F for details.
Our findings. First, we study the training dynamics of capacity value in RNNs with various initial weight rank (Figure 6c). In agreement with the previous finding in (Liu et al., 2024) using weight changes, we find that the capacity changes of the small initial weight rank RNNs are higher than those of the large initial weight rank RNNs. Inter- estingly, the capacity values at the final epoch are about the same for RNNs with different initial weight rank. is the difference in capacity value at initialization that dis- tinguishes the learning dynamics of RNNs with different initial weight rank. Namely, small initial weight rank RNNs are in the poorer-richer feature learning regime, while large initial weight rank RNNs are in the wealthier-lazz feature learning (Figure 6c).
Experimental setup. We use the neurogym package (Molano-Mazon et al., 2022) to simulate common cognitive tasks, e.g., perceptual decision making. To study how con- nectivity structure impacts learning strategies, we follow the setup in (Liu et al., 2024) and initialize recurrent neural net-
8

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
CIFAR-10
b DA performance improved when the learning became richer, while OOD performance dropoed when the learning became ultra-rich (VGG-11 trained CIFAR-10)
C
Geometric signatures from in-distri manifolds can reflect that from COD manifolds
Training
ship
acc
data
Denasct
CIFA--10
CIFAR-10 (in-dist)
Haday
Dimension
CIFAR-1OC
CIFAR
CIFAR-10C
Centersalca
Aes
(same classes)
k2y
eih
Demain-
adaptation
CIFAR-10C (DA)
(DA) data
0:0
CIFAR-100
(non-overiapping classes)
CIFAR-100
Out-of- distribution (OOD) data
tulp
house
(OOD)
Invers scale factor
Inverse scale factor
Figure 7. Out of distribution generalization. a. CIFAR-IOC domain adaptation (DA) dataset and CIFAR- 100 OOD dataset. b. Test accuracy improves for CIFAR- 10 and IUC as the training becomes richer and richer while the linear probe accuracy for CIFAR- 100 would drastically drop in the ultra-rich training regime. e, Effective manifold geometry of CIFAR-IO0 100 reveals that the expansion of manifold radius and the increase of center-axis alignment explain the failure of OOD generalization in the ultra-rich regime. The color is normalized for each row respectively.
Next, we find that manifold geometric organizations are quite different for different initial weight ranks (Figure 6d). For example, poorer-richer learning (i.e., small initial weight rank) ends up with a larger radius and smaller dimension, while it is the opposite for wealthier-lazier learning (i.e., large initial weight rank). This finding suggests that RNNs exhibit structural biases at the manifold geometry level.
capacity and test accuracy, which warrants the use of effec- tive geometry. Next, we find that the expansion of manifold radius and the increase of center-axis alignment in the ultra- rich regime explain the drop of capacity. Interestingly, we also see an architectural difference where it is the increment in dimension in the ultra-rich regime explaining the drop of capacity in ResNet- 18 (Figure 17). We leave it as fu- ture direction to extend our study, applying these geometric insights to improve OOD generalization performance.
5.2. Out-of-distribution generalization
Understanding how neural networks learn and represent features is crucial not only for task performance but also for their ability to generalize beyond the training distribu- tion. Out-of-distributiio (OOD) generalization remains a fundamental challenge, as models often struggle when en- countering data that deviates from their training set.
6. Conclusion and Discussion
Feature learning is feature in the study of neural net- works in computational neuroscience and machine learning, and it is much richer than the lazy versus rich dichotomy. Understanding the connection between feature learning and performance paves the way for designing network architec- tures and learning algorithms with greater reliability and transparency for practical applications.
Experimental setup. For each model pre-trained on CI,AR-10, we train a linear classifier (i.e., linear probe (Alain & Bengio, 2016)) on top of the last-layer rep- resentation with CIFAR- 100 train set, and then evaluate the linear probe's performance on CIFAR-100 test set (see de- tails in Appendix Section E.4). We also consider corrupted version of 100 the CIFAR-10C dataset (Hendrycks & Dietterich, 2018) as an example of domain adaptation (DA) task. Finally, we compute the manifold capacity and effective geometric measures on last-layer representations.
The primary contribution of this work is to demonstrate how the perspective of task-relevant manifold untangling (quantified by manifold capacity and delineated by mani- fold geometric measures) can enhance our understanding of feature learning at an intermediate level. We propose several promising future directions, including extending the theoretical analysis, exploring applications in other types of DNN (e.g., transformers) and addressing relevant scientific inquiries in neuroscience, such inferring plasticity mech- anisms from observed learning dynamics in neural data, and induced changes across brain regions. We believe that investigations in these intermediate-level understandings can be leveraged to design more robust, gen- eralizable, and safer deep neural networks, as well as more accurate models for neuroscience applications.
Our findings. We see that the test accuracy of the OOD dataset increases when the network enters the rich learning regime (n) around 0.1) but decreases drastically when the degree of feature learning is too rich (īj around 1.0). The failure in such ultra-rich feature learning regime is different from the test accuracy of both CIFAR- 10 and CIFAR- -OC ( (igurre 7b). Looking at the capacity and effective geometry ( (Figure 7c), we first see strong correlations between the
predicting learning-induced

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Acknowledgments
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., and Yang, G. High-dimensional asymptoties of feature learn- ing: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35: 37932-37946, 2022.
We thank William Yang, Yao-Yuan Yang, and the members of Chung lab for the discussion regarding many prelimi- nary results and early versions of the manuscript. We thank Ghana Shyam Bandi, Shiyu Ling, Shreemayi Sonti, and Zoe Xiao for helpful discussion during the piloting stage of this project. This work was supported by the Center for Computational Neuroscience at the Flatiron Institute, Si- mons Foundation, and by NIH grant RO1DA052220. S.C. was partially supported by a Sloan Research Fellowship, a
Bahri, Y., Dyer, E.,, Kaplan, J., Lee, J., and Sharma, U. Ex- plaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):221188881))
2024.
Belinkov, Y.. Durrani, N., Dalvi, F., Sajjad, H., and Glass, J. What do neural machine translation models learn about morphology? In Barzilay, R. and Kan, M.-Y. (eds.), Pro- ceedings of the 55th Annual Meeting of the Association for Computational Linguisties (Volume Long Papers), pp. 861-872 July 2017.
Klingenstein-Simont Simons Award, and the Samsung Advanced In- stitute of Technology project, "Next Generation Deep Learn- ing: From Pattern Recognition to AI." All experiments were performed using the Flatiron Institute's 's high-performance computing cluster.
Belkin, M., Hsu, D., Ma, S., and Mandal, S. Reconciling modern machine-learning practice and the classical bias- variance trade-off. Proceedings of the National Academy of Sciences, 116(32)155899111 15854, 2019.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Bernardi, S., Benna, M. K., Rigotti, M., Munuera, J., Fusi, S. and Salzman, C. D. The geometry of abstraction in the hippocampus and prefrontal cortex. Cell, 183(4): 954-967, 2020.
Code Availability
All code required reproduce the figures presented is avail- able under MIT License http://ggituub.co/huunggg reunaailabfeenur--eannngggeeom
Bordelon, and Pehlevan, C. The influence of learning rule on representation dynamics in wide neural networks. In The Eleventh International Conference on Learning Representations, 2022.
References
Bordelon, B., Canatar, A.. and Pehlevan, C. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learning, pp. 1024-1034. PMLR, 2020.
Alain, G. and Bengio, Y. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.0164,, 2016
Braun, L.. Dominé, C.,, Fitzgerald, J., and Saxe, A. Exact learning dynamics of deep linear networks with prior knowledge. Advances in Neural Information Processing Systems, 55:6615-6629, 2022.
Allen-Zhu, Z. Li, Y., and Liang, Y. Learning and generaliza- tion in overparameterrize neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
Britten, K., Shadlen, M., Newsome, W., and Movshon, J. The analysis of visual motion: comparison of neu- ronal and psychophysical performance. Journal of Neu- roscience, 1992.
Amit, D. J., Gutfreund, H., and Sompolinsky, H. Statistical mechanics of neural networks near saturation. Annals of physics, 173(1):30-67, 1987.
Ansuini, A., Laio, A., Macke, J. H., and Zoccolan, D. In- trinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Campagnola, L.,, Seeman, S. C.,, Chartrand, T., Kim, L., Hoggarth, A., Gamlin, C., Ito. S., Trinh, J., Davoudian, P., Radaelli, C., et al. Local connectivity and synaptic dynamics in mouse and human neocortex. Science, 375 (6585), 2022.
Arora, S. Du, S., Hu, W. Li, 7.., and Wang, R. Fine-grained analysis of optimization and generalization for overpa- rameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019.
Canatar, A., Bordelon, B., and Pehlevan, C. Spectral bias and model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1 2211)29144 2021.
10

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Cao, Y., Summerfield, C., and Saxe, A. Characterizing emergent representations in space of candidate learning rules for deep networks. Advances in Neural Information Processing Systems, 33:8660-86777 2020.
Ehrlich, D. B., Stone, J. T., Brandfonbrener, D., Atanasov, A., and Murray, J. D. Psychrnn: An accessible and flexi- ble python package for training recurrent neural network models on cognitive tasks. eneuro, 8(1), 2021.
Chizat, L., Oyallon, E., and Bach, F. On lazy training in differentiable programming. Advances in neural informa- tion processing systems, 32, 2019.
Engstrom, L., Ilyas, A., Santurkar, S.. Tsipras, D., Tran, B., and Madry, A. Adversarial robustness as a prior for learned representations. arXiv preprint 2019.
arXiv:1906.000
Chou, C.-N., Kim, R., Arend, L., Yang, Y. Mensh, B., Shim, W. M., Perich, M., and Chung, S. Geometry linked to untangling efficiency reveals structure and computation in neural populations. bioRxiv, 2025. doi: 10.1 100/2024. 28.2177..
Erhan, D., Bengio, Y., Courville, A., and Vincent, P. Visual- izing higher-layer features of a deep network. Technical Report, Univeristé de Moniréal, 2009.
Flesch, T., Juechems, K., Dumbalska, T., Saxe, A., and Summerfield, C. Rich and lazy learning task repre- sentations brains and neural networks. BioRxiv, pp. 2021- 2021.
Chung. S. and Abbott, L. F. Neural population geometry: approach for understanding biological and artificial neural networks. Current opinion in neurobiology, 70: 137-144, 2021.
Flesch, T., Juechems, K., Dumbalska, T., Saxe, A., and Summerfield, C. Orthogonal representations for robust context-dependent task performance in brains and neural networks. Neuron, 110())1258-1270) 2022.
Chung, S., Lee, D. D., and Sompolinsky, H. Classification and geometry of general perceptual manifolds. Physical Review X, 2018
Cohen, U., Chung, S., Lee, D. D., and Sompolinsky, H. Separability and geometry of object manifolds in deep neural networks. Nature commumicationss 2020.
Gao, P. and Ganguli, S. On simplicity and complexity in the brave new world of large-scale neuroscience. Current opinion in neurobiology. 32:148-155, 2015.
Cover, T. M. Geometrical and statistical properties of sys- tems of linear inequalities with applications in pattern recognition. IEEE transactions on electronic computers, pp. 326-334, 1965.
Gardner, E. The space of interactions in neural network models. Journal of physics A: Mathematical and general, 21(1):257, 1988.
Gardner, E. and Derrida, B. Optimal storage properties of neural network models. Journal of Physics A: Mathemat- ical and general, 21(1):271, 1988.
Damian, A.. Lee, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Conference on Learning Theory, pp. 5413-5452 5452. PMLR, 2022.
Geiger, M., Spigler, S., Jacot, A., and Wyart, M. Disentan- gling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020(11):1 2020.
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L.,, and Stephan, L. How two-layer neural networks learn, one (giant) step at time. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Ghosh, A., Mondal, A. K., Agrawal, K. K., and Richards, Investigating power laws in deep representation learning. arXiv preprint arXiv:2202.0588 05808, 2022.
DiCarlo, J. and Cox, D. D. Untangling invariant object recognition. Trends in cognitive sciences, 11(8)3333-41, 2007.
Girshick, R., Donahue, J., Darrell, T., and Malik, Rich fea- hierarchies for accurate object detection and semantic segmentation. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 580-587, 2014.
Du, S., Lee, J., Li, H., Wang. L.. and Zhai, X. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675- 1685. PMLR, 2019.
Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In Pro- ceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Du, S. S., Zhai, X., Poczos, B., and Singh, A. Gradient descent provably optimizes over- werrpparamterrized neural networks. In International Conference on Learning Rep- resentations, 2018

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Goudar, V., Peysakhovich, B., Freedman, D. J., Buffalo, E. A., and Wang, X.-J. Schema formation in neural pop- ulation subspace underlies learning-to-learn in flexible sensorimotor woblem--solving. Nature Neuroscience, 26 (5):879-890, 2023.
of task adaptation in language models. In Findings of the Association for Computational Linguisties: NAACL 2025, pp. 1855-1888. Association for Computational Linguistics, April 2025.
Kornblith, S., Norouzi, M., Lee, H., and Hinton, G. Sim- ilarity of neural network representations revisited. In International conference on machine learning, pp. 3519- 3529. PMLR, 2019.
Gumann, H. and Gajie, N. A. C. Signatures of task learning in neural representations. Current opinion in neurobiol- ogy, 83:102759, 2023.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Kriegeskorte, N. and Kievit, R. A. Representational geom- etry: integrating cognition, computation, and the brain. Trends in cognitive sciences, 17(8):401-411 2013.
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical report, University of Toronto, Toronto, Ontario, 2009.
Hendryeks, D. and Dietterich, Benchmarking neural network robustness common corruptions and perturba- tions. In International Conference on Learning Represen- tations, 2018
Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
Hofmann, T., Schölkopf, B., and Smola, A. Kernel meth- ods in machine learning. The Annals of Statistics, 36(3): 2008
Kumar, T., Bordelon, B., Gershman, S. J., and Pehlevan, C. Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, 2024.
Hong, H., Yamins, D. L., Majaj, N. J., and DiCarlo, J. J. Explicit information for category-orthogonal object prop- erties increases along the ventral stream. Nature neuro- science, 19(4):613-22 622, 2016.
Kuoch, M., Chou, C.-N., Parthasarathy, N., Dapello, J., DiCarlo, J. J., Sompolinsky, H., and Chung, S. Prob- ing biological and artificial neural networks with task- dependent neural manifolds. In Conference on Parsimony and Learning (Proceedings Track), 2024.
Hopfield, J. J. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7955)22554- 1982.
LeCun, Y., Bottou, L.,, Bengio, Y., and Haffner, Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(11 1) 2:277 2324, 1998.
Hu, B., Temiz, N. 7.,, Chou, C.-N., Rupprecht, P., Meissner- Bernard, C., Titze, B., Chung, S., and Friedrich, R. W. Representational learning by optimization of neural mani- folds in olfactory memory network. bioRxiv, pp. 2024 11, 2024.
Lee, J., Xiao, L., Schoenholz. S., Bahri, Y., Novak, R., Sohl- Dickstein, J., and Pennington, J. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent ker- nel: Convergence and generalization in neural networks. Advances in neural information processing systems. 31, 2018.
Liu, Y. H., Baratin, A., Cornford, J., Mihalas, S., SheaBrown, E. T., and Lajoie, G. How connectivity structure shapes rich and lazy learning in neural circuits. In The Twelfth International Conference Learning Representations, 2024.
Karp, S., Winston, E., Li, Y., and Singh, A. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances Neural Information Process- ing Systems, 34:24883-24899 2021.
Liu, Z., Kitouni, O., Nolte, N. S., Michaud, E.,, Tegmark, M., and Williams, M. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651- 34663, 2022.
Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Des- jardins, G., Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Grabska- Barsska-aa A., et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13)3521-3322 2017.
Mante, V., Sussillo, D., Shenoy, K. V., and Newsome, W. T. Context- dependtt computation by recurrent dynamics in prefrontal cortex. Nature, 503(7474):78 84, 2013.
Kirsanov, A., Chou, C.-N., Cho, K., and Chung, S. The geometry of prompting: Unveiling distinct mechanisms
12

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Mei, S. and Montanari, A. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667-766, 2022.
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfit- ting small algorithmic datasets. arXiv preprint arXiv:2201.02177 2022.
Miller, E. K., Erickson, C. A., and Desimone, R. Neural mechanisms of visual working memory in prefrontal cor- tex of the macaque. Journal of neuroscience, 16(16); 5154-5167, 1996.
Raghu, M., Unterthiner, T., Korblith, S., Zhang, C., and Dosovitskiy, A. Do vision transformers see like convolu- tional neural networks? In Advances in Neural Informa- tion Processing Systems, volume 34, pp. 12116-12128. Curran Associates, Inc., 2021.
Molano-Mazon, M., Barbosa, J., Pastor-Ciurana, J., Fradera, M., ZHANG, R.-Y., Forest, J., del Pozo Lerida, J., Ji-An, L., Cueva, C. J., de la Rocha, J., and et al. Neurogym: An open resource for developing and sharing neuroscience tasks, Feb 2022.
Rahaman, N., Baratin, A., Arpit, D., Draxler, F., Lin, M.,
Hamprecht, F., Bengio, Y., and Courville, A. On the spee- tral bias of neural networks. In International conference machine learning, pp. 5301-5310. PMLR, 2019.
Montanari, A., Ruan, F, Sohn, Y., and Yan, J. The gener- alization error of max-margin linear classifiers: Benign overfitting and high dimensional asymptotics in the over- parametrized regime. The Annals of Statistics, 53(2):822 - -853, 2025. doi: 10.1214/25-OS2249
Rahimi, A. and Recht, B. Random features for large-scale
kemel machines. Advances in neural information pro- cessing systems, 20, 2007.
Raman, D. V. and O'Leary, T. Frozen algorithms: How the brain's 's wiring facilitates learning. Current Opinion in Neurobiology, 67:207-214, 2021.
Nakkiran, P., Kaplun, G., Bansal, Y., Yang, T., Barak, B., and Sutskever, I. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechan- ics: Theory and Experiment, (121(12): 124003,, 2021.
Reinert, S., Hübener, M., Bonhoeffer, T., and Goltstein, P. M. Mouse prefrontal cortex represents learned rules for categorization. Nature, 593(7859):411-417, 2021.
Nanda, N., Chan, L.,, Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic inter- pretability. In The Eleventh International Conference Learning Representations, 2023.
Selvaraju, R. R., Cogswell, M., Das, A., Vedantam, R., Parikh, D., and Batra, D. Grad-cam: Visual explana- tions from deep networks via gradient- based localization. In Proceedings of the IEEE international conference on computer vision, pp. 618-626, 2017.
Nayebi, A., Srivastava, S., Ganguli, S., and Yamins, D. L. Identifying learning rules from neural network observ- ables. Advances in Neural Information Processing Sys- tems, 33:2639-2650, 2020.
Shi, Z., Wei, J., and Liang, Y. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations, 2022.
Niv, Y. Learning task-state representations. Nature neuro- science, 22(10):1544-1555 2019.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. In 3rd Inter- national Conference on Learning Representations (ICLR 2015), 2015.
Olah, C., Mordvintsev, A., and Schubert, L. Feature vi- sualization. Distill, 2017. doi: 223915/distill.0mm http::/diisill.pub/01///eeeee viterattrrevvuuuazzattoonnoo
Simonyan, K., Vedaldi, A., and Zisserman, A. Deep inside convolutional networks: Visualising image classification models and saliency maps. In Workshop International Conference on Learning Representations, 2014.
Olshausen, B. A. and Field, D. J. Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature, (6583)66076660 1996.
Paraouty, N., Yao, J. D., Varnet, I.., Chou, C.-N., Chung, S., and Sanes, D. H. Sensory cortex plasticity supports auditory social learning. Nature communications, 14(1): 5828, 2023.
Sorscher, B., Ganguli, S., and Sompolinsky, H. Neural repre- sentational geometry underlies few-shot concept learning. Proceedings of the National Academy of Sciences, 119 (43):e220080013 19, 2022.
Poort, J., Khan, A. G., Pachitariu, M., Nemri, A., Orso- lic, I., Krupic, J., Bauza, M., Sahani, M., Keller, G. B., Mrsic-Flogel, T. D., et al. Learning enhances sensory and multiple non-sensory representations in primary visual cortex. Neuron, (6(6):1478 100::-11899999 2015.
Stephenson, C., suchismita padhy, Ganesh, A., Hui, Y., Tang, H., and Chung, S. On the geometry of general- ization and memorization in deep neural networks. In International Conference on Learning Representations, 2021.
13

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Sundararajan, M., Taly, A., and Yan, Q. Axiomatic attribu- tion for deep networks. In International conference machine learning, pp. 3319-3328. PMLR, 2017.
Vershynin, R. High-dimensional probability: An introduc- tion with applications in data science, volume 47. Cam- bridge university press, 2018.
Wakhloo, A. J., Sussman, T. J., and Chung, S. Linear clas- sification of neural manifolds with correlated variability. Physical Review Letters, 2023.
Woodworth, B., Gunasekar, S., Lee, J. D., Moroshko, E., Savarese, P., Golan, I., Soudry, D., and Srebro, N. Ker- nel and rich regimes in overparametrized models. Conference Learning Theory, pp. 3635-3673, PMLR, 2020.
Xie, M., Muscinelli, S. P., Harris, K. D., and A. Task-dependent optimal representations for cerebellar learning. Elife, 12:e82914. 2023.
Yamins, D. L. and DiCarlo, J. J. Using goal- goal-driven deep learning models to understand sensory cortex. Nature neuroscience, 1933)335-336, 2016.
Yang, G. and Hu, J. Tensor programs iv: Feature learning in width neural networks. In International Con- ference on Machine Learning, pp. 11727-11737. PMLR, 2021.
Yerxa, T., Kuang. Y., Simoncelli, E., and Chung, S. Learn- ing efficient coding of natural images with maximum manifold capacity representations. Advances in Neural Information Processing Systems, 36:24103-24125 2023.
Zeiler, M. D. and Fergus, R. Visualizing and understand- ing convolutional networks. In Computer Vision-ECCV 2014 13th European Conference, Zurich, Switzerland, September 6-12. 2014, Proceedings, Part 13, pp. 818- 833. Springer, 2014.
14

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
A. More on Related Work
Visualization. Due to the black-box and complex nature of deep neural networks, various visualization techniques have been developed attempt to characterize the features that models learn during training (feature visualization) and identify which input pixel and or feature activation in the hidden layers contribute significantly to the final model outputs (feature attribution). Feature visualization techniques visualize features (e.g convolutional filter in the case of CNNs) by generating the input sample that maximizes the activation of that given feature via gradient descent (Olah et al., 2017) (Erhan et al., 2009) (Zeiler& & Fergus, 2014). With its vivid visualization, feature visualization provide good intuition about the qualitative characteristics of the features that DNNs learn across layers (Zeiler & Fergus, 2014) as well as different types of models (e.g. standard vs adversarially robust (Engstrom et al., 2019)). Feature attribution techniques generally identify how much each input and/or hidden features contribute to the final model prediction by computing the gradient of that input/hidden features to the output (some example techniques include saliency map (Simonyan et al., 2014), Grad-cam (Selvaraju et al., 2017), integrated gradient (Sundararajan et al., 2017)). Although both feature visualization and feature attribution offer intuitive understanding about the model's feature characteristies, the qualitative nature of visualization makes it difficult to quantify the degree of relevance of the learned features to a given task.
Kernel dynamics. Kernel methods (Hofmann et al., 2008) have been classic machine learning techniques, where the primary goal is to design an effective embedding that maps inputs to a feature space, thus facilitating efficient algorithms to find good solutions (e.g., linear classifier). While neural networks are inherently complex, seminal works (Jacot et al., 2018; Lee et al., 2019) have shown that in the infinite width limit, a network can be linearized by its neural langent kernel (NTK). Thus, studying the NTK of network allows an analytical understanding of various properties of neural networks, such as convergence to global minima (Du et al., 2018; 2019), generalization performance (Allen-Zhu et al., 2019; Arora et al., 2019), implicit bias (Bordelon et al., 2020; Canatar et al., 2021), and neural scaling laws (Bahri et al., 2024).
When a network is properly initialized (Chizat et al., 2019), gradient descent can converge to the NTK of the random initialization, a setting known as the kernel regime (a.k.a., lazy training or random anndoffeature regime). On the other hand, a network can also enter what is known as the feature learning regime (a.k.a., rich training or mean-field limit), where it deviates from the NTK of the initialization (Geiger et al., 2020). Extensive research has been conducted characterize lazy versus rich regimes (Geiger et al., 2020; Woodworth et al., 2020) and to demonstrate instances where feature learning outperforms lazy training (Yang & Hu, 2021; Ba et al., 2022; Dandi et al., 2023). It is important to note that even when a network undergoes feature learning, the NTK can still be defined at each epoch. Previous works also analytically characterized the dynamics of kernel in simpler models (Bordelon et al., 2020). Studying such kernel dynamics also provides a lens for exploring questions related to feature learning, such as grokking (Kumar et al., 2024).
Representational geometry. The visualization approaches mentioned above focus on studying the geometric properties of the feature map itself. Another fruitful direction is to examine the geometric properties of the neural representations of inputs (i.e., embedding vectors) and their connections to performance (Chung & Abbott, 2021; Gurnani & Gajic, 2023). Various dimensionality reduction methods (e.g., principal components analysis (PCA), Isomap, t-SNE, MDS, and UMAP) have been proposed to build intuitions about the organization of high-dimensional feature spaces. In addition, there are approaches that study lower-order statistics of embedding vectors, such as representational similarity (Kriegeskorte & Kievit, 2013) and spectral methods (Rahaman et al., 2019; Bahri et al., 2024; Ghosh et al., 2022). Methods for extracting higher-level geometric properties (e.g., dimension) have also been proposed (Chung al., 2018; Cohen et al., 2020; Chou et al., 2025; Ansuini et al., 2019), with wide applications in both machine learning (e.g., memorization (Stephenson et al., 2021), grokking of modular arithmetic (Liu et al., 2022; Nanda et al., 2023), in-context learning in LLM (Kirsanov et al., 2025), self self-supervised learning (Yerxa et al., 2023; Kuoch et al., 2024)), and neuroscience (e.g., perceptual untangling in object categorization (Chung et al., 2018), olfactory memory (Hu et al., 2024), abstraction (Bernardi et al., 2020), few-shot learning (Sorscher et al., 2022), social learning (Paraouty et al., 2023)).
A.1. Previous work on storage capacity
Storage capacity is defined as the information load for linear readouts and has been studied in several communities, including learning theory (Cover, 1965) and statistical physics of neural networks (Gardner & Derrida, 1988; Gardner, 1988). To enable a mathematical treatment, we focus on the proportional limit (a.k.a. the high- high-dimensional limit, the thermodynamic limit), i.e., N, P ->>> oo and - N/P- = O(1). For given network and input data, we denote the representation of
15

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
the i-th input 24 as (xxz) E IR where & is the (non-linear) feature map. The storage capacity of + is defined as.
a( ()) == lim max
Pr E R., ViE v. ((x,)) 2 0] 1 N
P
(2)
where E [+ is uniformly random sampled, @ is the linear classifier, and oN (1) denotes vanishing terms (i.e., ON (1) + 0 as N + x). One can also consider the setting where the distribution of y is biased toward some task direction (Montanari et al., 2025). Intuitively, a()) quantifies the number of patterns per neuron that a network can store and decode with linear readouts.
Recall that storage capacity is defined as the critical ratio between the number of stored patterns and the number of neurons (Equation (2)). Cover's theorem (Cover, 1965) shows that the success probability of having linear classifier for P points with random binary labels in general position 8 is p(N, P) = 21 >x k= 1 In particular, for P/N < 2 we have lim,- PNN, P) = 0 and for P/N > 2 we have lim -co P) = 1. Namely, the storage capacity of points in general position with random binary label is 2. See also Figure 8 for finite-size and numerical examples.
N
1.0
10-
0.-
Rrobabiityy
04
0.8
04
4.
P/N
Figure 8. Storage capacity of random points and labels. Storage capacity is defined as the critical ration P/N = 2 where the success probability undergoes a phase transition. Left: finite size success probability curves proved Cover's theorem. Right: numerical check for Cover's theorem.
In the seminal works of Gardner and Derrida (Gardner & Derrida, 1988; Gardner, 1988), the storage capacity for random points with non-zero margin is analytically characterized using replica method. the context of associative memory, the storage capacity of Hopfield networks (Hopfield, 1982) calculated by (Amit et al., 1987).
B. Manifold Capacity Theory and Effective Geometry
Manifold capacity theory (CCun et al., 2018; Chung & Abbott, 2021; Wakhloo et al., 2023; Chou et al., 2025) was originally developed for the study of manifold untangling (DiCarlo & Cox, 2007) in theoretical/compuuatiinaall neuroscience. Intuitively, manifold untangling refers to the increased separation of high-dimensional manifolds (c.g., point cloud manifolds) in the eyes of downstream readout. MCT quantifies this intuition via modeling a downstream neuron as a linear classifier, and uses the packing efficiency of the neural representational space evaluate the degree of manifold untangling. Mathematically, such packing efficiency coincides with support vector machine (SVM) in average-case setting.
B. Neural manifolds convex hulls of pre-readout representations
As we are studying feature learning, we are interested in the neural representations that correspond to activations obtained from the pre-linear readout layer neurons. The readers can refer to Appendix D and Appendix E for details activation extraction. Notation wise, let N be the number of neurons. Therefore, all neural representations live in R.W space.
Next, we group neural representations by their category labels assigned during training to obtain P data manifolds. For i E [1...., P). the i-th data manifold, denoted as M,, is a convex set in R.N To ensure convexity in practice, we take M; to be the convex hull of a collection of vectors M, = (xi...., } where M, is the number of points in the i-th manifold.
MMeningg that every N' < N points linearly independent. Note that random points in general position with probability o(1)
16

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Notice that the each data manifold lives in its own subspace of dimension D, N. Therefore, we can rewrite each data manifold in its own coordinate system:
M2
up +
S 3 s ($1, (s1,..., SD.) E S
(3)
Here, ug is the center of the i-th manifold and {u/) j= is orthonormal basis. The shape set S C IRD: is convex set denoting coordinates of the manifold points in its subspace. In practice, the manifold axes and shape sets S; are completely data driven.
B.2. A simulation definition for manifold capacity
Recall from Section 2 that the simulation version of manifold capacity is defined follows.
Definition B.1 (Simulated manifold capacity (Chou et al., 2025)). Let P, e N and M, S R.W be convex sets for each i [P] = (I....,P}. For each n E [AV, define
Pn = Pr FOER" (0(0,)) 2 O, Vi E (P), s SE)
yll.
where y is dichotomy vector randomly sampled" from y { and II, random projection from IN to R". Suppose pN 1 (which means that the input manifolds are linearly separable in R the simulated capacity of (Moheetr) is defined as
P
Orgim
En - Pn)
Intuitively, the simulated manifold capacity measures the packability (Chung et al., 2018) of manifolds by determining the smallest dimensional subspace needed to ensure they can be separated. Namely, manifolds that are more " (i.e., separable when projected to smaller dimensional subspaces) exhibit higher manifold capacity. Note that the simulated capacity can be estimated from data by empirically estimate Pn. and perform binary search to find the critical dimension minp.. 20.5 {n. This procedure is computationally expensive and requires some choices of hyperparameters (which makes the definition a little ad hoc). Nevertheless, Definition 2.1 provides good intuition on how to think about manifold capacity (and its connection to packing).
We remark that one don't necessarily need to sample y uniformly at random from {=1}". It is also reasonable fix choice of y sample y from subset of [+1] (e.g., all the 1-versus-rest dichotomies) as discussed in (Chou et al., 2025).
B.3. A mean-field definition for manifold capacity
To overcome the above- mentioned drawbacks of simulated manifold capacity, previous work (Chung al., 2018; Wakhloo al., 2023; Chou al., 2025) defined mean-field models enable mathematical definition of manifold capacity while providing good approximation to the simulated manifold capacity. In particular, in this paper use the GLUE (Geometry Linked to Untangling Efficiency) theory developed in (Chou et al., 2025).
Mean-field model from GLUE (Chou et al., 2025). Given collection of (finite) data manifolds (M.) A mean-field model is to generate infinitely many (Py) manifolds in an possible N such that these "mean-field" manifolds are separable. The key idea is that if this generating process nicely preserve the structure in the data manifolds, then the packing property of these mean feand-ieed manifolds will be very similar
(NM) space and characterizing the largest
Definition B.2 (Mean-field model from (Chou al., 2025)). Let (Mc)setP) be collection of data manifolds in R as defined in Equation 3. Let Rz0 and PM, NM be integers with the following properties: (1) PM, Ny + and (ii) PM/ = a < oc, and PM be divisible by P. We define the mean-field manifolds MM Nu) = {M JaMaaaae JaE[Pw/PM.REP)) follows.
"Typically we pick y to be the collection of all 1-versus-rest dichotomies. 1he reason why this is called "packing" is that projecting manifolds into smaller dimensional subspace like packing them into smaller neural representational space.
17

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Algorithm 1 Estimate simulated manifold capacity
Input: {M,}: point clouds, each containing M points in an N-dimensional ambient space. Output: O'gim: Simulated manifold capacity.
n" f BinarySearch( 1,
Otgim f- P/n".
return Otsim
% Binary search for the smallest such that Pa. 2 0.5. function (mmarySearch([M:: n,)
Thm + [(m + /22]
if nm = ny then
return nmm-
else
Pn... Esttrob([PMa)) nm).
if Pn... 0.5 then return BinarySearchi ( [M:), Ru, mm)-
else
return BinarySearch(( Rom, np).
end if end if end function
% Estimate the probability of linear separability after random projection. function (atPProb(.M,),.. cnt 4- 0
= 1000)
for from to m do + a random projection from IR- to R".
y f random vector from {+1}P.
(MD < (T...).
if {M,) is linearly separable then cnt < cnt + 1.
endie if end for
return cnt/ m.
end function
18

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
First, find orthogonal basis {e in R- for the basis vectors of all the data manifolds. Namely, for each E [P]. there exists a linear transformation Q' E ((D++I)xN such that uj. = Σx Q ej for each j {o,1,...,Dii},
Next, for each [PM/ Pl. generate vi..... Vy NN IN,) independently and let V be the Nu x matrix with vj on its columns.
Define MM (V-Q)) + Ei ,v(("*)), : s= (pp. E S.} as the i-th manifold the a-th cloud where (V"Q'): En vgQu for every a E [PM/P] P] and i E [P].
Now, we are ready to formally define the mean-field version of manifold capacity.
Definition B.3 (Mean-field manifold capacity (Chung et al., 2018; Chou et al., 2025)). Let {Mihejry be collection of data manifolds R as defined in Equation 3. The manifold capacity of {MM:Netp) is defined as
PM
ME PoE(P)/P).iE[P]) min AMMS Mi (0,6)
ox.(1)
OM
lim max N00 Ps
Pr y.Mu(Pu,NN))
NM
where and ON2( (1) + as NM + co.
Finally, previous work (Chung et al., 2018; Chou et al., 2025) derived a formula for mean-field manifold capacity as follows.
OM
P
E HiIE (O,()
max S.EM.
|Iproloone ((v...)) 113]
(4)
((-T-22 E Ayys:
1 A.yis:l12 {
max s.E.M.
--1#13" (~N(0.Ix)
2:20
where (n, Σ) denotes the multivariate Gaussian distribution with mean μ and covariance and cone(.) is the convex cone spanned by the vectors, i.e., [one([y:s:]) = {E; Asyss A 2 0}.
B.4. Effective geometric measures from capacity formula
The advantages of mean-field manifold capacity are: (i) ay can estimated via solving a quadratic program (Algorithm 2) and (ii) Equation connects manifold capacity to the structure of the manifolds (Mi). Specifically, for each y, t, (s, t)} = ys arg max [|projcone((uss)) tll as the anchor points with respect y and T. Intuitively, these anchor points are the support vectors with respect
random projection and dichotomy as in Definition 2.1. Specifically, these anchor points are analytically linked to manifold capacity via Equation I and are distributed over the manifolds (M,, This connection inspired the previous work (Chung et al., 2018; Chou et al., 2025) to define the following effective manifold geometric measures that capture the structure of manifolds while being analytically connected to capacity.
The first key idea of defining effective geometric measure is the segregation of anchor points into their center part and their axis part. Concretely, for each [P), define sQ := Ey. (y, (y, t.)] as the center of the i-th manifold and define is [y, := s, t) s 5e be the axis part of sy t) for each pair of (y,t) t). Next, (Chung al., 2018) used identity:
L and set
||rooconee (())) tl) and
(oro)con(((s!((... t rewrite the capacity formula (Equation as follows.
0'M
(px
y.t
(ss
IE P y,t
pro (1 (y.4)) t (Ipro)
1
|prooccoot +
19

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Then, they proceeded with the following approximation.
t Ey.t
(s (x).))))) +
(5)
Ey,t 1+ [ 2 (p)
(Chung et al., 2018; Chou et al., 2025) found that the above approximation empirically performs well. Furthermore, as the numerator mimics the notion of Gaussian width of a convex body and the denominator behaves like (normalized) radius of a sphere, they defined effective manifold dimension and radius as follows.
Definition B.4 (Effective manifold geometric measures (Chung et al., 2018; Chou et al., 2025)). For each E [P]. define sf = E (s: t]] as the center of the i-th manifold and define s! (y, t) s; -s. be the axis part of s; t) for each pair of (y, t).
Manifold dimension captures the degree of freedom of the noises/variations within the manifolds. Formally, is defined as DM == ly,, [[ll:rrrlloe(((e(tx))))).)) +1]).
Manifold radius captures the noise-to- signal ratio of the manifolds. Formally, it is defiend as Ry :=
proj fe, ()))))))
Fyy.
core (x))))) (3))
Center alignment captures the correlation between the center of different manifolds. Formally, it is defined as ;= = P(P-I) Eids Ks I-)--
Axis alignment captures the correlation between the axis of different manifolds. Formally, it is defined as Phr = P(P-T] Eizz [ysssss((y (y, t)>l}-
Center-axis alignment captures the correlation between the center and axis of different manifolds. Formally, is defined as Vhu == P) Eiz, Ey s) t))
capacity approximation formula by dimension and radius. Recall that Equation 5 previous work (Chung et al., 2018) used the identity a 1+ h-a to approximate the manifold capacity. After defining manifold dimension and radius, can then plug them back to Equation 5 and get the following approximation of manifold capacity via effective manifold dimension and radius.
1
R DM
O.M R
B.5. Connections between manifold capacity and its effective geometric measures
Here, we demonstrate the connections between manifold capacity and its effective geometric measures by synthetic manifolds. In particular, we consider isotropic Gaussian clouds parametrized by set of ground truth latent parameters: dimension Deround radius Rynund, center correlations ground, axis correlations Pground- and center-axis correlations grronnd- See Section D. for more details on the generative process. In this section, we focus on showing that the effective geometric measures DM, Ry, Pi Phg 1/M capture the corresponding ground truth parameter.
Effective manifold dimension and radius. We first set all the manifold correlations to be and vary the ground truth radius and dimension. Here we pick N = 1000 neurons, p = 2 manifold, M 200 points per manifold, varying the underlying dimension from 2 to 10, and varying the underlying radius from 0.8 to 2. In Figure 9, we vary the ground truth dimension in the x-axis, and in Figure 10, we vary the ground truth radius in the x-axis.
Effective alignment measures. Next, we fix the ground truth dimension to be Daround = and radius to be Rground = 1 and vary Peround Pe Pground Vground from 0 to 0.8. In Figure 11, we vary the center correlations, and in Figure 12, we vary the axis correlations.
B.6. Algorithms for estimating manifold capacity and effective geometric measure
We provide pseudocodes for estimating manifold capacity and effective geometric measure in Algorithm 2.
20

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Algorithm 2 Estimate manifold capacity and effective geometric measures
Input: {M,}: P point clouds, each containing M points in an N-dimensional ambient space; ny: number of samples for estimating the expectation.
Output: Og: Manifold capacity; Da: Effective dimension; Ry: Effective radius; phy Effective axis alignment; Phr Effective center alignment; t/y: Effective center-axis alignment.
% Step 1: Sample anchor points.
for from 1 to ni do
t - a vector sampled from isotropic N N-dimensional Gaussian distribution.
y + random dichotomy vector from (=1}".
A IN: q F -t; < ON-
G (y (M.) {Gij = yis is vector where is the j-th point in Mi.} output (p(A, q. G, h). {minx 1x fxxAxx Ax + x s.t. Gx s h.
Zdual f output[ "dual"] {The support vectors }
for from 1 to P do
[; < Σ, (Zinaa)) T G/ E/ (Zdma))
end for end for
% Step 2: Estimate (anchor) manifold centers.
for i from to P do
s + 1 E [k])).
end for G + s {Anchor center gram matrix.)
% Step 3: Separate the center and axis part of anchor points.
for 1 from 1 to n, do
for from to P do
[A] < s: [k]
s.-
(The axis part of the anchor poitn in the i-th manifold.)
end for
t [k] + Σ s tk]
G [M]
s (B)) {Anchor axis gram matrix.)
end for
% Step 4: Estimate manifold capacity and effective geometric measures.
OM
+ 1
P.P.
Σ The
k= (s([k)tt)) (s: [k] (s, (xs))))
DM n.P
[x:] G (R)
RM ni Dk=1 [ (G)+)
(Equivalent to the definition of radius after applying the Wood-
bury formula for numerical stabiltiy.}
Pi P(F-IJ EP'
|
PM f- P(P-1) P Σ
$ [M] T
j
k=
| 18]
1/u P[P-1) DP
Ejti mi
(e.)
return aM, DM, RM, Pirs Pig. V'M.
21

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Capacity
Radius
Dimension
Center Alignment
Axis Alignment
Conter-Axis Alignment
Radius
100
(Theory) Dimerstoo
1.0
1.0
1.0
10.0
1.0
0.6
0.6
0.0
0.6
1.4
0.4
1,8
2.5
0.2
0.2
2.0
0.0
0.0
0.0
0.0
0.0
10
10
10
10
(Ground trutch) Dimension
(Ground trutch) Dimension
(Ground trutch) Dimension
(Ground trutch) Dimension
(Ground trutch) Dimension
IGround trutch) Dimension
Figure 9. Effective manifold dimension tracks the ground truth dimension of uncorrelated isotropic Gaussian clouds. Note that the higher the dimension, the smaller capacity, as discussed in Figure 2c.
Capacity
Radius
Dimension
Center Alignment
Axis Alignment
Center-Axis Alignment
Dimension
2,0
0.8
1.5
0.8
0.8
0.8
4.0
0.6
0.6
0.6
(Thory)
0.4
0.4
8.0
10.0
0.2
0.5
0.2
0.2
0.2
0.0
0.0
0.0
1.0
1.5
1.5
(Ground trutch) Radius
(Ground trutch) Radius
(Ground trutch) Radius
(Ground trutch) Radius
(Ground trutch! Radius
(Ground trutch) Radius
Figure 10. Effective manifold radius tracks the ground truth radius of uncorrelated isotropic Gaussian clouds. Note that the higher the radius, the smaller capacity, as discussed in Figure 2c.
C. Theoretical Results
We consider the training of a fully-connected 2-layer network of the form f(x) = V/y o( x), where x E I2 an input, W Wdd is the hidden layer matrix, E RV is the readout weight, and σ R + is the (non-linear) activation function. To study feature learning in this setting, it common to consider W be randomly initialized (i.e.. random feature model (Rahimi & Recht, 2007)) and update via gradient descent with squared loss. Meanwhile, the readout weight a is randomly initialized and fixed to avoid lazy learning (where the network minimally adjusts the hidden layer and focuses on learning a good readout weight) as well as enable mathematical analysis (Ba et al., 2022). Input data and label (X), ) (X MPa..) were randomly generated by a fancher-student setting, where there is a hidden signal direction 3 that correlates with the label (see Setting C.2 for the full setting). As previously proved in (Ba et al., 2022) (see Proposition C.5), in the proportional asymptotic limit (i.e., Purann d, N - do at the same rate), the first-step gradient update can be approximated by rank-1 matrix that contains label information, resulting in the updated weight to be more aligned with the hidden signal 8r. Hence, in this setting, the learning rate 7) can be used as the ground-trui to measure the amount of task-relevant information (i.e., richness in learning) in the model representation after gradient updates.
We extend the previous results in (Ba et al., 2022) from a regression setting to a classification setting. Specifically, We prove that capacity correctly tracks the effective degree of richness after one gradient step"'. Moreover, we derive a monotone connection between capacity and prediction accuracy, thereby justifying the use of capacity as a measure of richness in feature learning within a well-studied theoretical setting. Here, we provide an informal statement of our results and leave the formal version and proof in Appendix C.
C.1. Formal theorem statement
Let be the input dimension and N E NN be the number of hidden units. Let Wo E Nxd be the weight matrix of a fully connected 2-layer neural network. The feature of an input vector defined as ox(x) (Wox) where a(.) IK R is non-linear activation function, e.g., ReLU tanh. The readout weight is denoted 2 2-layer NN is the sign of the readout, i.e., fx) = sgn(a" (xx)).
R. Finally, the output of the
Here we follow the convention in (Ba et al., 2022) and study only the first gradient step the key Gaussian equivalence step might not hold for more steps as remarked in footnote 2 of (Ba et al., 2022).
22

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Capacity
Radiuss
Dimension
Center Alignment
Axis Alignment
Conter-Axis Alignment
10
1.0
1.0
1.0
correlations
10.0
Center alighment
0.2
0.6
(Theary)
0.6
0.6
0.6
0.6
2.5
0.2
0.2
1,0
0.0
0.0
0.0
0.0
0.0
0.5
0.0
0.5
0.0
0.0
1.0
0.0
1.0
0.0
1.0
[Ground trutch) Center correlations
(Ground trutch) Center correlations
(Ground trutch) Center correlations
(Ground trutch) Center correlations
(Ground trutch) Center correlations
(Ground trutch) Center correlations
Figure 11. Effective manifold center alignment tracks the ground truth center correlations of isotropic Gaussian clouds. Note that the higher the center alignment, the smaller capacity, as discussed in Figure 2c. Also, in the large center correlations regime, the effective radius increases.
Capacity
Radius
Dimension
Center Alignment
Axis Alignment
Center-Axis Alignment
10
1.0
1.0
1.0
10.0
0.8
0.0
7.5
0.6
0.6
0.6
0.4
0.4
0.4
0.4
0.8
0.2
0.99
0.0
0.0
0.0
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
3.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
(Ground trutch) Axes correlations
(Ground trutch) Axes correlations
(Ground trutch) Axes correlations
(Ground trutch) Axes correlations
(Ground trutch) Axes correlations
(Ground trutch)
Axes correlations
Figure 12. Effective manifold axis alignment tracks the ground truth axis correlations of isotropic Gaussian clouds. Note that the higher the axis aligmment, the higher capacity, discussed in Figure 2c. Also, in the large axis correlations regime, the effective dimension decreases.
Let {(x1; [E))sE[Pmm be the collection of training data. We consider gradient descent over the mean square error (MSE) of the 2-layer NN, i.e., L(f) Pan Eie(pom ((f(x;;; ys) where ((2;, ys) } y) The gradient update with learning rate 17 > 0 Wo+1 W, + r)Cre where
G: = Purin sE(Pmm)
Σ
[(u - a To (W[xx)) (Wxxx) x
and d-(-) denotes the first order derivative of o(-). Assumption C.1. We adopt the following assumptions used in (Montanari et al., 2025; Ba al., 2022).
1. (Proportional limit) Pimin d, N + dc with ei = N/d, 4/2 = Prrain/d dd and 0 < 4/i, 1/2 < co.
2. (Gaussian initialization) [Wola] ~ NO0, 1/N) for each h [N] and j E [d].
3. (Gaussian readout) ag (O, 1/N) for each k E [N].
4. (Normalized activation) The non non-linear activation function a(-) has O(1)- first three derivatives almost surely. In addition, [[)( and (G)) # for G ~ (0,
5. (Non-degenerate label function) Let F R + [0, 1] be a continuous function satisfying
inf (a: Pr[T < Ø] > 0} = 0o and sup (a: Pr(T > x] > 0} =
where = G G and Pr[Y = || = 1- Pr(y = -1 G = F(G).
Setting C.2, We consider the following data generation process. Let F: RR- + 10, 1] be a function satisfying Assumption C.I. Let B. E d be a hidden vector with || = 1. The data distribution Dp(3_)) is defined by the following two steps: (i) sample x ~ (O, Ia), and (ii) sample y with Priy = = Pr(y = -1 F((,., (1((B..)) x)). Finally, the prediction accuracy of network defined as the expected accuracy of fresh sample, i.e., Pr(x.))~Do(a.) (f(x) 2 0).
23

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Parameter C.3. Given 11, 22 F, Ba from Assumption C.I and Setting C.2. We define the following parameters.
71
Go(())
G-N10,11
1/2
E
[o(G)²
E
[o(
G-
Gm
0
E
X~p01
7 + 12
7 X
02 = 1
E
7f + +72
Xmpp
03 E
[YG]
(G, DF
01 = V22
E (GYY) D
YY'GG"]
where Pdt is the Marchenko-Pastur distribution with the ratio parameter being thi and (G,Y) ~ Dr is defined as the sampling process: G (O, 1) and Pr[Y 1] = Pr|Y = -1] = FG)).
Theorem C.4. Given Assumption C.I and consider O < 1h 122, 1] < oc.
1. (Capacity tracks the degree of feature learning) The storage capacity of 2-layer network trained with synthetic data defined in Setting C.2 after one gradient step is Pron (1)1, th2, 11) and
Punse d.N-oo
Ok Prann N (4/1, m)
a a(vi, 122
Here the function oc(-) is defined as
((22,,22,71
min CER
E (2.G,)) 21
1 -2)6
where (Z, G) ~ Des wain vann is defined as the following sampling process
~ 1),, G ~ 1), Pr(Y
Pr(Y = =-1 ( (G)
and the scalar function f-(-) and T(V1, 22, 11) are defined as
f+(G) G'~N((011)
F(VI-PG+T00
and
T1P1, 1/2 v To(Vi, 22)2 Ta (1 1)
where TO(-) and T-(+) are scalar functions defined as
To(N1, 22)), = - -0
and
TA (1h,, the n) 1120 - 02)200
1 -//2 -2)00
where the parameters 0,'s are defined in Parameter C.3. In particular, < aldi,, v2,, a(1/1,3/2211 for all
O < n o<n...
2. (Capucity analytically links to prediction accuracy) The prediction accuracy of 2-layer network trained with synthetic data defined in Setting C.2 after one gradient step is ACCp (,, 122, 11) and
ACC P. d.N (1/, 122, 1) Paita -
Acc 1hi, 1/2, 11)
24

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Here the function Aco(-) is defined as
Acd(u11 42,,1)
YG
(G.Y)DDF
+ 7 7
In particular, there exists an increasing and invertible function g0r.v=
[0, 1] + RR such that
Acd(1/ 1/2 n) = 9wi. (1(1/1 (1,,,4121111) 4))
C.2. Proof for Theorem Theorem C.4
Step I: Rank-1 approximation of gradient descent in 2-layer networks by ref. (Ba et al., 2022). When the learning rate is constant, i.c., (1), ref. (Ba et al., 2022) shows that the gradient update matrix can be approximated by rank-1 matrix. In particular, the following is restatement of Proposition
(Ba et al., 2022).
Proposition C.5 (Proposition 2 in (Ba et al., 2022)). Given Assumption C.I and Setting C.2, there exist some constants
C' > O such that for all large Puiins N, d, the following holds
Go 71a
)
C Pirain /P
-|Gool
Prroin
with probability at least 1 P Piraint clog² Puse and
|| |||| denotes the operator norm.
Step 2: A formula for the storage capacity of Gaussian model by ref. (Montanari et al, 2025). The storage capacity of Gaussian model is proven in (Montanari et al., 2025). In particular, the following is restatement of the Proposition 5.1 in (Montanari et al., 2025).
Definition C.6 (Gaussian model). Let 02 R be some latent vector. A sample (x,, y.) E R x {+1} is ii.d. sampled as follows. First, sample x; from (0, Σ) where Σ is covariance matrix satisfying certain technical condition defined in Assumption 1-2 in (Montanari et al., 2025). Next, let y; +1 with probability f((,,, xx)) for some function f satisfying Assumption 3 in (Montanari et al., 2025).
Proposition C.7 (Theorem in (Montanari et al., 2025)). Consider a Gaussian model satisfying Definition C.6. As Priiin N, - oo, the storage capacity converges to
1(-crG- - 22)
a" min E (2,G,Y)DD,
where (Z, ~ Dy is defined as the following sampling process
2 ~ N(0,1, G ~ 1(0, Pr(Y 1| = =-1==(pp G).
where p is some scalar related to the Gaussian model as defined in Assumption 2 (Montanari et al., 2025).
Note that the capacity only depends on the alignment between data and task (as encoded in f) and does depend on the covariance structure. The dependence the covariance structure will appear when considers the non-zero margin version of capacity.
Step 3:
A Gaussian equivalent model for 2-layer NNs after one gradient step. Next, combine Gaussian equivalent model for random feature 2-layer NNs in (Montanari et al., 2025) (Theorem 3) and the rank- approximation of gradient step in Proposition C.5 to get a Gaussian equivalent model for 2-layer NNs after one gradient step.
Proposition C.8. Given Assumption C.I and 1/,, 122, 17 < oo. Let E N and (W1, Bo, F) be the weight matrix, hidden
25

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
vector, and label function from Setting C.2. Let a P. (1/1, 42 v) be the capacity of the following Gaussian model:
GM
Ed, wW + +221
0. d.y "d. ( + +121)-1W1,1
an,1 3/ WT (WT +
2
Td.1 -aå,
(I) = E [F(an + Ta.,
G)N1(,,1)
We have that
lim
a d, 22 7)) - d, (v1, 2/2, n))
QM
d. +00
and
a Pm N (1/1, 8/2, 1))
GM
Pinin d.N- a(4', $2,
Here the function ax-)) is defined as where the scalar function (+) and r(th1, 222, n) are defined as
alvi, 1/2.1)
min CER (2,.GY)~DD (2.0
1<-cGG- 2)
f+(G)= G',N((0,1)
E'
[FWII
and
T T(1,, 122
lim Td.
V Vo
To(111112) 2122
Ta (4 1)2
d
where (v vi2) = lim Td.0-
To derive the Gaussian equivalent model Proposition C.8 of the random features model after one gradient step defined in Setting C.2, we analyze the following random features and their associated labels:
(x(),) a (Wix,), Pr(y= 1 = Pr(ys -1 1xx = F((B., xi)), ||3. =1
where x; ~ (O, Ta) and W = Wo + Gr while Go satisfies the bound given in Proposition C.5. Given the assumptions in Assumption C.I, we can decompose the nonlinear activation function into Hermite polynomials. Following our parameters Parameter C.3, define the Gaussian equivalent features of our model as the linearization of Equation C.2:
g: 71 Wix; + h2
where h, (O, IN) are independent from everything else. Now, we wish to find a similar linearized Gaussian model for the labels yi given the Gaussian equivalent features g:. is easy check that the Gaussian features has the following covariance:
gi (O, Ea,y), Ean W.W. + 7
By matching covariance through Equation C.2, obtain
x; T1 2 1
where TN + W W1)-1 and h NO0 IN) independent of x,. Therefore, rewrite the label
function parameter as
(8., x,) = odd,, dd.... gi) +
where E; NN0 -2.,) are independent of g: Effectively, obtain equivalent label function
fan(x)) F
[F(aa, + G),
G~
such that Pr(y = = 1- Priy = - 1|x,] = fa, 8.)). It is easy to verify that this Gaussian model satisfies the assumptions in Definition C.6.
26

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Step 4: Analysis o. Finally, combine Proposition C.5 and Proposition C.8 to get the formula for the right hand side of Equation 7. From Proposition C.5, we approximate Wi as Wi Wo au where 7 E, vıx Purain- To rewrite the right hand side of Equation 7, we first deal with the matrix inverse term using the same trick as in ref. (Ba et al., 2022). Let Σ, = W," +72I.
Observe that
Σ = +
c]
L [
a
0
c
where c Wou. By Sherman- -Moriion- -ooddburyy formula, we have
ΣT = Eg - -1 a cl
1 0
+
2
a
2
a
c|
c
a
20
c
2g- Dan - Do0 + Aac + Aen
where
La L1
Aac
2
20 aa E
D
La
ACc
D
0
cc Ej¹
1 + LG 1
AaC
11
2g ac 20 221
D
2 1 + Lo 5- ca T
A 71
D
and
Lo 3. Wo Wo3.
L1 u T u
La u TT..
L3 7/a TaTTTp
LA cfc 2
Ls 3/c WoB.
La 7/a Enit
L a Tc
Ls 2ia T Eo WoB.
D (I4 In) - (1 + LG)²
Thus, can rewrite the right hand side of Equation follows.
Ta 1 - (Wo + au T (Wo + au )
'Wo
+ 7 (W, + au AT Aaa (Wo au )3a,
+ 7,:87 (W, + +Wu+au)) ) (W, + au T )
7.... + au ) A (W + au )3,
77.3. (W, + au Tu Tea (Wo + au ( )
- Lo - L - 21
L1 L1 (Iz + Lg)²
+
La (Ls + L228)²
+ La
(L2Læ La)(Ls L2T0).
27

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Similar to Proposition 29 in (Ba al., 2022), by Hanson-Wright inequality, have that Lg; Lgs L7
Lu + 2
L1 - rf01
L 1/03
La - 71 Xm-Po1
2
01
71-
72_
2
X
La + 1211
11
E
120000
Xmppi
!71 X +
7/
X
Lo + 99,003
11
E
100203
X
7
72
L8, Lr, L8 +
D + La([T4 - [1) - - - (02 1)0+
To sum up, we have
1 - 02 1120, 020,03((70), (02 - 1)04
lim Jd.
d-300
(02 1)0 -1
+ 1703(022 - 1)0301
0 (02 1)04 -
+ 01 010203
1 (02 1)0g
r)²0 0202
2
1 1)01 -
1²²0( (1 02)203
1 02 1 + 0/ (1 - 02)
This completes the proof for the first part of Theorem C.4.
Step 5: Analysis for prediction accuracy. Recall from Setting C.2 the definition of prediction accuracy of the network after gradient step is P((x))))DD(e.)lvva (ooeei/))aaToWWwW) (Wix) 2 0j. By Gaussian equivalence and Proposition C.5, have that the following.
Acc AcCPo (1 11)
Pr
[ya (WIxx) 2 0].
(x.)))Dr((.. a.Wi
By Proposition C.5, we can further approximate the equation as follows.
Pr (x.)) a.Wa.u
yya o(WW au T)x) 2 0] + o())
By Gaussian equivalence, we can further approximate the equation as follows.
Pr (x.)))Dx((L) W.
yya T (7 + au + 7. W.)x) 2 0] + o(1)
where W, E Nxd and ([.,k, ~ NO 1/N)) for each & E [N].j E [d]. Note that as a, Wo. W, are independent, we can
further simplify the equation as follows.
Pr (x.)) a. W.
W7 u x + V28 + 7 ya T W + o(1) 2 0 +(()
28

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
where W E Rgxx and ([W))a (N0,, 1/N)) for each k E [ E [d] Note that as a, W are independent, we can further simplify the equation as follows.
E loy' TTo x] + ynn + + + +(1 2 0 20 + t11)
Pr (x.))Dp((3) (0.1(
mi
(x
Note that by decomposing and x' to direction that's parallel to 3. and orthogonal to B., further simplify the equation follows.
mi
[YY'GG"] + v1/122)) + yn + 22 x2 ++ o =o1) + +o(1)
Pr (G 2, 2.01)
E 1G'Y"'>Dr
P
Pr (G.YinDF 2~A((011)
mig
+ 12 22 2+ o 2 0 +0(1)
422
ni0
YC
+0(1)- -
E (G.Y)~DF
o + y
Note that when fixing 1 2/2 and non-trivial F, both capacity formula and prediction accuracy formula are increasing and invertible with respect to 17. As consequence, the two quantities are also analytically connected by an increasing and invertible function. This completes the proof for the second part of Theorem C.4. We also provide numeric checks for the formulas in Figure 13.
B
3.1
U: 0/
Prediction Accuracy
Capacity
2.
27
0.50
0.43
3
Figure 13. Numerical checks for the formulas in Theorem C.4. We run the simulation with d = 2000, vi = 1, ReLU activation, and label function f(z) I=e-xæ for 50 repetitions. Left: numerical checks for the capacity formula. Right: numerical checks for the prediction accuracy formula.
D. 2-Layer Non-linear Neural Networks
In this paper, use 2-layer non-linear neural networks and Gaussian mixture models (for input data generation) convenient experimental setup to systematically explore different regimes in feature learning. Moreover, given its medium level of complexity, it might be possible to have an analytical characterization of our numerical findings, and we leave it as an interesting future direction.
D.1. Experimental setup
D. SYNTHETIC DATA GENERATION
We focus on point manifold, which consists of data points associated with the same label. As discussed in the previous section, we are particularly interested the effective radius, dimension, center alignment, axes alignment, and center-axes alignment of the representation manifolds. Therefore, we consider synthetic model to generate training and test data with relevant geometric interpretations. Namely, construct P synthetic data manifolds with radius R E RR,, intrinsic
29

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
dimension D N, size M N. The manifold layouts further determined by center correlation strength pC [0, 1), axes correlation strength Pa [0, and center-axes correlation strength vó [0, 1), all of which would detail in the following subsections.
Isotropic spherical manifolds. First, we consider the simplest case: manifolds with isotropic Gaussian center distribution and axes distribution with no correlations. This is the scenario considered in Section and Section 4.
Let E be the dimension the data. We consider point manifolds {Mi)ietp) with manifold size M E N and radius R that lies in subspace of dimension D. Each manifold defined as
M; (uo + R. Σ of u eVEckE]M)
where the axes uj (0, La/dd, the coordinates st ~ N (0, 1), the noise vectors V& ~ (0,, La/d), and e == 10-2. The pre-scaled points in the manifolds {Ej=1 of [e]M]]][]]]]] are well-normalized to unit norm.
Test manifolds share the same model except that the noise vectors V, are sampled again in the same distribution.
Isotropic Gaussian manifolds. In certain experiments, we drop in the intrinsic dimension D and directly consider manifolds defined as
M₁ (mo R kE(M)
where the noise vectors are V& N(, La/d). Test manifolds share the same model except that the noise vectors Va sampled again in the same distribution.
To generated correlated manifolds, we consider an auto-regressive model described by the covariance matrix C' (p) E TR. PxP, where E [0, 1) is either the center correlation strength PC or axes correla- tion strength PA The center covariance Cc is then mixed into the isotropic manifold centers {n} ~ (0, Ia/ /d) JE(M) The axes covariance matrices C² mixed into the isotropic axes (u: N(, Ia/ dd) IiE[M for each = 2, .... respectively. The mixing is performed through multiplying the column matrix Mc or M R of centers or each axes with the Cholesky decomposition of Cc or C²A, To incorporate center-axes correlation, we scale each center vector uo by a factor of (1 + &b q) where ~ (0,
Correlated spherical manifolds.
Labels. For P manifolds with manifold size M, the P labels are randomly sampled from a uniform distribution on (+1}. Each label is associated with M data points in the individual manifold. When learning with binary cross entropy, the labels reassigned as {0, 1} during loss and gradient computation.
D.1.2. 2-LAYER NEURAL NETWORK ARCHITECTURE
The model architecture we consider is similar to the architecture mentioned in Appendix C.
Let E N be the input data dimension, N E N be the number of hidden neurons, K E N be the number of linear readouts, + be the scaling factor of the readout weights.
Let W Wo e Nxd
x be the initial weight matrix of a fully connected 2-layer neural network. Let (agher be list of initial readout weights where of NNN. Let (+) IR IR be non-linear activation function, e.g. ReLU or tanh.
The feature of an input vector is defined as (xx) a (Wx). The 2-layer neural network parameterized by W and of is defined as
f(W, (W, = a a ((x)
where the label prediction for data point x is sgn(f(x)) when learning with the mean squared error loss function. When learning with binary cross entropy loss function, we use {0, 1} as labels and <xf(x)) as prediction instead, where ç the standard sigmoid function.
30

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
D. 1.3. LEARNING RULE
Loss funccion and gradient update. Let R+ be the learning rate of the weight matrix, 1R+, be the scaling factor of the readout learning rate, and let {(x, [e(aa be the collection of training data, where P is the number of manifolds and M is the manifold size.
We consider gradient descent over the loss function
C(f) = a² PM
Σ {(f(x)).
sE[PM]
where L R x {+1} + R is either the mean squared error (MSE)
1
(MSE((, y) 2 (2- --y) y)²
o x 10 - I is the binary cross entropy (BCE)
(BCE y) = y log(1 + =) +)+ (1 - y) log(1 e²)
Mean squared error.
For the weight matrix, the gradient update with learning rate > is W4+1 = W, + GG where
G. -
1 Σ
Σ
(us
VN at o(Wx))) a at 2oo x; VN
PM PM
K
1E|PM]
JE[K]
and d'(-) denotes the first order derivative of o()). For each linear readout, the gradient update is ac+1 aj eng. where
9t
1
PM
Σ iE[PM]
y/s
VN
a
(Wix,)
VN (Wix))
Note that the a multiplier on the loss function to ensure common convergence time when a - oo as mentioned in (Geiger et al., 2020).
Binary cross entropy.
For the weight matrix, the gradient update with learning rate > i W4+1 = W2 + n/G; where
G a PM
1 1 Σ Σ
Σ
(u cf af (WWx))))
VN af @o(())))
x T
K
(E[PM]
НЕ[K]
where denotes the standard sigmoid function and denotes the activation function. For each linear readout, the gradient update is 0++1 = + CT]9e where
1
Σ
a a (Wixi)] -
a VN (Wix.)
9t
a²
PM
Di
-
(E[PM]
If not otherwise noted, we conduct experiments with the MSE loss function and ReLU activation function by default.
A Note Learning rate. We define -1 as the normalized effective learning rate. During training, We implicitly scale the learning rate 7 by factor of VN in the experiments to enter the rich regime as mentioned in (Ba et al., 2022).
1.1.4. TRAINING
For each 2-layer neural network experiment conducted in the paper, forty random seeds are chosen from 0 to 39000 with an interval of 1000 to train forty models in parallel for 10a epochs. All training are conducted on the Flatiron Institute high performance computing clusters.
31

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
D.1.5. FEATURE EXTRACTION
During analysis, fifty epochs are sampled uniformly in log-scale. For each model checkpoint epoch t, we extract total P size M manifold representations (x,())) associated with labels (3 We perform conventional analysis and manifold capacity analysis described in Appendix A and Appendix B respectively. We will present more details in the following experiment sections.
D.2. Capacity is robust measure of feature learning across architecture, data, and learning rule variations
The purpose of this section is to support Section 3 by showcasing that capacity is able to quantify feature learning when model architecture, data distribution, and learning rule varies.
D.2.1. FEATURE ANALYSIS METHODS
Here, we briefly present the conventional feature analysis methods and capacity analysis method and how they are computed the experimental setup.
Representation level analysis. Activation stability is a representation level metric that intuitively captures how much neurons activated in hidden units. Formally, define it
ΣPM PM Σ5= (x,())(x))
PMN
Another conventional method to disentangle feature learning at representation level is tracking the norm of deviation from initial weights (Jacot et al., 2018)
[ Woll
Woll
On the other hand, the cosine similarity (Liu et al., 2024) can be used to study alignment at representation level
2
| Poll
where (De)is (x,) (x(x)) PMPPM is the gram matrix of features over the test data.
Kernel methods. The kernel methods for quantifying feature learning involves computing the Neural Tangent Kernel (NTK) (Jacot al., 2018) for each pair of test data points:
+ O x2) V w. f(1) V (xe)
where w., f denotes the total gradient of the neural network at epoch & with respect to the hidden weights W, and readout weights }. Note that we scale the readout contribution to the total gradient by the readout learning rate factor e R4 aforementioned. Hence,
K
1 K
V..((x)= Vw.f(x)
oe fxx)
After obtaining the gram matrix θ, = O,(x;, E PPPMPPM from the test data, we can compute the NTK change defined as
0 10,-00
which can be interpreted the relative deviation of the the kernel from initialization in the Frobenius norm metric. Conventionally studied, NTK change disentangles lazy and feature learning, as detailed in (Jacot et al., 2018). We present NTK change in Section 3 Figure to compare it with capacity as the metric to track feature learning.
The kernel alignment can be similarly defined as the cosine similarity of initial and current NTK gram matrices:
0,Op
||o.1|| 0
32

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
which can be interpreted the relative deviation of the kernel from initialization in terms of alignment. Kernel alignment also studied in (Liu et al., 2024) disentangle lazy and feature learning.
The centered kernel alignment (Kornblihh et al., 2019) is another approximation method to study kernel evolution when the gram matrices is large:
HSIC(O,, 9b) O.))CIQ,, 0b)
where
or(),
HSTC = (n 1)²
These kernel metrics can be readily computed from the trained models and extracted features.
Capacity and effective geometry. For more details on data-driven manifold capacity analysis, please refer to Appendix
Setup Figure 4a. In Figure 4a, we showcase that the degree of feature learning controlled by the effective learning rate 7j with the following standard setup:
Data: Isotropic Gaussian manifolds with R 0.5, .M- = 15.
Model: We set ReLU, N 1500, 1000, P 100, K K=1
Learning rule: We set [MSE,1) 50, c = 0 and
10/128, 10/ 112, 10/96, 10/80, 10/64, 10/16, 10/4, 10/1 that the normalized effective learning rates
fj 128, 112, 96, 80, 64, 16, 4,
which is computed by 7j n0 where the division by 5 normalizes the smallest n7 to be 1.
Training: We trained the models for 100000 epochs with 40 repetitions per parameter combination.
Plotting: We use sample mean and 95% confidence interval for each data point.
D.3. Effective geometry reveals distinct learning dynamics
1. LEARNING STRATEGIES
Compression strategy setup In Figure 5b where the networks performs the compression strategy, we use a difficult-task setup with higher data manifold radius and more readout tasks:
Data: Isotropic spherical manifolds with At = 1.0, D = M = 15.
Model: We set ReLU, N 300, d 200, 20, K 27.
Learning rule: we set f fMSE, a 1, 0 and
7 10, 30, 40, 50, 60, 70, 90, 100, 110, 120, 130, 140, 150
that the normalized effective learning rates
1, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150.
Training: We trained the models for 100000 epochs with 40 repetitions per parameter combination.
Plotting: We use sample mean for each data point.
33

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Flattening strategy setup. In Figure 5b where the networks performs the flattening strategy, we use easy-task setup with smaller data manifold radius and very few readout tasks:
Data: Isotropic spherical manifolds with R 0.5, D 8, M 15.
Model: We set = ReLU, = 300, d = 200, P = 20, K = 3.
Learning rule: we set l ÉMSE, a = 1, c= 0 and
1 80, 90, 100, 110, 120, 130, 140, 150, 160, 170
so that the normalized effective learning rates are
80, 90, 100, 110, 120, 130, 140, 150, 160, 170.
Training: We trained the models for 100000 epochs with 40 repetitions per parameter combination.
Plotting: We use sample mean for each data point.
Contour plot of learning strategies. In Figure 5b and c, contour plots to visualize the different learning strategies adopted by the network. We use Equation 34 in (Chung et al., 2018 8) to approximate capacity using effective radius and dimension:
R
Das
The scatter points with the same color correspond model trained with the same normalized effective learning ij over different epochs.
D.3.2. LEARNING STAGES
Setup. In Figure 5a, we adopt setup with moderate radius and number of readout tasks that shows clean learning stages:
Data: Isotropic spherical manifolds with It = D 8, M = 15.
Model: We set = ReLU, N 300, d 200, P = 20, K = 5.
Learning rule: set f fMSE, 1) 10, 1,c so that the normalized effective learning rate is 10.
Training: We trained the models for 100000 epochs with 40 repetitions per parameter combination.
Plotting: We use sample mean for each data point.
E. Deep Neural Networks
E.1. Experimental setup
In this section, provide detailed information about the experimental setup for deep neural networks, including model architectures, datasets, training procedure, and manifold capacity measurements.
E. 1. MODELS
We use the VGG-11 models (Simonyan & Zisserman, 2015) for experimental results in the main paper. We also repeat these experiments on ResNet-18 (He et al., 2016). The specific implementation follows a similar setting in (Chizat et al., 2019) and is adapted from https:// htopps///0
btttps://github... thubthb..... com/ http://doarrdollll///aaa--------
34

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Output rescaling. As previously studied in (Chizat et al., 2019), multiplying the model outputs by large scaling factor can induce lazy learning (we use the notation / instead of in (Chizat al., 2019) avoid confusion with the notation capacity in Equation (2) ). In this section, we use the inverse scaling factor B as the parameter control the degree of feature learning. We define the normalized effective learning rate πj g_1.. We also note several adjustments to the
common training framework to adapt to using the inverse scaling factor 9 as the parameter to control the degree of feature learning.
Rescaled loss function: To adjust for using the scaling factor 8, we use the rescaled loss function Lg = 1 with denotes the loss function to accommodate for the time parameterization of the loss dynamic for large 3 as previously indicated in (Chizat et al., 2019) and (Geiger et al., 2020).
3
Model's s initial outputs as 0: As mentioned in (Chizat et al., 2019), for the scaling factor 3 to be able to control the rate of feature learning, the model output as initialization f (W)) must be equal 0. To ensure this condition, we set s(W,) h(W,) hw)) with We be the model's weight at training step h be the output of the network, and f be the final adjusted network output.
Number of repetitions. All model measurements (train accuracy, test accuracy, activation stability, etc.) are reported the mean of 5 independently trained model (with different random seeds). The error bar indicates the bootstraped 95% confidence interval calculated using seaborn. lineplot (errorbar- ('ci', ).
1.2. DATASET
In this section, we list detailed information about the dataset used in the paper.
The CIFAR-10 dataset (Krizhevsky & Hinton, 2009) consists 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images.
CIFAR-10.
CIFAR-100. The CIFAR-100 dataset (Krizhevsky & Hinton, 2009) is similar to CIFAR-IO, except that has 100 classes containing 600 images each. There 500 training images and 100 testing images per class. Note that the images in CIFAR-10 and CIFAR-100 are mutually exclusive.
The COF dataset (Hendrycks & Dietterich, 2018) includes images from the CIFAR-10 10 evaluation set with common corruptions such as Gaussian noise, fog, motion blur, etc. The dataset has 15 different common corruption types, and 5 different severity levels for each corruption type.
CIFAR-1OC.
1.3. TRAINING PROCEDURE
Loss function: We follow the theoretical results and practice used in (Chizat al., 2019) to mcan-squared error loss to train all DNNs mentioned in the paper.
Optimizer:
We use Stochastic
Gradient Descent
with
momentum
as (implemented
torch. opt im. SGD (momentum=0.. 9) ) to train the models.
Data augmentation: We apply the following data augmentation during training: RandomCrop (32, padding=4), Randomtiorizznttal....
Learning rate and learning schedule: We follow the practice in (Chizat et al., 2019) and set initial learning rate noo 1.0 for VGG-11 and yo = 0.2 for R8sNer-18. The learning rate schedule is defined as 7)- no
Initialization: We follow the practice in (Chizat et al., 2019) to initialize the model's weight using Xavier initialization (Glorot & Bengio, 2010) and the bias to be 0.
Batch size: We use batch size of 128 during training and batch size of 100 during evaluation.
35

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
E. 1.4. MANIFOLD CAPACITY MEASUREMENTS
In this section, we provide detailed information about how we define object manifolds from the model's representations and measure the manifold capacity and geometric properties (Chung et al., 2018).
Features extraction: For each image, we extract the object representation from the last linear layer (dimension 512) before the classification layer (dimension 10).
Number of manifolds: We use object manifolds for each measurement.
Number of points per manifold: For each object manifold, we randomly sample 50 images from the interested class.
Number of repetitions: Every capacity and geometry measurement is repeated 10 times per model instance (50 times if we have 5 model repetitions) and we report the mean and the error bar as the bootstraped 95% confidence interval calculated using seaborn. 1 lineplop ot ot (errorbar- r= (* ci'i, 95) )
E,2. Capacity quantifies the degree of feature learning in deep neural networks
Capacity and manifold geometry for VGG-11 models. In Figure 4, we show manifold capacity along with other common metries used to identify feature learning such as train accuracy, test accuracy, relative weight norm change, and activation stability. In this section, we provide other manifold geometrie measurements along with manifold capacity in Figure 14.
Radius
Dimension
Center Alignment
Axes Alignment
Center-Axes Alignment
1.0
0.40
0.35
70
0.8
0.30
05
0,13
9.6
1.5
0.20
0.3
0.33
0.2
1.0
0.18
0.2
0.5
0.667
0.333
0.5
0.05
0.1
0.0
0.0
0.00
0.0
epoch
epoch
epoch
epoch
epoch
Figure 14. Manifold capacity and geometry for VGG-11 models trained with different 77
Capacity quantifies the degree of feature learning in ResNet-18 models. In section Section 3, we show that manifold capacity can capture the degree feature learning in DNNs, specifically in VGG models. this section, we empirically show this statement can also be extended other model architectures, specifically ResNets, in Figure 15.
Capacity quantifies the degree of feature learning in VGG-11 models trained with weight regularizer. While most theoretical work in the lazy vs rich learning literature are formulated with vanilla mean squared error (MSE) loss (Jacot al., 2018) (Chizat et al., 2019), in practice, MSE with weight regularizer (or weight decay) is used widely prevent over-fitting and improve model generalization. In Figure 16, explore the effect of weight decay to feature learning and demonstrate empirically that capacity can still quantify the degree of feature learning in models trained with 1.2-regularizer. We implemented L2-regularizer by setting torch. opt im. SGD (wei fethttddecyy-o 0002). We leave further study about the impact between the magnitude of weight regularizer and effective learning rate (and/or scaling factor) to the degree of feature learning as a potential future direction.
E.3. Manifold capacity and manifold geometry delineate learning stages in deep neural networks
In section Section 4.2, we have demonstrated the use of effective manifold geometry to uncover hidden learning stages in 2-layer neural networks. In this section, we showed that using similar technique, we can also discover geometric learning stages in deep neural networks as well.
36

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Manifold capacity captures the degree of feature learning in ResNet-18
Train Accuracy
Test Accuracy
Activation Stability
Relative Norm Change
Capacity
100
60
0.8
0:175
f
0.01
0.1
83
0.6
0:125
03333
0.
0.109
40
0.4
0.073
08
1.0
0.657
20
0.
0.000
7.0
1.353
3.0
100
10t
0.0
0.000
epoch
epoch
epoch
epoch
epoch
Manifold geometry for ResNet-18 models trained with different richness in feature learning
Radius
Dimension
Center Alignment
Axes Alignment
Center-Axes Alignment
12
0.6
π
0.01
10
0.5
611
02
0.333
0.6
0
05
0.6
0.3
0.4
0.4
0.2
1,313
2.0
0.1
5.0
10'
12a
10²
15a
epoch
epoch
epoch
epoch
epoch
Figure 15. Manifold capacity and geometry of ResNet-18 models trained with different scale factor.
Experiment setup We used similar setup mentioned in Section E.1. In this section, give a higher resolution into the learning dynamic, we extracted the model checkpoint each training step (after each training batch, with batch.size=100 instead of each training epoch (after whole train dataset iteration).
E.4. Feature learning and downstream task: out-of-distribution generalization
In this section, we measure the performance of the models trained with different degree of feature learning (quantified by effective learning rate 77) on the downstream tasks for OOD using CIFAR-100, a dataset with no overlap with CIFAR-1O, the dataset used to train the model.
1. EXPERIMENTAL SETUP
We use linear probe (Alain & Bengio, 2016) on representation from the last linear layer (dimension 512) to measure the performance of models trained on CIFAR-10 on the out-of-distribution dataset, CIFAR- 100. Linear probes are linear classifiers trained on top of the representation to probe how much information the representations encode about a particular task characteristic. This approach has been used widely in different fields including natural language processing (Belinkov et al., 2017) and computer vision (Raghu et al., 2021).
Here we provide detailed information about how we construct the linear probes.
Optimizer. We use Adam optimizer with initial learning rate 1ju- = 0.1 and learning rate schedule is defined as T)- Other parameters are default P Pytorch parameters.
710
1+ j+'
Number of epochs. method below.
The linear probe is trained for 50 epochs, unless it is stopped early, as described by the early stop
37

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Manifold capacity captures the degree of feature learning in VGG-11 models trained with L2-regularizer
Train Accuracy
Test Accuracy
Activation Stability
Relative Norm Change
Capacity
109
1.0
80
0:175
0.13
0.8
0.150
6,333
60
0:125
6.867
6.1
40
40
0.4
0.675
0.300
6.133
6.353
0.2
20
0.
02
0.000
03
0.005
107
10- epoch
104
10a
101 epoch
10
10'
104
100 10"
0.0
10t
10
0.000
107
1st epoch
10
epoch
epoch
Manifold geometry for VGG-11 models with L2-regularizer trained with different different richness in feature learning
Radius
Dimension
Center Alignment
Axes Alignment
Center-Axes Alignment
0.40
7
25
88
0:1
0.8
0.30
0.5
0.033
0,33
0.5
0.20
0.3
0.133
0.2
1.0
0.333
6.5
0.10
6.667
0.5
0.3
0.0
0.0
10a
0.00
00
epoch
epoch
epoch
epoch
epoch
Figure 16. Manifold capacity and geometry of VGG-11 models with L2- 12-regularizer trained with different scale factor.
Early stop. During training, if the validation loss is greater than the minimum validation loss so far for more than Npatience epoch, then training is stopped. We set Npatience =3 3.
E.4.2. OOD PERFORMANCE FOR RESNET-18
In Section 5.2, we demonstrate how capacity and effective manifold geometry can be used to characterize the OOD performance of VGG-11 models trained with different effective learning rate 1). In this section, we show OOD performance and effective geometry of ResNet-18 18 models trained with different effective learning rate π in Figure 17. Interestingly, unlike VGG-11, for ResNet-18, 18, the failure of models in the ultra-rich regime is characterized by the expansion of manifold dimension, not manifold radius.
F. Recurrent Neural Networks
F.1. Experimental Setup
In this section, we provide detailed information about the experimental setup for recurrent neural network in Section 5.1, including model architectures, datasets, training procedure, and manifold capacity measurements.
F. DATASET
We used the package neurogym (Molano-Mazon et al., 2022) to simulate common cognitive tasks. In this paper, we trained recurrent neural networks to perform the following cognitive tasks: perceptual decision making, context decision making, and delay match sample. We followed the task configuration used in (Liu et al., 2024). We list detailed information of task configuration and descriptions below.
38

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
a
Test Accuracy
Lazy
Rich
Ultra-rich
00
1.0
dataset
CIFAR-10
Test Accuracy
0.8
CIFAR-100
Capacity
40
CIFAR-10C
Radius
Dimension
Alignment
0.4
Alignmant
Center-Axes Alignment
10-2
10-1
0.0
Figure 17. OOD performance and effective geometric measure of ResNet- models trained with different scale factor.
Perceptual decision making (Britten et al., 1992) (documentation page)
Task description: each trial, given noisy stimulus, the agent needs integrate the stimulus over time to determine which stimuli has stronger signal.
Task configuration: We set up the task using the following parameters: {timing: (fixation: 0,
stimulus: 700, delay: 0, decision: 100}, dt: 100, seq.len: 8}
Context decision making (Mante et al., 2013) (documentation page)
Task description: In each trial, given two noisy stimulus, each has two modalities, the agent needs to integrate the stimulus in one specific modal while ignoring the other modal. The interested modal given by the context.
Task configuration: We set up the task using the following parameters: {timing: (fixation: 0,
stimulus: 200, delay: 500, decision: 100), dt: 100, seq 8)
Delay match sample (Miller et al., 1996) (documentation page)
Task description: In each trial, a sample stimulus is shown during the sample period, which followed by a delay period. Afterwards, the test stimulus is shown. The agent needs to determine whether the sample and the test stimuli are matched.
Task configuration: We set up the task using the following parameters: ftiming: [fixation: 0, sample: 100, delay:
500, test: 100, decision: 100), dt: 100, seqllen 8)
F. 1.2. MODELS
Model architecture We consider time-continuous recurrent neural networks (RNNs) architecture that are commonly used model neural circuits (Liu et al., 2024; Ehrlich et al., 2021). Specifically, we consider RNNs with hidden layer, ReLU activation, N,, input units, Nhidden hidden units, and Now output unit. Letx; E N., Vt TR be the corresponding input and output at time-step t. The model's hidden representation h and outputs at time step can be defined by the given equations:
h = phg + (1 ()(Waa((e) + Wax))
(8)
dt (h)(h))
(9)
39

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
dt In the above equation, W2 RR- Xaraedien Wh E RNNaadenn.Niiieee XNhiaden, Wo E R-Nisdden Naue, (.) is the non-linear activation function, which we used ReLU, and p is the decay factor which is defined by p= with time step dt and time constant т. We use Nhidden = 300 for all RNNs models.
Weight rank initialization Following the practice in (Liu et al., 2024), we initialize the recurrence weight Wh. by initializing an initial full-ranked random Gaussian matrix, and then use Singular Value Decomposition truncate the weight rank the desired rank. The truncated weight matrix is then re-scaled ensure that weight matrices with varying ranks have the same weight norm.
F. F.1.3. TRAINING PROCEDURE
Loss function:
Since all three tasks that we consider are classification tasks, we use cross entropy loss.
Optimizer: We use Stochastic Gradient Descent with momentum (implemented as torch. omtim.e. SGD (1r=0.003) momentum=O.. 9) to train the models.
Batch size: We use batch size of 32 for each training step.
The models are trained for 10000 iterations and all models being compared achieved similar loss and accuracy after training (see Figure Figure 18, Figure 19, Figure 20 for more details).
F.1.4. MANIFOLD CAPACITY MEASUREMENTS
In this section, we provide detailed information about how we define object manifolds from the model's representations and measure the manifold capacity and geometric properties (Chung et al., 2018).
Features extraction: We extract the representation he (in Equation Equation 8) from the hidden layer (dimension 300) with t being the decision period of the trial.
Number of manifolds: The number of possible choices in the decision period of all the three tasks that consider is 2, so the number of manifolds are 2.
Number of points per manifold: For each task-relevant manifold, we randomly sample 50 trials of the corresponding ground truth choices.
Number of repetitions: Every capacity and geometry measurement is repeated 50 times and we report the mean and the error bar as the bootstraped 95% confidence interval calculated using seaborn. lineplot (errorbar- (' ci', 95) )
E.2 Additional results on other cognitive tasks
In section Section 5.1, we present the results on how the initial structural connectivity bias (initialized by varying the rank of the weight matrix) affects the feature learning regime and representational geometry of a given model in the perceptual decision making task (also called the two-alternative forced choice task) (Britten al., 1992). In this section, we show more detailed results (including accuracy and loss) on the perceptual decision making task in Figure Figure 18. along with two other cognitive tasks, which are context decision making task (Mante et al., 2013) in Figure Figure 19 and delay match sample task (Miller et al., 1996) in Figure Figure 20.
40

---

Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry
Perceptual Decision Making Task
Stage
Stage Untralned Trained
Untrainne
Trained
d
Accuracy
Activaliion zim lonny
Hignieent
Werght change G C
E
Aoes
&
10"
10"
10"
UP
weight
Initial weight rank
Initial weight: rank
Capacity
Initial weight rank
Initial weight rank
1.00
Kernel alignmmet
Dimersie
Centor Mignnent
10
127
100
102
107
spt
tal
tel
101
weight rank
Initial woight rank rank
Inital woight rank
Intiaa weight rank
Initial weight rank
Initial weight rank
Figure 18. Structural connectivity bias in the two-alternative forced choice task. a. Model train and loss accuracy b. Weight change and alignment measurements Manifold capacity measurements d. Effective manifold geometry measurements.
Context Decision Making Task
Stage
Stage Untrained Trained
Untrained
Trained
d
1.00
Weight change
Activesor limilartty
0.4
Radius
2.3
Axes
Alig menntt
10'
10
20t
20
30t
30
301
Initial weight rank
Initial weight rank
Initial weight rank
Capacity
Initial weight rank
Initial weight rank
1.0
Representation algnnment
Kernel all algnimant
Dimenssoo
Center Alignmant
0.7
604
Initial weigitt rank
weight
weight Initial weight rank
Initial weight
weight
Figure Structural connectivity bias context decision making task a. Model train and loss accuracy b. Weight change and alignment measurements c. Manifold capacity measurements d. Effective manifold geometry measurements.
Delay Sample Matching Task
Stage
Stage
Untrained
Trained
Untrained Trained
d
Accuracy
3.6
Walgitt change
Activation
Rediues
Auccs Aigriment
2
2.6
127
301 sel
101
101
301
301
tel
3et
321
tel
Initial weight rank
Initial weight rank
initial weight rank
Capacity
Initial weight rank
Initial weight rank
Repressental ation alignment
1.
Karnel Aignmentt
Dimners on
Center Alggmment
LOOS
2.
19
Initial weight rank
Initiel
Initial weight rank
Initial weight rank
Initial weight rank
Initial weight rank
Figure 20. Structural connectivity bias in the delay mataching sample task. a. Model train and loss accuracy b. Weight change and alignment measurements Manifold capacity measurements d. Effective manifold geometry measurements.


Reviews:
[{"review_id": "2aEpqHRLZY", "reviewer": "Reviewer_guAP", "summary": "This paper uses manifold capacity measures to assess neural representations and learning in the rich and lazy learning regimes. The authors show, both numerically and, in some cases, analytically, that these manifold capacity measures provide a deeper understanding of learning dynamics and neural representations that are beyond traditional metrics like accuracy, weight changes, label alignment, and representational alignment. Novel insights include, that representations differ when the network starts with either a beneficial (wealthy) or poor (decremental) initial weight structure; and that learning unfolds in four distinct phases: clustering, structuring, separating, and stabilizing, which vary depending on the learning regime. The authors further extend their analysis to recurrent neural networks, exploring how hidden-layer representations evolve under high- and low-rank initializations and further investigate out-of-distribution generalization, showing how the the rich and lazy regimes influences the performance depending on the underlying out-of-distribution-tasks' complexity.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer **guAP** for spending time and effort to thoroughly read, evaluate, and provide detailed comments and suggestions for our manuscript!\n\nWe greatly appreciate the reviewer for recognizing that our works provide `a deeper understanding of learning dynamics and neural representations that are beyond traditional metrics like accuracy, weight changes, label alignment, and representational alignment` and the claim is supported by `strong (analytical and) numerical evidence`!\n\nWe also thank the reviewers for excellent advices on (1) adding the missing references to Woodworth, Blake, et al., (2) improving the legibility of text, labels, and ticks in Figures 2 and 4, and (3) providing more context on how to relate “exponentially fast learning” in the lazy regime with the “poorly organized manifold” concept. We really values the reviewer’s suggestions and have added action items to incorporate these points to our updated version!\n\nWe thank reviewer **guAP** once again for their time and actionable feedback!", "claims_and_evidence": "The authors claim that manifold capacity measures provide an improved and more detailed understanding of the rich and lazy learning regimes. All specific claims (outlined in the summary) are supported by strong (analytical and) numerical evidence. The analysis is both thorough and comprehensive. The figures are clear and effectively visualize the evidence, providing strong support for the claims.", "methods_and_evaluation_criteria": "The methods and evaluation criteria, including the choice of network architectures and datasets, are well justified. They cover a range of cases, from the simpler and more tractable 2-layer nonlinear ANNs and point clouds to more realistic architectures like ResNet and datasets such as CIFAR. Additionally, the tasks and scope of the recurrent neural network studies are well justified and align with established standards.", "theoretical_claims": "I did not verify the theoretical claim in Appendix C.", "experimental_designs_or_analyses": "See Methods And Evaluation Criteria", "relation_to_broader_scientific_literature": "The presented work uses manifold capacity measures to study representation learning in artificial neural networks. In theory, this approach could also be applied to examine interindividual differences in neural representations in neuroscience, suggesting that individuals may differ in their operation within the rich and lazy regimes. Furthermore, the analysis pipelines and measures presented could enhance our understanding and comparison of representation learning both in neural networks and in the brain during learning.", "essential_references_not_discussed": "Woodworth, Blake, et al. \"Kernel and rich regimes in overparameterized models.\" Conference on Learning Theory, PMLR, 2020. This paper is one of the first to systematically study the relationship between the rich and lazy regimes and generalization.", "other_strengths_and_weaknesses": "The paper is very dense but well written.", "other_comments_or_suggestions": "The text, labels, and ticks in Figures 2 and 4 are often too small and should be made more legible.\n\nThere is a duplication of text: \"We adopt the setting from previous work (Liu et al., 2024) on investigating how differences in connectivity initialization affect the learning process.\" and \"To study how connectivity structure impacts learning strategies, we follow the setup in (Liu et al., 2024)...\". This repetition could be avoided by consolidating the statements.\n\nI think the following paragraph:  \"In a network that does not learn task-relevant features (e.g., lazy learning, random features, Figure 1b, left), the manifolds are poorly organized, making them harder to distinguish (e.g., smaller margin, smaller solution volume). In contrast, when a network learns task-relevant features (e.g., rich learning Figure 1b, right), the manifolds become well-organized and easier to separate (e.g., larger margin, larger solution volume).\" may be a little bit counterintuitive as the lazy and rich learning regimes are generally associated with fast exponential and slow step-like learning dynamics. Maybe the authors can elaborate on that, e.g. it is easy to find a linear searation in a high-dimensional random projection of the data which leads to exponentially fast learning, however also leads to poorly organised (or not even really to any) manifolds.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "suL8JCN03o", "reviewer": "Reviewer_sxs9", "summary": "Numerous studies in representation learning have been conducted to evaluate the quality of features learned by DNNs, particularly in determining whether a neural network functions within the lazy or rich regime. In this paper, the authors presented theoretical foundations grounded in manifold capacity theory to address the Lazy vs. Rich dichotomy issues, examining feature learning through the geometric properties of task-relevant manifolds. Additionally, the paper highlighted that the training of neural networks evolves through distinct learning stages, as reflected by the dynamics of manifold geometry. It also identified emerging learning strategies as networks demonstrate varying levels of richness in their learning. With robust theoretical underpinnings and empirical support, the proposed geometric features provide a valuable tool for evaluating the depth of feature learning.\n\n## Update after rebuttal\nThank you to the authors for their clarifications. I look forward to reviewing the updated manuscript and will maintain my current score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer **sxs9** for their thorough evaluation of our paper’s motivation, methods, and results, as well as the valuable feedback and insightful questions! We greatly appreciate the reviewer for recognizing that our work “offers substantial potential as an interpretable tool for elucidating the model's learning dynamics” and “the experimental setups clearly demonstrate the effectiveness of the proposed method”. Below we address the reviewer questions about simulated manifold capacity and missing details in the Appendix.\n\n1. `According to Equation 1 and Section 2.1, the authors employed a random projection \\Phi from R^N to R^n. Further details about this process are needed.` \n    \n    We thank the reviewer for pointing out the missing details about random projection in Equation 1, which specifies the definition of simulated manifold capacity. Intuitively, simulated manifold capacity measures the smallest dimensional subspace such that the projected manifolds is linearly separable with probability ≥ 0.5, under the distribution of random binary label dichotomy and random projection. The random projection matrix can be sampled from the standard normal distribution (then normalized to ensure unit norm). Therefore, simulated capacity can be computed numerically by performing a bisection search to find the smallest dimensional subspace such that the probability of manifold separability ≥ 0.5. Further details can be found in [Cohen et al. 2020], p11, section “Measuring capacity numerically from samples”). Below we provide the pseudocode to compute simulated capacity, which will also be included in the final version of the paper\n    \n    **Pseudocode:**\n    \n    Step 1: Set `min_dim` (minimum subspace dimension for random projection, usually 2) `max_dim` (maximum subspace dimension, usually the original dimension),  `num_repetition` (number of random repetitions to sample the binary dichotomy and random projection matrix), `tolerance` (error tolerance between the target, which is the 0.5 probability value, and the found value), and `max_iteration` .\n    \n    Step 2: Set `mid_point = min_dim + np.floor((max_dim-min_dim)/2)` . Estimate the probability of linearly separable for `mid_point` , or `f(mid_point)`\n    \n    Step 2.1: To estimate the linearly separable probability for `mid_point` , we sample random projection matrix `M` and binary label `y` for `num_repetition` times. For each repetition, we can use quadratic optimization to determine whether the current projected manifold sample is linearly separable (0 or 1). The returned estimated probability is the ratio between the number of linearly separable samples over the total number of repetitions.\n    \n    Step 3: While `abs(f(mid_point) - 0.5) > tolerance` and `current_iteration < max_iteration` , update `mid_point`, and repeat the computation of `f(mid_point)` until reach the value within tolerance, or the maximum number of iterations.\n    \n    Step 3.1: If `f(mid_point) > 0.5` , update `max_dim = mid_point`, else update `min_dim = mid_point` . Store tuple `(mid_point, f(mid_point))`\n    \n    Step 4: Return `mid_point` if within `tolerance`, else use interpolation to estimate the value of `mid_point` such that `f(mid_point) = 0.5` \n    \n    References: \n    \n    1. Cohen, Uri, et al. \"Separability and geometry of object manifolds in deep neural networks.\" *Nature communications* 11.1 (2020).\n2. Supplementary Material: `There were a few minor missing details, such as the type of initialization for random Fourier features and the specifics of the teacher-student setting in Section C.`\n    \n    We thank the reviewer for pointing out the minor missing details. We follow the same setting as in [Montanari et al. 2019] and [Ba et al., 2022]. Specifically, the initial weights of the 2-layer network were sampled independently from isotropic Gaussians (i.e., the random features are orthogonal with each other with high probability), this is also described in item 2 of Assumption C.1 in the appendix. As for the teacher-student setting, the teacher is modeled as a hidden direction $\\beta^*$ and examples data $x_1,\\dots,x_{n_{\\text{train}}}$ are generated indenpedently from isotropic Gaussians with labels being $y_1,\\dots,y_{n_{\\text{train}}}$ where $y_i=1$ with probability $F(\\langle\\beta^*,x_i\\rangle)$ for some monotone function $F(\\cdot)$.  This is also discussed in Setting C.2 in the appendix and we will provide more details in the updated version of the paper.\n    \n    References:\n    \n    1. Montanari, Andrea, et al. \"The generalization error of max-margin linear classifiers: High-dimensional asymptotics in the overparametrized regime.\" *arXiv preprint (2019).\n    2. Ba, Jimmy, et al. \"High-dimensional asymptotics of feature learning: How one gradient step improves the representation.\" *Advances in Neural Information Processing Systems* 35 (2022).\n\nWe thank reviewer **sxs9** once again for their time, insightful questions, and actionable feedback!", "claims_and_evidence": "- C1. The paper is well-motivated, and the experimental setups clearly demonstrate the effectiveness of the proposed method. Drawing on the theoretical definitions of feature manifolds, the authors provide empirical evidence, including variations in manifold capacities in relation to the Lazy vs. Rich regime, compared to traditional measures (Figure 2). Additionally, the comparison of input dimensions and their impact on manifold capacities is particularly intriguing (Figure 3).\n- C2. The use of manifold geometry to illustrate learning strategies and stages is effectively demonstrated (Figs. 4 and 6). This approach offers substantial potential as an interpretable tool for elucidating the model's learning dynamics, with the added benefit of enhancing the model's generalizability and robustness.", "methods_and_evaluation_criteria": "- M1. All definitions are grounded in manifold capacity theory, including the computation of dimension, radius, and various alignments, and are thoroughly evaluated within Lazy vs. Rich regime frameworks on both synthetic and image datasets. The experiments involving RNNs, as well as those focused on domain adaptation and out-of-distribution generalization, are intriguing and offer great insights.", "theoretical_claims": "- T1. One of the questions concerns the definitions of simulated manifold capacity and packability presented in the main text. According to Equation 1 and Section 2.1, the authors employed a random projection \\Phi from R^N to R^n. Further details about this process are needed. Specifically, is this projection intended for dimensionality reduction? What type of random weights were utilized? Additionally, does the use of random weights always guarantee the identification of the same feature manifold? (This question addresses the identifiability of manifolds.)", "experimental_designs_or_analyses": "- E1. Overall, the authors' experimental designs and analyses were well-structured and rigorously evaluated. For instance, the experiments measuring the degree of feature learning (Figs. 2 and 3) effectively demonstrate a strong alignment between manifold capacity and the Lazy vs. Rich regimes, particularly when compared to other conventional metrics, which is persuasive. Furthermore, Figure 4 presents consistent results using manifold radius and dimension across different regimes.", "relation_to_broader_scientific_literature": "Feature learning is a fundamental aspect of neural network research in machine learning, extending far beyond the simplistic lazy-versus-rich dichotomy. Understanding the relationship between feature learning and performance is essential for designing network architectures and learning algorithms that offer high reliability and transparency for practical applications. The authors' proposed method, grounded in manifold capacity, holds significant potential as an interpretable tool that aligns well with the current research direction.", "essential_references_not_discussed": "In terms of comparison with other conventional measures, I believe the authors have sufficiently addressed the relevant references, and no critical omissions require attention.", "other_strengths_and_weaknesses": "- S1. Overall, the paper is well-structured and presents the experimental results in an intuitive and accessible manner. The figures illustrating the correlations between geometric characteristics and capacity are particularly valuable in enhancing understanding.", "other_comments_or_suggestions": "I have no other comments or suggestions.", "questions_for_authors": "I have listed my questions in each section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rzvqDx3vcQ", "reviewer": "Reviewer_NZt7", "summary": "The authors revisit the dichotomy of *lazy learning*, where neural networks do not learn data-dependent features and instead act essentially as a kernel machine, and *feature learning*, where they do. They argue that there are in fact several distinct learning regimes in the rich regime, which they untangle by studying the the geometric properties of task-relevant manifolds, for example the point cloud of neural activations corresponding to stimuli in a given class for a classification task.\n\nThey study experimentally two-layer and deep neural networks, tuning the amount of feature learning by varying the inverse scale factor $\\eta$ à la Chizat 2019. They show that various geometric measures of manifold geometry are connected to the prediction error (but unfortunately do not even state the result in the main text). They then go on to show experimentally that various measures of manifold geometry correlate well with feature learning on synthetic data (Sec 3) and that different measures of manifold geometry vary at different stages in training, which the authors call the clustering, structuring stage, separating stage and stabilizing stage. \n\nFinally, they apply their methodology to recurrent neural networks trained on tasks studied in theoretical neuroscience. They confirm previous work showing that the rank of the initial connectivity governs whether the learning dynamics are lazy or rich, but find that RNNs achieve roughly the same manifold capacity either way.\n\n## After the discussion...\n... I increased my score, as I explained in my rebuttal comment.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer NZt7 for their thoughtful reviews and suggestions. We appreciate the reviewer recognized that our works “presents an interesting and really well-written exploration of how various quantities describing the geometry of the task manifold evolve during the training dynamics of neural networks”. We greatly value the reviewer’s insightful questions about (1) the robustness of our learning stage results, (2) the presentation of the theoretical results, (3) the relationship between our different geometric measurements, as we agree that addressing these concerns would greatly strengthen the paper. Below we address these well-thought questions and concerns:\n\n1.  `An interesting observation that the authors make is that they identify four different stages of the rich regime (line 340ff) However, this appears to be a one-time observation in an experiment, and unfortunately, the authors do not explore it further, try to reproduce it in another setting`\n    \n    We appreciate that reviewer recognize that our geometric measures can offer a deeper understanding on learning dynamic. About the concern about the robustness of the “learning stages” findings in Fig. 4c, we want to clarify that:\n    \n    1. **Number of random seeds repetitions:** All our empirical experiments contain results from 5 different models repetitions, initialized from different random seeds, and the heat map in Fig. 4c is the results averaged from 5 repetitions. While we did include the information about number of repetitions in Appendix section E.1, we admit that this is an important information and should have been mentioned in the caption figure, and we will include this information in the final version of the paper.\n    2. **Different models configurations:** While Fig. 4c focuses on a single model configuration (VGG-11 with $\\bar\\eta=0.2$) to demonstrate a specific example of how our methods can offer deeper insights on learning stages, we indeed have experimental results that similar hidden stages can be observed across various richness degrees (across $\\bar\\eta$ values) and model architectures (VGG-11 and ResNet-18) at https://imgur.com/a/learning-stages-figures-f5pZEWa. Since reviewer NZt7 mentioned that it is worthwhile to explore this observation further, we sincerely hope that the reviewer can take a look at these experiment results, which shows that our geometric measurements can reveal hidden stages in various settings of richness degrees and model architectures.\n2. `The authors report a theoretical result which sounds intriguing - connecting various geometric quantities to prediction error. However, the authors do not even state this result in the main text`\n    \n    We thank the reviewer for the careful reading of our theoretical results. We fully agree that an analytical characterization of capacity throughout training offers valuable insights that help justify our approach. Due to page limitations, we made the difficult decision to place the detailed theoretical analysis in the appendix, as we wanted to emphasize the breadth of our empirical findings in the main text. In the final version of the paper, we will include more details on the theoretical results in the main text to better highlight their significance.\n    \n3. `It is hard to track all the different quantities relating to manifold geometry; for me for example it is not clear how independent these measures are, and to what extent one can explain the other.`\n    \n    Previous work on manifold capacity and its effective geometric measures [Chou et al., 2024] offers comprehensive explanations, covering theoretical foundations, numerical analyses, and intuitive insights, on how to interpret these measures and their relationship to capacity in the supplementary material. We have incorporated some relevant parts into our own appendix. In the final version of the paper, we will include additional examples to better illustrate the independence among these measures. In particular, we will provide a mathematical explanation on how effective dimension interact with axis alignment (higher axis alignment would effectively reduce the manifold dimension, because the variation subspaces become more overlapped); how effective radius interact with center alignment (higher center alignment would effectively increase the radius, because the manifolds become closer to each other).\n    \n    References:\n    \n    1. Chou, Chi-Ning, et al. \"Neural manifold capacity captures representation geometry, correlations, and task-efficiency across species and behaviors.\" *bioRxiv* (2024).\n\nWe thank reviewer **NZt7** once again for their time, insightful questions, and actionable suggestions that help us to strengthen both the presentation of our theoretical results and the robustness of our experimental results! We hope our responses address your concerns. Please let us know if there’s any other details that we can further clarify. Thank you very much for your time and consideration!", "claims_and_evidence": "See Strengths & Weaknesses.", "methods_and_evaluation_criteria": "See Strengths & Weaknesses.", "theoretical_claims": "The authors report a theoretical result which sounds intriguing - connecting various geometric quantities to prediction error. However, the authors do not even state this result in the main text, for reasons that are not clear to me. If it is important, it should be stated there!", "experimental_designs_or_analyses": "See Strengths & Weaknesses.", "relation_to_broader_scientific_literature": "Broadly speaking, the paper does a good job of relating to the broader scientific literature. A couple of points should be improved:\n\n- The approach of the paper is really a (well-executed!) application of the manifold untangling hypothesis that originates in neuroscience. The authors mention DiCarlo & Cox (2007) in a footnote, but I think this should be stated more prominently in the main text, for example at the beginning of section 1.1.\n- There were clearer and earlier demonstration of the advantage of feature-learners over lazy models, including Ghorbani et al. NeurIPS 2019 and 2020; Daniely & Malach NeurIPS 2020; and Refinetti et al. ICML 2021 (start of p. 2)", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "The article presents an interesting and really well-written exploration of how various quantities describing the geometry of the task manifold evolve during the training dynamics of neural networks. These geometric quantities show promise in that they are related to the accuracy via a theoretical result; however, since this result is not even stated in the main text, and the discussion in the supplementary material spans several additional pages, I do not want to consider it too heavily in this review out of respect for the page limit.\n\nAnother strength of the paper is that it is packed with different concepts to analyse neural networks - but that is also a weakness. It is hard to track all the different quantities relating to manifold geometry; for me for example it is not clear how independent these measures are, and to what extent one can explain the other. Finally, it is not clear at all how they relate to accuracy; I guess the theoretical result should fill in that void, but then it should be discussed in the main text.\n\nAn interesting observation that the authors make is that they identify four different stages of the rich regime (line 340ff) However, this appears to be a one-time observation in an experiment, and unfortunately, the authors do not explore it further, try to reproduce it in another setting, or further analyse what to me looks like the main result of this study (and it is indeed one of the three main results mentioned by the authors in their introduction) I therefore find the overall amount of results a bit lacking to recommend acceptance, but I think it is worthwhile to explore this observation further. I would make acceptance contingent on further expansion of this observation. While it is obviously not for myself to decide, I wonder whether this paper in a longer form published in a good journal like JMLR wouldn't do more justice to its content.", "other_comments_or_suggestions": "- This is a small point, but I would not describe the (great!) paper of Chizat, Oyallon & Bach '19 as demonstrating \"that neural networks can perform well even when there are negligible changes in the weights of the networks\" (p. 1) Instead, they showed that two-layer neural networks whose first-layer parameters move only a little are instead severely limited on high-dimensional tasks compared to feature learners in the same way that kernels are, therefore coining the term \"lazy\".\n- When citing books, please give a more specific reference to a subsection or a theorem, otherwise the reference is useless (for example when citing the book by Vershynin.", "questions_for_authors": "No further questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JXVuN4jz6q", "reviewer": "Reviewer_EQUL", "summary": "This paper introduces a framework to study subtypes of the rich regime during training of natural or biological neural networks.\n\nSpecifically, the authors use the known concept of manifold capacity to distinguish different phases of training and extract insights for ML and neuroscience.\n\n## update after rebuttal\n\nI thank the reviewers for their rebuttal, which partially answered my questions. The motivation of manifold capacity over manifold topology (and geonetry/dimension) is not convincing enough, as the authors state it is because it task-dependent, yet topology, geometry and dimension are also task-dependent. Hence, I maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer EQUL for their thorough evaluation and insightful questions. Below we address the reviewer questions on (1) our method, (2) experimental results. About the theoretical result, please refer to our response to reviewer NZt7, section 2, which address similar question.\n\n1. Method and experiment setup\n\na. `It seems that the method applies only to the last layer? Similarly, it seems that the analysis only applies to classification tasks`\n\nOur method is applicable to any hidden layer, we chose to focus on the last layer because it's typically where final learned features are examined. We’ll note in the paper that investigating intermediate layers is an interesting future direction. Our manifold capacity method focuses on classification tasks, building upon the classic perceptron storage capacity [Cover, 1965]. Extending the capacity notion to other tasks, such as regression, remains a promising area for future works.\n\nb. `The choice of using manifold capacity could be further motivated in terms of what alternatives exist to study feature learning (e.g, dimensionality and frames, curvature, topology)`\n\nWe thank the reviewer for providing relevant geometric measures and will include these in our “Relevant works” section. Compared to other geometric approach, our method plays a unique role. Because these geometric measures are task-relevant, they allow us to more directly track the relationship between task-level performance and corresponding representational geometry (so the manifold dimension is extrinsic dimension, as it depends on the labels). In contrast, other geometric metrics, such as raw dimensionality, curvature, or topology, provide natural ways to characterize representations, but their connection to network performance remains more elusive.\n\nc. `Can you explain what is the scale factor and why it represents the degree of feature learning?` \n\nWe use the standard setup in [Chizat, 2019] to use the scaling factor to tune the degree of feature learning (as also explained in reviewer NZt7 summary). For example, the original formula for MSE loss is $\\frac{1}{2}(f(x)-y)^2$, with the output scaling factor $\\alpha$, the loss would be $\\frac{1}{2}(\\alpha f(x) - y)^2$, which mean that we scale the network output by a scaling factor $\\alpha$. The scaling factor can be used to tune the degree of feature learning as mentioned in [Chizat, 2019] section “Rescaled models” in p3 and Theorem 2.3 in p5. Intuitively, as the scaling factor grows, the network weights only need to change minimally while still achieving a big decrease in the objective loss, leading to lazy learning, in which the learned weight can be linearly approximated from the random initialized weight and not contain task-relevant features.\n\nd. `The “conventional methods” to which the manifold capacity method is compared are not introduced in the related works. Why choosing these methods?` \n\nWe thank the reviewer for pointing this out! Weight changes, test accuracy are two measures used in the first lazy-rich paper [Chizat 2019]. NTK-label alignment and Representation-label alignment are measures used in follow-up works [Geiger 2020] and [Kumar 2024]. We will put these citations in our updated version.\n\n2. Experimental results:\n\na. `Fig 4c: distinguishing different stages of learning during training` \n\nFig. 4c shows that while the performance (e.g capacity) monotonically increases (1 learning stage), the geometric measures can capture more subtle change in the representations during training, and reflect distinct learning stages (as acknowledge by reviewers guAP, sxs9, NZt7). About the concern on robustness of this finding, please refer to our response to reviewer NZt7, section 1.\n\nb. `Empirical justification in standard settings: I’m not sure that training two neural networks is enough to justify manifold capacity` \n\nWe showed empirically that capacity can correctly track the ground-truth of feature learning (measured by the inverse scaling factor $\\bar\\eta$) across different settings, including simple 2-layer models (Fig. 3) and deep nets (VGG-11: Fig. 2a,b, Fig. 13, ResNet18: Fig 14).\n\nWe want to clarify that each reported data point resulted from 5 models initialized with different random seeds. We also vary the feature learning rate $\\bar\\eta$ from 0.01 to 1 with 10 different values to train our models, therefore including both rich and lazy regimes. For each model architecture, our results are reported from training 50 neural networks.\n\nc. `In Section 5, bold claims are made that geometric signatures can reflect DA and OOD behavior. Fig 6.c only show *one* instance where that is the case. How many neural networks were trained here?` \n\nFor the DA and OOD experiments, we used two distinct architectures VGG-11 (Fig 6) and ResNet18 (Fig 16), and indeed we've observed capacity can capture OOD performance and distinct geometric signatures with different architectures. About the number of trained neural nets , we refer to the answer on 2b.", "claims_and_evidence": "Until the method section, it seems that the analysis applies to any neural representation. However, it seems that it applies only to the last layer’s? Can the authors comment on the scope of the analysis? Similarly, it seems that the analysis only applies to classification tasks: can the authors comment on this too?\n\nThe authors mention distinguishing different stages of learning during training, but Figure 4c on that topic is hard to understand: it seems that capacity slowly and steadily increases during learning: how can the authors say that the four stages are “evident”? What would be an algorithm to automatically find these stages?", "methods_and_evaluation_criteria": "On Fig 3, if we zoom on the last epochs: could test accuracy differentiate between the regimes?\n\nSection 4 only evaluates on two-layer neural networks. How about *evaluating* the framework on larger NNs, and comparing capacity to the other existing measures of NTK-label alignment, etc? (I understand that the framework is later *applied* to large NNs, but the values of the other measures -- such as NTK-label alignment -- are not given).", "theoretical_claims": "Theorem C.4 seems very important, since it could justify the use of capacity as a measure of feature learning. It should be in the main text. Generally, a lot of content is in appendix such that the main text is hard to read without referring to it constantly.", "experimental_designs_or_analyses": "In paragraph: Empirical justification in standard settings: I’m not sure that training two neural networks is enough to justify manifold capacity. Esp since both were trained in the lazy regime, which is not the point of this paper? Did I miss something?\n\nLikewise, in Section 5, bold claims are made on the fact that geometric signatures can reflect DA and OOD behavior: but I’m not sure whether there are enough experiments to be able to make this point. Fig 6.c only show *one* instance where that is the case. How many neural networks were trained here? Do you have quantitative evidence that manifold capacity reflects DA and OOD across the models (and if so, could you comment on this in the main text?)?", "relation_to_broader_scientific_literature": "The choice of using manifold capacity could be further motivated in terms of what alternatives exist to study feature learning (and that could have potentially found similar learning stages?).\n\nEg dimensionality and frames: Kvinge et al. Internal Representations of Vision Models Through the Lens of Frames on Data Manifolds\n\nEg. Curvature: Acosta et al. Quantifying Extrinsic Curvature in Neural Manifolds.\n\nEg topology: Yoon et al. Tracking the topology of neural manifolds across populations", "essential_references_not_discussed": "The  “conventional methods” to which the manifold capacity method is compared are not introduced in the related works in the main text. \n\nSee: “here we compare our method with several common measures for feature learning: accuracy curves, weight changes, and alignment methods (Table 1)” . *Why* choosing these methods for the comparison?", "other_strengths_and_weaknesses": "Figures are excellent, but often hard to read.\n\nFig 2, eg, is gorgeous but *packed.* \n\nFig 4b is too small.", "other_comments_or_suggestions": "Is manifold dimension the extrinsic or intrinsic dimension?\n\nAre the definitions of wealthy and poor regimes classical definitions? If so, could the authors add a reference?\n\nCan you explain in more details what is the scale factor and why it represents the degree of feature learning? And unpack the link with the learning rate?", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a compelling and novel framework for understanding feature learning dynamics in neural networks via the geometry of task-relevant manifolds. It moves beyond the standard lazy vs. rich dichotomy by characterizing learning stages through manifold capacity and related geometric descriptors.

The reviewers found the work both theoretically insightful and empirically well-supported, noting its relevance across domains, including neuroscience and machine learning. In particular, the identification of learning stages (clustering, structuring, separating, stabilizing) via representational geometry is a notable contribution. The experiments are diverse and carefully designed, covering synthetic and image datasets, and RNNs, and include settings involving OOD generalization.

While several reviewers raised valid concerns, such as the placement of theoretical results in the appendix, the robustness of learning stages, and the motivation for choosing manifold capacity over alternative measures, these were well addressed in the rebuttal. 

Given the paper’s originality, technical depth, and relevance to both theoretical and applied ML communities, I support acceptance to ICML.
Decision: Accept (spotlight poster)