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

技術要約：学習ステップ数とバッチサイズを考慮したスケーリング則

問題提起
Chinchilla則（チンチラ則）に代表される、大規模言語モデル（LLM）における既存のスケーリング則は、モデルサイズ（$N$）および総学習データ量（$D$）の関数としてテスト損失をモデル化している。これらは、$N$と$D$への計算資源の最適な割り当てを決定するには有効であるが、学習ハイパーパラメータ、具体的にはバッチサイズ（$b$）および学習ステップ数（$K$）を明示的に考慮していない。ハイパーパラメータのスケーリングをモデル化しようとする過去の試みは、損失ベースのスケーリング則との統合に失敗しているか、あるいは「クリティカル・バッチサイズ」という概念に依存している。しかし、一部の理論的定式化におけるクリティカル・バッチサイズは、最適なバッチサイズが1であることを示唆しており、これはトークン予算に応じて最適なバッチサイズがスケールするという経験的な知見と矛盾している。さらに、最適なハイパーパラメータのためのスケーリング則を適合させるには、様々なトークン予算に対してバッチサイズを網羅的にスイープする必要があり、膨大な計算コストを要する。

手法
著者らは、学習データ予算 $D$ を、バッチサイズにおけるトークン数（$M = b \times s$、ここで $s$ はシーケンス長）と学習ステップ数（$K$）に明示的に分解する、**三項スケーリング則（3TL）**を提案している。すなわち、$D = MK$である。提案される損失$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)$ は、訓練ランから適合（フィット）されるパラメータである。

著者らは、この 3TL を、モデルサイズ $N$ を固定し、各 $N$ に対して $L(M, K) = E + B/M^\beta + C/K^\gamma$ という形式を用いて独立して適合させた 二項スケーリング則（2TL） アプローチと比較している。

主な手法ステップは以下の通りである：

1. 最適バッチサイズの導出： $D=MK$ という制約の下で 3TL を $M$ に関して最小化することにより、著者らは、最適バッチサイズ $M^\star$ が $D$ の冪乗としてスケールすること（$M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$）を導出した。
2. データの活用： 各構成における単一の最適バッチサイズを特定する必要がある従来の手法とは異なり、3TL は 劣最適なバッチサイズ からのデータを用いて適合される。これにより、3TL は大幅に少ない訓練ランでロバストに適合することが可能となる。
3. 劣最適レジームのモデリング： 理論的な最適バッチサイズがハードウェア制約により実行不可能な場合に対処するため、著者らは二段階の適合手順を提案している。まず、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 の予測をクリティカル・バッチサイズに関する経験的な知見と比較している。

主な結果
著者らは、2つのデータセット（Li (Li et al., 2025a) および OpenEuroLLM）を用いて彼らの手法を評価した。

• 最適バッチサイズのスケーリング： 3TL は、最適バッチサイズ $M^\star$ のスケーリングを正常に回収した。Li データセットにおいて、導出されたスケーリング $M^\star \approx 0.667 \cdot D^{0.566}$ は、Li et al. (2025a) による「ステップ則」（$M^\star \approx 0.58 \cdot D^{0.571}$、フルバッチサイズ・スイープから導出されたもの）と密接に一致している。
• 計算効率： 主要な知見は、3TL は従来の手法で必要とされる訓練ランのわずか 28% を用いて適合できることである。$(N, D)$ の組み合わせにつき2つのバッチサイズのみを含むようにデータセットをマスクした場合でも、3TL は $M^\star$ の正しいスケーリング則を回収できた。対照的に、利用可能な「最良の」バッチサイズのみに直接パワーローを適合させると（損失曲面をモデル化せずに）、スケーリング則は極めて不安定で不正確なものとなった。
• 劣最適バッチサイズ： 二段階の適合手順は、最大5%の計算資源の無駄を招く劣最適バッチサイズの区間を正常にモデル化した。著者らは、この区間の幅がバッチサイズにおいておよそ $2^2$（4倍の係数）であり、そのスケーリング挙動は異なるモデルサイズやデータセット間で一貫していることを見出した。
• クリティカル・バッチサイズ： 3TL は、バッチサイズの増加に対して目標損失に達するためのステップ数が、ある臨界点を超えると感度が低下するという現象を正しくモデル化している。決定的なことに、このモデルは、このクリティカルな挙動がトークン予算 $D$ と共にスケールする一方で、モデルサイズ $N$ にはほとんど依存しないことを示しており、これは Zhang et al. (2025) による経験的な観察と一致している。
• Chinchilla との比較： 3TL を標準的な Chinchilla 形式（最適 $M^\star$ を損失方程式に代入することで）に簡略化すると、モデルサイズ $N$ に対する示唆される計算最適化スケーリングは、Chinchilla 則の直接的な適合とは異なる。3TL は $D$ に対するより小さなスケーリング指数を示唆しており、これはスケーリング則の定式化の選択が、導出される最適構成に影響を与えることを浮き彫りにしている。

意義と主張
本論文は、3TL が Chinchilla 型のスケーリング則（モデルサイズ vs データ）とハイパーパラメータのスケーリング則（バッチサイズ）の間の溝を埋める統一的なフレームワークを提供すると主張している。その意義は主に以下の3点にある：

1. 疎なデータに対する堅牢性： 劣最適なバッチサイズを用いて法則を適合できる能力は、スケーリング則を導出するための計算コストを劇的に削減し、網羅的なスイープなしに最適なバッチサイズを決定することを可能にする。
2. 理論的一貫性： その定式化は確率的最適化の収束境界と構造的に類似しており、経験的なスケーリング分析と最適化理論の間の理論的な架け橋を提供する。
3. 実用的な適用可能性： 最適バッチサイズが1であることを示唆したり、劣最適レジームを記述できなかったりする従来のモデルとは異なり、3TL は非自明な最適バッチサイズを提供し、ハードウェア制約下での実用的なトレーニングにおける性能を明示的にモデル化する。

著者らは、パラメータ値（例：既約損失項 $E$）がデータセット間で一貫しない可能性や、依然として各構成に対して最適な学習率を必要とする点などの限界についても言及している。また、3TL が示唆する計算最適化スケーリングが直接的な Chinchilla 適合とは異なることから、それが最適なモデルサイズを決定するための唯一の手段ではない可能性も認めており、むしろステップ、バッチサイズ、およびデータの相互作用を理解する上で優れていることを示唆している。