Table of Contents
Fetching ...

Rejection Mixing: Fast Semantic Propagation of Mask Tokens for Efficient DLLM Inference

Yushi Ye, Feng Hong, Huangjie Zheng, Xu Chen, Zhiyong Chen, Yanfeng Wang, Jiangchao Yao

TL;DR

ReMix (Rejection Mixing), a framework that introduces a novel Continuous Mixing State as an intermediate between the initial masked state and the final decoded token state, mitigates combinatorial contradictions by enabling continuous-space refinement during discrete diffusion decoding.

Abstract

Diffusion Large Language Models (DLLMs) promise fast non-autoregressive inference but suffer a severe quality-speed trade-off in parallel decoding. This stems from the ''combinatorial contradiction'' phenomenon, where parallel tokens form semantically inconsistent combinations. We address this by integrating continuous representations into the discrete decoding process, as they preserve rich inter-position dependency. We propose ReMix (Rejection Mixing), a framework that introduces a novel Continuous Mixing State as an intermediate between the initial masked state and the final decoded token state. This intermediate state allows a token's representation to be iteratively refined in a continuous space, resolving mutual conflicts with other tokens before collapsing into a final discrete sample. Furthermore, a rejection rule reverts uncertain representations from the continuous state back to the masked state for reprocessing, ensuring stability and preventing error propagation. ReMix thus mitigates combinatorial contradictions by enabling continuous-space refinement during discrete diffusion decoding. Extensive experiments demonstrate that ReMix, as a training-free method, achieves a $2-8 \times$ inference speedup without any quality degradation.

Rejection Mixing: Fast Semantic Propagation of Mask Tokens for Efficient DLLM Inference

TL;DR

ReMix (Rejection Mixing), a framework that introduces a novel Continuous Mixing State as an intermediate between the initial masked state and the final decoded token state, mitigates combinatorial contradictions by enabling continuous-space refinement during discrete diffusion decoding.

Abstract

Diffusion Large Language Models (DLLMs) promise fast non-autoregressive inference but suffer a severe quality-speed trade-off in parallel decoding. This stems from the ''combinatorial contradiction'' phenomenon, where parallel tokens form semantically inconsistent combinations. We address this by integrating continuous representations into the discrete decoding process, as they preserve rich inter-position dependency. We propose ReMix (Rejection Mixing), a framework that introduces a novel Continuous Mixing State as an intermediate between the initial masked state and the final decoded token state. This intermediate state allows a token's representation to be iteratively refined in a continuous space, resolving mutual conflicts with other tokens before collapsing into a final discrete sample. Furthermore, a rejection rule reverts uncertain representations from the continuous state back to the masked state for reprocessing, ensuring stability and preventing error propagation. ReMix thus mitigates combinatorial contradictions by enabling continuous-space refinement during discrete diffusion decoding. Extensive experiments demonstrate that ReMix, as a training-free method, achieves a inference speedup without any quality degradation.
Paper Structure (20 sections, 7 equations, 7 figures, 5 tables, 1 algorithm)

This paper contains 20 sections, 7 equations, 7 figures, 5 tables, 1 algorithm.

Figures (7)

  • Figure 1: The ReMix decoding process and its impact on the quality-speed trade-off. (Left) By introducing an intermediate Continuous State ($\mathbf{C}$), ReMix extends the standard $\mathbf{M}\rightarrow T$ transition by introducing a flexible loop involving Mixing $\mathbf{M} \rightarrow \mathbf{C}$, self-transition $\mathbf{C} \circlearrowright$ and Rejection $\mathbf{C} \rightarrow \mathbf{M}$, enabling in-process refinement. (Right) On GSM8K, ReMix demonstrates superior performance, turning the baseline's 7.05% accuracy loss (when accelerating 3.39x) into a 2.80% accuracy gain while achieving similar acceleration.
  • Figure 2: ReMix resolves combinatorial contradictions by allowing continuous-space refinement. Using the poker hand example, this figure compares the decoding processes. (Top) Vanilla Parallel Decoding: Succumbs to the contradiction. It independently decodes "Full" and "Pair" at positions 11 and 12, resulting in the semantically incorrect "Full Pair" in 8 steps. (Middle) ReMix: Correctly generates "Full House" in only 4 steps. (Bottom) ReMix Output Distribution: This visualizes the internal mechanism. In Steps 1 & 2, both positions are in the Continuous Mixing State (State $\mathbf{C}$, gradient). Their output distributions are iteratively refined, allowing the model to recognize the "Full" $\leftrightarrow$ "House" dependency. This mutual coordination before discretization prevents the error and leads to the correct response. The Ġ symbol (e.g. in ĠFlush indicates a token that begins with a whitespace.
  • Figure 3: Comparison of GSM8K responses (LLaDA-based) score distributions addressing combinatorial contradictions on naive parallel decoding (left) and ReMix (right)
  • Figure 4: Experiment results on fully diffusion-based decoding (generation length = block length).
  • Figure 5: Ablation of the mixing module on GSM8K (left) and MathVista-mini (right).
  • ...and 2 more figures