Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry

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.

1. Method and experiment setup

a. It seems that the method applies only to the last layer? Similarly, it seems that the analysis only applies to classification tasks

Our 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.

b. 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)

We 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.

c. Can you explain what is the scale factor and why it represents the degree of feature learning?

We 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.

d. The “conventional methods” to which the manifold capacity method is compared are not introduced in the related works. Why choosing these methods?

We 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.

2. Experimental results:

a. Fig 4c: distinguishing different stages of learning during training

Fig. 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.

b. Empirical justification in standard settings: I’m not sure that training two neural networks is enough to justify manifold capacity

We 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).

We 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.

c. 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?

For 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.

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:

1. 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

   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:

   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.
   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.

2. 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

   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.

3. 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.

   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).

   References:

   1. Chou, Chi-Ning, et al. "Neural manifold capacity captures representation geometry, correlations, and task-efficiency across species and behaviors." bioRxiv (2024).

We 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!

We thank reviewer guAP for spending time and effort to thoroughly read, evaluate, and provide detailed comments and suggestions for our manuscript!

We 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!

We 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!

We thank reviewer guAP once again for their time and actionable feedback!

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.

1. 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.

   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

   Pseudocode:

   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 .

   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)

   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.

   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.

   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))

   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

   References:

   1. Cohen, Uri, et al. "Separability and geometry of object manifolds in deep neural networks." Nature communications 11.1 (2020).

2. 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.

   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.

   References:

   1. Montanari, Andrea, et al. "The generalization error of max-margin linear classifiers: High-dimensional asymptotics in the overparametrized regime." *arXiv preprint (2019).
   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).

We thank reviewer sxs9 once again for their time, insightful questions, and actionable feedback!