[논문 리뷰] Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers

Artem Riabinin, Andrey Veprikov외 3명|arXiv (Cornell University)|2026. 02. 05.

Stochastic Gradient Optimization Techniques인용 수 0

한 줄 요약

이 논문은 규범-제약 최적화기에 대한 이론적 프레임워크와 실용적 적응형 워밍업 스케줄러를 제공합니다. 이는 일반화된 매끄러움 가정을 도입해 로컬 곡률이 비최적도와 함께 증가함을 보이고, 워밍업이 자연스럽게 발생함을 증명하며, 적응형 워밍업이 추가 하이퍼파라미터 탐색 없이 LLM 사전학습에서 성능을 향상시킵니다.

ABSTRACT

We study adaptive learning rate scheduling for norm-constrained optimizers (e.g., Muon and Lion). We introduce a generalized smoothness assumption under which local curvature decreases with the suboptimality gap and empirically verify that this behavior holds along optimization trajectories. Under this assumption, we establish convergence guarantees under an appropriate choice of learning rate, for which warm-up followed by decay arises naturally from the proof rather than being imposed heuristically. Building on this theory, we develop a practical learning rate scheduler that relies only on standard hyperparameters and adapts the warm-up duration automatically at the beginning of training. We evaluate this method on large language model pretraining with LLaMA architectures and show that our adaptive warm-up selection consistently outperforms or at least matches the best manually tuned warm-up schedules across all considered setups, without additional hyperparameter search. Our source code is available at https://github.com/brain-lab-research/llm-baselines/tree/warmup

연구 동기 및 목표

• 경험적 휴리스틱을 넘어서 규범-제약 최적화기에 대한 워밍업의 동기를 제시하고 정당화한다.
• 곡률과 부분 최적성 사이의 연결을 갖는 일반화된 매끄러움 모델을 도입한다.
• 이 모델 하에서 워밍업-후 감소(warm-up-then-decay) 학습률로 수렴을 증명한다.
• 표준 하이퍼파라미터에 의존하는 실용적인 적응형 워밍업 스케줄러를 개발한다.
• 대형 언어 모델 사전학습에서 스케줄러를 검증하고 경쟁력 있는 성능을 보인다.

제안 방법

• LMO 기반 업데이트 x^{t+1}=x^{t}+ eta^{t} LMO(g^{t})를 정의하고 이를 이차 손실 근사와 관련짓는다.
• Assumption 2: (\rho, K0, K1, K\rho)-매끄러움으로 제시하고 곡률 한계가 f(x)-f^{*}에 의존하도록 한다.
• 학습률 \neta^{t}= (\filled{Δ^{t}})/(D·K(x^{t}))이면 부분 최적성 Δ^{t}가 감소하고 K(x^{t})도 감소한다(정리 1).
• 가중치 감소를 확장: 업데이트 x^{t+1}=(1−\lambda η^{t})x^{t}+η^{t} LMO(g^{t})와 수렴성 증명(정리 2).
• 확률적 설정으로 확장: 보정된 그래디언트 정규화와 보간 가정(정리 3).
• 제약 조건 하에서 η(Δ)=Δ/(K0+K1Δ+K2Δ^2)의 3-매개변수 모델을 피팅하여 실용적 적응형 워밍업 스케줄러를 도출하고 Δ′ 이후에 감소로 전환한다(Algorithm 5).

실험 결과

연구 질문

• RQ1 LMO 기반 최적화기의 학습률 워밍업이 순전히 휴리스틱이 아니라 이론적으로 정당화될 수 있는가?
• RQ2 학습 시작 시 수동 조정 없이 워밍업 기간을 자동으로 적응시킬 수 있는가?
• RQ3 부분 최적성 의존 매끄러움 모델이 워밍업과 감소를 최적화의 자연스러운 동역학으로 설명하는가?
• RQ4 적응형 워밍업 스케줄러가 하이퍼파라미터 탐색 없이 대규모 LLM 사전학습에서 잘 작동하는가?

주요 결과

• 일반화된 매끄러움 모델이 최적화 경로에서 부분 최적성 간격에 대한 곡률에 의존함을 실증적으로 지지한다.
• 모델 하에서 워밍업에 이어 감소가 LMO 기반 최적화기의 수렴 증명에서 자연스럽게 등장한다.
• 표준 하이퍼파라미터를 사용하고 training 시작 시 K0, K1, K2 및 Δ′를 추정하는 실용적 적응형 워밍업 스케줄러를 도출한다.
• 적응형 워밍업은 Muon, Lion, 및 normSGD에서 LLaMA 사전학습에 대해 추가 하이퍼파라미터 탐색 없이 최상의 수동으로 조정된 워밍업 스케줄과 같거나 더 우수한 성능을 보인다.
• 모델 크기와 배치 규모에 관계없이 접근 방식이 견고하며, 특히 작은 배치 설정에서 이점이 크다.