AdaLRS: Loss-Guided Adaptive Learning Rate Search for Efficient Foundation Model Pretraining¶

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

AdaLRS is a plug-in-and-play online learning rate search algorithm that monitors training loss descent velocity to adaptively adjust the learning rate toward its optimum within a single pretraining run. It exploits the empirically and theoretically demonstrated convexity of pretraining loss and loss slope w.r.t. learning rate, and their shared optimum, to replace costly multi-run hyperparameter search. Applied to both LLM and VLM pretraining, AdaLRS corrects suboptimal LRs efficiently and improves downstream model performance.

Problem¶

Existing methods for finding optimal learning rates in foundation model pretraining either (a) derive scaling laws that lack generalizability across model structures and data compositions, or (b) transfer hyperparameters from proxy models via µTransfer-style methods, which still require expensive multi-run search on the proxy. No method can perform optimal LR identification in a single training run without prior assumptions about model size or data.

Method¶

AdaLRS maintains a sliding window estimate of the training loss curve slope (loss descent velocity) using least-squares regression. At each adjustment point it tentatively upscales the current LR by factor α′ = max(λᵗα, 1) and runs k trial steps. It keeps the upscaling if the new slope exceeds the old by more than 2e (the estimation error bound); otherwise it reverts model state (backtracking) and downscales by β′ = max(1/(λᵗβ), 1). The scaling factors α and β are multiplicatively independent reals and decay via λ ∈ (0,1), ensuring the search neighborhood contracts over time. LR adjustment is restricted to an early fraction of training (e.g., steps [10%, 35–40%] of total), after which the base scheduler runs unmodified. The algorithm is scheduler-agnostic (demonstrated with cosine and WSD) and adds negligible compute beyond the trial k-step windows.

Theoretical grounding: under SGD with Lipschitz-gradient assumptions, expected loss descent per step is −η‖∇L‖² + (C_Lip/2)η²‖∇L‖², giving a convex optimum at η* = 1/C_Lip. This extends heuristically to show that loss and loss-slope share the same optimal LR, verified empirically across 2B/7B LLM and 2B VLM experiments spanning LRs in [2e-6, 8e-3] (LLM) and [2e-5, 0.2] (VLM).

Convergence: Theorem 2.1 proves {ηt} converges almost surely to the e-neighborhood of η; Theorem 2.4 proves geometric error decay |ηt+k − η| ≤ γ|ηt − η*|, giving O(log R) adjustments for search range R.

Key Contributions¶

Theoretical and empirical demonstration that foundation model pretraining loss and its slope are jointly convex in LR and share the same optimum.
AdaLRS algorithm: single-run, online LR search via loss descent velocity optimization with convergence proof and geometric error-decay bound.
Backtracking downscaling strategy that restores model state before failed upscale attempts, preventing parameter disruption.
Validation across 2B/7B LLMs and 2B VLM with cosine and WSD schedulers, multiple model sizes, continual pretraining, and varied α/β/λ hyperparameters.

Results¶

LLM pretraining (Qwen2.5-1.5B, 7B; ~120B tokens, SlimPajama): - Small LR (2e-5): AdaLRS reduces 2B validation PPL from 32.69 → 12.66 (baseline vs. AdaLRS); 7B from 15.30 → 10.61. - Large LR (2e-3): 2B validation PPL 204.97 → 183.94; 7B 165.70 → 158.22. - Fit LR: 2B validation PPL 12.42 → 13.51 (marginal degradation), 7B 10.72 → 10.67 (marginal gain). - AdaLRS from small LR (7B) surpasses Fit LR baseline on Alpaca-Gen (21.15 vs. 20.68) and KNIGHT-Gen (13.53 vs. 12.29). - Convergence acceleration: AdaLRS surpasses baseline final loss at ~50% of steps for small LR; reduces >30% training cost for large LR.

VLM pretraining (2B, InternViT-300M + Qwen2.5-1.5B; ~120B tokens): - Small LR: average benchmark score 57.34 (AdaLRS) vs. 53.77 (baseline); TextVQA 64.85 vs. 57.86. - Large LR: average 48.96 (AdaLRS) vs. 47.67 (baseline); TextVQA 60.32 vs. 58.08. - Fit LR: average 56.16 vs. 55.80 (AdaLRS comparable or slightly better).

Continual VLM pretraining: AdaLRS reduces training loss from 0.8851 → 0.8286 under small LR; large LR remains problematic (catastrophic forgetting not resolved).

Limitations¶

Large initial LR disrupts model parameters before AdaLRS can correct, leaving residually elevated training loss even after adjustment.
Aggressive upscaling/downscaling factors can introduce oscillatory LR behavior; requires tuning of α, β, λ.
LR search is restricted to early training steps; if the window is too short, AdaLRS may not converge to the true optimum.
In continual pretraining with large LR, AdaLRS downscaling fails to recover from catastrophic forgetting.
Hypothesis 1 (joint convexity and shared optimum) is verified empirically but rests on a simplified SGD/constant-LR theoretical derivation; behavior under modern optimizers (AdamW) with complex schedules is assumed rather than proven.
Experiments use proprietary Huawei 910B NPU infrastructure; reproducibility on standard GPU clusters is not directly confirmed.

Relevance to Vision-Language Models¶

AdaLRS directly addresses the learning rate sensitivity problem in VLM pretraining, where connector/projector training from scratch and continual fine-tuning of vision encoders are notoriously sensitive to LR choices. By enabling reliable single-run LR correction, it reduces the iteration overhead in VLM pretraining pipelines that typically require costly pilot studies (as practiced in SAIL-VL, InternVL, Qwen-VL). The demonstrated gains on OCR and visual understanding benchmarks (DocVQA, TextVQA, OCRBench) are directly relevant to VLM capability development, and the approach generalizes across training paradigms (from-scratch projector training vs. full-parameter continual pretraining), making it applicable to the dominant two-stage VLM training recipes in the literature.