Title: Learning quadratic neural networks in high dimensions: SGD dynamics and scaling laws
Venue: NeurIPS 2025
Accepted: yes
Keywords: scaling laws

stochastic gradient descent

shallow neural network

multi-index model
Abstract: We study the optimization and sample complexity of gradient-based training of a two-layer neural network with quadratic activation function in the high-dimensional regime, where the data is generated as $f_*(\boldsymbol{x}) \propto \sum_{j=1}^{r}\lambda_j \sigma\left(\langle \boldsymbol{\theta_j}, \boldsymbol{x}\rangle\right), \boldsymbol{x} \sim \mathcal{N}(0,\boldsymbol{I}_d)$, $\sigma$ is the 2nd Hermite polynomial, and $\lbrace \boldsymbol{\theta}_j \rbrace _{j=1}^{r} \subset \mathbb{R}^d$ are orthonormal signal directions. We consider the extensive-width regime $r \asymp d^\beta$ for $\beta \in [0, 1)$, and assume a power-law decay on the (non-negative) second-layer coefficients $\lambda_j\asymp j^{-\alpha}$ for $\alpha \geq 0$. 
We present a sharp analysis of the SGD dynamics in the feature learning regime, for both the population limit and the finite-sample (online) discretization, and derive scaling laws for the prediction risk that highlight the power-law dependencies on the optimization time, sample size, and model width. Our analysis combines a precise characterization of the associated matrix Riccati differential equation with novel matrix monotonicity arguments to establish convergence guarantees for the infinite-dimensional effective dynamics.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores the optimization and sample complexity of training two-layer quadratic neural networks in high dimensions, focusing on the extensive-width regime where the teacher network width scales as $ r \\asymp d^\\beta $ $(\\beta \\in (0,1)$) and second-layer coefficients exhibit power-law decay $\\lambda_j \\asymp j^{-\\alpha}$. It derives scaling laws for prediction risk, highlighting power-law dependencies on optimization time, sample size, and model width. Using matrix Riccati ODEs and novel matrix monotonicity techniques, it analyzes SGD dynamics, establishing convergence guarantees for feature learning in population and finite-sample settings, offering insights into neural scaling laws.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**: The paper delivers a robust theoretical framework, elegantly connecting empirical scaling laws to rigorous mathematics through innovative matrix monotonicity and operator norm control, effectively addressing the complex extensive-width regime. It surpasses previous sample complexity bounds, especially in isotropic settings, and validates the additive model hypothesis, deepening understanding of feature learning dynamics. Its clear problem setup, thorough proofs, and relevance to practical neural network training underscore its value. \n\n**weakness**:\n1. The paper relies on idealized assumptions—uncorrelated Gaussian inputs, perfectly orthogonal teacher weights, and a strict power-law decay for second-layer coefficients—that often don’t match real-world data. I understand the challenges of theoretical research; given the complexity of analyzing practical models, it’s nearly impossible to directly characterize scaling laws for large-scale models. This isn’t a flaw in theoretical work. However, I believe the authors should delve deeper into the challenges of extending their findings to more realistic scenarios and candidly address the limitations of their current theoretical progress.\n\n2. It’s also uncertain whether the matrix monotonicity approach applies to other activation functions like ReLU or optimization methods like Adam. The paper doesn’t explore the method’s scalability or broader applicability, which restricts its relevance to larger neural network problems.\n\nQUESTIONS:\nNo\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-25"}

comments: SUMMARY:
The authors study GD of the population loss on two layer networks with $x^2-1$ activation and orthonormal weights in the first layer, learning a similar 2-layer architecture. They then proceed to use the closed form solution of the dynamic to decompose the risk into a sum of independent contributions, each given by a hidden unit. This result allows the authors to derive some neural scaling laws. These results hold both for gradient flow and discrete-step GD

STRENGTHS AND WEAKNESSES:
The paper is clear and well written. The proof is elegant, although not very technically innovative. The scaling laws, which constitute in my opinion the main result of this paper, are intuitively reasonable, but I would have appreciated extensive numerical simulations to verify them empirically. While the scaling laws itself are not particularly surprising, there are not exact results on learning linear width networks, so I believe the contribution is valid. I think the paper would be stronger if it also tackled the case of generic covariance and label noise, but even as is it's in my opinion a strong paper.

