技術要約: AURORA-LM

問題提起

言語生成は、現代の生成モデリングにおいて依然として例外的な存在である。画像、ビデオ、オーディオは、拡散プロセスを用いた連続的な潜在空間でモデル化されることが増えているが、テキスト生成は依然として主に離散的なトークンに依存している。連続的な拡散を言語に適用しようとする既存の試みは、根本的な設計上の妥協に直面している：

1. 継承された表現: Diffusion-LMやTEncDMのような手法は、トークン埋め込みや事前学習済みエンコーダーから継承された埋め込み空間を使用している。これらの空間は、正確なテキスト再構成と効果的な連続生成という両方の要求を明示的に設計したものにはなっていない。
2. 圧縮のトレードオフ: オートエンコーダベースの手法（LD4LG、COSMOSなど）は、テキストを連続的な潜在変数にマッピングするが、多くの場合、拡散分布の学習を容易にするためにこれらの表現を圧縮する。この圧縮はトークンレベルの忠実度を犠牲にし、生成の扱いやすさを向上させるとデコーディングの精度が低下するという、ボトルネックを生み出す。

核心となる課題は、正確なデコーディングをサポートする忠実な連続テキスト表現を構築すると同時に、その表現能力を損なうことなくその分布を学習できる生成モデルを設計することである。

手法: AURORA-LM フレームワーク

AURORA-LM（Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling）は、表現の構築と分布のモデリングを2つの明確な段階に分離することで、これらの課題に対処する。

1. 連続テキスト潜在変数の構築（オートエンコーダ）

表現を継承するのではなく、AURORA-LMはクエリベースのエンコーダ・デコーダを用いて、高容量でデコード可能なテキスト潜在変数を明示的に構築する：

• クエリベース・エンコーダ: 離散的なテキストを高容量で接頭辞（prefix）に整合した潜在シーケンスへと整理する。これは、$N$ 個の潜在クエリを使用して、トークンの接頭辞を段階的に長いものへと集約する。クエリのインデックス $i$ が増加するにつれ、エンコーダはより長いトークン接頭辞（$w_{1:\lceil iL/N \rceil}$）に注意を向ける。これにより、潜在位置がテキストの接頭辞と整列する因果的順序を誘導する。
• クエリベース・デコーダ: 潜在接頭辞をトークンのロジットへとマッピングする。決定的なのは、デコーダが対応する潜在接頭辞のみに注意を向けるよう制限されていることであり、これにより潜在シーケンスとトークンシーケンスの間の厳格な因果関係が保証される。
• 高容量設計: 生成の容易さのために潜在変数を圧縮する先行研究とは異なり、AURORA-LMは高いチャネル幅（$D$）と最小限のシーケンス圧縮（$c \approx 1$）を維持する。これにより、潜在空間を単純化された事前分布としてではなく、テキストへの忠実なインターフェースとして扱い、正確なトークン復元に必要な情報を保持する。
• 学習: オートエンコーダはトークンレベルのクロスエントロピーを通じて学習され、その後凍結される。これにより、ターゲットとなる潜在分布が固定され、生成モデルを表現学習から独立して研究することが可能になる。

2. ブロック因果的潜在拡散（生成モデル）

**ブロック因果的拡散トランスフォーマー（Block-Causal Diffusion Transformer）**が、**フロー・マッチング（Flow Matching）**を用いて、凍結された潜在シーケンスの分布を学習する：

• ブロック因果的因数分解: 因果的な生成と並列的なデノイジングのバランスを取るため、潜在シーケンスをブロックに分割する。モデルは左から右へとブロックを生成する（前のブロックに対して因果的な条件付けを行う）が、各ブロック内のすべての位置については並列にデノイジングを行う。
• フル幅潜在変数予測: モデルは、圧縮されたバージョンではなく、デコーダが必要とするフル幅のクリーンな潜在変数（$x_0$）を予測する。
• ノイズ入力のボトルネック: ノイズを含む入力から高次元の分布を学習することの困難さに対処するため、モデルはノイズを含む入力パスウェイに対してのみ低ランク投影（$W_{down} \in \mathbb{R}^{D_b \times D}$）を適用し、出力ヘッドはフル幅のクリーンなターゲットを予測する。これにより、モデルは破損した入力を効率的に処理しながら、デコーダに提示する容量を維持できる。
• ノイズ較正: ノイズレベルの分布は、潜在変数の幅に合わせて較正される。経験的な結果によれば、より広い潜在表現は、高ノイズの入力に重点を置くノイズスケジュールから恩恵を受ける。
• 自己軌跡一貫性（Self-Trajectory Consistency）: 独立してサンプリングされたノイズ状態での学習と、反復的なデノイジングによる推論との間のギャップを埋めるため、モデルは一貫性損失を導入する。これは、デノイジングの軌跡における隣接する状態でのクリーンな潜在変数予測を整合させ、生成プロセスを安定させ、少ないサンプリングステップで高品質な結果を実現する。

