지평선 없는
학습률 스케줄

대규모 언어모델은 점점 더 연속적이고 개방된 환경에서 학습된다. 전체 학습 길이를 미리 알 수 없는 상황이 일반적이 되었다.

그럼에도 대부분의 사전학습 레시피는 anytime이 아니다. Loshchilov와 Hutter가 2016년에 제안한 이래 사실상의 표준이 된 코사인 감쇠는 학습 종료 시점을 미리 알아야 스케줄을 정할 수 있다. 학습 길이가 정해지면 손실 곡선의 전체 궤적이 그 시점에 의해 암묵적으로 결정된다. 데이터가 계속 유입되고 섞이는 연속 학습 환경에는 근본적으로 맞지 않는다.

이 논문은 스케줄러에 두 가지 요건을 부과한다.

1. 스케줄은 계획된 학습 스텝 수에 의존하지 않아야 한다.
2. 임의의 중간 지속시간 T에 대해, T 스텝으로 잘 조율된 코사인 스케줄과 대등해야 한다. 즉 긴 학습 안의 각 체크포인트에서 조율된 코사인들이 그리는 코사인 포락선(envelope)을 따라가야 한다.

이 포락선 관점이 핵심이다. 하나의 종료 지평선에 맞춰 조율한 코사인 스케줄은 중간 체크포인트에서 평가하면 최적에서 크게 벗어난다. 저자들이 아는 한, 이 포락선을 명시적 평가 기준으로 삼은 선행 연구는 없다.

아래는 각 스케줄이 시간에 따라 학습률을 어떻게 정하는지 직접 그린 것이다. 코사인은 종료 시점을 알아야 곡선이 닫히지만, 상수와 1/√t는 어느 시점에서 멈추든 같은 궤적을 유지한다.

WSD: 거의-anytime 대안

Warmup-Stable-Decay(WSD)는 선형 워밍업, 일정 구간 상수 유지, 마지막 10%에서 학습률을 원래의 10%까지 선형 감쇠하는 3단계로 구성된다. 엄밀히 horizon-free는 아니다. 감쇠를 언제 시작할지 결정하려면 체크포인트 저장과 판단이 필요하기 때문이다. 다만 그 오버헤드가 작고 성능은 코사인에 필적한다.

가중치 평균: 진짜 열쇠

학습률 감쇠의 대안은 가중치 평균(모델 병합)이다. 최근 반복들의 평균을 유지하거나 EMA로 근사하며, 학습에는 마지막 반복만 쓰되 평가에는 평균 모델을 쓴다. 파라미터 β ≈ 1 − 1/N로 제어한다. 실험에서는 EMA의 반감기가 현재 시점의 일정 비율 f가 되도록 τ_t = 1/2^f/t를 쓰고, f ∈ {0, 6.25, 12.5, 25, 50, 100}의 여러 EMA를 동시에 유지한다. 1/t^γ 스케줄이 SGD의 미니맥스 최적률에 도달하려면 이 평균이 반드시 필요하다.

저자들은 과매개변수 선형 회귀에서 anytime 스케줄이 실제로 존재함을 증명한다. 가우시안 데이터, 배치 크기 1, 독립 가법 잡음 σ²를 가정한다.

이 결과가 상수 학습률과 대비되는 지점이 핵심이다. Zhang 등(2024a)은 상수 학습률 + 가중치 평균도 같은 최적률에 도달할 수 있음을 보였으나, 특정 source·capacity 지수에서는 학습률을 학습 지평선의 함수로 1/√N처럼 조정해야 한다. 그러면 종료 시점에 의존하므로 더 이상 anytime이 아니다. 반면 1/t^γ는 그 의존성 없이 최적률을 얻는다.

최적 지수 γ*는 스펙트럼이 정한다

Corollary 1은 데이터 공분산이 멱법칙을 따를 때를 다룬다. capacity 지수 a와 source 지수 b를 λ_i ≂ i^−a, 𝔼λ_i(w*_i)² ≂ i^−b로 정의하면 최적 지수는 다음과 같다.

여기서 b는 신호가 공분산 H의 상위 고유방향에 얼마나 집중되는지를 나타낸다. 이 최적 γ가 실현되려면 a < b가 필요하다.

WSD도 이론에 편입된다

Theorem 2는 상수 구간 + 선형 감쇠로 이뤄진 WSD를 멱법칙 스펙트럼(a ∈ (1,2), b > 1)에서 분석한다. b = a일 때 두 항 모두 지수 1 − 1/b로 감쇠하며, 이는 Corollary 1의 결과와 일치한다. 결국 WSD는 상수 학습률 + 가중치 평균과 점근적으로 같은 수렴률을 가지며, 두 스케줄을 교체해도 무방하되 상수 쪽이 anytime이라는 이점을 얻는다.

