Towards Engineering Scaling Laws with Pretraining Data Composition

本論文は、高忠実度シミュレータによって安価な合成データの生成が可能である素粒子物理学において、事前学習データセットを戦略的にキュレーションすることにより、ハドロンジェット分類モデルのスケーリング挙動を、モデルのサイズよりもデータの多様性と整合性を優先するように設計できることを示している。

要約

あなたは、学生に街中のさまざまな種類の乗り物を識別する方法を教えようとしていると想像してください。学習を助けるための主な方法は2つあります。一つは、より「大きな脳」（より大きなモデル）を与えること。もう一つは、より多くの「練習問題」（より多くのデータ）を与えることです。

長い間、人工知能（AI）を研究する科学者たちは、ある「黄金律」が存在すると信じてきました。もし、決まった計算資源（コンピューティング・バジェット）があるなら、最も賢い学生を作るためには、リソースを「大きな脳を作る」ことと「より多くの練習問題を与える」ことに、およそ50対50の割合で分配するのが最善である、という考え方です。

しかし、この新しい論文は、素粒子物理学の世界において、「何を最初に学ぶか」を変えることで、より優れたルールを**設計（エンジニアリング）**できることを示唆しています。

設定：物理学の教室

研究者たちは「ジェット」を扱っています。素粒子物理学では、微小な粒子が衝突すると、他の粒子による流れである「ジェット」が噴出します。それは花火が爆発したようなものですが、火花ではなく、亜原子粒子の流れが現れるのです。

目標は、AIにこれらの流れを見せ、「ああ、これは特定の種類の爆発から来たものだ！」と言わせるように教えることです。

実験：教科書を変える

研究者たちは、学習ルールがどのように変わるかを確かめるために、2種類の「教科書」（事前学習データセット）をテストしました。

1. 退屈な教科書（QCDのみ）： この本には、「標準的な」粒子の爆発の例しか含まれていませんでした。それは、標準的なセダン（普通乗用車）の運転だけを教える自動車学校のようなものです。
2. 多様な教科書（BSM拡張版）： この本には、標準的な例に加えて、私たちの通常の宇宙では起こりえない、複雑で稀な、エキゾチックな爆発（「標準模型を超えた（BSM）」物理学のシミュレーション）が含まれています。それは、セダンの運転だけでなく、レーシングカーやトラック、さらには空飛ぶ乗り物まで教える自動車学校のようなものです。

発見：ルールの書き換え

研究者たちが退屈な教科書を使ってAIを訓練したとき、従来の50/50のルールが成立しました。より良い結果を得るためには、脳を大きくすることと、より多くの練習問題を与えることのバランスを取る必要がありました。

しかし、多様な教科書を使用してAIを訓練したとき、ルールは完全に変わりました。AIは、「より多くの練習問題」の方が、「より大きな脳」よりもはるかに価値があることを学んだのです。

• 比喩： 多様な教科書で訓練されたAIは、あらゆる種類の乗り物をすでに見たことのある学生のようなものです。新しいテストを与えられたとき、彼らは新しい車を理解するために大きな脳を必要とするのではなく、ただその例をもっと多く見るだけで完璧になれるのです。彼らの「経験」が豊かであるため、「脳」を成長させるスピードはそれほど速くなくてもよいのです。

結果：新しい「データ優先」戦略

この論文は、初期訓練に多様でエキゾチックなデータを用いた場合、以下のことが分かったと述べています。

• 「大きな脳」戦略の重要性が低下した。
• 「より多くのデータ」戦略が勝者となった。

実際、研究者たちは、計算資源を1単位使うごとに、約78%をデータの取得（増やすこと）に使い、わずか22%をモデルを大きくすることに充てるべきであることを見出しました。これは、従来の50/50の分割から大きな転換です。

なぜこれが物理学にとって重要なのか

この論文は、物理学のユニークな利点を強調しています。それは、**「私たちは自分たちでデータを作ることができる」**ということです。

医学や言語の分野では、新しいデータを得ることは困難で、コストがかかり、あるいは不可能です（新しい人間の患者を単に「シミュレート」することはできません）。しかし、素粒子物理学では、強力なコンピュータを使用して粒子の衝突をシミュレートします。シミュレーションさえ走っていれば、高品質で多様なデータを無限に、無料で生成できるのです。

まとめ：
もしあなたが物理学のための超スマートなAIを作ろうとしているなら、単に可能な限り大きな脳を作ろうとするのではなく、AIが最初に学ぶための、より多様で優れた「カリキュラム」を設計することに時間と資金を投じるべきです。一度、AIが多様な「エキゾチックな」例を目にするようになれば、与えられた特定のタスクからもより速く、より良く学ぶことができます。そして、モデルを大きくするよりも、より多くのデータを供給する方が、より良い結果が得られるのです。

要するに： よく選ばれた、多様な内容のトレーニングデータの「食事」は、より大きな脳よりも強力なのです。

技術概要

技術要約：事前学習データの組成によるエンジニアリング・スケーリング則に向けて

問題提起
ニューラル・スケーリング則は、計算量、モデルサイズ、およびデータセットサイズに対して、モデルの性能がいかにべき乗則に従って向上するかを記述するものである。大規模言語モデル（LLM）においては十分に確立されているが、これらの関係性は素粒子物理学においても現れつつある。基礎物理学における重要な相違点は、大規模なモデルをトレーニングするための計算コストと比較して、シミュレーターを介した高忠実度な合成データを比較的低コストで生成できる能力である。これは、事前学習データ自体をエンジニアリングすることで、スケーリング挙動に影響を与えるという独自の機会を生み出す。本研究が取り組む中心的な問いは、事前学習データの組成（具体的にはその多様性とダウンストリーム・タスクへの適合性）によって、計算量を「より大きなモデル」よりも「より大きなデータセット」へと優先させる計算最適スケーリング・レジームへとシフトさせることができるか、という点である。

