FLARE: Diffusion for Hybrid Language Model¶

会议: NeurIPS 2026
arXiv: 2606.01774
代码: https://github.com/yuchen-zhu-zyc/FLARE
领域: LLM效率
关键词: 扩散大语言模型, 混合注意力架构, Gated DeltaNet, 状态调度, 投机解码

一句话总结¶

FLARE 提出了首个面向混合注意力（Softmax 与 Gated DeltaNet 线性注意力）大模型的扩散转换框架，通过因果/双向 token 平衡双流目标、寄存器级硬件状态调度核函数以及统一推理架构，仅用 10B token 微调就在 2B/4B/9B 尺度上保留了 Qwen3.5 原生能力，单卡并发吞吐较开源 dLLM 基线最高提升 4.8 倍。

研究背景与动机¶

自回归（Autoregressive, AR）大语言模型虽然在各类推理与生成基准上取得了巨大的实际成功，但逐 token 串行解码（sequential decoding）的固有范式成为了低延迟与实时交互部署的核心瓶颈。针对推理效率的优化长期沿着两个平行的技术路线展开：其一是架构侧优化，即通过线性注意力（Linear Attention）、状态空间模型（SSM）以及软最大-线性混合注意力（Hybrid Attention，如 Qwen3.5、Nemotron-H、MiniMax 等）压缩 KV 缓存并降低单步前向传播开销；其二是生成侧优化，即利用非自回归扩散语言模型（dLLM）打破串行因果限制，通过分块并行迭代去噪（Iterative Denoising）在更少的前向调用轮次中输出更多 token。

然而，将这两项互补的效率红利结合在一起面临严重的理论与工程壁垒。一方面，以往的自回归迁移扩散模型（AR-to-dLLM）往往伴随着严重的基模能力崩塌（例如代码生成与长链推理大幅退化），且业界对于退化根源是损失函数设计、掩码机制还是数据分布一直缺乏隔离诊断；另一方面，混合注意力骨干网络中存在定长的循环隐藏状态（如 Gated DeltaNet 的 \(S_t\)）以及用于局部特征提取的因果卷积（ShortConv），无法像标准 Softmax Transformer 那样直接通过 2D 稠密注意力掩码实现块级双向扩散能见度，强行展开掩码会引发 \(O(L/B)\) 的显存爆炸与核函数调度崩溃。

针对上述瓶颈，本文从数据、算法设计、硬件底层核函数到服务系统进行了全栈重构。通过控制变量排查发现迁移微调数据的质量与分布匹配是决定能力保留的最关键因素，进而提出了硬件感知的混合注意力状态调度方案与统一的双解码机制。核心 idea：构建因果干净流与块双向噪声流 token 平衡的联合训练目标，通过寄存器级重构跨步检查点解决线性注意力递归状态继承难题，并在单套模型权重上统一支持扩散并行去噪与自回归投机验证双解码路径。

方法详解¶

整体框架¶

FLARE 的整体流水线涵盖输入样本的文档打包、双流前向传播、底层混合状态调度以及统一推理 Serving。在训练阶段，输入序列的提示词（Prompt）作为纯因果干净上下文，而回复词（Response）被切分为大小为 \(B\)（默认 \(B=4\)）的连续块。模型构造出两条互补的信息流：黄色因果干净流负责预测下一个 token 并提供 \(L_{\mathrm{AR}}\) 监督，橙色块级双向噪声流则基于两组互补掩码视图 \(\mathcal{M}_b\) 与 \(\mathcal{M}_b^c\) 进行掩码去噪，提供 \(L_{\mathrm{diff}}\) 监督；在 Softmax 层通过文档阻断的因果-双向组合掩码实现可见性，在线性注意力层则通过硬件状态重构使每个噪声块精准继承前序干净边界状态；在推理阶段，同一套权重可自由切换以扩散置信度为核心的 Diffusion-Trust 模式或以因果 logits 为验证基准的 AR-Trust 投机解码模式。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入样本<br/>文档打包与 Prompt/Response 切块"] --> B["Token 平衡干净/噪声双流监督<br/>全序列因果流与互补双视图噪声流"]
    B --> C["混合注意力状态调度核函数<br/>寄存器级重构前序干净边界状态"]
    C --> D["硬件感知前向与反向计算<br/>GDR 融合两流与 ShortConv 边界防泄漏"]
    D --> E{"统一推理解码引擎"}
    E -->|Diffusion-Trust| F["并行块迭代去噪<br/>置信度门限提交与因果回放"]
    E -->|AR-Trust| G["自回归投机验证<br/>噪声草稿并行生成与干净流精确检验"]

