一句话总结

南京大学、南洋理工大学和帝国理工学院的研究人员提出AURORA-LM，一种连续潜在扩散语言模型，采用基于查询的编码器-解码器构建高容量、前缀对齐的潜在序列，并利用带流匹配、噪声水平校准和自我轨迹一致性的分块因果扩散Transformer，在OpenWebText自由生成和XSum条件摘要任务上，在所有评估的连续和基于扩散的语言模型中取得了最强性能，同时可扩展至1B参数。

核心贡献

• 基于查询的编码器-解码器构建一个前缀对齐、高容量的连续文本潜在表示，并在扩散训练前冻结以解耦表征学习；更宽的潜在表示在受到损坏时仍能稳健保留token级别的信息。
• 分块因果扩散Transformer通过流匹配对整个潜在分布建模，从左到右生成块，同时并行地对块内位置进行去噪。
• 为了学习高维分布而不降低解码能力，仅在噪声输入上应用低秩瓶颈，并根据潜在宽度校准噪声水平分布，同时通过自我轨迹一致性使训练与迭代去噪对齐；这使其在OpenWebText自由生成和XSum条件摘要任务上，在所有评估的连续和基于扩散的语言模型中取得了最强性能，并且一个1B参数模型超越了一个更大的公开发布的潜在扩散语言模型。

引言

连续潜在空间是现代生成模型的重要支柱，使视觉和音频领域的高质量合成成为可能。然而，语言建模仍依赖于离散token，这在多模态系统中造成了不对称。此前在连续空间中生成文本的尝试要么复用token嵌入或预训练编码器特征，这些方法并未同时为准确重构和可处理的生成而显式优化。基于自编码器的方法可以学习专门的潜在表示，但面临一种折中：压缩表示简化了生成任务却牺牲了解码保真度，而保留细节又使潜在建模更复杂。作者引入AURORA-LM框架，将文本表征的构建与其分布的学习解耦。他们首先利用基于查询的编码器-解码器构建一个高容量、因果有序的潜在空间，该空间能够忠实地支持token恢复；然后将其冻结，并通过流匹配训练一个分块因果扩散模型。为处理由此产生的全宽度潜在分布而不进行压缩，他们结合了低秩输入投影、宽度感知的噪声分配和自我轨迹一致性正则项。该设计实现了强大的生成和条件摘要性能，在有效扩展的同时超越了先前的连续和离散基线。

方法

AURORA-LM框架围绕两个主要学习目标组织：构建一个可准确解码回离散token的连续文本表征，以及对该连续潜在空间的生成分布进行建模。其架构包含一个基于查询的编码器-解码器和一个分块因果去噪器，后者是一个使用流匹配训练的Transformer。

第一阶段侧重于连续文本潜在表征的构建。作者利用基于查询的编码器-解码器将离散文本序列组织成一个高容量、因果有序的潜在接口。给定长度为 $L$L 的token序列，编码器产生一个连续潜在序列 $z_{\mathrm{enc}} \in \mathbb{R}^{N \times D}$zenc​∈RN×D，其中 $D$D 是通道宽度，$N = \mathrm{round}(cL)$N=round(cL) 是由潜在保留比率 $c$c 决定的潜在位置数。为构建该有序序列，使用 $N$N 个潜在查询聚合变长token序列。每个查询由一个共享的可学向量初始化，并通过其位置可见的上下文进行区分。在第 $\ell$ℓ 层，查询 $i$i 关注此前的潜在状态和对应的token前缀，并使用RoPE编码位置。层的更新计算为：

$$\widetilde{z}_{\mathrm{enc}, i}^{(\ell)} = z_{\mathrm{enc}, i}^{(\ell)} + \mathrm{MHA}^{(\ell)} \Big( z_{\mathrm{enc}, i}^{(\ell)}, [ z_{\mathrm{enc}, 1: i}^{(\ell)}; E [ w_{1: \lceil i L / N \rceil} ] ] \Big),$$zenc,i(ℓ)​=zenc,i(ℓ)​+MHA(ℓ)(zenc,i(ℓ)​,[zenc,1:i(ℓ)​;E[w1:⌈iL/N⌉​]]), $$z_{\mathrm{enc}, i}^{(\ell + 1)} = \widetilde{z}_{\mathrm{enc}, i}^{(\ell)} + \mathrm{FFN}^{(\ell)} \Big( \widetilde{z}_{\mathrm{enc}, i}^{(\ell)} \Big).$$zenc,i(ℓ+1)​=zenc,i(ℓ)​+FFN(ℓ)(zenc,i(ℓ)​).

