AURORA-LM：面向连续潜空间扩散语言建模的自编码统一表示

用高容量、可解码的连续文本潜变量桥接扩散生成与离散解码

前置知识

流匹配 (Flow Matching)

流匹配是连续生成模型的一种训练范式，把生成过程看作学习一个“速度场”，把高斯噪声分布沿直线（或曲线）连续地搬运到复杂数据分布，再用 ODE 求解器从噪声反向采样。本文采用最简的线性插值路径 $x_t=(1-t)x_0+t\epsilon$（$t$ 从 0 到 1 表示从干净样本到纯噪声），让网络直接预测干净终点 $x_0$。它比传统扩散更稳定、无需估计复杂方差调度。

AURORA-LM 的生成器正是用流匹配学习文本潜变量的分布，理解它才能看懂第二阶段的训练目标和采样过程。

潜在扩散 (Latent Diffusion)

潜在扩散不直接在原始数据空间（像素、token）上做扩散，而是先用编码器把数据压到一个连续潜空间，再在该潜空间训练扩散/流匹配模型，最后用解码器还原，图像领域的 Stable Diffusion/LDM 是经典代表。好处是扩散模型只需学一个更紧凑、结构更友好的分布。AURORA-LM 把这一思想用到语言：用自编码器把离散 token 映射成连续潜变量序列，再在潜空间做流匹配。

这是 AURORA-LM 的基本范式，论文全部创新都围绕“潜空间怎么设计”和“潜分布怎么学”两个问题展开。

基于查询的编码器-解码器 (Query-based Encoder-Decoder)

用一组可学习的“查询向量”(query) 抽取信息的架构，类似 Perceiver/Q-Former。AURORA-LM 用 $N$ 个查询聚合变长 token 序列：第 $i$ 个查询只能看到前 $\lceil iL/N\rceil$ 个 token 前缀和已有潜状态，后生成的潜位置包含越来越多上下文，天然形成从左到右的因果前缀顺序。解码器要求第 $j$ 个输出 token 只依赖对应潜前缀 $z_{1:\lfloor(j-1)N/L\rfloor+1}$，把“文本前缀”和“潜变量前缀”对齐。

这个对齐的因果结构是后续“分块从左到右生成”的基础，没有它块因果扩散就无法实现。

块因果分解 (Block-causal Factorization)

介于纯自回归（逐位置生成，慢但上下文完整）和一次性全并行去噪（快但丢掉因果顺序）之间的折中。把长度 $N$ 的潜序列切成 $B=\lceil N/Q\rceil$ 个大小为 $Q$ 的块，块间按因果顺序从左到右依次生成（每块依赖前面已生成块），块内所有位置同时并行去噪，写成 $p_\theta(\alpha)=\prod_{b=1}^{B}p_\theta(\alpha^{(b)}\mid\alpha^{(<b)})$。块大小 $Q$ 越小上下文越完整、质量越高但顺序步数越多。

它决定生成质量与效率的权衡，是 AURORA-LM 把扩散并行性与语言因果性结合起来的关键机制。

无分类器引导 (Classifier-free Guidance, CFG)

一种推理阶段的采样技巧：用同一个模型分别做“有条件”和“无条件”两次预测再外推 $\hat{x}_{\text{CFG}}=\hat{x}_{\text{uncond}}+w(\hat{x}_{\text{cond}}-\hat{x}_{\text{uncond}})$，权重 $w>1$ 强化条件方向，权衡质量与多样性。AURORA-LM 在条件任务用标准 CFG（$w=2\sim3$），无条件任务则发明 SC-CFG，用“带自条件”和“不带自条件”两次预测替代，且引导只在高噪声区间 $\sigma\in[0.9,1.0]$ 生效。

论文最终指标都依赖引导调参，理解 CFG 才能看懂系统对比里的采样配置。

研究动机

语言在现代生成建模中是个异类：图像、视频、音频越来越多地在连续潜空间里建模，而文本生成仍主要依赖离散 token。早期的连续语言模型走了两条路，但都有硬伤。第一条是直接在已有表示上做扩散——Diffusion-LM、PLAID 直接对词嵌入做扩散，TEncDM 用预训练编码器的上下文特征——但这些表示的潜宽度、维度、序列组织都不是为“同时支持精确重建和有效生成”而设计的。第二条是先自编码再做扩散（LD4LG、COSMOS、Cola-DLM），把文本映射到连续潜空间再学流匹配，但为了让潜分布更好学，往往会去压缩、平滑这个潜变量。问题恰恰在于：解码器要用的正是这个被压缩的潜变量。压缩得越狠，生成越好学，但恢复精确词语、句法、局部顺序所需的区分度也丢得越多；反过来保留更多信息，重建更准，但先验要面对一个更宽、更复杂的分布。论文的消融很直观：潜宽度 $D=128$ 时，只要给潜变量加一点高斯扰动，token 重建准确率就迅速崩塌，而 $D=1024$ 能扛住大得多的扰动范围。也就是说，“潜生成难度”和“文本重建保真度”被同一个表示瓶颈死死耦合在一起。

