MLE-World、AI研究エージェントに機械学習実験を事前予測させる

導入

AI研究エージェントは、機械学習プログラムを作成し、修正し、実験結果を見ながら改善できる。しかし、候補を評価するには通常、実際にモデルを学習させて検証しなければならない。候補が増えるほど計算コストも膨らむため、エージェントが一度に探索できるアイデアの数には限界がある。MLE-Worldは、この問題に対して実験結果の「事前予測」を導入する。

主なポイント

• 実行結果を先に予測。 MLE-WorldはQwen3.5-4Bを基盤に微調整され、プログラムのコードとデータ説明から、エラー追跡、学習ログ、検証スコアを予測する。
• 信頼度も推定。 予測値だけでなく、その予測が正しい可能性を示す信頼度スコアも出力する。信頼度が低い候補は実際に実行し、予測を検証する。
• 5種類の探索方針。 順次修正、Best-of-N、ペア比較、融合、そして信頼度に基づく探索ボーナスを加えたUCB型の方針を構築した。
• 大規模モデルを上回る予測精度。 17個のMLE-bench Liteタスクにおいて、留保されたプログラムの実際の終端出力との一致率は77.9%。比較対象の、より大きなモデルをプロンプトだけで使う方法は41.1%だった。
• 実行回数を削減。 MLE-Worldを使ったテスト時スケーリングは、世界モデルを使わないエージェントより観測平均性能を最大22.1%高め、実際の実験も少なくした。

意義と影響

MLE-Worldは、学習や評価を完全に置き換えるものではない。まず安価な予測で候補をふるいにかけ、計算資源を有望な案や不確実な案の検証に集中させる仕組みである。この構成により、エージェントは各ステップでより多くの候補を比較しながら、実行コストの増加を抑えられる。

重要なのは、信頼度を単なる補助情報ではなく、行動を決める信号として使っている点だ。予測されたログやスコアは実験の事実ではないため、誤りを前提にして実行へ戻る経路が必要になる。高信頼度の予測は探索を広げるために使い、低信頼度の予測は実行による確認につなげる。

この研究は、AI研究エージェント向けの世界モデルが、コードの結果を予測するだけでなく、限られた実行予算の配分にも利用できることを示す。今後は予測精度だけでなく、不確実性をどれだけ適切に見積もれるかが、実用性を左右するだろう。

出典：Hugging Face Daily Papers