AURORA-LM: 面向连续潜变量扩散语言建模的自编码统一表示

TL;DR AURORA-LM 构建高容量可解码的连续文本潜变量，并用块因果扩散 Transformer 直接建模其分布，在连续扩散语言模型中取得最强性能。

问题

问题背景

主流生成式 AI 正从离散表示转向连续潜在空间：图像、视频、语音均已广泛采用潜在扩散模型 （Latent Diffusion），而文本生成仍固守离散 token 序列（如自回归 LLM）。这种不一致阻碍了多模态生成框架的统一，也限制了语言模型本身对连续空间的利用潜力。

现有方法局限

现有连续语言模型主要存在两类方案，均未妥善解决表示与生成的平衡：

• 复用预训练嵌入：直接使用 LLaMA 等模型的 embedding 作为连续潜在，其空间并未为生成任务的解码优化，导致去噪目标与最终解码失配，生成质量下降。
• 压缩潜在表示：为降低扩散建模难度，对文本自编码器的潜在进行有损压缩（如 PCA 降维），但牺牲了 token 级保真度，造成字词错漏或语法损伤，无法满足高精度生成场景。

技术难点与重要性

文本的离散 token 与连续空间的平滑分布之间存在根本张力。保留高容量、可完美解码的文本潜在（例如每个 token 对应一个高维向量），意味着其分布高噪声、多模态，对扩散模型的分布拟合能力要求极高。同时，语言本有的自回归顺序约束使得并行去噪难以直接复用自回归架构。业界之所以高度关注，是因为解决该问题可以：

1. 将所有生成模态统一到连续潜在扩散框架下，简化训练与多任务融合；
2. 为多模态大模型提供更自然的内部共享表示，促进图文理解与生成的深层对齐；
3. 为可控生成与编辑提供平滑的潜在插值空间。

行业类比

如同视频生成通过像素级扩散 → 潜在空间扩散 的转变平衡了质量与效率，AURORA-LM 正尝试在文本连续潜在的高保真解码与可建模性之间寻找工程折衷，这或将为通用多模态生成器铺平道路。

核心洞察

• AURORA-LM 的核心创新在于分离高容量可解码的连续潜在表示与扩散建模，挑战了“降低潜在维度以适配扩散”的常规思路。与现有连续语言模型（如 D3PM、MDLM）不同，它保留全宽度的潜在表示，仅限制噪声输入通路，从而维持了精确的词元级解码能力。这种设计使得扩散模型可以直接学习高维分布的复杂细节，而无需牺牲表示的信息量，为追求高保真文本生成提供了新方向。
• 块因果扩散 Transformer 结合流匹配，实现了自回归块生成与块内并行去噪的混合范式。不同于传统自回归模型逐词生成，也不同于全序列扩散一次性去噪，该方法按块顺序生成，每块内部并行去噪，在推理效率与生成质量间取得了更优的平衡。自轨迹一致性损失进一步缩小了训练时独立采样噪声与推理时迭代去噪之间的分布差距，增强了模型稳定性。

方法

输入与表示构建

文本通过 Query-based Encoder-Decoder 转换为连续潜在序列。编码器使用一组可学习的 query 向量，将离散 token 序列压缩为前缀对齐的高容量潜在（prefix-aligned latent），每个输出位置对应一个 patch 或 token 组，保持可解码性。该表示不降低维度或信息量，直接作为后续扩散建模的目标。

关键模块：Block-causal Diffusion Transformer

潜在序列的分布由 Block-causal Diffusion Transformer 通过 流匹配（flow matching）学习。该架构将序列划分为块，以块因果方式从左到右生成，每个块内所有位置并行去噪。这种设计平衡了自回归的长程依赖与扩散的并行生成优势。

训练策略的适配

连续高容量潜在比离散 token 或压缩特征更难被扩散模型直接拟合。AURORA-LM 做了三项关键调整：

• 噪声路径限制：仅在输入侧注入噪声，保留完整干净潜在作为预测目标，避免解码器能力因潜在维度的压缩而下降。
• 噪声水平校准：根据潜在宽度动态调整噪声调度，使训练更稳定。
• 自轨迹一致性（self-trajectory consistency）：在训练中模拟推理时的迭代去噪路径，弥合独立噪声采样与多步推理之间的差距。

输出与生成流程

推理时，从随机噪声开始，按块从左到右顺序生成，块内并行去噪，最终得到完整潜在序列，经由预训练的解码器还原为文本。

与同类方法的差异点：现有连续文本扩散模型往往通过压缩潜在或沿用非生成设计的嵌入空间来降低建模难度；AURORA-LM 直接保留高保真、可解码的完整潜在，并改造扩散过程以适应这类更难建模的表示，从而在 token 级保真度和生成质量上取得显著提升。

实验

实验设计

