Title: Stepsize anything: A unified learning rate schedule for budgeted-iteration training
Venue: NeurIPS 2025
Accepted: yes
Keywords: learning rate schedule

budgeted-iteration training

networks optimization
Abstract: The expanding computational costs and limited resources underscore the critical need for budgeted-iteration training, which aims to achieve optimal learning within predetermined iteration budgets. While learning rate schedules fundamentally govern the performance of different networks and tasks, particularly in budgeted-iteration scenarios, their design remains largely heuristic, lacking theoretical foundations.   In addition, the optimal learning rate schedule requires extensive trial-and-error selection, making the training process inefficient. In this work, we propose the Unified Budget-Aware (UBA) schedule, a theoretically grounded learning rate schedule that consistently outperforms commonly-used schedules among diverse architectures and tasks under different constrained training budgets.  First, we bridge the gap by constructing a novel training budget-aware optimization framework, which explicitly accounts for the robustness to landscape curvature variations.  From this framework, we derive the UBA schedule, controlled by a single hyper-parameter $\varphi$ that provides a trade-off between flexibility and simplicity, eliminating the need for per-network numerical optimization.  Moreover, we establish a theoretical connection between $\varphi$ and the condition number, adding interpretation and justification to our approach.   Besides, we prove the convergence for different values of $\varphi$. We offer practical guidelines for $\varphi$ selection via theoretical analysis and empirical results.   Extensive experimental results show that UBA $\textit{consistently surpasses}$ the commonly-used schedules across diverse vision and language tasks, spanning network architectures (e.g., ResNet, OLMo) and scales, under different training-iteration budgets.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a theoretical framework for a unified learning rate schedule that is budget-iteration aware. The resulting schedule adapts to the loss landscape of the model being trained and better leverages the available training budget across different training conditions. This is done by addressing the worst-case loss function and approximating the optimization trajectory across the loss landscape by approximating it with a sequence of local optima. The authors include a comprehensive theoretical analysis of the proposed framework and the related learning rate schedule. The experimental section extensively explores a wide range of datasets and domains.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The proposed method and framework are well motivated and based on a strong theoretical analysis and an intuitive concept. \n- The paper has extensive theoretical and experimental justifications and analysis of the proposed approach.\n\nWeaknesses:\n- The practicality of the proposed method is not explored or discussed. The authors mention that practitioners can easily leverage this method, but do not provide a discussion about the implementation.\n- The paper focuses on training efficiency, but the authors do not discuss the additional computational cost that the proposed learning rate schedule brings.\n- The abstract and the introduction of the paper address the work from a high-level perspective and include scarce details of the proposed approach. The readers would benefit from more introductory details of the proposed theoretical framework and the proposed learning rate schedule in the initial parts of the paper.\n- The text should be reviewed for typos and grammar. See the formatting section for more details on this.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nAs listed in the weaknesses, the authors do not analyze the computational overhead from the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a quadratic framework for deriving a family of learning rates schedules parameterized by a $\\varphi$ parameter. The schedules show strong results for ResNet training.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper has some interesting ideas but is very limited by a compressed presentation and the lack of proper citations to relevant work.\n\nThe idea of deriving a learning rate schedule numerically optimized for a quadratic has been explored before by Pan et. al (2021) in their Eigencurve paper. This Eigencurve paper is cited in the appendix, so the authors are aware of it. It should be in the body of the paper, and given the extreme similarity in the development framework and concept, the Eigencurve work needs to be properly credited. They also show that cosine-like schedule shapes arise from their framework.\nDefazio et. al. (2023)'s work on theoretically motivated linear decay schedules should also be mentioned, as it's the only other completely theory driven framework for learning rate scheduling, and is currently the largest comparison of learning rate scheduling approaches.\n\nThe experimental results in this paper are strong and encouraging, but it's hard to trust the results given the following limitations:\n - No error bars or standard errors are reported. ResNet experiments are famously noisy, and the differences between different schedules are often at the same level as the noise. Highlighting UBA as the winning method for 75 epoch ResNet-50 training with accuracy 76.00 compared to cosines schedule's 75.99 is not appropriate as the run-to-run variability is more than 0.2%!\n - The use of a fixed learning rate for all schedules doesn't give fair comparisons. This is a major limitation. Generally each schedule will have a different optimal learning rate in practice. The choice of this learning rate is not motivated, and may favor certain schedules over others.\n- Is the budgeted iteration method listed actually the linear decay schedule? this should be labelled as linear decay if so.\n- The introduction is too long and and largely restates the abstract. It could be shortened significantly to make more room for discussion of related work (I see you have a longer related work section in the Appendix that you could use some of).\n\nOverall I like the ideas in this paper and I think it's a good extension of the work of Pan et. al, and I would generally support acceptance if the authors commit to the improvements mentioned above, and a proper phrasing of their work within the context of the existing literature.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nIntroducing an additional tunable parameter into your learning rate schedule is a clear limitation. I think it's discussed in reasonable depth.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a unified budget-aware learning rate scheduling strategy Unified Budget-Aware schedule, which aims to solve the problem of how to efficiently train models under a fixed training iteration budget. The author constructs a theory-driven optimization framework, designs a scheduling function that only requires a hyperparameter $\\varphi$ to control the worst case, and establishes the connection between it and the optimization condition number. Experiments on a large number of vision and language tasks have shown that UBA outperforms existing mainstream scheduling strategies under different models, tasks, and budgets, and has good generalization and robustness.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1: This paper derives a UBA schedule based on a well-motivated min-max optimization framework with a rigorous approach.\n\n2: This paper proves the proposed scheduler is an exact solution to a specific min-max problem and analyzes its convergence in noisy environments.\n\n3: Experiments in vision and language tasks show that the method in this paper outperforms baseline models under various settings.\n\nWeaknesses:\n\n1: The method relies on a local quadratic approximation of the loss surface, which is valid for models with smooth structures, but its applicability is still questionable in transformers with small batch sizes in the early training stage.\n\n2: Although the authors achieved good results by fixing $\\varphi$, there is still a lack of robustness analysis on how to adaptively select $\\varphi$ or when $\\varphi$ is selected incorrectly.\n\nQUESTIONS:\nQuestions:\n\n1: Have the authors evaluated whether the local quadratic approximation still applies in the early stages of training?\n\n2: $\\varphi$ is fixed and related to the condition number. Have the authors considered adapting $\\varphi$ based on the optimization state?\n\n3: Although the authors have verified that it is effective on both SGD and AdamW, the scheduling strategy itself is currently independent of the optimizer. Have the authors considered combining the $\\varphi$ or $η_t$ design with the state term in adaptive gradient methods?\n\nLIMITATIONS:\nPartially. While the limitations are addressed in the conclusion, more discussion on the applicability of the quadratic approximation would improve the clarity and transparency of the work. No societal risks are apparent.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a Unified Budget-Aware (UBA) learning rate schedule designed to address training under limited computational resources and varying model scales. UBA requires tuning only a single hyperparameter, φ, which controls the learning rate step size and balances adaptivity with simplicity. Despite its minimal tuning cost, the method consistently yields marginal but robust improvements across a range of models, tasks, and training iteration budgets.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1. The paper provides a clear and elegant theoretical analysis of step size selection under finite-iteration training budgets, culminating in a unified learning rate schedule governed by a single hyperparameter, φ. This formulation is both principled and practical, offering a clean balance between theoretical rigour and ease of application.\n\nWeaknesses\n\n2. The performance improvements over baseline learning rate schedules are marginal across tasks and architectures, raising questions about the practical impact of the proposed method.\n\n3. Experimental results are reported using single-value metrics without accompanying variance or standard deviation. Given the small performance gains, statistical measures (e.g., mean ± std over multiple runs) are necessary to assess whether improvements are significant and reproducible.\n\n4. While the authors claim UBA generalises across model scales, the experimental validation spans a relatively narrow range—from ResNet-18 to ResNet-50—limiting the strength of the claim. Broader evaluation on larger or more diverse model families would strengthen this argument.\n\nQUESTIONS:\n1. Given the marginal improvements reported, how does the proposed UBA schedule compare to baseline schedules under more challenging or large-scale tasks where training budgets are more constrained?\n\n2. Can the authors provide statistical significance analyses (e.g., mean and standard deviation across multiple runs) to validate that the reported improvements are consistent and not due to random fluctuations?\n\n3. While UBA is claimed to generalize across model scales, have the authors considered evaluating it on a broader set of architectures (e.g., ResNet101, ViT, or Transformer-based models) to better support this claim?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}