手法
本研究は、高エネルギー粒子衝突によって生成されるハドロン・ジェットの分類タスクに焦点を当てている。著者らは、ジェットデータをポイントクラウドとして処理する汎用的なトランスフォーマー・アーキテクチャを使用し、深さとアテンション・ヘッドの次元を固定したまま、モデルサイズを約3,000から1,0,500万パラメータまで変化させた。

実験デザインには、以下の二段階のトレーニング・プロトコルが含まれる：

1. 事前学習： モデルは、188クラスのシミュレーション・ジェットを含むJetClass-IIデータセットのサブセットを用いて事前学習される。著者らは、多様性と適合性を操作するために、4つの異なる事前学習サブセットを定義している：
   • QCD: 軽いクォークまたはグルオンによって生成されたジェットのみ（17クラス）。
   • QCD + res2p: QCDジェットに加え、新標準模型（BSM）レゾナンスの二体崩壊から生成されたジェット。
   • QCD + res34p: QCDジェットに加え、BSMレゾナンスの三体または四体崩壊から生成されたジェット。
   • QCD + res2p + res34p: すべてのBSMレゾナンス崩壊を含むフル・データセット。
   • 注記： BSMサブセットは、QCDのみのデータと比較して、より高い多様性（より多くのプロセス・クラス、より広い位相空間のカバー範囲）と、ダウンストリーム・タスクに対する優れた適合性（マルチ・プロング・トポロジー）を導入する。
2. ファインチューニング： 事前学習されたモデルは、10クラスのジェット分類タスク（軽いクォーク/グルオン、トップクォーク、W/Zボソン、およびヒッグス粒子の識別）のために、オリジナルのJetClassデータセットでファインチューニングされる。このタスクは、プロングの多重度と質量スケールの識別を必要とするが、これらはBSM拡張事前学習データにはよく表現されているものの、QCDのみのデータには乏しく表現されている。

スケーリング指数は、総計算量（$C$）の関数としての計算最適モデルサイズ（$N^*$）およびデータセットサイズ（$D^*$）に対して、べき乗則をフィッティングすることで抽出される。本研究では、「スクラッチ」トレーニング（事前学習なし）と、各種事前学習構成の間でこれらの指数を比較する。

主な結果
本研究は、事前学習データの組成が計算最適スケーリング指数を著しく変化させることを示している：

• スクラッチ・トレーニング： スクラッチからのトレーニングでは、指数の値は $a \approx 0.52$（モデルサイズ）および $b \approx 0.48$（データセットサイズ）となり、モデルサイズとデータの間の計算リソースの割り当てがほぼ均衡していることを示しており、これはLLMにおける知見と一致する。
• QCDのみの事前学習： QCDジェットのみでの事前学習では、指数の変化はわずかであり（$a \approx 0.53, b \approx 0.47$）、特定の適合性や多様性を伴わない事前学習単体では、スケーリング・レジームを根本的に変えることはできないことを示唆している。
• BSM拡張事前学習： 事前学習コーパスにBSMレゾナンス崩壊を含めると、劇的な変化が生じる。フルセットのBSM拡張データセットを用いると、指数は $a \approx 0.22$ および $b \approx 0.78$ へとシフトする。
  • これは、計算最適戦略が、モデルサイズの増大よりもデータセットサイズの増大を強く支持するレジームであることを示している。
  • このシフトは、スクラッチのベースラインと比較して、モデルサイズのスケーリング指数が2.3倍減少したことを意味する。
  • ファインチューニングの損失曲線は、BSM強化事前学習がすべてのモデルサイズにおいて一貫して損失を低下させ、その恩恵はより大きなモデルほど大きくなることを裏付けている。

主要な貢献

1. スケーリング則のエンジニアリング： 本論文は、事前学習データの組成をエンジニアリングすることで、基礎物理学におけるスケーリング指数をシフトさせられることを示す、初の体系的な研究を提供している。事前学習コーパスにおける多様性とダウンストリームへの適合性が、スケーリング・レジームをデータ重視の戦略へと移動させられることを実証している。
2. 定量的シフト： 本研究は、BSM物理学を事前学習に組み込むことで、均衡したスケーリング・レジーム（$a \approx b \approx 0.5$）から、強力なデータ重視のレジーム（$a \approx 0.22, b \approx 0.78$）へのシフトを定量化している。
3. 基盤モデルへの示唆： これらの結果は、多様で適合性の高い合成データで事前学習された基盤モデルは、より小さなパラメータ数で最適な性能を達成できる可能性を示唆しており、これにより節約された計算予算を、さらなるデータの生成へと振り向けることが可能になる。

意義と主張
著者らは、本研究が科学的機械学習における新しい設計空間、すなわち「基盤モデルのトレーニングのための物理学的入力」を特定したと主張している。データの可用性が制限される自然言語や画像ドメインとは異なり、基礎物理学では、安価で高忠実度なシミュレーターを活用して、スケーリング則を能動的に形成する事前学習コーパスを構築できる。

本論文は、組成の優れたコーパスを用いた事前学習によって、ダウンストリームの計算量をより多くのデータに投じるのが最適となるレジームが可能になることを示しつつも、これらの結果が異なるファインチューニング・タスク、より大きなモデルスケール、および異なるデータセットサイズに対して一般化するかどうかを検証するには、さらなる研究が必要であると控えめに結論付けている。本研究は、スケーリングの課題をすべて解決したと主張するものではなく、科学的基盤モデルの発見ポテンシャルを最大化するための、未探索のレバーとしての「事前学習組成エンジニアリング」を強調している。