PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimization Hyper-parameter Laws for Large Language Models

Xingyu Xie, Kuangyu Ding, Shuicheng Yan, Kim-Chuan Toh, Tianwen Wei

arXiv:2409.04777v4cs.LGmath.OC

TL;DR

Large-scale LLM training makes hyper-parameter choices costly, while existing scaling laws offer limited guidance on time-varying learning-rate schedules. Opt-Laws predicts final loss from small-scale experiments using SDE-motivated features across model and data scales. It achieves a 94% held-out Top-2 hit rate, detects divergence with F1=0.92, and identifies the best schedule family in all five evaluated out-of-family groups.

  • Problem

    Scaling laws guide model-size and data choices but few studies address time-dependent hyper-parameters such as learning-rate schedules, leaving practitioners reliant on heuristics.

    Source: §1 · p. 2

  • Method

    Opt-Laws uses small-scale data and SDE-based convergence and escape analyses to construct features that predict final loss across model sizes, data budgets, and learning-rate schedules.

    Source: §1 · pp. 2–3

  • Results

    94% Top-2 hit rate on the held-out benchmark, F1=0.92 for divergence detection, and correct identification of the best schedule family in all five out-of-family groups.

    Source: §7 · p. 34

  • Takeaways & Limitations

    Opt-Laws supports pre-selecting learning-rate schedules for pre-training, continual training, and fine-tuning before full-scale runs.

    Source: §1 · p. 3

    Bridging the discrete fixed-β framework and continuous coupled-β SDE framework remains an open problem beyond the scope of this work.

    Source: §6.3.2 · p. 33

Abstract

from arXiv · show

Large Language Models have driven significant AI advancements, yet their training is resource-intensive and highly sensitive to hyper-parameter selection. While scaling laws provide valuable guidance on model size and data requirements, they fall short in choosing dynamic hyper-parameters, such as learning-rate (LR) schedules, that evolve during training. To bridge this gap, we present Optimization Hyper-parameter Laws (Opt-Laws), a framework that predicts final training loss as a function of LR schedule, model size, and data size. Grounded in SDE-based convergence and escape analyses, Opt-Laws yield interpretable convergence and escape features that predict final training loss across model scales, enabling schedule pre-selection from small-scale experiments. Empirically, Opt-Laws achieve a 94% Top-2 hit rate for identifying near-optimal schedule candidates on held-out configurations, correctly identify the best-performing schedule family in all five evaluated out-of-family settings, and detect training divergence with F1 = 0.92.

1 Introduction

Opt-Laws addresses the limited guidance scaling laws provide for time-varying training hyper-parameters by predicting final loss from small-scale experiments. It combines SDE-motivated features with empirical evaluation of schedule selection, divergence detection, and generalization.

1 Introduction

  • Existing scaling laws cover model size and data volume, but few address time-dependent hyper-parameters such as learning-rate schedules.Practitioners therefore often rely on heuristics, while fine-tuning and continual training add distribution shift and a retention-versus-adaptation trade-off.Source: §1 · p. 2
  • Opt-Laws also offers qualitative interpretations of training patterns and is validated on MoE models with up to 4 billion trainable parameters and over 450 billion tokens.The reported scope includes pre-training, continual training, and fine-tuning.Source: §1 · pp. 2–3
  • Opt-Laws fits small-scale experiments to predict final training loss across model sizes, data budgets, and learning-rate schedules, supporting schedule pre-selection before full-scale runs.The framework uses 15 SDE-motivated features plus an intercept, based on convergence and escape analyses for SGD and Adam.Source: §1 · pp. 2–3
  • 94% Top-2 hit rate for near-optimal schedule candidates, F1=0.92 for divergence detection, and correct identification of the best schedule family in all five out-of-family groups.Evaluations also cover held-out prediction and schedule ranking in continual training and fine-tuning.Source: §1 · p. 3

2 Related Work

Related work develops scaling laws for model and data size, dynamical-systems analyses of optimization convergence, and SDE-based accounts of escape from local minima. Opt-Laws extends these strands toward schedule-dependent final-loss prediction and time-varying optimization analysis.

