Beyond the Lazy versus Rich Dichotomy: Geometry Insights in Feature Learning from Task-Relevant Manifold Untangling

Authors: Chi-Ning Chou, Hang Le, Yichen Wang, SueYeon Chung
Primary area: applications to neuroscience & cognitive science
Keywords: Computational neuroscience, storage capacity, neural manifolds, representational geometry, rich and lazy learning, training dynamics, feature learning

Abstract:
The ability to integrate task-relevant information into neural representations is a fundamental aspect of both human and machine intelligence. Recent studies have explored the transition of neural networks from the *lazy* training regime (where the trained network is equivalent to a linear model of initial random features) to the *rich* feature learning regime (where the network learns task-relevant features). However, most approaches focus on weight matrices or neural tangent kernels, limiting their relevance for neuroscience due to the lack of representation-based methods to study feature learning. Furthermore, the simple lazy-versus-rich dichotomy overlooks the potential for richer subtypes of feature learning driven by variations in learning algorithms, network architectures, and data properties.

In this work, we present a framework based on representational geometry to study feature learning. The key idea is to use the untangling of task-relevant neural manifolds as a signature of rich learning. We employ manifold capacity—a representation-based measure—to quantify this untangling, along with geometric metrics to uncover structural differences in feature learning. Our contributions are threefold: First, we show both theoretically and empirically that task-relevant manifolds untangle during rich learning, and that manifold capacity quantifies the degree of richness. Second, we use manifold geometric measures to reveal distinct learning stages and strategies driven by network and data properties, demonstrating that feature learning is richer than the lazy-versus-rich dichotomy. Finally, we apply our method to problems in neuroscience and machine learning, providing geometric insights into structural inductive biases and out-of-distribution generalization. Our work introduces a novel perspective for understanding and quantifying feature learning through the lens of representational geometry.

Decision: reject
Reviewer scores: 6, 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730437929260

thread_type:
official_review

<thread_object>
creator: reviewer_d4ag

<content>
summary:
This paper proposes a novel framework for understanding neural network feature learning beyond the “lazy” versus “rich” dichotomy by focusing on task-relevant manifold untangling using representational geometry. The authors introduce manifold capacity as a key metric to quantify feature learning richness, and explore how manifold geometric measures reveal distinct learning stages and strategies. This framework is applied in contexts ranging from standard machine learning tasks to neuroscience, offering insights into structural inductive biases and challenges in out-of-distribution (OOD) generalization.

strengths:
1. The authors present both theoretical analysis and empirical evidence demonstrating that manifold capacity effectively quantifies the degree of feature learning. Comparisons with other metrics, such as accuracy and weight changes, further underscore the efficacy of manifold capacity in this context.
2. The authors offer insightful empirical findings on the relationship between manifold geometry and learning stages, with applications in neuroscience and OOD detection. Robust experiments substantiate these findings, providing a solid foundation for their conclusions.

