Learning to Couple: Context-Adaptive Joint Denoising for Diffusion Language Models

Zilei Shao; Ian Li; Heng Zhao; Kewei Lian; Guy Van den Broeck; Zhe Zeng; Anji Liu

cs.LGOct 19, 2026

Learning to Couple: Context-Adaptive Joint Denoising for Diffusion Language Models

Authors: Zilei Shao, Ian Li, Heng Zhao, Kewei Lian, Guy Van den Broeck, Zhe Zeng, Anji Liu

Organizations: University of California, Los Angeles · University of California, San Diego · University of Virginia · National University of Singapore

Abstract

Diffusion language models accelerate generation by predicting multiple tokens in parallel, but factorized denoisers treat simultaneous predictions as conditionally independent. Decoding heuristics reduce the resulting errors by controlling which tokens are committed, but leave the underlying distribution factorized. Joint modeling addresses the source of the problem, yet existing methods attach separately learned layers to frozen denoisers, leaving the backbone unadapted to the joint objective and the coupling parameters fixed across contexts. We propose Codpula, a context-parameterized denoiser combining a Transformer with a tractable probabilistic inference layer. Given a partially masked input, the Transformer produces token potentials that shape per-position predictions and compact gates that modulate the layer's coupling parameters, yielding a context-dependent joint distribution while retaining exact conditional inference. We jointly optimize both components under a conditional denoising objective using a dual-flow multiplicative update that preserves valid probabilistic parameters. At inference, Codpula supports standard heuristics and uses conditional queries to adapt commitments to dependencies among proposed tokens without extra Transformer evaluations within a step. Codpula sets new state-of-the-art results across a broad suite of reasoning benchmarks while advancing the accuracy-latency Pareto frontier in mathematical reasoning and code generation.

Figures & tables

Figure 1: Factorized versus context-dependent joint prediction. Left: Given a masked return statement, a factorized denoiser predicts the sorting flag and index independently, assigning equal probability to all four completions regardless of the comment. Middle: Codpula assigns context-dependent joint probabilities, placing its mass on the completions consistent with each comment. Right: Codpula improves accuracy–latency trade-offs, achieving $2.95\times$ and $4.02\times$ speedups at matched accuracy on AIME24 (LLaDA-2.0; Bie et al. (2025) ) and MBPP (Dream; Ye et al. (2025) ).

Figure 2: Block-diffusion NLL on Dream. Negative log-likelihood (nats/token) versus the fraction of masked tokens in the current block. Both Codpula variants substantially reduce NLL when most of the block is masked.

Figure 3: Accuracy-latency trade-offs on LLaDA-1.0 and Dream. Accuracy versus wall-clock time per question for adaptive Codpula, the base backbone, and DUS on math and code benchmarks.

Model	Method	GSM8K (512)				MATH500 (1024)
Model	tokens/step	2	4	8	16	2	4	8	16
LLaDA 2.0	Base	70.3	68.6	36.7	2.4	65.6	55.6	28.2	1.4
	ReMDM	71.0	66.6	39.0	2.4	63.2	56.8	24.4	1.4
	DUS	N/A	N/A	60.3	8.5	N/A	N/A	41.8	4.4
	Tolerator	64.4	56.2	19.9	2.8	58.6	51.6	16.6	2.6
	CE-SFT	80.7	69.2	19.9	1.1	70.0	56.6	13.8	2.0

Table 1: Parallel Decoding Performance (LLaDA 2.0). Accuracy (%) under low-confidence decoding as the number of tokens decoded per diffusion step increases. Parentheses after each benchmark denote the maximum generation length; column headers denote tokens decoded per step. Shaded indicates the best and underline indicates the second-best performance in each column.

