画像生成や動画生成の分野では、拡散モデルが主流となっています。しかし、テキスト生成に適用しようとすると、文字化けや単語の繰り返しといった問題が発生しやすいのはなぜでしょうか。
その理由は、テキストの本質が「離散的なトークン」であるのに対し、拡散モデルは連続的なデータを扱うことを得意とするからです。これまで、拡散モデルをテキスト生成に利用するために、研究者たちは主に2つのアプローチを試みてきました。
1. 離散拡散言語モデル:離散的なトークン空間で直接、拡散過程を定義します。具体的には、MASKトークンで単語を隠して段階的に復元したり、トークンをほぼ一様分布になるまで攪拌し、徐々に修正していくといった手法です。この方法は近年主流となっており、総合的な性能も高い傾向にあります。
2. 連続拡散言語モデル:まずトークンを連続的な埋め込みベクトル(Embedding)に変換し、連続空間でノイズ除去を行い、最後に再び離散トークンに戻すというものです。このアプローチは理論的には自然で、画像の拡散モデルに近い方法ですが、実際の性能は長らく離散方式に劣っていました。
この問題を解決するために、MIT准教授でありGoogle DeepMindの著名な科学者でもある何愷明(Kaiming He)氏の研究チームが、「Embedded Language Flows(ELF)」を発表しました。これは、連続時間フローマッチングに基づき、連続的な埋め込み空間で動作する新しい種類の拡散モデルです。
既存の拡散言語モデルと異なり、ELFはほとんどの時間ステップで連続埋め込み空間にとどまり、最終ステップでのみ、共有重みネットワークを用いて離散トークンにマッピングします。この形式により、画像拡散モデルで培われた成熟した技術を直接的に応用することが可能になります。
論文リンク: https://arxiv.org/abs/2605.10938
研究結果によると、連続拡散言語モデルは、離散化処理を最小限に抑えるだけでも、非常に高い競争力を持ちうることが示されました。ELFは、蒸留技術を用いることなく、より少ないサンプリングステップ数で、より低い生成パープレキシティ(困惑度)を達成しました。さらに、必要な訓練トークン数は、従来手法のわずか10分の1で済みます。
図|ELFは蒸留なしで、より少ないサンプリングステップ数で、既存のDLM(離散拡散言語モデル)よりも低い生成パープレキシティを実現。同時に、訓練トークン数も10分の1に削減されています。
まず連続で生成し、その後離散デコード
ELFの中核的なアプローチは、まず離散トークンを連続的な埋め込み空間にマッピングし、この空間で連続時間フローマッチング(Flow Matching)を用いて、ガウスノイズからクリーンな埋め込みへの「ノイズ除去」軌跡をモデル化することです。そして、最終の時間ステップでデコードモードに切り替わり、結果を離散トークンにデコードします。
図|ELFの概念図。オレンジ色の点は連続埋め込み空間でのデータ表現、紫色の線はガウシアンノイズからクリーンな埋め込みへのノイズ除去軌跡を示しています。離散化は最終ステップ(t=1)でのみ、共有重みネットワークを介して行われます。
訓練段階では、研究チームは事前学習済みのT5エンコーダーを使用して、テキストトークンを文脈情報を含んだ連続的な埋め込みに変換します。各埋め込みは1つのトークンに対応しますが、それ自体は語彙中の特定の単語ではなく、文脈におけるそのトークンのベクトル表現です。その後、ELFは連続埋め込み空間でノイズ除去プロセスをモデル化し、ノイズからクリーンな埋め込みに至る連続的なフローの経路を学習します。
推論段階では、ELFはエンコーダーを使用しません。モデルは連続埋め込み空間でテキスト表現を段階的に生成し、最終の時間ステップでデコードモードに切り替え、共有重みネットワークと学習可能な逆埋め込み行列を通じてトークンを出力します。
ELFの設計上の重要なポイントは、1つのネットワークでノイズ除去とデコードの両方の機能を兼ね備え、二値の「モードトークン」によって区別する点です。モデルは、80%の確率でノイズ除去ブランチに、20%の確率でデコードブランチに振り分けられ、それぞれ平均二乗誤差(MSE)損失とクロスエントロピー損失で訓練されます。
さらに研究チームは、自己条件付け(Self-conditioning)メカニズムも導入しました。推論時、モデルは1つ前のステップの予測結果を、次のノイズ除去の条件として利用し、ゼロから予測することはしません。これにより生成品質が向上するだけでなく、分類器フリーガイダンス(CFG)のための既成の条件信号源を提供し、計算負荷をほとんど増やしません。
図|訓練中、離散トークンはまずクリーンな埋め込み x にエンコードされ、次に z_t に摂動されます。ELFは z_t から x̂ を予測します。モデルは、ノイズ除去損失 L_MSE またはトークンごとのクロスエントロピー損失 L_CE のいずれかで訓練されます。推論中、ELFはガウスノイズ z_0 から開始し、埋め込みを z_t から z_{t+1} へと反復的にノイズ除去します。最終ステップでのみ、ELFはデコードモードに切り替わり、逆埋め込み層を介して最終的な埋め込みを離散トークンに投影します。
より少ないサンプリングステップ、より低い訓練コスト
研究チームは、ELFを3つのタスクでテストしました。OpenWebText(OWT)での無条件テキスト生成、WMT14の独英翻訳タスク、XSumでのニュース要約タスクです。
無条件生成において、ELF-Bのメインモデルの規模は105Mパラメータです。OWTでのシステムレベル比較では、追加の蒸留を行わない場合、ELF-Bはわずか32サンプリングステップで生成パープレキシティを24まで低減し、比較対象の他の離散および連続拡散言語モデルのベースラインを上回りました。訓練コストに関しては、ELFは約452億の有効訓練トークンを使用しました。対して、MDLM、Duo、LangFlowなどのベースラインは約5243億トークン、蒸留版のMDLM+SDTTとDuo+DCDは5505億トークン、FMLMは5767億トークンを使用しています。
図|システムレベルでの比較。ELF-Bは、同様の実験設定下で、離散型および連続型の拡散言語モデルを凌駕しています(a)。追加の蒸留訓練を必要とするベースラインモデルに対しても、匹敵する競争力を示し(b)、同時に使用する訓練トークンが大幅に少ないことがわかります(c)。
条件付き生成において、ELF-BはWMT14独英翻訳タスクでBLEUスコア26.4を達成しました。XSum要約タスクでは、ROUGE-1、ROUGE-2、ROUGE-Lがそれぞれ36.0、12.2、27.8に達しました。同規模の自己回帰モデルや拡散言語モデルと比較して、ELF-Bは両方のタスクで最高の結果を達成しました。
図|機械翻訳と要約タスクの結果。研究チームはWMT14独英翻訳とXSum要約タスクでELF-Bを評価し、同程度のパラメータ数を持つベースラインモデルと比較しました。†は既存研究から直接引用した結果(独英タスクのデフォルト)、‡は公開コードを用いてチームが再現した結果(XSumタスクのデフォルト)を示します。XSumでは、可能な場合に異なる評価サンプルでの標準誤差も報告しています。ELFは両方のタスク設定で最高性能を達成しました。
さらに、アブレーション実験(要因分析)で明らかになったこととして、事前学習済みエンコーダーによる文脈依存埋め込みは、通常のトークン埋め込みや学習可能な埋め込みよりも優れた性能を示しました。共有重みのノイズ除去/デコーダーは、デコーダーを別途訓練する方法と同等の性能を、よりシンプルなプロセスで実現します。サンプリング方式では、SDE(確率微分方程式)に着想を得たサンプラーが、少ないステップでの生成においてODE(常微分方程式)サンプラーより優れていました。研究チームは、モデルを105Mから342M、652Mへとスケールアップすると、同等の多様性の下で生成パープレキシティが低下すること、そして同等の生成パープレキシティの下では、テキストの多様性が向上することを指摘しています。
図|主要な設計選択に関するアブレーション実験。
課題と今後の方向性
研究チームは、現時点でのELFモデルには依然として限界があると指摘しており、主な点は以下の通りです。
1. モデルの規模がまだ限定的
現在評価されているモデルの規模は、主に105M、342M、652Mであり、ELFをGPT-4やClaude、Llamaといった大規模命令モデルと直接比較したわけではありません。したがって、ELFが証明したのは、同類の拡散言語モデルの中での競争力であり、主流の自己回帰型大規模モデルに対する全面的な代替を意味するものではありません。
2. タスクの範囲がまだ限定的
研究実験において、OpenWebTextでの生成パープレキシティは代理的な指標であり、実際のユーザー嗜好を直接代表するものではありません。WMT14とXSumは翻訳と要約の性能を示せますが、複雑な推論、長文脈の対話、コード生成、多ターンのインタラクションをカバーするものではありません。
3. 連続空間は事前学習済みエンコーダーに依存
研究チームはゼロから訓練したエンコーダーや非文脈埋め込みもテストしましたが、事前学習済みの文脈埋め込みが依然として最も優れた結果を示しました。この結果は、ELFの効果の一部が既存の事前学習済みエンコーダーに由来しており、連続的な言語空間を完全にゼロから学習したわけではないことを示しています。
4. 実際のデプロイコストは未検証
研究チームはサンプリングステップ数、訓練トークン予算、自動評価指標を報告していますが、実際のサービスにおけるエンドツーエンドの遅延、スループット、GPUメモリ消費量を報告しておらず、成熟した自己回帰モデルのデプロイ手法との直接比較も行われていません。したがって、ELFモデルがサンプルステップや訓練トークンを削減できるかどうかは、実際のデプロイ環境で検証される必要があります。