論文解説: AURORA-LM — テキストを連続潜在空間で生成する拡散言語モデル、「圧縮せずに学び切る」という選択

テキストを離散トークンではなく連続潜在空間で生成する拡散言語モデルAURORA-LMを論文本文から解説。潜在を圧縮して拡散を楽にする従来路線を捨て、幅1024の復元可能な潜在をそのまま、低ランク入力・高ノイズ校正・自己軌道一貫性で学び切る設計を追う。

AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling

テキストだけが「離散」に取り残されている

画像も動画も音声も、いまの生成AIは連続的な潜在空間の中で拡散モデルが作るのが主流です。ところがテキストだけは、いまだに語彙表から1トークンずつ選ぶ離散生成が支配的です（§1）。マルチモーダルモデルでも、画像や音声を離散トークン列に変換して言語モデルの土俵に持ち込むのが普通で、論文はその逆を問います——言語のほうを連続空間で表現して生成できないか？（§1）

引っ越しに例えてみます。文章をいったん連続ベクトルの「荷物」に詰め（エンコード）、生成モデルという運送屋がその荷物を一から作れるように学び、届いた荷物を開けて文章に戻す（デコード）。これが連続潜在方式の言語生成です。問題は荷造りにあります。箱を小さく軽くすれば運送屋は楽ですが、開けてみると食器が欠けている——つまり元の単語が正確に戻らないのです。

AURORA-LM（南京大学・南洋理工大学・インペリアルカレッジの共同研究、arXiv:2608.02602）の答えは明快です。箱は大きいまま、運送屋を鍛える。

先行研究のジレンマ: 圧縮すれば運べるが、開けると壊れている

論文は既存の連続言語生成を2つの路線に整理します（§1, §2.3）。

1つ目は埋め込み空間系。Diffusion-LMやPLAIDは単語埋め込みの列を、TEncDMは事前学習エンコーダの特徴をそのままノイズ除去の対象にします。既にある表現を借りるので手軽ですが、その幅や並び方は「生成しやすさ」と「復元しやすさ」の両立を狙って設計されたものではありません（§1）。

2つ目はオートエンコーダ系。LD4LG・COSMOS・Cola-DLMはテキストを専用の潜在空間に写してから拡散モデルを学びます。潜在を圧縮・平滑化できるので拡散側は楽になりますが、その同じ潜在にデコーダも依存しているため、圧縮は正確な単語・語順の復元に必要な情報まで削ります。生成のしやすさと復元の忠実さが、1つの表現ボトルネックで綱引きになっている——これが論文の現状診断です（§1）。

この「非可逆圧縮で細部が消える」感覚は、JPEGで体感できます。圧縮率を上げるほど運びは軽くなりますが、細部は戻ってきません。テキスト潜在で同じことをすると、消えるのは「正確なその単語」です。

AURORA-LMの答え: 表現づくりと分布学習を完全に分ける

AURORA-LMは仕事を2段階に切り離します（§3）。

第1段階では、Query-based Encoder-Decoder（クエリ型エンコーダ・デコーダ）がテキストを高容量の潜在列に写し、そこから元のトークンを復元できることだけを目標に学習します。生成のしやすさへの妥協はしません。

第2段階では、このオートエンコーダを凍結し、その出力分布をBlock-causal Diffusion Transformer（ブロック因果な拡散Transformer）がflow matchingで学びます。ブロックを左から右へ順に生成し、ブロック内の位置は並列にノイズ除去する構成です。表現側を固定してから生成側を設計するので、「拡散が難しいから表現を削る」という逆流が起きません（§2.3, §3）。

部品1: 前から読める潜在列を作るオートエンコーダ

エンコーダは、長さ $L$ のトークン列を $N = \operatorname{round}(cL)$ 個・幅 $D$ の潜在ベクトル列に変換します。$c$ は保持率で、$c=1$ なら圧縮なしです（§3.2）。仕掛けはクエリの見える範囲にあります。$i$ 番目の潜在クエリは、先行する潜在と先頭から $\lceil iL/N \rceil$ 番目までのトークンしか見られません。つまり潜在列の前置き（プレフィックス）が、そのままテキストの前置きに対応します。デコーダ側も対称で、$j$ 番目のトークンを復元するとき対応する潜在前置きまでしか参照できません。この対応関係が、後段の「左から右へのブロック生成」を支えます（§3.2）。

学習はトークン復元のクロスエントロピーだけ。入力埋め込みを確率0.3で丸ごと落とすトークン埋め込みdropoutと、潜在の各座標を確率0.6で落とす潜在dropoutで頑健化します（§6.1.1）。学習後は凍結し、チャネルごとに平均・分散を揃えて標準化してから第2段階に渡します（§3.3）。

部品2: 左から右へ、塊ごとに拡散する

生成側の土台はflow matchingです（背景は拡散モデル入門へ）。きれいな潜在 $x_0$ とガウス雑音 $\varepsilon$ を直線で結びます（§3.1）。

$x_0$ がきれいな潜在、$\varepsilon$ が純粋な雑音、$t$ が混合比です。$t=0$ で元データ、$t=1$ で完全な雑音になり、モデルはこの直線を逆向きにたどる方法——具体的には任意の途中地点からきれいな端点 $x_0$ を当てる方法——を学びます。つまりこの式は、本物と砂嵐をつまみひとつで混ぜているだけです。つまみを0にすれば元の文章の潜在そのまま、1にすれば完全な砂嵐、その中間はどこかが壊れた途中の状態で、モデルはそれを元に戻す係を任されます。

潜在列は $Q$ 個ずつのブロックに区切られ、確率は自己回帰的に分解されます: 各ブロックは、確定済みの前ブロック列を条件に、ブロック内の全位置を同時にノイズ除去して生成されます（§3.3）。$Q$ は「逐次の文脈」と「並列の速さ」のダイヤルです。論文の実測では $Q=4$ でMAUVE 0.909、$Q=64$ で 0.631 まで落ち、採用値 $Q=16$ は 0.816 を保ちつつ逐次段数を $Q=4$ 比で 1/4 にします（§4.1.3）。推論時は完成した前置きをKVキャッシュして次ブロックの条件に使います（図3c。仕組みはKVキャッシュの解説へ）。

幅1024の分布をどう学び切るか: 3つの工夫

復元のために潜在を太くしたぶん、その分布は拡散モデルにとって手強くなります。論文はここに3つの設計を重ねます（§3.4, §4.1.2）。

1. 低ランク入力ボトルネック。 ノイズまみれの入力をフル幅で読む価値は薄い——高ノイズ時の変動はほぼガウス雑音だからです。そこで入力経路だけを2枚の線形層で細くします。

$$W_{\mathrm{in}} = W_{\mathrm{up}} W_{\mathrm{down}}, \qquad W_{\mathrm{down}} \in \mathbb{R}^{D_b \times D}$$

$W_{\mathrm{down}}$ が幅 $D$ のノイズ入力を細い幅 $D_b$ に落とし、$W_{\mathrm{up}}$ がTransformerの隠れ幅へ持ち上げます。つまり、太い入力層1枚を「いったん細い戸口を通してから広げ直す」2枚に分けただけの式です。モデルはノイズまみれの入力を $D$ 個そのまま読むのではなく、要約として読みます。予測目標は幅 $D$ のままなので、デコーダに渡す情報量は削れません——読むのは要約でも、返すのは $D$ 個ぜんぶ、というわけです。実験では $D_b=128$ が最良でした（平均MAUVE 0.808、§4.1.2）。

2. ノイズ配分の校正。 学習時にどのノイズ強度を重点的に見せるかを、潜在の幅に合わせて調整します。tan-$d$ スケジュールの掃引では、高ノイズ側に質量を寄せるほど（$m_{0.7}$ が大きいほど）MAUVEが一貫して改善し、この傾向はlogit-normal系でも同じでした。最終系は $d=7$ を採用しています（§4.1.2）。

3. きれいな潜在を直接当てる。 予測目標と損失空間の4通り比較では、$x_0$ 予測＋$x_0$ 空間損失が両方の入力幅で最良。速度（$v$）空間で損失を測ると誤差が $1/\sigma^2$ 倍に重み付けされて低ノイズ偏重になり、上の高ノイズ校正と衝突して崩壊します（$D_b=128$ でMAUVE 0.815 対 0.059、表1・§6.1.2）。

さらに推論と学習のズレを埋める2つの補強が入ります。self-conditioningは前ステップのきれい潜在推定を次ステップの入力に戻す仕組み。self-trajectory consistency（自己軌道一貫性）は、モデル自身の1ステップ更新の前後で予測がぶれないよう、EMA（指数移動平均）版モデルの予測に揃える補助損失です（§3.4）。効果は少ステップ生成で劇的で、8ステップ生成のMAUVEが 0.116→0.672、16ステップで 0.310→0.786 に跳ねます（§4.1.3、表16）。

生成ループを擬似コードにするとこうなります（§3.5の要約）。

z = encode(prompt)                  # 空ならば無条件生成
for b in range(num_blocks):
    x = randn(Q, D)                 # ブロック分のガウス雑音
    for t in schedule(steps):       # 高ノイズ → 低ノイズ
        x0_hat = denoiser(x, t, prefix=z)  # 幅Dのまま予測
        x = euler_step(x, x0_hat, t)       # 雑音を1段抜く
    z = concat(z, x0_hat)           # 完成ブロックを文脈に追加
text = frozen_decoder(unstandardize(z))

実験: 同一条件の比較で連続系・拡散系の最高値

OpenWebTextの1,024トークン無条件生成（表2）では、AURORA-LM-S（デノイザ130M）がGen-PPL 23.56・MAUVE 0.890で評価対象中の最良でした。連続系の先行ELF-BはPPLこそ近い（24.11）ものの、MAUVEは0.229と分布の類似性で大差がつきます。自己回帰ベースラインは39.40/0.851でした。XSum要約（表3）でもROUGE-1/2/L = 36.6/13.4/28.9で全指標最良です。

スケーリング検証では、約1,500 EFLOPs・294.7Bトークンの混合コーパスで学習した1B版AURORA-LM-Lが、公開済みの1.8B潜在拡散モデルCola-DLMを9つの言語ベンチマークすべてで上回りました（マクロ平均32.6対25.1、表4・§4.3）。実験はすべてAscend NPU上です（§4）。

限界と読むときの注意

論文の数値をそのまま受け取る前に、線を引いておくべき点がいくつかあります。

• 絶対性能はまだ研究スケールです。1B版のMMLUは22.2%で、これは4択の偶然水準に近い値です（表4）。「連続潜在方式どうしの比較で勝った」ことと「実用LLMに並んだ」ことは別問題です（この段落の評価は筆者の解釈です）。
• パラメータ表記はデノイザのみで、オートエンコーダを除きます。総計はSで267.8M、Lで約1.34Bです（§4、表5）。
• Cola-DLM比較は公開チェックポイントの評価であり、訓練条件はマッチしていないと論文自身が明記しています（§6.2.4）。OWTのベースラインは自前再現、XSumのベースライン値はELF論文からの引用（†）です（§4.2、表3）。
• 文脈長はSで1,024・Lで512トークンと短く、指示チューニングもなし（表6）。長文脈化は今後の課題として挙げられています（§5）。

現場ではこう使う

この論文が直接効くのは、生成モデルの研究エンジニアと、テキストと画像・音声を1つの生成基盤に載せたいマルチモーダル基盤の設計者です。読みどころはアーキテクチャそのものより「高次元潜在を拡散で学ぶときのチューニング手順」にあります。

論文で実際に回しているダイヤルは、ノイズスケジュールのシフト（tan-$d$ の $d$）、入力ボトルネック幅 $D_b$、ブロックサイズ $Q$、サンプラー（Euler / SDE-DPM++）とステップ数、ガイダンス強度（無条件はSC-CFG $w_{sc}=4$、条件付きはCFG $w=2.0$、しかも高ノイズ帯 $\sigma\in[0.9,1.0]$ 限定）です（§4.2、§6.2.4）。

事故になりやすい落とし穴も論文の表から読めます。第一に、$v$ 空間の損失をそのまま流用しないこと。画像拡散の定番設定でも、幅広の潜在ではMAUVE 0.059まで崩壊します（表1）。第二に、ノイズスケジュールを表現の幅と独立に決めないこと。最適な高ノイズ配分は $D$ に依存して動きます（§4.1.2）。第三に、少ステップ推論を狙うなら一貫性系の補助損失が前提になること。素のflow matchingは8〜16ステップで大きく劣化します（§4.1.3）。実務上は、この3点は連続潜在をテキスト以外（例えば独自トークナイザの時系列やコード）に適用するときも最初に疑うべきチェックリストとして使えます——ここは論文の主張ではなく筆者の実務的な読み替えです。

まとめ

• AURORA-LMは「表現を生成に合わせて削る」従来路線を反転し、復元可能な幅1024の潜在を固定してから、拡散側を表現に合わせて設計した（§1, §3）
• 鍵は3点セット——低ランクのノイズ入力経路＋フル幅の予測目標、幅に合わせた高ノイズ校正、自己軌道一貫性（§3.4）
• 同一評価条件では連続系・拡散系の中で最良、1Bスケールでも公開潜在拡散モデルを一貫して上回る。ただし絶対性能は研究スケールで、長文脈と実用性能は次の宿題（§4, §5）

土台の仕組みから追いたい方は、拡散モデル入門とKVキャッシュの解説、非可逆圧縮の直感にはJPEGの仕組みをどうぞ。