Polestar: Drift-Aware Cache Calibration and Token Commitment for Efficient Inference of Diffusion LLMs¶

会议: NeurIPS 2026
arXiv: 2607.14107
OpenReview: https://openreview.net/forum?id=AsdJ0sPsgT
领域: 模型压缩
关键词: 扩散语言模型, KV Cache复用, 表示漂移, 并行解码, 稀疏校准

一句话总结¶

针对扩散大语言模型（dLLM）因双向注意力机制导致 KV 缓存难以复用且并行解码置信度不稳定的双重瓶颈，Polestar 首次提出以注意力分布的 KL 散度表征「表示漂移（Representation Drift）」作为统一物理信号，联合驱动局部稀疏 KV 缓存动态校准与漂移突跃前瞻 Token 提交，在无需重训练的前提下实现最高 3.7 倍吞吐加速与 3.67 步均解码并行度。

研究背景与动机¶

自回归大语言模型（AR-LLM）由于严格依赖因果掩码进行单向自回归预测，在长文本生成与复杂推理场景中面临严苛的串行延迟瓶颈。近年来，掩码扩散大语言模型（dLLM，如 LLaDA、Dream）将文本生成重塑为离散状态空间上的多步去噪扩散过程，理论上允许在单个前向传播中并行解码多个 Token，因而被视作打破自回归生成瓶颈的关键突破口。然而，现有 dLLM 在实际部署中面临严重的 Pareto 权衡：一方面，最高精度的生成往往依赖每步仅解除一个掩码（TPF=1）的极慢采样；另一方面，全序列的双向注意力机制使得历史 Context 不断受到新生成 Token 的逆向影响，破坏了传统 AR 模型中 KV 缓存的只读不可变假设，导致每一步推理都必须全序列重算 KV，单步计算复杂度随序列长度二次方激增。

为兼顾吞吐与精度，学术界涌现了两条正交的优化路径：降低单步延迟的 KV 缓存复用方案（如 Fast-dLLM、Elastic-Cache、d2Cache）与提升并行度的启发式采样策略（如固定置信阈值解码、DAWN 依赖图、KLASS）。然而，现有缓存方法要么死板地假定块内表征静态不变，要么仅凭极少数被高频注意 Token 的余弦距离作为粗糙代用指标，忽视了全序列在双向注意力下的协同演化，导致过期缓存引发置信度暴跌；而并行解码方法采用静态阈值，极易在文本演化不确定阶段误提交幻觉 Token，或在候选词收敛阶段过于保守。这两大痛点看似割裂，其深层物理根源其实完全同源：随着未掩码 Token 的逐步确定，双向注意力的全局上下文重加权会引起未提交及已提交 Token 隐藏状态的持续演进。

本文的研究切入点正是将这一被以往工作视为「缓存误差副产物」的现象，重新定性为 dLLM 推理过程的内生核心机制——「Token 表示漂移（Token Representation Drift）」。通过建立一阶摄动理论与 Pinsker 不等式界，作者证明了注意力权重的 KL 散度漂移不仅能精确圈定需要刷新 KV 状态的过期时空坐标，还能直接表征 Token 表征走向最终收敛的突跃临界点。核心 idea：以 Token 注意力分布的 KL 散度漂移为统一物理指引，构建 Polestar-Cache 局部稀疏缓存刷新与 Polestar-Commit 漂移敏感前瞻提交的双轮驱动架构，辅以 NVFP4 残差 CPU 异步卸载与 CUDA 流重叠，全方位释放扩散语言模型的并行推理潜力。

方法详解¶

整体框架¶

