Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization

We thank the Reviewer for taking the time to review and give feedback on our manuscript. We appreciate the positive comments regarding the paper's clarity, the topic's significance, the originality of our ideas, and the supportive experimental results. We are particularly grateful for the reviewer's recognition of Theorem 4.5's contribution regarding feature learning under $\mu$P.

Regarding the reviewer's main concern about the global convergence result, we believe there is a significant misunderstanding about what our result delivered and why it is non-trivial. The reviewer states (in the comments on "Theoretical Claims" and in Question 1) that

Any neural network enjoys zero training loss and global convergence if it will no longer be updated after some training step.

We respectfully but strongly disagree with this premise as a general statement and wish to clarify this crucial point, as it seems to underlie the reviewer's doubt.

Clarifying the Core Misunderstanding Regarding Global Convergence

While it is true that training algorithms like SGD stop updating when they reach a point where the gradient is zero (a stationary point), this absolutely does not guarantee that this point is a global minimum. In the complex, non-convex loss landscapes of neural networks, such stationary points can very often be suboptimal local minima or saddle points, where the training loss is significantly higher than the lowest possible value (the global minimum) and the model has not achieved the best possible fit to the training data.

The core contribution of our Corollary 4.6 is precisely to demonstrate rigorously that, for the specific setting we study (infinitely wide networks with $\mu$P trained using SGD), the optimization dynamics avoid getting trapped in these suboptimal stationary points. Our proof establishes that the training process is guaranteed to converge to a global minimum of the training loss function.

Therefore, proving global convergence is far from trivial; it requires showing that suboptimal stationary points (such as local minima with loss higher than the global minimum and relevant saddle points) are escaped or avoided, leading specifically to the desired global minimum (the lowest possible training loss). This property is not inherent to any network training process that stops; it relies critically on the theoretical properties of $\mu$P in the infinite-width limit that we analyze in our paper. Our work provides a theoretical foundation for how $\mu$P facilitates finding the global optimum, which is a key aspect of understanding its effectiveness.

Response to Specific Questions:

Q1. In my opinion, any neural network enjoys zero training loss and global convergence if it will no longer be updated after some training step.

A1. As detailed above, this statement reflects the core misunderstanding we wish to clarify. Halting updates only signifies reaching a stationary point (zero gradient), which is not necessarily a global minimum. It could be a suboptimal local minimum or a saddle point with higher loss. Our key result in Corollary 4.6 is the non-trivial proof that, under $\mu$P and infinite width, SGD specifically converges to a global minimum of the training loss, overcoming the challenge of potentially getting stuck in other stationary points with higher loss values. This is our specific contribution.

Q2. Is the t in Theorem 4.5 a finite constant w.r.t. width like that in the TP series (e.g., TP4)?

A2. Our current analysis focuses on the behavior of infinitely wide neural networks. Therefore, investigating the dependence of constants like t on finite network widths is beyond the scope of this paper. This remains an interesting open question for exploring finite-width corrections.

We trust this clarification addresses the reviewer's primary concern by highlighting the crucial distinction between halting training updates and achieving mathematically proven global convergence. We believe our theoretical analysis provides a rigorous justification for the observed empirical success of µP and clarifies the optimization dynamics in this important regime.

We thank the reviewer for their support and constructive feedback. We address each question below:

Q1: Can the authors directly measure correlations between different features during training to support their analysis?

A1: Yes, we have directly measured the correlations between different features during training. Our analysis confirms that features remain largely independent throughout the training process. Specifically, we compute the Gram matrix G of the centered features, where $G_{ij} = (h(x_i) - \bar{h})^{\top} (h(x_j) - \bar{h})$ represents the similarity between features of data points i and j. The set of N centered feature vectors $\\{h(x_i) - \bar{h}\\}\_{i=1}^N$ is linearly independent if and only if the minimum eigenvalue λ_min(G) > 0. Our experiments track this minimum eigenvalue throughout training, demonstrating that $\mu$P maintains feature diversity while IP shows catastrophic eigenvalue collapse. These results directly validate our theoretical claims. We have additional experimental results here: https://anonymous.4open.science/r/mup_rebuttal-0E0B/Rebuttal.pdf.