关键设计¶

1. Token 平衡干净/噪声双流目标：无损保留因果语义与扩散去噪能力

以往纯块扩散微调会导致模型丧失自回归因果先验，引发基准测试断崖式下跌。FLARE 设计了单位权重的 token 平衡双流目标。设训练样本序列切为 \(K\) 个大小为 \(B\) 的块，对每个块采样掩码集合 \(\mathcal{M}_b\)，构造互补的两个噪声视图 \(\tilde{\mathbf{x}}_b\) 和 \(\tilde{\mathbf{x}}_b^c\)。模型在前向传播中同时保留未加噪声的因果流与被掩码的噪声流：

\[\mathcal{L}_{\mathrm{FLARE}}(\theta) = \mathcal{L}_{\mathrm{AR}}(\theta) + \mathcal{L}_{\mathrm{diff}}(\theta) = \sum_{\ell=1}^L \log p_\theta(x^\ell \mid \mathbf{x}^{<\ell}) + \sum_{b=1}^K \left[ \sum_{\ell \in \mathcal{M}_b} \log p_\theta(x^\ell \mid \tilde{\mathbf{x}}_b, \mathbf{x}^{<b}) + \sum_{\ell \in \mathcal{M}_b^c} \log p_\theta(x^\ell \mid \tilde{\mathbf{x}}_b^c, \mathbf{x}^{<b}) \right]\]

因 \(\mathcal{M}_b \cup \mathcal{M}_b^c\) 恰好对每个块完成无交集的完全分割，每个 token 在全局严格贡献 1 次自回归因果监督和 1 次扩散去噪监督。同时，噪声流的预测输出引入了向左 logit shift，精确对齐自回归语义与目标词位置，避免在块边界浪费预测能力。

2. 混合注意力状态调度核函数：寄存器级重构化解显存爆炸

在混合架构（如 Qwen3.5）中的 Gated DeltaNet（GDN）层，隐状态更新遵循门控 Delta 规则：\(\mathbf{S}_t = \alpha_t(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)\mathbf{S}_{t-1} + \beta_t \mathbf{k}_t \mathbf{v}_t^\top\)。为了让噪声块 \(b\) 看到前序因果干净上下文同时在块内保持双向可见，噪声流状态必须初始化自干净边界状态 \(\tilde{\mathbf{S}}_{e(b-1)} := \mathbf{S}_{e(b-1)}\)，且噪声块内所有位置读取块尾状态 \(\tilde{\mathbf{o}}_\ell = \tilde{\mathbf{S}}_{e(b)}^\top \tilde{\mathbf{q}}_\ell / \sqrt{d_k}\)。若采用朴素的 Chunk-then-Refine（Route I），需要在高带宽内存（HBM）中固化全序列所有块边界状态，产生随 \(L/B\) 扩展的巨大张量，在 \(B=4\) 时导致显存溢出。FLARE 提出了融合双流调度方案（Route II）：在前向传播中仅保存步长为 \(S\) 的干净状态跨步检查点（Strided Checkpoints），其余状态在 GPU 寄存器中直接通过小步演化实时重构；反向传播中将块级梯度直接注入因果递归主链，使 \(B=1\) 时的峰值显存从 18.14 GiB 骤降至 0.45 GiB，训练吞吐达到与纯自回归基线一致的 24.81% MFU。

3. 短卷积边界防泄漏与文档打包硬隔离：杜绝 cross-document 污染

混合骨干网络在进入线性注意力前依赖宽度为 \(W\)（通常 \(W=4\)）的一维因果短卷积（ShortConv）提取局部特征。噪声流在块起始处如果直接执行卷积，其滑动窗口会跨越边界读取前序噪声数据而非干净历史。FLARE 编写了融合 ShortConv Triton 核函数，当卷积滞后量跨过块起点 \(\ell_{\mathrm{blk}}\) 时，动态从干净流输入 \(x^{\mathrm{c}}\) 取数，在块内部则读取噪声流输入 \(\tilde{x}\)。此外，针对变长文档打包（Document Packing），FLARE 强制要求所有文档边界与扩散块大小 \(B\) 对齐，并在核函数内植入 cu_seqlens 警戒条件，在文档首块硬置零噪声初始状态并阻断跨文档梯度累加，确保打包微调与单样本训练数学等价。

