대규모 사전 학습을 위한 배치 크기는 어떻게 설정해야 하는가?

How to Set the Batch Size for Large-Scale Pre-training?

OpenAI가 개척한 임계 배치 크기(Critical Batch Size) 개념은 오랫동안 대규모 사전 학습의 기반 원칙으로 여겨져 왔다. 그러나 Warmup-Stable-Decay (WSD) 학습률 스케줄러로의 패러다임 전환과 함께, 기존의 이론적 프레임워크와 그 기저 메커니즘이 새로운 사전 학습 역학(dynamics)과는 더 이상 부합하지 않는다는 사실이 관찰되었다. 이론과 실제의 간극을 좁히기 위해, 본 논문은 WSD 스케줄러에 맞춘 수정된 E(S) 관계식을 도출하여 사전 학습 중 데이터 소모량 E와 학습 단계 수 S 간의 상충 관계(trade-off)를 규명한다. 우리의 이론적 분석은 WSD 기반 사전 학습의 두 가지 핵심 속성을 밝혀낸다. 첫째는 목표 손실(loss)을 달성하기 위한 최소 배치 크기 임계값인 B_min이며, 둘째는 총 토큰 수를 최소화하여 데이터 효율성을 극대화하는 최적 배치 크기인 B_opt이다. 이러한 속성들을 바탕으로 우리는 동적 배치 크기 스케줄러(dynamic Batch Size Scheduler)를 제안한다. 광범위한 실험 결과, 우리의 수정된 공식이 대규모 사전 학습의 역학을 정확하게 포착하며, 이에 따른 스케줄링 전략이 학습 효율성과 최종 모델 품질을 크게 향상시킨다는 점을 입증하였다.

The concept of Critical Batch Size, as pioneered by OpenAI, has long served as a foundational principle for large-scale pre-training. However, with the paradigm shift towards the Warmup-Stable-Decay (WSD) learning rate scheduler, we observe that the original theoretical framework and its underlying mechanisms fail to align with new pre-training dynamics. To bridge this gap between theory and practice, this paper derives a revised E(S) relationship tailored for WSD scheduler, characterizing the trade-off between training data consumption E and steps S during pre-training. Our theoretical analysis reveals two fundamental properties of WSD-based pre-training: 1) B_min, the minimum batch size threshold required to achieve a target loss, and 2) B_opt, the optimal batch size that maximizes data efficiency by minimizing total tokens. Building upon these properties, we propose a dynamic Batch Size Scheduler. Extensive experiments demonstrate that our revised formula precisely captures the dynamics of large-scale pre-training, and the resulting scheduling strategy significantly enhances both training efficiency and final model quality.