Q2. Regarding SP and NTP at large learning rates and feature learning

A2. Thank you for pointing out the paper "Scaling Exponents Across Parameterizations and Optimizers." It is relevant and we will add a citation and brief discussion about it in the revision. This work empirically characterizes when feature learning occurs.

Our work provides a rigorous mathematical framework guaranteeing global convergence during feature learning. The referenced paper shows that SP and NTP can empirically exhibit feature learning at large learning rates but lack theoretical guarantees. The existing NTK analysis framework cannot directly analyze feature learning in SP and NTP at large learning rates, creating a theoretical gap.

While the precise theoretical characterization of SP and NTP at large learning rates remains an open problem in the community, our analysis provides new insights that could potentially be extended to study this challenge in future work.

Q3. Construction of feature matrix from joint space-time features in Figure 2.

A3.In Figure 2, we analyze the combined space-time feature diversity by constructing a feature matrix that simultaneously captures both initial and final representations. This approach provides a comprehensive assessment of how features evolve during training while maintaining linear independence.

For each parameterization scheme, we construct a combined representation matrix as follows:

Let $h_0(x_i) \in \mathbb{R}^n$ represent the feature vector at initialization for input $x_i$, and $h_T(x_i) \in \mathbb{R}^n$ represent the feature vector after training completion. For our dataset with $N$ samples, we form the combined feature matrix $H_{\text{combined}} \in \mathbb{R}^{n \times 2N}$ by concatenating both initial and final representations:

$H_{\text{combined}} = \begin{bmatrix} h_0(x_1) & h_0(x_2) & \cdots & h_0(x_N) & h_T(x_1) & h_T(x_2) & \cdots & h_T(x_N) \end{bmatrix}$

We then compute the Gram matrix $G = H_{\text{combined}}^T H_{\text{combined}}$, where each element $G_{ij}$ captures the similarity between combined representations. The minimum eigenvalue of this Gram matrix quantifies the linear independence across both spatial dimensions (different inputs) and temporal dimensions (initialization versus final state).

This combined analysis provides a stronger test of feature diversity than analyzing features in isolation. As demonstrated in Figure 2, μP maintains significantly higher minimum eigenvalues compared to other parameterizations as width increases, confirming its unique capability to preserve feature diversity throughout training while enabling substantial feature learning.

For example, if a network doesn't learn meaningful features (e.g., NTP at large widths) and $h_T(x_i) \approx h_0(x_i)$ for all inputs, the combined matrix would have linearly dependent columns, yielding a minimum eigenvalue near zero. The higher minimum eigenvalues for μP confirm it learns new features that are linearly independent from initialization.

Q4. Realistic nature of Assumption 4.1 for practical settings beyond toy models

A4: Thank you for this question. Assumption 4.1 requires distinct inner product magnitudes between data points. Our empirical verification (sampling 5,000 random triplets from each dataset) shows:

These results definitively confirm that Assumption 4.1 is not merely a theoretical convenience but reflects geometric properties inherent to real-world datasets. The perfect compliance across standard benchmarks validates that our theoretical framework directly applies to practical settings beyond toy models.

Thank you for your thorough review. Due to space constraints, our responses are necessarily concise while addressing all key points:

Q1. How does your SGD analysis account for mini-batch randomness when your theoretical focus is on full-batch GD?

A1. Our analysis based on the Tensor Program framework already accounts for mini-batch randomness in SGD. Definition 3.1 models this with the indicator function $\mathbf{1}\\{i ∈ \mathcal{B}_t\\}$, and our derivations incorporate this sampling process. Under SGD, the induced Gaussian processes maintain the same covariance structure as in GD, with the additional randomness from batch selection not affecting the feature independence property central to our global convergence proof.

