提出梯度引导的反向熵策略,提升掩码扩散模型生成质量与效率。
TABES: Trajectory-Aware Backward-on-Entropy Steering for Masked Diffusion Models
- 用单次反向传播近似无限时域前瞻,避免局部错误累积。
- 在多个数据集上实现更优的生成质量-速度权衡,优于现有方法。
- 适合追求高效高质非自回归生成的研究者与开发者。
掩码扩散模型(MDMs)作为生成任务中的非自回归范式,具备并行解码和双向上下文利用的优势。然而,现有采样方法依赖简单的置信度启发式规则,忽略局部决策的长期影响,导致早期幻觉引发全局不一致。虽然搜索类方法可缓解此问题,但计算成本极高(每步需O(K)次前向传播)。本文提出反向熵(BoE)引导推理框架,通过一次反向传播近似无限时域前瞻。我们从轨迹代价函数的一阶展开中推导出词元影响得分(TIS),证明未来熵关于输入嵌入的梯度是降低不确定性的最优控制信号。为保证可扩展性,引入 exttt{ActiveQueryAttention},一种基于掩码目标结构的稀疏伴随算子,显著降低反向传播复杂度。实验表明,BoE在推理时间扩展性上达到更优的帕累托前沿,验证了梯度引导推理在数学上严谨且高效的路径。代码将公开。
原文摘要 · Abstract (English)
Masked Diffusion Models (MDMs) have emerged as a promising non-autoregressive paradigm for generative tasks, offering parallel decoding and bidirectional context utilization. However, current sampling methods rely on simple confidence-based heuristics that ignore the long-term impact of local decisions, leading to trajectory lock-in where early hallucinations cascade into global incoherence. While search-based methods mitigate this, they incur prohibitive computational costs ($O(K)$ forward passes per step). In this work, we propose Backward-on-Entropy (BoE) Steering, a gradient-guided inference framework that approximates infinite-horizon lookahead via a single backward pass. We formally derive the Token Influence Score (TIS) from a first-order expansion of the trajectory cost functional, proving that the gradient of future entropy with respect to input embeddings serves as an optimal control signal for minimizing uncertainty. To ensure scalability, we introduce \texttt{ActiveQueryAttention}, a sparse adjoint primitive that exploits the structure of the masking objective to reduce backward pass complexity. BoE achieves a superior Pareto frontier for inference-time scaling compared to existing unmasking methods, demonstrating that gradient-guided steering offers a mathematically principled and efficient path to robust non-autoregressive generation. We will release the code.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。