DIP: Dynamic In-Context Planner For Diffusion Language Models
Yang Li, Han Meng, Chenan Wang, Haipeng Chen
TL;DR
This work tackles the high cost of in-context learning in diffusion language models by leveraging the diffusion generation paradigm to dynamically adjust context. The authors propose Dynamic In-Context Planner (DIP), a training-free plug-in that ranks in-context examples with Maximal Marginal Relevance and progressively inserts them during block-wise, KV-cache-accelerated decoding. DIP achieves substantial throughput improvements—up to $12.9\times$ faster than standard LLaDA decoding and $1.17\times$ faster than KV-cache methods—while maintaining competitive accuracy on GSM8k. The approach offers a practical pathway to scalable, efficient diffusion-based ICL across suitable architectures and tasks.
Abstract
Diffusion language models (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 language models (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.