1차원 평균 추정 예시가 이 등가성을 명료하게 드러낸다. 상수 학습률 + 반복 평균과, 평균 없이 감쇠하는 학습률은 같은 암묵적 표본 가중을 서로 다른 방식으로 구현할 뿐 수학적으로 동일하다.

OLMo 코드베이스 기반의 150M·300M 모델을 C4 데이터셋으로 학습한다. AdamW, 시퀀스 길이 1024, 데이터 반복 없는 완전 온라인 설정이다. Chinchilla 규모(20 TPP)를 기준으로 150M은 1×–32×(3.3B~), 300M은 1×–16×(6.6B~)까지 학습한다.

실험 설계의 공정성이 중요하다. 코사인은 각 종료 배수마다 독립적으로 조율해 별도 학습한다. 반면 1/√t와 상수 + 평균은 최대 배수까지 한 번만 학습하고 중간 체크포인트에서 평가한다. WSD는 각 배수의 90% 지점에서 저장한 체크포인트로부터 선형 감쇠해 구현한다.

그 결과, anytime 스케줄들은 전 구간에서 코사인 감쇠에 근접하게 일치한다. 학습 초반과 종료 부근에서만 미미한 성능 손실이 있을 뿐이다.

임계 배치 크기(CBS) 설정

150M은 배치 256, 300M은 √N 스케일링으로 512를 임계 배치로 근사한다. 이 설정에서 1/√t는 스케줄러를 η√(α/(t+α))로 매개변수화하고 α를 조율한다. α가 총 스텝에 약하게 의존하지만, 1×–32× 포락선 전체에 걸쳐 거의 최적인 단일 α를 찾을 수 있다.

대형 배치 설정 (배치 4096)

임계 배치를 크게 넘는 배치 4096에서는 다른 양상이 나타난다. 큰 배치는 경사 잡음을 억제하므로 분산 항을 제어할 학습률 감쇠의 필요가 줄어든다. 실제로 상수 학습률 + 평균이 1× 이후 모든 지평선에서 코사인을 확연히 능가한다. 긴 학습에 거의 최적인 학습률이 학습 내내 거의 최적으로 유지된다. 1/√t는 여전히 경쟁력이 있고 코사인을 개선하지만, 이 영역에서는 불필요한 감쇠 탓에 상수보다는 뒤진다.

합성 선형 회귀 검증

차원 d = 5×10⁵, 최대 N = 5×10⁴ 표본의 정확한 위험 점화식으로 이론을 뒷받침한다. source 지수 a ∈ {1.1, 1.5, 1.9}, capacity b = a 및 b = 2a 조합에서 1/√t가 다양한 스펙트럼에 걸쳐 anytime을 유지하며 거의 최적 성능을 낸다. √(α/(t+α))의 α를 조율하면 α* = Θ(T)일 때 상수 학습률을 모사할 수 있어, b = a 영역의 최고 성능을 맞추고 b = 2a 영역에서는 개선된 수렴률을 얻는다.

이 연구는 지평선을 모르고도 잘 조율된 코사인에 필적하는 학습이 가능함을 이론과 실험 양쪽에서 보인다.

• 과매개변수 선형 회귀에서 anytime 학습 스케줄의 존재를 증명한다. 1/t^γ는 임의의 스펙트럼 지수에 대해 로그·상수 인자를 제외하고 최적률을 달성하며, 1/√t가 그 대표 사례다.
• 미니맥스 최적률 달성에서 가중치 평균이 차지하는 중심적 역할을 규명한다. 평균 없이는 다항 감쇠 스케줄이 일반적으로 최적률에 이르지 못한다.
• 150M·300M 모델을 1×–32× Chinchilla 규모로 학습해, 상수 + 평균과 1/√t + 평균이 지평선 지식 없이도 코사인 포락선을 근사하게 추종함을 확인한다.
• WSD를 이론에 편입해, 특정 멱법칙 스펙트럼에서 상수 학습률 + 가중치 평균과 점근적으로 등가임을 보인다.
• 대형 배치 영역에서는 상수 + 평균이 코사인을 능가하여, 학습률 감쇠가 불필요해지는 조건을 드러낸다.

종합하면, 가중치 평균과 단순하고 horizon-free한 스텝 크기의 결합은 대규모 언어모델 사전학습에서 코사인 학습률 스케줄을 대체할 실용적이고 효과적인 anytime 대안을 제공한다. anytime 사전학습은 연속 학습으로 가는 필수적인 디딤돌이다.