D5P4: Partition Determinantal Point Process for Diversity in Parallel Discrete Diffusion Decoding

Abstract: Discrete diffusion models are promising alternatives to autoregressive approaches for text generation, yet their decoding methods remain under-studied. Standard decoding methods for autoregressive models, such as beam search, do not directly apply to iterative denoising, and existing diffusion decoding techniques provide limited control over in-batch diversity. To bridge this gap, we introduce a generalized beam-search framework for discrete diffusion that generates candidates in parallel and supports modular beam-selection objectives. As a diversity-focused instantiation, we propose D5P4, which formulates the selection step as MAP inference over a Determinantal Point Process. Leveraging a scalable greedy solver, D5P4 maintains multi-GPU compatibility and enables an explicit trade-off between model probability and target diversity with near-zero compute overhead. Experiments on free-form generation and question answering demonstrate that D5P4 improves diversity over strong baselines while maintaining competitive generation quality.

• The paper introduces D5P4, a partition-constrained Determinantal Point Process method that selects diverse descendants from parallel diffusion beams using model-internal quality and representation signals.
• D5P4 delivers a wider quality–diversity trade-off than temperature scaling, diverse beam search, and best-of-k, including a MAUVE* score of 0.9615 versus 0.9134 and over 10× faster greedy selection than diverse beam search.
• The method mitigates classifier-free-guidance mode collapse in LLaDA while preserving comparable task quality, reducing average answer similarity and Self-BLEU with near-zero additional inference cost.

D5P4 addresses a specific gap in the decoding of masked discrete diffusion LLMs (MDLMs): while such models refine all sequence positions in parallel, they lack a principled beam-search analogue that controls diversity across parallel hypotheses. The paper proposes D5P4, which casts each intermediate selection step as MAP inference over a partition-constrained Determinantal Point Process (DPP), using only signals already computed by the diffusion model. The method is evaluated on open-ended generation with MDLM and question answering with LLaDA, showing improved quality–diversity trade-offs relative to temperature scaling, diverse beam search, and best-of-$k$ baselines at near-zero compute overhead.

Motivation and problem setting

The authors motivate their work by two observations. First, discrete diffusion models such as MDLM [2411.xxxx-style recipes] and LLaDA achieve competitive performance with autoregressive models but rely on simple sampling procedures; autoregressive decoding machinery like beam search does not transfer because diffusion trajectories are non-monotonic and parallel rather than left-to-right. Second, output diversity is degrading across modern LLM pipelines: reinforcement learning and supervised fine-tuning saturate pass@$k$ coverage in ARMs, and strong classifier-free guidance (CFG) induces mode collapse in both continuous and discrete diffusion. The paper argues that diversity should be reasoned about at the set level during decoding, not left to independent sampling.

A notable structural feature of the paper is that it does not claim novelty on the diffusion backbone itself. It builds directly on MDLM's Rao-Blackwellized training objective under a linear schedule $\alpha_t = 1-t$ and LLaDA's remasking-based inference, unifying both under a projection operator $\Pi_{t,s}$ that maps denoising logits to the next noisy state.

Beam-style decoding for discrete diffusion

The framework maintains $k$ beams with branching factor $w$, producing $n = k \cdot w$ candidates per diffusion step by applying the stochastic projection operator to each retained beam. Because diffusion models admit no monotonic prefix likelihood, candidate quality is scored with sequence-level proxies: token-level entropy and self-certainty (KL divergence between logits and uniform). The authors then impose a transversal partition constraint: candidates are grouped by parent beam, and exactly one descendant per group must be selected, preventing lineage collapse toward a single ancestry — the failure mode identified for standard beam search in autoregressive settings.

Partition DPP selection

Selection is formulated as MAP inference over an $\mathbf{L}$-ensemble DPP, where $\mathbb{P}(S) \propto \det(\mathbf{L}_S)$. The kernel combines per-sequence quality scores $\mathbf{Q}$ with pairwise similarities computed from hidden representations immediately preceding the unembedding layer — representations that are available without extra forward passes:

$k$0

Here $k$1 is a cosine or RBF kernel over normalized sequence embeddings, and $k$2 provides an interpretable quality–diversity knob. A useful observation made explicit in the paper: standard top-$k$3 beam search is the MAP solution of a diagonal-only kernel, so DPP-MAP strictly generalizes conventional beam search.

Exact MAP is NP-hard, so the authors extend the fast greedy solver of Chen et al. (2018) to handle the transversal partition constraint, adding multi-initialization from each group's argmax run in parallel. The extension raises asymptotic complexity from $k$4 to $k$5, but GPU parallelization keeps wall-clock cost essentially unchanged. No closed-form solution exists for DPPs under partition constraints, so greedy approximation is the only practical route here.

Alignment of internal signals