Model	Method	HumanEval+				MBPP+
Model	tokens/step	1	2	4	8	1	2	4	8
SDAR	Base	68.9	51.2	28.1	10.4	62.2	41.5	24.6	8.5
	ReMDM	65.2	50.6	28.7	14.0	58.2	43.7	23.3	6.6
	DUS	N/A	N/A	37.8	12.2	N/A	N/A	35.2	12.7
	CE-SFT	64.6	52.4	32.9	11.0	61.9	42.6	23.0	6.6
	CoDD	72.6	55.5	32.9	11.6	65.1	42.9	25.1	9.8

Table 2: Parallel Decoding Performance (SDAR & Dream). Accuracy (%) on HumanEval+ and MBPP+ as the number of tokens decoded per step increases. SDAR uses block size 16, max generation length 512, and low confidence decoding; Dream uses full diffusion with generation length 128 and entropy decoding. Shaded indicates the best and underline indicates the second-best performance within each model block.

Figure 4: Accuracy-latency trade-offs on LLaDA-2.0 and SDAR. Accuracy versus wall-clock time per question for Codpula with uniform and adaptive decoding, base backbone, dynamic-confidence, and EB-sampler decoding on math (LLaDA-2.0) and code (SDAR) benchmarks.

Model	Method	GSM8K (256)			MATH500 (512)			MBPP (512)			HumanEval (512)
Model	tokens/step	2	4	8	2	4	8	2	4	8	2	4	8
LLaDA 1.0	Base	69.1	65.4	37.8	41.0	30.2	10.0	41.3	25.7	10.1	33.5	25.0	11.6
	CE-SFT	66.2	56.0	25.6	37.8	27.6	10.2	41.6	28.0	12.8	34.2	22.0	11.6
	ReMDM	69.6	64.3	37.2	40.2	31.8	11.2	38.1	24.9	8.2	34.8	25.0	12.2
	CoDD	71.5	66.5	50.3	38.6	34.8	22.8	43.2	29.2	9.7	34.8	25.0	12.8
	Codpula	71.7	69.0	52.4	41.6	35.2	20.8	48.3	31.9	15.6	37.2	29.3	16.5

Table 3: Parallel Decoding Comparison across Math and Code Benchmarks (LLaDA-1.0 & Dream). Accuracy (%) as the number of tokens decoded per step increases. Parentheses after each benchmark denote the maximum generation length; column headers denote tokens decoded per step. Math benchmarks use low-confidence decoding; for code, LLaDA uses margin decoding on MBPP and low-confidence decoding on HumanEval, while Dream uses entropy decoding. Shaded indicates the best and underline indicates the second-best performance within each model block.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Total params.	Active params.	Layers	Architecture	Diffusion structure	Block
LLaDA-1.0-8B-Instruct ( Nie et al., 2025 )	8B	8B	32	Dense	Full-sequence masked diffusion	32
LLaDA-2.0-mini ( Bie et al., 2025 )	16B	1.4B	20	MoE	Block-causal masked diffusion	32
Dream-v0-Instruct-7B ( Ye et al., 2025 )	7B	7B	32	Dense	Full-sequence masked diffusion	32
SDAR-8B-Chat-b16 ( Cheng et al., 2025 )	8B	8B	36	Dense	Semi-autoregressive block diffusion	16

Appendix

Table 4: Backbone models used in our experiments. Parameter counts are the nominal sizes reported by the model developers; for dense models, all parameters are active during each forward pass. Block denotes the decoding block length used in our experiments.

Data and objective
	LLaDA-1.0-8B		LLaDA2.0-mini	SDAR-8B	Dream-v0-7B
	Math	Code	Math	Code (b16)	Math	Code
Training scheme	Two-stage	Two-stage	Two-stage	Joint	Joint	Joint
Training examples	1,910	98,660	521,776	101,648	521,776	98,660
Block accum. threshold	64	64	400	32	400	400
Inv.-mask weighting	$10^{-2}$	$10^{-2}$	$10^{-2}$	No	$10^{-2}$	$10^{-2}$

Appendix

