Alpha Diffusion Language Models: Factorization Alone Is Not the Problem

Nikita Gushchin; Dmitry Baranchuk; Alexander Korotin

cs.LGSep 29, 2026

Alpha Diffusion Language Models: Factorization Alone Is Not the Problem

Authors: Nikita Gushchin, Dmitry Baranchuk, Alexander Korotin

Organizations: Applied AI Institute, Moscow, Russia · AXXX, Russia · Yandex Research

Abstract

Discrete diffusion language models can generate multiple tokens in parallel, but reducing the number of denoising steps can lead to inconsistent predictions. Standard cross-entropy training fits conditional token marginals, whereas parallel generation requires consistent joint predictions. We introduce Alpha Diffusion Language Models (AlphaDLM), trained with a sequence-level alpha loss that recovers cross-entropy in the limit of vanishing alpha and has a joint-mode optimum at alpha one. Our analysis characterizes how the objective and factorization jointly determine the fitted distribution. We identify conditions under which intermediate alpha preserves multiple valid completions while excluding invalid token combinations. Trained on TinyGSM, our method achieves 34.6% accuracy on GSM8K with only four model evaluations. We further scale the method to SDAR-1.7B and evaluate it on code and mathematics benchmarks. These results show that changing the training objective can improve the accuracy-computation trade-off of factorized diffusion language models.

Figures & tables

Figure 1: GSM8K accuracy with 1–16 model evaluations. A single AlphaDLM (ours) checkpoint with

k=16

reaches 34.6% in four evaluations. AlphaDLM and MDLM (confidence) use fixed-NFE confidence sampling with the token-reveal schedule of FMLM+. Other methods use their respective samplers. Sources and protocols : Table 1 and Appendix C .

Figure 2: Marginal fitting, token-mode fitting, and joint-mode fitting. Cross-entropy assigns probability

0.48

to ungrammatical completions. At

\alpha=1

, token-wise alpha loss predicts I is , while sequence-level alpha loss fits the joint mode I am . Each fitted distribution minimizes its population loss within the same factorized family. Green and red mark grammatical and ungrammatical outcomes with positive probability.

Figure 3: Interpolating from cross-entropy to mode fitting. Rows show globally optimal distributions for Figure 2 ’s target: sequence-level loss without factorization, sequence-level loss with factorization, and token-wise loss with factorization. Columns use the same

\alpha

, with

\alpha=0

denoting cross-entropy. At

\alpha=0.5

, the unrestricted optimum

q_{T}

supports all three original completions, while the sequence-level factorized optimum supports only he is and she is . Both sequence-level rows reach I am at

\alpha=1

, whereas token-wise fitting reaches I is . Here

p=p(x_{1},x_{2})

and

p_{i}=p_{i}(x_{i})

. Green marks grammatical combinations, red ungrammatical ones. Entries are rounded.

Method / NFE	4	8	16	32	64
MDLM (ancestral) ( Sahoo et al., 2024 )	1.7	5.3	13.0	22.1	28.2
DUO ( $T=0.1$ ) ( Sahoo et al., 2025 )	5.2	14.0	23.0	30.8	31.2
$\mathbb{S}$ -FLM (exact, $T=0.1$ ) ( Deschenaux and Gulcehre, 2026 )	7.6	12.6	15.1	16.9	16.9
$\mathbb{S}$ -FLM (top-1)	7.7	12.4	14.9	16.5	17.4
FMLM+ ( Agarwal et al., 2026 )	2.9	8.7	13.4	19.0	19.1
FMLM+ (Distill)	3.9	10.3	18.7	21.6	23.4

Table 1: GSM8K accuracy (%)

\uparrow

at fixed NFE. Recomputed MDLM, DUO, and

\mathbb{S}

-FLM results and published FMLM+ and DBTM values. IDLM values are reevaluated at 4–16 NFE and taken from the publication at 32–64 NFE. MDLM (confidence) and AlphaDLM use the same fixed-NFE confidence sampler. Other baselines use their original samplers. Bold values are the column maxima. Protocols are detailed in Appendix C .