Q2. Why aren't your experiments more comprehensive in validating theory and testing assumption necessity?

A2. CIFAR-10 provides an ideal balance between complexity and tractability for testing our theoretical claims across multiple parameterizations. In the revised manuscript, we expanded our experimental section with comprehensive evaluations across varying network widths and depths.

Q3. Which layer features do your plots show, $h_2$? Could you add x2 results to validate both pre/post-activation feature independence?

A3. Yes, it is $h_2$. We've analyzed both pre/post-activation features across all layers. Our supplementary analysis [https://anonymous.4open.science/r/mup_rebuttal-0E0B/Rebuttal.pdf] includes complete eigenvalue spectrum plots demonstrating linear independence is maintained throughout training in all layers. These results validate Theorem 4.5 across feature types and network depth. The contrast between $\mu$P and IP becomes more significant in deeper layers, where IP shows catastrophic eigenvalue collapse while $\mu$P maintains robust feature diversity.

Q4. Why do Tanh and ReLU show similar results despite not meeting your theoretical assumptions?

A4: Tanh satisfies the GOOD function properties in Assumption 4.3. ReLU was included for completeness despite its non-smoothness. As we already discussed in Appendix A.1, extending our analysis to non-smooth activations is a promising future direction. ReLU's similar empirical behavior suggests our theory's core mechanisms may be more general than the specific assumptions in our proofs.

Q5. The authors claim to follow [1], but their setup is actually closer to [2].

A5. Both schemes give the same dynamic. [1] appears to scale learning rate uniformly. But this difference arises because [1] expresses weights as $w_{\ell}$ with scaling $a_{\ell}$ in $W^{\ell}=n^{-a_{\ell}}w^{\ell}$, while [2] directly uses $W_{\ell}$. Such equivalence is justified in Tensor Program V (arXiv:2203.03466).

Q6. No explicit constraints on learning rate for global convergence. Paper only specifies $\eta = \Theta(1)$ without clear convergence conditions

A6. We prove that when convergence occurs under $\mu$P parameterization, the non-degeneracy of features ensures it will be at a global minimum. We didn't provide explicit convergence conditions because our focus was on characterizing the convergence point, not the convergence process itself. Determining explicit convergence condition of $\eta$ for $\mu$P remains an important open question.

Q7. The paper omits some recent theoretical works on $\mu$P and NTK for infinite-depth neural networks.

A7. Thank you for pointing this out. We already discussed Tensor Program 6 in our paper and will discuss the other suggested works in our revision.

Q8. Definition of $\alpha_k$ in Lemma C.1

A8. They are arbitrary real numbers representing linear combinations in our independence proof.

Q9. Does a finite-width FFN always converge to this infinite-width case? ... Can the results be extended to the finite-width FFN case? What's the challenge?

A9. Yes - convergence is formally established by Theorem 7.4 in Tensor Program IV. Our empirical results show networks with widths larger than 128 nearly align with infinite-width predictions (Figures 1-4). The main challenge is determining precise convergence rates due to the complexity of tracking higher-order interactions in finite-width networks.

Q10. Can you explicitly write out the covariance expressions for equations (5.3) and (5.4)?

A10. The expressions are provided in Definition 3.1 (Line 187) and Line 212 and restated after Equations (5.3) and (5.4) in Lines 310-311.

Q11. What's the intuition behind GOOD functions, and why is Condition 4 necessary?

A11. GOOD functions maintain feature diversity and ensure non-trivial gradient flow by preventing constant decomposition. Condition 4 helps prevent feature collapse, allowing sigmoid, tanh, and SILU to satisfy our analysis.

Q12. Does your paper assume W and $W^\top$ are independent during training?

A12. No, we consider standard backpropagation.

We thank the reviewer for their thoughtful feedback and valuable suggestions.

Q1. Lack of experiments across different architectures, layers, and seeds.

