Title: Dimension-adapted Momentum Outscales SGD
Venue: NeurIPS 2025
Accepted: yes
Keywords: power-law random features

scaling laws

momentum

random matrix theory

acceleration

stochastic optimization

compute-optimal
Abstract: We investigate scaling laws for stochastic momentum algorithms on the power law random features model, parameterized by data complexity, target complexity, and model size. When trained with a stochastic momentum algorithm, our analysis reveals four distinct loss curve shapes determined by varying data-target complexities. While traditional stochastic gradient descent with momentum (SGD-M) yields identical scaling law exponents to SGD, dimension-adapted Nesterov acceleration (DANA) improves these exponents by scaling momentum hyperparameters based on model size and data complexity. This outscaling phenomenon, which also improves compute-optimal scaling behavior, is achieved by DANA across a broad range of data and target complexities, while traditional methods fall short. Extensive experiments on high-dimensional synthetic quadratics validate our theoretical predictions and large-scale text experiments with LSTMs show DANA's improved loss exponents over SGD hold in a practical setting.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study the scaling laws for stochastic momentum algorithms for random feature models with the input \ndistribution and ground truth direction satisfying power laws. They show via both theoretical analysis (under certain \nsimplifying assumptions) and experiments that (a) SGD with momentum yields the same scaling law exponent as SGD, and \n(b) the scaling law exponent can be improved if we allow the momentum learning rate to depend on the dimension \n(or the time-dependent effective dimension).\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\nThis seems to be an important paper, as it shows that at least in certain cases, changing the algorithm, in particular, \nletting the momentum parameter depend on the (effective) dimension, can potentially change the exponent of the scaling law,\ninstead of just the constant. Moreover, the authors conduct extensive theoretical analysis and experiments, covering a\nwide range of possible power law random feature parameters, algorithms, and training regimes, to support their claims. \n\nThe framework (Gen-Mom-SGD) the authors use covers many existing algorithms as special cases and might be of \nindependent interest. In particular, this framework/strategy allows continuous-analysis of momentum-based algorithms, \nwhich is often quite tricky.\n\n### Weakness\n\nOne issue of this paper is the lack of a proof sketch section. The main text is mostly about the results, experiments,\nand interpretation. Without a proof sketch section (which can be put in the appendix), it is extremely hard to check the \nsoundness of the proof of a 194 pages long paper.\n\nQUESTIONS:\n* In Assumption 1, you implicitly assume that the input distribution and the ground truth direction share the same \n  eigenvectors. Is this essential? What will happen if they are not completely aligned?\n* Is your rates optimal in any sense? (Is there a way to define optimal in this setting, that is more fine-grained than the minimax rate?)\n\nLIMITATIONS:\nThe authors suggest in the introduction that the theoretical analysis is rigorous (line 30-32), while they actually \nneed to ignore some higher-order terms and use the deterministic equivalent (which is only mentioned in one sentence \naround line 131) at certain point of the analysis. To be fair, the authors do provide experimental results to justify \nthese approximations, but they should be more clear about this non-rigorousness.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper offers a theoretical analysis of momentum-based stochastic optimization in high-dimensional settings, particularly when only a single pass over the data is permitted. The authors study random feature models with power-law spectra (PLRF), employing tools from random matrix theory, continuous-time dynamics, and Volterra integral equations to characterize generalization error. Their main contribution is a dimension-adaptive momentum schedule (DANA), given by \\$\\mu(d) = 1 - c/d^{1/3}\\$, which yields provable improvements over standard SGD.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1. **Originality**: The proposed dimension-dependent momentum schedule is novel and well-justified, especially when contrasted with heuristic or fixed hyperparameter choices commonly used in SGD.\n\n2. **Theoretical Depth**: The analysis draws on a sophisticated toolkit (Volterra integral equations, high-dimensional asymptotics, and random matrix theory) to derive interpretable bounds on generalization error.\n\n3. **Clarity of Exposition**: Despite the technical complexity, the paper maintains a coherent high-level narrative. Key ideas are well-motivated, and derivations are presented in a readable form.\n\n4. **Empirical Support**: Though primarily theoretical, the work includes experiments on PLRF models and linear settings that reinforce its claims.\n\n**Weaknesses:**\n\n1. **Excessive Length**: The paper reads more like a dissertation than a conference submission. Its length far exceeds the typical scope of a NeurIPS paper, and the sheer volume of technical material renders comprehensive verification by reviewers infeasible. While key lemmas appear sound, I did not attempt to check every derivation.\n\n2. **Restrictive Momentum Design**: The analysis assumes a scalar momentum coefficient (\\$\\gamma\\_3\\$) applied uniformly across all dimensions. This simplification may limit DANA’s relevance in anisotropic, high-dimensional settings. The paper compares DANA to Adam, but does not explore vector-valued or coordinate-adaptive alternatives. A discussion of their feasibility could broaden DANA’s appeal.\n\n3. **Limited Noise Modeling**: The authors model gradient noise through deterministic equivalents (e.g., \\$\\mathcal{K}\\_{pp}(t)\\$) and rely on ODE approximations. This sidesteps the stochastic dynamics that drive implicit regularization. A more direct SDE-based approach, as explored in \\[a] and \\[b], could better capture generalization phenomena like the preference for flatter minima \\[c].\n\n4. **Restricted Applicability**: The analysis centers on fixed-feature PLRF models, which limits insights into DANA’s performance in deep networks with learned representations. The LSTM experiments (Section L) offer some evidence, but serve mainly as a proxy. Experiments on models like CNNs or Transformers would better demonstrate practical utility.\n\n**References**\n\n[a] Mandt, Stephan, Matthew D. Hoffman, and David M. Blei. \"Stochastic gradient descent as approximate bayesian inference.\" Journal of Machine Learning Research 18.134 (2017): 1-35.\n\n[b] Li, Qianxiao, and Cheng Tai. \"Stochastic modified equations and adaptive stochastic gradient algorithms.\" International Conference on Machine Learning. PMLR, 2017.\n\n[c] Yang, Ning, Chao Tang, and Yuhai Tu. \"Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions.\" Physical Review Letters 130.23 (2023): 237101.\n\nQUESTIONS:\n1. The paper compares DANA to Adam, but why is the momentum coefficient scalar (\\$\\gamma\\_3\\$)? Would a coordinate-wise or adaptive scheme be more effective in anisotropic, high-dimensional settings?\n\n2. The authors model gradient noise via deterministic quantities (e.g., $K\\_{pp}(t)$, Appendix C, Eq. 7) and reduce the analysis to ODEs. How might an SDE-based approach, in the spirit of \\[a], \\[b], or \\[c], alter the results or insights?\n\n3. LSTM experiments in Section L are interesting, but serve only as indirect evidence. Could DANA be tested on architectures like CNNs or Transformers to assess its relevance in more realistic settings?\n\n4. The paper is too long. Consider trimming the main text and moving technical appendices to a companion or journal version. The current length exceeds what is practical for peer review.\n\n5. The Volterra formulation is elegant but dense. Would it be possible to include plots of bias terms (\\$F\\_{pp}\\$, \\$F\\_{ac}\\$, \\$F\\_0\\$) and variance (\\$K\\_{pp}\\$) over time, for varying \\$\\alpha\\$, \\$\\beta\\$, and \\$d\\$? Visualizing these could aid interpretation.\n\n**References**\n\n[a] Mandt, Stephan, Matthew D. Hoffman, and David M. Blei. \"Stochastic gradient descent as approximate bayesian inference.\" Journal of Machine Learning Research 18.134 (2017): 1-35.\n\n[b] Li, Qianxiao, and Cheng Tai. \"Stochastic modified equations and adaptive stochastic gradient algorithms.\" International Conference on Machine Learning. PMLR, 2017.\n\n[c] Yang, Ning, Chao Tang, and Yuhai Tu. \"Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions.\" Physical Review Letters 130.23 (2023): 237101.\n\nLIMITATIONS:\nThe paper does not fully acknowledge its limitations. In particular, the assumptions of fixed features, a global momentum parameter, and deterministic dynamics constrain the analysis. The authors are encouraged to expand the limitations section and clarify to what extent DANA generalizes to modern, feature-learning models.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