Figure 4: How the objective changes parallel prediction. (a) Sequence-level sweep. (b) Token-wise comparison. All runs continue the same 200k-update

k=1

checkpoint to 250k using fixed-NFE confidence sampling. Values: Table 3 . Panels (c,d) continue on page 4.1 .

Figure 5: Ancestral sampling at 32 NFE. Reveal positions are random and independent of confidence. Tokens are selected by argmax (

T=0

) or sampled (

T=1

). Values: Table C .

Figure 6: Accuracy–TPF trade-offs with SDAR-1.7B. (a–d) Objective comparisons on mathematics and code. AlphaDLM (ours) uses

k=0.8

for mathematics and

k=0.4

for code. (e,f) Mathematics parameter sweep. All panels use adaptive confidence sampling with argmax tokens. Points average accuracy and TPF across runs at each confidence threshold. Rings in (a–d) mark threshold 0.99. Appendix D provides standard-deviation bands (Figures D – D ) .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Valid groups and a predictor that mixes them. The target is the two-token example of Figure 3.2 . Rows are the first token and columns the second, as in that figure. Green cells are valid and red cells are invalid. In (b), row heights are $p_{1}$ and column widths are $p_{2}$ , so cell areas equal joint probabilities. The two valid groups receive total probability $0.16+0.36=0.52$ . The remaining $0.48$ falls on invalid combinations.

Panel of Figure 8 — Figure 8: Why the threshold is $1/m$ . For $m=2$ , assigning group 1 probability $s$ at both positions gives $W_{\alpha}(s)=\sum_{g}p(S_{g})^{\alpha}=s^{2\alpha}+(1-s)^{2\alpha}$ . This is the total weight in Equation 22 , not the objective itself. It is below one for $0<s<1$ when $\alpha>1/2$ . At $\alpha=0.55$ , (a) gives $W\approx 0.933$ , while (b) gives $W=1$ .

Panel of Figure 9 — Figure 9: Conditioning improves the running predictor at $\alpha=0.55$ . Both conditioned predictors outperform (a). The comparison in Step 3 gives $F_{\alpha}(p)\approx 0.382\leq F_{\max}\sum_{g}p(S_{g})^{\alpha}\approx 0.388<F_{\max}\approx 0.415$ . Conditioning preserves the relative token probabilities within each group. It does not yet optimize them.

Panel of Figure 10 — Figure 10: Selecting a group preserves all of its completions. At $\alpha=0.55$ , start with a predictor that always generates he is and give she probability $\varepsilon$ at the first position. The objective becomes $0.35(1-\varepsilon)^{0.55}+0.25\varepsilon^{0.55}$ , which increases for small positive $\varepsilon$ . Its maximum assigns positive probability to both valid completions. This best predictor within $S_{2}$ also outperforms the only predictor within $S_{1}$ .

Component	Setting
Denoiser	Bidirectional DiT with 12 blocks, 12 attention heads, hidden width 768, and conditioning width 128
Regularization	Dropout 0.1 and zero weight decay
Sequence length	512 tokens including prompt, response, and padding
Optimizer	AdamW, learning rate $3\times 10^{-4}$ , $(\beta_{1},\beta_{2})=(0.9,0.999)$ , $\epsilon=10^{-8}$
Learning-rate schedule	2,500 warmup updates followed by a constant rate
Optimization	Global batch size 512, gradient clipping 1.0, and training seed 1

Appendix

Table 2: Shared TinyGSM training recipe. Sequence-level and token-wise continuations use the same architecture, data processing, and optimizer settings.