2.1 Scaling Laws

  • Scaling laws chiefly relate model size, data volume, and loss, while recent schedule-aware predictors target per-step loss; Opt-Laws instead predicts final loss across scales and training stages.Earlier work also extends power laws to other hyper-parameters and downstream settings.Source: §2.1 · pp. 3–4

2.2 Convergence Analysis via Dynamical Systems

  • The paper situates its SDE convergence analysis among work on ODEs, differential inclusions, and SDEs, noting that general non-convex convergence bounds for arbitrary learning-rate policies remain limited.It presents its bounds as among the first for general non-convex problems without requiring a specific learning-rate pattern.Source: §2.2 · p. 4

2.3 Escaping analysis via SDEs

  • For escape under time-inhomogeneous SDEs, the paper uses a Gaussian approximation because the classic Eyring-Kramers approach is not directly applicable.Prior SDE escape analyses largely estimate exit times from sharp local minima.Source: §2.3 · pp. 4–5

3 Optimization Hyper-parameter Laws

The simplified Opt-Laws model predicts final loss at fixed model size from the learning-rate schedule, while the section illustrates its fit and uses its convergence and escape terms to interpret schedule effects. The qualitative analysis emphasizes that schedule sensitivity depends on data scale and schedule configuration.

3 Optimization Hyper-parameter Laws

  • The fixed-model section presents the simplified law as a schedule-to-final-loss relationship, with the generalized multi-scale extension and held-out evaluation deferred to Section 4.This marks the scope boundary for the qualitative fixed-size analysis.Source: §3 · p. 5, §3.2 · p. 7

3.1 Opt-Laws with Fixed Model Size

  • The fixed-size law relates final training loss to the learning-rate schedule, under a setup that sets the initial and final learning rates to zero for simplicity.Its constants and power terms are fitted on a small model with similar architecture and a small data subset before large-scale use.Source: §3.1 · p. 5
  • The equation’s schedule features distinguish convergence speed from ease of escaping a local region, with schedule derivatives contributing to the escape terms.The model’s constants depend on the model and training data.Source: §3.1 · pp. 5–6
  • At sufficiently large token budgets, the low-loss region broadens and warmup sensitivity weakens, although very short warmup remains harmful.The figure description compares warmup and normalized peak learning rate across four token budgets.Source: §3.1 · p. 6, Figure 1 · p. 6

3.2 Illustrative Fixed-Model Fit

  • As training data increases, the strong-performance region expands and becomes less sensitive to warmup, while excessively short warmup remains harmful; regression error on convergent runs is below 0.5%.The illustration fits an 8 × 0.1B MoE model across 3B, 6B, 10B, and 30B tokens.Source: §3.2 · p. 7

3.3.1 Influence of Warmup Steps on Training Loss

  • When token count exceeds the Chinchilla ratio of 25.6 tokens per parameter, effective warmup durations broaden and final loss becomes insensitive to warmup; fine-tuning at lower ratios is more sensitive.The section relates this pattern to prior continual-training findings and to narrower optimal warmup ranges in fine-tuning.Source: §3.3.1 · pp. 7–8

3.3.2 Insights into LR schedule Effects through Opt-Laws

  • At smaller training horizons, a constant-learning-rate phase can yield lower predicted loss than cosine decay under the stated warmup and cooldown settings, while the ordering reverses without that phase.The examples set warmup to 0.01S and the constant phase endpoint to 0.85S; the no-constant-phase comparison sets a = ac.Source: §3.3.2 · p. 9
  • As data size grows, cosine and constant-phase schedule predictions converge to a schedule-independent asymptotic loss determined by model capacity and data properties.This interpretation reconciles negligible schedule sensitivity at large data scales with schedules that reach performance limits faster during fine-tuning or continual training.Source: §3.3.2 · pp. 8–9

4 Generalized Opt-Law across Scales

Generalized Opt-Laws extends schedule-based loss prediction across model and data scales, using a divergence criterion and convergence-escape features. On the held-out training grid, it outperforms all baselines on five metrics and achieves Top-2 94% and Regret 0.003.

