NI Sampling++: Optimizing Token Order for Fast Discrete Diffusion Sampling with Reinforcement Learning
Enshu Liu · Theodore Zhengde Zhao · Xuefei Ning · Zinan Lin · Yu Wang
We propose a reinforcement learning framework to optimize token ordering for efficient sampling of diffusion language models.
摘要
Recently, discrete diffusion language models (dLLMs) have attracted increasing attention as an alternative to autoregressive language models, as they naturally allow tokens to be generated in flexible orders and support parallel decoding. However, early heuristic sampling strategies remain inefficient, as they update only a small subset of tokens at each step. NI Sampling addresses this issue by introducing token order optimization for fast sampling, where a lightweight neural indicator determines which tokens to update at each step. Nevertheless, the training objective of NI Sampling is only a proxy for the ultimate goal (i.e., achieving high sampling efficiency without sacrificing generation quality), and its empirical speedup still falls short of the validated upper bound. To bridge this gap, we reformulate token order optimization as a Markov decision process and optimize the neural indicator using reinforcement learning. Experimental results demonstrate that our method achieves further sampling acceleration compared to NI Sampling.