Polestar 建立在半自回归块状解码（Semi-autoregressive Block-wise Decoding）范式之上。对于总生成长度 $L_G$，序列被划分为若干长度为 $B=32$ 的连续块。在当前活动块 $\mathcal{J}_b$ 的去噪迭代循环中，Polestar 的推理管线主要由三个紧密协同的模块构成： 1. Polestar-Cache：维护活动块周边的局部窗口 $\Omega_b$（包含前缀与后缀），利用球形 K-means 对输入隐藏状态聚类后投影代理注意力，通过 KL 散度漂移筛选出 top-$k$ 最过期的聚类，在对应层发起稀疏前向更新包（Update Packet），实现低开销的逐层 KV 缓存定向校准； 2. Polestar-Commit：监测最终输出层的漂移增量 $\Delta_j^{(t)}$，对当前块结合动态置信阈值提早提交处于收敛突跃期的 Token；同时复用缓存校准的最终层更新包，对后缀块高漂移且高置信的 Token 执行跨块安全预先提交； 3. 软硬件协同优化（System Optimization）：GPU 仅保留紧凑的聚类质心，将隐藏状态残差经 NVFP4 量化后异步卸载至 CPU 固定内存，并在多 CUDA 流上将代理计算、漂移判定与主前向传播完全流水线化掩盖。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Prompt 与掩码序列<br/>初始化半自回归块划分"] --> B["Polestar-Cache：局部稀疏 KV 缓存校准<br/>质心投影代理注意力 + KL 漂移筛选 top-k 聚类更新"]
    B --> C["Polestar-Commit：漂移突跃与动态阈值置信提交<br/>历史漂移增量动态判决 + 当前块与后缀块协同早解掩码"]
    C --> D["软硬件协同优化：CPU 异步卸载与多流延迟隐匿<br/>NVFP4 残差异步传输 + 主副 CUDA 流重叠计算"]
    D -->|仍有未解掩码 Token| B
    D -->|全部块解掩码完成| E["输出完整去噪生成文本"]

关键设计¶

1. Polestar-Cache：基于表示漂移的局部稀疏 KV 缓存校准
针对全序列重算计算代价过大、而粗暴复用历史 KV 会导致严重语义陈旧的矛盾，作者从理论与实证分析发现：KV 缓存累积误差的绝大部分高度聚集在当前解码块紧邻的局部上下文窗口内（包括紧邻的 $2B$ 长度前缀与 $B$ 长度后缀），且误差大小与注意力分布的 KL 漂移呈强正相关。为此，Polestar-Cache 在块入口处缓存局部窗口 $\Omega_b$ 内所有层的输入隐藏状态 $\mathbf{H}_{\Omega_b}$。为避免逐 Token 跟踪的计算膨胀，方法在特征空间使用球形 K-means 将隐藏状态凝聚为 $K=8$ 个质心。在每一步去噪解码中，质心被投影为代理 Query，与当前活跃块 KV 联合计算代理注意力分布： $$D_{i,\ell}^{(t)} = \mathrm{KL}\left(\mathbf{A}_{i,\ell}^{(t)} \,\|\, \mathbf{A}_{i,\ell}^{(t-1)}\right)$$ 系统挑出漂移量最大的 top-$k$（默认 $k=4$）个聚类，仅针对这些聚类对应的隐藏状态在当前层执行稀疏前向传播，生成更新包 $U^{(t,\ell+1)}$。下一层 $\ell+1$ 依据该更新包局部刷新陈旧的 KV 缓存与对应质心。此外，更新调度采用动态事件驱动机制：仅当单步解掩码数超过 $\tau_{\mathrm{upd}}=3$ 或累积解掩码数达到该阈值时才触发刷新，并在每隔一个块入口执行一次全序列刷新以校准全局低漂移区，大幅消除冗余计算。

2. Polestar-Commit：基于漂移突跃与动态阈值的置信提交
现存方法使用固定的静态置信度阈值（如 $\tau_s = 0.9$），在长尾或多模态语义下往往导致两种极端：要么门槛过高导致 TPF 停滞在 1 左右，要么强行调低阈值引发不可逆的级联错词。Polestar-Commit 揭示出 Token 表征朝向最终真实预测的演进并非线性渐变，而是伴随着强烈的「漂移突跃（Drift Spike）」，即表征在某一瞬大幅偏离历史漂移均值并迅速锁定在 PCA 吸引子轨迹上。针对当前块，算法维护过去 $h=5$ 步的滑动漂移均值，计算突跃增量 $\Delta_j^{(t)} = D_j^{(t)} - \bar{D}_j^{(t)}$，并定义置信度调制的非线性动态门限： $$\tau_{d,j}^{(t)} = \alpha \left(\tau_s - c_j^{(t)}\right)^2$$ 其中 $\alpha=10$。当模型置信度 $c_j^{(t)}$ 接近 $\tau_s$ 时，门限趋近于 0，退化为标准置信判定；而对于中等置信度（如 0.7~0.8）但出现极端漂移突跃的 Token，该动态阈值允许其安全提前提交，有效识别出置信度滞后但表征已然收敛的关键 Token。更进一步，Polestar-Commit 将最终层生成的更新包顺延用于后缀块（Suffix Block）：当未进入当前块的后缀 Token 被聚类判定为高漂移且其在更新包计算下的置信度满足 $c_j^{(t)} \ge \tau_s$ 时，允许跨块提前解掩码，突破了传统块状解码必须逐块推进的硬性屏障。