4 Generalized Opt-Law across Scales

  • The multi-scale training grid reveals non-monotonic loss structure and divergence boundaries that the simplified unimodal law cannot capture.Divergent runs are assigned loss 7 in the grid.Source: §4 · p. 9, §4.1 · p. 10

4.1 Predicting Training Divergence

  • The divergence criterion compares below-threshold warmup area with above-threshold LR exposure; R > 1 indicates likely failure, while R < 1 indicates hyper-parameters unlikely to cause divergence.The threshold ηL is inversely related to model size and directly related to data size.Source: §4.1 · pp. 10–11

4.2 Generalized Opt-Laws

  • The generalized feature family combines schedule-dependent convergence and escape bases with model- and horizon-scale terms and restricted interactions.The compact, theory-motivated basis is fitted by linear regression; its scale-aware features expose a convergence-escape trade-off.Source: §4.2 · pp. 11, 13

4.3 Evaluation on the Training Grid

  • Held-out predicted surfaces closely match actual losses and correctly rank configurations within each model-size and token-budget block.The predicted surface is smooth and does not reproduce local empirical irregularities.Source: §4.3 · pp. 15–16
  • Top-2 94% and Regret 0.003; generalized Opt-Laws improves on all baselines across five held-out benchmark metrics.The selected configuration is best or second-best in nearly all blocks, with expected loss sub-optimality below 0.3%.Source: Table 1 · p. 15, §4.3 · pp. 16–17
  • The divergence criterion achieves F1=0.92 on 260 training-grid configurations, outperforming a peak-LR threshold (F1=0.38) and logistic regression (F1=0.54).Divergent runs are excluded from loss regression but retained to fit and evaluate the divergence criterion.Source: §4.3 · pp. 15, 17

5 Evaluation beyond the Training Grid

The evaluation tests Opt-Laws beyond its fitting grid, spanning schedule families, model and data scales, continual training, and fine-tuning. Predictions and rankings transfer across these settings, with boundaries varying by schedule and training scenario.

5.1 Effects of Learning Rate Schedules

  • At 3B tokens, constant-plus-linear yields the lowest loss; family gaps narrow at 10B and 100B as convergence dominates, while linear slightly beats cosine at equal budgets.The larger token budget increases convergence contributions and reduces the relative contribution of escape terms that distinguish schedules.Source: §5.1 · pp. 17–19
  • Gap narrowing is slower for 8×0.6B than 8×0.1B; at 10B tokens cosine and linear nearly meet for the smaller model but still differ visibly for the larger.Opt-Laws attributes this sensitivity to convergence terms scaling with model size at the same token budget.Source: §5.1 · pp. 17–18

5.2 Pre-training Extrapolation

  • Without re-fitting, errors stay below 0.2% for 300B simple schedules and below 0.5% for 100B multi-phase schedules; inverse design yields nearly equal final losses across distinct schedules.The tested schedules extend beyond the small-scale fitting grid.Source: Table 2 · p. 19, §5.2 · pp. 19–20, Figure 7 · p. 20
  • At 300B tokens, generalized Opt-Laws reports 0.1% mean relative error versus 8.9–63.8% for baselines, and maintains errors below 0.5% across all three groups.Methods are fitted on the small-scale training grid (Table 7) and evaluated on configurations in Table 2 without re-fitting.Source: Table 3 · p. 20, §5.2 · p. 21

Schedule-family generalization.

  • Opt-Laws identifies the best schedule family in all five evaluated out-of-family groups, with prediction errors below 2%; the authors do not claim validity for arbitrary schedules.Tests cover cosine and constant-plus-linear schedules, while the fitting grid contains linear warmup-plus-cooldown schedules.Source: Schedule-family generalization · p. 21

5.3 Continual-Training Ranking

  • History-aware ranking carries the pre-training convergence budget into continual training, while the escape features use only the current schedule and normalize its peak learning rate.The method accounts for optimization history because continual training starts from a pre-trained checkpoint.Source: §5.3 · pp. 21–22
  • Under strong shift, Opt-Laws recovers the exact 1–2–3 ranking; under weak shift, losses differ by 0.005 and the near-tied top two limit ranking resolution.The top two weak-shift schedules differ by only 0.001 in observed loss; rankings are reported as correct where observed differences are resolvable.Source: §5.3 · pp. 22–23

