기술 요약: ScheduleFree+ 대규모 언어 모델로 스케일링된 학습률 프리 및 스케줄 프리 학습

1. 문제 제기

전통적인 머신러닝 최적화는 수렴을 보장하기 위해 학습률 스케줄 (예: 선형 감쇠, 코사인, 워머업 - 안정 - 감쇠) 에 크게 의존합니다. 이러한 방법들은 효과적이지만 상당한 튜닝이 필요하며, 종종 불안정하거나 예측 불가능한 학습 곡선을 초래합니다. 최근 "스케줄 프리" 학습 방법들은 스케줄을 공격적인 반복체 평균화 (iterate averaging) 로 대체하여 소규모 벤치마크에서 가능성을 보였으나, 역사적으로 대규모 언어 모델 (LLM) 학습으로 확장하는 데 실패했습니다. 구체적으로 이전 Schedule-Free 구현체들은 다음과 같은 문제점을 겪었습니다:

• 대규모 배치 크기에서의 발산: 배치 크기가 소규모 한계를 초과하여 증가함에 따라 성능이 크게 저하되었습니다.
• 기울기 노름 (Gradient norm) 드리프트: 이 방법은 기능을 수행하기 위해 비정상적으로 작은 가중치 감쇠 (weight decay) 값을 필요로 하여, 불안정한 가중치 노름과 장기 학습에서의 낮은 성능을 초래했습니다.
• 장기 실행에서의 비최적 수렴: 이 방법은 긴 학습 기간 (예: 파라미터당 1000 토큰) 에 걸친 평균화의 이점을 완전히 활용하지 못했습니다.

본 논문은 LLM 사전 학습을 위한 최첨단 (SOTA) 스케줄 기반 방법의 성능과 일치하거나 초과하는 스케일링 가능한 Schedule-Free 학습의 과제를 해결하여, 학습률 튜닝이 필요 없는 안정적이고 "언제나 사용 가능한 (anytime)" 학습 절차를 목표로 합니다.

2. 방법론: ScheduleFree+

저자들은 LLM 에 대해 견고하게 만들기 위해 기본 Schedule-Free 알고리즘에 몇 가지 특정 수정을 통합한 복합 방법인 **ScheduleFree+**를 제안합니다. 핵심 구성 요소는 다음과 같습니다:

A. 내부 모멘텀의 재도입

원래 Schedule-Free 작업은 메모리를 절약하기 위해 내부 모멘텀 (예: Adam 의 $\beta_1$) 을 제거했으나, 저자들은 이것이 대규모 배치 크기에서 해롭다는 것을 발견했습니다.

• 해결책: 기본 옵티마이저에 내부 모멘텀 (예: $\beta_1 = 0.75$ 또는 $0.9$) 을 재도입합니다.
• 효과: 이는 더 큰 배치 크기에 대한 견고성을 복원하고 발산 없이 더 높은 학습률을 허용하여, 표준 Schedule-Free AdamW 에서 관찰된 "스케일링 절벽 (scaling cliff)"을 해결합니다.

B. 역기울기 노름 가중치 (Polyak 스텝 크기)

표준 Schedule-Free 학습은 종종 증가하는 기울기 노름을 보상하기 위해 작은 가중치 감쇠를 필요로 하여 불안정성을 초래했습니다. 저자들은 Polyak 스텝 크기에 기반한 이론적 보정을 유도했습니다.

• 메커니즘: 스텝 크기 $\gamma_t$는 기울기의 L1 노름에 반비례하도록 스케일링됩니다 ($\gamma_t \propto 1/\|g_t\|_1$).
• 구현: 기울기 L1 노름의 지수 이동 평균 (EMA) 을 사용하여 분모를 추정하며, 이는 볼록 후회 (convex regret) 경계에서 유도된 최적 가중치를 근사합니다.
• 결과: 이 보정은 정규화 계층으로 인한 유효 학습률 드리프트를 상쇄하여, 가중치 노름을 안정화하고 장기 실행을 불안정하게 만드는 "축소되는 가중치 노름" 현상을 방지하는 표준적이고 더 큰 가중치 감쇠 값 (예: 0.05–0.5) 사용을 가능하게 합니다.

C. Polyak 스텝을 통한 학습률 프리 최적화

역기울기 가중치에 기반하여, 저자들은 수정된 Polyak 스텝 크기를 사용하여 완전히 학습률이 없는 변형을 구현합니다:
$$\gamma_t = \frac{f(y_t) - f^* + \beta \langle \nabla f(y_t), z_t - x_t \rangle}{\sqrt{\pi/2} \| \nabla f(y_t) \|_1}$$

