Multi-Stream Residual Computation in Diffusion Language Models

Ekaterina Alimaskina; Suren Kazarian; Gleb Molodtsov; Aleksandr Beznosikov

cs.LGOct 13, 2026

Multi-Stream Residual Computation in Diffusion Language Models

Authors: Ekaterina Alimaskina, Suren Kazarian, Gleb Molodtsov, Aleksandr Beznosikov

Organizations: BRAIn Lab, Yandex Research

Abstract

Diffusion language models reconstruct masked tokens in parallel, offering an alternative to autoregressive generation. Yet their Transformer backbones typically carry visible evidence and evolving masked predictions through a single residual stream. We equip masked diffusion models with multiple residual streams that persist through depth and adaptive routing between them, then investigate how they support these complementary roles. On the pretraining task, two streams reduce reconstruction perplexity by 5.71% on validation data and 5.13% on average across external corpora compared with standard residual connections. We observe that, without predefined roles, visible and masked positions learn distinct routing patterns. Causal interventions show that this distinction organizes read and write routing, while token-specific write strength remains important within each role. Guided by these findings, a model trained with role-conditioned routing recovers half the validation gain from adding a second stream. State interventions explain the benefit of the additional stream: it preserves early token representations for later reuse while evidence accumulates at masked positions. This mechanism recurs at a larger model scale, showing how persistent streams support context preservation alongside prediction refinement.

Figures & tables

Figure 1: Prepare–preserve–reuse within a denoising pass. Early blocks prepare position-specific states, one stream preserves them while the other accumulates evidence, and late blocks reuse the retained information. Arrows show selected read and write routes; weak connections are omitted.

Figure 2: Two streams capture most gain from additional streams. Validation-CE reduction relative to learned

n=1

. The

n=2

model captures 86.6% of the largest additional gain in the matched stream-count sweep.

Interface	OWT	WT103	LAMBADA	C4	Mean
Standard residual	0.00	0.00	0.00	0.00	0.00
$n=1$	1.44	1.17	$-0.15$	1.77	1.06
$\bm{n=2}$	4.61	7.61	3.02	5.29	5.13
$n=3$	5.51	6.15	2.27	5.17	4.78
$n=4$	5.14	3.30	2.67	5.39	4.13

Table 1: Two streams transfer best on average. Masked-reconstruction perplexity reduction (%) over the standard residual. Higher is better; column bests are bold.

Figure 3: Masked and visible tokens use distinct residual routes. Role-conditioned means remove most degradation from global averaging of read and write destination. Swapping either setting performs worse than the standard residual. A shared mixer nearly preserves performance, while write volume retains substantial within-role dependence. Lower masked-target CE is better. Individual rows replace one control, leaving the others state-dependent; the joint row replaces all four. Estimates and confidence intervals appear in Appendix D .

Figure 4: The functional role of token-specific write volume shifts through depth. (a) At

t=0.5

, within-role shuffling is more damaging than role-mean replacement while preserving the current volume multiset within each sequence and role at every intervened interface. (b) Single-block interventions reveal an early-visible / late-masked sensitivity profile; effects use separate zero-anchored axes. (c) Sensitivity to within-role volume variation decreases near full masking. Error bars and bands show paired 95% sequence-bootstrap intervals;

n=1{,}024

in (a–b) and 128 per corruption level in (c).

Figure 5: The preserved stream has little read/write activity at middle depth and reopens late. Columns show the $32$ attention/MLP interfaces ( $a_{k}$ , $m_{k}$ ); rows show per-stream read share, write share $\overline{g}_{\ell i}/\sum_{j}\overline{g}_{\ell j}$ , and role-mean physical write gain $\overline{g}_{\ell i}$ . Mixer off-diagonal mass is shown once per role, with visible positions above and masked positions below. Scales are $0$ – $1$ for shares, $0$ – $2$ for gains, and $0$ – $0.25$ for mixer mass. Across $a_{1}$ – $a_{8}$ , $s_{\mathrm{P}}$ receives $6\%$ of read mass and $3\%$ of write gain.

Figure 6: Masked targets gather visible evidence in

s_{\mathrm{E}}

during preservation. Mean recovery falls when restoring both source streams (left) and rises when restoring target

s_{\mathrm{E}}

(right). Bands show

95\%

bootstrap intervals; shading marks

s_{\mathrm{P}}

preservation. Details are in Appendix F .

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Intervention	337M	1B
Role mean, masked	97.48	99.51
Role mean, visible	98.22	99.58
Within-role shuffle, masked	93.33	97.09
Within-role shuffle, visible	93.15	97.04
Global mean, both roles	97.59	99.48
Role swap, both roles	95.15	99.10

Appendix

Table 2: Mean interface-wise acceptance coverage (%) for volume interventions at

r=0.5

. Global mean and role swap target both roles; the other rows target the specified role.

Model	Intervention	Masked	Visible	All
337M	Destination, global mean	93.84	98.45	96.15
337M	Destination, role means	97.00	97.61	97.31
337M	Destination, role swap	91.04	98.48	94.76
337M	Joint, global mean	100.00	100.00	100.00
337M	Joint, role means	100.00	100.00	100.00
337M	Joint, role swap	100.00	100.00	100.00