4. 统一推理系统与状态回滚：一套参数并行兼顾两类解码机制

在部署侧，FLARE 基于 SGLang 构建了统一Serving框架。针对以扩散置信度为主的 Diffusion-Trust，模型执行多步块级并行去噪；去噪过程中的中间前向不写入全局因果隐状态池，待整个块达到置信度门限 \(\gamma_s\) 并完全确定后，执行一次轻量级的因果重放（Causal Replay），将最终状态写入池中。针对投机解码模式 AR-Trust，噪声流充当 Draft 生成器，干净流负责并行验证；针对部分 token 接受（接受前 \(r\) 个 token，\(r < K\)）的场景，传统线性注意力需要重新进行矩阵乘法回放，而 FLARE 提出了 Cache-and-Scatter 机制，在验证前向阶段将 \(K\) 个候选步的中间状态以访存写写入缓冲区，验证判定后通过单个 Triton 核函数执行聚合分散（Gather-Scatter），实现零算力开销的即时状态回滚。

实验关键数据¶

主实验¶

主实验在 12 个涵盖数学推理、通用知识、指令遵循及代码编写的基准测试上全面评估了 FLARE-9B，并与参数量相近甚至更大的主流开源扩散语言模型（LLaDA 系列、SDAR）及商业扩散模型 Mercury-2 进行对比。

基准测试	评估领域	Mercury-2 (商业 API)	SDAR-30B-A3B (MoE)	LLaDA-2.1-flash (100B-A5B)	FLARE-9B (Diffusion-Trust)	FLARE-9B (AR-Trust)	Qwen3.5-9B (源自回归基线)
ARC-C	知识推理	—	93.20	—	95.65	96.33	97.70
MMLU	通用知识	—	82.80	—	80.75	84.80	88.21
MMLU-Pro	复杂知识	—	61.50	75.31	74.73	77.39	81.39
GPQA-D	专家级推理	73.00	36.70	66.67	64.65	71.21	80.30
IFEval	指令遵循	—	60.60	83.36	63.22	71.35	91.31
GSM8K	小学数学	90.62	91.40	—	93.10	93.33	89.16
MATH-500	竞赛数学	81.00	77.80	—	93.60	95.20	96.60
AIME-24	高难度竞赛	51.10	16.70	—	60.00	63.33	65.56
AIME-25	高难度竞赛	—	10.80	63.33	53.33	54.44	60.00
HumanEval	代码生成	—	87.20	—	82.32*	92.07	95.12
MBPP	代码生成	—	71.60	—	82.10*	91.05	89.11
LCB v6	竞技代码	67.30	21.70	44.05	5.71*	49.71	49.71

(注：Diffusion-Trust 在部分严格格式的代码提取基准中受限于块级非因果语法验证与脚本截断，标 * 项需保守解读)

在更小尺寸（2B/4B）上，FLARE 同样展现出碾压同量级开源扩散模型的性能：FLARE-2B 在 MATH-500 上达到 84.40%（SDAR-1.7B 仅 63.20%），AIME-24 达到 31.11%（SDAR-1.7B 仅 10.00%）；FLARE-4B 在 MATH-500 上达到 94.20%，远超 30B 级别的 SDAR-30B-A3B（77.80%）。

消融实验¶

针对 AR-to-dLLM 算法组件的累加消融实验（在 Qwen3-1.7B 基模与 Long-CoT 数据集上进行）清晰揭示了各项设计对保留自回归能力的作用：

实验配置	Math + Reasoning (5 项均值)	Knowledge + IF (3 项均值)	Code (3 项均值)	说明
(a) AR-SFT 基线	51.93	57.29	54.81	纯自回归下词微调基线
(b) 纯块扩散微调 (Block Diffusion)	40.98	31.29	26.47	去除干净流，全域大幅崩塌 (-21.8 分)
(c) + 因果干净上下文掩码	43.09	48.42	49.19	恢复因果干净流可见性，回血 14.0 分
(d) + 干净流 NTP 自回归损失	52.70	54.53	51.41	引入干净流监督，数学推理完全恢复
(e) + 噪声流 Logit Shift (FLARE)	51.00	56.91	51.93	对齐块边界自回归语义，达到饱和最优