3. 软硬件协同优化：CPU 异步卸载与多流延迟隐匿
引入中间层隐藏状态缓存与聚类虽然减少了计算量，但若全量驻留显存会造成显存爆炸并与 KV Cache 争抢带宽。Polestar 采用算法与系统协同设计：显存仅存储极精简的 $K=8$ 个 FP16 质心，局部窗口内各 Token 相对所属质心的残差特征则被压缩至 NVFP4（1 位符号、2 位指数、1 位尾数，配合块大小为 16 的 FP8 E4M3 缩放因子），内存开销骤降 75% 以上，并异步卸载至 CPU Pinned Memory。在每一步解码执行时，利用专用 CUDA 辅助流根据当前选出的 top-$k$ 聚类，通过 PCIe 异步预取并使用融合 Triton 算子实现「反量化-更新包计算」一体化。同时，主 CUDA 流全速执行当前解码块的自注意力和 FFN 计算，辅助流并发执行代理注意力投影、KL 漂移计算与下一层更新，将 Polestar 带来的算法额外延迟近乎 100% 隐匿于主前向计算流水线之中。

损失函数 / 训练策略¶

Polestar 属于严格的测试时无训练（Training-Free）推理加速框架，不修改原始扩散模型的预训练权重，无需任何 LoRA 或额外微调。核心调控参数完全在推理时配置：活动块长 $B=32$；局部窗口规模 $|\Omega_b|=3B$（前缀 $2B$，后缀 $B$）；球形 K-means 聚类数 $K=8$，刷新选定聚类数 $k=4$；事件刷新阈值 $\tau_{\mathrm{upd}}=3$；基准置信度 $\tau_s=0.9$，漂移滑动窗口 $h=5$，动态缩放系数 $\alpha=10$。

实验关键数据¶

主实验¶

在 NVIDIA A100 80GB GPU 上对 LLaDA-8B-Instruct 与 Dream-7B-Instruct 进行全面评估，覆盖 GSM8K（5-shot）、MATH（4-shot）、HumanEval（0-shot）等权威基准。对比基线涵盖原始单步解码 Baseline、固定置信度并行解码 Baseline+Parallel、Fast-dLLM、Dynamic-dLLM 以及 Elastic-Cache。

数据集	生成长度	模型 / 方法	准确率 (%)	步均 Token (TPF)	吞吐量 (TPS)	相对基线加速比
GSM8K	256	LLaDA-8B (Baseline)	79.30	1.00	6.58	1.0×
GSM8K	256	Fast-dLLM	77.88	2.72	47.88	7.3×
GSM8K	256	Elastic-Cache	77.58	2.93	55.01	8.4×
GSM8K	256	Polestar (Ours)	78.33	3.67	87.57	13.3×
GSM8K	512	LLaDA-8B (Baseline)	77.50	1.00	3.93	1.0×
GSM8K	512	Fast-dLLM	76.52	2.58	38.07	9.7×
GSM8K	512	Elastic-Cache	76.97	2.61	38.44	9.8×
GSM8K	512	Polestar (Ours)	78.18	3.59	80.60	20.5×
MATH	256	LLaDA-8B (Baseline)	33.50	1.00	9.10	1.0×
MATH	256	Fast-dLLM	30.18	2.30	38.13	4.2×
MATH	256	Polestar (Ours)	32.77	2.91	70.05	7.7×
HumanEval	256	LLaDA-8B (Baseline)	43.35	1.00	17.60	1.0×
HumanEval	256	Fast-dLLM	38.63	2.64	54.19	3.1×
HumanEval	256	Elastic-Cache	41.67	3.08	72.75	4.1×
HumanEval	256	Polestar (Ours)	42.42	3.38	87.45	5.0×
HumanEval	512	LLaDA-8B (Baseline)	44.10	1.00	8.57	1.0×
HumanEval	512	Polestar (Ours)	47.73	3.23	77.67	9.1×
GSM8K	256	Dream-7B (Baseline)	75.85	1.00	9.85	1.0×
GSM8K	256	Fast-dLLM	66.67	1.71	39.86	4.0×
GSM8K	256	Polestar (Ours)	72.40	2.39	52.80	5.4×

消融实验¶

