# CDLM - Consistency Diffusion Language Models for Faster Sampling > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 4 (May 21 / Thu)、Grand Ballroom 2、08:45 - 09:00 PDT > - **セッション名:** Research Track Oral -- LLM Serving 5（Moderator: Xiang Song） > - **登壇者:** Minseo Kim (Seoul National University) ※発表者は筆頭著者から推定 > - **URL:** https://mlsys.org/virtual/2026/oral/3785 > - **共著者:** Minseo Kim^1, Chenfeng Xu^{2,3}, Coleman Hooper^2, Harman Singh^2, Ben Athiwaratkun^3, Ce Zhang^3, Kurt Keutzer^2, Amir Gholami^2（^1 Seoul National University, Seoul, South Korea / ^2 University of California, Berkeley, CA, USA / ^3 Together AI, San Francisco, CA, USA）。Correspondence: Amir Gholami <[email protected]> > - **関連研究 URL:** https://github.com/SqueezeAILab/CDLM > - **OpenReview:** https://openreview.net/forum?id=eB8yjR6alL > [!abstract] 概要（論文 PDF Abstract -- 忠実日本語訳） > 拡散言語モデル（DLM）は有望な並列生成パラダイムだが、推論が遅い。原因はリファインメントステップ数の多さと、標準的な KV キャッシングが使えないことである。本研究では CDLM（Consistency Diffusion Language Models）を提案する。これは訓練ベースの高速化手法であり、両方のボトルネックを同時に解消する。CDLM はコンシステンシモデリングを統合し、必要なサンプリングステップ数を大幅に削減する。これはマルチトークンの確定を可能にし、さらにファインチューニング時にブロック単位の因果的アテンションマスクを適用することで、モデルを KV キャッシングと完全に互換にする。実験の結果、CDLM は数学・コーディングタスクで競争力のある精度を維持しつつ、3.6 倍から 14.5 倍のレイテンシ削減を達成した。訓練・評価コードは https://github.com/SqueezeAILab/CDLM で公開されている。 ## 背景と動機: DLM の 2 つのボトルネック - **DLM の利点**: 自己回帰（AR）モデルと異なり、DLM は全トークン位置を各ステップで並列に更新する。マスクされたトークン列を反復的にアンマスクし、テキスト埋め込み・コード生成など双方向文脈が有利なタスクに強い。クローズドソースモデル（Gemini Diffusion 等）は AR モデル比で最大 10 倍のスループットを報告している（論文 Section 1）。 - **ボトルネック (1) キャッシュ非互換**: 標準 DLM は完全双方向アテンションを使うため、すべてのデノイジングステップで $O(L^2)$ のアテンション再計算が必要となり、標準的な KV キャッシュが適用できない（スライド4）。 - **ボトルネック (2) 過剰なリファインメントステップ**: 高品質な生成には系列長に匹敵する数のデノイジングステップが必要であり、推論が非効率になる。オープンソース DLM（LLaDA-8B-Instruct 等）は同規模 AR モデル（LLaMA-3-8B-Instruct）より大幅に遅い（スライド3: プロンプト長 128、生成長 256 で LLaDA は約 26,000 ms に対し LLaMA-3 は約 7,500 ms）。 ## CDLM の全体設計 CDLM は**訓練ベース**の手法であり、コンシステンシモデリング（consistency modeling）を DLM に導入することでステップ削減とキャッシュ互換を同時に実現する（スライド5）。 全体パイプライン: 1. **教師 DLM** が双方向アテンションでブロック単位のデコーディング軌跡（trajectory）を生成 2. 軌跡と隠れ状態を $(y, y^*, \mathbf{H})$ の教師信号ペアとして収集 3. **ブロック因果的（block-causal）学生モデル** を 3 つの損失関数で訓練 4. 推論時は厳密な KV キャッシュ + 信頼度閾値処理による並列デコーディングを実行 ### ブロック因果的アテンションマスク - 教師 DLM は完全双方向アテンションを使用するが、学生 DLM はブロック単位の因果マスク（block-wise causal mask）で訓練される（論文 Figure 2、スライド5）。 - プロンプトおよび確定済みの過去ブロックに対してはアテンションが許可され、現在のデコーディングブロック内ではブロック内双方向アテンションが保持される。 - この設計により、確定済みブロックの KV キャッシュを再利用でき、ブロック境界での早期終了（early stopping）も可能になる（論文 Section 4.1）。 ## 軌跡収集（Trajectory Collection） - 教師 DLM（双方向アテンション）をブロック単位デコーディングで実行し、各ステップで最も信頼度の高い 1 トークンを確定する（$N = L_g$ ステップ、$B = 32$ トークン/ブロック）（論文 Algorithm 1、スライド7）。 - 各トークン確定時の最終層隠れ状態を $d \times L_g$ のバッファに記録し、ホワイトボックス蒸留に使用する。ロジットの直接保存（語彙サイズ $|V|$）と比較して約 30 倍のストレージ削減を実現する（論文 Appendix A.1）。 - **温度拡張（temperature augmentation）**: プロンプトごとに複数の温度 $\tau \in \{0.0, 0.5\}$ で軌跡を生成し、多様性を確保する。$\tau = 1.0$ は推論チェーンを不安定化させるため除外する（論文 Appendix A.1、スライド7）。 ## 訓練: 3 つの損失関数 CDLM は以下の 3 つの損失を同時最小化する（論文 Section 4.2、Eq. 7）: $\mathcal{L}(\phi) = w_{\text{distill}} \, \mathcal{L}_{\text{Distillation}} + w_{\text{cons}} \, \mathcal{L}_{\text{Consistency}} + w_{\text{dlm}} \, \mathcal{L}_{\text{DLM}}$ ### (a) 蒸留損失（Distillation Loss） - 軌跡上の中間状態 $y$ とそのブロック完了状態 $y^*$ の間で**新たにアンマスクされた位置** $\mathcal{U}_y$ に対し、教師の分布と学生の分布の間の順方向 KL ダイバージェンスを計算する（論文 Eq. 4、スライド8）。 - 教師の隠れ状態バッファ $\mathbf{H}_x$ から `lm_head` を適用して教師分布を再構築する。 - **直観**: 学生に複数トークンの同時確定を教える。蒸留損失は訓練のアンカーとして機能し、単独でも高速に収束する（論文 Table 3 第1行）。 ### (b) コンシステンシ損失（Consistency Loss） - 同一軌跡上の中間状態 $y$ とブロック完了状態 $y^*$ の間で**まだマスクされている位置** $\mathcal{S}_y$ に対し、学生自身の予測分布間の順方向 KL ダイバージェンスを最小化する。ターゲットはストップグラディエント（stop-gradient）で切り離す（論文 Eq. 5、スライド9）。 - **直観**: 情報量の少ない状態の予測を、より情報量の多い状態の予測に揃える。これによりデコーディング軌跡上の安定したマルチステップジャンプが促進される。 - コンシステンシ損失**単独**では崩壊する（教師監督なしの自己一致性最適化は失敗する）。蒸留損失と結合することで相乗効果が生まれ、収束が高速かつ安定になる（論文 Table 3、スライド15）。 ### (c) DLM 損失 - 正解テキスト $\hat{y}$ をランダムにマスクし、マスクされたトークンの対数尤度を最大化する標準的なマスク復元損失（論文 Eq. 6、スライド10）。 - **直観**: 元のマスク予測能力を保存する補助損失。これを除去すると HumanEval は上昇するが GSM8K は低下するため、タスク能力の維持に寄与する（論文 Table 3 第4--6行）。 ### 損失重みの設定 - CDLM-Dream: $(w_{\text{distill}}, w_{\text{cons}}, w_{\text{dlm}}) = (1.0, 0.5, 0.01)$ - CDLM-LLaDA: $(w_{\text{distill}}, w_{\text{cons}}, w_{\text{dlm}}) = (1.0, 0.5, 0.1)$（LLaDA は DLM 損失の絶対スケールが小さいため $w_{\text{dlm}}$ を大きく設定） ## 推論 - ブロック因果マスクの下でブロック単位にデコードし、プロンプトおよび確定済みブロックの KV キャッシュを再利用する（スライド11）。 - **信頼度閾値処理による並列確定（confidence-thresholded parallel finalization）**: 各ステップで信頼度が閾値 $\tau_{\text{conf}}$ を超えるトークンを一括確定する。既定 $\tau_{\text{conf}} = 0.90$（論文 Section 4.3）。 - **早期終了**: 現在のブロック内で `<endoftext>` トークンが生成されると、そのブロックのデコードを終了する。 - ブロック間並列化（inter-block parallelism）等の追加ヒューリスティクスは意図的に避けている（タスク・ドメイン依存のハイパーパラメータを増やさないため）。 ## 実験設定 - **対象モデル**: Dream-7B-Instruct（Ye et al., 2025）と LLaDA-8B-Instruct（Nie et al., 2025b）の 2 つのオープンソース DLM。 - **訓練データ**: Bespoke-Stratos-17k（Labs, 2025）から 7.5k プロンプトを抽出（プロンプト長 512 以下にフィルタ）。LLaDA 向けには DParallel（Chen et al., 2025）から 7.5k の数学系プロンプトを追加。教師の回答は Qwen2.5-7B で生成。 - **ファインチューニング**: LoRA（Hu et al., 2022）をアテンションおよび MLP モジュールに適用。ブロック因果マスク、$B = 32$、$L_g = 256$。 - **訓練時間**: CDLM-Dream は約 **8 時間**、CDLM-LLaDA は約 **16 時間**（4x NVIDIA A100 80GB GPU）。 - **ベンチマーク**: GSM8K、GSM8K-CoT、MATH（数学的推論）、HumanEval、MBPP（コード生成）。すべて greedy decoding（温度 0.0）、$L_g = 256$、$B = 32$、$\tau_{\text{conf}} = 0.9$。 - **計測環境**: 4x NVIDIA A100 (80 GB) GPU、バッチサイズ 1、データ並列。レイテンシ・ステップ数・生成長はサンプル平均。 ## 主要結果 ### Dream-7B-Instruct に対する結果（論文 Table 1） | ベンチマーク | 手法 | TPS | レイテンシ (s) | ステップ数 | スコア | |---|---|---|---|---|---| | GSM8K-CoT (8-shot) | Dream-7B-Instruct | 4.1 (x1.0) | 23.5 (x1.0) | 256.0 (x1.0) | 79.1 | | | Fast-dLLM (Par.+D.C.) | 36.6 (x8.9) | 2.5 (x9.4) | 60.8 (x4.2) | 77.3 | | | **CDLM-Dream** | **51.7 (x12.6)** | **2.1 (x11.2)** | **44.1 (x5.8)** | **78.8** | | HumanEval-Instruct (0-shot) | Dream-7B-Instruct | 15.5 (x1.0) | 13.4 (x1.0) | 256.0 (x1.0) | 48.2 | | | Fast-dLLM (Par.+D.C.) | 79.9 (x5.2) | 2.5 (x5.4) | 71.6 (x3.6) | 46.3 | | | **CDLM-Dream** | **43.3 (x2.8)** | **2.2 (x6.1)** | **49.6 (x5.2)** | **50.0** | | MATH (4-shot) | Dream-7B-Instruct | 8.3 (x1.0) | 21.9 (x1.0) | 256.0 (x1.0) | 38.0 | | | **CDLM-Dream** | **53.8 (x6.5)** | **2.9 (x7.6)** | **63.2 (x4.1)** | **32.4** | | MBPP-Instruct (0-shot) | Dream-7B-Instruct | 2.3 (x1.0) | 21.7 (x1.0) | 256.0 (x1.0) | 51.8 | | | **CDLM-Dream** | **48.1 (x20.9)** | **1.5 (x14.5)** | **33.2 (x7.7)** | **53.0** | ### LLaDA-8B-Instruct に対する結果（論文 Table 2） | ベンチマーク | 手法 | TPS | レイテンシ (s) | ステップ数 | スコア | |---|---|---|---|---|---| | GSM8K (4-shot) | LLaDA-8B-Instruct | 8.2 (x1.0) | 28.3 (x1.0) | 256.0 (x1.0) | 77.1 | | | **CDLM-LLaDA** | **54.3 (x6.6)** | **3.3 (x8.6)** | **57.7 (x4.4)** | **73.9** | | HumanEval (0-shot) | LLaDA-8B-Instruct | 7.4 (x1.0) | 11.3 (x1.0) | 256.0 (x1.0) | 37.8 | | | **CDLM-LLaDA** | **50.9 (x6.9)** | **1.9 (x5.9)** | **32.3 (x7.9)** | **40.2** | | MATH (4-shot) | LLaDA-8B-Instruct | 8.8 (x1.0) | 25.7 (x1.0) | 256.0 (x1.0) | 24.1 | | | **CDLM-LLaDA** | **50.2 (x5.7)** | **4.2 (x6.1)** | **75.3 (x3.4)** | **28.3** | | MBPP-Instruct (0-shot) | LLaDA-8B-Instruct | 17.7 (x1.0) | 11.4 (x1.0) | 256.0 (x1.0) | 40.8 | | | **CDLM-LLaDA** | **60.6 (x3.4)** | **3.2 (x3.6)** | **58.0 (x4.4)** | **38.4** | ### 主要な知見 - **レイテンシ削減**: CDLM は全ベンチマークで最大のレイテンシ削減を達成する。CDLM-Dream は GSM8K-CoT で最大 11.2 倍、MBPP-Instruct で最大 14.5 倍のレイテンシ削減。ステップ削減とブロック単位 KV キャッシングの組み合わせが効果的。 - **ステップ削減**: CDLM-Dream はリファインメントステップを約 4.1--7.7 倍削減する。素朴なステップ切り詰め（naive truncation）は品質を著しく劣化させるのに対し、CDLM はコンシステンシ訓練により同等ステップ数で品質を維持する（論文 Table 4: Dream を 48 ステップに切り詰めるとスコア 41.8 だが CDLM-Dream は 44.1 ステップで 78.8）（スライド14）。 - **精度**: CDLM の精度は DLM バックボーンの能力に制約される。CDLM-Dream は GSM8K-CoT で 78.8（AR ベースライン Qwen2.5: 42.6 の TPS を上回る 51.7 TPS）、MBPP-Instruct で 53.0（ナイーブ Dream の 51.8 を超える）。 ### AR モデルとのスループット比較（論文 Figure 3、スライド13） - CDLM はナイーブ DLM のスループットを 3--21 倍向上させ、同規模 AR ベースラインをも上回る。 - CDLM-Dream は GSM8K-CoT で Qwen2.5-7B-Instruct（42.6 TPS）比 1.2 倍の 51.7 TPS、MBPP-Instruct で Dream（2.3 TPS）比 48.1 TPS。 - CDLM-LLaDA は HumanEval で LLaMA-3.1-8B-Instruct（13.2 TPS）比 4.2 倍の 55.1 TPS。 - 各 CDLM ステップは行列積演算のため AR ステップより計算コストが高いが、1 ステップで複数トークンを確定するため（例: CDLM-Dream は GSM8K-CoT で 1 ステップあたり約 2.4 トークン）、総スループットで AR を凌駕する。 ## アブレーション ### 損失重みの構成（論文 Table 3、スライド15） - **蒸留のみ**（$w_{\text{cons}} = 0$）: 約 46.7 ステップで高速に収束するがスコアは弱い（GSM8K 73.2）。 - **コンシステンシのみ**（$w_{\text{distill}} = 0$）: 崩壊する（スコア 0.0、ステップ 124.3）。教師監督なしの自己一致性最適化は不安定。 - **蒸留 + コンシステンシ**（$w_{\text{dlm}} = 0$）: 良好かつ高速だが、DLM 損失なしでは一部タスクで不安定。 - **3 損失すべて**: 最も強力で安定（GSM8K 75.1、HumanEval 45.7）。 ### ステップ削減の有効性（論文 Table 4、スライド14） - ナイーブにステップ数を同数に揃えた場合（Dream-7B-Instruct を 48 ステップ）: レイテンシ 4.4 s、スコア 41.8。 - CDLM-Dream（44.1 ステップ）: レイテンシ 2.1 s、スコア 78.8。 - コンシステンシ訓練により安定したマルチトークン確定が可能になり、ステップ削減が品質を損なわない。さらに KV キャッシングにより同等ステップ数でもレイテンシが約半減する。 ### 信頼度閾値 $\tau_{\text{conf}}$ の感度（論文 Table 7） - $\tau_{\text{conf}} = 0.95$: 保守的、TPS 低下（GSM8K-CoT: 42.7 TPS / 2.5 s / 78.8）。 - $\tau_{\text{conf}} = 0.90$: バランスが良い（GSM8K-CoT: 51.7 TPS / 2.1 s / 78.8）。 - $\tau_{\text{conf}} = 0.85$: 攻撃的、TPS は最大だがスコアが若干低下（GSM8K-CoT: 57.7 TPS / 1.9 s / 78.4）。 ## システムレベルのスケーラビリティ分析（Arithmetic Intensity） 論文 Section 5.4 および Figure 4（スライド16）では、AR モデル・ナイーブ DLM・ブロック因果的 DLM（CDLM）の arithmetic intensity（AI; 演算量対メモリトラフィック比）をルーフラインモデルで比較している。 - **AR モデル**: バッチサイズ 1 で AI が 1 近辺と強くメモリ律速。バッチサイズ 128 でも AI 71.3 でリッジポイント以下にとどまる。 - **ナイーブ DLM（完全アテンション）**: バッチサイズ 1 の時点で AI 438.9 とリッジポイントを大幅に超え、計算律速に飽和。バッチサイズ増加の恩恵は限定的。 - **ブロック因果 DLM（CDLM）**: バッチサイズ 1 で AI 4.0--31.1（$B \in \{4, 16, 32\}$）と、AR と DLM の中間に位置する均衡的な領域を占める。ブロック内の並列処理により低バッチ設定でも AR より高い演算利用率を達成しつつ、即座に計算飽和しない。$B = 32$ の場合、バッチサイズ約 8 でリッジポイントを超える。 この分析は、CDLM がメモリ律速な AR モデルと計算飽和したナイーブ DLM の中間に位置し、低バッチでの高スループットとスケーラビリティを両立する設計であることを示している。 ## 議論と限界 - **表現力 vs 効率のトレードオフ**: 完全双方向アテンションの $O(L^2)$ を厳密な KV キャッシュで緩和するが、ブロック内局所的なリファインメント能力（インフィリング等）は保持される。左から右への弱い帰納バイアスは推論スタイルの生成と整合的とみなされる（論文 Section 6）。 - **推論専用高速化との直交性**: CDLM は推論専用手法（D2F のブロック間並列化、dLLM-Cache の適応的キャッシング等）と組み合わせ可能。現在はブロック単位 KV キャッシングと信頼度閾値のみを使用しており、さらなるレイテンシ削減の余地がある。 - **教師の上限**: ブロック因果学生に蒸留された DLM は教師の知識を超えられない。より強力な AR 教師からの蒸留はこの上限を引き上げる自然な方向である（論文 Appendix C）。 - **限界 -- オフライン静的軌跡**: 訓練は現在オフラインの静的軌跡に依存する。学生がオンラインで軌跡を生成し教師が検証するオンポリシー学習は、訓練・推論間の不一致を低減する可能性があるが、DLM の生成速度が制約となる。 - **限界 -- データ規模**: 訓練データは約 15k サンプルと小規模で、主に数学系に偏る。MATH での精度低下（Dream: 38.0 -> 32.4）は訓練データの難易度範囲の狭さと 256 トークンの生成長制限に帰せられる。データの多様化・スケーリングにより改善が見込まれる。 ## 今後の方向性（スライド17） - 軌跡コーパスを広いタスクカバレッジにスケールする - より強力な DLM 教師からの蒸留 - より長い生成バジェットへの拡張 ## 関連研究の位置づけ - **コンシステンシモデル（Song et al., ICML 2023）**: 連続拡散モデルの ODE 軌跡上で隣接状態の予測を一致させるフレームワーク。CDLM はこれを離散トークン空間の DLM デノイジング軌跡に適応させた経験的一般化。 - **ブロック因果 DLM（Wang et al., 2025; Wu et al., 2025a）**: 事前学習済み DLM にブロック単位の因果性を導入して KV キャッシュを有効化する先行研究。CDLM はこれに加えてコンシステンシ誘導蒸留によりステップ数の削減も達成。 - **Fast-dLLM（Wu et al., 2025b）**: 信頼度閾値による並列デコーディングと近似デュアルキャッシュ KV を組み合わせた推論専用手法。CDLM は訓練ベースであり、直交的に組み合わせ可能。 - **知識蒸留**: CDLM はホワイトボックス自己蒸留（同一サイズ・同一タイプだがアーキテクチャが異なるモデル間）と見なせる。Deschenaux & Gulcehre (2025) や Ma et al. (2025b) と同様に、マルチステップ軌跡を非連続状態間のジャンプに蒸留するが、ブロック因果学生でのコンシステンシ目的関数の適用が新規。