同时，在微调数据配比的对比实验中，涵盖长链思考、竞赛数学与指令遵循的综合数据集 Mix 4（Long-CoT+Math+IF）取得了各领域最平衡的保留效果；而仅使用短链思考的配方（Mix 2）会导致代码能力断崖式下跌。

此外，单卡并发 Serving 吞吐评测（1×A100-80GB, 并发数 \(C=8\), 生成长度 2048）显示： - GSM8K 吞吐: FLARE-2B 达到 2,087.0 tokens/s，是 LLaDA-2.1-mini（963.0 tokens/s）的 2.17 倍，是 SDAR-1.7B（438.3 tokens/s）的 4.76 倍。 - GPQA-D 吞吐: FLARE-2B 达到 1,440.8 tokens/s，是 LLaDA-2.1-mini（401.2 tokens/s）的 3.59 倍。

关键发现¶

迁移数据质量是能力保留的第一主因：在消融实验中，纯算法层面的掩码与损失形式固然重要，但在损失-掩码方案对齐后，残余性能差距完全由微调数据的覆盖面与分布决定。更重要的是，AR-SFT 的微调表现是评估迁移数据质量极其可靠且低成本的代理指标。
混合线性注意力的硬件调度存在两极性：在扩散微调必须的小块区（\(B=1\) 或 \(B=4\)），Route II 寄存器重构方案在吞吐和显存上占据压倒性优势（MFU 从 13.8% 提升至 24.8%）；而在大块区（\(B \ge 16\)），张量核心对稠密 GEMM 的高利用率使 Route I 反超。
代码生成对去噪模式存在语法脆弱性：Diffusion-Trust 在高结构化代码任务中显著弱于 AR-Trust（如 LCBv6 仅 5.71% vs 49.71%），主要因扩散并行去噪缺失严格从左到右的代码语法依赖检验，微小的括号或缩进扰动导致评测脚本执行失败。

亮点与洞察¶

算法与系统协同解决线性注意力状态回滚：投机解码中最耗时的“部分接受重算”在线性注意力中被 Cache-and-Scatter 机制优雅化解，在前向中顺带写入寄存器状态切片，一次 Scatter 搞定状态对齐，避免了额外的层级前向重演。
双解码路径兼顾两端场景：FLARE 打破了“自回归”与“扩散”模型的互斥对立，单套权重既可在对代码等极度敏感的任务中使用 AR-Trust 获得数学等价的高精度加速，又能在数学推理等长序列生成中直接利用 Diffusion-Trust 释放最大吞吐。
数据筛选的低成本代理启示：AR-to-dLLM 耗费算力，但实验证明数据配比对两者的能力影响严格单调同向。研究者可先用极小代价做 AR-SFT 验证数据组合，再投入扩散迁移。

局限与展望¶

2x 计算与显存开销：双流目标拼接了干净流与噪声流，前向输入长度翻倍（\(2L\)），在长上下文训练中显卡压力依然较大。未来需探索在单流前向中共享特征的轻量化微调技术（如参考 Orthrus 或 DiffuMamba）。
与源模型仍存分布漂移：使用外部蒸馏的 CoT 轨迹（来自 DeepSeek-R1、GPT-OSS 等）会导致微调后的模型分布偏离原 Qwen3.5 教师，未来应探索同源自生成数据以弥合指令遵循差距。
模型尺度与 RL 探索：目前验证限于 10B 以内的稠密模型，尚未覆盖 MoE 架构（如混合专家的路由掩码交互）；此外，扩散目标的不可积似然函数为强化学习（RL）对齐带来挑战，ELBO 代理梯度的引入将是后续重点。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将扩散生成与线性-Softmax混合注意力成功融合，软硬件结合深度极高。
实验充分度: ⭐⭐⭐⭐⭐ 从数据配比、算法消融、内核微基准到真实 Serving 吞吐，维度极其严密完备。
写作质量: ⭐⭐⭐⭐⭐ 论证脉络极清晰，不仅讲清了怎么做，更深入剖析了 AR-to-dLLM 失败的根本成因。
价值: ⭐⭐⭐⭐⭐ 为大模型端侧低延迟部署开辟了“算法并行度+架构高效度”融合的标准落地范式。