comments: SUMMARY:
This paper provides a comprehensive study of the scaling laws of SGD with momentum in linear regression.

STRENGTHS AND WEAKNESSES:
## Strengths

1. This paper extends the scaling law in linear regression optimized with SGD to SGD with momentum acceleration.
2. The findings that SGD with heavy-ball momentum does not outscale vanilla SGD but Nesterov momentum with dimension adaptation can are impressive.
3. The authors perform extensive experiments.
4. The authors discuss the comparison of DANA against the most important variants of SGD-M extensively in the appendices.

## Weaknesses

1. The technical terms "population bias", "model capacity", "embedding bias" and "variance" are not clearly defined, although they already appear in the literature.
2. It is understandable for such an extensive study, but a little unpleasant to see that the paper has only 4 sections and no subsections at all. Section 3 is also very long, which makes the understanding of the main idea a bit more difficult.

QUESTIONS:
1. How does the variance error occur in the problem setting? From my understanding, the variance error should stem from the label noise, but there is no label noise in this setting.
2. How does DANA-decay compare with the hyperparameter selection in Li et al. (2023) (Ref. [59])?

LIMITATIONS:
This paper does not have any outstanding limitations or negative societal impacts.
TITLE: Official Review
CLARITY: 3
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces DANA, a momentum algorithm that adapts its momentum and learning rate hyperparameters based on model size and data complexity. The core claim is that DANA achieves a superior scaling law exponent compared to SGD and classical SGD-M in several regimes, particularly when $2\\alpha > 1$. The authors leverage a continuous-time analysis of the risk dynamics, combined with RMT, to theoretically prove these advantages. These claims are supported by extensive experiments on the synthetic PLRF model and small LSTM language models.\n\nSTRENGTHS AND WEAKNESSES:\nMy evaluation of this paper is based on limited expertise in the area. Some parts of the paper may be beyond my full understanding, and my review reflects this limitation.\n\nStrengths:\n1.  The paper is well-written and clearly structured.\n2.  The primary contribution, a theoretical proof that DANA can improve scaling law exponents is novel. The theoretical framework appears solid, and the claims are supported by extensive empirical validation.\n\nWeaknesses:\n1.  The theoretical analysis is limited to the PLRF model, which has a quadratic objective. While the LSTM experiments suggest it is a useful proxy, the extent to which these results generalize to the non-convex landscapes of more complex architectures remains unknown.\n2.  The practical guidance is limited by the fact that the optimal DANA schedule requires knowing the exponents $\\alpha$ and $\\beta$, which are unknown for real-world tasks. The proposed heuristic of sweeping a new hyperparameter, $\\kappa_3$, while practical, makes it less clear how DANA's tuning process compares in difficulty and outcome to other advanced optimizers that also require tuning.\n3.  The theoretical proof of outscaling is established for the $2\\alpha > 1$ regime. There is no direct evidence as whether this applies across modern machine learning applications.\n\nQUESTIONS:\n1. Although I have not thoroughly verified the correspondence between the ODE and the discrete-time algorithm, I noticed that the SDE model (Equation (7)) is not a straightforward continuous-time counterpart of DANA, i.e., DANA does not appear to be a direct discretization of the SDE. How should I interpret their relationship? What are the errors between the continuous-time models and DANA?\n2. Given that DANA shares some similarities with the algorithms listed in Table 3, I’m curious about its empirical performance on real-world tasks, particularly when hyperparameter tuning is permitted.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-14"}