Appendix

Table 3: Acceptance coverage (%) for destination and joint interventions using normalized write destination

q_{i}=g_{i}/(g_{0}+g_{1})

Interface	Parameters	MDLM CE	AR CE
Standard residual	336.1M	3.605514	3.018551
Learned $n=1$	336.2M	3.580209	3.014518
$n=2$	336.6M	3.546667	2.986202
$n=3$	337.4M	3.541492	2.966527
$n=4$	340.4M	3.542203	2.957888

Appendix

Table 4: Final in-domain validation at update 53,399. Learned

n=1

retains token-adaptive scalar gates. Lower CE is better.

Corpus	Streams	$\Delta$ CE	95% interval
OpenWebText	$n=1$	$-.0145$	[ $-.0190,-.0100$ ]
	$n=2$	$-.0472$	[ $-.0521,-.0423$ ]
	$n=3$	$-.0566$	[ $-.0618,-.0516$ ]
	$n=4$	$-.0528$	[ $-.0583,-.0476$ ]
WikiText-103	$n=1$	$-.0118$	[ $-.0178,-.0058$ ]
	$n=2$	$-.0791$	[ $-.0848,-.0735$ ]

Appendix

Table 5: External masked-reconstruction CE change from the standard residual, with paired 95% sequence-bootstrap intervals. All

n\geq 2

models improve on every corpus;

n=2

has the largest mean percentage reduction across corpora.

Interface	$r=.1$	$r=.3$	$r=.5$	$r=.7$	$r=.9$	Mean
Residual	1.7060	2.2986	3.2581	4.5212	6.2436	3.605514
Learned $n=1$	1.6760	2.2593	3.2229	4.4990	6.2439	3.580209
FRR	1.6518	2.2550	3.2054	4.4759	6.2288	3.563375
Full $n=2$	1.6212	2.2361	3.1823	4.4649	6.2289	3.546667

Appendix

Table 6: Final MDLM validation CE by corruption ratio. The mean weights the five ratios equally.

Control	Setting	$\Delta$ CE	95% interval
Read	Global	$+.9493$	[ $+.9406,+.9575$ ]
	Role	$+.0405$	[ $+.0392,+.0419$ ]
	Swap	$+1.8108$	[ $+1.7970,+1.8241$ ]
Write destination	Global	$+1.2551$	[ $+1.2391,+1.2712$ ]
	Role	$+.0171$	[ $+.0162,+.0181$ ]
	Swap	$+2.9771$	[ $+2.9471,+3.0062$ ]

Appendix

Table 7: Increase in masked-target CE over intact

n=2

, with paired 95% intervals. Global means discard role; role means retain masked/visible assignment; swaps exchange it.

Positions changed	Role mean	Within-role shuffle
Masked	$+.0684$ [ $+.0665,+.0703$ ]	$+.1322$ [ $+.1300,+.1345$ ]
Visible	$+.1242$ [ $+.1206,+.1276$ ]	$+.2939$ [ $+.2886,+.2992$ ]

Appendix

Table 8: Volume interventions at

r=0.5

: paired CE increases over the intact checkpoint, with 95% sequence-bootstrap intervals on 1,024 held-out sequences.

Block	Masked-only	Visible-only
0–9	100.00	100.00
10	100.00	98.79
11	100.00	99.19
12	100.00	100.00
13	95.13	99.37
14	86.29	100.00

Appendix

Table 9: Mean acceptance coverage (%) for single-block volume interventions in the 337M model at

r=0.5

. Blocks are indexed from zero; the first row applies to each block individually.

Mask ratio	Masked positions	Visible positions	Both roles
0.1	$+.0481$	$+.1220$	$+.1769$
0.3	$+.0556$	$+.1376$	$+.2014$
0.5	$+.0621$	$+.1256$	$+.1972$
0.7	$+.0547$	$+.0882$	$+.1501$
0.9	$+.0302$	$+.0349$	$+.0672$

Appendix

Table 10: Network-wide role-mean volume replacement at different mask ratios. Each entry is a paired CE increase over its intact checkpoint; 128 sequences are used per ratio.

Mask ratio	Masked-only	Visible-only	Both roles
0.1	97.54	98.50	98.45
0.3	97.45	98.27	98.16
0.5	97.36	98.20	97.99
0.7	97.28	98.03	97.77
0.9	97.45	98.38	97.79

Appendix

Table 11: Mean acceptance coverage (%) for network-wide role-mean volume replacement in the 337M model across mask ratios.

Figure 7: A first-block state persists through the middle network. The preserved stream receives little read or write mass through middle depth, while its cosine similarity to its own post-

m_{0}

state remains high at masked and visible positions. Later read and write activity coincides with transformation of this sample-specific state.

Role	Both set to $s_{\mathrm{E}}$	Both set to $s_{\mathrm{P}}$
Visible	$+.1475$ [ $+.1356,+.1604$ ]	$+.0104$ [ $+.0087,+.0121$ ]
Masked	$+.0772$ [ $+.0741,+.0804$ ]	$+.2304$ [ $+.2184,+.2432$ ]

