언제든지 가능한 사전학습: 학습률 무한 스케줄과 가중치 평균의 힘

초록

본 논문은 전체 학습 기간을 사전에 알 수 없는 연속형 사전학습 환경을 위해, 학습률 스케줄을 사전 지식 없이도 최적에 가깝게 유지할 수 있는 “anytime” 방식을 제안한다. 이론적으로 과다 파라미터화된 선형 회귀에서 다항식 형태(ηₜ∝1/t^γ)의 학습률과 가중치 평균(EMA 또는 Tail‑Averaging)이 최소극대(minimax) 수렴 속도를 달성함을 증명하고, 실제 대규모 언어 모델(150M·300M 파라미터) 실험에서 1/√t 스케줄·상수 스케줄·Warm‑up‑Stable‑Decay(WSD)와 가중치 평균을 결합한 방법이 기존 코사인 디케이와 거의 동일한 최종 손실을 보이며, 학습 단계마다 “코사인 엔벨로프”를 따라가는 것을 확인하였다.

상세 분석

이 논문은 두 가지 핵심 질문에 답한다. 첫째, 학습 단계 수(T)를 미리 알 수 없을 때도 최적에 가까운 학습률 스케줄을 설계할 수 있는가? 둘째, 이러한 스케줄이 실제 대규모 언어 모델 사전학습에 적용될 때, 기존에 널리 쓰이던 코사인 디케이와 비교해 손실 및 효율성 면에서 경쟁력을 유지할 수 있는가?

이론적 분석에서는 과다 파라미터화된 선형 회귀 문제를 모델링한다. 데이터 공분산의 스펙트럼이 파워‑law 형태(λ_i∝i^{-α})를 따르고, 소스와 용량 지수(β,γ)가 주어졌을 때, SGD의 평균화된 iterate(Polyak‑Judd 평균)와 학습률 η_t=1/t^γ (0<γ<1) 조합이 최소극대 위험(minimax risk)와 동일한 수렴 속도 O(t^{-γ})를 달성함을 정리 1(비공식)으로 제시한다. 여기서 γ는 스펙트럼의 기하학적 특성에 의해 결정되며, γ=½(즉 1/√t) 가 실험적으로 가장 안정적인 선택으로 나타난다.

또한, 상수 학습률(η_t=η₀)과 가중치 평균을 결합한 경우에도, 평균화 윈도우 길이 N을 적절히 조정하면 동일한 최적 속도를 얻을 수 있음을 보인다. 그러나 기존 연구(예: Zhang et al., 2024a)와 달리, 이 논문은 “anytime”이라는 제약 하에서 학습률이 사전 지식 없이도 정의될 수 있음을 강조한다. 즉, 학습률이 학습 종료 시점에 의존하지 않으며, 모든 중간 단계 T에서도 코사인 스케줄을 별도로 튜닝한 경우와 거의 동일한 위험을 보인다.

실험 부분에서는 OLMo 기반의 트랜스포머 모델(150M, 300M 파라미터)을 Chinchilla 기준 1×~32×(150M) 및 1×~16×(300M) 토큰 규모로 사전학습한다. 코사인 디케이 베이스라인은 각 규모마다 별도 튜닝된 스케줄을 사용하고, 제안된 anytime 스케줄은 가장 큰 규모(32× 또는 16×)에서 한 번만 학습한 뒤 중간 체크포인트에서 평가한다. 결과는 다음과 같다.

1. 손실 곡선: 모든 체크포인트에서 constant + averaging, 1/√t + averaging, WSD + averaging이 코사인 엔벨로프와 겹치며, 특히 중간·후반 구간에서 손실 차이가 0.1% 이하로 미미했다.
2. 학습 효율성: 별도 코사인 튜닝이 필요 없으므로 하이퍼파라미터 탐색 비용이 크게 절감된다. 또한 EMA를 여러 β값으로 동시에 유지해도 메모리 오버헤드가 파라미터 복사 한 번 수준에 불과했다.
3. 스케줄 안정성: 1/√t 스케줄은 초기 단계에서 약간의 과학습(over‑fit) 현상을 보였지만, 평균화가 이를 완화시켰다. 상수 학습률은 초기 손실이 다소 높았지만, 장기 학습에서는 거의 동일한 수렴을 보였다.

이러한 실험 결과는 “anytime” 스케줄이 실제 LLM 사전학습에 적용 가능함을 강력히 뒷받침한다. 특히, 지속적인 데이터 스트리밍이나 멀티‑태스크 연속 학습 상황에서 학습 종료 시점을 사전에 정의할 수 없는 경우, 제안된 스케줄은 코사인 디케이와 동일한 성능을 유지하면서도 구현 복잡도와 튜닝 비용을 크게 낮춘다.

마지막으로 논문은 향후 연구 방향으로 (i) 비선형 모델(예: Transformer)에서의 이론적 수렴 분석, (ii) 다양한 EMA 윈도우 전략과 메타‑학습을 결합한 자동 스케줄링, (iii) 실제 다운스트림 태스크(예: 질문‑응답, 번역)에서의 일반화 효과 검증 등을 제시한다.