5.4 Fine-Tuning Ranking

  • Across nine fine-tuning configurations after pre-training and continual training, predicted and actual rankings agree perfectly, including best and worst candidates at 60B tokens.At 3B and 30B tokens, all three candidates are correctly ranked; at 60B, one candidate uses a two-phase schedule.Source: Table 5 · p. 23, §5.4 · p. 24

6 SDE Analysis behind Opt-Laws

This section develops SDE-based convergence and escape analyses for non-convex optimization with time-varying learning rates, motivating Opt-Laws’ interpretable features. It relates learning-rate schedules and model and data size to progress toward stationarity, basin selection, and diffusion, while noting the analyses’ assumptions and an open bridge between coupled-β SDEs and fixed-β discrete analyses.

6 SDE Analysis behind Opt-Laws

  • SGD and Adam SDE bounds have similar forms, linking learning-rate schedules, data size, and model size to convergence and escape behavior.This similarity supports generalizing Opt-Laws to relate these factors to final training loss.Source: §6 · p. 24
  • For general non-convex smooth problems with time-varying learning rates, the analysis extends beyond prior SDE results focused on convex or quadratic objectives or fixed rates.The resulting nonlinear, time-inhomogeneous SDEs motivate new theoretical analysis for LLM optimization.Source: §6 · p. 24

6.1 Optimization Methods and SDEs

  • SGD and Adam dynamics are modeled as SDEs with time-varying learning rates, but the analysis assumes Gaussian gradient noise because its exact training distribution remains unresolved.For Adam, the SDE analysis uses a single-time-scale parameterization of its moment coefficients.Source: §6.1 · p. 25

6.2 Inspiration from Convergence Guarantee

  • Under smoothness and noise assumptions, the section establishes covariance bounds that support the subsequent convergence analysis.The analysis fixes batch size at B = 1 for simplicity and states that results adapt to any fixed positive batch size.Source: §6.2 · pp. 26–27

6.2.1 SGD Convergence Analysis

  • The SGD convergence theorem bounds the average squared gradient norm using training hyper-parameters, with its order scaling as O(N divided by the integrated learning-rate schedule).The averaged squared gradient norm serves as a trajectory-level convergence measure; the feature construction uses its relationship to the schedule without requiring a PL condition.Source: §6.2.1 · pp. 27–28

6.2.2 Adam Convergence Analysis

  • Under its stated assumptions, Adam’s momentum and gradient bounds share the same convergence rate in the learning-rate schedule, paralleling the SGD relationship.The discussion expresses the shared bound as O(N^γ divided by the integrated schedule), with γ ≥ 1; Opt-Laws uses data-driven constants to account for practical rates.Source: §6.2.2 · p. 29

6.3 Inspiration from Escaping Capacity

  • The escape analysis quantifies how learning-rate choices affect the probability that SGD-SDE and Adam-SDE leave a local region.It is motivated by the use of large learning rates followed by cooldown to help move away from sharp local minima.Source: §6.3 · p. 30

6.3.1 Linearization Approximation of SDEs

  • Because the time-dependent Fokker–Planck equation is difficult to solve directly, the analysis locally linearizes its drift and diffusion, yielding a Gaussian approximation governed by mean and covariance ODEs.The approximation enables analytical estimation of escape probability from a local region.Source: §6.3.1 · pp. 30–31

6.3.2 Escape Probability from Local Minima

  • Theorems 3 and 4 provide lower bounds on the probability of leaving an ε-radius ball around x*, with larger bounds indicating stronger escape capacity.A smaller probability that the final iterate remains in the local region after the LR cools indicates greater capacity to escape suboptimal minima.Source: §6.3.2 · p. 33
  • Escape features capture how LR cooldown abruptness constrains basin selection, complementing convergence features that measure progress toward stationarity.Gradual cooldown sustains noise-driven diffusion and allows migration toward better basins, whereas abrupt cooldown locks the iterate into its current region.Source: §6.3.2 · p. 33
  • Top-2 rose from 66% to 89% when escape features were added to convergence-only features, confirming their independent predictive signal.The ablation attributes this gain to escape features that capture cooldown abruptness and its effect on basin selection.Source: §6.3.2 · p. 33
  • In Adam-SDE, coupling β1,k and β2,k with ηk is necessary for a non-degenerate SDE; with fixed β as η approaches zero, the dynamics reduce to deterministic normalized gradient flow (Malladi et al., 2022).The SDE provides escape-related insights inaccessible from discrete analysis alone, while bridging the coupled-β SDE and fixed-β discrete frameworks remains open.Source: §6.3.2 · p. 33

