技術要約: A2D – 拡散言語モデルのための任意の順序、任意のステップにおける安全性アライメント

1. 問題提起

拡散大規模言語モデル（dLLM）は、決定的な左から右へのシーケンスではなく、任意の位置で反復的にトークンを生成する**任意の順序デコーディング（any-order decoding）**を可能にすることで、自己回帰（AR）モデルとは大きく異なる性質を持っています。この柔軟性は複雑な推論や構造化されたコンテンツ合成をサポートしますが、従来の安全性アライメント手法を無効にする独自の安全性の脆弱性を導入します。

特定された核心的な問題は、**浅いアライメント（shallow alignment）**です。ARモデルと同様に、現在のdLLMは生成の初期ステップにおいてのみ強い安全性信号を示し、デコーディングが進むにつれてその信号が急速に減衰するという特性があります。この制限は、以下のような広大な攻撃表面を生み出します：

• 任意の順序に対する脆弱性（Any-Order Vulnerability）: 有害なコンテンツがデコーディングの軌道上の任意の位置に注入されることで、初期の拒絶を回避できてしまいます。
• 任意のステップに対する脆弱性（Any-Step Vulnerability）: DIJA (Wen et al., 2025) のようなテンプレートベースの攻撃は、部分的にプリフィルされたプロンプト内に有害なテキストと [MASK] トークンを交互に配置することで、この減衰を悪用します。これらの攻撃は、安全性信号がすでに薄れてしまった数ステップ後のデコーディング段階を標的とし、レスポンスレベルのアライメントでは防げない不適切な補完を誘発します。

RT（拒絶トレーニング）、SFT（教師あり微調整）、VRPO（分散減少型選好最適化）を含む既存のベースラインは、レスポンスレベルの監督に依存しています。その結果、多様なデコーディング戦略（ランダム、信頼度ベースなど）の下で安全性を維持することに苦戦し、初期トークンをバイパスする攻撃を防ぐことができません。

2. 手法: A2D (Any-Order, Any-Step Defense)

著者らは、dLLMの柔軟なデコーディングプロセスに特化して設計されたトークンレベルのアライメント手法であるA2Dを提案しています。フルレスポンスレベルでのみ拒絶を監督する従来のアプローチとは異なり、A2Dは、有害なコンテンツが任意のマスク位置で検出された場合に、特定の [EOS] (End of Sequence) トークンを出力するようにモデルをアライメントします。

コアメカニズム

• トークンレベルの監督: 学習中、ターゲットシーケンスに有害なスパンが存在する場合、モデルはマスクされたトークンを元の有害なトークンではなく、[EOS] に置き換えるよう監督されます。
• 学習目的関数: この手法は、修正されたターゲット分布を用いた標準的なマスク拡散目的関数を利用します：
  • 有害セット ($D_{harm}$): 有害なレスポンス内のマスクされたトークンは、[EOS] を予測するように監督されます。これにより、部分的なコンテキストに関わらず、モデルが有害な継続を停止することを学習します。
  • 保持セット ($D_{retain}$): 安全なレスポンス（または有害なプロンプトに対する安全な補完）内のマスクされたトークンは、元のトークンを再構成するように監督され、一般的な有用性を維持します。
• ランダム化されたマスキング: 学習プロセスでは、破損率 $\lambda$ を一様にサンプリングし、モデルを初期および後半の両方のデコーディング段階にさらします。これにより、アライメントが生成軌道全体にわたって堅牢であることを保証し、「浅いアライメント」現象を防ぎます。
• 拒絶トークン: [EOS] トークンが選択された理由は、それがdLLMにおいて標準的なパディングおよび終了マーカーであり、性能を低下させる可能性のある分布外（OOD）トークンの導入を必要とせずに、モデルがその使用に慣れているためです。

リアルタイム・モニタリングと早期拒絶

A2Dは本質的にリアルタイムの安全性信号を提供します。モデルは有害なコンテンツが存在する場合に [EOS] への高い確率を割り当てるよう訓練されているため、この確率質量はオンラインの安全性指標として機能します。