Table 5: Codpula training configurations by backbone and task, including parameter conditioning. Two-stage runs first train the PC with the backbone frozen and then train the adapter with the PC fixed. Joint runs optimize the PC and adapter together, and additionally optimize the router for Codpula (Both). Router settings for Input Only and Both are listed separately; all other SDAR settings are shared between the two variants. A dash in a router row indicates that no router is trained. Paired stage-specific entries denote stage-1/stage-2 values, and accumulation thresholds count response blocks. An EM schedule

f

–

e

remains flat for

f

updates and then decays cosinely through update

e

Arm	Trained component	LR / EM step size	Schedule	Batch	Updates	Selected
LLaDA-1.0-8B-Instruct, math numina_hard, 1,910 examples
CE-SFT	LoRA 512/128/0.05	$2{\times}10^{-4}$	Constant	64	673	112
CoDD †§	PC $32{\times}1{,}024$	$0.1{\to}0.01$	67–end	64	673	final
LLaDA-1.0-8B-Instruct, code code_gen, 98,660 examples
CE-SFT	LoRA 128/128/0.05	$10^{-5}$	Constant	64	100	75
CoDD §	PC $32{\times}1{,}024$	$0.3{\to}0.03$	120–end	64	673	final

Appendix

Table 6: Baseline training configurations by backbone and task. CE-SFT trains only the LoRA adapter using masked-diffusion cross-entropy, whereas CoDD trains only the PC with the backbone frozen. LoRA configurations report rank/scale/dropout, and PC configurations report block length

\times

latent states. Batch sizes count sequences for CE-SFT and response blocks for CoDD.

Dataset	Evaluated split	# Samples
Mathematical reasoning
MATH500	Test	500
GSM8K	Test	1,319
AIME 2024	AIME I and AIME II	30
AIME 2025	AIME I and AIME II	30
Code generation

Appendix

Table 7: Datasets and evaluation-set sizes.

Benchmark	Model	$n$	Metric key	System prompt	Canvas
GSM8K	LLaDA-8B, Dream-7B	1,319	exact_match,flexible-extract	none	256
	LLaDA2.0-mini	1,319	exact_match,flexible-extract	none	512
MATH500	LLaDA-8B, Dream-7B	500	math_verify,none	task default	512
	LLaDA2.0-mini	500	math_verify,none	task default	1,024
MBPP (sanitized)	LLaDA-8B, Dream-7B	257	pass_at_1,extract_code	task default	512
HumanEval	LLaDA-8B, Dream-7B	164	pass@1,create_test	task default	512

Appendix

Table 8: Evaluation protocol. Sets are always evaluated in full.

Decoding	Method	GSM8K (256)					MATH500 (512)				MBPP (512) †				HumanEval (512)
Decoding	tokens/step	2	4	8	16	32	2	4	8	16	2	4	8	16	2	4	8	16	32
Random	Base	57.5	42.9	22.1	4.8	1.7	30.8	16.2	7.6	4.0	24.1	12.5	3.9	1.6	18.9	8.5	6.1	–	–
	CE-SFT	52.2	41.1	21.2	–	–	23.4	17.8	9.6	7.6	24.9	13.6	–	–	17.7	9.8	–	–	–
	CoDD	58.5	45.5	24.4	–	–	26.6	21.4	14.0	8.4	27.6	12.5	–	–	17.7	12.2	–	–	–
	Codpula	60.1	49.0	27.9	12.3	5.6	25.4	19.6	12.2	9.0	26.1	17.1	9.7	4.7	21.3	12.8	8.5	–	–
Margin	Base	70.4	64.4	44.3	10.1	1.7	39.2	31.8	18.0	4.6	40.9	24.5	9.0	2.0	35.4	25.0	11.0	4.9	2.4

Appendix

Table 9: Alternative Decoding Rules (LLaDA-1.0). Accuracy (%) under alternative token-selection rules, complementing the main-text results. Parentheses after benchmark names denote maximum generation lengths; column headers denote tokens decoded per step. † For MBPP, the second block uses low-confidence decoding; the other benchmarks use margin decoding. Cells marked – were not run. Shading and underlining indicate the best and second-best accuracy, respectively, in each column within each decoding block.