7 Conclusion

Opt-Laws predicts final LLM training loss from the learning-rate schedule, model size, and data size. Its schedule features separate convergence from escape dynamics, and the conclusion reports strong held-out and out-of-family results alongside bounded empirical validation.

7 Conclusion

  • Opt-Laws predicts final training loss as a function of learning-rate schedule, model size, and data size.Its 15 interpretable features are fitted from small-scale experiments.Source: §7 · p. 34
  • Convergence features derived from integrated learning rate capture approach to stationarity, while escape features capture which basin the optimizer ultimately settles into.Together, the features bridge SDE-motivated analysis and data-driven prediction.Source: §7 · p. 34
  • 94% Top-2 hit rate and Spearman correlation of 0.84 on the held-out benchmark; divergence detection reached F1=0.92, and the best schedule family was identified in all five out-of-family groups.The history-aware extension also produced correct rankings across continual-training and fine-tuning pipelines.Source: §7 · p. 34
  • Empirical validation covers MoE models up to 4B trainable parameters and piecewise-linear, cosine, and constant-plus-linear schedules; broader validation is a stated next step.Appendix A supplements the main experiments with setup, evaluation protocol, exponent selection, ablations, divergence prediction, and schedule-family generalization details.Source: §7 · p. 34, Appendix A · p. 34

A.1 Model Configurations and Shared Training Settings

The experiments use MoE model configurations with shared training settings. The reported setup includes architectural and training parameters, with a zero minimal learning rate chosen to support schedule comparisons.

A.1 Model Configurations and Shared Training Settings

  • Table 6 covers architecture parameters and shared training settings, including model dimensions, parameter counts, token length, batch size, optimizer, weight decay, and gradient clipping.The settings are described as following standard LLM training practices without special adjustments.Source: §A.1 · p. 34, Table 6 · p. 35
  • The minimal learning rate is set to 0.0 to maximize comparability across schedules and avoid an artificial performance floor from a non-zero residual rate.Source: §A.1 · p. 34
  • Table 7 inventories the small-scale training grid by model-size and token-budget block, with 35–40 hyper-parameter configurations per block and divergence defined as final loss above 6.Source: Table 7 · p. 35

A.2 Fitting Setup and Evaluation Protocol

The held-out evaluation uses a shared training grid to fit and test Opt-Laws across configurations. Regression, divergence prediction, and schedule-family generalization are evaluated as separate tasks.

A.2 Fitting Setup and Evaluation Protocol

  • The grid contains 260 configurations across seven model-size and token-budget blocks, varying peak learning rate from 10−3 to 1.5×10−2 and warmup from 128 to 6000 steps.Runs use linear warmup followed by linear cooldown; 234 configurations converge and 26 are labeled divergent when final loss exceeds 6.Source: §A.2 · p. 35
  • Five-fold cross-validation partitions configurations within each block, fits on combined training folds, and reserves held-out folds for evaluation.Constant-plus-linear cooldown runs are excluded from fitting and reserved for schedule-family generalization.Source: §A.2 · p. 36
  • The grid supports simplified-law regression, generalized-law regression, and divergence prediction, which are evaluated independently.Regression uses convergent configurations, while divergence prediction includes all 260 configurations.Source: §A.2 · pp. 35–36

A.2.1 Metric Definitions

The appendix defines held-out metrics and describes how generalized Opt-Laws exponents and coefficients are fitted. It also specifies normalization and the fitted divergence criterion.

