Different Tasks, Different Timesteps: Temporal Task-Aware Quantization for Diffusion Language Models
Ofek Avivi, Eliron Rahimi, Avi Mendelson, Chaim Baskin
cs.AI
2026-10-22
T2Quant 用无标签的散度乘影响力信号,在固定 4 bit 预算下按任务分配去噪步精度,三个模型六个基准平均最多提高 6.7 个点。
扩散语言模型(DLM)每一步都在双向上下文里改掉一批被 mask 的位置,句子要走完去噪轨迹才出来。调用次数等于步数,激活量化直接决定每步的显存和算力。DLLMQuant、STaR-Quant、FAIR-Calib 把激活比特宽沿整条轨迹钉死。图像扩散做过按时间步分精度,自回归 LLM 做过按层和按权重电路的任务感知量化。缺的是任务这根轴:同一条轨迹上,不同任务怕掉精度的步并不一样。
诊断在 LLaDA-8B-Instruct 上做,512 token、128 步,16 个互不重叠的 8 步窗口,权重 4 bit GPTQ 加 QuaRot,激活动态 per-token 4 bit。统计的是评分函数会读到的已提交 token(这一步已经把 mask 写成具体词,往往不可逆)上,全精度分布和量化分布的 KL 散度,用来表示两个预测差多远。均匀分摊时每窗 6.25%。GSM8K 第一窗占 27%,约 4.4 倍,前四分之一合计 55%。HumanEval 早期不敏感,第 32 到 71 步的五个窗口占 52%。MMLU-Pro 最摊开,峰谷仍近 2 倍。一条不顾任务的时间表对不上这三种形状。
T2Quant 不新做量化器。它接在训练后量化(PTQ,不重训、只靠少量数据定参数)上,决定每个去噪窗口用哪一档激活精度,平均比特与均匀 Atarget 相同。校准只要该任务的无标签 prompt,不看答案,也不跑准确率。
每个窗口的分数从模型内部读,不看标签。Kw 是该窗已提交位置上全精度对量化的 KL,看加精度会不会改决定。Iw 是窗前窗后,全精度模型对剩余 mask 位置的预测变了多少,看这些决定还会不会往后传。代理是两项期望的乘积。只看散度会把 HumanEval 最后一窗排到最前,那些 token 提交后影响不了后文。
调度时每个候选比特宽都在全精度轨迹上重放保存的状态,输入固定,测到的是局部量化误差。窗口收益等于平均影响力乘以参考精度与该精度的散度差。参考档记 0 分,降档得负分。动态规划在窗口和累计预算上求精确解,各窗代价之和等于均匀参考的总预算。整数代价下,它对这个加性代理全局最优。真实准确率一般不可加,不在这个保证里。推理按窗口下标查表。
16 条无标签 prompt 的代理估计,和全量 prompt 的相关不低于 0.96。GSM8K 与 MMLU-Pro 上,代理找对收益最大的单窗,收回全精度差距的 90% 和 100%,窗口级相关 0.98 和 0.96。HumanEval 上代理和实测增益都挤在前两窗,噪声把相关压到 0.63。诊断里代码任务的散度峰在中段,单窗准确率增益却在开头,影响力项负责解释这个错位。
主表是 W4A4:权重 4 bit,激活平均 4.00 bit。三个模型,Dream-7B-Instruct、LLaDA-1.5-8B、LLaDA-8B-Instruct,六个基准,评测 prompt 与校准集不相交。三种量化器原样跑一遍,再叠 T2Quant。同预算的随机时间表用来对照混合精度本身。
| 模型 | 最强原方法 | 最好的 +T2Quant | FP16 | 最大平均增益 |
| Dream-7B | STaR 38.55 | STaR+T2 42.10 | 44.05 | DLLMQuant 33.99→40.69 |
| LLaDA-1.5 | STaR 44.09 | STaR+T2 45.22 | 47.73 | DLLMQuant 39.09→42.47 |
| LLaDA-8B | STaR 43.36 | FAIR+T2 44.76 | 46.32 | DLLMQuant 37.35→40.88 |
九组组合的平均分全部上升,Dream、LLaDA-1.5、LLaDA-8B 最多加 6.7、3.4、3.5 个点。54 个基准级对比里 42 个变好。最好的量化配置距 FP16 还差 2.0、2.5、1.6 个点。GSM8K 回血最大:Dream 上 DLLMQuant 35.55→53.52,STaR-Quant 46.09→58.79。LLaDA-1.5 最大约 9.6 个点(35.94→45.51),LLaDA-8B-Instruct 约 10.4 个点(31.84→42.19)。朴素 RTN 几乎全线归零。
Dream 上随机时间表比原基线低 14 到 19 个点。比特挪到哪一段,比有没有混合精度更决定结果。
消融换到 LLaDA-8B-Instruct 的 GSM8K、均匀 W4A4(n=512):均匀 44.1,T2Quant 49.4,FP16 52.9。16 条校准收回约 75% 的增益,时间表稳定偏向早期步。每步 4 token 上得到的表,迁到每步 2 token 仍大约多 5 个点。每步 8 token 时准确率整体垮掉,增益一起没了。平均只给 3.75 bit,T2Quant 还有 48.4,高于均匀 4 bit 的 44.1。超过 4 bit 再加精度,收益可以忽略。
量化器已经有了,缺的是同一笔比特预算沿去噪时间怎么切。T2Quant 不改量化器,也不改采样器,上线时按窗口查表。任务之间不能共用一条曲线。数学题的精度堆在开头,代码题的单窗增益也在前两窗,尽管评分 token 上的散度峰位于中段。
这是渐进改进。平均激活仍是 4 bit,换的是比特落在哪些步。消融里 3.75 bit 还能打过均匀 4 bit,预算还有往下拧的空间。
速度、能耗,以及量化对安全相关回答的影响,论文都留到以后。精度沿时间不均匀,偏移有可能集中在某几步,这一点没有测量。
动态规划只保证加性代理最优。单窗升级对得上全精度差距,多窗联立时排序会不会漂,没有单独验证。HumanEval 相关只有 0.63。表里有局部回退:LLaDA-1.5 上 STaR-Quant 的 HumanEval 从 20 降到 17,LLaDA-8B-Instruct 上从 25 降到 21。平均分仍在涨:GSM8K 和 ARC-Challenge 九组全涨,HumanEval 这类局部回退没能把平均拉下来。
校准默认有一批像目标任务的无标签 prompt。请求类型事先未知,时间表就无表可查。每步提交的 token 数一变,增益会消失。实验只覆盖 7B 到 8B 的 masked diffusion,没有把权重比特和激活放进同一次调度。