技術的サマリー：SMOOTHIE（トークン埋め込み上の平滑化拡散）

問題定義

拡散モデルは、画像や音声などの連続領域において最先端の結果を達成してきましたが、言語の離散性により、テキスト生成への適用には重大な課題が存在します。既存のアプローチは一般的に 2 つのカテゴリに分類され、それぞれに明確な限界があります：

1. 潜在空間におけるガウス拡散：これらの手法は連続的なトークン埋め込みに対してガウスノイズを適用します。これらは（ユークリッド意味空間仮説に基づき）意味構造を保持しますが、離散的なデコーディング段階で困難に直面し、連続ベクトルからトークンへの明示的なマッピングを必要とすることが多く、不安定であったり情報損失を伴ったりします。
2. カテゴリカル/単体（Simplex）拡散：これらの手法は、語彙確率単体またはカテゴリカル分布上で直接動作します。これらはテキストの離散性を尊重し、明示的なデコーディングの必要性を排除しますが、ノイズ付加プロセス中にトークン間の意味的関係を通常無視します。これにより、情報の劣化がより不安定で、意味の薄いものとなります。

核心的な問題は、テキストの離散性を尊重しつつ、前方プロセス中に意味的一貫性を維持する形で意味情報を段階的に除去する拡散フレームワークが存在しないことです。

手法：SMOOTHIE

著者は、意味的類似性に基づいてトークン埋め込みを平滑化することで、連続拡散と離散拡散の間のギャップを埋める新しいフレームワークであるSMOOTHIE（トークン埋め込み上の平滑化拡散）を提案します。

1. 前方拡散プロセス（ノイズ付加）

埋め込みやカテゴリカル確率に直接ノイズを加えるのではなく、SMOOTHIE は対象シーケンス内の各トークン $w_i$ を、その埋め込みと語彙内のすべてのトークンの埋め込みとの間の負の二乗ユークリッド距離のベクトルとして表現します：
$$D_0 = \left( -\|E_{w_i} - E_j\|^2 \right)_{i,j=1}^{m,V}$$
ここで、$E$ は固定された埋め込み行列（例：BERT からのもの）、$m$ はシーケンス長、$V$ は語彙サイズです。

前方プロセスでは、これらの距離ベクトルにガウスノイズを加えることで徐々に攪乱します：
$$q(D_t | D_0) = \mathcal{N}\left(D_t \mid \frac{1}{\sigma_t^2}D_0, \delta^2 I\right)$$
ここで、$\sigma_t$ は帯域幅を制御するノイズスケジューラーであり、$\delta$ は確率性を制御します。ノイズの付いたベクトル $D_t$ は、softmax 関数を用いて語彙上の確率分布に変換されます：$p_t = \text{softmax}(D_t)$。

重要な洞察：ノイズが増加し（$\sigma_t$ が大きくなる）、初期に正しいトークンに集中していた確率質量は、他のトークンへと徐々に分散します。重要なのは、入力がユークリッド距離に基づいているため、確率質量は遠くのトークンに到達する前に、まず意味的に類似したトークン（ユークリッド距離が小さいもの）へと広がることです。これは「熱の放散」の原理を模倣し、意味情報が無秩序ではなく段階的に除去されることを保証します。

2. 逆拡散プロセス（ノイズ除去）

逆プロセスは、ニューラルネットワーク $f_\theta$ がノイズの付いた入力 $p_t$ から元の距離ベクトル $D_0$（またはそのスケーリング版）を予測するマルコフ連鎖としてモデル化されます。モデルは、予測された距離ベクトルと実際の距離ベクトルの間の平均二乗誤差（MSE）を最小化するように訓練されます：
$$L(\theta) = \mathbb{E}[\|D_0 - g_\theta(p_t, t)\|^2]$$
著者は、距離ベクトルの予測が期待される埋め込みの予測と等価であることを証明しており、これによりモデルはクリーンな埋め込み推定値を出力し、それを最近傍探索または argmax を介してトークンに変換できます。

3. サンプリングとソルバー

著者は、標準的な DDPM ソルバーのように以前の潜在状態に依存するのではなく、モデルの予測に対して前方プロセスを直接適用して潜在状態を更新する特定のサンプリング方式（式 9）を導入します。
$$D_{t-1} = \frac{1}{\sigma_{t-1}^2}D_0(f_\theta(p_t, t)) + \tilde{\delta}\epsilon$$
このアプローチは、更新ステップ中に以前の隠れ状態を無視します。著者は、テキスト埋め込みの疎なデータ分布においてこれが有益であり、潜在表現がデータ多様体から逸脱するリスクを低減すると主張します。

