DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Diffusion Language Models
Abstractの概要
本論文は、拡散言語モデル向けのプラグイン型トークン順序付けモジュールであるDPRM（Doob h変換プロセス報酬モデル）を提案する。DPRMはホストモデルのアーキテクチャ、デノイジング目的関数、および教師信号を変更せず、順序付けポリシーのみを修正する。DPRMは信頼度ベースの漸進的順序付け（訓練と推論を整合）から開始し、オンラインバケット化報酬推定とショートリストベースのSoft-BoN再重み付けを用いて、Doob h変換に着想を得たプロセス報酬誘導型の順序付けへと段階的に移行する。理論解析として、厳密なDPRMポリシーを報酬傾斜ギブス開示則として特徴付け、Soft-BoN近似のO(1/N)収束を証明し、経験的ベルンシュタインレートでのオンライン追跡保証を確立し、所定の仮定の下でのサンプル複雑度の優位性を示している。自然言語事前学習、推論事後学習、テスト時スケーリング、および科学的離散拡散タスク（タンパク質、単一細胞、分子、DNA）にまたがる7つのホスト設定での実験により、複数のベンチマークでの改善とドメイン固有のトレードオフが実証されている。
新規性
主な新規性は、ホストモデルや訓練目的関数を変更することなく既存システムに挿入可能な、Doob h変換に着想を得た拡散言語モデル向けトークン順序付けコントローラである。特徴的な点は段階的設計であり、訓練・テスト整合型の信頼度順序付けから開始し、バケット化報酬推定とSoft-BoN再重み付けによるオンラインプロセス報酬誘導へと移行する。これには、厳密な報酬傾斜ギブスポリシー、そのO(1/N) Soft-BoN近似、オンライン追跡収束、および所定の仮定下でのサンプル複雑度分離に関する理論的保証が伴う。
成果
自然言語設定では、DPRM-PUMAがGSM8K検証平均を29.34から34.27に改善し、DMPO-DPRMがMATH Hardを44.3から47.9に、Countdown Hardを29.6から33.4に改善し、DPRM-Prismが投票付きGSM8K精度を82.41から83.85に改善した（ただしNFEは609から1,071に増加）。科学的ドメインでは結果がより混在しており、DPRM-DCMはトークン回復率（63.97%から75.92%）およびゼロ発現精度（78.39%から99.90%）で大幅な改善を示す一方、タンパク質・分子・DNA実験では、順序付け対応バリアントが選択的な指標（フォワードフォールディングRMSD、リンカー妥当性、HepG2スコアなど）を改善できるものの、全ての品質指標で一様に優位とはならないことが示された。
論文の注目点
- DPRMはホストアーキテクチャ、損失関数、データパイプラインを保持しつつ順序付けコントローラのみを変更するプラグイン型トークン順序付けモジュールであり、バケットレベルの報酬推定が信頼できるようになるにつれて、信頼度ベースの順序付けからオンライン報酬誘導型の順序付けへと移行する。
- 理論的枠組みはDPRMを報酬傾斜ギブス開示則として特徴付け、段階的Soft-BoN近似のO(1/N)収束を証明し、経験的ベルンシュタインレートでのオンライン追跡を確立し（バケット粗粒化、ウォームアップ、非定常性によるバイアスを伴う）、扱いやすいが間接的にのみ検証された仮定の下でサンプル複雑度の優位性を示している。
- 7つのホスト設定にわたり、DPRMは複数の自然言語推論ベンチマークを改善し、単一細胞遺伝子発現拡散で大幅な改善をもたらす一方、タンパク質・分子・DNA実験では全指標での普遍的な改善ではなくドメイン固有のトレードオフが明らかになった。