← 返回 PaperDaily
大模型与智能体
苹果谷歌联手开源：16B扩散模型6层线性化，推理提速1.73倍
大模型推理贵、响应慢，一直是落地最扎心的痛点。苹果、谷歌DeepMind、哈佛这次联手给16B扩散语言模型LLaDA 2.1动了场小手术：6层注意力换成线性版本，只训练约60小时就换来最高1.7倍解码加速，代码已开源。扩散模型也能线性化，这波操作值得细品。
龙哥读论文
发布于 2026-08-14 21:47:13
阅读 3
查看原文
原论文信息如下：
扩散语言模型的速度瓶颈在哪里？ 扩散语言模型（dLLM）生成文本时，从一个被「掩码」遮住的序列出发，逐步揭开被遮住的token。这种「任意位置同时揭开」的机制天然适合并行解码，因此被视为大模型推理加速的重要路线。
近期大规模落地的扩散语言模型，普遍采用「块状半自回归解码」策略：在生成一个激活块内部的多个token时并行去噪，块与块之间按顺序生成。这样已生成的前置块可以借助KV缓存复用历史信息，兼顾并行吞吐和上下文连贯性。
但问题也随之而来：哪怕用上KV缓存，每个去噪步骤仍需和前面所有已提交的块做一次完整注意力计算。每生成一个新token，都要重新「翻旧账」，把历史块从头读一遍。随着序列变长，前缀注意力成本线性增长，成为推理过程中的算力黑洞。
正是瞄准这个痛点，来自Apple、Google DeepMind和Harvard University的研究团队提出疑问：既然线性注意力能在自回归模型里把每个token的注意力成本压缩成常数级别，能否把这套思路移植到扩散语言模型上来？线性注意力最早由Katharopoulos等人在2020年提出，核心是用核特征映射替换softmax中的指数核，把前缀信息压缩成固定大小的循环状态。这个想法在自回归模型上已被反复验证，但在扩散语言模型上却一直没人啃下来。
为什么没人做？因为扩散语言模型是双向注意力，每个位置都能看到序列前后所有token，这种左右互看的结构导致隐藏状态永远在变化，键和值无法像自回归那样稳定缓存。但苹果和谷歌的研究员偏不信邪，硬是把这件「不可能」的事办成了。
块混合注意力：鱼与熊掌兼得的设计 论文提出的核心方案叫「块混合注意力」。在当前正在去噪的激活块内部，保留精确的双向softmax注意力；在跨块上下文上，把所有已提交的块压缩成一个固定大小的线性注意力循环状态。相当于在一栋楼里，本单元住户互相串门走楼梯，去其他单元则直接坐电梯。
具体来说，假设当前要生成的块是Bₙ₊₁，前序已提交的块是B₀到Bₙ。对于激活块内的查询向量q，先计算块内softmax注意力的分子和分母，保留精确的注意力权重。同时，跨块读取使用一种叫Hedgehog的线性注意力核，它使用带符号的可学习特征映射来逼近softmax核的行为：把输入x分别通过可学习的权重矩阵W映射后取softmax，再把x和-x两条路径的结果拼接，形成维度为2F的正值特征向量。这样既能模拟softmax核的指数行为，又能表达点积相似度的正负结构。
当一个块被提交后，块内所有键和值通过特征映射累加到一个固定大小的循环状态里，即论文公式中的S和Z。S代表「特征的加权值累加」，Z代表「特征的累加计数」，维度固定，与已提交的token数量无关。查询时只拿当前块的q和这个固定状态做一次矩阵运算，即可一次性读取所有历史块的信息。相当于把所有历史写进一张小抄，查询时只需扫一眼，不用再翻整本书。
最终的注意力输出，是块内softmax分支和跨块线性分支的加权融合。论文没有把两个分支单独归一化再相加，而是共用一个分母。对每个查询头学习一个标量门控w，用w控制块内分支和跨块分支的注意力权重配比。共享归一化让两个分支争夺同一个「注意力预算」，避免独立归一化导致的总量失衡。这个设计让模型能灵活决定当前查询更信任本块上下文还是更依赖全局历史信息。
这种块混合注意力带来的直接收益是：跨块读取的耗时和内存占用不再跟序列长度挂钩，变成固定开销的O(Fd)矩阵乘法。无论生成2048个token还是100万个token，线性化后的注意力层在访问历史上下文时，时间和空间复杂度都是常数。这在长文本生成场景下是决定性的性能差异。
两阶段后训练：轻量改造预训练模型 有了结构设计，接下来是怎么把它塞进已预训练的扩散语言模型里。从头训练成本太高，论文借鉴了LoLCATs的两阶段线性化方案，并适配到扩散模型和块混合注意力的设定下。
第一阶段叫「注意力转移」，核心是让每个新替换的块混合注意力模块模仿原softmax注意力模块的输出行为。骨干网络完全冻结，只训练新引入的块混合参数——每个查询头的特征映射矩阵W和标量门控α。损失函数用均方误差，最小化块混合注意力输出和原始softmax注意力输出之间的差距。关键操作是：每个块混合层接收的隐藏状态不是前一个线性化层的输出，而是原始softmax教师模型产出的隐藏状态。这种「跨层教师强制」策略能隔离每层的近似误差，单独优化，避免误差在层间叠加干扰。
第二阶段叫「端到端适配」。局部注意力转移只保证单层输出靠谱，但多个混合层串联后层间误差可能累积放大。所以第二阶段使用原始掩码扩散目标对整个混合模型微调，在预训练骨干和块混合参数全部冻结的情况下，只训练挂载在注意力和MLP投影上的低秩适配器（LoRA）。冻结块混合参数可保留第一阶段练好的注意力行为，LoRA适配器负责修正深层分布偏移。第二阶段的损失函数是标准掩码扩散训练目标——对每个掩码位置，最小化预测干净token分布和真实token分布之间的交叉熵。
论文选用的底座模型是LLaDA 2.1-mini，一个16B级别的开源扩散语言模型。在20层注意力中，选择跨层均匀分布的子集进行线性化——第0、4、8、12、16、18层，共6层，占总层数30%。剩下14层保留softmax注意力。线性化层数越多，虽能压缩更多注意力计算，但也会引入更多近似误差。论文在消融实验中反复权衡后认为，6层是质量和效率的甜点区域。
训练数据和优化方案方面，两个阶段都使用公开的Tulu SFT数据集。每个样本截断到2048个token以内，动态批处理。掩码扩散的损坏过程只施加在助手回复的token上，提示词保持不动，与推理时的条件生成设置对齐。整个训练约60小时：阶段1约24小时，阶段2约6小时，在两张NVIDIA L40S GPU上完成。
这个训练成本在当下大模型优化语境下「便宜得不像话」。对比从头预训练16B模型动辄数百万美元的计算账单，用不到三天的单机训练换来1.7倍推理加速，这笔账怎么算都划算。
实验结果：速度与质量的权衡 付出了6/20层近似误差的代价，换来的是否仅仅是吞吐量提升？准确率到底掉没掉？论文在编码、数学推理、科学推理三个维度的基准测试上给出了答案。
从表1看，LLADA-HYBRID在HumanEval上拿下72.0%，对比教师模型的75.6%有3.6个百分点回落。HumanEval+差距3.1个百分点。CMATH回落1.6个百分点，数学推理能力相对稳定。MBPP基本持平，MBPP+反而逆势上涨5.3个百分点，可能和LoRA微调带来的正则化效应有关。最明显的滑铁卢出现在GPQA-Diamond上，从38.9%跌到30.8%，掉了8.1个百分点。
龙哥认为这个结果基本符合预期：日常编码和数学任务对注意力细微变化不敏感，主要靠模式匹配和局部语义理解；但GPQA这种研究生级别的科学推理，需要在多个长距离事实之间交叉对照和深层次推理，对注意力权重精确性要求极高，线性近似在这里就露怯了。
再看表2的端到端吞吐量。LLADA-HYBRID在每一个并发级别上都实现了对softmax注意力教师的全面碾压。16路并发时，吞吐量从1845.7 token/s提升到2761.7 token/s，加速比1.50倍。128路并发时达到峰值，从2310.2提升到3994.4 token/s，加速比1.73倍。256路并发时双方都开始遇到系统瓶颈，加速比依然维持在1.60倍。
进一步深入内核层的孤立测试，固定状态的优势体现得更加淋漓尽致。论文附录提供Per-block内核延迟对比表，展示当前缀长度P从短到长变化时，教师模型和混合模型的单块内核延迟差异。当P=256时，教师模型延迟4.42微秒，混合模型3.22微秒，差距不大。但当P增长到2048时，教师模型延迟飙升至13.66微秒，而混合模型因只读取固定大小的「小抄」状态，延迟只增加到3.67微秒。教师模型开销随前缀长度线性膨胀，混合模型几乎原地不动。
内存方面，论文还专门测了固定内存预算下的最大并发批次。由于线性化层的循环状态是固定大小的，每个请求的显存占用被大幅压缩，因此在同样显存预算下，混合模型可以塞进更多并发请求。这张表直接证明了固定大小循环状态在大规模推理服务中的「内存红利」。KV Cache不再随每个已提交token数量线性膨胀，服务端就能在同样显存限制下服务更多用户，在高并发API推理场景中就是实打实的降本增效。
消融实验部分，论文直接回答了「为什么要线性化6层而不是8层或者4层」。结果显示，当线性化层数从6层增加到8层时，HumanEval的pass@1成绩会从72.0%快速下滑到更低水平。这说明在当前的训练配方下，6层已经是大规模线性化的一个临界点。再往上加层，速度和内存优势还会增加，但质量损失将变得不可接受。论文在质量和效率的十字路口，选择了稳妥的右侧路线。
局限与未来展望 客观来讲，这篇论文还有很多未尽的事情。最明显的局限在于，它只线性化了20层中的6层，剩下14层仍保留序列依赖的注意力计算和KV缓存开销。因此，虽然跨块注意力访问变成常数级，但块内attention的复杂度仍和块大小相关，长期来看仍是一个不可忽视的瓶颈。
论文目前的评测集中在中等生成长度（2048 token以内），恰好和训练时使用的截断长度一致。如果在更长的上下文（比如10万token）下做推理，线性化层的固定状态是否还能稳定工作、注意力转移是否依然有效，都是未知数。论文自己也承认了这一限制，并指出需要专门的注意力转移策略或更长序列的训练方法来应对。
另外，当前模型假设推理时的块大小是固定的，如果实际推理时块大小变化，模型的注意力行为可能出现不稳定。让模型在训练时就适应多种块大小，或者发展出对块大小不敏感的架构，是未来可以探索的方向。更激进的层选择策略，或者自适应地决定哪些层线性化、哪些层保留全注意力，同样值得研究。
还有一点值得注意的是，论文的推理服务实现是在SGLang框架中，通过一个Triton内核把块内softmax分支、线性读出、门控融合和状态更新融合到一次内核启动中完成。这个工程实现层面的优化是论文能获得1.5-1.7倍实际端到端加速的重要原因。但这也意味着，这套方法的核心代码深度耦合了SGLang的调度器和分页KV缓存机制。如果其他推理框架想复现这个收益，需要做相应的工程移植，这本身也是一笔不小的成本。
从更宏观的视角看，这篇论文的研究思路其实和目前AI圈子里流行的「事后优化」趋势一脉相承。与其花天价从头预训练一个新架构，不如研究怎么把已经训练好的超大模型低成本地改造成更高效的形态。LoRA微调、稀疏化、蒸馏、量化，以及这次的线性注意力改造，本质上都是在这个思路上做文章。苹果和谷歌的这次合作清楚地传递了一个信号：扩散语言模型的推理加速，不一定要靠堆硬件，算法层面的「小手术」同样可以起到立竿见影的效果。
龙迷三问 这篇论文到底在解决什么问题？ 苹果、谷歌DeepMind与哈佛大学联合提出块混合注意力，把线性注意力成功植入预训练16B扩散语言模型LLaDA 2.1。这篇工作最值得看的点是什么？ 在6/20层线性化配置下，HumanEval从75.6%降至72.0%，MBPP+从57.7%提升至63.0%，CMATH从88.3%降至86.7%，GPQA-Diamond从38.9%降至30.8%；解码吞吐量提升1.50-1.73倍，最大并发批处理量从544提升至704。这篇工作的边界或风险在哪里？ 优点：提出新颖的块混合注意力机制，有效结合softmax和线性注意力的优势；两阶段后训练策略高效，仅需约60小时和少量数据即可完成改造；推理加速显著，内存占用大幅降低。缺点：仅线性化6/20层，仍有大量层保留序列长度相关的注意力计算；在GPQA-Diamond等困难推理任务上性能下降明显；当前评估限于中等生成长度。如果你还有哪些想要了解的，欢迎在评论区留言或者讨论~
龙哥点评 论文创新性分数： ★★★☆☆
块混合注意力的思路简洁有效，将线性注意力引入扩散语言模型，是领域内首次尝试。但从技术层面看，主体思想是对LoLCATs和Hedgehog的组合适配，原创性更多体现在「怎么把已有工具安到新场景」上。实验合理度： ★★★★☆
实验设计覆盖了质量、吞吐量和内核延迟三个维度，基准选择也较为全面。训练数据只用Tulu SFT单数据集，缺乏多数据集对比。消融实验对层数的影响做了探究，但对门控机制的贡献没有单独剥离分析。学术研究价值： ★★★★☆
论文证明了一个重要的可行性边界：预训练的扩散语言模型可以通过低成本后训练实现线性化加速。这个结论对后续研究者有很强的启示，打开了「事后改造」的新方向。稳定性： ★★★☆☆
在中短长度生成上表现稳定，但长序列和复杂推理任务上的表现有明显下滑。GPQA-Diamond上8.1个百分点的性能损失是不可忽视的短板，说明在高度复杂的推理场景中仍不够稳定。适应性以及泛化能力： ★★★☆☆
目前只在LLaDA 2.1-mini单一模型上验证，且训练和评测都集中在中短序列。对于其他架构（如MoE类扩散模型）是否同样有效，缺乏证据。推理块大小也是预设固定的，实际场景中动态变化时性能尚未验证。硬件需求及成本： ★★★★☆
训练只需两张L40S GPU约60小时，对工业界来说几乎是零门槛。推理阶段虽然需要Hopper架构的Tensor Core支持，但这是目前主流数据中心GPU的标配，不算苛刻条件。复现难度： ★★★★☆
代码已在GitHub开源（https://github.com/Diuven/LLaDA-Hybrid），且训练数据用的是公开的Tulu SFT。不过，实现细节深度耦合SGLang框架和Triton内核，复现时需要对这两个框架有深入的工程理解。产品化成熟度： ★★★☆☆
对于中短长度的指令跟随生成场景（如聊天机器人、代码补全），已经具备初步产品化条件。在复杂推理场景（如科学问答、多跳推理）中，质量损失还不足以支撑直接落地。长文本生成场景需要进一步验证。可能的问题： 6/20层的线性化比例仍偏保守，更激进的比例是否会带来更大幅度的加速值得探索。GPQA-Diamond上8.1个百分点的脱落说明近似误差对复杂推理的伤害不容忽视。训练仅用Tulu SFT单数据集，数据的多样性不足，可能导致模型在特定领域上的泛化能力受到影响。论文对门控机制的贡献、对块大小的敏感性等关键设计维度缺少系统的消融分析。
主要参考文献
[1] Kim J, Roh Y, Kim J. Retrofitting Linear Attention into Diffusion Language Models[J]. arXiv preprint arXiv:2608.06628, 2026. [2] Bie T, Cao M, Chen K, et al. LLaDA2.1: Speeding up text diffusion via token editing[J]. arXiv preprint arXiv:2602.08676, 2026. [3] Zhang M, Bhatia K, Kumbong H, et al. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry[C]. ICLR, 2024. [4] Zhang M, Arora S, Chalamala R, et al. LoLCATs: On low-rank linearizing of large language models[C]. ICLR, 2025. [5] Katharopoulos A, Vyas A, Pappas N, et al. Transformers are RNNs: Fast autoregressive transformers with linear attention[C]. ICML, 2020. [6] Lambert N, Morrison J, Pyatkin V, et al. TULU 3: Pushing frontiers in open language model post-training[J]. arXiv preprint arXiv:2411.15124, 2024. [7] Zheng L, Yin L, Xie Z, et al. SGLang: Efficient execution of structured language model programs[C]. NeurIPS, 2024. [8] Arriola M, Gokaslan A, Chiu J T, et al. Block diffusion: Interpolating between autoregressive and diffusion language models[C]. ICLR, 2025. *本文仅代表个人理解及观点，不构成任何论文审核或者项目落地推荐意见，具体以相关组织评审结果为准。欢迎就论文内容交流探讨，理性发言哦～ 想了解更多原文细节的小伙伴，可以点击 "阅读原文"， 查看更多原论文细节哦！
