Multi-Stream Residual Computation in Diffusion Language Models
cs.LG, cs.CL
Submitted: 2026-10-13
Updated: 2026-10-13
Comments: 24 pages
Code: https://github.com/brain-lab-research/diffusion-residual-streams
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion language models reconstruct masked tokens in parallel, offering an alternative to autoregressive generation.
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.
Sources
- Analyzing Stream Collapse in Hyper-Connections: From Diagnosis to Mitigation
- Structured Denoising Diffusion Models in Discrete State-Spaces
- Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
- Attention Residuals
- Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
- Fixed-Point Masked Generative Modeling
- DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging
- The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale
- Ablate and Rescue: A Causal Analysis of Residual Stream Hyper-Connections
- Masks Can Be Distracting: On Context Comprehension in Diffusion Language Models
- Qwen2.5 Technical Report
- Simple and Effective Masked Diffusion Language Models
- mHC: Manifold-Constrained Hyper-Connections
- STaR-Quant: State-Time Consistent Post-Training Quantization for Diffusion Large Language Models
- mHC-lite: You Don't Need 20 Sinkhorn-Knopp Iterations
- Dream 7B: Diffusion Large Language Models
- How Does mHC Use Its Residual Streams? Selective Routing and Near-Identity Mixing
- Hyper-Connections
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks