📊 statistics

Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling

이 논문은 손실 역학(loss dynamics)을 보존하기 위해 배치 크기를 확장하는 동시에 학습률을 조정함으로써 대규모 언어 모델의 사전 학습을 가속화하고, 표준 코사인 감쇠 스케줄 대비 동일한 FLOPs에서 성능을 유지하면서도 실제 학습 시간(wall-clock training time)을 약 36% 단축하는 원칙적인 프레임워크인 Seesaw를 소개한다.

원저자: Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade

게시일 2026-07-28

📖 2 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, Sham Kakade

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대하고 아주 똑똑한 로봇에게 책 한 산더미를 먹여 세상을 이해하도록 가르치려 한다고 상상해 보세요. 이 과정을 "학습(training)"이라고 부르며, 이는 오늘 우리가 보고 있는 AI 챗봇과 도구들의 뒤에 숨겨진 비법입니다. 잘 배우기 위해서 로봇에게는 두 가지 주요 요소가 필요합니다. 하나는 단 한 페이지를 읽은 후 자신의 생각을 얼마나 바꿀지를 결정하는 "학습률(learning rate)"이고, 다른 하나는 생각을 멈추고 조절하기 전까지 몇 페이지를 읽을지를 결정하는 "배치 크기(batch size)"입니다.

오랫동안 과학자들은 한 번에 더 많은 페이지를 제공하는 것(더 큰 배치 크기)이 학습 속도를 높이는 최선의 방법이라고 생각했습니다. 하지만 여기에는 함정이 있습니다. 만약 생각하는 방식을 바꾸지 않은 채 너무 많은 페이지를 한꺼번에 제공하면, 로봇은 혼란에 빠져 효율적으로 학습하는 것을 멈추게 됩니다. 이는 백과사전 한 권을 한자리에서 통째로 읽으며 언어를 배우려는 것과 같습니다. 사실 관계는 알게 될지 모르지만, 문법은 이해하지 못할 것입니다. 연구자들이 던져온 핵심 질문은 이것입니다: 어떻게 하면 로봇이 길을 잃지 않으면서도, 더 많이 읽는 것과 신중하게 생각하여 더 빨리 배우는 것 사이의 균형을 맞출 수 있을까?

이 논문은 이 균형 잡기 문제를 해결하기 위해 **시소(Seesaw)**라고 불리는 영리하고 새로운 전략을 소개합니다. 연구진은 수학적인 규칙이 완벽한 시소처럼 작동한다는 것을 발견했습니다. 즉, 로봇이 한 번에 읽는 페이지 수(배치 크기)를 두 배로 늘릴 때마다, 단순히 학습률을 그대로 유지하거나 절반으로 줄여서는 안 된다는 것입니다. 대신, 학습률을 매우 특정한 양만큼, 즉 2의 제곱근(약 1.41)으로 나누어 조정해야 합니다.

이렇게 생각해 보세요. 만약 당신의 스터디 그룹 규모(배치 크기)를 두 배로 키운다면, 대화 속도(학습률)를 생각보다 그렇게 많이 늦출 필요는 없습니다. 이 "시소" 리듬을 사용함으로써, 로봇은 더 적은 단계만으로도 동일한 양의 정보를 처리할 수 있습니다. 저자들은 단순한 학습 과제에 대해 이를 수학적으로 증명했으며, 이후 1억 5천만, 3억, 6억 개의 파라미터를 가진 거대 언어 모델들을 대상으로 테스트했습니다. 그 결과, 시소 방식을 사용했을 때 동일한 수준의 지능을 유지하면서도 실제 구동 시간(wall-clock time) 기준으로 모델을 약 36% 더 빠르게 학습시킬 수 있다는 것을 발견했습니다.

또한 이 논문은 지나친 탐욕을 경계하라고 명시적으로 경고합니다. 만약 학습률을 올바르게 조정하지 않은 채 배치 크기를 너무 공격적으로 늘리려고 한다면, 로봇의 학습 과정은 불안정해지고 성장을 멈추게 됩니다. 저자들은 수학적 법칙이 무너지는 "임계점(tipping point)"이 존재함을 보여주었으며, 그들의 시소 방식은 안전하게 그 선의 오른쪽에 머물러 있음을 입증했습니다. 요컨대, 시소는 단순한 추측이 아닙니다. 이는 AI 모델이 동일한 교훈을 훨씬 적은 시간에 배울 수 있게 해주는 수학적 근거가 있는 레시피이며, 이를 통해 우리는 학습이 끝나기를 몇 달씩 기다리지 않고도 더 똑똑한 AI를 만드는 데 더 가까워질 수 있습니다.

기술 요약: Seesaw: 학습률과 배치 크기 스케줄링의 균형을 통한 학습 가속화

문제 정의

대규모 언어 모델(LLM)의 사전 학습은 점점 더 벽시계 시간(wall-clock time)에 의해 제약을 받고 있으며, 최첨단 모델의 경우 종종 수개월까지 연장되기도 합니다. 이를 단축하기 위한 주요 전략은 배치 크기를 늘리는 것인데, 이는 병렬 계산을 활용하여 전체 순차적 최적화 단계 수를 줄이는 방식입니다. 그러나 단순히 배치 크기를 무한히 늘리는 것은 샘플 효율성이 떨어지는 "임계 배치 크기(critical batch size, CBS)" 너머에서는 수확 체감의 법칙에 직면하게 됩니다.

최근의 대규모 학습 실행(예: LLaMA, Nemotron, OLMo)에서는 "배치 램프(batch ramp)" 스케줄(학습 과정 중에 배치 크기를 점진적으로 증가시키는 방식)을 채택해 왔으나, 이러한 전략들은 현재 휴리스틱하게 튜닝되고 있습니다. 특히 Adam과 같은 적응형 옵티마이저(adaptive optimizer)를 사용할 때, 학습률 감소(learning rate decay)와 배치 크기 증가(batch size ramp-up)를 어떻게 최적으로 균형 잡아야 하는지에 대한 이론적 근거가 부족합니다. 핵심 질문은 다음과 같습니다: 모델 성능을 희생하지 않으면서 직렬 실행 시간을 최소화하기 위한 최적의 배치 크기 스케줄은 무엇인가?

방법론

이론적 프레임워크

저자들은 학습률 감소와 배치 크기 증가 사이의 등가성을 확립하기 위해 노이즈가 포함된 선형 회귀(noisy linear regression)에 기반한 원칙적인 프레임워크를 개발했습니다.

SGD 등가성: 본 논문은 노이즈가 포함된 선형 회귀 상의 확률적 경사 하강법(SGD)에 대해, 매 단계마다 배치 크기가 두 배로 늘어나고 학습률은 고정되는 프로세스가, 배치 크기는 고정되고 매 단계마다 학습률이 절반으로 줄어드는 프로세스와 (과잉 리스크 측면에서 상수 인자 차이를 제외하면) 동일한 총 데이터 포인트를 소비한다는 것을 보여주는 최초의 유한 샘플(finite-sample, 비점근적) 증명을 제공합니다.
정규화된 SGD(NSGD)로의 확장: Adam과 같은 적응형 옵티마이저와의 간극을 메우기 위해, 저자들은 Adam의 다루기 쉬운 분석적 대리 모델인 정규화된 SGD(NSGD)를 분석합니다. 저자들은 **가정 2(Assumption 2)**를 도입하며, 이는 실제 관찰되는 분산 지배 영역(variance-dominated regime)에서 기대 제곱 그래디언트 노름(expected squared gradient norms)이 가산 노이즈(additive noise, $O(1/B)$ 로 스케일링됨)에 의해 지배된다는 가설을 담고 있습니다. 이 가정하에 NSGD는 학습률이 재조정된 SGD로 효과적으로 축소됩니다.
3.러 등가 조건: 두 파라미터를 조정할 때 손실 역학(loss dynamics)을 유지하기 위해 필요한 특정 관계를 도출합니다. 학습률 감소 계수를 $\alpha$ , 배치 크기 증가 계수를 $\beta$ 라고 할 때, $\alpha\sqrt{\beta}$ 가 일정하게 유지된다면 두 프로세스는 동일한 상태를 유지합니다.

Seesaw 알고리즘

이러한 이론적 등가성에 기반하여, 저자들은 학습률과 배치 크기를 동적으로 균형 잡는 스케줄러인 Seesaw를 제안합니다.