解码器镜像该结构，使用第二个共享向量初始化 $L$L 个输出查询。输出位置 $j$j 关注潜在前缀 $z_{\mathrm{enc}, 1: \lfloor (j - 1) N / L \rfloor + 1}$zenc,1:⌊(j−1)N/L⌋+1​ 和此前的解码器状态以重构token logits。自编码器通过token级交叉熵损失联合训练，并由token嵌入和潜在dropout进行正则化。训练完成后，冻结自编码器以固定潜在空间，供后续生成建模阶段使用。

第二阶段涉及连续文本潜在表示的分块因果建模。冻结编码器的输出通过逐通道仿射归一化进行标准化，以形成经验分布。分块因果去噪器通过将潜在序列因子分解为大小为 $Q$Q 的连续块来建模该分布。这种分块因果分解可以在块间进行从左到右的因果生成，同时在块内对所有位置联合去噪。对于干净块 $\alpha^{(b)}$α(b) 和高斯噪声 $\varepsilon^{(b)} \sim \mathcal{N}(0, I)$ε(b)∼N(0,I)，线性概率路径定义为：

$$\alpha_{t}^{(b)} = (1 - t) \alpha^{(b)} + t \varepsilon^{(b)}, \quad t \in [0, 1].$$αt(b)​=(1−t)α(b)+tε(b),t∈[0,1].

去噪器被训练为从噪声状态、噪声水平 $t$t 和干净前缀 $\alpha^{(<b)}$α(<b) 预测干净的端点 $\alpha^{(b)}$α(b)。

为高效学习全宽度潜在分布，作者引入了一个噪声潜在输入瓶颈。虽然冻结解码器需要维度为 $D$D 的全宽度干净块，去噪器通过一个瓶颈维度为 $D_b < \min(D, H)$Db​<min(D,H)（$H$H 为Transformer隐藏宽度）的低秩通路处理噪声输入。这迫使模型在进一步处理前提取紧凑的表示。并行分块训练通过使用来自冻结编码器的干净前缀并应用双流注意力掩码实现，从而在单次前向传播中评估所有块条件。干净端点流匹配损失定义为：

$$\mathcal{L}_{\mathrm{FM}} = \mathbb{E}_{z \sim q_{\mathcal{E}}, t \sim \pi, \varepsilon \sim \mathcal{N}(0, I)} \left[ \frac{1}{D|\mathcal{J}|} \sum_{i \in \mathcal{J}} \| \hat{z}_{\theta, i} - z_{i} \|_{2}^{2} \right].$$LFM​=Ez∼qE​,t∼π,ε∼N(0,I)​[D∣J∣1​i∈J∑​∥z^θ,i​−zi​∥22​].

为弥合独立训练样本与推理时的顺序性质之间的差距，该框架融入了自条件和自我轨迹一致性。自条件将来自上一采样步的干净潜在预测在训练时作为额外输入反馈给模型。自我轨迹一致性通过最小化当前模型的预测与指数移动平均模型在更低噪声水平 $t'$t′ 下的预测之间的距离，使采样轨迹上相邻状态的干净潜在预测对齐：

