How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

기술 요약: 학습 단계 및 배치 크기에 따른 스케일링 법칙

문제 정의
Chinchilla law와 같은 기존의 거대 언어 모델(LLM) 스케일링 법칙은 모델 크기($N$)와 총 학습 데이터($D$)의 함수로 테스트 손실을 모델링한다. 이러한 방식은 $N$과 $D$ 사이의 계산량 최적 배분을 결정하는 데는 효과적이지만, 학습 하이퍼파라미터, 구체적으로 배치 크기($b$)와 학습 단계 수($K$)를 명시적으로 고려하지는 않는다. 하이퍼파라미터 스케일링을 모델링하려는 이전의 시도들은 손실 기반 스케일링 법칙과 통합하는 데 실패하거나, 이론적 구성에서 최적 배치 크기가 1임을 암시하는 "임계 배치 크기(critical batch size)" 개념에 의존하는데, 이는 토큰 예산에 따라 최적 배치 크기가 확장된다는 경험적 발견과 모순된다. 더욱이, 최적 하이퍼파라미터를 위한 스케일링 법칙을 피팅하는 작업은 다양한 토큰 예산에 대해 배치 크기를 전수 조사(exhaustive sweeps)해야 하므로 막대한 계산 비용이 발생한다.

방법론
저자들은 학습 데이터 예산 $D$를 배치 크기(토큰 단위, $M = b \times s$, 여기서 $s$는 시퀀스 길이)와 학습 단계 수($K$)의 곱($D = MK$)으로 명시적으로 분해하는 **3항 스케일링 법칙(3TL)**을 제안한다. 제안된 손실 함수 $L$의 형태는 다음과 같다:

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

여기서 $(E, A, B, C, \alpha, \beta, \gamma)$는 학습 실행으로부터 피팅된 파라미터들이다.

저자들은 모델 크기 $N$을 고정하고 각 $N$에 대해 $L(M, K) = E + B/M^\beta + C/K^\gamma$ 형태의 **2항 법칙(2TL)**을 사용하여 3TL을 비교한다.

주요 방법론적 단계는 다음과 같다:

1. 최적 배치 크기 유도: 제 constraint $D=MK$하에서$M$에 대해 3TL을 최소화함으로써, 저자들은 $D$의 거듭제곱 법칙으로 스케일링되는 최적 배치 크기 $M^\star$를 유도한다: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. 데이터 활용: 모든 구성을 위해 전체 스위프(full sweeps)가 필요한 기존 방식과 달리, 3TL은 차적(suboptimal) 배치 크기로부터의 데이터를 사용하여 피팅된다. 이를 통해 훨씬 적은 수의 학습 실행만으로도 법칙을 견고하게 피팅할 수 있다.
3. 차적 영역 모델링: 이론적 최적 배치 크기가 하드웨어 제약으로 인해 불가능한 경우를 다루기 위해, 저자들은 2단계 피팅 절차를 제안한다. 첫째, 3TL을 사용하여 단순화된 함수 형태 $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ ($\gamma \approx \beta$ 가정)를 도출한다. 둘째, $\epsilon$ (여기서 $\epsilon$은 5%의 계산 낭비 임계값) 이내의 차적 배치 크기 구간에 대해 거듭제곱 법칙을 피팅한다.
4. 임계 배치 크기 분석: 저자들은 목표 손실 $\bar{L}$에 도달하는 데 필요한 단계 수 $K_{\bar{L}}(b)$를 배치 크스의 함수로서 분석하며, 3TL 예측치를 임계 배치 크기에 관한 경험적 발견과 비교한다.

주요 결과
저자들은 Li (Li et al., 2025a) 및 OpenEuroLLM 두 데이터셋에 대해 자신들의 접근 방식을 평가했다.

• 최적 배치 크기 스케일링: 3TL은 최적 배치 크기 $M^\star$의 스케일링을 성공적으로 복구한다. Li 데이터셋에서 유도된 스케일링 $M^\star \approx 0.667 \cdot D^{0.566}$은 Li et al. (2025a)의 "Step-Law" ($M^\star \approx 0.58 \cdot D^{0.571}$)와 밀접하게 일치하며, 이는 전체 배치 크기 스위프를 통해 유도된 것이다.
• 계산 효율성: 주요 발견 중 하나는 3TL이 전통적인 방식에 필요한 학습 실행의 **28%**만을 사용하여 피팅될 수 있다는 점이다. $(N, D)$ 조합당 2개의 배치 크기만을 포함하도록 데이터셋을 마스킹하더라도, 3TL은 여전히 $M^\star$에 대한 올바른 스케일링 법칙을 복구한다. 반면, (손실 표면을 모델링하지 않고) 가용한 몇 개의 "최적" 배치 크기에 직접 거듭제곱 법칙을 피팅하는 것은 매우 불안정하고 부정확한 스케일링 법칙을 초래한다.
• 차적 배치 크기: 2단계 피팅 절차는 최대 5%의 계산 낭비를 초래하는 차적 배치 크기의 구간을 성공적으로 모델링한다. 저자들은 이 구간의 폭이 배치 크기에서 약 $2^2$ (4배) 정도이며, 그 스케일링 동작이 서로 다른 모델 크기와 데이터셋에 대해 일관적임을 발견했다.
• 임계 배치 크기: 3TL은 배치 크기 증가에 따라 목표 손실에 도달하는 단계 수가 특정 지점 이후로 덜 민감해지는 현상을 정확히 모델링한다. 결정적으로, 이 모델은 이러한 임계 동작이 토큰 예산 $D$에 따라 스케일링되지만 모델 크기 $N$에는 거의 불변함을 보여주며, 이는 Zhang et al. (2025)의 경험적 관찰과 일치한다.
• Chinchilla와의 비교: 3TL을 표준 Chinchilla 형태로 축소하면(즉, 최적의 $M^\star$를 손실 방정식에 대입하면), 암시된 모델 크기 $N$에 대한 계산 최적 스케일링은 Chinchilla 법칙을 직접 피팅한 것과 다르다. 3TL은 $D$에 대한 더 작은 스케일링 지수를 제안하며, 이는 스케일링 법칙의 공식 선택이 유도된 최적 구성에 영향을 미침을 강조한다.

의의 및 주장
본 논문은 3-term law가 Chinchilla 스타일의 스케일링 법칙(모델 크기 vs 데이터)과 하이퍼파라미터 스케일링 법칙(배치 크기) 사이의 간극을 메우는 통합된 프레임워크를 제공한다고 주장한다. 이 연구의 의의는 다음 세 가지 주요 영역에 있다:

1. 희소 데이터에 대한 견고성: 차적 배치 크기를 사용하여 법칙을 피팅할 수 있는 능력은 스케일링 법칙을 도출하기 위한 계산 비용을 획기적으로 줄여주며, 전수 조사를 수행하지 않고도 최적의 배치 크기를 결정하는 것을 가능하게 한다.
2. 이론적 일관성: 이 공식은 확률적 최적화의 수렴 경계(convergence bounds)와 구조적으로 유사하며, 경험적 스케일링 분석과 최적화 이론 사이의 이론적 가교를 제공한다.
3. 실용적 적용 가능성: 최적 배치 크기가 1임을 암시하거나 차적 영역을 설명하지 못하는 이전 모델들과 달리, 3TL은 비자명한(non-trivial) 최적 배치 크기를 제공하며 하드웨어 제약 하에서의 실질적인 훈련에 중요한 차적 영역의 성능을 명시적으로 모델링한다.

저자들은 데이터셋 간의 파라미터 값(예: 기약 손실 항 $E$)에서의 잠재적 불일치와, 이 법칙이 여전히 각 구성에 대한 최적의 학습률을 필요로 한다는 점을 한계로 언급한다. 또한 3TL이 암시하는 계산 최적 스케일링이 직접적인 Chinchilla 피팅과 다르다는 점을 인정하며, 이것이 최적 모델 크기를 결정하기 위한 유일한 도구는 아닐 수 있지만 단계, 배치 크기, 데이터 간의 상호작용을 이해하는 데 있어 더 우월하다고 시사한다.