weaknesses:
1. The structure needs refinement, as it currently feels like the authors have packed too much content into the main paper, resulting in diminished clarity. The main paper covers a wide range of topics—from manifold capacity to manifold geometry, from theory to experiments—and the relationship between manifold capacity and various manifold geometry measures is weakly explained, relying primarily on Figure 2c with minimal analysis. I suggest that the authors avoid treating manifold geometry as a separate section, even though some interesting findings are presented. An alternative approach would be to frame geometry as an extension of capacity (or to present capacity itself as a facet of manifold geometry) and to provide necessary analysis (I notice there is analysis in the appendix. It's better to consolidate relevant analysis from the appendix into a concise summary in the main paper).
2. The algorithm is not clearly presented, which also appears to be a side effect of the structural issue noted above. Given that the paper aims to provide practical insights into feature learning, with applications in neuroscience and machine learning, it is important to ensure readability for readers unfamiliar with the manifold background. Thus, a clear algorithm outlining the process for computing capacity during training would be essential.

questions:
1. Could the authors provide an example of the lazy regime? Specifically, how can a neural network be trained without actually modifying its internal features?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645072510

thread_type:
official_review

<thread_object>
creator: reviewer_k6x2

<content>
summary:
This paper advocates the use of manifold capacity as a way of classifying as a measure of the regime of neural networks. Manifold capacity is a measure from neuroscience literature that quantifies the performance of a linear classifier over features for classification as a measure of the richness of classifiers. The authors advocate this measure over other measures such as "lazy regimes" during which features show limited change in learning. They use this to analyze 2-layer RELU networks

strengths:
I think the use and introduction of manifold capacity as a measure is interesting for the ML community as a metric for representation learning. They relate manifold capacity to test accuracy in 2-layer networks and deeper networks empirically. They also relate to other measures such as weight changes and alignment, though this is mostly in the appendix.

weaknesses:
The paper is haphazardly written. The difference between previous work and their problem statement is not well delineated. Their goals arent succinctly mentioned. The main theorem statement being nearly vacuous in the main paper. "In 2-layer neural networks trained with gradient descent in the rich regime the changes in capacity track the underlying degree of richness in feature learning." Since richness itself is defined as capacity, this is either circular or vacuous. 

A deeper read of the result in the appendix shows that they approximate the result of gradient descent and use a gaussian model to approximate a 2-layer network. This is perhaps the most interesting result in the paper but is not at all covered in the main paper. The authors should state that they prove that using SGD, models with high capacity converge to high accuracy or something less vacuous and offer the proof. 

Other features such as manifold radius and "alignment" are also not measured. The connection to "untangling" is also not mentioned. Weight changes are used mainly as a strawman to invoke connections to NTK and need not be done. 

Figure 2 is very confusing with a nonsensical caption, "Higher capacity means that a higher number of
manifolds per neuron can be packed in the neural state space." What does it mean to "pack manifolds into neurons." This kind of shoddy language obfuscates the message to the reader. 

As the current writing stands, this seems to contribute no more (and likely less by way of confusion) to the work in 2018 by Chung et al. which describes the full intuition of manifold capacity and outlines its use both in neuronal and neural networks.

questions:
What does "untangling" specifically mean? 

What are the implications of the 2-layer results on deeper neural networks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730679558969

thread_type:
official_review

<thread_object>
creator: reviewer_v9qp

<content>
summary:
This paper proposed that manifold capacity is a better way to identify the process of feature learning than commenly used rich-lazy dichotomy. This claim is theoretically proved in a limited setting, and empirically verified.

strengths:
- The idea of catching the richness by manifold capacity is potentially useful.
- The experiments verifies that manifold capacity captures different stages of learning.

weaknesses:
- The presentation is very confusing. For example, in the experiments, there are several critical measures, such as effective radius, dimension, center alignment etc. but they are not defined (at least not in the main paper).
- The main paper claimed that the Theorem 1 is proved for a 2-layer NN. However, the NN used in the proof is actually different from what people would expect to be a "2-layer NN", since based on Assumption 1: 1) the second layer is not trained but set to random; and 2) the choice of activation function is very limited as there is a rather strong condition on the activation function.

questions:
In Section 2.1, what does "i-th input category" mean? In eq. (1), $P$ is the variable of the $\max$ operator, and for any $i \in [P]$, $\mathcal M_i$ is defined, and $\mathcal M_i$ is defined by $\mathcal X_i$. This somehow implies that $\mathcal X_i$ is also a part of the variable of the $\max$ operator, instead of pre-defined. Is it true?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714001075

thread_type:
official_review

<thread_object>
creator: reviewer_cjit

<content>
summary:
In this paper, the authors focus on the feature learning in deep learning. They use manifold capacity as a metric to quantify the degree of richness of feature learning. Experiment results show that such capacity can reveal different learning stages in different settings. They also apply this to problems in neuroscience and out-of-distribution tasks.