$$\mathcal{L}_{\mathrm{ct}} = \mathbb{E} \left[ \frac{1}{D|\mathcal{J}|} \sum_{i \in \mathcal{J}} \| \hat{z}_{\theta}(\alpha_{t}, t)_{i} - \mathrm{sg}(\hat{z}_{\theta_{\mathrm{EMA}}}(\widetilde{\alpha}_{t'}, t')_{i}) \|_{2}^{2} \right].$$Lct​=E[D∣J∣1​i∈J∑​∥z^θ​(αt​,t)i​−sg(z^θEMA​​(αt′​,t′)i​)∥22​].

完整的训练目标结合了这些损失：$\mathcal{L}_{\mathrm{train}} = \mathcal{L}_{\mathrm{FM}} + \lambda_{\mathrm{ct}} \mathcal{L}_{\mathrm{ct}}$Ltrain​=LFM​+λct​Lct​。

在推理过程中，潜在生成从左到右逐块进行。对于带提示的条件生成，提示被编码并标准化以初始化前缀。其余块通过从 $t=1$t=1 到 $t=0$t=0 去噪高斯噪声生成，并以已完成的前缀为条件。生成的标准化潜在序列随后被去标准化并通过冻结的查询解码器恢复出输出token序列。作者在推理时还使用了引导机制，例如用于无条件生成的自条件无分类器引导和用于提示条件生成的标准无分类器引导，以控制生成质量与多样性的权衡。

实验

实验首先在OpenWebText上进行控制变量消融以验证关键设计选择，结果表明宽潜在表征（D=1024）、窄噪声输入瓶颈（128）、高噪声训练调度、x₀预测和x₀空间损失，以及结合自我轨迹一致性的分块生成共同优化了性能。系统级比较随后表明，AURORA-LM-S在无条件文本生成和条件摘要任务上均优于自回归、离散扩散和其他连续生成基线，证实了对学习到的连续表征进行建模的优势。最后，将AURORA-LM-L扩展到1.01B参数，在九个不同语言任务上取得了相比当前更大规模最先进潜在扩散模型的一致提升，证明了该设计在更大规模上的有效性。

直接回归干净潜在表征（目标x0，损失x0）在两种噪声输入宽度下均达到最高的MAUVE分数，大幅超越所有基于速度的方案。干净目标与速度损失结合导致质量近乎为零，而速度目标无论搭配哪种损失都表现中等但较低。当噪声输入被瓶颈化至128维或保持全1024维宽度时，这些趋势保持一致。干净潜在回归（x0目标，x0损失）在两种瓶颈宽度下均使MAUVE超过0.80，远超次优配置。将速度损失与干净潜在目标搭配导致MAUVE低于0.06，表明预测目标与损失空间之间存在根本性的不匹配。速度预测目标（v目标）搭配x0损失在128宽度下达中等MAUVE（0.73），但在全宽度下骤降至0.33。目标-损失组合的相对排序在噪声输入宽度间保持不变，但在全1024维宽度下，干净目标与速度目标之间的性能差距扩大。

在OpenWebText的1024 token无条件生成任务上，连续潜在AURORA-LM-S取得了最佳整体质量，在所有比较系统中获得最低的Gen-PPL（23.56）和最高的MAUVE（0.890）。它超越了此前的连续模型ELF-B（后者尽管困惑度具有竞争力，但MAUVE极低），并大幅优于自回归和离散扩散基线。AURORA-LM-S将Gen-PPL相对于自回归基线降低约40%，相对于离散扩散模型降低超过70%，同时使MAUVE提升至0.890。ELF-B获得了相近的Gen-PPL（24.11），但其MAUVE塌缩至0.229，表明AURORA-LM-S独特地将低困惑度与对真实文本的高分布相似性结合在一起。离散扩散方法（SEDD、MDLM）的Gen-PPL高于119，且MAUVE低于0.7，表明其长文本生成质量远不如连续潜在模型。

在XSum条件生成任务上，AURORA-LM-S取得了最高的ROUGE-1、ROUGE-2和ROUGE-L分数，超越了包括此前最佳模型ELF-B在内的所有评估基线。这一强大性能将模型的连续潜在范式从无条件生成扩展到了提示条件生成。AURORA-LM-S在XSum上的所有三个ROUGE指标均领先于所有比较模型。在ROUGE-2和ROUGE-L上，它比次优模型ELF-B均提高了一个点以上。

更小的1B参数AURORA-LM-L在所有九个评估基准上均优于更大的1.8B参数Cola-DLM，并取得了更高的宏平均，证明该模型的优势在更大规模和不同语言任务上仍然保持。AURORA-LM-L的宏平均为32.6，而Cola-DLM为25.1，尽管参数约为后者的一半，却在每个任务上均表现更优。绝对差距最大的任务包括故事和常识推理，例如HellaSwag（18.4 vs. 5.7）和StoryCloze（54.8 vs. 33.8）。

关于潜在回归目标和损失的消融研究证实，直接以干净目标损失预测干净潜在可获得最高的生成质量，而基于速度的公式会剧烈降低性能，尤其是在不匹配的情况下。在无条件长文本生成中，连续潜在AURORA-LM模型取得了低困惑度和高分布相似性的最佳组合，远超自回归和离散扩散基线——后者要么连贯性差，要么相似性崩塌。AURORA-LM将这种强大的生成能力扩展到了提示条件摘要（取得最高ROUGE分数）和少样本推理基准上；在这些基准中，一个更小的1B参数版本在所有任务上均优于一个1.8B的离散扩散方案，证实了其在不同语言评估中的高效扩展性。