AE Studioは最近、Modalプラットフォーム上で強化学習を用いて言語モデルに数学の定理証明を学習させる研究を実施しました。彼らは従来のGroup Relative Policy Optimization（GRPO）と、自然選択に触発された進化戦略（ES）の2つの手法を比較しました。研究では、Lean形式検証言語を報酬信号として使用しました。モデルがLeanコードを生成し、コンパイラがその正しさを検証することで、明確な報酬が得られます。
実験を効率的に実行するため、AE StudioはModalプラットフォームを選択しました。Modalは並列GPU計算、サンドボックス分離、永続的ボリュームストレージなどの重要な機能を提供します。具体的なワークフローは3つの部分からなります：証明生成はvLLMを用いてGPU上で実行、証明検証はLeanコンパイラをCPU上でModalのサンドボックス内で実行、調整プロセスがトレーニングループ全体を管理します。ESの方法では、わずかに異なるモデルコピーの集団を作成し、それらを評価して元のモデルを更新します。一方、GRPOはグループ内の相対的なパフォーマンスに基づいて勾配更新を行います。
実装の詳細では、AE StudioはModalの特徴を最大限に活用しました。各計算ロールに独立したイメージを使用し、GPUワーカーは推論、Leanサンドボックスは検証、オーケストレーターは軽量に保ちました。並列GPUファンアウトはModalの.map()を使用し、各摂動評価は現在のチェックポイント、定理バッチ、1つの摂動シード、生成パラメータのみを必要とします。ESでは重み摂動がシードによって完全に決定されるため、更新ステップはシードと報酬のみで集団全体を集約でき、GPU間の高価な重み転送を回避できます。各ワーカーは完全なシード/報酬履歴を受け取り、基本重みから現在のモデルを再構築してから自身の摂動を適用します。
検証部分では分離が最も重要です。証明の試行はハングしたりチェッカーをクラッシュさせたりする可能性があるため、Modalのサンドボックスを使用して各検証バッチごとに独立したLeanサーバーを起動し、処理後にシャットダウンしました。1回の反復で3,840件の証明試行が生成され、64件ずつのバッチに分割されて並列検証されました。
ESの優れた特性として、モデル状態全体が基本モデルと（シード、報酬）ペアのリストで記述できる点があります。オーケストレーターはシード/報酬エントリの実行リストを維持し、各ワーカーに渡します。GPU側では、各ワーカーがボリュームから元の基本モデルをロードし、完全な履歴を再生して現在の重みを再構築してから自身の摂動を適用します。再生は過去の反復のシードと報酬を順にたどり、決定論的シードを使用してGPU上で同じノイズを再生成し、重み付き更新をその場で適用します。基本モデルはボリュームに保存されるため、ワーカーは毎回Hugging Faceから再ダウンロードする必要がありません。
パフォーマンス面では、Modalの実装はわずか約250行のプラットフォーム設定コードで済み、他のプラットフォームの典型的な600行の半分以下でした。これにより、プロジェクト開始から2日以内にトレーニング実行が完了し、代替プラットフォームより60%高速でした。ランタイム効率も大幅に向上し、生成ステップ（GPU必要）は平均147秒でしたが、検証ステップ（CPUのみ）のために完全ループは538秒かかりました。Modalの弾力性により、GPUが実際に使用されている時間のみ支払うため、無駄なGPU時間は弾力性の低いプラットフォームと比較して約3.7倍削減されました。これらのタイミングに基づき、Modalでの全実行コストは122ドルと推定され、同じハードウェアを使用した代替ソリューションでは180〜480ドルかかるとされています。
初期の結果は有望でした。複数の定理証明実行において、ESは反復ごとの検証済み証明数でGRPOベースラインに匹敵またはそれを上回り、特にトレーニングデータが限られている場合にサンプル効率が高いことが示されました。ただし、他の設定では利得は小さく、ハイパーパラメータの選択、データセットの微妙な違い、または集団サイズなどの要因による変動の程度はまだ分離できていません。この領域でのESのスケーリング動作はまだ十分に理解されておらず、一貫して優位性を発揮する条件や、大規模トレーニング設定でGRPOと競合できるかどうかを確認するには、さらなる実験が必要です。
研究上の疑問はまだ解決していませんが、Modalを使用したインフラストラクチャは今後も継続して使用し、検証器をバックエンドに持つトレーニングループを実行、検査、反復するために、独自のインフラを構築する時間を費やすことなく活用していく予定です。次の実験は、トレーニングループ、検証パス、チェックポイントの受け渡しが機能しているため、より簡単になっています。分散分析（シグマ、集団サイズ、定理選択の体系的なスイープ）と、より大きなモデル（オープンな7B蒸留Kiminaモデル）のテストが計画されています。この時点で、未解決の質問は主に実験設計に関するものであり、システムの信頼性に関するものではありません。
このインフラストラクチャ設定は定理証明に固有のものではありません。多くのMLシステムで同じ形状が見られます：GPUで出力を生成し、外部検証器、コンパイラ、またはテストハーネスに対して実行し、それらの結果をトレーニング信号に変換し、多くの独立した候補に対して繰り返します。実験の重要な技術的側面の1つは、GPU推論の並列化、検証の分離、トレーニング状態の移植性をすべて1つのワークフローで実現したことです。.map()が証明生成を分散し、サンドボックスが検証器の障害を封じ込め、ボリュームが反復間でチェックポイントを利用可能に保ち、ローカルマシンに戻す必要をなくしました。
Modalを使用することで、チームは3つの非常に異なるランタイム間で同じスパース報酬強化学習ワークフローを、毎回セットアップを再構築することなく実行できました。ローカルマシンから実行を開始し、基本モデルをリモートに保持し、MLflowにログ記録し、実験の進行に伴って集団サイズ、定理バッチサイズ、検証並列度を変更しました。Modalプラットフォームの機能と使いやすさのおかげで、プロジェクト開始から数日以内に実験結果を見ることができました。他のインフラプラットフォームでは、何らかのシグナルを得るまでに数週間のテストと反復が必要だったでしょう。一度セットアップが完了すれば、チームはシステムの信頼性を心配することなく、実験設計に完全に集中できました。
実験の詳細を掘り下げたり、自身のワークフローに適応させたい場合は、完全な実験コードがGitHub（github.com/agencyenterprise/modal-rl-theorem-case-study）で公開されています。