strengths:
-	Understanding feature learning in deep learning is an important and interesting research problem.
-	The paper proposes a metric called manifold capacity to measure the feature learning progress, which seems to be new.
-	Several experiments in different domains are presented to support the claim.

weaknesses:
See questions section below.

questions:
-	There seems to be no explanation for Figure 2c. I don’t know what does those operations mean.
-	I’m a bit confused about eq (1), the definition of model capacity. Are the manifold $\{\mathcal{M}_i\}$ predefined or they are changing when $N$ is increasing. If they are changing, how do you choose them? Also, what values should $y_i$ take?
-	When looking at Appendix C for the way to compute manifold capacity, several notions seem to be not defined, such as $T$, $\lambda$.

-	I believe the way of computing manifold capacity should be mentioned in the main text, or at least mentioned under what conditions are those values computed (I believe they cannot be exactly computed unless some conditions are assumed).

-	For Theorem 1, when looking at the actual statement in appendix, these results are proved only under the setting that with one gradient step update. Though it is understandable that going beyond this is technically difficult, I feel the statement in the main text gives the reader a wrong impression.

-	In section 3.2/figure 3b, it’s not clear to me what the definition of wealthy and poor regime are and how they are related to the input dimension. Also, I’m not sure why at initialization there will be task-relevant features (and I don’t know what are task-relevant features in this setting). Moreover, the purple line in capacity in figure 3b changes only a little bit throughout the training, I’m not sure how to interpret this (feature learning only changes a little?).

-	It’s a bit hard for me to follow section 4, presumably because the definition of these metric (e.g., dimension, radius,…) are missing.


Typo:

-	Line 173, $o_N(1) \to 1$ -> $o_N(1) \to 0$
</content>
</thread_object>

</thread>
<thread>
creation date:
1731130652304

thread_type:
official_review

<thread_object>
creator: reviewer_ncob

<content>
summary:
This paper presents a geometric framework to examine the manifold of neural network representations, providing insights into the distinctions between lazy and rich training regimes in feature learning. Specifically, it revisits the manifold capacity concept introduced in Chung et al. (2018), theoretically demonstrating that manifold capacity can serve as an indicator of the underlying richness in feature learning. Based on empirical studies using synthetic data and two-layer neural networks, observations are made regarding the relationship between manifold capacity and the degree of feature learning, as well as the stages of feature evolution. The proposed geometric measures are further applied to neural networks in neuroscience and out-of-distribution generalization tasks to explore the broader implications of this approach.

strengths:
- The novel application of manifold capacity and other effective geometric measures to investigate the lazy-vs-rich dichotomy is intriguing. It shows better alignment with the degree of feature learning compared to other metrics, such as weight changes or NTK-label alignments.
- Most of the derivations seem correct, though I could not verify every detail.

weaknesses:
- The theoretical derivation relies on a one-step gradient argument, but the fact is not mentioned in the manuscript. Moreover, the link between Theorem 2’s results and the increase in feature learning degree is not entirely clear, and additional commentary could enhance clarity.
- It would be helpful to discuss the generalizability of the observations, such as those in Figure 4. Various hyperparameters (e.g., the choice of optimization algorithms, weight initialization methods, batch size, learning rate, and scheduling) could influence implicit biases in the algorithm, affecting neural representations, geometric metrics, and even the stages of learning.

questions:
- Regarding the proposed effective geometric measures to explain capacity changes, are there standard reference lengths compared to the radius? A simple manifold radius magnitude may not accurately capture the problem's complexity when the differences between manifold means scale identically.
- Additionally, the definitions and implications of axes alignment, center alignment, and center-axes alignment are less discussed compared to radius and dimensionality (e.g., in Figures 5b, 6c).
- When using geometric measures, which layer(s) should be analyzed? Are the results consistent across different layers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732038840317

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Responses Round 1 (1/4)

comment:
We thank all the reviewers for their time and effort. Your valuable feedback has greatly helped us improve the clarity and quality of our paper.

