DIP: Dynamic In-Context Planner For Diffusion Language Models

Abstract: Diffusion LLMs (DLMs) have shown strong potential for general natural language tasks with in-context examples. However, due to the bidirectional attention mechanism, DLMs incur substantial computational cost as context length increases. This work addresses this issue with a key discovery: unlike the sequential generation in autoregressive LLMs (ARLMs), the diffusion generation paradigm in DLMs allows \textit{efficient dynamic adjustment of the context} during generation. Building on this insight, we propose \textbf{D}ynamic \textbf{I}n-Context \textbf{P}lanner (DIP), a context-optimization method that dynamically selects and inserts in-context examples during generation, rather than providing all examples in the prompt upfront. Results show DIP maintains generation quality while achieving up to 12.9$\times$ inference speedup over standard inference and 1.17$\times$ over KV cache-enhanced inference.

• The paper introduces DIP, a framework that dynamically inserts in-context examples to boost Diffusion LM inference throughput by up to 12.9×.
• The method employs a two-stage process using MMR-based ranking and a Bernoulli policy driven by generation confidence and process penalty.
• Experiments demonstrate that DIP maintains accuracy within 0.5 points of static baselines while optimizing context for speed and efficiency.

DIP: Dynamic In-Context Planning for Efficient Diffusion LLM Inference

Introduction

Diffusion LLMs (DLMs) represent a divergence from the traditional left-to-right decoding paradigm of Autoregressive LLMs (ARLMs), leveraging bidirectional attention and iterative refinement to enable more flexible and globally-consistent text generation. Despite recent progress in DLMs, their inference efficiency has lagged due to quadratic attention scaling with context length and the inefficiency of existing in-context learning (ICL) strategies, which naively encode all examples into the prompt. This work addresses these challenges with the Dynamic In-Context Planner (DIP), a training-free, plug-in framework that dynamically selects and inserts in-context examples during generation, capitalizing on the unique properties of DLMs.

The primary hypothesis underpinning DIP is that, unlike ARLMs—which require all context at the outset due to the causal nature of generation—DLMs support the progressive, context-aware insertion of examples as the generation process unfolds. This operational distinction enables optimization of both computational throughput and context utility, yielding significant improvements in inference speed without degradation in quality.

Problem Formulation and Background

The task focuses on inference throughput optimization in DLMs under an ICL setting: given a query and a pool of candidate in-context examples, the objective is to select, at each block of the block-wise decoding process, a subset of examples to maximize downstream performance per unit computational cost. DIP frames this as a policy learning problem, where the policy $\pi$ determines the example set for each block based on real-time confidence signals and progression through the generation process.

Masked DLMs, such as those used in [sahoo2024simple] and [ye2025dream], operate via an absorbing-state diffusion process, iteratively masking and unmasking tokens based on bidirectional attention. Block-wise decoding and recent advances such as Fast-dLLM caching [wu2025fast] have partially alleviated inefficiencies, but remain constrained by static context.

Figure 1: Throughput of DLMs using Fast-dLLM demonstrates strong negative correlation with prompt context length, motivating context optimization.

Dynamic In-Context Planner: Architecture and Design

DIP decomposes context optimization into two decoupled stages:

1. Example Ranking (Stage 1): The candidate in-context examples are ranked using Maximal Marginal Relevance (MMR), balancing relevance to the current query with diversity relative to previously selected examples. This ensures that the dynamically inserted context maintains both utility and minimal redundancy as the input evolves.
2. Example Insertion Policy (Stage 2): The core innovation leverages DLMs’ unique generation paradigm. At each decoding block, DIP applies a policy that stochastically decides whether to insert the next-highest ranked example. The policy is governed by two terms:
   • Generation Confidence: Based on the average confidence of generated tokens in the current block versus the running average, measuring generation stability and quality.
   • Generation Process Penalty: A scheduling function penalizing early insertion to encourage minimal context length in early generation stages, thus maximizing throughput gains.

The policy output is a Bernoulli variable governing insertion action, enabling efficient, block-wise prompt updates that align with the performance characteristics of DLMs with KV-cache architectures.

Figure 2: The DIP workflow, illustrating (1) MMR-based example ranking and (2) dynamic, progressive, generation-aware context insertion at each block during DLM decoding.

Experimental Evaluation

DIP is evaluated on LLaDA-Instruct using the 5-shot GSM8k benchmark under the Fast-dLLM framework. The main comparison metrics are match accuracy and throughput (tokens/sec), with baselines including static-context decoding (LLaDA standard) and KV cache-optimized decoding (Fast-dLLM).

Key findings include:

• Throughput: DIP achieves up to 12.9× speedup over LLaDA’s standard inference and 1.17× over Fast-dLLM.
• Accuracy: The best configurations of DIP (e.g., $\lambda{=}0.1$, $\epsilon{=}0.2$) maintain accuracy within 0.5 points of the best KV-optimized baselines, demonstrating little degradation as context is dynamically inserted rather than statically fixed.
• Ablation: Increasing the penalty parameter $\epsilon$ improves throughput but slightly reduces accuracy, indicating a trade-off between early context insertion (benefitting quality) and delayed insertion (benefitting speed).

Implications and Future Directions

DIP concretely demonstrates the practical implications of DLMs’ non-autoregressive, bidirectional architecture: context can be fluidly optimized online during generation, unlocking inference acceleration unattainable in ARLMs. This approach sets a precedent for adaptive, utility-driven inference policies in large language understanding systems.

On the theoretical side, DIP calls attention to the under-explored flexibility in DLM context formation and encourages subsequent work in online learning and reinforcement learning-based strategies for context selection. Practically, DIP holds promise for retrieval-augmented generation pipelines and memory-limited deployment scenarios, where inference cost is tightly coupled to context size.

Further research opportunities include:

• Expanding DIP evaluation across widely varying tasks and larger DLM architectures.
• Incorporating richer confidence estimation and utility prediction schemes for even more fine-grained control.
• Extending DIP to other forms of context, such as knowledge snippets retrieved in real-time, for open-domain QA and reasoning.

Conclusion

DIP (Dynamic In-Context Planner) introduces a context-optimized inference protocol for DLMs, leveraging their generative flexibility to achieve substantial acceleration in inference throughput without quality loss. The approach exemplifies how dynamic, feedback-driven context management can overcome the fundamental limitations of static ICL, with strong implications for both the deployment efficiency and theoretical modeling of next-generation LLMs.