QUESTIONS:
1. Would you be able to consider the case of arbitrary covariance using the same proof technique?
2. Would the derivation change if you were to consider an activation $x + {\rm He}_2(x)$?
3. Do you expect the power-law to change if you were allowed to use the data for more than one GD iteration?
4.Minor issue: could you explicitly define $\land$ and $\lor$ as min and max somewhere in the paper?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is an interesting paper with delicate theoretical results on feature learning and risk. The set-up assumes that the teacher model is a two-layer network with a quadratic activation function and with orthogonal features and algebraically decaying coefficients. It is about the extensive-width regime when the number of student model neurons is growing in specific ways with the dimension d of the data and the sample or observation size from the teacher model.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nNon-trivial mathematical results\n\nThe feature-learning setting is indeed more interesting than the lazy regime.\n\nBetter understanding of SGD dynamics and scaling laws is important\n\nThe \"matrix-monotone comparison framework\" might be of independent interest. However, I don’t have the expertise to evaluate this contribution and I don’t see a convincing story-line for whom this might be of interest.\n\nWeaknesses:\n\nThe theoretical results seem to be solid, but they are asymptotic as d tends to infinity. More clear motivations for the coefficient decay assumption are necessary. How about other decay rates?  It would be insightful to see some simulation studies about the convergence rates of these asymptotic results, and under the scenario that the teacher model is observed with noise.\n\nA key weakness of the paper is its reliance on a series of strong and specific assumptions that, while enabling the theoretical analysis, distance the results from practical applications. These assumptions include the following: Quadratic activation function (polynomial activation functions are not very popular activations functions in practice, partially because they violate the universal approximation theorem and because their gradients can explode). Also the theory only covers shallow neural networks with only 1 hidden layer. Additionally the gradient dynamics were simplified (as commonly done in theoretical literature). Additionally infinitely large training datasets are assumed. Additionally the limit $d\\to\\infty$ is taken. (Additionally the result does not cover transformers or convolutional NNs.) Additionally you assume that the true data generating process and the model you are training have an extremely similar structure. Additionally you assume orthogonal features, specific coefficient decay, etc.. These are multiple very specific assumptions which make the result very niche.  \nThere are no empirical experiments with real-world architectures using ReLU or, GeLU or any non-polynomial activation function. In practice these assumptions are (almost) never met, so the paper is mainly motivated by the hope that this theory can give insights beyond its very restricting assumptions.\n\nThe presentation needs to be improved so the main take-home messages are clearer and more crisp, especially for practitioners. What does a practitioner learn from this paper that they can use to guide their practice?  How does one empirically check that the assumed teacher model captures reality for any real data sets? The paper is hard to follow. There are multiple notation-heavy mathematical results, but not a very clear storyline. \n\nMoreover, the connections and differences between the submission and other prior works need to be made concise and clear. What precise and insightful advances have been made in this paper relative to existing works? E.g., Emergence and scaling laws in SGD learning of shallow neural networks by Yunwei Ren, Eshaan Nichani, Denny Wu, Jason D. Lee (https://arxiv.org/abs/2504.19983) seems highly related and was not cited. This preprint was only uploaded a few days before the submission deadline, so if you did not follow oral talks of this paper multiple months ago, not citing this paper is OK. Under this very specific assumptions, you might be the first to have proven this very specific result, but the question is if you have derived any impactful new insights.  The main novelty of your work seems to be the information exponent =2. A general result for an information exponent larger than 2, seems much more interesting to me than a result which only holds under your extremely restricting sets of assumptions.\n\nQUESTIONS:\nQuestions:\n“equivalently” in line 147 does not imply that the gradient dynamics are equivalent? The authors should clarify that this refers to an equivalent representation of the student network, not necessarily equivalent gradient dynamics.\n\n“which is negligible relative to the feature learning cost.” in line 207 only holds for d>>r_s? This should be emphasized by the authors.\n\nComments: \n\nYou mention scaling laws, but you could be more explicit about how the paper attempts to provide a theoretical underpinning for the empirically observed \"neural scaling laws\" in large language models. This is a central motivation for the work.\n\nConclusion:\nThe paper could be very interesting for a small number of highly specialized scientists working on very specific types of theory for shallow neural networks with quadratic activation functions that learn data generated from shallow neural networks with a quadratic activation function. However, this paper does not meet the required broad impact for NeurIPS and suffers from poor presentation. We highly recommend submitting this work to a more specialized venue where people working on similar topics appreciate this technically impressive but very niche result.\n\nLIMITATIONS:\nThis work does not discuss limitations. It would be useful to add a discussion section to the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-07"}