不同任务不同时间步:面向扩散语言模型的时序任务感知量化 Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models
量化敏感度随任务与去噪步而变,T2Quant按时序分配激活精度,同预算最高提升6.7个百分点。
前置知识
扩散语言模型(DLM)
以掩码扩散方式生成文本的语言模型(如 LLaDA、Dream):从全 [MASK] 序列出发,每步基于双向上下文并行预测所有位置,按置信度解锁若干 token(低置信重掩码采样),迭代去噪直至生成完毕。生成 512 token 通常需 128 步去噪、每步提交 4 个 token,推理要反复跑前向。
本文所有分析都建立在去噪轨迹这条时间轴上:量化误差正是沿这 128 步非均匀分布的,理解逐步提交 token 的机制是读懂时序敏感性诊断的前提。
后训练量化(PTQ)与 W4A4
不重新训练、直接把训练好的模型权重和激活值映射到低位宽的技术。W4A4 表示权重 4 bit、激活 4 bit。常见组件包括 GPTQ(逐层误差补偿的权重量化)和 Hadamard 旋转(QuaRot 风格,把激活分布旋转变匀以抑制离群值),激活通常逐 token 动态量化并做范围裁剪。
T2Quant 不发明新量化器,而是叠加在现有 PTQ 方法(DLLMQuant、FAIR-Calib、STaR-Quant)之上,只改变激活位宽的时间分配,不了解 PTQ 基本流程就无法理解它的插入位置。
KL 散度
衡量两个概率分布差异的指标 $D_{\mathrm{KL}}(p\|q)=\sum_x p(x)\log\frac{p(x)}{q(x)}$,非负且不对称,散度越大说明量化让模型在该位置的预测偏移越严重。本文用它度量量化前后提交位置的预测差异 $K_w$、升精度后提交分布的变化 $J_w$、以及窗口前后 FP 模型对剩余掩码位置信念的位移 $I_w$。
论文的无标签敏感性信号($K_w$、$\kappa$、$\iota$)全部由 KL 散度构造,命题 1 的理论界也由它给出,不熟悉它就无法理解方法的核心公式。
混合精度量化
在同一模型内给不同部分分配不同位宽,在固定平均位宽预算下换取更高精度:重要的部分多给比特、不重要的部分少给。分配对象可以是层、权重(空间维度,如 TAQ、TaCQ),也可以是扩散模型去噪时间步(时间维度,如图像扩散的 TMPDiff),通过求解预算约束下的分配问题得到调度。
本文的关键动作就是把混合精度从空间维度(层)搬到时间维度(去噪窗口)并叠加任务依赖,整个方法就是一个时间维混合精度的分配算法。
背包问题与动态规划
在总预算约束下从候选集合中选取组合使总收益最大的经典优化问题,可用动态规划精确求解:定义 $F_k(z)$ 为前 $k$ 个物品/窗口用预算 $z$ 能取得的最大收益,逐阶段递推并回溯得到最优解,复杂度随物品数、预算和选项数线性增长。
T2Quant 的调度器就是一个预算约束 DP:为 16 个去噪窗口各选一个位宽、总成本等于均匀 4-bit,命题 2 证明该递推在代理目标上全局最优。
研究动机
扩散语言模型靠迭代去噪生成文本,512 个 token 要跑 128 步前向,推理成本远高于自回归模型,量化因此格外重要。但现有 DLM 量化方法(DLLMQuant、STaR-Quant、FAIR-Calib、Quant-dLLM)虽然针对扩散状态做了校准或补偿,激活位宽却在整条去噪轨迹上恒定不变。实验显示这个设计的代价很大:W4A4 下 Dream-7B-Instruct 的最强基线 STaR-Quant 六基准平均准确率只有 38.55%,距 FP16 的 44.05% 还差 5.5 个百分点,RTN 更是全线崩溃到 0%。作者的时间敏感度诊断(Figure 1)进一步表明量化误差沿轨迹的分布极不均匀:128 步切成 16 个 8 步窗口后,均匀假设是每窗承担 6.25% 的任务相关 KL 散度,但 GSM8K 的第一个窗口就占了 27%(4.4 倍于均匀值),前四分之一轨迹累计 55%;HumanEval 却对早期不敏感,敏感峰出现在第 32–71 步的五个中部窗口(合计 52%);MMLU-Pro 虽是最均匀的任务,最敏感与最不敏感窗口之间仍有近 2 倍差异。也就是说,"在哪里多花比特"是一个被现有方法完全忽略的自由度。
本文的目标是本文的目标分三层。第一,搞清楚 DLM 激活量化敏感性在去噪时间轴上的分布规律,并验证它是否随下游任务而变(RQ1、RQ2)。第二,回答能否在不使用任何标签、不跑任何任务评测的前提下,预测"给哪个窗口追加精度最划算"(RQ3),从而把精度分配变成一个只用少量无标注校准 prompt 就能完成的校准问题。第三,在此之上构建 T2Quant 框架:给定固定平均激活位宽预算(如严格 4.00 bit),为每个任务求解一条时间维混合精度调度,使其能即插即用地叠加在 DLLMQuant、FAIR-Calib、STaR-Quant 等现有量化器上,把 W4A4 量化模型的平均准确率显著拉近 FP16(目标把差距压到 2 个百分点左右),并给出调度求解的全局最优性保证。
与已有工作不同的是,本文的独特之处在于同时占据"时间维"和"任务维",且全程无标签。已有的 DLM 量化工作虽然承认时间动态重要——DLLMQuant 做扩散感知校准、FAIR-Calib 针对脆弱去噪状态、STaR-Quant 用状态引导变换做时间补偿——但激活位宽在整个轨迹上保持恒定;Quant-dLLM 的敏感度混合精度也不按时间步分配。任务感知量化(TAQ、TaCQ)确实证明了"预算应该按任务分配",但它们作用于自回归 LLM 的空间维度(层、权重电路),没有时间概念。时间维混合精度此前只在图像扩散模型中被研究过(Q-Diffusion、TMPDiff),其假设是各时间步的潜变量误差近似可加;DLM 在三点上本质不同:生成由离散且常常不可逆的 token 提交构成,潜空间重构误差不直接对应下游正确性,一个时间步的重要性取决于它提交了哪些任务相关的 token。T2Quant 由此提出"量化散度 × 下游影响"的无标签信号,并用精确动态规划在预算约束下求任务专属调度,填补了这一空白。
核心方法
直觉:量化真正伤害任务的时刻,是某个去噪窗口"提交了会被后续步骤依赖的 token"、而量化恰好让这些提交发生偏移的时刻。因此一个窗口值不值得升精度,取决于两个因素的组合——量化在该窗口造成的预测分布偏移(散度),以及这些提交对下游尚未解锁位置的撬动作用(影响)。技术路线上,T2Quant 把 128 步去噪轨迹切成 $K=16$ 个 8 步窗口,用少量无标注任务 prompt 让 FP16 模型生成一条参考轨迹并保存每步状态与提交位置;再把每个保存状态在候选位宽 $\mathcal{B}=\{3,4,6,8\}$ 下重放,得到每窗每位宽的平均提交位置散度 $\kappa_{\tau,k}(b)$;同时在参考轨迹上计算每窗的下游影响 $\bar{\iota}_{\tau,k}$(窗口前后 FP 模型对剩余掩码位置信念的 KL 位移);二者相乘得到收益分 $g_{\tau,k}(b)$。最后把收益分当价值、激活算力当成本,用精确动态规划在"平均位宽 = 4.00 bit"的硬约束下求全局最优调度。推理时每个去噪步只需按窗口序号查表选择激活位宽,没有任何额外诊断前向,对底层量化器和采样器零改动。
核心创新是无标签敏感性信号 $S_w=\mathbb{E}[K_w]\,\mathbb{E}[I_w]$ 及其背后的"扰动×敏感度"分解。命题 1 证明:把窗口 $w$ 从 $A_{\mathrm{target}}$ 升到 $A_{\mathrm{up}}$ 引起的任务分数变化满足 $|\Delta M_w|^2 \le \tfrac{1}{2}\,\mathbb{E}[J_w]\,\mathbb{E}[R_w^2]$,其中 $J_w$ 是升精度导致的提交结果(token 值与位置)分布的 KL 散度(扰动项),$R_w$ 是期望最终分数对各种可能提交的取值范围(敏感度项)。由于从单一无标签轨迹观测不到这两个理论量,作者用两个可直接计算的代理替代:$K_w$(式 1,量化前后该窗口各提交位置预测分布 KL 之和)对应扰动,$I_w$(式 2,窗口前后 FP 模型对剩余掩码位置信念的 KL 位移)对应下游杠杆。关键区别在于:只看散度会把 HumanEval 的最后一个窗口误判为高价值——该窗口的提交根本影响不到任何后续 token;只有乘上影响项才能区分"扰动大"与"扰动且致命"。与已有方法的本质区别:DLLMQuant 等用时间信息改校准与补偿、位宽不变;TAQ/TaCQ 用任务信号分配空间精度;T2Quant 是第一个用任务信号分配时间精度的框架。
方法步骤详情
完整流程五步。(1) 分窗与参考轨迹:用 FP16 模型对任务 $\tau$ 的 $N$ 条无标注 prompt 各生成一条轨迹(512 token、128 步、每步 4 token),保存每步状态 $x_{i,t}$、提交集合 $C_{i,t}$ 及提交位置的 FP 预测分布,并把轨迹切成 $K$ 个不重叠窗口(实验中 $K=16$,每窗 8 步)。(2) 局部散度:把每个保存状态在候选位宽 $b\in\mathcal{B}$ 下重放(输入固定,隔离量化对轨迹的扰动),按式 (3) 对窗口内所有提交位置求平均 KL,得 $\kappa_{\tau,k}(b)$。(3) 下游影响:取每窗前后状态 $x^-_{i,k}, x^+_{i,k}$ 及其后仍掩码的位置 $M_{i,k}$,按式 (2) 算 FP 信念位移 $\iota_{i,k}$,对 prompt 取平均得 $\bar{\iota}_{\tau,k}$;无剩余掩码则影响为 0。(4) 收益打分:$g_{\tau,k}(b)=\bar{\iota}_{\tau,k}[\kappa_{\tau,k}(b_0)-\kappa_{\tau,k}(b)]$(式 4),正值支持升位宽、负值衡量降位宽代价。(5) 预算求解:以均匀 4-bit 成本 $C_0$ 为约束,DP 递推 $F_k(z)=\max_{b:\,c_k(b)\le z}F_{k-1}(z-c_k(b))+g_{\tau,k}(b)$(式 6),回溯得最优调度,若 $F_K(C_0)\le 0$ 返回均匀调度。复杂度 $O(KC_0|\mathcal{B}|)$ 时间、$O(KC_0)$ 内存;推理时按窗口序号查表选位宽,零额外开销。
技术新颖性
新颖性体现在四点。其一,问题表述新:首次把"激活位宽"当作沿去噪轨迹逐窗分配的资源,并实证其最优分配是任务相关的——GSM8K 前置、HumanEval 中段、MMLU-Pro 宽平台,且该模式跨模型稳定(LLaDA-1.5 与 LLaDA-8B-Instruct 的逐窗画像相关系数达 0.99/0.98/0.92,Dream 上仍任务相关但峰值位置不同)。其二,信号设计有理论支撑:不是启发式选指标,而是从命题 1 的"扰动×敏感度"界推导出乘积形式代理 $S_w$;Table 3 显示它优于仅散度(HumanEval 上 $r=-0.08$)、仅影响($r=0.66$)和几何均值形式($r=0.61$)。其三,求解是精确的:预算约束 DP 给出全局最优(命题 2),而非贪心或随机搜索,且复杂度只随窗口数线性增长;实验还发现调度可压缩为"前 $m$ 窗 8-bit、后 $4m$ 窗 3-bit"的简单先验族。其四,工程上是量化器无关的加载层:不改 DLLMQuant/FAIR-Calib/STaR-Quant 的量化器与采样器,只叠加时间调度,九种"基线×模型"组合全部获益(54 个基准项中 42 项提升)。与图像扩散的 TMPDiff 相比,它不依赖潜误差可加假设,直接面向离散 token 提交的语义。
实验结果
主实验(Table 1)覆盖三个 DLM × 六基准 × 三个量化基线,W4A4、平均激活严格 4.00 bit、逐题配对比较。结论一:九个"基线×模型"组合平均准确率全部提升,54 个基准项中 42 项改善,最大增益 Dream +6.7(DLLMQuant 33.99→40.69)、LLaDA-1.5 +3.4、LLaDA-8B +3.5 分。结论二:每模型最强 T2Quant 配置超最强基线 3.55/1.13/1.40 分,与 FP16 差距缩至 2.0/2.5/1.6 分(LLaDA-8B:FP16 46.32 vs FAIR-Calib+T2Quant 44.76)。结论三:增益集中在语言任务——Dream 的 GSM8K 从 35.55(DLLMQuant)升至 53.52、从 46.09(STaR-Quant)升至 58.79;LLaDA-1.5 与 LLaDA-8B 的 GSM8K 分别 +9.6 和 +10.4 分;部分项超 FP16(LLaDA-1.5 DLLMQuant+T2Quant 的 ARC-C 83.79 > 82.62)。结论四:同预算随机调度平均掉 14–19 分(Dream 上仅 19.61%),证明收益来自"分配对了"。消融(Figure 4,LLaDA-8B/GSM8K:均匀 44.1→T2Quant 49.4,FP16 52.9):16 条校准 prompt 恢复约 75% 增益;调度从 4 迁移到 2 token/步仍保持约 5 分、8 token/步失效;3.75 bit 的 T2Quant(48.4%)胜过 4 bit 均匀(44.1%)。代理信号与实测窗口增益相关 $r=0.98/0.96/0.63$(GSM8K/MMLU-Pro/HumanEval),前两者分别找回 90% 和 100% 的 FP 差距。极限测试(Table 4):W4A3 下均匀量化崩溃(Dream GSM8K 5.3%),T2Quant 在 Dream ARC-C 恢复 32.0 分、LLaDA-1.5 GSM8K 恢复 8.6 分;W8A8 下调度基本无害(仅一项 +3.5,$p=0.04$)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| W4A4 六基准平均准确率(Dream-7B-Instruct) | 准确率 (%) | 42.10(STaR-Quant + T2Quant) | 38.55(STaR-Quant,最强基线;FP16 为 44.05) | +3.55 个百分点,距 FP16 仅 1.95 分 |
| W4A4 六基准平均准确率(LLaDA-1.5-8B) | 准确率 (%) | 45.22(STaR-Quant + T2Quant) | 44.09(STaR-Quant) | +1.13 个百分点,距 FP16 47.73 仅 2.51 分 |
| W4A4 六基准平均准确率(LLaDA-8B-Instruct) | 准确率 (%) | 44.76(FAIR-Calib + T2Quant) | 43.36(STaR-Quant) | +1.40 个百分点,距 FP16 46.32 仅 1.56 分 |
| GSM8K(Dream + DLLMQuant) | 准确率 (%) | 53.52 | 35.55 | +17.97 个百分点 |
| GSM8K(Dream + STaR-Quant) | 准确率 (%) | 58.79 | 46.09 | +12.70 个百分点 |
| GSM8K(LLaDA-8B-Instruct + DLLMQuant) | 准确率 (%) | 42.19 | 31.84 | +10.35 个百分点 |
| W4A3 3.00 bit 极限预算(Dream ARC-C) | 准确率 (%) | 77.5–78.9(T2Quant,两种位宽集合) | 46.9(均匀 W4A3) | +30.7 ~ +32.0 个百分点 |
| 无标签代理 vs 实测逐窗增益(GSM8K / MMLU-Pro / HumanEval) | Pearson 相关系数 r | 0.98 / 0.96 / 0.63($S_w=\mathbb{E}[K_w]\mathbb{E}[I_w]$) | 仅散度 $K_w$:0.97 / 0.68 / −0.08 | 乘上影响项后在 HumanEval 上从负相关扭转为有效排序 |
局限与改进
作者明确承认两点局限:一是未测部署相关指标——推理速度、能耗等硬件效率没有验证,所有量化实验都是 bf16 算术上的模拟;二是任务/时间感知的精度分配对模型安全性、鲁棒性的影响完全留待未来研究。我的补充观察:第一,调度是任务专属的,线上服务若同时处理多任务需要任务路由或为每任务缓存一套调度,文中未讨论未知任务或新任务场景;第二,代理信号在小基准上噪声明显——HumanEval 仅 100 题、$r=0.63$,窗口上采样实验需要在 GSM8K/HumanEval 上平均 5/10 个量化种子才能稳定,单种子结论不够可靠;第三,命题 2 的最优性只对加性代理目标成立,不保证真实任务指标最优,且并非处处正收益——42/54 意味着仍有约 22% 的基准项受损(如 LLaDA-8B DLLMQuant+T2Quant 的 ARC-C 从 34.96 降到 29.88,LLaDA-1.5 STaR+T2Quant 的 BBH 从 42.19 降到 39.45);第四,只在掩码扩散架构(LLaDA 系列、Dream)上验证,未覆盖半自回归块扩散或其他 DLM 变体;第五,校准 prompt 与评测 prompt 从同一评测集抽样划分,真实分布偏移下调度的稳健性未知。
独立分析的弱点
弱点一:无真实加速证据。混合精度意味着不同窗口要用不同位宽的反量化/算子路径,kernel 切换与显存管理可能吃掉理论收益;改进方向是与推理引擎(如 vLLM 类系统)集成,报告端到端吞吐和时延。弱点二:任务依赖性是双刃剑——不知道任务就无法选调度;改进方向是训练轻量任务路由器,或利用论文发现的先验结构(前 $m$ 窗 8-bit、后 $4m$ 窗 3-bit,Figure 8b)构造任务无关的默认调度,再少量数据在线微调 $m$。弱点三:代理信号统计脆弱。HumanEval 上 $r=0.63$、仅散度信号甚至为 $-0.08$,说明百题量级下方差过大;改进方向是 Bootstrap 置信区间、跨种子聚合或贝叶斯收缩,并探索更便宜多样的辅助信号(文中已试过熵、top-2 margin 等六种变体)。弱点四:调度粒度与采样器耦合,8 token/步的高并行场景调度失效(Figure 4b),说明每窗影响信号在高并行下饱和;改进方向是把精度调度与重掩码策略、并行度联合设计,而非独立优化。弱点五:与基线叠加仍有约 22% 基准项下降,缺少"何时该信任调度"的事前判据;可把现有的 $F_K(C_0)\le 0$ 退回机制升级为逐窗口的统计显著性检验,不显著时自动回退均匀位宽,避免个别基准受损。
未来方向
作者提出的方向:考察任务/时间感知精度分配对安全性与鲁棒性的影响(量化已被证明会改变拒答等安全相关行为,Hong et al. 2024、Egashira et al. 2024);补齐硬件效率(推理速度与能耗)评估。基于本文成果可延伸的方向:其一,时间×空间联合分配——把 $g_{\tau,k}(b)$ 推广到(窗口 × 层/权重电路)的二维资源网格,DP 升级为多维背包或近似求解器;其二,把同一套"散度×影响"信号用于其他逐窗决策,如每窗重掩码比例、每步并行 token 数、投机解码的触发时机,做成统一的轨迹级资源调度器;其三,在线自适应调度:生成过程中用已观测的熵/置信度实时更新窗口位宽,摆脱对任务先验和离线校准的依赖;其四,扩展到更多 DLM 形态(半自回归块扩散、AR-扩散混合)与更大规模模型,检验画像的跨规模稳定性——本文已证跨对齐配方稳定(LLaDA-1.5 与 LLaDA-8B 画像相关 0.92–0.99),但 Dream 的峰值位置漂移说明架构因素还需系统研究;其五,理论上为确定性解码情形($J_w$ 在点质量分布下发散、命题 1 的界失效)补出量化保证,或将代理界转化为对真实任务指标的真正上界。
复现评估
复现条件总体友好。模型全部公开:Dream-v0-Instruct-7B、LLaDA-1.5、LLaDA-8B-Instruct 均有官方权重与聊天模板;基准公开且论文给出明确划分(Table 5:GSM8K/ARC-C/MMLU-Pro/DROP/BBH 各 512 题评测 + 256 条校准,HumanEval 100 + 50),校准池与评测池不相交、所有行逐题配对。方法计算量不大:校准只需对数百条 prompt 的 FP 参考轨迹做 4 个位宽的状态重放,DP 复杂度 $O(KC_0|\mathcal{B}|)$ 可忽略;7–8B 模型单卡 A100 级 GPU 即可承载,量化采用 bf16 模拟协议,无需低比特硬件。主要障碍:论文未提到发布官方代码(正文只有占位的 "Project Page"),三个基线 DLLMQuant/FAIR-Calib/STaR-Quant 均为作者自行重实现并统一到底座,是最大的复现工作量与不确定性来源;量化种子也影响结果(窗口上采样实验需 5–10 个种子平均)。综合判断:只复现 T2Quant 主张(Hadamard 旋转 + GPTQ W4 底座 + 16 窗时间调度)难度中等;从零复现全部对照实验难度中等偏上。
论文图表
三个任务(GSM8K、HumanEval、MMLU-Pro)在 128 步去噪轨迹上按 8 步窗口汇总的任务相关 KL 散度曲线:GSM8K 呈前置陡峰(首个窗口占 27%,是均匀值 6.25% 的 4.4 倍,前四分之一轨迹累计 55%);HumanEval 早期不敏感、中部五窗(第 32–71 步)合计占 52%;MMLU-Pro 为宽平台、窗口间仍有近 2 倍差异。
这是全文问题的直接可视化证据,确立了"量化敏感性既随去噪时间也随下游任务变化"这一核心发现,是后续一切方法设计的出发点。
LLaDA-1.5(经偏好优化对齐的 LLaDA-8B)的逐窗敏感度画像,与 LLaDA-8B-Instruct 的画像高度一致:逐任务 Pearson 相关为 0.99(GSM8K)、0.98(HumanEval)、0.92(MMLU-Pro),各任务的峰值位置与形状几乎复现。
证明"任务-时间敏感模式"是对齐训练之前的稳定性质而非微调产物,增强了诊断信号的可迁移性结论。
Dream-v0-Instruct-7B(不同架构与初始化、自研采样器)的画像仍呈任务依赖,但峰值位置与 LLaDA 不同:GSM8K 集中在第一个窗口(答案约 6 token 直接输出)、HumanEval 迟峰(第 96–104 步)、MMLU-Pro 首窗尖峰后稳步上升;EOS 处理规则显著影响画像(Dream 默认先提交 padding 会把答案挤进最后几步,需降低 EOS 预测位置的置信度)。
跨架构验证了任务相关性这一核心命题,同时揭示采样器与 EOS 规则对画像的影响,是把方法用于新模型时必须注意的实践细节。