本文的目标是本文要同时解决两个相互纠缠的挑战：第一，构造一个连续文本表示，既能忠实保留语言语义、又携带足够信息让 token 能被精确恢复；第二，设计一个生成模型，能在不牺牲解码保真度的前提下有效学到这个潜分布。核心主张是把“表示构造”和“分布建模”彻底解耦成两个阶段。可量化的目标包括：在 OpenWebText 自由生成和 XSum 条件摘要上成为已评测连续/扩散语言模型里最强的；把扩散模型扩到约 1B 参数、总计算量约 1500 EFLOPs，并在统一评测协议下超过一个更大的公开潜在扩散语言模型（Cola-DLM，约 1.8B）。

与已有工作不同的是，已有潜在扩散语言模型（COSMOS、Cola-DLM 等）的通行做法是“把表示弄简单来迁就生成模型”——压缩潜序列、加 KL 正则让分布更平滑。AURORA-LM 反其道而行：它抓住了“表示首先要当好文本的接口，其次才考虑好不好学”这个被忽视的点，坚持保留高容量、全宽度的可解码潜变量，转而改造扩散模型让它直接去学这个复杂分布。具体表现为三个独到设计：只对噪声输入通路施加低秩瓶颈（$D_b=128\ll D=1024$）、输出头仍预测全宽干净潜变量；把噪声调度标定到潜宽度；并引入自轨迹一致性。这把传统的“压缩表示换可学性”的权衡整个翻转了过来。

核心方法

直觉上，AURORA-LM 像是请了两位互不干扰的专家：一位“翻译官”（自编码器）负责把离散文本转写成一种高保真、可精确逆推的连续“速记”；一位“雕塑家”（扩散模型）负责从噪声里直接塑造出这种速记，却绝不染指翻译官的手艺。技术路线分两阶段。第一阶段用一个基于查询的编码器-解码器，把长度 $L$ 的 token 序列组织成 $N$ 个因果有序、宽度为 $D$ 的连续潜位置，靠 token 级交叉熵训练它精确重建文本，训完即冻结。第二阶段冻结自编码器、对其输出做逐通道仿射归一化，再用一个块因果扩散 Transformer 通过流匹配去学这些全宽潜变量的分布；推理时从左到右逐块去噪，再把生成的潜变量反归一化、送回冻结解码器还原 token。两阶段数据流在图 3 里完整画出。

全文最关键的一点是：不要为了迁就生成模型而去压缩解码器面对的潜变量，而要让生成模型主动适配高容量表示。具体做法是“低秩输入 + 全宽目标”的不对称设计：冻结解码器要求一个宽度为 $D$（最终配置 $D=1024$）的干净潜块，但扩散 Transformer 不必以同样宽度去读噪声输入——它通过两层线性投影构成的低秩瓶颈 $W_{\text{in}}=W_{\text{up}}W_{\text{down}}$（瓶颈宽度 $D_b=128$）先把噪声块压成紧凑表示再送入网络，而输出头仍预测全宽 $D$ 维干净块。这样既保留了用于精确解码的全部通道容量，又让高维潜分布变得可学。再叠加两个配套设计——把训练噪声按潜宽度标定（越宽越偏向高噪声）、以及自轨迹一致性——前者解决“该在哪个噪声段多放监督”，后者解决“独立采样噪声态训练、却要沿单条轨迹迭代去噪”的失配。三者共同把“高容量潜变量 = 难学”的成见打破。

方法步骤详情

