技術的概要：DiHAL – 幾何学的に導かれた隠れ状態の置換

問題定義

連続拡散言語モデルは、歴史的に自己回帰型トランスフォーマーに遅れをとってきました。拡散モデルは反復的なノイズ除去を通じて画像生成において卓越していますが、トークンの離散性により、これをテキストに適用することは困難です。離散トークンの汚染、マスク拡散、トークン埋め込みおよび学習された潜在空間における連続拡散など、言語への拡散適用を試みた先行研究は、しばしば「連続から離散」への回復ボトルネックに直面します。連続表現空間におけるわずかな誤りが、誤ったトークン回復につながり、生成品質を低下させます。

著者らは、性能格差が単に離散性によるものではなく、不適切な幾何学を有する連続空間における拡散の適用に起因すると仮説を立てています。トランスフォーマー言語モデルは連続的な隠れ状態上で動作しますが、これらの状態は層によって幾何学的性質が異なります。提示された中心的な問いは、*「事前学習済みトランスフォーマーの内部のどこに、拡散ベースのノイズ除去と回復に効果的な幾何学を有する表現空間が存在するか？」*というものです。

手法：DiHAL

本論文は、言語生成のための拡散 - トランスフォーマーハイブリッドアーキテクチャであるDiHALを提案します。これは、単一の拡散モデルを訓練するのではなく、事前学習済みトランスフォーマーに拡散ブリッジを統合する「特定と置換（Locate-and-Replace）」フレームワークです。

1. 理論的動機

著者らは、そのアプローチをランジェヴィン力学と集中理論に基づいています。彼らは、「拡散に友好的な」表現空間が以下の 3 つの幾何学的性質を満たす必要があると主張します。

• 高速収束: 空間は、ノイズサンプルが目標分布へ迅速に収束することを可能にするべきです（強い対数凹性/曲率によって支配されます）。
• 安定性: ノイズ除去プロセスは、スコア推定誤差に対して頑健であるべきです（これも曲率によって支配されます）。
• 低い有効複雑性: 表現の変動は低次元多様体に集中しているべきであり、これにより拡散モデルの学習が容易になります。

2. 特定：幾何学的に導かれた層選択

隠れ状態の真の密度とヘッセ行列が利用できないため、DiHAL は経験的な幾何学的代理指標を用いて各トランスフォーマー層をスコアリングします。

• 局所曲率 ($\hat{m}_{curv}$): 局所近傍の共分散の最大固有値の逆数を通じて推定されます。高い値はコンパクトな近傍を示します。
• 大域単調性 ($\hat{m}_{mono}$): 正則化精度行列のレイリー商を通じて推定されます。高い値は大域的な方向剛性を示します。
• 有効ランク ($\hat{k}$): 共分散のトレースをそのスペクトルノルムで割った値 ($\text{tr}(\Sigma)/\|\Sigma\|$) として推定されます。低い値は低次元多様体への集中を示します。

これらの代理指標を組み合わせる（曲率を報酬とし、有効ランクを罰則とする）ことで選択スコアを計算し、最適な挿入層 $\ell^*$ を特定します。

3. 置換：隠れ状態拡散ブリッジ

最適な層 $\ell^*$ が特定されると、以下の処理が行われます。

• 下部トランスフォーマープレフィックス（層 1 から $\ell^*$）が、条件付き拡散ブリッジに置換されます。
• ブリッジは、ソースモデルの入力から導出された埋め込みベースの条件に基づき、隠れ状態 $h_{\ell^*}$ を再構成するように訓練されます。
• 上部トランスフォーマー層（$\ell^*+1$ から $N$）および元の言語モデル（LM）ヘッドは保持され、凍結（または最小限の微調整）されます。
• 推論時、ブリッジは条件をノイズ除去して $\hat{h}_{\ell^*}$ を生成し、これが保持された上部層に渡されてトークン予測が行われます。これにより、直接トークンの回復を回避し、問題を隠れ状態の再構成として再定義します。

主要な貢献

1. 拡散挿入の定式化: 本論文は、事前学習済みトランスフォーマーへの拡散の挿入を、幾何学的に導かれたインターフェース選択問題として定式化し、拡散に友好的な空間を特定するための実用的な層別代理指標（局所コンパクト性、大域剛性、有効ランク）を提案します。
2. 固定幾何学スコア: 著者らは、8B スケールのバックボーン全体でブリッジ可能性（再構成品質）と強く相関する固定幾何学スコアを導入しました。これにより、すべての候補層に対してブリッジを網羅的に訓練することなく、効果的な挿入層を特定できます。
3. DiHAL アーキテクチャ: 下部トランスフォーマー層を条件付き拡散ブリッジに置換し、上部層と LM ヘッドを再利用するハイブリッドアーキテクチャです。診断的なマッチド・バジェット設定において、DiHAL は、埋め込みベースおよび潜在ベースの連続拡散ベースラインを上回る隠れ状態回復性能を示しました。

実験結果

著者らは、Dolma v1.7 コーパスを用いて、Llama-3.1-8B-InstructおよびQwen3-8Bバックボーン上で DiHAL を評価しました。

• 層選択: 幾何学スコアは効果的な挿入層を成功裡に予測しました。Llama-3.1-8B ではスコアが層 3 を選択し、Qwen3-8B では層 2 を選択しました。これらの層は埋め込みインターフェースの近くに位置しており、初期の隠れ状態が拡散に適した埋め込みのような幾何学的構造を保持していることを示唆しています。
• ブリッジ可能性との相関: 固定幾何学スコアは、検証用ブリッジ損失と強いスピアマン相関（$\rho \approx 0.91$–$0.93$）を示し、幾何学的代理指標が、拡散ブリッジの成功した訓練が可能である層を安価に特定できることを確認しました。
• 診断的比較: 一致したバジェットでのベースライン（Diffusion-LM、SED、LD4LG、CoDAR）との比較において、DiHAL は最低の生成パープレキシティ（Gen.PPL）と最高の多様性を達成しました。具体的には、Qwen3-8B において、DiHAL は Gen.PPL を 144.83（CoDAR）から 136.02 に改善し、多様性を 0.4777 から 0.5913 に向上させました。
• 完全訓練: 完全に訓練された場合、幾何学で選択された層は、検証損失オラクル（すべてのブリッジを明示的に訓練することで選択された層）と競争力があり、"最悪の層"コントロールおよび CoDAR ベースラインを大幅に上回りました。

意義と主張

本論文は、隠れ状態の幾何学が、言語モデル内で拡散が効果的に動作する場所を決定する重要な要因であると主張しています。結果は、言語における連続拡散の失敗が単なる離散性の結果ではなく、ノイズ除去要件と表現空間の幾何学とのミスマッチにも起因することを示唆しています。

主な教訓は以下の通りです。

• 問題の再定義: 言語のための連続拡散は、トークンを直接回復することでテキストを生成する必要はありません。事前学習済みトランスフォーマーが既にデコード方法を知っている内部表現を再構成することができます。
• ハイブリッドモデルの実現可能性: 挿入点が幾何学的基準に基づいて選択される場合、事前学習済みトランスフォーマーの一部を拡散ブリッジに置換することは可能です。
• 限界: 著者らは控えめに、DiHAL はまだ単一の拡散言語モデルではなく、トークン予測のために保持されたトランスフォーマー接尾辞に依存していると指摘しています。さらに、ノイズ除去ステップにより、推論レイテンシーは現在、元のバックボーンよりも高いです。この研究は、この規模での自己回帰生成の完全な代替というよりも、拡散が既存モデルにどこから進入できるかを理解するための一歩として提示されています。