A.2.1 Metric Definitions

  • R2 measures held-out log-loss fit, while relative error is the mean absolute relative error over convergent configurations.Source: §A.2.1 · p. 36
  • Spearman measures rank agreement within each block, Top-k Hit Rate tests whether the predicted best is among the two actual best, and regret measures excess loss.The reported ranking metrics average across block-fold evaluations.Source: §A.2.1 · p. 36

A.3 Exponent Selection and Fitted Coefficients

  • Three SDE-derived exponents are fixed at ±1.0; remaining exponents are selected from {0.15, 0.20, 0.25, 0.50, 1.0} using a two-level search.Group-level combinations are searched first, then the most sensitive term within each group is adjusted, with held-out folds reserved for final evaluation.Source: §A.3 · p. 37

A.3.1 Fitted Coefficients

  • The generalized law fits coefficients for 15 features, including correlated interaction terms, without constraining coefficients to be positive.Individual coefficient signs represent partial effects conditional on the other features.Source: §A.3.1 · p. 37
  • Model size, training steps, step-related variables, and learning rate are normalized before estimating model loss.Model size is expressed in billions, step counts are converted to token counts in billions, and learning rate is divided by 1.5 × 10−2.Source: §A.3.1 · p. 37

A.3.2 Divergence Criterion Coefficients

  • The divergence criterion is fitted on small-scale training-grid runs labeled by whether final loss exceeds 6, with normalized step and warmup values squared for fitting stability.Source: §A.3.2 · pp. 37–38

A.4 Feature Ablation and Sensitivity Analysis

The ablations show that convergence, escape, and interaction features add predictive signal beyond scale terms alone, with the full feature set reaching Top-2 94%. The default convergence integral excludes the decay phase, and schedule-dependent exponent perturbations are robust while scale-power perturbation is more sensitive.

A.4 Feature Ablation and Sensitivity Analysis

  • Top-2 reaches 94% and Spearman reaches 0.84 with mixed interactions added to the feature set.The caption describes the final regression model as 15 non-bias features plus one intercept.Source: §A.4 · p. 38, Table 9 · p. 39
  • Adding convergence terms raises Top-2 from 51% to 66% and halves regret from 0.027 to 0.016, while escape terms raise Top-2 from 66% to 89%.The escape terms also raise Spearman from 0.61 to 0.77; scale terms alone capture moderate within-block ranking.Source: §A.4 · p. 38
  • Including the decay phase slightly increases pre-training mean prediction error from 0.15% to 0.23%, while continual-training rankings are identical under strong shift.Under weak shift, Spearman falls from 0.90 to 0.30, but the actual losses lie within a 0.005 range; overall, the default exclusion performs slightly better.Source: §A.4 · p. 39
  • Perturbing schedule-dependent exponents preserves Top-2 at 94% and R2 at least 0.998, whereas setting the scale-power exponent to α=0.50 lowers Top-2 to 80%.The authors conclude that feature structure supplies the primary predictive signal, with exact exponent values playing a refinement role.Source: §A.4 · p. 39

A.5 Divergence Prediction Benchmark

The appendix evaluates Opt-Laws as a divergence classifier and tests schedule-family transfer, then supplies derivations and proof steps for the theoretical framework. The reported benchmark results include strong divergence detection and correct identification of the best schedule family in all tested groups.

A.5 Divergence Prediction Benchmark

  • F1 = 0.92 and balanced accuracy = 0.94 for the R criterion, exceeding the three reported divergence-detection baselines.The peak-ηmax threshold and logistic-regression baselines have F1 = 0.38 and F1 = 0.54, respectively; the grid has 234 convergent and 26 divergent configurations.Source: §A.5 · pp. 39–40, Table 12 · p. 41

A.6 Schedule-Family Generalization

  • The fitted law identifies the best schedule family in all five groups, with prediction errors below 2%.The tested cooldown shapes include cosine and constant-plus-linear schedules outside the piecewise-linear fitting family.Source: §A.6 · pp. 40–42

Appendix B. Theory Supplement

  • The theory supplement collects proofs and derivations for the theoretical results stated in the main text.Its listed material includes SDE derivations and convergence and escape analyses.Source: Appendix B · p. 42