A preliminary analysis justifies using diffusion-internal signals instead of external evaluators. Diffusion entropy estimates correlate strongly with autoregressive log-likelihoods (Spearman $k$6 for both MDLM/GPT-2 and LLaDA/Llama-3; Pearson $k$7 for MDLM Monte-Carlo likelihood vs. GPT-2), and CKA between diffusion latents and Jina embeddings reaches 0.821 for MDLM. This matters practically: the entire selection procedure runs without external scorers or reward-model calls. An ablation shows flattened sequence embeddings yield the best CKA alignment (0.821 for MDLM, 0.667 for LLaDA), and entropy scoring correlates far better with reference perplexity ($k$8) than self-certainty ($k$9).

Open-ended generation results

On FineWeb-driven generation with MDLM (batch of 32, 8 groups of 4), systematic parameter sweeps show that all search-based methods dominate naive independent sampling on the perplexity–cosine-similarity Pareto front. The key differentiator is failure behavior: categorical temperature scaling exhibits abrupt perplexity blow-up beyond a narrow operating range, diverse beam search (implemented as transversal MMR) degrades earlier, while both D5P4 variants delay perplexity degradation furthest. The additive kernel dominates in the high-diversity regime; the multiplicative kernel preserves likelihood structure over a wider cosine range. Parameter correlations confirm controllability — e.g., $\alpha_t = 1-t$0 correlates at $\alpha_t = 1-t$1 with PPL and $\alpha_t = 1-t$2 with cosine similarity for the multiplicative variant. MAUVE experiments reveal an optimal intermediate $\alpha_t = 1-t$3: MAUVE* peaks at 0.9615 ($\alpha_t = 1-t$4) versus a baseline of 0.9134, indicating that moderate diversification actually improves distributional fidelity rather than merely trading quality for coverage.

Question answering and CFG collapse mitigation

With LLaDA on TruthfulQA and CommonsenseQA, increasing CFG strength monotonically reduces lexical and semantic diversity (Distinct-2, Self-BLEU, cosine similarity), consistent with prior findings on guidance-induced mode collapse. D5P4 counteracts this: even at high CFG values it preserves substantially higher diversity while maintaining comparable F1, Wasserstein distance, and perplexity. Under FLOP-matched comparison against best-of-$\alpha_t = 1-t$5, D5P4$\alpha_t = 1-t$6 improves average in-batch cosine similarity from 0.963 to 0.946 (TruthfulQA) and 0.969 to 0.92 (CommonsenseQA), lowers Self-BLEU from 47.1 to 40.4 and 52.6 to 43.0 respectively, and slightly improves perplexity (17.4 → 15.7 and 27.5 → 26.0). Accuracy metrics are essentially flat or marginally lower (F1 0.212 → 0.184 on TruthfulQA), which the authors present as competitive quality. Combining D5P4$\alpha_t = 1-t$7 with partial CFG (guidance applied only during the first half of denoising steps) yields the best overall profile on several metrics, including the lowest average cosine similarity (0.918 / 0.859) and highest EAD.

One caveat worth noting: absolute accuracy numbers are low across all methods (F1 ≤ 0.24, BLEU ≤ 5.0), so the QA evaluation primarily demonstrates trade-off control rather than large correctness gains.

Solver efficiency

On synthetic kernels with 32 groups of 32 items, the greedy MAP solver achieves the highest normalized subdeterminant (1.0214) at 0.0023 s — over 10× faster than diverse beam search (0.6645 at 0.0295 s), while exact DPPy sampling takes 0.55 s and cannot enforce transversality. Scaling studies up to 64×64 configurations confirm the advantage grows with combinatorial complexity, and a Triton-optimized variant further reduces runtime. The overhead relative to random selection is dominated by GPU synchronization, supporting the claim of negligible inference cost.

Limitations and open questions

Several limitations are acknowledged or evident. The greedy MAP is an approximation with no optimality guarantee under partition constraints, and no closed-form alternative exists. Quality and diversity estimation rely on proxy signals whose validity was established empirically via correlation with external evaluators on two specific model/dataset pairings; generalization to other models is assumed rather than proven. Notably, an appendix fragment indicates that applying the approach to autoregressive models "doesn't work because the internal representation is not rich enough," bounding the method's scope to diffusion architectures. Comparisons with Particle Gibbs inference-time scaling are explicitly avoided as computationally incomparable, leaving open how D5P4 compares against reward-guided MCMC methods under matched budgets. Evaluation covers two models and a small set of tasks; whether the Pareto advantages hold for larger-scale diffusion LLMs or other modalities remains untested.

Conclusion

D5P4 contributes a modular beam-selection framework for discrete diffusion in which diversity-aware selection is cast as partition-DPP MAP inference solvable greedily on GPU at negligible cost. Its empirical contribution is a demonstrably wider and more predictable quality–diversity Pareto front than temperature scaling or MMR-style search, plus a concrete mitigation of CFG-induced mode collapse in question answering. The main open questions are the theoretical properties of greedy MAP under transversal constraints and the portability of the entropy/embedding kernel design beyond the evaluated MDLM and LLaDA setups.