Don't Commit Alone: Joint Token Commitment in Diffusion Large Language Models

Abstract: Diffusion LLMs (dLLMs) commit multiple tokens per denoising step by decoding each selected position independently from the shared context; when those positions are dependent, the resulting factorization error is captured by conditional total correlation, which confidence-based selection cannot observe from marginals alone. We propose CoCommit, a marker-gated coordination pass that briefly defers commitment: after the usual bundle selection, a learned marker announces the commit set and the backbone's last-$n$ layers are re-applied so marked positions coordinate -- approximating joint-mode decoding -- before greedy argmax writes tokens. The method reuses existing weights with one extra partial forward pass and no auxiliary model. On LLaDA2.1-mini with LoRA adapters and matched greedy inference, joint commitment improves accuracy on all six benchmarks we evaluate, with the largest gains on reasoning and exact-answer tasks.

• The paper introduces CoCommit, a marker-gated coordination mechanism that refines token commitment in diffusion LLMs.
• It identifies conditional total correlation as the key source of parallel decoding error and quantifies the factorization gap in token commitments.
• Empirical results show consistent gains in reasoning and math tasks, with increased coherence and reduced mode-averaging error.

Joint Token Commitment in Diffusion LLMs: CoCommit

Conditional Total Correlation as the Source of Parallel Decoding Error

Diffusion LLMs (dLLMs) achieve computational efficiency by committing multiple tokens per denoising step during text generation, as opposed to the strictly sequential commitment inherent in autoregressive decoding. However, this parallelism introduces systematic errors: tokens committed in the same step are predicted independently from a shared context, lacking explicit joint dependency modeling. The paper identifies this factorization gap as conditional total correlation $TC(x_S \mid ctx)$ of the commit bundle: parallel decoding replaces the true joint posterior $p(x_S \mid ctx)$ with the product of marginals $\prod_{i \in S} p(x_i \mid ctx)$. $TC$ quantifies the failure to account for inter-token dependencies, which can result in mutually inconsistent outputs.

Crucially, confidence-based selection strategies that determine which tokens to commit rely on peakedness of individual marginals and ignore total correlation—high confidence does not indicate low dependency, as marginal distributions may average over disparate joint modes. This structural blind spot renders marginal-based commitment error estimation inadequate.

Figure 1: Factorized vs. joint commitment—single-layer factorized logits at every position, versus marker-gated coordination that enables within-step token dependence.

CoCommit: Marker-Gated Coordination Mechanism

To address the within-step factorization error, the authors introduce CoCommit: a marker-gated coordination pass that leverages existing model capacity without modifying the backbone architecture or deploying auxiliary models. Each denoising step consists of two sequential stages sharing prefix computation:

• Stage A: The last $n$ layers of the transformer run as usual to select the commit bundle $S$ based on confidence metrics, but commitment is deferred.
• Stage B: A learned marker is added to the hidden states of positions in $S$, signaling them as the current decision set. The last $n$ layers are re-applied, enabling attention-based coordination among marked positions, so their logits are determined jointly, approximating a joint-mode decode. The final tokens are committed greedily per coordinated argmax.

Thus, CoCommit transforms the model's commitment rule: instead of drawing independent argmax from marginalized logits, it computes a tuple from a single joint mode, reducing mode averaging and inconsistency. Under greedy decoding, this deterministic coordination ensures that simultaneously committed tokens resolve to a coherent joint completion.

Figure 2: Split forward pass: training as parallel branches for factorized and coordination; inference as serial stages—commit selection then coordination before writing.

Training Paradigm and Implementation

Empirically, the paper tests CoCommit using LoRA adapters and marker vectors on top of a frozen LLaDA2.1-mini backbone, trained on generic English prose with self-generated token corruption. The training comprises two branches on doubled batches:

• Selection branch: Unmarked, maintains full supervision for factorized output logits as in standard training.
• Assignment branch: Randomly marks bundles of masked positions, passes marker through the last $n$ layers, and applies cross-entropy loss for mask-to-token (M2T) predictions on these bundles. Attention layers are the only channel for bundle coordination.

This setup is not aimed at producing optimal checkpoints, but at stress-testing whether within-step coordination signal can be extracted via lightweight adapter training.

Experimental Results and Benchmark Analysis

The evaluation protocol is matched across conditions, using LLaDA2.1-mini with block-oriented denoising, fixed inference hyperparameters, and greedy decoding. Six benchmarks are used:

Category	Benchmark	Base	CoCommit	$\Delta$
Knowledge	TriviaQA (EM)	43.71	44.63	+0.92
	MMLU-Pro	58.78	58.93	+0.15
Reasoning	DROP (EM)	53.98	58.36	+4.38
Math	CMATH	82.33	84.34	+2.01
	AIME 2025	30.00	33.33	+3.33
	GSM Plus	67.54	68.04	+0.50

CoCommit consistently outperforms the baseline, with marked gains in reasoning and exact-answer tasks (DROP, CMATH, AIME 2025), as these rely on coherent multi-token commitment. The improvement is minimum on knowledge benchmarks and GSM Plus, confirming that within-step coordination is particularly crucial when the output requires a joint-mode resolution. All differences are positive; the largest positive $p(x_S \mid ctx)$0 is +4.38 (DROP).

Practical and Theoretical Implications

This work advances dLLM decoding methodology by explicitly targeting the structural paralellization error. Its principled framing of commitment error via conditional total correlation and marker-gated coordination at inference is a direct intervention at the source of the error, as opposed to downstream repair or remasking strategies.

Practically, CoCommit is minimal: it reuses backbone computation, introduces no new heads or external models, and incurs only a single extra partial forward pass per denoising step. Its effectiveness for reasoning and math tasks suggests broader adoption in domain-specific dLLMs requiring coherent multi-token outputs, such as program synthesis, formal proof generation, and structured answer extraction.

Theoretically, CoCommit connects to conditional joint modeling via deterministic coordination, opening refinements in block-wise joint-mode selection. The mechanism is extensible: iterating the coordination pass would offer adjustable depth of within-step jointness, potentially approximating full block-level conditional diffusion. Marking visible tokens could further unify commitment and revision. The mode-averaging diagnosis generalizes across dLLMs with multi-token parallel commitment, subject to validation in other model families.

Figure 1: Factorized vs. joint commitment in one denoising step—independent argmax per visible/masked position versus marker-enabled joint decode.

Conclusion

CoCommit addresses parallel decoding error in diffusion LLMs by introducing marker-gated coordination, enabling deterministic joint-mode commitment via attention in the backbone's last $p(x_S \mid ctx)$1 layers. The approach yields consistent gains in tasks requiring coherent multi-token outputs, substantiating the impact of within-step coordination. Its minimal training and compute footprint, as well as extensibility to deeper coordination or unified revision, render it a promising candidate for further research and practical deployment in structured generation scenarios.