Decoding	Method / Tokens per Step	MBPP (512)				HumanEval (512)
Decoding	Method / Tokens per Step	2	4	8	16	2	4	8
Low Confidence	Base	56.0	40.1	26.8	8.6	48.8	28.7	11.6
	CE-SFT	52.1	33.9	15.6	2.3	37.8	31.1	14.0
	CoDD	62.3	48.6	33.5	8.9	50.0	37.8	19.5
	Codpula	63.8	49.0	33.1	13.6	50.6	39.6	17.1
	Codpula (Both)	62.3	51.4	33.9	13.6	54.9	33.5	17.7

Appendix

Table 10: Alternative Decoding Rule (Dream, block diffusion). Accuracy (%) on MBPP and HumanEval under low-confidence decoding as the number of tokens decoded per step increases. Parentheses after each benchmark denote the maximum generation length; column headers denote tokens decoded per step. Shaded indicates the best and underline indicates the second-best performance in each column.

Dataset	Model	Tokens per step	Generated tokens	Wall time (s/q)
AIME24	LLaDA2.0-mini	2	7388.2	743.6
	Codpula	2	5334.8	471.1
	LLaDA2.0-mini	4	7870.7	411.2
	Codpula	4	4080.8	163.9
AIME25	LLaDA2.0-mini	2	7137.3	710.5
	Codpula	2	4358.6	359.9

Appendix

Table 11: Generation Length and Runtime with Early Stopping. Mean statistics on the full AIME24 and AIME25 datasets. All other decoding settings are held fixed: maximum generation length 10,240, block length 32, response-anchored block grids, and EOS-based early stopping.

Figure 5: SDAR Accuracy versus Denoising Steps per Block. Accuracy on MBPP+ and HumanEval+ as a function of the mean measured number of denoising steps per block. The results use the same evaluation settings as the SDAR panels in Figure 4 . This view complements the runtime comparison by separating per-block denoising effort from the number of blocks generated before early stopping.

Decoding	Operating point	Overhead
Uniform	Any step budget	$+2.9\%$
Adaptive	NFE $53.8$ , $dw=16$	$+7.7\%$
	NFE $74.9$ , $dw=16$	$+8.2\%$
	NFE $103.0$ , $dw=8$	$+7.7\%$
	NFE $131.0$ , $dw=5$	$+6.9\%$
	NFE $154.2$ , $dw=4$	$+6.8\%$

Appendix

Table 12: Codpula decoding overhead on MATH500 (LLaDA-1.0). Additional time relative to backbone-only decoding at the same NFE. Measurements use batch size 1, merged LoRA adapters, generation length 512, block length 32, and top-

1{,}024

evidence candidates. Adaptive results correspond to five Pareto-front operating points;

dw

denotes the deferral width.

Model	Method	MBPP			MBPP+			HumanEval
Model	tokens/step	2	4	8	2	4	8	2	4	8
Dream	Base	50.3	42.8	32.7	46.0	38.9	31.2	34.1	26.8	24.4
	CE-SFT	46.3	36.6	29.2	43.4	30.4	13.0	39.0	26.8	25.0
	Codpula	52.1	49.4	39.3	51.6	45.2	35.7	34.8	32.9	31.7
	Codpula (Both)	55.6	48.2	38.1	49.2	47.9	33.3	45.1	41.5	29.9

Appendix

Table 13: Parallel Decoding on Code Benchmarks (Dream, full diffusion). Accuracy (%) under full diffusion with generation length 128 and entropy decoding. Column headers denote tokens decoded per step. Shading and underlining indicate the best and second-best accuracy, respectively, in each column.