流程分五步。①自编码器训练：token 序列经查询编码器聚合成 $N=\text{round}(cL)$ 个因果有序潜位置 $z_{\text{enc}}\in\mathbb{R}^{N\times D}$（第 $i$ 个查询只见前 $\lceil iL/N\rceil$ 个 token 前缀），解码器把潜前缀映射到 logits $o_j=h_jE^\top$，用交叉熵 $\mathcal{L}_{\text{AE}}$ 训练，辅以 token 嵌入与潜 dropout，训完即冻结。②逐通道归一化 $z=(z_{\text{enc}}-\mu)\oslash s$，把 $N$ 个位置切成 $B=\lceil N/Q\rceil$ 个块（$Q=16$）。③并行块训练：每块独立加噪 $\alpha^{(b)}_t=(1-t)\alpha^{(b)}+t\epsilon^{(b)}$，双流注意力掩码保证块 $b$ 只看自身噪声与干净前缀 $\alpha^{(<b)}$，一次前向训完全部块，得 $\mathcal{L}_{\text{FM}}$。④加自条件（喂回上一步干净估计）与自轨迹一致性（EMA 模型在相邻态 $t'=\text{Prev}_S(t)$ 给目标，最小化 $\mathcal{L}_{\text{ct}}$），总损失 $\mathcal{L}_{\text{train}}=\mathcal{L}_{\text{FM}}+\lambda_{\text{ct}}\mathcal{L}_{\text{ct}}$（$\lambda_{\text{ct}}=1.0$，$S=16$）。⑤推理：逐块去噪并缓存前缀 KV，反归一化 $z_{\text{dec}}=\mu+s\odot\hat{z}$ 后送冻结解码器得 token，引导在高噪声窗口 $\sigma\in[0.9,1.0]$ 生效。

技术新颖性

和已有连续语言模型相比，AURORA-LM 的本质区别有三点。对 Diffusion-LM/PLAID 这类直接在词嵌入上扩散的方法，它不沿用现成嵌入，而是显式构造一个为生成和解码共同设计的自编码潜空间。对 LD4LG/COSMOS/Cola-DLM 这类潜在扩散方法，它不靠压缩潜序列换可学性，而是保留全宽干净目标、只压噪声输入通路——低秩输入借鉴自 JiT 和 ELF，但 AURORA-LM 首次把“冻结自编码器 + 噪声按宽度标定 + 自轨迹一致性”组合到语言潜扩散里。对 ELF（冻结 T5 编码器、去噪器与解码器共享权重、只在最后一步做 token 监督），AURORA-LM 把两者完全解耦，用独立查询解码器只做文本重建监督，让去噪器专心学潜分布。

**Figure 3**: AURORA-LM 总览流程。(a) 基于查询的自编码器，(b) 块因果潜扩散，(c) 推理策略

**Figure 4**: 自轨迹一致性总览：沿模型自身更新轨迹对齐相邻的干净潜变量预测

实验结果

论文从消融、系统、扩展三层递进。①潜容量（OWT128）：$D=1024$ 在加噪下 token 重建鲁棒性远超 $D=128$，配好噪声调度后 MAUVE 随宽度单调上升，$D=1024$ 最高达 0.839（tan-d $d=7$）；保留 90% 潜位置 MAUVE 仅微降到 0.807，70% 才跌到 0.671。②建模全宽潜分布：低秩噪声输入瓶颈 $D_b=128$ 最优（均值 MAUVE 0.808），过窄过宽都不好；高噪声占比 $m_{0.7}$ 越高 MAUVE 越好；表 1 里 $x_0$ 预测+$x_0$ 空间损失最高 0.815，而 $v$ 空间损失因 $1/\sigma^2$ 加权跌到 0.059。③高效块生成：自轨迹一致性在每个步数预算上都提升 MAUVE，8 步从 0.116 飙到 0.672；块粒度 $Q=4$ 最优(0.909)，综合选 $Q=16$(0.815)，顺序阶段从 32 降到 8。④系统对比（AURORA-LM-S，130M 去噪器）：OpenWebText 自由生成（表 2）拿最低 Gen-PPL 23.56 和最高 MAUVE 0.890，全面超过 AR(39.40/0.851)、MDLM(121.36/0.668)、ELF-B(24.11/0.229)、Duo(86.57/0.704) 等；XSum（表 3）ROUGE-1/2/L 36.6/13.4/28.9 三项最优，超 ELF-B(36.0/12.2/27.8)。⑤扩展到 1B（约 1500 EFLOPs）：表 4 九项宏平均 32.6，反超更大的 Cola-DLM(1.8B) 25.1 且九项全胜，如 HellaSwag 50.3 vs 45.3、SQuAD EM 38.2 vs 25.7。

**Table 1**: 预测目标与损失空间分析：直接干净潜变量回归在两种噪声输入宽度下都拿最高 MAUVE

**Table 2**: OpenWebText（1024 token）无条件生成对比

**Table 4**: 九项基准评测：AURORA-LM-L（1B）对 Cola-DLM（1.8B）

**Figure 1**: AURORA-LM 在不同生成设置与模型规模上的评估。(a) 系统级对比，(b) 扩展评估

查看结构化数据

