Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models

要約

拡散言語モデルの推論フレームワークであるラテラル思考の拡散チェーン（DCOLT）}を紹介します。
DCOLTは、逆拡散プロセスの各中間ステップを潜在的な「思考」アクションとして扱い、結果ベースの強化学習（RL）で最終回答の正しさの報酬を最大化するために推論軌跡全体を最適化します。
因果関係のある線形思考プロセスに従う従来のチェーン（COT）方法とは異なり、DCOLTは、中間の思考のステップの中で文法的正しさに関する厳格なルールなしで双方向の非線形推論を可能にします。
2つの代表的な拡散言語モデル（DLMS）にDCOLTを実装します。
まず、SEDDを代表的な連続時間離散拡散モデルとして選択します。このコンクリートスコアは、中間拡散ステップのシーケンス全体にわたってRL報酬を最大化する確率的ポリシーを導き出します。
さらに、離散時間マスクされた拡散言語モデルであるLladaを検討し、トークンを予測してマスクする順序が、ランキングベースのアンマスキングポリシーモジュール（UPM）から生じるRLアクションを最適化するために不可欠な役割を果たしていることを発見しました。
数学とコード生成の両方のタスクでの実験では、パブリックデータと16 H800 GPUのみを使用して、DCOLT強化DLMがSFTまたはRL、またはその両方でトレーニングされた他のDLMを上回ることが示されています。
特に、DColt-Reinforced Lladaは、GSM8K、Math、MBPP、およびHumanevalで、推論の精度を +9.8％、 +5.7％、 +11.4％、 +19.5％増加させます。

要約(オリジナル)

We introduce the \emph{Diffusion Chain of Lateral Thought (DCoLT)}, a reasoning framework for diffusion language models. DCoLT treats each intermediate step in the reverse diffusion process as a latent ‘thinking’ action and optimizes the entire reasoning trajectory to maximize the reward on the correctness of the final answer with outcome-based Reinforcement Learning (RL). Unlike traditional Chain-of-Thought (CoT) methods that follow a causal, linear thinking process, DCoLT allows bidirectional, non-linear reasoning with no strict rule on grammatical correctness amid its intermediate steps of thought. We implement DCoLT on two representative Diffusion Language Models (DLMs). First, we choose SEDD as a representative continuous-time discrete diffusion model, where its concrete score derives a probabilistic policy to maximize the RL reward over the entire sequence of intermediate diffusion steps. We further consider the discrete-time masked diffusion language model — LLaDA, and find that the order to predict and unmask tokens plays an essential role to optimize its RL action resulting from the ranking-based Unmasking Policy Module (UPM) defined by the Plackett-Luce model. Experiments on both math and code generation tasks show that using only public data and 16 H800 GPUs, DCoLT-reinforced DLMs outperform other DLMs trained by SFT or RL or even both. Notably, DCoLT-reinforced LLaDA boosts its reasoning accuracy by +9.8%, +5.7%, +11.4%, +19.5% on GSM8K, MATH, MBPP, and HumanEval.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google