Model	Method / Tokens per Step	HumanEval+						MBPP+
Model	Method / Tokens per Step	1	2	4	8	16	32	1	2	4	8	16	32
SDAR	Base (static)	72.6	59.2	31.7	12.8	5.5	3.1	59.0	39.7	15.9	3.7	0.3	0.0
	CE-SFT	63.4	45.7	40.9	15.2	4.9	0.0	53.2	34.9	20.4	5.3	0.5	0.0
	Codpula	68.9	59.2	40.9	18.9	10.4	6.7	63.2	50.3	35.2	19.1	5.6	2.7

Appendix

Table 14: Parallel Decoding Performance (SDAR, block size 32). Accuracy (%) on HumanEval+ and MBPP+ as the number of tokens decoded per step increases. Column headers denote tokens decoded per step. Shaded indicates the best and underline indicates the second-best performance in each column.

Figure 6: Probability mass discarded by top-

1{,}024

evidence truncation. Mean and maximum backbone probability mass outside the retained candidates across

46{,}080

positions on LLaDA-1.0, as a function of evidence temperature

\tau

. Shading spans the mean and maximum; the horizontal dashed line marks

1\%

mass. At the training temperature

\tau=1

, the discarded mass averages

0.53\%

and reaches

13.2\%

. At the evaluated decoding temperatures

\tau\leq 0.3

, the mean residual is numerically negligible, on the order of

10^{-8}

Train	Decode	Uniform		Adaptive
Train	Decode	$\tau=0.05$	$\tau=0.1$	$\tau=0.05$	$\tau=0.1$
Top- $k$	Full	63.5	68.0	65.0	64.5
Top- $k$	Top- $k$	63.5	68.0	65.0	64.5
Full	Full	50.5	50.5	64.0	63.0

Appendix

Table 15: Evidence truncation ablation. Accuracy (%) on GSM8K (200 examples), with

k=1{,}024

. Train and Decode indicate the evidence candidates retained at each stage. The first two rows use the same circuit checkpoint. Results are shown for one checkpoint; all eight matched decoding comparisons yield identical accuracy.

Figure 7: Confidence Source for Token Commitment. Accuracy–runtime trade-offs on MBPP+ and HumanEval+ for SDAR with Codpula (Input Only) and Codpula (Both). Both confidence variants use Codpula’s token proposals and dependence-based deferral, with candidate positions ranked using either backbone confidence ( base conf. ) or the refined distribution ( Codpula conf. ). The unmodified backbone is included as a reference.

Arm	MBPP					HumanEval
tokens/step	2	4	8	16	32	2	4	8	16	32
Base	41.3	25.7	10.1	3.9	0.0	33.5	25.0	11.6	3.7	2.4
CE-SFT (adapter only)	41.6	28.0	12.8	4.3	0.0	34.2	22.0	11.6	7.3	1.8
Stage-1 circuit (no adapter)	43.6	27.2	13.2	4.3	1.6	37.8	26.8	13.4	7.9	3.1
CoDD circuit (no adapter)	43.2	29.2	9.7	2.3	0.8	34.8	25.0	12.8	7.9	3.7
Codpula (circuit + adapter)	48.3	31.9	15.6	5.8	1.2	37.2	29.3	16.5	7.3	4.3

Appendix

Table 16: Circuit and Adapter Ablation on Coding Benchmarks (LLaDA-1.0-8B-Instruct). Accuracy (%) for five arms decoded on the same canvas, block grid, and rule; MBPP uses margin and HumanEval uses low-confidence decoding, with

\tau

swept per arm per budget. Stage-1 circuit is our stage-1 circuit decoded without the adapter; CoDD circuit is a circuit trained with CoDD’s External-additive objective, also without an adapter; Codpula is the chosen stage-1 circuit + stage-2 adapter checkpoint. The two difference rows isolate one component each:

\Delta_{\text{adapter}}=\text{Codpula}-\text{Stage-1 circuit}

holds the circuit fixed and measures what the adapter adds, while