任务	指标	本文	基线	提升
OpenWebText 自由生成	Gen-PPL（GPT-2 Large 外评，越低越好）	23.56	次优连续方法 ELF-B 24.11；自回归 AR 39.40	比次优 ELF-B 低 0.55；较 AR 降低约 40%
OpenWebText 自由生成	MAUVE（分布相似度，越高越好）	0.890	AR 0.851；MDLM 0.668；ELF-B 0.229	比次优 AR 高 +0.039，远超其余基线
XSum 条件摘要	ROUGE-L	28.9	ELF-B 27.8；MDLM 25.8	ROUGE-1 36.6(+0.6)、ROUGE-2 13.4(+1.2)、ROUGE-L 28.9(+1.1)，三项均最优
九项语言基准宏平均	Avg（越高越好）	32.6（AURORA-LM-L, 1B）	Cola-DLM 25.1（1.8B）	+7.5，九项基准全部胜出

局限与改进

作者自己承认的局限不多，但有几处值得注意。其一，自编码潜宽度与系统配置存在不一致：消融明确选了 $D=1024$ 为最优，而最终部署的 AURORA-LM-S 自编码潜通道却是 $D=768$（见表 5），说明受维度约束并未把最优设置完全搬上去，潜力还没榨干。其二，绝对分数仍偏低：1B 模型在 MMLU 只有 22.2、ARC-C 21.2、SIQA 30.2，与同参数量级自回归模型相比仍有明显差距，连续语言模型在知识和推理上尚未真正“打平”。其三，推理仍需 16~64 步去噪加引导，相比自回归单步生成在延迟上不占优。其四，自轨迹一致性的 $S=16$、$\lambda_{\text{ct}}=1.0$ 等超参是手调的，跨任务/跨规模是否稳定未充分验证。我的观察：系统对比里 AURORA-LM-S 的去噪器 130M 比基线 92M 更大，还多带一个冻结的自编码器，参数量并非完全公平；Cola-DLM 是用其默认采样协议评测、未重训，匹配度有限。

独立分析的弱点

我独立看到的弱点与改进方向。①效率-质量的两难：块因果 $Q$ 越小质量越高但顺序阶段越多，扩散步数越少越需要自轨迹一致性救场；改进方向是把块级 KV 缓存与少步蒸馏（如一致性蒸馏到 1~4 步）结合，或用非自回归的并行块调度。②自编码器与去噪器的“分阶段冻结”割裂了联合优化：冻结后去噪器无法反过来影响表示，万一潜空间对扩散不友好只能靠噪声调度补救；改进方向是引入可微的轻量联合微调，或对潜空间加“扩散友好”的正则。③评估口径偏生成流畅度：核心指标 Gen-PPL/MAUVE 偏语义流畅与分布相似，缺少严格的负对数似然或下游任务稳健性；改进方向是补 token 级 NLL 与更多条件任务。④引导依赖重：最终系统靠 SC-CFG $w=4$、CFG $w=2\sim3$ 才打出好分数，说明基础分布仍有缺陷；改进方向是把引导强度训练化（如引导蒸馏）。⑤数据/算力门槛高：1B 模型用了 294.7B token 和约 1500 EFLOPs，对小团队不友好。

未来方向

作者明确提到两个延伸方向：把连续语言表示推广到更长上下文和更广语言能力，以及发展跨语言与其他连续模态的统一生成模型（让语言不再被迫离散化）。基于成果可延伸的方向包括：把“低秩噪声输入 + 全宽目标”这一不对称设计迁移到多模态（如图文统一潜空间）；用强化学习或偏好优化提升块因果扩散在推理/代码任务上的表现（已有离散扩散走 RL 的先例）；把自轨迹一致性进一步发展成真正的少步/单步一致性模型以逼近自回归延迟；探索与自回归模型混合的投机解码。

复现评估

复现友好度中等偏上。论文公开了代码（GitHub: fyv587/AURORA-LM）和项目主页，附录给出了详尽的架构表（表 5）、训练超参表（表 6，含优化器、学习率 $4\times10^{-4}$、批大小 1024、自条件概率 $p_{\text{sc}}$、一致性 $\lambda_{\text{ct}}=1.0$、$S=16$ 等）和完整采样协议（表 7、表 8 的引导/步数扫描）。数据方面，小规模用公开的 OpenWebText 和 XSum，可直接复现；但 1B 版本的 294.7B token 混合（含合成数据）未完全公开，且全部实验在华为昇腾 NPU 上做，GPU 用户需适配。算力门槛高：130M 版去噪器训练约 155 EFLOPs，1B 版约 1500 EFLOPs，复现完整扩展实验对小团队不现实，但复现系统级 130M 对比实验是可行的。

AURORA-LM：面向连续潜空间扩散语言建模的自编码统一表示 AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling