SLIM-RL: 軌跡スライシング不要のリスク予算付きランダムマスキングRL for 拡散LLM
SLIM-RLは、拡散大規模言語モデル(dLLM)のための新しい強化学習手法です。τバジェットデコーダを用いて各ロールアウトステップのコミットリスクを制限し、軌跡の再構築をせずに総リスクを低減します。分散低減手法を組み込んだトレースフリーのランダムマスキング目的関数を使用し、SDAR-4BにおいてTraceRLの最高MATH500精度を0.46倍の訓練サンプルで達成し、数学・コードベンチマークでTraceRLを上回ります。
近年、拡散大規模言語モデル(dLLM)は自然言語処理分野で大きな可能性を示していますが、その強化学習訓練には効率性の課題があります。現在の最先端手法であるTraceRLは、ランダムマスキングがモデルの推論軌跡と一致しないと主張し、訓練時に各ロールアウトを最大K/s個の軌跡整合サンプルにスライスすることで軌跡を再構築しており、そのコストはブロックサイズKに比例して増大します。この問題に対処するため、複数の研究機関の研究者らはSLIM-RLという新しい手法を提案しました。これは軌跡を再構築せずに不一致を緩和し、訓練効率を大幅に向上させます。
SLIM-RLの核心は、τバジェットデコーダを用いて各ロールアウトステップのコミットリスクを制限し、訓練データ全体の総リスクを低減することです。最適化時には、トレースフリーのランダムマスキング目的関数を使用し、シーケンスレベルの重要度サンプリング、マスキングレベルに関する決定論的求積法、そして平均保存かつ単調減少するブロック単位のマスクスケジュールを導入します。これらの設計により、性能を維持しつつ必要な訓練サンプルを大幅に削減できます。
実験結果は、SDAR-4Bモデルにおいて、ブロックサイズ16でSLIM-RLがTraceRLの最高MATH500精度をわずか0.46倍の訓練サンプルで達成することを示しています。動的サンプリングを一致させた条件下では、MATH500で6.32%、GSM8Kで11.05%の改善を達成しました。ブロックサイズ4では、4BパラメータのSLIM-RLが数学タスクにおいて、より大規模なLLaDA-8BやDream-7Bを上回り、MATH500でLLaDA-8Bを10.76%上回りつつ、自己回帰モデルQwen2.5-7Bには及びませんでした。コードタスクでは、MBPPで4.20%、HumanEvalで3.65%の改善を記録しました。さらに、τバジェットデコーダはLLaDA、Dream、SDAR間で訓練不要で転送可能であり、優れた汎化能力を示しています。
本研究はarXivに発表され、ソースコードはGitHub(https://github.com/laolaorkkkkk/SLIM-RL)で公開されています。SLIM-RLは拡散大規模言語モデルの強化学習において、より効率的で経済的な代替手段を提供し、この分野の実用化とさらなる発展に貢献することが期待されます。