Appendix

Table 12: Late stream-content intervention after

m_{12}

. Both streams at positions of the selected role are set to either their current

s_{\mathrm{E}}

or their current

s_{\mathrm{P}}

before the remaining network is rerun. Values are paired CE increases over intact

n=2

, with 95% intervals on the same 128 sequences (blocks 256–383) and corruption masks as the preservation tests.

Objective	Streams and role	Middle read share	Late read share
MDLM	$n=2$ , masked / visible	.062 / .059	.386 / .617
MDLM	$n=3$ , masked / visible	.062 / .086	.598 / .639
MDLM	$n=4$ , masked / visible	.051 / .075	.301 / .368
AR	$n=2$	.473	.793
AR	$n=3$	.167	.632
AR	$n=4$	.117	.560

Appendix

Table 13: Read share of the candidate stream selected by minimum middle-depth reading. Diffusion candidates have particularly low middle-depth read share and are read more heavily near output.

Figure 8: Diffusion shows stronger early-state preservation and lower middle-depth reading. For each residual stream count, the top row gives the candidate stream’s read share and the bottom row its cosine similarity to the post-

m_{0}

state. Diffusion is separated by visible and masked roles; AR positions are not partitioned by a corruption mask. Shading marks the middle-depth preservation interval.

Figure 9: Preparation leads to position-specific preservation. (a) Later early-state anchors are less costly to reuse at 1B. (b) Own-state reuse versus same-role shuffling, measured across

m_{0}\!\to a_{8}

(337M) and

a_{8}\!\to a_{16}

(1B). The right axis is logarithmic.

Figure 10: Source-to-target transfer and late visible-state reuse recur at 1B. (a) Reference target-state patches recover nearly all source influence at 75% depth in both models, while source-state patches recover much less. Circles/dashes denote 337M; squares/solid lines denote 1B. (b) At

m_{12}

(337M) and

m_{24}

(1B), duplicating

s_{P}

at visible positions is much less damaging than duplicating the other stream. The figures show reported point estimates. The legend’s “Only” labels denote copying the selected state into both streams.

Block	Masked-only	Visible-only
0	100.00	100.00
31	86.88	99.12

Appendix

Table 14: Mean acceptance coverage (%) for first- and last-block role-mean volume replacement in the 1B model at

r=0.5

. Blocks are indexed from zero.

Figure 11: Role organization and depth-dependent volume assignment recur. (a) Fraction of global-mean damage recovered by role means,

100(1-\Delta\mathrm{CE}_{\mathrm{role}}/\Delta\mathrm{CE}_{\mathrm{global}})

, computed from reported loss increases. (b) First/last-block role-mean volume replacements reproduce the early-visible / late-masked contrast. First- and last-block interventions are shown.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 466 papers

Masked Diffusion Decoding as $x$ -Prediction Flow
Jun 27, 2026Weitian Wang, Lianlei Shan, Shubham Rai +2Masked Diffusion Models Diffusion Language Model Decoding
Robert Bosch GmbH, Germany · University of the Chinese Academy of Sciences, China · Ruhr University Bochum, Germany
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Diffusion Language Model Distillation
NVIDIA
Edit-Based Refinement for Parallel Masked Diffusion Language Models
May 10, 2026Houxing Ren, Mingjie Zhan, Zimu Lu +5Masked Diffusion Models Parallel Decoding
CUHK MMLab · SenseTime Research · Shenzhen Loop Area Institute +1

Jun 27, 2026J/K move · Enter open · S save

Masked Diffusion Decoding as $x$ -Prediction Flow

Weitian Wang, Lianlei Shan, Shubham Rai, Cecilia De La Parra+1

Robert Bosch GmbH, Germany · University of the Chinese Academy of Sciences, China · Ruhr University Bochum, Germany

Masked Diffusion Models Diffusion Language Model Decoding Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens, but their standard decoder reduces each step to a binary action: a position is either committed to a single token or left fully masked, with no representation of partial belief in between. This all-or-nothing regime discards rich predictive information and forces premature, irrevocable commitments, leading to poor performance under a limited decoding budget. In this paper, we reinterpret mask prediction as clean-state prediction (

x

-prediction) and show that it can be used to induce a continuous flow in input embedding space. Building on this view, we propose a continuous decoding framework for MDLMs where tokens can accumulate partial progress at each diffusion step and remain revisable. To match the uneven contextual constraints across positions in language, we replace the globally synchronous schedule in image diffusion with a confidence-based asynchronous update in which the diffusion progress is token-wise accumulated. Additionally, we introduce a lightweight policy network and formulate its training as a reinforcement learning problem. Applied to pretrained LLaDA, our continuous decoder reaches 97% of its performance on the HumanEval dataset with 25% of decoding budget.

Abstract

Diffusion Language Models

Explore similar work

Masked Diffusion Decoding as xxx-Prediction Flow

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Edit-Based Refinement for Parallel Masked Diffusion Language Models

Masked Diffusion Decoding as $x$ -Prediction Flow