Predict before you train: Scaling Laws for particle physics foundation models

完璧なケーキを焼こうとしていると想像してください。しかし、あなたは小麦粉や砂糖、あるいはどれくらいの熱が必要なのかを知りません。素粒子物理学の世界では、科学者たちは常に、粒子衝突から生じる混沌とした破片を分析するために「デジタル脳（機械学習モデル）」を構築しており、自然界の新しい法則を明らかにする隠れたパターンを見つけ出そうとしています。これらのデジタル脳は巨大なオーブンのようなものです。オーブンが大きければ大きいほど（モデルが大きく）、そして投入する材料が増えれば増えるほど（データが多く）、通常、ケーキの味は良くなります。しかし、これらのケーキを焼くことは非常に高価です。膨大な電力とスーパーコンピュータを必要とするからです。長年、科学者たちは、焼き始める前に自分のオーブンがどのくらいの大きさであるべきかを推測しなければなりませんでした。彼らは何百万ドルもの費用を投じて巨大なオーブンを用意しましたが、後になって、中くらいのサイズのオーブンで十分だったことが判明したり、あるいは望む結果を得るためにはもっと大きなオーブンが必要だったことが分かったりしたのです。大きな疑問はこうでした：オーブンを点火する前に、ケーキの味を予測することはできるのだろうか？

「Predict before you train（訓練する前に予測せよ）」と題されたこの論文は、まさにその問題に取り組んでいます。研究者たちは、特定の計算量（コンピューティング・パワー）をどれだけ必要とするかを、実際に巨大なモデルを構築することなく正確に教えてくれる「レシピの規則（スケーリング則）」を見つけられるかどうかを検証したいと考えました。彼らは、通常はテキストを読むために使われる「トランスフォーマー」と呼ばれる特定の種類のデジタル脳に焦点を当て、ここでは衝突器から飛び出す粒子を読ませるように教えました。一連の小さくて安価なモデルを先に訓練することで、彼らはもしモデルを100倍大きくした場合に何が起こるかを予測するグラフ上の線を引こうと試みました。また、より優れた「事前学習（一般的なパターンの学習）」を行うことが、後にトップクォークの特定や、クォークとグルーオンの識別といった特定の物理学的タスクにおいて、モデルを向上させるのに役立つのかどうかについても調べました。

素粒子物理学のための「水晶玉」

スタンフォード大学とSLACの研究者を中心とするチームは、大胆なアイデアをテストすることに決めました。それは、小さな、安価なモデルの結果を見るだけで、大規模な機械学習モデルの性能を予測できるのか？というものです。これを行うために、彼らは約10億個のシミュレーションされた粒子ジェット（粒子の噴流）を含む「OmniLearned」というデータセットを使用しました。彼らはParticleViTと呼ばれるモデルのファミリーを構築しました。これらのモデルをデジタル探偵のセットと考えてください。最も小さいものは新人探偵であり、最も大きいものは膨大な記憶を持つ超一流の探偵です。

研究者たちは、まず少量の計算量（$10^{19}$ FLOPs未満）を用いて新人探偵を訓練することから始めました。彼らは、これらの小さな探偵がどのように学習するかを観察しました。その後、数学的な公式（スケーリング則）を使用して、記憶容量が100倍大きい探偵を訓練した場合に何が起こるかを予測するグラフ上の線を引きました。

予測は的中した
その結果は、まるで水晶玉を持っているかのようでした。実際に巨大なモデルを訓練したとき（最大$3.4 \times 10^{20}$ FLOPsを使用）、結果は予測された場所とほぼ正確に一致しました。実際の性能は予測の1パーセント以内に収まりました。これは非常に重要なことです。なぜなら、これまでは、素粒子物理学において、まだ訓練していないモデルの性能を予測することに成功した人は誰もいなかったからです。これは、科学者が高価な巨大な訓練を実行する前に、どの程度の計算量を購入する必要があるかを判断するために、小さくて安価なモデルを使って「予行演習」ができることを意味します。

「情報の密度」に関する驚き

