Archer: Adaptive Reuse of Cached Hidden States for Efficient Rollback in Diffusion Language Models

2026-08-08

AI Summary

Diffusion language models incur substantial computational overhead during rollback inference due to repeated recomputation of prompt and response states caused by global context updates. This work proposes a training-free key-value caching method that, for the first time, aligns bounded prompt state reuse with rollback mechanisms. By employing an asymmetric strategy to synchronize mutable response states, reusing fixed-prompt key-value caches within a bounded neighborhood, and integrating state-aware cache refreshing with decoder marginal decision preservation, the approach significantly reduces rollback costs. Empirical results demonstrate an average speedup of 2.57× (up to 2.95×) on the main test suite, a Pass@1 improvement of up to 3.05 percentage points, and achieves a best-in-class average performance of 33.63%.

Abstract

Diffusion language models (DLMs) iteratively refine a sequence, allowing earlier predictions to be revised as context evolves. This rollback capability distinguishes them from irreversible autoregressive generation, but makes inference costly. Every denoising update alters the global context, forcing both prompt and response states to be recomputed even though only response tokens are revisable. Key-value (KV) caching could reduce this cost, yet conventional caching assumes immutable historical states and is therefore difficult to reconcile with rollback.In this paper, we introduce Adaptive Reuse of Cached Hidden States for Efficient Rollback (Archer), a training-free KV caching method for rollback-capable DLMs. Archer asymmetrically keeps the mutable response synchronized with the current hypothesis while reusing prompt K/V within a bounded state neighborhood. Although prompt representations also change under bidirectional attention, their token identities remain fixed; bounded reuse therefore amortizes repeated prompt computation without caching mutable response states. It also delays feedback from tentative tokens, reducing premature reinforcement of transient high-confidence errors and giving rollback more opportunity to correct them. Our analysis characterizes prompt reuse as a reversibility-aligned cache boundary, bounds its state-dependent approximation error, and gives a decoder-margin condition for preserving full-refresh decisions.Existing DLM acceleration often trades quality for speed. Archer shifts this frontier, attaining the best mean performance of 33.63% together with a 2.57x mean speedup on the main suite. Across evaluated settings, it improves Pass@1 by up to 3.05 points and reaches up to 2.95x speedup. Controlled analyses connect the quality gain to delayed prompt feedback and validate state-aware refresh. Our code is available at https://github.com/Hxnng/Archer.

Problem

• diffusion language models
• rollback
• KV caching
• efficient inference
• hidden state reuse

Innovation

• Diffusion Language Models
• KV Caching
• Rollback
• Adaptive Reuse
• Training-Free Acceleration