A1. We focused on MLPs as the conventional building blocks widely used in theoretical studies. While a full investigation across diverse architectures is beyond the scope of this theoretical paper, our core contributions establish the mathematical principles governing feature learning under different parameterizations. To improve experimental breadth, we've expanded our analysis with comprehensive eigenvalue spectrum plots across all three layers for both pre-activation and post-activation features [https://anonymous.4open.science/r/mup_rebuttal-0E0B/Rebuttal.pdf]. These visualizations show μP's advantageous properties becoming more significant in deeper layers.

Layers: While our current plots focus on the second layer, we add results demonstrating the minimum eigenvalue trend across different depths to show that the effect persists. [https://anonymous.4open.science/r/mup_rebuttal-0E0B/Rebuttal.pdf]

Seeds: Robustness across initializations was already considered in our experiments. As detailed in Appendix A (Experimental Details), the presented results are mean values computed over 10 independent trials using different random seeds (seeds 42-51).

Q2. Questions on empirical verification of linear independence and choice of minimum eigenvalue vs. entire spectrum analysis.

A2. Regarding the crucial empirical verification of linear independence among the feature vectors corresponding to different input data points, this is precisely achieved in our work by analyzing the minimum eigenvalue ($\lambda_{min}$) of the $N \times N$ Gram matrix computed from the centered features. This matrix, capturing the inner products between feature representations, is fundamentally related to kernel methods and the Neural Tangent Kernel (NTK). We compute this Gram matrix $G$. If $X$ is the $n \times N$ matrix whose columns are the centered feature vectors $h(x_i) - \bar{h}$, then the Gram matrix is $G = X^{\top} X$. (Its element $G_{ij} = (h(x_i) - \bar{h})^{\top} (h(x_j) - \bar{h})$represents the similarity between features of data points $i$and $j$). Crucially, the set of $N$ centered feature vectors $\\{h(x_i) - \bar{h}\\}\_{i=1}^N$ is linearly independent if and only if the minimum eigenvalue $\lambda{min}(G) > 0$. Linear dependence among these$N$vectors occurs precisely when $\lambda_{min}(G) = 0$. Therefore, observing $\lambda_{min}(G)$ provides a definitive and direct verification of whether the features representing different inputs maintain linear independence.

While focusing on the minimum eigenvalue provides the most direct test for linear independence, we have now included comprehensive eigenvalue spectrum analysis [https://anonymous.4open.science/r/mup_rebuttal-0E0B/Rebuttal.pdf] that examines the entire distribution of eigenvalues. These plots reveal that $\mu$P maintains significantly higher eigenvalues throughout the entire spectrum compared to other parameterizations. Notably, IP exhibits catastrophic eigenvalue collapse at higher percentiles (e.g., eigenvalues dropping to $10^{-7}$ in layer 3), while $\mu$P maintains eigenvalue orders of magnitude larger across all percentiles. This full-spectrum evidence further strengthens our claims about feature diversity under $\mu$P.

Q3. Definition of GOOD function arbitrary? How about $r_1=r_2=0$?:

A3. Regarding Assumption 4.3 / Condition 4: Here, we require that for any real numbers $r_1$ and $r_2$, the function $(r_1 + \phi(x))(r_2 + \phi'(x))$ is not almost everywhere constant. It is worth noting that this condition involves "any real number," not the existence of some real numbers satisfying a property.

The GOOD function definition, including this condition, serves as a sufficient condition that guarantees our theoretical results. For an activation function that satisfies our definition, we can show that feature learning maintains linear independence. This condition allows for a broad class of activation functions, including most commonly used ones (e.g., tanh, sigmoid).

Q4. Previous works investigate covariance structure of neural networks and differentiate them from the current study i.e. they don't explicitly study cross-layer interactions.

A4. Thank you for suggesting these valuable references. As recommended, we will add the related works "A Rainbow in Deep Network Black Boxes" and "Structured random receptive fields enable informative sensory encodings" in our revision. We will also highlight, as you noted, that these works don't explicitly study cross-layer interactions, which is one of the key differentiating aspects of our contribution.