Sumi：从头开始的开放均匀扩散语言模型 Sumi: Open Uniform Diffusion Language Model from Scratch
首个大规模开源均匀扩散语言模型，7B参数训练1.5T token
前置知识
扩散模型
扩散模型通过逐步去噪过程生成数据，从纯噪声开始，每步预测并去除部分噪声，最终得到清晰的输出。语言扩散模型将这个思想应用于文本生成，与自回归模型逐个生成token不同，扩散模型可以同时更新所有token位置。这个过程通过定义前向扩散过程（添加噪声）和反向去噪过程（预测噪声并去除）来实现，训练目标通常是去噪得分匹配或变分下界优化。
理解扩散模型是理解本文的基础，因为Sumi正是基于扩散范式而非传统的自回归范式。扩散模型的核心思想是逐步从噪声中恢复有意义的内容，这与自回归模型按顺序生成完全不同，掌握这个范式转换对于理解Sumi的设计动机和训练目标至关重要。
均匀扩散语言模型（UDLM）
UDLM是扩散语言模型的一种，其特点是允许任何token在任何去噪步骤被更新。这与掩码扩散语言模型（MDLM）不同——MDLM一旦填充了某个掩码位置，该位置就永远固定无法修改。UDLM的这种灵活性理论上支持更灵活的生成和自我修正，因为模型可以在任何步骤重新考虑任何位置的内容，而不是一旦做出决定就不可撤销。UDLM的训练目标通常最小化ELBO，通过在不同噪声水平下预测原始token来实现。
这是本文的核心范式，Sumi就是UDLM，理解它与MDLM和AR模型的区别是读懂本文的关键。均匀扩散代表了扩散语言模型的一种理想形式，它提供了最大的灵活性，但这种灵活性是否真正转化为更好的性能是一个开放问题，本文通过大规模实验开始回答这个问题。
广义插值离散扩散（GIDD）
GIDD是离散数据的扩散框架，它通过在干净数据和纯噪声之间进行插值来构造扩散过程。关键参数log-SNR控制每步的信号噪声比，限制在[-9, 9]范围内，这对应于从几乎纯噪声（-9）到几乎无噪声（9）的扩散程度。GIDD提供了一个统一的框架来参数化不同的扩散调度，训练目标是最小化ELBO（证据下界），这可以通过去噪得分匹配等价实现。GIDD的改进形式直接用SNR参数化ELBO，提供了更好的数值稳定性和解释性。
Sumi基于GIDD框架训练，理解其数学形式和目标函数有助于理解训练过程。GIDD为离散数据提供了一个优雅的扩散框架，通过SNR参数化可以更好地控制扩散过程，这对于训练稳定性和生成质量控制都很重要。理解GIDD的数学基础有助于理解Sumi的训练细节和超参数选择。
Canvas长度 vs 生成长度
Canvas长度是开始生成时分配的token位置总数，即模型填充的缓冲区大小。生成长度是实际进行评分的位置数。对于likelihood任务，超出答案的位置填充随机token；对于生成任务，只允许在生成长度内更新token，而超出部分保持在随机初始化状态。论文在生成区域和尾部随机token之间插入EOS/BOS边界，以匹配训练时看到的打包分布。Canvas长度的选择对生成质量有重要影响，论文发现模型在其训练的序列长度范围内表现最佳，超出范围会明显退化。
论文发现生成流畅度对canvas长度敏感，理解这个概念是理解实验结果的关键。Canvas长度是扩散模型特有的概念，与自回归模型的上下文长度不同，它不仅限制了最大输出长度，还影响了整个扩散过程的动力学。论文的探索性分析表明，canvas长度需要与训练序列长度匹配才能获得最佳性能，这对于实际使用扩散语言模型有重要指导意义。
研究动机
当前扩散语言模型领域存在明显空白。掩码扩散语言模型如LLaDA已经扩展到8B参数和2T以上训练token，达到了与强AR基线竞争的性能。但均匀扩散语言模型尚未在参数规模和token预算上都达到大规模水平。最大的现有模型是在相对较小token预算上训练的计算最优检查点，唯一在数据丰富环境下训练的模型只有1.7B参数。DiffusionGemma虽然应用了均匀扩散，但它是通过适配预训练的AR模型而非从头训练。这意味着大规模均匀扩散模型在数据丰富环境下的行为仍然未被探索，社区缺乏研究其缩放行为、生成动力学、可控性以及与AR和掩码扩散模型权衡的参考点。更重要的是，没有完全开源的大规模UDLM，限制了社区对这一范式的深入研究。
本文的目标是本文的目标是填补这个空白：开发并发布一个完全开源的7B参数均匀扩散语言模型，从头开始预训练，使用1.5T token。这个模型将作为研究大规模原生均匀扩散行为的干净参考点，帮助社区研究其缩放行为、生成动力学、可控性以及与已建立的AR和掩码扩散模型的权衡。除了模型本身，论文还报告了一组探索性推理时间分析，研究生成任务的采样策略、canvas长度敏感性、并行解码效果等，为理解原生训练的均匀扩散模型提供初步证据。论文承诺完全开源：模型权重、中间检查点、完整训练配方，包括公开语料库上数据混合的完整规格，以便社区可以重建功能等价的训练语料库。
与已有工作不同的是，本文的独特切入角度是原生训练——不是适配AR模型，而是从零开始用扩散目标训练。这能研究真正的均匀扩散行为，而非受AR预训练遗留影响。AR预训练的模型可能隐式学习了顺序偏好，这会混淆对均匀扩散固有行为的理解。Sumi的原生训练避免了这个问题，提供了一个干净的研究对象。此外，本文承诺完全开源：模型权重、中间检查点、完整训练配方，包括公开语料库上数据混合的完整规格。这与只发布部分或完全不发布训练数据的AR模型形成对比。完全开放使得社区可以重建训练数据，研究数据组成的影响，这对于理解模型行为和推动可重复研究至关重要。
核心方法
Sumi采用GIDD框架的SNR重参数化版本，在纯均匀噪声下训练。模型架构基于LLaMA风格：36层、隐藏维度4096、SwiGLU MLP（FFN大小12288）、分组查询注意力（32头，8KV组，头维度128）、RMSNorm、未绑定的输入/输出嵌入、无偏置或dropout。注意力中使用off-by-one softmax来缓解attention sink问题，这是一个常见的长上下文模型问题，模型倾向于过度关注初始token。训练分为三个阶段：预训练约1.3T token（序列长度1184），然后两个中期训练阶段共250B token（第一个130B在原长度1184，第二个120B扩展到4864以支持更长上下文）。整个流程使用单一WSD学习率调度器：2000步预热到2e-4，然后恒定，最后2000步线性冷却到2e-5，覆盖所有三个训练阶段。这种调度确保了训练稳定性，论文报告训练损失和基准监控分数都保持稳定。
核心创新是首次实现大规模原生均匀扩散语言模型。与DiffusionGemma适配AR模型不同，Sumi完全从零开始用扩散目标训练。这避免了AR预训练带来的隐式顺序偏差，让模型真正学习均匀扩散的动态。技术关键点包括：使用GIDD的SNR重参数化形式，log-SNR限制在[-9, 9]以覆盖从纯噪声到无噪声的完整范围；未绑定的输入/输出嵌入，允许模型学习更好的token表示；off-by-one softmax缓解attention sink，这是长上下文模型的重要技术；训练时无注意力掩码（为优化吞吐量，使用最快的可用内核不支持自定义掩码），生成时用EOS/BOS边界分隔生成区域和尾部随机token作为workaround。这些设计选择都针对训练稳定性和效率，同时保持生成质量。
方法步骤详情
训练分为三个阶段。第一阶段（预训练）：在约1.3T token上训练，序列长度1184，全局批量大小4608序列（每步约5.5M token），使用bfloat16。数据来源是llm-jp-corpus-v4的英语部分（排除FineWeb子集）、StarCoder代码子集、swallow_code_v2的Python代码，以及en_fineweb-rescored（用Qwen3-32B蒸馏的轻量级分类器按教育分数重排序，保留高教育分数的文档）。优化器使用AdamW，参数(0.9, 0.95)，权重衰减0.1，梯度裁剪1.0，辅助z-loss系数1e-5。第二阶段（中期训练1）：130B token，保持原序列长度1184。第三阶段（中期训练2）：120B token，扩展序列长度到4864，批量大小调整为1152序列以保持每步约5.6M token的token批量大小。数据混合调整：downsample nemotron_pretraining_code_v2从491B到63B避免代码过权重，添加en_megamath-web-pro-max-oss（按math_score过滤到48.9B token，保留最高分数的数学文档），推理数据只保留约4096 token以内的样本以匹配模型上下文长度。硬件使用288个NVIDIA H100 GPU，总计算43,308 GPU-hours。
技术新颖性
技术新颖性体现在多个方面。首先是规模：这是首个在参数规模（7B）和token预算（1.5T）上都达到大级别的原生均匀扩散模型。之前的UDLM要么参数小，要么训练数据少，Sumi填补了这个空白，提供了研究大规模原生均匀扩散行为的第一个参考点。其次是完全开源：不仅模型权重和检查点，还包括完整训练配方和数据混合规格，这与大部分AR模型形成对比。完全开源使得社区可以复现和扩展工作，这对于推动研究至关重要。第三是探索性推理时间分析：论文研究canvas长度敏感性、置信度采样对token提交顺序的影响、并行解码效果、以及显式修订预算是否能带来自我修正。这些分析为理解原生均匀扩散的行为提供了初步证据，虽然样本量小（每个任务30个问题）且是方向性而非结论性的，但它们指出了值得进一步研究的问题。
实验结果
论文在13个基准测试上评估Sumi，分为四类：通用知识（MMLU、RACE、TruthfulQA）、推理与数学（ARC-Easy/Challenge、GPQA、BIG-Bench Hard、GSM8K）、编码（HumanEval、MBPP）、常识（WinoGrande、PIQA、HellaSwag）。在相同协议下，与Falcon-7B（1.5T token）、Llama 2-7B（2T token）、OLMo-7B（2.5T token）比较。Sumi在通用知识和编码上达到最佳分数：MMLU 51.1%（vs Falcon 27.2%、Llama 2 46.0%、OLMo 28.0%），HumanEval 22.6%（vs Falcon 0%、Llama 2 12.8%、OLMo 13.4%），MBPP 26.6%（vs Falcon 12.4%、Llama 2 23.2%、OLMo 21.4%）。在推理和数学上，Sumi与Llama 2-7B竞争，大多领先Falcon-7B：GSM8K 32.8%（vs Falcon 5.3%、Llama 2 13.5%、OLMo 3.8%），ARC-Easy 70.0%（vs Falcon 70.8%、Llama 2 73.8%、OLMo 68.8%），ARC-Challenge 43.0%（vs Falcon 43.2%、Llama 2 45.1%、OLMo 40.3%），BBH 31.8%（vs Falcon 27.1%、Llama 2 39.6%、OLMo 29.8%），GPQA 26.1%（vs Falcon 24.6%、Llama 2 24.3%、OLMo 24.8%）。在常识上，Sumi是评估模型中最弱的：PIQA 66.4%（vs Falcon 80.5%、Llama 2 78.7%、OLMo 79.8%），HellaSwag 60.0%（vs Falcon 76.3%、Llama 2 76.2%、OLMo 75.6%），WinoGrande 60.0%（vs Falcon 71.6%、Llama 2 74.7%、OLMo 71.3%）。论文认为教育密集的数据混合是这种差距的可能原因，虽然教育过滤被观察到改善知识和推理密集型基准，但会降低像HellaSwag和PIQA这样的常识基准。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MMLU | accuracy (5-shot) | 51.1 | Llama 2-7B: 46.0, Falcon-7B: 27.2, OLMo-7B: 28.0 | vs Llama 2: +11.1%, vs Falcon: +23.9% |
| HumanEval | accuracy (0-shot) | 22.6 | Llama 2-7B: 12.8, Falcon-7B: 0.0, OLMo-7B: 13.4 | vs Llama 2: +9.8%, vs OLMo: +9.2% |
| GSM8K | accuracy (4-shot) | 32.8 | Llama 2-7B: 13.5, Falcon-7B: 5.3, OLMo-7B: 3.8 | vs Llama 2: +19.3%, vs Falcon: +27.5% |
| PIQA | accuracy (0-shot) | 66.4 | Llama 2-7B: 78.7, Falcon-7B: 80.5, OLMo-7B: 79.8 | 落后于所有基线，vs Llama 2: -12.3% |
局限与改进
作者承认的局限性：推理时间分析是方向性而非结论性的，每个任务只采样30个问题，使用Falcon-7B困惑度作为粗糙流畅度代理，没有运行与掩码扩散和自回归模型的匹配比较来归因这些行为到均匀扩散范式本身而非Sumi特定。常识差距过大，不能完全由教育和代码密集的数据混合解释，且没有直接测试这种数据组成归因。Sumi作为预训练基础模型发布，没有经过指令调优、对齐或安全过滤，因此继承了此类模型的常见风险：对抗性或粗心提示可以引出有害、冒犯或其他敏感文本，也可非预期地出现此类输出，例如反映预训练语料库中的偏见。Sumi没有确保事实准确性的机制，可能以明显置信度陈述虚假信息。作者将其发布以支持研究而非直接部署，鼓励任何基于它构建的用户权衡自己设置的风险并独立验证事实声明。这些局限性需要在未来的工作中解决，特别是探索不同数据组成的影响和添加安全过滤。
独立分析的弱点
常识性能明显落后是主要弱点，这可能与过度教育过滤的数据混合有关。改进方向：测试不同数据组成（如增加更多对话式、日常文本）对常识基准的影响，或针对常识任务进行针对性微调。论文没有直接测试数据组成归因，这是一个明显的实验缺口。GSM8K对并行解码极度敏感（每步2个token就立即损失精度），而编码任务可以承受4个token/步，这表明算术推理对顺序高度依赖。改进方向：为算术任务设计特殊的采样策略，如强制顺序生成数值部分，或者混合策略——对顺序敏感的部分强制顺序，其他部分保持并行。修订预算没有带来自我修正（最多1%的最终token与第一次提交不同），主要是A到B再到A往返而非定向编辑。改进方向：设计针对可能错误的修订设置，而非通用过去噪，或引入明确的自我验证机制，如让模型检查自己的输出或提供置信度评分。Canvas长度敏感性：GSM8K在训练范围外困惑度急剧上升，其他任务在短canvas上下降更温和。改进方向：训练时使用更多样化的序列长度，或动态调整canvas长度，或者开发自适应canvas长度选择策略。
未来方向
作者提出：准备指令调优版本的Sumi并在未来更新中发布。指令调优可以显著提升模型的实用性和用户体验，使其更适用于实际任务。缺乏通用修订预算下的自我修正留下了开放问题——针对可能错误设计的设置是否能恢复它。这需要设计更智能的修订策略，可能需要额外的训练信号或目标。更广泛地说，匹配评估协议下的控制比较将澄清Sumi的哪些生成行为是均匀扩散内在的，哪些是Sumi特定的，这是本文提出但未解决的问题。基于成果可延伸的方向：(1) 研究置信度采样诱导的自组织提交顺序是否适用于更广泛任务，或可设计更显式的顺序策略。(2) 探索canvas长度敏感性的根本原因——是训练分布偏差还是扩散范式固有，这可能需要对训练过程进行更细致的分析。(3) 为需要强顺序的任务（如算术）设计混合策略，在部分区域强制顺序而在其他区域保持并行，这可能显著提升算术推理的性能。(4) 研究如何有效利用均匀扩散的自我修正能力，可能需要不同的目标或训练信号，如引入显式的错误检测和修正模块。
复现评估
复现性评估：完全开源，模型权重、中间检查点、完整训练配方都已发布，包括公开语料库上数据混合的完整规格。算力需求高：预训练35,776 GPU-hours，两个中期训练阶段共7,531 GPU-hours，总计43,308 GPU-hours在288个NVIDIA H100 GPU上。数据来源都来自公开语料库：llm-jp-corpus-v4、StarCoder、swallow_code_v2、en_fineweb-rescored、llm-jp-corpus-midtraining-v2、en_megamath-web-pro-max-oss。论文详细描述了数据选择过程和结果混合（图1），包括每个数据源的token数量和比例。评估使用lm-evaluation-harness，修改只支持基于扩散的评分。技术细节充足：架构参数（层数、隐藏维度、注意力配置）、优化器设置（AdamW参数、权重衰减、梯度裁剪、z-loss系数）、学习率调度（WSD调度、预热、冷却）、训练序列长度、批量大小等都有详细说明。复现难度中高：需要大量算力（4万多GPU-hours），但所有必要信息都已提供，且数据都是公开的，理论上可以重建功能等价的训练语料库。完全开源的承诺使得复现比大多数闭源模型容易得多。
论文图表
