继何恺明团队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 参数