• 動的な介入: モデルはレスポンスを開始できますが、不適切な継続が現れた場合には自動的に終了します。
• 早期拒絶: 最初のデコーディングステップの最も左側のマスク位置における [EOS] 確率を監視することで、システムは閾値ベースの早期拒絶メカニメントを実装できます。確率が閾値 $\tau$ を超えた場合、モデルは有害なプロンプトに対して一切のトークンを生成することなく即座に停止し、推論レイテンシを大幅に削減します。

3. 主な貢献

1. 任意の順序/任意のステップにおける脆弱性の特定: 本論文は、dLLMがデコーディングステップを通じて急速に減衰する「浅いアライメント」に苦しみ、デコーディング順序に対しても堅牢ではないことを示し、DIJAのような攻撃に対して脆弱であることを実証しました。
2. トークンレベルのアライメントフレームワーク: A2Dは、レスポンスレベルの拒絶から [EOS] を用いたトークンレベルの抑制へとパラダイムを転換し、生成プロセスのあらゆるステップにおける堅牢な防御を実現する新しいアライメント手法を導入しました。
3. リアルタイム安全性モニタリング: 本手法は [EOS] 確率を内部の安全性信号として活用し、動的な介入と有害なプロンプトの早期拒絶を可能にします。
4. 包括的な評価: 著者らは、有害なレスポンスの中に単一のマスクされた文章が埋め込まれた極限状態でのアライメントをストレス・テストするために、FITS (Fill-in-the-Sentence) 攻撃を導入しました。

4. 実験結果

著者らは、3つのインストラクション・チューニング済みdLLM（LLaDA-8B-Instruct、LLaDA-1.5、Dream-v0-Instruct-7B）を用いて、RT、SFT、VRPOのベースラインと比較しながらA2Dを評価しました。

• 脱獄（Jailbreak）に対する防御: A2Dは、すべてのベンチマーク（Zeroshot, PAIR, ReNeLLM, Prefilling, DIJA）において、一貫して最も低い平均攻撃成功率（ASR）を達成しました。
  • DIJAの性能: A2Dは、ベースラインモデルにおけるDIJAの成功率を80%以上に抑え込んでいたところから、LLaDA-8Bでは 1.3%、LLaDA-1.5では 3.5%、Dreamでは 0.0% まで低減させました。
  • FITS攻撃: 極端なFITS攻撃（有害なレスポンスから1文だけが欠落している状態）において、A2Dは全モデルでASRをほぼ0%に抑えましたが、他の手法は60%を超えるASRで失敗しました。
• デコーディング戦略への堅牢性: A2Dは、5つのデコーディング戦略（左から右、右から左、信頼度、エントロピー、ランダム）において低いASRを維持しましたが、ベースラインモデルは特定の順序に対して高い脆弱性を示しました。
• 有用性の維持: A2Dは過剰な拒絶を回避し、XSTest ベンチマーク（有害に見える無害なプロンプト）において 0% の偽陽性 を達成しました。また、大幅な能力低下を招く他のアライメント手法を上回り、一般的な能力ベンチマーク（MMLU, GSM8K, HumanEval）において競争力のある性能を維持しました。
• 効率性の向上: 早期拒絶メカニズムにより、有害なプロンプトに対して（$\tau=0.8$ のとき）最大 19.3倍高速 な安全な終了を実現しつつ、高いアライメントの忠実度を維持しました。

5. 重要性と主張

本論文は、dLLMの柔軟なデコーディングと、従来の厳格なレスポンスレベルのアライメントとの間の根本的なミスマッチに対処することで、A2Dが拡散ベースのテキスト生成における安全性の新基準を確立すると主張しています。

• 深いアライメント: A2Dは「深いアライメント」を実現し、安全性信号が生成の開始時だけでなく、全生成軌道を通じて持続することを保証します。
• 順序への依存性がない: デコーディング戦略に関わらず堅牢な防御を提供します。これは、dLLMの実用的な展開において重要な要件です。
• 実用的な展開: リアルタイム・モニタリングと早期拒絶を可能にすることで、A2Dはレイテンシと安全性が極めて重要となる現実世界のアプリケーションにおいて、dLLMを安全に展開するための道筋を示しています。

著者らは、このアプローチがdLLMの生成の柔軟性と、現実世界での使用に求められる安全性の制約との間の溝を埋め、信頼性の高い拡散ベースの言語システムの実現に向けた道を切り開くものであると結論付けています。