B.1.1 Proof of Proposition 1

  • As model size S grows, the Opt-Laws predictions for cosine and constant schedules both approach b, so their difference tends to zero..Source: §B.1.1 · p. 42

B.1.2 Derivations for SGD-SDE and Adam-SDE

  • The Euler–Maruyama scheme recovers the SGD sequence under the stated time-step setting, while the Adam derivation recovers its momentum update and aligns its parameter and variance updates with their corresponding equations.The SGD condition is Δtₖ₊₁ ≡ η₀.Source: §B.1.2 · p. 43

B.1.3 Proof of Proposition 2

  • The Proposition 2 proof uses Gaussian gradient-noise vectors and Bernstein’s inequality to establish a concentration bound.The bound uses the operator norm of the gradient-noise covariance matrix.Source: §B.1.3 · p. 44

B.1.4 Proof of Proposition 3

  • The Proposition 3 proof standardizes gradient deviations and constructs an empirical covariance estimate from the resulting variables.The derivation invokes eigenvalue variance bounds for covariance matrices.Source: §B.1.4 · p. 44

B.2.1 Proof of Theorem 1

  • The Theorem 1 proof applies Itô’s formula to the SGD-SDE, integrates and takes expectations, then rearranges terms to obtain its bound..Source: §B.2.1 · p. 45

B.2.2 Proof of Proposition 4

These proofs establish boundedness and convergence-related bounds for Adam-SDE quantities, then develop covariance and Gaussian anti-concentration tools for the escape-probability results. Theorem 4’s proof follows the approach used for Theorem 3.

B.2.2 Proof of Proposition 4

  • The Proposition 4 proof establishes boundedness of the expected squared Adam momentum and of the variance variable over time.It first bounds the momentum mean and then derives a solution for the variance variable.Source: §B.2.2 · p. 46

B.2.3 Proof of Theorem 2

  • The Theorem 2 proof applies Itô’s formula to Lyapunov functions, including one that links the noiseless gradient with momentum.The proof then integrates and takes expectations to derive bounds.Source: §B.2.3 · pp. 47–48

B.3.1 Proof of Proposition 5

  • For Proposition 5, vectorizing the covariance matrix converts its ODE into a vector-valued ODE using Q = I ⊗ H + H ⊗ I.The proof considers the trajectory at a fixed point x* with Hessian H.Source: §B.3.1 · p. 49

B.3.2 Proof of Proposition 6

  • Proposition 6 uses a linearized Fokker–Planck approximation to derive differential equations for the state mean and covariance.The proof substitutes the specified Adam initial state into the general Gaussian approximation.Source: §B.3.2 · pp. 49–50

B.3.3 Anti-concentration for Gaussian

  • The escape-probability proofs invoke a general anti-concentration inequality for Gaussian variables.The lemma assumes x is Gaussian with mean μ and covariance Σ.Source: §B.3.3 · p. 50

B.3.4 Proof of Theorem 3

  • Theorem 3’s proof diagonalizes the Hessian and combines a covariance-concentration event with Gaussian anti-concentration to derive its bound.The proof defines an event bounding the trace of the noise covariance and an event that the terminal iterate is near x*.Source: §B.3.4 · pp. 51–52

B.3.5 Proof of Theorem 4

  • The Theorem 4 proof follows the proof of Theorem 3, with details omitted..Source: §B.3.5 · p. 52

Appendix C. Limitations and Scope of Validity

Opt-Laws is formulated for time-varying learning-rate schedules that start and end at zero. Its escape analysis and empirical validation have corresponding scope boundaries.

Appendix C. Limitations and Scope of Validity

  • Fixed learning-rate schedules have zero derivative terms, so Opt-Laws retains only convergence and scale terms for them.The framework is most informative for schedules with nontrivial warmup and cooldown phases.Source: Appendix C · p. 52
  • The SDE escape analysis assumes the learning rate eventually decays to zero; a nonzero terminal learning rate requires extending the escape bounds..Source: Appendix C · p. 52
  • Empirical validation covers piecewise-linear, cosine, constant-plus-linear, and polygon schedules, but not cyclic schedules or warmup restarts.The authors identify those untested schedule types as future work.Source: Appendix C · p. 52
Loading 2409.04777v4…