3. 推論

推論時、潜在ブロックは左から右へと逐次的に生成される。生成された標準化された潜在変数は、元のエンコーダ空間へと逆変換され、凍結されたデコーダに渡されて最終的なトークンシーケンスが復元される。本フレームワークは、無条件生成およびプロンプト条件付き生成（分類器フリー・ガイダンスまたは自己条件付けガイダンスを使用）の両方をサポートする。

主な貢献

1. 統一されたフレームワーク: AURORA-LMは、テキスト表現の学習と生成分布のモデリングを分離し、トークン埋め込みから継承するのではなく、生成モデルが学習すべき連続的な表現を明示的に構築する。
2. 高容量インターフェース: 正確なトークンデコーディングのために潜在表現の全チャネル容量を保持し、表現を単純化するのではなく、特定のアーキテクチャ的選択（ブロック因果的因数分解など）を通じて、拡散事前分布がその複雑な分布を学習するように適応させる。
3. ブロック因果的フロー・マッチング: 左から右へとブロックを生成しながら、ブロック内の位置を並列にデノイジングする、フロー・マッチングで訓練されたブロック因果的デノイザーを導入した。
4. アーキテクチャの革新: 低ランクのノイズ入力パスウェイ、潜在幅へのノイズレベル較正、および自己軌跡一貫性の組み合わせにより、フル幅の潜在分布を効果的に学習することを可能にした。

実験結果

実験はAscend NPU上で実施された。

• OpenWebText (無条件生成): AURORA-LM-S（パラメータ数130M）は、評価された自己回帰型、離散拡散型、および連続埋め込み型のベースラインの中で、最も低い生成パープレキシティ（Gen-PPL: 23.56）と最も高いMAUVE（0.890）を達成した。
• XSum (条件付き要約): AURORA-LM-Sは、ROUGE-1、ROUGE-2、およびROUGE-Lの指標において、すべてのベースライン（ELF-B、AR、MDLMを含む）を上回った。
• スケーリング: AURORA-LM-L（パラメータ数1.01B、約1,500 EFLOPsの計算量）は、一貫した性能向上を示し、一致した評価プロトコルの下で、より大きな公開潜在拡散モデル（Cola-DLM、約1.8B）を凌駕した。
• アブレーション研究:
  • 潜在幅: より広い潜在変数（$D=1024$）は、破損に対する堅牢性と高い生成性能を提供した。
  • 入力ボトルネック: ノイズ入力に対する中程度のボトルネック幅（$D_b=128$）が最適であり、破損した入力をフル次元で処理することなく、フル幅の予測を可能にした。
  • 軌跡一貫性: 自己軌跡一貫性は、数ステップのサンプリング品質を大幅に向上させた。

意義と主張

本論文は、連続的な言語生成が、拡散ベースの生成モデリングと離散的なトークンデコーディングを効果的に橋渡しできることを主張している。著者らは、生成モデルに合わせて表現を単純化するという従来の設計上の妥協は不要であると述べている。代わりに、高容量でデコード可能なテキスト潜在変数を保持し、特定のアーキテクチャ的選択（ブロック因果的因数分解など）を通じて、拡散モデルがその分布を直接学習するように設計することで、生成モデルの利点を活用しながら、高いセマンティックな忠実度と構文的な規則性を維持したまま、連続的な空間を言語の一般的なインターフェースとして利用できる。結果は、連続的な空間が、拡散モデリングの恩恵を受けつつ、言語のセマンティックな忠実度と構文的な規則性を保持する汎用的なインターフェースとなり得ることを示唆している。