Objective	1 NFE	2 NFE	4 NFE	8 NFE	16 NFE
Sequence $k=1$	0.91	0.91	8.11	26.23	40.56
Sequence $k=2$	0.99	1.29	10.61	28.58	43.06
Sequence $k=4$	1.82	2.58	16.68	35.33	46.47
Sequence $k=8$	2.73	7.73	28.89	43.97	49.20
Sequence $k=16$	6.29	18.04	34.65	39.04	40.56
Token $\alpha=0.25$	0.83	1.74	12.66	33.28	44.35

Appendix

Table 3: GSM8K accuracy (%)

\uparrow

with a shared initialization and fixed-NFE confidence sampling . All checkpoints use 250k total updates, including the common

k=1

prefix through 200k. Each value is an observed correct fraction over 1,319 examples.

Training parameter	Ancestral ( $T=1$ )	Ancestral ( $T=0$ )
k=1	8.19	26.00
k=2	10.46	26.61
k=4	15.09	29.80
k=8	21.30	34.95
k=16	32.52	37.07

Appendix

Table 4: Ancestral accuracy (%)

\uparrow

at 32 NFE. Checkpoints share the

k=1

initialization and 250k total updates. At

T=0

, clean-token predictions are argmax while reveal positions remain stochastic.

Figure 11: Accuracy–NFE trade-offs under adaptive confidence sampling . Left: AlphaDLM (ours) parameter sweep. Right: sequence-level and token-wise objectives. All checkpoints share the

k=1

initialization and 250k total updates. Points vary the confidence threshold. Higher accuracy and fewer evaluations are preferable.

Objective	$\tau=0.6$	$\tau=0.7$	$\tau=0.8$	$\tau=0.9$
Sequence $k=1$	28.58 / 7.64	40.41 / 15.54	42.61 / 28.42	43.06 / 57.59
Sequence $k=2$	26.76 / 6.94	40.41 / 13.09	44.43 / 24.12	45.87 / 49.54
Sequence $k=4$	26.00 / 5.61	38.21 / 9.57	45.87 / 16.72	48.82 / 35.32
Sequence $k=8$	24.72 / 4.06	35.33 / 5.69	42.15 / 8.83	49.43 / 17.83
Sequence $k=16$	17.74 / 2.79	26.31 / 3.37	32.15 / 4.03	40.03 / 5.70
Token $\alpha=0.25$	18.20 / 5.21	32.52 / 7.82	42.30 / 12.86	47.31 / 32.00

Appendix

Table 5: Adaptive confidence sampling from a common initialization. Each entry gives accuracy (%)

\uparrow

/ average NFE

\downarrow

on all 1,319 examples. The checkpoints are evaluated at 250k updates after the common sequence-level

k=1

initialization at 200k.

Figure 12: SDAR accuracy–TPF trade-offs. AlphaDLM (ours) uses

k=0.8

for mathematics and

k=0.4

for code. Points show means at each confidence threshold, with bands of

\pm 1

sample standard deviation in accuracy across runs. Rings mark threshold 0.99. Axes match Figure 4.2 (a–d).

Figure 13: SDAR parameter sweep for AlphaDLM (ours) and CE. Panels (e,f) use the same axes as Figure 4.2 (e,f). Bands show

\pm 1

sample standard deviation in accuracy.

Figure 14: Objective comparisons under two initializations. Colors identify the continuation objective. Solid lines start from

k=1

and dashed lines from CE, both at 200k updates and evaluated at 250k. Only parameters available under both initializations are shown.

Objective	Initialization	1 NFE	2	4	8	16
Sequence $k=2$	$k=1$	0.99	1.29	10.61	28.58	43.06
Sequence $k=2$	CE	0.99	1.14	9.40	25.93	41.17
Sequence $k=4$	$k=1$	1.82	2.58	16.68	35.33	46.47
Sequence $k=4$	CE	1.59	2.12	14.71	31.99	43.37
Sequence $k=8$	$k=1$	2.73	7.73	28.89	43.97	49.20
Sequence $k=8$	CE	2.27	5.91	25.02	40.49	47.61

Appendix

