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

1. Problem Statement

The paper addresses a critical bottleneck in Diffusion Large Language Models (DLLMs): the severe trade-off between inference speed and generation quality.

• The Context: Unlike autoregressive (AR) models that generate tokens sequentially, DLLMs generate entire sequences in parallel, offering the potential for massive speedups.
• The Core Issue: The authors identify a phenomenon they term "Combinatorial Contradiction." In standard parallel decoding, tokens are sampled independently at each step. Because they cannot perceive one another during the sampling process, they often form semantically inconsistent combinations (e.g., generating "High House" instead of "Full House" in a poker context).
• Limitations of Existing Solutions: Previous methods either rely on complex verification blocks (introducing overhead), enforce left-to-right generation (defeating the purpose of parallelism), or fail to resolve the fundamental lack of inter-token dependency in discrete sampling.

2. Methodology: Rejection Mixing (ReMix)

The authors propose ReMix, a training-free decoding framework that bridges the gap between discrete and continuous representations to resolve combinatorial contradictions.

Core Concept

ReMix introduces an intermediate Continuous Mixing State ($C$) between the initial Masked State ($M$) and the final Token State ($T$). This allows the model to refine token representations in a continuous space before committing to a discrete token, enabling mutual coordination among positions.

The Three-State Transition Mechanism

The decoding process for each token position involves three states and specific transition rules:

1. Initialization ($M$): All positions start as [MASK].
2. Decoding Rule ($M/C \to T$): If the model's confidence for a token exceeds a threshold ($\tau_{conf}$), the position transitions directly to the discrete token state ($T$).
3. Mixing Rule ($M \to C \leftrightarrow C$): If a position is not confident enough to be decoded, it transitions to the Continuous State ($C$).
   • Instead of being re-masked, the position retains a continuous embedding derived from the model's output distribution.
   • This embedding is iteratively updated (refined) across decoding steps, allowing the model to "sense" dependencies with other tokens in the sequence.
   • The update formula blends the predicted distribution with the [MASK] embedding:
     $$y^{emb}_i \leftarrow \beta W^\top p_\theta(\hat{y}_i | X, Y) + (1 - \beta)Emb[\text{MASK}]$$
   • An Adaptive Top-p Sampling strategy is used to ensure numerical stability, consolidating low-probability mass back into the [MASK] embedding.
4. Rejection Rule ($C \to M$): To prevent error propagation from unstable continuous states, a rejection mechanism is applied.
   • If the output distribution of a position in the continuous state fluctuates significantly between steps (measured by Jensen-Shannon Divergence $> \tau_{rej}$), the position is rejected and reset to the [MASK] state ($M$).
   • This acts as a regularizer, pruning divergent paths early and forcing the model to re-evaluate the token from scratch.

3. Key Contributions

• Identification of "Combinatorial Contradiction": The paper formally defines the root cause of quality degradation in parallel DLLM decoding as the inability of discrete, independent sampling to maintain semantic consistency across positions.
• ReMix Framework: A novel, training-free decoding strategy that integrates continuous representations into the discrete diffusion loop. It enables iterative refinement of token intents before discretization.
• Stability Mechanism: The introduction of a Rejection Rule ensures that uncertain continuous representations do not corrupt the generation process, maintaining the stability of the diffusion model trained on discrete tokens.
• Efficiency: The method requires no retraining or fine-tuning of the backbone model, making it immediately applicable to existing DLLMs.

4. Experimental Results

The authors evaluated ReMix on diverse benchmarks across language and multimodal domains.

• Language Domain (LLaDA-8B):
  • Tested on 8 benchmarks including GSM8K, MATH-500, HumanEval, and ARC.
  • Speedup: Achieved 2.4× to 4.6× end-to-end inference speedup (reducing decoding steps by 2.5×–5.0×).
  • Quality: Contrary to typical speed-accuracy trade-offs, ReMix improved accuracy on all benchmarks (e.g., +2.65% on GSM8K, +14.05% on ARC-C).
• Multimodal Domain (MMaDA-8B):
  • Tested on 6 vision-language tasks (Flickr30k, MathVista, MMMU, etc.).
  • Speedup: Achieved 3.75× to 7.52× speedup.
  • Quality: Consistent accuracy improvements across all tasks (e.g., +3.00% on MathVista-mini).
• Ablation Studies:
  • Removing the mixing module (setting $\beta=0$) degraded performance, confirming the necessity of continuous refinement.
  • The method remains robust across different generation lengths and block sizes.
  • Computational overhead for ReMix-specific operations is negligible (~9% of total runtime).

5. Significance

• Breaking the Trade-off: ReMix successfully decouples the traditional inverse relationship between inference speed and generation quality in diffusion models. It proves that parallel decoding can be both fast and accurate without sacrificing the model's training distribution.
• New Paradigm for Decoding: By introducing a "soft" continuous intermediate state, the paper suggests a new direction for generative modeling where discrete and continuous representations are dynamically integrated during inference.
• Practical Impact: As a training-free method, ReMix offers an immediate, low-cost solution for deploying efficient DLLMs in real-world applications, potentially narrowing the performance gap between diffusion models and autoregressive models while retaining the parallel inference advantages of diffusion.