在 LLaDA-8B-Instruct、GSM8K（生成长度 256）上拆解各模块对推理性能的逐步增益（以扩展自 Fast-dLLM 的基线为起点）：

配置	准确率 (%)	步均 Token (TPF)	吞吐量 (TPS)	相对基线加速比	说明
Baseline (LLaDA)	79.30	1.00	6.58	1.0×	基础逐步无缓存去噪
Baseline + Parallel	78.57	2.95	17.65	2.7×	基础静态置信度并行
Fast-dLLM 基准	77.88	2.72	47.88	7.3×	块静态 KV 缓存
+ Polestar-Cache	78.89	2.87	60.54	9.2×	引入漂移引导局部稀疏刷新，精度回升 +1.01%
+ Current-Block Commit	78.57	3.48	82.14	12.5×	引入漂移突跃动态阈值，TPF 提升至 3.48
+ Suffix-Block Commit (完整)	78.33	3.67	87.57	13.3×	开启后缀块跨块前瞻提交，TPS 达到 87.57

关键发现¶

双向驱动的协同增益：单纯引入 Polestar-Cache 不仅将生成精度提升 1.01%（修复陈旧缓存导致的模型困惑），还使 TPF 提高 0.15，说明校准缓存直接增强了解码置信度；而叠加上 Polestar-Commit 后，步均解码 Token 跃升至 3.67，吞吐量相较于 Fast-dLLM 提升 1.83 倍。
反超基线的异常现象：在长文本场景（如 GSM8K-512 与 HumanEval-512）中，Polestar 的精度甚至超越了全量重算的 Baseline（GSM8K 上 78.18% vs 77.50%，HumanEval 上 47.73% vs 44.10%）。其本质在于传统强制每步解码固定数量 Token 容易把错误猜测固化，而 Polestar-Commit 依靠表征收敛轨迹动态截断，过滤了大量低置信度噪点。
难并行测试的鲁棒性：在专为检测并行退化设计的 ParallelBench-easy 上，Polestar 在维持 74.18% 高精度的同时实现 2.43 TPF 与 82.41 TPS，显著优于 Fast-dLLM（68.09% / 1.88 TPF）和 Elastic-Cache（69.84% / 1.96 TPF）。

亮点与洞察¶

物理视角的范式转换：以往加速工作将双向注意力引起的 KV 变化视为单纯的「误差扰动」，而 Polestar 敏锐地将其重塑为反映上下文演化与表征收敛的统一「物理信号」，一举打通了缓存失效判定与提前解码两个看似分离的子问题。
优雅的轻量代理设计：利用球形 K-means 质心作为代理 Query 探测注意力 KL 漂移，仅需极小代价值便能准确定位局部窗口中最需要刷新的前 50% 关键特征，既规避了逐 Token 探测的高昂算力，又打破了单 Token 启发式统计的盲目性。
软硬件深度融合的系统落地：将隐藏状态残差量化至极窄精度的 NVFP4 并卸载至主机固定内存，结合双 CUDA 流全流水线掩盖，成功将算法创新转化为货真价实的硬件吞吐提升，未在显存和通信上交税。

局限与展望¶

长序列全量刷新开销：Polestar-Cache 目前在每隔一个块入口仍需触发一次全序列前向刷新，以弥补非局部窗口的低频累积误差。当序列扩展至数万 Token 级别时，这一全量刷新步骤可能逐渐成为吞吐放大的主要瓶颈。
自适应超参的跨架构迁移：当前质心数 $K=8$、动态缩放因子 $\alpha=10$ 以及局部窗口大小 $3B$ 是在密集 Transformer 结构上经验调优设定的。面对 MoE 稀疏激活扩散模型或超深架构时，各层漂移规律差异巨大，可能需要更自适应的层级超参策略。
未来改进方向：可进一步探索将后缀块的漂移突跃信号与投机采样（Speculative Decoding）相结合，甚至在预训练阶段引入轻量级的漂移感知蒸馏头。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 颠覆性地将双向注意力的表征漂移重塑为统一驱动信号，理论推导严密且兼顾缓存与解码。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖数学、代码及多模态基准，深入对比主流 KV 缓存与并行解码方案，消融分析细致扎实。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰紧凑，理论与工程设计环环相扣，图表详实且极具说服力。
价值: ⭐⭐⭐⭐⭐ 彻底打破了扩散大语言模型「高精度低吞吐」的核心落地枷锁，代码友好且具备极强的工程部署价值。