In this paper, we propose using manifold capacity and its geometric measures to investigate the lazy vs. rich dichotomy in feature learning. We appreciate that the reviewers found our approach using manifold capacity to measure feature learning to be "novel application" (**reviewer nCob, cJiT**), have "better alignment with the degree of feature learning compared to other metrics" (**reviewer nCob**), with "robust experiments substantiate these findings" (**reviewer D4ag**). **Reviewer v9qp, and D4ag** also noted that our findings in learning stages and OOD generalization being insightful.

The main concerns and questions to our paper surrounded the clarity of presentation, specifically due to the lack of (a) definition of manifold capacity and manifold geometric measurements, (b) algorithm to empirically compute manifold capacity, and (c) interpretation of the theoretical results . We have taken all the reviewers’ suggestions seriously and incorporated them into the revised manuscript.

Specifically, in the revision for the main text, we have incorporated (a) formal definition for manifold capacity and all geometric measurements, (b) formula and pseudocode to empirically compute manifold capacity, (c) more detailed context, interpretation and a more precise version of our theoretical results. Below in the thread we provide a more detailed response for each of the above general concern. Detailed responses to individual reviews will be replied separately in the threads. 

The following is a summary of the changes we made in our updated manuscript to address these issues:
* Definition of manifold capacity:
    * Simulated version (more intuitive but not analytically tractable): Definition 1 in line 168-177.
    * Mean-field version (more complicated but analytically tractable): Definition 2 and 3 in line 979-1020.
    * Capacity formula: equation (1) in line 186-190.
* Definition of manifold geometric measures:
    * Intuition: Section 2.2, line 198-206.
    * Definitions: Section 2.2, line 207-220 . Appendix B.4.
    * How effective geometric measures connect to the changes in capacity value: Fig 2c and Appendix B.5.
* Overview of our theoretical results and interpretations: Section 3.1, line 277-314.
* An updated informal theorem statement: Theorem 1, line 299-309.

We believe these revisions significantly enhance the paper's accessibility and contributions, and we appreciate the reviewers’ re-evaluation of our work. Thanks again for the thoughtful feedback and time!

## References
* [Chung et al., PRX 2018] Chung, SueYeon, Daniel D. Lee, and Haim Sompolinsky. "Classification and geometry of general perceptual manifolds." Physical Review X 8.3 (2018): 031003.
* [Ba et al., NeurIPS 2022] 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): 37932-37946.
* [Chou et al., 2024] Chou, Chi-Ning, et al. “Neural Manifold Capacity Captures Representation Geometry, Correlations, and Task-Efficiency Across Species and Behaviors” bioRxiv 2024.02.26.582157; doi: https://doi.org/10.1101/2024.02.26.582157
* [Montanari et al., 2019] Montanari, Andrea, et al. "The generalization error of max-margin linear classifiers: High-dimensional asymptotics in the overparametrized regime." arXiv preprint arXiv:1911.01544 7 (2019).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the concept of manifold capacity (as defined in Chung et al. 2018) for quantifying the untangling of task-relevant neural manifolds, providing a framework for understanding lazy vs rich feature learning. Theoretically, the paper establishes connections between manifold capacity and feature learning in two-layer neural networks trained under gradient descent, supported by empirical evidence on real settings such as VGG / ResNet on CIFAR-10. This connection revealed new geometric insights into subtypes of feature learning. Additionally, the paper demonstrates the framework’s utility in both neuroscience tasks and out-of-distribution generalization studies.

The reviewers generally acknowledged the novelty of employing manifold capacity for the study of feature learning. However, several reviewers highlighted a significant clarity issue concerning several key technical components. While the authors made significant efforts during the rebuttal phase to address these issues through extensive rewriting, the limited discussion period may not have been sufficient for the reviewers to fully reassess the manuscript’s technical quality. Therefore, I recommend reconsideration for future submissions after a careful rewriting to improve clarity of the presentation.