\Delta_{\text{circuit}}=\text{Stage-1 circuit}-\text{CoDD circuit}

holds the adapter absent and measures circuit quality alone. Shaded indicates the best and underline indicates the second-best arm in each column.

Figure 8: Additional AIME25 results with LLaDA-2.0-mini. Accuracy (%) versus wall-clock time per question. Adaptive Codpula achieves a

3.84\times

speedup over the base model at matched

10\%

accuracy.

Tokens/step	Steps	$\tau=0.05$	$\tau=0.1$	$\tau=0.3$	$\tau=0.5$
2	256	39.6	36.2	39.4	37.6
4	128	29.4	34.8	31.8	29.4
8	64	17.8	19.4	15.4	17.4
16	32	6.8	10.0	10.8	11.2

Appendix

Table 17: Evidence-temperature sensitivity on MATH500. Accuracy (%) with LLaDA-1.0 under uniform low-confidence decoding, using a fixed adapter and circuit, generation length 512, and block length 32. The table reports the temperature grid shared by all four budgets. Bold indicates the best accuracy among the displayed temperatures in each row.

Figure 9: Accuracy–time trade-offs on AIME24 and AIME25. Results use LLaDA-2.0-mini. Codpula (Input Only) is evaluated with uniform decoding (solid purple) and dependence-adaptive decoding (dashed purple). Each marker denotes an evaluated decoding configuration; lines connect configurations of the same method.

Diffusion Language Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 165 papers

CoDAR: Continuous Diffusion Language Models are More Powerful Than You Think
Mar 3, 2026Junzhe Shen, Jieru Zhao, Ziwei He +1Continuous Diffusion Language Models Diffusion Language Model Decoding
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · School of Computer Science, Shanghai Jiao Tong University · Shanghai Innovation Institute +1
Hierarchical Continuous Diffusion Language Models
Oct 1, 2026Hui Ren, Zihan Li, Chang Liu +2Structured Reasoning in Language Models Diffusion Language Model Decoding
University of Illinois Urbana-Champaign · Amazon.com, Inc.
DA-DLM: Explicitly Modeling Token Dependencies in Diffusion Language Models
Sep 14, 2026Pengyu Ji, Zichen Zhang, Xiang Hu +1Non-Autoregressive Text Generation Diffusion Language Model Decoding
‡School of Information Science and Technology, ShanghaiTech University · §Shanghai Engineering Research Center of Intelligent Vision and Imaging · ¶Tencent

Mar 3, 2026J/K move · Enter open · S save

CoDAR: Continuous Diffusion Language Models are More Powerful Than You Think

Junzhe Shen, Jieru Zhao, Ziwei He, Zhouhan Lin

LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · School of Computer Science, Shanghai Jiao Tong University · Shanghai Innovation Institute+1

Continuous Diffusion Language Models Diffusion Language Model Decoding

Open paper AskPDF arXiv

We study why continuous diffusion language models (DLMs) have lagged behind discrete diffusion approaches despite their appealing continuous generative dynamics. Under a controlled token--recovery study, we identify token rounding, the final projection from denoised embeddings to tokens, as a primary bottleneck. Building on these insights, we propose CoDAR (Continuous Diffusion with Contextual AutoRegressive Decoder), a two--stage framework that keeps diffusion entirely continuous in an embedding space while learning a strong, context--conditional discretizer: an autoregressive Transformer decoder that cross--attends to the denoised embedding sequence and performs contextualized rounding to tokens. Experiments on LM1B and OpenWebText demonstrate that CoDAR substantially improves generation quality over latent diffusion and becomes competitive with strong discrete DLMs, while exposing a simple decoder--temperature knob to navigate the fluency--diversity trade off.

Abstract

Diffusion Language Model Inference

Explore similar work

CoDAR: Continuous Diffusion Language Models are More Powerful Than You Think

Hierarchical Continuous Diffusion Language Models

DA-DLM: Explicitly Modeling Token Dependencies in Diffusion Language Models