评估聚焦连续潜空间扩散语言模型的自由生成与摘要能力：在 OpenWebText 上进行无条件生成，在 XSum 上进行文本摘要，对比现有离散与连续语言模型。AURORA-LM 扩展到 1B 参数，总训练计算量约 1500 EFLOPs，所有实验在昇腾 NPU 上完成。采用一致的评估协议，将结果与公开发布的更大规模潜在扩散语言模型直接对比。

关键发现

• AURORA-LM 在已评估的连续与扩散语言模型中取得最佳性能，证明保留高容量、可解码的文本潜空间并直接建模其分布，比压缩潜空间以适应扩散的做法更有效。
• 该模型维持令牌级保真度，未因扩散建模难度而牺牲解码器侧容量；通过仅约束噪声输入路径、保留完整干净潜变量预测目标，弥合了训练与推理的行为差异。
• 在相同评估协议下，AURORA-LM 超越了一个参数规模更大的公开潜在扩散语言模型，显示出更好的参数效率与算力利用率。

与基线的深度对比

与常见的连续语言模型对比，AURORA-LM 不使用为联合生成与解码设计的嵌入空间，也不压缩潜空间来降低扩散难度，而是直接学习高容量潜表示的分布。这使其在生成文本的连贯性与事实一致性上优于以往方法。相较于更大的扩散模型，AURORA-LM 通过前缀对齐的查询式编解码器与块因果扩散 Transformer，实现了更好的结构归纳偏置，在匹配评估设置下性能反超，印证了“表征保真度优于简化的生成假设”这一设计理念。

行业影响

落地场景

AURORA-LM 作为连续潜空间扩散语言模型，特别适合需要并行解码或非自回归生成的场景。在 长文本生成、摘要、对话系统 等任务中，其分块去噪的机制能够平衡生成速度与质量，有望应用于：

• 内容平台: 实时生成新闻摘要、视频脚本或商品描述，替代逐 token 生成的自回归模型。
• 企业服务: 自动撰写邮件、报告草稿，加速文档工作流。
• 搜索引擎与问答: 并行生成候选答案，提升首 token 延迟敏感的交互体验。

商业价值

• 降本: 扩散模型的并行去噪特性，在推理时可通过少量优化步骤（如一致性蒸馏）大幅减少计算量，相比自回归模型能节省 GPU/ NPU 资源，尤其对需要批量生成的长序列任务。
• 增收: 更快的生成速度可以直接改善实时交互产品的用户留存率，例如聊天机器人、实时翻译。
• 体验提升: 连续潜空间保留了更丰富的文本特征，可能生成更具多样性和流畅性的文本，减少“模式坍塌”。

与现有产品/工作流的接口

AURORA-LM 可封装为标准推理 API，与现有基于 Transformer 的生成栈并行或替换：

• 模型侧: 将 Encoder-Decoder 部分作为通用文本潜空间编码器，扩散 Transformer 作为生成核，通过 ONNX/ TensorRT 部署，支持动态 batch 和 KV-cache。
• 服务侧: 兼容 OpenAI API 规范，用户只需修改 model 参数即可调用，或作为 Hugging Face Pipeline 组件集成。
• 数据与微调: 支持标准 tokenizer 输出，可沿用现有数据集和微调框架（如 LoRA），但需注意连续潜空间的适配层。

具体落地 Use Case

1. 全球电商平台的商品描述生成: 用户输入简要属性，AURORA-LM 快速并行生成多语言、多风格的营销文案，节省人工撰写成本，并支持 A/B 测试不同风格对转化率的影响。
2. 在线教育领域的自动问答与解答: 对于学生提问，模型在边界约束下并行生成多个候选解答，由置信度模型排序后返回最佳答案，缩短等待时间，提升学习平台活跃度。

局限

• **连续扩散模型固有的生成速度劣势**：尽管 AURORA-LM 通过块因果结构实现了块内并行去噪，但每个块的生成仍需多次前向传播（例如数十步流匹配 ODE 求解），推理速度远慢于同等规模的自回归离散模型。论文未提供吞吐量或延迟数据，也未讨论蒸馏或一步生成等加速手段，这限制了其在实时或高吞吐场景的可用性。
• **隐空间容量与训练复杂度的权衡**：为保留高容量可解码隐空间，模型引入了 Query-based Encoder-Decoder 和 Block-causal Diffusion Transformer，整体架构比典型连续扩散 LM（如使用 VAE 瓶颈的 D3PM）更为复杂。训练需要精心设计的噪声级别分布校准和自我轨迹一致性策略，这可能增加了超参敏感性和计算开销，论文未提供消融实验证明各组件的单独贡献。
• **评测范围与主流离散模型的差距**：实验仅在 OpenWebText 和 XSum 上与同类连续/扩散 LM 对比，未与当前主流的离散自回归模型（如 GPT-2、LLaMA）在同一评测协议下对比困惑度或生成质量。由于连续扩散语言建模仍属新兴方向，其与成熟技术的实际差距尚不明确，这削弱了该工作在推动整体语言建模领域发展的定位价值。