Kimina-Prover: 大規模形式推論モデルにおけるテスト時RL検索の適用

Hugging Face は Kimina-Prover-72B をリリースしました。これは Qwen2.5-72B をベースにし、Kimi k1.5 RL パイプラインで学習された最先端の形式定理証明モデルです。テスト時強化学習 (Test-Time Reinforcement Learning, TTRL) 検索フレームワークと統合されたエラー修正機構を導入することで、miniF2F ベンチマークで記録的な 92.2% の合格率 を達成しています。

72B のフラッグシップモデルに加えて、2 つの蒸留バリアントが提供されています:Kimina-Prover-Distill-8B(Qwen3-8B ベース)と Kimina-Prover-Distill-1.7B(Qwen3-1.7B ベース)。

miniF2F における最先端性能

Kimina-Prover-72B は、複数のサンプリング予算にわたって従来の形式推論モデルを上回ります。提供されたベンチマークによると、モデルは pass@1 が 63.9%、pass@32 が 84.0%、pass@1024 が 87.7% を達成しています。完全な TTRL 検索フレームワークを適用すると、合格率は 92.2% に上昇します。

Model pass@1 pass@32 pass@1024
Kimina-Prover-1.7B 46.7 73.4
DSP+ 52.5 71.3 80.7
DeepSeek-Prover-V2-7B 58.6 75.6 79.9
Kimina-Prover-8B 61.1 78.3
DeepSeek-Prover-V2-671B 61.9 82.4 86.6
Kimina-Prover-72B 63.9 84.0 87.7

テスト時強化学習 (TTRL) 検索

長期的な推論が必要な複雑な問題を解くために、Kimina-Prover は TTRL 検索フレームワークを利用し、モデルが中間補題を自律的に発見・組み合わせ・再利用できるようにします。

Lemma-Enabled Pattern

RL 訓練中、モデルは「lemma-enabled pattern」と呼ばれる手法にさらされます。これは、問題コンテキストの先頭に 1〜3 個の形式的補題をランダムに付加するものです。モデルが実際にこれらの補題を使用することを保証するため、好みベースの報酬整形戦略を採用しています:提供された補題を活用した解答には高い報酬が与えられ、無視した解答はペナルティが課せられます。その結果、補題利用率は 30〜40% に達しました。

再帰的検索とフィルタリング

TTRL は以下のメカニズムを通じて、ランダム探索から戦略的検索へと移行します:

  • Dynamic Scoring: システムは各候補補題に対して「補題利用スコア」を追跡します。各 RL イテレーションで、入力の 60% は上位スコアの補題を使用し、残りの 40% は探索促進のためにランダム補題を使用します。
  • Pruning: 50 回の試行後に利用スコアが $\tau=0.10$ 未満の補題は除外されます。
  • Recursive Decomposition: 定理または補題が $N=128$ 回の試行で解決できない場合、システムは新たな候補サブ補題を生成し、問題をより小さなサブ問題へと再帰的に分解します。
  • Negation Filtering: 論理的に矛盾する補題の利用を防ぐため、システムは新しい補題の否定を証明しようとします。否定が証明可能であれば、その補題は破棄されます。

統合エラー修正機能

Kimina-Prover は Lean 4 のエラーメッセージを解釈し、対象を絞った修正案を提示できます。これは、証明を最初から再生成するよりもサンプル効率が高いです。

SFT データとバッチ失敗リプレイ

一般的な LLM は Lean のエラーメッセージに苦戦するため、チームは (誤った証明、Lean のフィードバック、正しい証明) の三つ組からなる特化した教師あり微調整(Supervised Fine-Tuning, SFT)データセットを作成しました。これに Claude 3.7 Sonnet が生成した推論チェーンを加えて拡張しています。

トレーニングの安定化のため、Batched Failure Replay 戦略を実装しました。失敗した試行はイテレーション $N$ で収集され、次のイテレーション $N+1$ で標準問題と混合されます。これにより、モデルはエラー修正タスクに継続的に曝露されます。

効率向上

エラー修正は、固定された計算予算下での成功率を大幅に向上させます。59 件の難易度の高い MiniF2F 問題に対するテストでは、"16+16 試行・修正" 戦略(最初の 16 回の試行+16 回の修正)が 35.6% の成功率を達成し、32 回の独立試行による "32x1 暴力的" 戦略の 28.8% を上回りました。

追加の技術的強化

  • Continuous Pre-training (CPT): モデルは 60 億トークンのデータセットで CPT を実施しました。その中には GitHub からの 2.6 億トークンと、コンパイラ検証済みロールアウトデータ 55 億トークンが含まれます。
  • Random Proof Cut Augmentation: 人手で注釈された証明を活用するため、チームは "Proof truncation"(証明の末尾を削除)と "Proof infilling"(内部ブロックを sorry に置換)を用いて、モデルに論理的ギャップの埋め方を学習させました。
  • Prompt Set Curation: 訓練セットは約 9 万件のコンペティション志向問題に絞り込まれ、動的フィルタリングで容易すぎる問題は除外し、難しすぎる問題は分解されました。
  • Non-proof Problem Solving: モデルは最終的な数値回答が求められる問題にも対応でき、まず答えを導き出し、その後形式的証明で正当性を示します。

Sources