シーケンス重み付けのスケーリングに関する研究
A study of sequence weighting at scale (blog.janestreet.com)
要約
本研究では、大規模な言語モデル(LM)におけるデータ重み付けのスケーリング則を、社内モデルとオープンウェイトモデルの両方で調査しました。その結果、モデルの規模が大きくなるにつれて、シーケンスの重みに対するモデルの損失削減効果が非単調に変化する挙動が明らかになりました。具体的には、小規模モデルはデータ重みに関係なく一般的なパターンを学習し、中規模モデルはデータ重みに比例して特定のパターンを学習し、大規模モデルは再びデータ重みに関係なく全てのパターンを学習することが示されました。
全文翻訳
TL;DR: インハウスおよびオープンウェイトLLMにおけるデータ重み付けのスケーリング則を研究し、スケール全体にわたる非単調な振る舞いを発見しました。トレーニング中にシーケンスに割り当てられる重みを変化させ、モデルのシーケンスに対する損失削減量がそのシーケンスの重みにどれだけ依存するかを測定しました。全体として、我々の結果は一般的な傾向と一致しています。モデルが小規模から中規模へと移行するにつれて、データ重みに依存しない一般的なパターンの学習から、データ重みに比例したデータ固有パターンの学習へと移行します。その後、モデルが中規模から大規模へと移行すると、再びデータ重みに依存せず、データ中に存在する全てのパターンを学習できるようになります。
ニューラルネットワークをトレーニングする際、モデルに何を学習させたいかを正確に考えることに多くの時間を費やします。全てのコーディングタスクにわたる一般的な知能を持つことと、OCamlに特化して優れていることのどちらをどれだけ重視するか?古い市場レジームからのデータと最近収集されたデータを比較して、どちらがどれだけ有用か?一般的に、高品質なデータに重みを付けることと、モデルがあらゆるデータから学習することを許可することのどちらをどれだけ重視したいか?これらの問いは、最終的にはデータミキシングの問い、すなわちモデルの表現能力とFLOPsをどのように、どこで費やすかという問いです。
さらにこれらの問いを複雑にしているのは、実際には実行したい実験の多くが法外に高価であるということです。中規模スケールでさえ、ハイパーパラメータの密な多次元グリッドサーチを実行することは現実的ではありません。ましてや最大スケール(タスクごとに1つのモデルしかトレーニングできない場合がある)ではなおさらです。これに対する標準的なアプローチは、スケーリング則アプローチを通じてデータミックスウェイトのようなハイパーパラメータを適合させることです。このアプローチでは、小規模スケールでハイパーパラメータを適合させ、それらを大規模スケールに外挿します(この典型的な例はChinchillaスケーリングです)。
[画像: スケール全体にわたる学習率のChinchillaスタイルの適合例。]
図1: スケール全体にわたる学習率のChinchillaスタイルの適合例。
スケーリング則における中心的な課題は、外挿が成り立つためには、大規模モデルが(a)スケールに対して不変である(例:多くのハイパーパラメータの最適値がスケール全体で安定しているMuPスタイルの結果)か、(b)小規模での振る舞いの外挿によって予測可能である(例:損失がスケールと共に予測可能に減少するKaplanスタイルの結果)かのいずれかの振る舞いを持つ必要があることです。残念ながら、全ての振る舞いがこれらの特性を持つわけではありません。一部の振る舞いはスケールで出現し、スケーリング則では予測できません。本稿の残りの部分では、このような予測不可能なスケーリング振る舞いを異常と呼びます。
我々の内部実験では、データミキシングが異常なスケーリング振る舞いを特に引き起こしやすい設定であることが示されています。これはMAI-Thinking-1技術レポート(セクション2.5.2)での発見と並行しており、コード中心のミックスとSTEM中心のミックスのトレーニング品質の相対的な順序がモデルサイズの増加とともに逆転しました。したがって、我々は言語モデルがデータミックスでそのデータに与えられた重みに比例して実際にデータ中のパターンをどれだけ学習するか、そしてその振る舞いがスケール全体でどのように変化するかを正確に特徴付けることに着手しました。
データミックス実験は、コーパスにおけるデータの重みの効果と、他の重要だが異なる考慮事項を混同する可能性があります。特に、データミックス実験の結果は、異なるソース間のデータ品質のばらつきや、任意のデータソースにおける周辺トークンのユニークさの違いによって大きく影響されます。ミックスウェイト単独の効果を分離するために、本稿ではデータ重み付けを特に分析します。これはデータミキシングのバリアントであり、トレーニングデータセット内の個々のシーケンスが損失において異なる重みを受け取ります。与えられた重みを持つトレーニングシーケンスをモデルがどれだけ学習したかを測定するために、まずデータセット内の各シーケンスにランダムなシーケンスウェイトを割り当て、データセットでモデルをトレーニングし、その後同じトレーニングデータセットで再評価します。
次に、シーケンスの重みの何乗が、そのシーケンスに対するモデルの期待損失削減量に最も比例するかを測定する指標、実効シーケンスウェイト指数((p^)と表記)を定義します(例:(p^ = 1)は、モデルのシーケンスに対する期待損失削減量がその重みに直接比例することを示し、(p^* = 0)は、モデルのシーケンスに対する期待損失削減量が全ての可能なシーケンスウェイトに対して同じであることを示します)。
以前の研究は、スケールがどのように影響するかについて混合した予測をしています。(Byrd & Lipton 2019)と一般的な大規模補間の現象は、大規模モデルは(p^* \to 1)に漸近すべきであると予測していますが、Li et al. 2026は、パラメータあたりのトークン数が固定の場合、価値のあるドメインの最適なトークン繰り返し回数はモデルサイズとともにわずかに増加することを発見しています。我々は問います:実効シーケンスウェイト指数はスケールとともに滑らかなべき乗則なのか、それとも異常なのか?
我々は、内部のテキストベースのベンチマークにおけるシーケンス重み付けの効果を評価します。モデルファミリーとスケール全体をテストするために、数十億から数百億パラメータの範囲の2つのインハウス事前学習済みLLMファミリーと、5億から720億パラメータの範囲の1つのオープンウェイトモデルファミリー(Qwen 2.5)を評価します。
我々は、いくつかの異常なスケーリング振る舞い(スケール全体でのデータミキシングに関する推論は複雑です!)といくつかの顕著な傾向を発見しました。全体として、我々の結果は、実効シーケンスウェイト指数における非単調な上昇→下降の傾向と一致しています。小規模モデルは小さな実効シーケンスウェイト指数に適合し、データ重みに依存せずにデータセット全体にわたるパターンを学習します。中規模モデルはより大きな実効シーケンスウェイト指数に適合し、データ重みに比例したデータ中のパターンを学習します。大規模モデルは再び小さな実効シーケンスウェイト指数に適合し、重みに関係なくデータ中に存在する全てのパターンを学習します。エポッキングは、実効シーケンスウェイトのピークをより小さなモデルの方へシフトさせます。
異常なスケーリング則は、我々のモデルトレーニング方法論にとって不便ですが、それらを特定することで悪い外挿を避けることができます。我々の科学では、そのような異常な振る舞いに対して常に警戒しており、それらを発見した際には、それらを予測可能にするようにスケーリング実験と評価を再設計します。1 我々の実験は、十分に大規模なモデルからのみデータミックスの結果を評価・外挿すること、または観測された(p^*)を補償するためにトレーニングウェイトを調整することなど、いくつかの可能な解決策を示唆しています。2 そして、我々はモデルが予測からどれだけ外れているかを測定し続け、スケーリング則が破綻する次の場所を捉え、より大きく、より良いモデルをさらに自信を持ってトレーニングできるようにしています。
もしこのようなことが興味深いなら、応募を検討してください。あなたは、数万台のGPU、ペタバイト級のトレーニングデータ、そして最高のアイデアに投資するための機敏性とリソースを活用する、緊密なグループの優秀で協力的同僚たちに加わることになります。
方法論
データセット。 我々は、内部のテキストベンチマークにおけるシーケンス重み付けの効果を評価します。このデータセットのシーケンスには、0.01から10の範囲で対数一様分布からサンプリングされたシーケンスウェイトを使用します。3エポックトレーニングし、各エポック後に評価します。
モデル。 このデータセットで3つの異なるモデルファミリーを評価します。
- JS-dense:9つのモデルからなる、インハウス事前学習済み密結合LLMのファミリー
- JS-sparse:8つのモデルからなる、インハウス事前学習済みMoE(Mixture of Experts)LLMのファミリー
- Qwen 2.5:5億パラメータから720億パラメータの範囲の、オープンウェイト密結合LLMのファミリー。
JS-denseおよびJS-sparseモデルは、数千万から数百億パラメータの範囲です。全ての文脈において、これらのモデルではホールドアウト性能はスケールとともに増加していることに注意してください。
シーケンス重み付けの効果の測定。 まず、モデルをトレーニングデータセット(各シーケンスが(w_i)に対応する重みを持つ)でトレーニングし、トレーニングデータセットでモデルを再評価し、各シーケンスについてトレーニングによって達成された損失削減量を測定します。次に、(p^*)を、(適切に正規化された)シーケンス(i)に対する損失削減量が、そのシーケンスの重み(w_i)によって最もよく説明されるように適合させます。
この適合を、付随するノートでより正確に行いますが、異なる(p^*)値が何を意味するかをよりよく理解するために、いくつかのシナリオを考えてみましょう。
- (p^* = 0):シーケンスに対する期待損失削減量は、その重みに影響されません(つまり、低重みシーケンスと同じくらいしか高重みシーケンスの損失を削減しません)。
- (0 < p^* < 1):高重みシーケンスの損失は低重みシーケンスよりも多く削減されますが、その削減量はべき乗(p^)に従ってサブ線形にスケールします((p^ = 0.5)の場合、シーケンスの損失削減量はそのトレーニングウェイトの平方根に比例します)。
- (p^* = 1):高重みシーケンスの損失は、シーケンスウェイトが示唆するものよりも集中した方法で、低重みシーケンスよりも多く削減されます(例:モデルが最も重みのあるシーケンスを記憶し、他の全ての行に対してランダムに予測をサンプリングした場合、(p^* \approx 1)が見られるでしょう)。
ハイパーパラメータ。 検証セットに対して損失を最小化するようにハイパーパラメータを調整します。3 最大のインハウスモデルは経済的にハイパーパラメータを調整するには大きすぎるため、Chinchillaスタイルのべき乗則スケーリング法を用いてハイパーパラメータを設定し、スケールの関数として最適なハイパーパラメータを適合させます。
結果と分析
[画像: モデルサイズ(ラダーの段)とトレーニングデータの各エポック数に対する(p^*)(3つのモデルファミリーそれぞれ)。]
図2: モデルサイズ(ラダーの段 #)とトレーニングデータの各エポック数に対する(p^*)。各モデルファミリーの3次元サーフェスは以下にあります。参考として、3つのラダー全体で最も小さいモデルはJS-denseの段1であり、最も大きいモデルはJS-sparseの段8です。
以下に、各モデルファミリーのインタラクティブな3Dサーフェスを示します。
特定のモデルファミリーとエポック数内では、(p^*)はしばしば小規模(特に数十億パラメータまで)で増加し、大規模(数百億パラメータのレジームに近づくにつれて)で減少します(参考として、図示されている最大のJS-denseの段は1000億パラメータ未満です)。
より多くのエポックでトレーニングすると、ピークはより小さなモデルスケールの方へシフトします。モデルファミリー間では、(p^)の実際の値との間にはほとんど関係が見られません。JS-sparseの最小モデルはJS-denseの最小モデルよりも実質的に大きいにもかかわらず、(p^)は小さく、JS-sparseの最大モデルはQwen 2.5の最大モデルよりも実質的に大きいにもかかわらず、(p^*)は大きいです。
なぜ(p^*)は上昇してから下降するのでしょうか?我々は、モデルが異なる速度でデータ中に存在する異なるパターンを学習するためだと仮説を立てています。一部のパターンは全てのシーケンスにわたってうまく一般化します(例えば、英語の理解)、一方、他のパターンは少数のシーケンスのグループに固有です。このデータセットでは、一般的なパターンよりも固有のパターンによる損失削減の可能性の方が大きいです。したがって、小規模モデルは限られた容量を使って一般的なパターンを学習し、固有のパターンを表現できないと仮説を立てています。中規模モデルは、一般的なパターンと最も重要な固有のパターン(シーケンスウェイトによる)の両方を学習します。大規模モデルは、全ての重みを持つ全ての固有のパターンを学習する能力を持っています。エポッキングは、モデルが一般的なパターンと固有のパターンの両方をより多く学習することを可能にします。
より広範には、小規模でのシーケンス重み付けの振る舞いが、最大スケールでの振る舞いをきれいに予測しないことがわかります。これは、あまりにも一般的な結果をうまくカプセル化しています。大規模モデルは、小規模モデルが示すものとは種類が異なる振る舞いを持つことがあります。したがって、我々のスケーリング研究では、曲線から外れた理由を理解しようとすることに、曲線自体を適合させるのと同じかそれ以上の時間と注意を費やす必要があります。
-
このような異常なスケーリング則に対する修正の古典的な例は、多肢選択式質問応答に見られます。小規模モデルはサンプリングされた応答でノイズフロアの周りをしばしばホバリングし、モデルが指示に従うことを学習するまで離散的なジャンプがありますが、バイトあたりのビットのような連続的なメトリックを回答で測定すると、滑らかで異常でないスケーリングが明らかになります。↩
-
特に、実効シーケンスウェイトが(p^)としてスケールする場合、目標損失削減ウェイトを達成するには、減衰されたウェイト、例えば(\tilde{w}_i = w_i^{p^})でトレーニングする必要があるかもしれません。この仮説はテストしていません。↩
-
全てのケースで重い正則化を適用し、結果として得られるモデルがホールドアウトセットで性能をほぼ最大化するようにします。非常に小さいアーキテクチャを除くほとんど全てのアーキテクチャのハイパーパラメータを最初のパスで入力データを完全に記憶するように調整できるかもしれませんが、これは我々が関心のあるシナリオを代表するものではありません。↩