最も興味深い発見の一つは、実際にどれほどの「脳の力（モデルサイズ）」が必要かについてでした。言語モデル（エッセイを書いたりチャットしたりするもの）の世界では、大量のデータを扱うために巨大なモデルが必要であるというのが定石です。しかし、素粒子物理学については、研究者たちは異なる結果を見出しました。

彼らは、ジェット内の単一の粒子が運ぶ「情報」は、文章内の単一の単語よりもはるかに少ないことを発見しました。単語は文脈に応じて百万通りの意味を持ち得ますが、粒子はいくつかの特定の数値が付随しているだけの、ただの粒子です。データが情報の「密度」において低いため、この10億粒子データセットに対する最適なモデルサイズは、言語の専門家が予想するよりもはるかに小さくなります。論文は、この特定の仕事においては、巨大な脳は必要ではなく、ただ脳に多くの例を与えればよいのだと示唆しています。モデルサイズの「スイートスポット」は、言語モデルの標準的なルールよりもはるかに小さいことが判明し、素粒子物理学においては、モデルを大きくするよりも、より多くのデータに予算を割く方が良いことを示唆しています。

より優れた「一般知識」は役に立つのか？

チームはまた、一般的な事前学習タスク（すべての粒子についての学習）で優れた成績を収めることが、後に特定の物理学的タスクにおいてモデルを向上させるのに役立つかどうかについても確認しました。彼らは2つの一般的なタスクをテストしました：

1. トップ・タギング（Top Tagging）: ジェットの粒子の集まりが、重い「トップクォーク」から来たものかどうかを識別する。
2. クォーク/グルーオン・タギング（Quark/Gluon Tagging）: クォークによって作られたジェットと、グルーオンによって作られたジェットを区別する。

彼らは明確な直線を見出しました。事前学習中に（つまり、一般的なパターンをより良く学習して）「損失（Loss）」が低かったモデルは、微調整（ファインチューニング）後のこれらの特定のタスクにおいても高い性能を発揮しました。言い換えれば、もしあなたのデジタル探偵が一般的な素粒子物理学において天才であれば、トップクォークを見つける際にも天才になるということです。これは、「計算予算」が期待される物理学的性能に直接翻訳できることを裏付けています。利用可能な計算量を知っていれば、自分のモデルが新しい物理学を発見するためにどれほど優れているかを、今や予測できるのです。

最終決戦

最後に、研究者たちは、彼らの新しい汎用的な「ParticleViT」モデルを、物理学の特定のルール（相対性理論など）が設計に組み込まれた、有名な高度に専門化されたモデルであるOmniLearnedと比較しました。結果は驚くべきものでした。特別な物理学のルールを持っていない汎用的なParticleViTモデルは、ほとんどの指標において、専門化されたOmniLearnedモデルと同等の性能を示しました。

専門化されたモデルがわずかな優位性を示したのは、非常に困難で極端なケースである「トップ・タギング（高純度テール）」においてのみであり、そこでは背景ノイズを排除する能力が約5%高くなっていました。しかし、それ以外のほぼすべてのケースにおいて、ゼロから学習したシンプルな汎用モデルは、同等の性能を発揮しました。これは、多くの物理学的タスクにおいて、必ずしも物理学に特化した複雑なアーキテクチャを構築する必要はなく、十分なデータで訓練された標準的なスケーリング済みのトランスフォーマーで、効果的に仕事をこなせることを示唆しています。

未来にとっての意味

論文は、これら5つのモデル（小規模から超大規模まで）と完全なコードを公開して締めくくられています。これにより、他の科学者たちもこれを利用できるようになります。主な教訓は、この分野における「推測と希望」の時代は終わったということです。科学者は今、小さくて安価な実験を用いて、大規模で高価な実験の結果を予測することができます。特定の目標に到達するためにどれだけの計算量が必要かを正確に計算でき、時間とリソースを節約できます。この論文は、すべての物理学を解決したと主張しているわけではありませんが、機械学習モデルの訓練という高価な風景をナビゲートするための信頼できる地図を提供しました。時には、最も小さなことから始めることが、未来を見るための最善の方法であることを証明したのです。