Explorative Modelingが生成モデルの事前学習にbest-of-K探索を追加

Why it matters

XM reframes generative scaling around training-time search. If its efficiency gains survive larger models, labs could trade extra pretraining work for cheaper inference and better data use.

Alexi Gladstone (@AlexiGlad)は7月31日に、新しい生成モデルの訓練手法を詳述した。この手法は各例について複数の候補出力を探索し、最も近い一致に対して学習することで計算を事前学習に移し、より良いモデルと安価な推論を目指す。

GladstoneはNSF Graduate Research Fellowであり、University of Illinois Urbana-Champaignの計算機科学博士課程の学生で、Explorative ModelingプロジェクトをUIUCの教授Heng Ji (@hengjinlp)およびHarvardの助教Yilun Du (@du_yilun)と共同で主導した。UIUCのBLENDER LabはGladstoneの研究領域をワールドモデル、システム2思考、自己教師あり学習、およびマルチモーダル学習と記載しており、MetaやPalantirでの以前のインターン経験もある。

研究者らは、探索がモデルパラメータとデータに並ぶ第3の事前学習軸になり得ると主張している。彼らの証拠は、画像生成、動画生成、masked diffusion language models、ロボットポリシー、目標条件付きワールドモデルにまたがる実験から得られている。結果は7月29日のarXivプレプリントで発表され、GitHubでPyTorchのコードが公開されている。

探索を訓練に取り込む

生成モデルは基本的な問題に直面している：プロンプトや入力には多くの妥当な出力があり得るということだ。再構成損失のもとで一つの直接的な予測を行うよう訓練されたモデルは、それらの可能性の平均に収束してしまい、どの可能性とも一致しない結果を生成してしまう。Diffusionや自己回帰型システムは、生成を小さなノイズ除去ステップやトークン予測に分割することでその平均化の多くを回避している。

Explorative Modeling（またはXM）は別の方向を取る。各訓練ステップでモデルはK個の候補を生成し、それぞれをターゲットと比較して損失の最も小さい候補を通じてバックプロパゲーションを行う。最も単純な実装はbest-of-Kループである。公開リポジトリでは研究者は--xm_best_of_k Kフラグでこの手法を有効にでき、K=1が探索を伴わないベースラインとなる。

各追加候補は計算コストを伴う。論文は、Forward XMにおける追加候補は別の完全な逆伝播を伴わない別のフォワードパスを必要とするため、標準的なトランスフォーマー訓練ステップのおおよそ3分の1のコストになると見積もっている。候補はバッチ次元に折り込んで並列処理することもできる。したがってこの手法は直接的なトレードオフをもたらす：より良い潜在ノイズとデータの組み合わせを見つけるために訓練時に追加の計算を使い、その支出をより速い収束や生成ステップの削減で回収する可能性がある。

Gladstoneはプロジェクトエッセイで、この研究が生成モデリングが難しい理由を第一原理から説明しようとしたことから生じたと書いている。「我々は生成モデルの規模や訓練データ量を拡大している…それならば彼らが生成できるものをなぜ拡大してこなかったのか？」と述べている。

報告された改善

著者らは、画像生成のレシピに探索を追加したところ、ベースラインの最良結果に到達するのに必要な訓練サンプル数が6.2倍少なく、総FLOPsが4.1倍少なくなったと報告している。大規模モデルが5つの候補を探索した場合、パラメータが47%多い超大規模ベースラインを上回ったとも論文は述べている。

これらの数字は著者らが行った実験を測ったもので、他のアーキテクチャに一般的な削減が当てはまることを意味するものではない。論文は、データをスケールした際に探索の相対的利得が7%から36%に増加し、モデルサイズの増加に伴って13%から23%に増加したと報告している。テスト範囲内ではKの増加に伴って画像と動画の指標が改善し、探索的masked diffusion language modelはベースラインよりも優れたパープレキシティと多様性のトレードオフを示した。

推論結果は、より大規模で再現されればより重要になる可能性がある。Explorative Policyは、論文におけるDiffusion Policyのベースラインに対して5つのロボット操作タスクで同等かそれ以上の性能を示しながら、100回ではなく1回のネットワークフォワードパスで済ませた。Maze2Dのプランニングタスクでは、研究者らはDiffuserベースラインより高い平均スコアを報告しており、推論ステップは16倍から256倍少なくなっている。

スケーリングに関する主張はより大きな検証を待っている

論文の中心的主張は、現在の実験規模を超えている。画像実験は256×256ピクセルのクラス条件付きImageNetを用いており、動画実験は128×128解像度で10フレームをモデル化している。研究者らはより高解像度の動画実験は計算予算を超えたと述べている。独立した再現実験やより大きなモデルでのテストによって、探索が訓練コストが最先端規模に達した場合でも計算効率的であり続けるかどうかが決まるだろう。

言語は別の境界である。論文はmasked diffusion language modelsでより明確な結果を報告しており、自己回帰型言語モデルは検索可能な潜在変数を導入することが自然ではないため改善が難しかったと述べている。著者らは自己回帰型の結果を控えめだと表現しており、より完全な実験は後の研究に委ねている。この区別は重要だ。というのも支配的な大規模言語モデルのアーキテクチャは依然として自己回帰型だからである。

完全なエンドツーエンドの画像生成もコストの問題を生む：複雑な分布におけるモード数に応じて候補数を増やす必要があるかもしれない。研究者らは、生成ごとにデータターゲットを低い計算コストで探索するXMの逆バージョンを提案しているが、そのアプローチは崩壊を防ぐためのエントロピーまたはカバレッジ制約を必要とする。

Explorative Modelingは引き続きUIUCおよびHarvardに所属する学術的かつオープンソースのプロジェクトである。Laude Instituteは6月にそれを第3回Slingshots助成のコホートに含めたほか、論文はFlapping Airplanes fellowshipとNational Science Foundation Graduate Research Fellowship Programからの支援に感謝を示している。

短期的な価値は検証可能性にある。XMは基礎となるアーキテクチャを置き換えることなく既存のdiffusionやflowの訓練パイプラインに追加でき、研究チームに対して事前学習中に追加の候補探索を行うことが、計算コストに見合うだけの収束、品質、あるいは推論節約をもたらすかどうかを計測する具体的な手段を提供する。