继何恺明团队ELF后，南京大学基于昇腾算力同步提出连续扩散语言模型

最近，何恺明老师团队提出 ELF（Embedded Language Flows），在连续语义空间中建模语言模型，在文本生成类任务上取得了相当优秀的结果。这使得连续扩散语言模型成为了讨论焦点：语言能否像其他连续模态一样，在连续空间里完成生成？

近日来自南京大学模式识别实验室（PRLab）的梁嘉骏、廖宇程，在司晨阳教授的指导下，联合新加坡南洋理工大学、英国帝国理工学院，提出连续扩散语言模型 AURORA-LM（Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling），对连续空间下的语言建模作出了更多的思考和发现。AURORA-LM 的全部实验均在昇腾 NPU 上完成，并进一步扩展至约 10 亿参数规模，验证了连续扩散语言模型在国产 AI 算力平台上的训练与扩展的可行性

• 论文标题：AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling
• 论文：https://arxiv.org/abs/2608.02602
• 项目主页：https://aurora-lm-project.github.io/
• 代码地址：https://github.com/fyv587/AURORA-LM

连续空间下的语言建模难在哪里？

语言由离散 token 组成，即使生成过程转到连续空间，最终结果仍需从连续表达（latent）序列映射回 token。于是，这条序列既要保留足够的信息，供解码器准确地映射回完整的文字；又不能让模型面对过于难以学习的分布。究竟如何构造连续空间下的 latent 序列映射，便成为连续空间语言生成中的关键问题。

AURORA-LM 将这个问题拆解成两步：

• 第一步，通过一个完全重新训练的自编码器为文本构造语义信息足够充分、可解码的连续向量序列；
• 第二步，由一个特别设计的 “分块因果扩散模型” 来学习这条编码后的向量序列的生成分布。

同时，AURORA-LM 针对带噪状态输入、训练过程和少步采样等进行设计，保障在更大语义空间下的稳定高效的建模。

模型结构

AURORA-LM 将连续语言生成拆成两个阶段。第一阶段，AURORA-LM 先为文本构造一条按前缀顺序组织的连续 latent 序列。随着 latent 位置向后移动，它能够读取的 token 前缀逐步变长；反过来，解码某个 Token 时，模型也只能使用相应的 latent 前缀。这样，连续 latent 序列保留了与文本从左到右展开相对应的结构。

编码器 - 解码器以恢复原始 token 为训练目标：

其中，是第个位置输出的 token 概率分数，是对应的真实 token。训练过程中，模型随机丢弃部分 token 来保障解码鲁棒性。

第二阶段，通过分块因果扩散模型建模 latent 序列的生成分布。AURORA-LM 将 latent 序列划分为连续的多个块，块与块之间采用从左到右的因果注意力掩码， 而块内通过双向注意力机制进行建模。对每个带噪 latent 块，模型直接预测其对应的干净 latent；所有块的预测拼接后，在完整 latent 空间中计算训练误差：

其中，是编码器产生的干净 latent，是扩散模型的预测，为参与训练的有效位置。模型始终以完整干净 latent 为目标。

推理时，块与块间从左到右按顺序推理，并通过 kv cache 复用前缀计算结果；块内不同位置可以采用双向注意力机制、进行联合去噪，同步输出：

其中，是当前要生成的 latent 块，是此前已经生成的前缀块。

此外，AURORA-LM 对当前带噪 latent 块的预测可以利用两类来源不同的附加信息。对于无条件生成，模型将当前 latent 块在上一步所预测的干净 latent 回传给下一步，为后续预测提供了自身的轨迹信息；对于有条件生成，给定的提示文本先被编码器转换为干净的 latent 前缀，提供了待生成文本的外部语义条件，从而强化模型对上下文的遵循。

为什么要保留高容量 latent？

这里所说的「高容量」，主要体现为每个 latent 向量有更大的维度。更高维的向量为文本信息保留了更多丰富的语义，使词语、语法和局部顺序在噪声扰动下仍更容易被解码。

如图（a）所示，增加 latent 向量的维度后，其在噪声扰动下仍能保留更多可供解码的文本信息；图（b）展示了，不同宽度校准训练噪声分配后，生成质量也随之提升。

但更高的 latent 容量也带来代价：扩散模型需要从带噪状态中恢复完整的干净 latent 也同样维度更高，学习目标随之更加复杂。

通过实验，研究团队发现两个关键策略对模型训练效率有较大的影响：如图（a）所示，研究团队发现适度收窄带噪输入维度（bottleneck），极大程度提升最终模型文本生成的质量，图（b）则显示，提高在训练过程中，输入高噪声状态的概率可进一步提升生成结果。

少步去噪时，如何避免误差累积？

此外，研究团队进一步探索了 AURORA-LM 在少步去噪时如何抑制误差沿去噪轨迹累积。训练时，模型面对的是独立采样的带噪 latent；实际生成时，却需沿着自身预测连续推进。若相邻去噪状态下的判断不一致，误差便可能逐步累积。

为此，AURORA-LM 引入自轨迹一致性：在训练时，模型先根据当前带噪 latent 得到对应的干净 latent 估计，再据此推进至相邻的更低噪声状态，并约束相邻状态下的干净 latent 估计保持一致。

这一更新过程可写为：

其中，是相邻的更低噪声状态，表示停止梯度。更新后的状态会交给 EMA 模型再次预测干净 latent。模型将当前状态与更新后状态得到的两次干净 latent 估计之间的差异，作为自轨迹一致性损失，以约束两次预测保持一致。

训练时，AURORA-LM 将 flow-matching 损失与这一一致性损失共同优化：

其中，控制一致性损失的权重。

实验显示，自轨迹一致性在所有评测的生成步数下均能提升结果，其中少步生成时的改善最为明显。

更多数据、更大模型规模下的效果验证

研究团队先对齐 ELF-B 的测试设置，在 130M 规模模型下，采用 OpenWebText 与 Xsum 作为训练数据，通过自由生成与条件摘要两个任务形式进行对比，再将分块因果扩散模型扩展至 1B 参数，并基于更丰富的开源数据训练，来验证当下策略在更大的模型规模下的效果。

• 自由生成： 在 OpenWebText 上，AURORA-LM 与自回归、离散扩散和其他连续生成模型进行比较。AURORA-LM 的 Gen-PPL 降至 23.56，MAUVE 达到 0.890，在两项指标上均取得表中最佳结果。
• 条件摘要：在论文展示的 XSum 对比中，AURORA-LM 的 ROUGE-1、ROUGE-2 和 ROUGE-L 分别达到 36.6、13.4 和 28.9，三项结果均为表中最高。
• 1B 规模验证： 研究团队进一步将扩散结构扩展至约 1B 参数，基于开源预训练数据累计训练计算量约为 1500 EFLOPs。如图（b）所示，在涵盖常识推理、阅读理解等九项语言基准任务的评测中，AURORA-LM-L 的平均得分为 32.6；与一个公开可用、参数规模更大的连续语言模型相比，AURORA-LM-L 在论文表中的九项任务上均取得更高分数。

结语

AURORA-LM 先为文本学习高容量、可映射回 token 的连续 latent，再让分块因果扩散模型对其分布进行建模。围绕这条 latent 所做的输入、训练与采样设计，使连续扩散生成能够同离散 token 解码衔接起来，也为更长上下文和更大规模的连续语言模型提供了一条可继续探索的路径。

文章来自于微信公众号 “机器之心”，作者 “机器之心”