Chinchillaのスケーリング：再現の試み

要約

Hoffmann et al. (2022) は、一定の計算予算のもとでTransformer言語モデルを学習させる際の最適なモデルサイズとトークン数を調べた。著者らは400を超える言語モデルを学習させ、計算最適な学習では、モデルサイズと学習トークン数を同じ割合で増やすべきだと結論づけた。つまり、モデルサイズを2倍にするごとに、学習トークン数も2倍にすべきだということである。著者らはこの結果に従って計算最適となるように、Chinchillaという700億（70B）パラメータのモデルを1.4兆トークンで学習させた。パラメータあたり20トークンの比率である。このため、彼らが提案したスケーリング則はしばしば「Chinchillaのスケーリング則」と呼ばれる。

Hoffmannらは論文の中で、最適なスケーリング方針を導くために三つの異なる手法を用いている。そのうちの一つでは、次のパラメトリックなスケーリング則を推定している。

L = E + \frac{A}{N ^{α}} + \frac{B}{D ^{β}}

この式は、交差エントロピー損失Lが、モデルのパラメータ数Nと学習トークン数Dの関数としてどう減少するかを予測する。このスケーリング則は大きな影響力をもち、その推定値はこのテーマの研究で広く使われてきた。私たち自身の研究も例外ではない。

私たちはHoffmannらの論文の図からデータを抽出してその一部を再構成し、同じパラメトリックモデルを当てはめた。分析の結果、Hoffmannらのスケーリング則のパラメータ推定値には、いくつかの問題がある可能性が明らかになった。

データ再構成のノイズの可能性を考慮し、外れ値のモデルを除外しても、Hoffmannらの推定したモデルは再構成したデータにうまく当てはまらない。
Hoffmannらが報告している信頼区間は、彼らがもっていたと考えられるデータ点の数（約400）からすると、ありえないほど狭い。これほど狭い区間を得るには、数十万の観測値が必要になる。
Hoffmannらの推定したモデルが示すスケーリング方針は、彼らのほかの手法とも、Chinchillaモデルの学習に使われた「パラメータあたり20トークン」という経験則とも矛盾している。

これに対して、私たちのパラメータ推定値（下の式を参照）では、データのスケーリング指数 $β$ と削減不能な損失 $E$ の値がより大きい。このモデルは再構成したデータによく当てはまり、そこから導かれるスケーリング方針は、Hoffmannらのほかの手法とも、Chinchillaの学習に使われたトークン対パラメータ比とも整合する。本研究の結果は、規模をさらに拡大したときの見返りや、最適なスケーリング方針がどこまで不確かなのかについての理解に影響する。

L (N, D) = 1.8172 + \frac{482.01}{N ^{0.3478}} + \frac{2085.43}{D ^{0.3658}}

結果をまとめたプレプリントを公開した。結果はこのColabノートブックで再現できる。Hoffmannらから再構成したデータセットはこちらにあり、データの再構成はこちらで再現できる。

本研究の意味するところ

これまでの研究は、Chinchillaのパラメトリックなスケーリング則についてHoffmannらが当てはめた推定値に頼ってきた。Epochの研究の一部もそうである。今回の発見からは、こうした結果の多くを見直す必要があるかもしれないと考えられる。

私たちは、Hoffmannらが示したスケーリング則のうち、パラメトリックなスケーリング則（同論文の手法3）が最も有用で、最も幅広く適用できると考えている。手法1と2は計算最適な（つまり「Chinchilla最適な」）モデルを学習させる場合にしか当てはまらない。これに対してパラメトリックなスケーリング則は、モデルサイズ（N）とデータ（D）をどう選んでも、データセット上の損失を直接推定できる。そのため、テスト損失を最適化するようにNとDを選ぶあらゆる場面で役に立つ。

パラメトリックなスケーリング則の有用性を理解するために、大規模言語モデルは学習後、一般に本番環境に投入されてユーザーに提供されることを考えてみよう。推論実行のコストを最小にするには、通常、計算最適とされるよりも（少なくとも事前学習の計算量という意味で）小さなモデルを長く学習させたほうがよい。したがってパラメトリックなスケーリング則は、推論実行と学習の総コストの予算が与えられたとき、モデルサイズと学習時間の最適なトレードオフを決める助けになる。

このトレードオフはSardanaらが示している。彼らは推論実行のコストを考慮して、Chinchillaのスケーリング則に代わる独自のスケーリング則を構築した。しかし、自分たちで事前学習の実験を行わなかったため、Sardanaらの研究は、Hoffmannらのパラメトリックなスケーリング則のうち当てはまりの悪いものに頼っていた。同様に、Faizらは、Hoffmannらの当てはめた値を使って、モデルの学習中の二酸化炭素排出量とテスト損失のトレードオフを計算した。Corrêaらは、これらの当てはめた値を使って、TeenyTinyLlamaと名づけた小さなオープンソースLLMの適切なサイズと学習時間を見積もった。

さらに、nostalgebraistによる人気のブログ記事chinchilla’s wild implicationsは、Hoffmannらの当てはめた値を使い、より大きな言語モデルを学習させるには「データが尽きつつある」かもしれないという考えを説明した。

パラメトリックなスケーリング則についてHoffmannらの値に頼っている点は、Epoch自身の研究の一部にも共通している。たとえば、Epochの研究者たちは、研究開発のタスクで人間の研究者の代わりになれるモデルを学習させるのに必要な計算量の上限を見積もるために、彼らの値を使った。また、Sardanaらの手法と同じように、Epochの研究者Pablo VillalobosとDavid Atkinsonは以前、Hoffmannらの推定値を使って、モデルの規模拡大と推論実行のコスト削減とのトレードオフを分析した。David Owenも、モデルサイズとデータ量からLLMのベンチマーク結果を予測する試みの中でこれらの値を引用していた。

したがって、このスケーリング則の修正は、Epoch自身の研究を改善するうえで有益な作業だった。その影響は私たちの研究にとどまらないかもしれない。