4. アーキテクチャと条件付け

• アーキテクチャ：UNet 風のスキップ接続を備えたトランスフォーマーデコーダーレイヤーに基づいています。
• 条件付け：シーケンス間タスクの場合、追加のトランスフォーマーエンコーダーがソースシーケンスを処理し、クロスアテンションを介して統合されます。
• 自己条件付け：モデルは自己条件付けを利用し、安定性と品質を向上させるために、自身の以前の予測を後続のステップの入力として供給します。

主要な貢献

1. 新しい拡散フレームワーク：SMOOTHIE は、ユークリッド距離に基づいてトークン埋め込みを平滑化することで、テキストの離散性を尊重しつつ意味情報を段階的に除去する、初のテキスト拡散手法です。
2. 単体拡散の一般化：著者は、距離ベースの潜在空間が単体ベースの潜在空間（SSD-LM や TESS で使用されるもの）を理論的に一般化することを示しています。単体拡散は、距離計量が自明（同一トークンでは 0、それ以外は 1）である SMOOTHIE の特殊な場合であることが示されています。
3. 改良されたサンプリングソルバー：本論文は、更新中に以前の隠れ状態を無視するソルバーを提案しており、これはシーケンス間タスクにおいて標準的な DDPM および DDIM ソルバーよりも経験的に優れていることが示されています。
4. 経験的優位性：広範な実験により、SMOOTHIE が既存の拡散ベースモデル（埋め込みベースおよび単体ベースのバリアントを含む）を上回り、BART や GPT-2 などの自己回帰モデルと同等の品質を複数のタスクで達成することが示されました。

実験結果

著者は、SMOOTHIE を 1 つの無条件生成タスク（ROCStories）と 4 つのシーケンス間タスク（Quora Question Pairs、Quasar-T、XSum、ParaDetox）で評価しました。

• パフォーマンス：SMOOTHIE は、他の拡散ベースのベースライン（DiffuSeq、TESS、SSD-LM、TEncDM など）および同じフレームワーク内で訓練された代替潜在空間実装（埋め込みベースおよび単体ベースの拡散）を一貫して上回りました。
• 指標：
  • XSum：BERTScore で 68.8、ROUGE-L で 26.0 を達成し、他の拡散手法を上回り、自己回帰ベースラインに近づきました。
  • ROCStories：拡散モデルの中で最も高い MAUVE スコア（80.7）を達成し、多様性を維持しつつ GPT-2（78.9）とほぼ同等の性能を示しました。
  • ParaDetox：テキスト解毒において最も高い J-Score（51.7）を達成しました。
• アブレーション研究：
  • $\tilde{\delta}$（確率性）：シーケンス間タスクでは、$\tilde{\delta}$ の値を低く（例：0.1）設定することで、パープレキシティと一貫性が向上しましたが、無条件生成では多様性を維持するためにやや高い値が最適でした。
  • ソルバー：提案されたソルバー（式 9）は、すべてのタスクで DDPM および DDIM を一貫して上回りました。
  • 埋め込み：モデルは事前学習された BERT 埋め込みに依存しています。アブレーションにより、GPT-2 埋め込みは同様の性能を示しましたが、GloVe 埋め込みは品質が低下することが示され、埋め込みの品質の重要性が浮き彫りになりました。

意義と主張

本論文は、SMOOTHIE が意味保持と離散構造の間のトレードオフを解決することで、テキスト拡散における重要な前進を表すと主張しています。ユークリッド距離を介してノイズ付加プロセスに意味的関係を直接組み込むことで、情報が段階的かつ意味的に除去される「自然なデコーディングプロセス」を可能にします。

著者は、この手法が本質的な類似性を持つ他のカテゴリカルドメイン（グラフ、タンパク質配列など）にも適用可能であると控えめに指摘しつつも、現在は MSE 損失の安定性の課題により、エンドツーエンドの訓練ではなく事前学習された埋め込みに依存していることを認めています。また、各ステップでペアワイズ距離を計算する計算上のボトルネックを認めつつも、性能の向上がコストを正当化すると主張し、高い多様性と制御性が求められるタスクにおいて、自己回帰モデルに対する堅牢な代替手段を提供すると述べています。

この研究は、潜在空間の選択と特定の拡散ソルバーが、モデルアーキテクチャ自体と同様に重要であることを示唆しており、SMOOTHIE は連続拡散と離散拡散の両方の強みを活用する統合フレームワークを提供しています。