💬 NLP

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

Polestar is a training-free inference framework that leverages token representation drift as a unified signal to enable efficient KV-cache reuse and reliable token commitment, thereby significantly improving both the accuracy and throughput of diffusion large language models.

Original authors: Mingyu Lee, Akshat Ramachandran, Souvik Kundu, Tushar Krishna

Published 2026-07-17

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Mingyu Lee, Akshat Ramachandran, Souvik Kundu, Tushar Krishna

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive jigsaw puzzle, but instead of looking at the whole picture at once, you are forced to place one piece at a time, waiting for the previous piece to settle before you can even look at the next one. This is how most modern AI chatbots work today; they are incredibly smart but move slowly because they are so cautious. A newer, faster type of AI called a "diffusion model" tries to solve the puzzle by guessing many pieces at once, like a painter filling in a whole section of a canvas in a single brushstroke. However, this new method has a tricky problem: as the AI fills in more of the picture, the context of the whole image changes, making its earlier guesses about the other pieces suddenly outdated. It's like if you were painting a sunset, and every time you added a cloud, the color of the ocean you painted five minutes ago suddenly looked wrong. To fix this, the AI usually has to stop and repaint the entire ocean every time it adds a cloud, which is incredibly slow and wasteful.

Scientists have been trying to figure out how to let these fast AI models keep their speed without making mistakes. The big question is: How can we tell the AI when it's safe to lock in a piece of the puzzle and move on, and when it needs to go back and fix its old work? If we get this wrong, the AI either gets stuck repainting the same thing over and over (slow) or it locks in the wrong pieces and the final picture looks weird (inaccurate). This is the challenge that researchers at Georgia Tech and Intel set out to solve with a new method called "Polestar."

The Drifting Compass

The researchers discovered that the secret to fixing this problem lies in something they call "token representation drift." In plain English, imagine the AI's understanding of a word as a tiny, glowing compass needle. When the AI first guesses a word, the needle points in a certain direction. But as the AI figures out more words around it, the context shifts, and that needle starts to wobble or "drift" to a new direction. The team realized that this drifting isn't just a mistake; it's a signal.

Previous methods tried to guess when to update the AI's memory by looking at static rules, like "update every 10 steps" or "update if the confidence score is high." The Polestar team found these approaches were like trying to drive a car by only looking at the speedometer; they missed the fact that the road itself was changing. Instead, Polestar watches the compass needles directly. If a needle starts wobbling wildly, it means the AI's memory of that part of the sentence is getting stale and needs a quick refresh. If a needle suddenly stops wobbling and locks into a steady direction, it means the AI has figured it out and can safely "commit" to that word, allowing it to move on to the next one faster.

Polestar: The Smart Painter

The paper introduces Polestar, a system that acts like a super-smart art director for these AI painters. It has two main jobs, which the authors call "Polestar-Cache" and "Polestar-Commit."

First, Polestar-Cache is the memory manager. Instead of repainting the entire ocean every time a cloud is added, it watches the compass needles. It only goes back to repaint the specific spots where the needles are wobbling the most. This is a huge efficiency boost. The researchers found that by using this "drift" signal, they could update the AI's memory much more precisely than before. They don't just guess; they measure how much the AI's understanding of a word has changed using a mathematical tool called KL-divergence (which is just a fancy way of measuring how much a probability distribution has shifted). By focusing only on the "drifting" spots, they save a massive amount of computing power.

Second, Polestar-Commit is the decision-maker. Usually, the AI waits until it is 100% sure before locking in a word. But Polestar noticed that sometimes, a word's compass needle stops drifting and becomes steady before the AI's confidence score gets high enough to trigger the usual "lock-in" signal. Polestar-Commit spots these "sharp drift events"—moments where the needle suddenly stabilizes—and says, "Hey, this one is done! Let's lock it in now." This allows the AI to process multiple words at once (parallelism) without losing accuracy.

The Results: Faster and Smarter

The team tested Polestar on several difficult tasks, including math problems (GSM8K), coding challenges (HumanEval), and complex reasoning (MATH). The results were impressive. In these tests, Polestar managed to make the AI up to 3.7 times faster (measured in tokens per second) compared to the standard baseline, while achieving up to 10.73% accuracy improvement over existing baselines in specific scenarios.