Table 6: GSM8K accuracy (%)

\uparrow

under two initializations. Each paired row uses the same continuation objective and decoding protocol. The final row reports the CE training run continued to 250k, whose 200k checkpoint initializes the CE-start experiments. It is a reference, not a matched-objective comparison with the

k=1

continuation.

Figure 15: GSM8K pass@

K

with ancestral sampling. Top: coverage versus number of answers. Bottom: the same results versus total NFE,

KR

. Colors identify

k

, and line styles identify NFE per answer

R

. Estimates use 32 attempts per problem. Reveal positions remain random at

T=0

$R$	$k$	Pass@1 $\uparrow$	Pass@2 $\uparrow$	Pass@4 $\uparrow$	Pass@8 $\uparrow$	Pass@16 $\uparrow$	Pass@32 $\uparrow$
4	1	2.56	3.62	5.03	6.86	9.33	12.74
4	8	8.51	11.52	15.23	19.91	25.60	32.07
4	16	14.73	19.14	24.17	29.89	36.34	43.37
16	1	16.83	26.23	37.45	48.77	59.04	68.39
16	8	27.33	38.14	48.95	58.56	66.85	74.22
16	16	32.24	42.35	51.62	59.55	66.39	72.71

Appendix

Table 7: GSM8K pass@

K

(%)

\uparrow

estimated from 32 answers per problem.

T

is the sampling temperature,

R

is NFE per answer, and

k

is the training parameter. Bold entries mark the highest value across training parameters for the same

T

R

, and

K

Explore similar work

Factorization-Error-Free Discrete Diffusion Language Model via Speculative Decoding
May 14, 2026Xun Fang, Yunchen Li, Hang Yuan +1Speculative Decoding for Diffusion Language Models Diffusion Language Model Inference
East China Normal University · Beijing Zhongguancun Academy
Representation-Space MMD for Diffusion Language Models
Oct 5, 2026Ilya Drobyshevskiy, Ilia Sudakov, Maksim Semenov +7Diffusion Model Alignment Maximum Mean Discrepancy
Yandex Research · HSE University · T-Tech +3
ALoDLM: Adaptively Looped Diffusion Language Models
Oct 3, 2026Liancheng Fang, Zhuowei Li, Youngeun Kim +10Diffusion Language Models Diffusion Language Model Inference
University of Illinois Chicago · Amazon AGI · Korea University

May 14, 2026J/K move · Enter open · S save

Factorization-Error-Free Discrete Diffusion Language Model via Speculative Decoding

Xun Fang, Yunchen Li, Hang Yuan, Zhou Yu

East China Normal University · Beijing Zhongguancun Academy

Speculative Decoding for Diffusion Language Models Diffusion Language Model Inference Diffusion Language Model Decoding Diffusion Language Models

Open paper AskPDF arXiv

Discrete diffusion language models improve generation efficiency through parallel token prediction, but standard

X_0

prediction methods introduce factorization errors by approximating the clean token posterior with independent token-wise distributions. This paper proposes Factorization-Error-Free Discrete Diffusion Language Modeling (FeF-DLLM), which replaces independent clean-token prediction with an exact prefix-conditioned factorization of the clean posterior to better preserve token dependencies. To reduce the sequential cost introduced by prefix conditioning, FeF-DLLM further incorporates speculative decoding within diffusion denoising, accelerating inference while maintaining the parallel prediction and re-masking properties of DLLMs. Theoretically, we prove that FeF-DLLM generates from the true joint distribution and derive its expected acceleration ratio. Experiments on GSM8K, MATH, HumanEval, and MBPP demonstrate that our method improves accuracy by an average of 5.04 percentage points while achieving an average inference speedup of

3.86\times

Abstract

Explore similar work

Factorization-Error-Free Discrete Diffusion Language Model via Speculative Decoding

Representation-Space MMD for Diffusion Language Models

ALoDLM: Adaptively Looped Diffusion Language Models