메커니즘: 표준 스케줄러(예: 코사인 감소)가 학습률을 $\alpha$ 비율로 줄여야 할 때, Seesaw는 대신 학습률을 $\sqrt{\alpha}$ 만큼 줄이고 배치 크기를 $\alpha$ 만큼 증가시킵니다.
제약 조건: 저자들은 발산을 방지하기 위한 이론적 한계를 식별했습니다. 안정성을 유지하면서 가장 공격적인 램프업 스케줄은 $\alpha = \sqrt{\beta}$ 로 정의됩니다. 이를 초과하는 경우(즉, 학습률 감소에 비해 배치 크기를 너무 공격적으로 증가시키는 경우) 발산이 발생합니다.
구현: Seesaw는 기존 학습률 스케줄러의 드롭인 교체재(drop-in replacement)로 설계되었습니다. 실제로 Seesaw는 코사인 감소를 스텝 감소(step decay)로 근사하며, 학습률이 삭감되어야 하는 특정 토큰 수 시점에 맞춰 배치 크기 증가를 트리거합니다.

주요 기여

이론적 등가성: 본 논문은 SGD에 대한 학습률 감소와 배치 크기 증가 사이의 최초의 비점근적(non-asymptotic) 등가성을 확립하고, 이를 분산 지배 영역 하의 NSGD로 확장했습니다.
Seesaw 스케줄러: 손실 역학을 보존하도록 이론적으로 근거가 마련된, 학습률 감소를 결합된 학습률 및 배치 크기 스케줄로 대체하는 원칙적인 알고리즘을 소개합니다.
가속도 분석: 저자들은 직렬 실행 시간 단축에 대한 이론적 상한선을 도출합니다. 가장 공격적인 안정적 램프업( $\alpha = \sqrt{\beta}$ )을 코사인 감소 베이스라인과 비교하여 사용함으로써, 최대 이론적 가속도는 $(1 - 2/\pi) \approx 36.3\%$ 로 계산됩니다.

실험 결과

저자들은 OLMo 코드베이스를 사용하여 Chinchilla 규모( $D=20N$ )로 학습된 150M, 300M, 600M 파라미터 모델에 대해 Seesaw를 평가했습니다.

성능 동일성: Seesaw는 임계 배치 크기에서 학습되었을 때 모든 모델 크기에 대해 표준 코사인 어닐링(cosine annealing) 스케줄의 검증 손실 역학과 일치하는 성능을 보였습니다.
실행 시간 단축: Seesaw는 상당한 실행 시간 단축을 달야냈으며, 벽시계 시간을 약 36% 감소시켜 분석에서 도출된 이론적 한계치에 근접했습니다.
견고성: Seesaw는 AdamW 및 다양한 가중치 감쇠(weight decay) 설정에서도 효과적으로 작동합니다.
공격성의 한계: $\alpha$ 와 $\beta$ 파라미터를 변화시킨 실험은 이론적 제약 조건을 확인해 줍니다. $\alpha = \sqrt{\beta}$ 보다 더 공격적인 스케줄(예: 학습률을 유지하거나 너무 느리게 감소시키면서 배치 크기를 두 배로 늘리는 경우)은 발산이나 최적 미달의 성능을 초래합니다.
임계 배치 크기 한계: 저자들은 Seesaw 전략이 임계 배치 크기까지 효과적임을 관찰했습니다. 노이즈 항이 그래디언트 노름을 더 이상 지배하지 않는 지점(즉, 임계치를 넘어서는 지점)에서는 등가성이 깨지며, Seesaw가 표준 코사인 감소보다 성능이 떨어지게 됩니다.

의의 및 주장

본 논문은 "배치 램프" 관행에 대해 엄밀한 이론적 토대를 제공함으로써, 이를 휴리스틱한 튜닝에서 원칙적인 최적화 전략으로 격상시켰다고 주장합니다. 학습률 감소와 배치 크기 증가가 특정 조건 하에서 서로 교환 가능하다는 것을 입증함으로써, Seesaw는 모델 품질을 저해하지 않고 LLM 사전 학습을 가속화할 수 있는 실용적인 방법을 제시합니다.

저자들은 Seesaw가 적응형 옵티마이저에서 기존 스케줄러를 대체할 수 있는 드롭인 교체재임을 강조합니다. 주요 의의는 하드웨어 개선 속도가 학습 요구량을 따라잡지 못하는 하드웨어 병목 현상을 해결하기 위해, 대규모 모델 학습의 벽시계 시간을 약 1/3 정도 줄일 수 있다는 데 있습니다. 또한 본 연구는 이 전략의 경계를 명확히 하며, 이 방법이 분산 지배 영역에 의존하며 평균 그래디언트 항이 지배적인 임계 임계치를 초과하는 배치 크기에는 적용되지 않을 수 있음을 밝힙니다.