For example, on the GSM8K math benchmark, the standard method could only process about 1 token per forward pass (meaning it was very slow). Polestar managed to process 3.67 tokens per forward pass while achieving a high score of 78.33% accuracy. While the baseline model in that specific run scored slightly higher at 79.30%, it was significantly slower. Polestar's true power lies in its ability to set a new state-of-the-art on the accuracy-throughput trade-off, often beating other fast methods by a wide margin. In another test on the HumanEval coding benchmark, Polestar achieved a speedup of 9.1 times over the baseline while maintaining high accuracy.

The researchers also showed that their method works without needing to retrain the AI models. It's a "training-free" upgrade, meaning it can be plugged into existing AI systems like LLaDA or Dream-7B to make them instantly faster and more reliable. They even built a system optimization that moves some of the heavy lifting to the computer's CPU to keep the graphics card (GPU) from getting bogged down, ensuring the speed gains are real and not just theoretical.

What Polestar Isn't

It's important to note what Polestar does not do. The paper explicitly argues against the idea that token drift is just a "KV-cache error" (a glitch in the memory system) that needs to be ignored or averaged out. Instead, they prove that drift is a fundamental, natural part of how these models work. They also rule out the idea that simply lowering the confidence threshold (making the AI less picky) is a good way to speed things up; their tests showed that doing so without watching for drift causes the AI to make too many mistakes. Polestar doesn't just guess; it uses the specific behavior of the model's internal "compass" to make decisions.

Why This Matters