• 안정성: 분자는 노이즈가 있는 현재 반복체 대신 함수 값 간격 ($f(z_t) - f^*$) 의 추정을 사용하며, 분자는 평활화된 L1 노름을 사용합니다. 이는 확률적 Polyak 스텝에서 흔히 발생하는 소멸하는 분모 문제를 피합니다.
• 적응성: 이 접근 방식은 배치 크기에 따라 자동으로 학습률을 스케일링 (AdamW 의 경우 $\sqrt{b}$ 스케일링 따름) 하고, 수동 학습률 그리드 검색의 필요성을 제거합니다.

D. 워밍업 및 어닐링 전략

• C-워밍업: 평균화 버퍼가 초기 급속 수렴 단계를 희석하는 것을 방지하기 위해, 저자들은 초기 기간 동안 (예: 학습률 워밍업 기간의 $2\times$) 평균화 업데이트 ($c_t = 1$) 를 지연시킵니다.
• 베타 어닐링: 현재 반복체 $z_t$와 평균 $x_t$를 혼합하는 보간 매개변수 $\beta$는 학습 실행 동안 낮은 값 (예: 0.9) 에서 높은 값 (예: 0.965) 으로 어닐링됩니다. 이는 빠른 초기 수렴을 포착하면서 장기 손실 감소를 최적화합니다.
• 가중 평균 ($r=1$): 장기 실행의 경우, 평균화 가중치는 기본값인 $r=0$ 대신 $r=1$(선형적으로 증가하는 가중치) 로 설정되어 비강볼록 문제의 수렴을 가속화합니다.

3. 주요 결과

저자들은 1 억 2 천만에서 10 억 파라미터까지의 LLM 에 대해 최대 파라미터당 1000 토큰으로 훈련된 광범위한 스케일링 실험을 수행했습니다.

• 장기 학습 (1000 TPP): 이 규모에서 ScheduleFree+ 는 워머업 - 안정 - 감쇠 (WSD) 및 선형 감쇠 스케줄 모두를 크게 능가합니다.
  • 효율성: 최종 손실은 31% 더 짧은 학습 시간 (목표 손실에 도달하는 데 필요한 토큰 수 45% 감소) 으로 선형 감쇠 스케줄과 동일하게 달성합니다.
  • 안정성: 이 방법은 이론적 $1/\sqrt{t}$ 감쇠율을 따르는 매끄럽고 예측 가능한 손실 곡선을 생성하여, 초기 학습 데이터로부터 최종 손실을 정확하게 예측할 수 있게 합니다.
• 단기 학습 (20–100 TPP): ScheduleFree+ 는 모든 모델 크기에서 WSD 스케줄을 능가합니다. 선형 감쇠에 비해 매우 짧은 기간 (20 TPP) 에는 정상 상태에 도달하는 데 필요한 시간으로 인해 비슷하거나 약간 뒤처지지만, 전반적으로 경쟁력 있는 성능을 보입니다.
• 배치 크기 스케일링: 내부 모멘텀을 통해 ScheduleFree+ 는 8M 배치 크기까지 견고하게 확장되어, 선형 감쇠를 사용하는 표준 AdamW 와 동등한 성능을 냅니다.
• 가중치 노름 안정성: Polyak 스텝 크기와 더 큰 가중치 감쇠를 사용하여, 이 방법은 수정되지 않은 Schedule-Free 방법이 요구하는 작은 가중치 감쇠 값과 관련된 발산 및 양자화 문제를 피하면서 안정적인 가중치 노름을 유지합니다.

4. 중요성 및 주장

본 논문은 ScheduleFree+ 가 사전 학습 중 모델 평균화 및 체크포인트 병합을 위한 이론적 기초를 제공하여, 온라인 최적화와 오프라인 모델 수프 (model souping) 간의 격차를 해소한다고 주장합니다.

• 언제나 사용 가능한 학습 (Anytime Training): 학습 말미의 특정 어닐링 단계를 필요로 하여 이득을 실현하는 WSD 스케줄과 달리, ScheduleFree+ 는 진정한 "언제나 사용 가능한" 기능을 제공합니다. 모델은 전용 감쇠 단계 없이도 중단 시점에 언제든지 사용 가능하고 최적입니다.
• 튜닝 제거: 이 방법은 학습률 그리드 검색의 필요성을 제거하고, 대신 Polyak 스텝 크기와 역기울기 가중치에 의존하여 배치 크기와 학습 기간에 자동으로 적응합니다.
• 예측 가능성: 저자들은 Schedule-Free 학습의 손실 곡선이 경험적으로 안정적일 뿐만 아니라 이론적 근거가 있으며, 비볼록 LLM 설정에서도 볼록 최적화 이론이 예측하는 $1/\sqrt{t}$ 수렴율을 엄격하게 따른다고 주장합니다.

요약하자면, 본 논문은 내부 모멘텀, 역기울기 가중치, 적응형 어닐링과 같은 특정 수정을 통해 Schedule-Free 학습이 대규모 언어 모델 훈련을 위한 전통적인 학습률 스케줄에 비해 우월하고, 안정적이며, 튜닝이 필요 없는 대안으로 확장될 수 있음을 보여줍니다.