[논문 리뷰] Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful

Martin Marek, Sanae Lotfi외 3명|ArXiv.org|2025. 07. 09.

Topic Modeling인용 수 3

한 줄 요약

이 논문은 배치 크기가 1을 포함한 작은 배치 크기가 언어 모델을 안정적이고 견고하게 학습시킬 수 있으며, 대개 큰 배치 및 최적화기 중심의 접근 방식과 대등하거나 이를 능가하는 경우가 많고, 그라디언트 축적을 피할 것을 권고한다.

ABSTRACT

Conventional wisdom dictates that small batch sizes make language model pretraining and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes all the way down to batch size one, and we propose a rule for scaling Adam hyperparameters to small batch sizes. In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, we show that a small batch size combined with an optimizer with a small state size can provide the performance benefits of full fine-tuning while maintaining a similar memory footprint to LoRA.

연구 동기 및 목표

• 사전학습과 미세조정 전반에 걸쳐 작은 배치 크기가 언어 모델 학습을 안정화시킬 수 있는지 조사한다.
• 다양한 배치 크기에서 서로 다른 최적화 기법과 하이퍼파라미터의 강건성을 검토한다.
• 처리량과 메모리 제약에 따라 Adam 하이퍼파라미터를 확장하고 배치 크기를 선택하는 실용적 지침을 개발한다.
• 작은 배치 크기를 사용하는 것과 그라디언트 축적을 선택할 때의 메모리 및 하드웨어 고려사항을 탐구한다.

제안 방법

• 배치 크기를 1에서 4096까지 다양하게 하여 SGD, Adam, Adafactor, Muon에 대한 포괄적인 그리드 탐색을 수행한다.
• 각 배치 크기별로 학습률과 Adam 감쇠 인자(beta1, beta2)를 조정하여 검증 손실을 최소화한다.
• 두 번째 모멘트의 반생명 t2 개념을 도입하고 토큰 기반 반생명에 따른 beta2의 스케일링 규칙을 제시한다.
• 소배치 구간에서 기본 SGD(모멘텀 없음)와 Adafactor와 같은 메모리 효율적 최적화기를 비교한다.
• 더 큰 모델(GPT-2 124M 및 GPT-3 1.3B)에서 확장 휴리스틱을 검증하고 미세조정 시나리오에서도 확인한다.
• 메모리 영향 평가와 메모리가 제약된 학습에 대한 실용적 권고를 제공한다.

실험 결과

연구 질문

• RQ1 아주 작은 배치 크기(최저 1까지)가 모멘텀이나 복잡한 최적화기를 사용하지 않아도 언어 모델을 안정적으로 학습시킬 수 있는가?
• RQ2 작은 배치 크기에서 성능을 유지하기 위해 배치 크기에 따라 Adam 하이퍼파라미터를 어떻게 확장해야 하는가?
• RQ3 작은 배치 크기가 큰 배치 크기와 비교하여 하이퍼파라미터 잘못 설정에 대한 강건성을 제공하는가?
• RQ4 언어 모델 학습에서 작은 배치 사용과 그라디언트 축적 간의 메모리 및 하드웨어 영향은 무엇인가?
• RQ5 발견이 미세조정 및 더 큰 모델 규모에도 확장되는가?

주요 결과

• 적절히 하이퍼파라미터를 스케일링하면 작은 배치 크기가 모든 최적화기에서 대형 배치를 일관되게 따라가거나 능가한다.
• 작은 배치 크기에서는 모멘텀이 덜 필요하며, 기본 SGD가 경쟁력 있게 작동한다.
• 토큰 기준 두 번째 모멘트 반생명(t2)을 고정하고 Adam을 확장하는 것이 beta2를 고정하는 것보다 배치 크기에 관계없이 더 나은 성능을 낸다.
• 그라디언트 축적은 자주 불필요하며 Adafactor나 SGD와 같은 더 간단한 최적화기를 사용한 작은 배치에 비해 메모리 비효율적이다.
• 매우 큰 모델의 경우 배치 크기 1의 SGD가 최소한의 조정으로 AdamW 기초선과 맞먹고, Adafactor가 메모리-성능의 균형을 우수하게 제공한다.
• 권장되는 관행은 처리량을 극대화하는 가장 작은 배치 크기를 사용하는 것이며, 다중 디바이스 설정이 아닌 한 그라디언트 축적을 피하는 것이다.