The beauty of Polestar is that it turns a problem (the AI's memory getting outdated) into a feature (using the drift to know exactly when to update). By treating the AI's changing understanding as a useful signal rather than noise, the researchers have created a way to make these powerful, fast AI models actually live up to their potential. They aren't just making the AI faster; they are making it smarter about when to move forward and when to look back, setting a new standard for how efficient AI inference can be. As the paper suggests, this approach could be the key to unlocking the full speed of diffusion-based language models without sacrificing the quality of the answers they give.

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

1. Problem Statement

Diffusion Large Language Models (dLLMs) offer a promising alternative to autoregressive LLMs (AR-LLMs) by enabling parallel token decoding through iterative denoising. However, their inference efficiency is currently constrained by two interrelated challenges:

KV-Cache Inefficiency: Unlike AR-LLMs, dLLMs utilize bidirectional attention, causing token representations to evolve (drift) as context is updated at every decoding step. This renders standard Key-Value (KV) caching ineffective, as cached representations quickly become stale. Existing solutions either assume static representations within blocks (leading to accuracy degradation) or rely on coarse proxy signals (e.g., most-attended tokens) that fail to capture the joint evolution of token representations across the sequence.
Parallelism vs. Quality Trade-off: Increasing decoding parallelism (decoding multiple tokens per forward pass) typically requires static confidence thresholds for token commitment. These fixed thresholds often fail to adapt to varying decoding steps and inputs, leading to either premature commitment (degrading quality) or conservative commitment (limiting throughput).

The authors observe that both challenges stem from a shared phenomenon: token representation drift. As tokens are decoded, their contextual integration via bidirectional attention causes their hidden state representations to evolve. This drift is not merely an error artifact but an inherent characteristic of dLLM inference.

2. Methodology: Polestar Framework

Polestar is a training-free inference framework that leverages token representation drift as a unified signal to address both KV-cache staleness and token commitment. It consists of three core components:

A. Polestar-Cache (Drift-Aware KV Refresh)

Instead of refreshing the entire KV cache or relying on static intervals, Polestar-Cache identifies stale positions based on drift magnitude.

Drift Measurement: Drift is quantified using KL-divergence between a token's attention distribution at consecutive steps ( $D^{(t)}_i = KL(A^{(t)}_i \| A^{(t-1)}_i)$ ). This captures contextual shifts more accurately than cosine dissimilarity on subsets.
Local Window & Clustering: To minimize overhead, the system maintains a local window (prefix and suffix blocks) around the current decoding block. It clusters hidden states within this window using Spherical K-Means and quantizes residuals using NVFP4 to create compact proxy representations.
Sparse Updates: At each step, the system computes "proxy attention" using cluster centroids. It identifies the top- $k$ drifting clusters and performs a sparse forward pass only on these positions to generate an "update packet." This packet refreshes the corresponding KV-cache entries and hidden states in subsequent layers, avoiding full-sequence recomputation.
Schedule: Updates are triggered selectively when decoded context causes significant misalignment (e.g., after decoding $\tau_{upd}$ tokens) or via periodic full-sequence passes for non-local positions.

B. Polestar-Commit (Drift-Aware Token Commitment)

Polestar-Commit detects "sharp drift events" to identify when a token's representation has converged, allowing for earlier and more reliable commitment than static confidence thresholds.

Current Block Commitment: It monitors the drift delta ( $\Delta^{(t)}_j$ ), defined as the deviation of current drift from the historical mean. A sharp positive delta indicates rapid contextual integration. A token is committed if its drift delta exceeds a confidence-conditioned dynamic threshold ( $\tau^{(t)}_{d,j} = \alpha(\tau_s - c^{(t)}_j)^2$ ). This allows lower-confidence tokens to be committed if they exhibit strong convergence signals.
Suffix Block Commitment: The system also analyzes the final-layer update packet. If suffix-window centroids are selected as top-drifting clusters, it indicates strong contextual updates from the current block. Suffix tokens are committed only if they meet a high confidence threshold ( $\ge \tau_s$ ), ensuring safety while increasing parallelism.

C. System Optimization

To ensure high throughput, Polestar employs:

CPU Offloading: Quantized hidden-state clusters are offloaded to pinned CPU memory, while only centroids remain on the GPU. Top- $k$ clusters are fetched asynchronously.
Latency Hiding: CUDA streams overlap auxiliary operations (clustering, quantization, proxy attention, update packet generation) with the main dLLM forward pass, effectively hiding the overhead of drift calculation.

3. Key Contributions

Unified Signal: The identification of token representation drift as a principled signal to jointly solve KV-cache staleness and parallel commitment, moving beyond treating drift as a mere caching error.
Polestar-Cache: A novel sparse refresh scheme using layer-wise drift detection and cluster-based proxies to maintain cache validity with minimal overhead.
Polestar-Commit: A mechanism to detect sharp drift events for reliable early token commitment, enabling higher decoding parallelism without the accuracy collapse associated with lowering static confidence thresholds.
System Co-Design: An optimized execution flow utilizing CPU offloading and latency hiding to minimize the runtime cost of drift monitoring.

4. Experimental Results

The authors evaluated Polestar on multiple dLLM families (LLaDA-8B, LLaDA-1.5, Dream-7B) across mathematics (GSM8K, MATH), coding (HumanEval, MBPP), and parallelization stress-test (ParallelBench) benchmarks.

Accuracy-Throughput Pareto Frontier: Polestar sets a new state of the art, achieving up to 10.73% accuracy improvement over existing baselines while simultaneously increasing throughput.
Throughput Gains: It achieves up to 3.7× higher throughput (TPS) compared to baseline models and up to 20.5× over the baseline in specific configurations (e.g., GSM8K-512).
Decoding Parallelism: Polestar achieves a high decoding parallelism of 3.67 tokens per forward pass (TPF), significantly outperforming baselines like Fast-dLLM (2.72 TPF) and Elastic-Cache (2.93 TPF) while maintaining or improving accuracy.
Ablation Studies: Experiments confirm that KL-divergence based drift measurement outperforms cosine similarity or uncertainty-based selection. Furthermore, the combination of drift-aware commitment and suffix-block commitment is critical for achieving high TPF without accuracy degradation.

5. Significance

The paper claims that Polestar fundamentally shifts the efficiency paradigm for dLLMs. By recognizing that token representations inherently drift due to bidirectional attention, the framework transforms this "drift" from a liability into a control signal. This allows for:

Efficient Cache Reuse: Enabling the reuse of KV caches in dLLMs without the accuracy penalties seen in prior methods.
Reliable Parallelism: Decoupling decoding parallelism from rigid static confidence thresholds, allowing the model to commit tokens based on the actual convergence of their representations.

Ultimately, Polestar demonstrates that training-free inference optimizations can simultaneously maximize accuracy and throughput, establishing a new benchmark for the practical deployment of diffusion-based language models.