コード生成AIを「正しい」だけでなく「速い」プログラムへ導く強化学習

導入

コード生成モデルの強化学習は、これまで主に「正しいコードを書けるか」を対象にしてきた。モデルがプログラムを生成し、隠れテストを通れば報酬を与える、という仕組みである。しかし次の目標である「正しく、かつ速いコード」を学ばせるには、単に実行時間を報酬に足すだけでは足りない。論文「Reinforcement Learning for Code Optimization」は、実行時間を学習信号にすると、計測ノイズ、報酬の疎さ、GRPO の不安定性が容易に信号を壊してしまうと指摘する。

重要なポイント

• 速度報酬は扱いが難しい。 実行時間は入力規模、サンドボックス環境、計測誤差に左右される。制御されていない計測では、モデルが本当のアルゴリズム改善ではなく偶然の揺らぎを学ぶ危険がある。
• DMC-Optim を構築。 著者らは、大きな最適化テストと校正済みの実行サンドボックスを備えたベンチマークを用意し、速度差をより信頼できる形で評価しようとしている。
• 正しさと速度を同時に報酬化。 強化学習環境では、正解性と実行速度を組み合わせる。さらにオフラインシミュレータを使い、有望な設定を事前に予測することで、試行錯誤の不安定さを抑えている。
• GRPO の調整。 timed execution は通常の合否判定より疎でノイズが大きい。そのため、学習と評価の方法もこの条件に合わせて変更されている。

結果と意味

DMC-Optim では、最も強い最適化対応設定により、Qwen 2.5 7B の厳格な top-50% pass@1 が 18.0% から 31.3% に、CWM 32B が 30.7% から 50.4% に向上した。より厳しい top-30% では、CWM 32B の相対改善は 125% に達し、純粋な正解性能は維持された。タイミング用サンドボックスの品質を落とした条件でも、ロバストな最適化 RL は標準 RLVR に対して評価基準に応じて 100% から 200% の改善を示した。

LCB では、CWM 32B が標準 RLVR に対し、中央値サンプルの速度比較で最大 83% 勝利している。また、問題ごとの最速の正解人間提出と比べると、計算量クラス改善の頻度は人間の約半分、14% 対 28% に到達した。

この研究の意義は、コード生成AIの目標を「動くコード」から「効率よく動くコード」へ広げた点にある。将来のコーディング支援は、関数を補完するだけでなく、データ構造や計算量を意識した書き換えまで担う可能性がある。一方で、速度最適化の強化学習には、信頼できるテスト、安定した計測、慎重な報酬設計が不可欠であることも示している。

出典：Hugging Face Daily Papers