DIP: Dynamic In-Context Planner For Diffusion Language Models

The repo contains the official implementation for the paper: DIP: Dynamic In-Context Planner For Diffusion Language Models published in Findings of the 2026 Conference on Empirical Methods in Natural Language Processing.

Yang Li, Han Meng, Chenan Wang, Zhenyu Bi, Xuan Wang, Haipeng Chen

DIP is a context-optimization algorithm based on average verified confidence that dynamically ranks and inserts in-context examples during DLMs' generation.

Overview

Inference process of fixed context generation baseline and dynamic in-context planner inference using Fast-dLLM inference framework. (a) Illustration of the Fast-dLLM generation process under a fixed context setting, which is commonly used in 3-shot in-context learning scenarios. (b) Illustration of a dynamic context policy π based generation process. During decoding, each block may use a different number of in-context examples, depending on the policy.

Repository Layout

DIP/
├── eval_llada_DIP.py           # Entry point: registers the `llada_dist` lm-eval model and runs DIP evaluation
│
├── eval_DIP.sh                 # Example launch script
│
├── llada/                      # Diffusion LM backbone and DIP-aware decoding
│   ├── generate_DIP.py         # DIP block-wise generation loop (`generate_DIP`)
│   └── model/                  # LLaDA model definition
│       ├── configuration_llada.py
│       └── modeling_llada.py
│
├── dynamic_context/            # The DIP algorithm itself
│   ├── policy.py               # `AddPolicy`: average verified confidence + time penalty decides when to add an example
│   ├── MMR.py                  # Maximal Marginal Relevance re-ranking of candidate in-context examples
│   ├── gsm8k_DC.py             # Dynamic prompt builder for GSM8K
│   ├── math_DC.py              # Dynamic prompt builder for MATH-500
│   └── MBPP_DC.py              # Dynamic prompt builder for MBPP
│
├── data_utils/                 # Dataset loading and few-shot pool construction
│   ├── gsm8k_data.py
│   ├── math_data.py
│   └── mbpp_data.py
│
├── task/                       # Custom lm-evaluation-harness task configs
│   ├── math500/                # minerva_math500.yaml, minerva_math_algebra.yaml, utils.py
│   └── mbpp/                   # mbpp.yaml, utils.py
│
└── figs/                       # Figures used in the paper and this README

Citation

If you find our work useful, please cite our paper:

@misc{li2026dipdynamicincontextplanner,
      title={DIP: Dynamic In-Context Planner For Diffusion Language Models}, 
      author={Yang Li and Han Meng and Chenan Wang and Zhenyu Bi and Xuan Wang and Haipeng Chen},
      year={2026},
      eprint={2601.03199},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2601.03199}, 
}