Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stepsize anything: A unified learning rate schedule for budgeted-iteration training

Anda Tang, Yiming Dong, Yutao Zeng, zhou Xun, Zhouchen Lin

TL;DR

Budgeted-iteration training aims to maximize model performance under fixed iteration budgets. The authors propose the Unified Budget-Aware (UBA) schedule, derived via a budget-aware min–max optimization that accounts for curvature variation and links a single hyper-parameter phi to the condition number. They prove convergence and demonstrate that UBA outperforms common learning-rate schedules across vision and language tasks with negligible overhead, across architectures such as ResNet and OLMo. The work provides practical guidelines for selecting phi and shows robustness to optimizers, presenting a unified, theory-grounded scheduling approach for budgeted training.

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 its selection via theoretical analysis and empirical results. Extensive experimental results show that UBA 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.

Stepsize anything: A unified learning rate schedule for budgeted-iteration training

TL;DR

Budgeted-iteration training aims to maximize model performance under fixed iteration budgets. The authors propose the Unified Budget-Aware (UBA) schedule, derived via a budget-aware min–max optimization that accounts for curvature variation and links a single hyper-parameter phi to the condition number. They prove convergence and demonstrate that UBA outperforms common learning-rate schedules across vision and language tasks with negligible overhead, across architectures such as ResNet and OLMo. The work provides practical guidelines for selecting phi and shows robustness to optimizers, presenting a unified, theory-grounded scheduling approach for budgeted training.

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 its selection via theoretical analysis and empirical results. Extensive experimental results show that UBA 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.

Paper Structure

This paper contains 46 sections, 4 theorems, 45 equations, 11 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Budgeted training
  4. Learning rate schedule
  5. Budgeted-iteration training
  6. Finite optimization under limited training iterations
  7. Theoretical analysis
  8. Experiment results
  9. Baselines
  10. Experiments for vision classification tasks
  11. Experiments for language models
  12. Ablation study
  13. Performance across different optimizers
  14. Parameter analysis of $\varphi$
  15. Performance across different periods
  16. ...and 31 more sections

Key Result

Proposition 1

The fit function fitted-function is the exact closed-form solution to the min-max optimization problem: when the hyper-parameter $\varphi$ are determined by $\lambda_{l}^{(k)}$ and $\lambda_{u} ^{(k)}$ through the relation $\varphi=2\frac{\lambda_{u}^{(k)}}{\lambda_{l}^{(k)}}$ and $\eta_{\max}=1, \quad \eta_{\min}=0$.

Figures (11)

  • Figure 1: Evolution of the learning rate in UBA schedule across training iterations.
  • Figure 2: Training dynamics and performance for language tasks under 150B tokens on 300M OLMo. We present the training loss, validation loss, and downstream performance on HSWAG and ARC-E, demonstrating that UBA schedule achieves superior performance.
  • Figure 3: Conceptual illustration: visualization of the motivations behind the proposed modeling approaches. (a) shows the rationale for the modeling approach in \ref{['eq1']}, while (b) explains the modeling approach for \ref{['eq2']}. The solid line represents the optimization trajectory, while the dashed circles indicate the local approximation around corresponding minima.
  • Figure 4: The curve fitting of numerical solutions.
  • Figure 5: Visualization of relation between parameter $a, b$ and $c$.
  • ...and 6 more figures

Theorems & Definitions (8)

  • Proposition 1
  • Proposition 2
  • Theorem 1
  • proof
  • proof
  • Lemma 1
  • proof
  • proof