Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

[2605.22967] Learned Relay Representations for Forward-Thinking Discrete Diffusion Models Skip to main content Search Submit Donate Log in Computer Science > Machine Learning arXiv:2605.22967 (cs) [Submitted on 21 May 2026 ( v1 ), last revised 28 Sep 2026 (this version, v4)] Title: Learned Relay Representations for Forward-Thinking Discrete Diffusion Models Authors: Benjamin Rozonoyer , Jacopo Minniti , Dhruvesh Patel , Neil Band , Avishek Joey Bose , Tim G. J. Rudner , Andrew McCallum View PDF HTML (experimental) Abstract: When Masked Diffusion Models (MDMs) generate sequences through iterative refinement, the rich internal computation over masked positions is discarded, forcing every subsequent refinement step to recompute the valuable internal information stored as model representations. To avoid a hard reset between denoising rounds, we propose Learned Relay Representations (Relay), a method that allows MDMs to be forward-thinking when denoising by explicitly learning how to propagate latent information for the benefit of future denoising steps. Relay introduces a differentiable per-token channel that passes information between forward passes and is trained via truncated backpropagation through time (BPTT). We show that this framework can be scaled to state-of-the-art Diffusion Language Models (DLMs), and is seamlessly compatible with techniques like block diffusion and KV caching. We first provide a thorough justification of the design choices in Relay on a challenging Sudoku-based planning task. We then scale Relay to Fast-dLLM v2, a state-of-the-art DLM, outperforming standard supervised finetuning on coding tasks while reducing inference latency by up to 32%. Our empirical results demonstrate that state-of-the-art DLMs can be explicitly trained to relay latent information forward across decoding steps, advancing the performance-latency Pareto frontier. We provide code for all our experiments. Comments: 18 pages, 3 figures. Equal contribution: Benjamin Rozonoyer, Jacopo Minniti, and Dhruvesh Patel. Code: this https URL Subjects: Machine Learning (cs.LG) ACM classes: I.2.6; I.2.7 Cite as: arXiv:2605.22967 [cs.LG] (or arXiv:2605.22967v4 [cs.LG] for this version) https://doi.org/10.48550/arXiv.2605.22967 Focus to learn more arXiv-issued DOI via DataCite Submission history From: Benjamin Rozonoyer [ view email ] [v1] Thu, 21 May 2026 18:53:22 UTC (89 KB) [v2] Fri, 29 May 2026 16:01:35 UTC (89 KB) [v3] Fri, 28 Aug 2026 02:33:55 UTC (89 KB) [v4] Mon, 28 Sep 2026 17:23:42 UTC (92 KB) Full-text links: Access Paper: View PDF HTML (experimental) TeX Source view license Current browse context: cs.LG < prev | next > new | recent | 2026-05 Change to browse by: cs References & Citations NASA ADS Google Scholar Semantic Scholar export BibTeX citation Bookmark Bibliographic Tools Bibliographic and Citation Tools Bibliographic Explorer Toggle Bibliographic Explorer ( What is the Explorer? ) Connected Papers Toggle Connected Papers ( What is Connected Papers? ) Litmaps Toggle Litmaps ( What is Litmaps? ) scite.ai Toggle scite Smart Citations ( What are Smart Citations? ) Code, Data, Media Code, Data and Media Associated with this Article alphaXiv Toggle alphaXiv ( What is alphaXiv? ) Links to Code Toggle CatalyzeX Code Finder for Papers ( What is CatalyzeX? ) DagsHub Toggle DagsHub ( What is DagsHub? ) GotitPub Toggle Gotit.pub ( What is GotitPub? ) Huggingface Toggle Hugging Face ( What is Huggingface? ) ScienceCast Toggle ScienceCast ( What is ScienceCast? ) Demos Demos Replicate Toggle Replicate ( What is Replicate? ) Spaces Toggle Hugging Face Spaces ( What is Spaces? ) Spaces Toggle TXYZ.AI ( What is TXYZ.AI? ) Related Papers Recommenders and Search Tools Link to Influence Flower Influence Flower ( What are Influence Flowers? ) Core recommender toggle CORE Recommender ( What is CORE? ) IArxiv recommender toggle IArxiv Recommender ( What is IArxiv? ) Author Venue Institution Topic About arXivLabs arXivLabs: experimental projects with community collaborators arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website. Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them. Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs . Which authors of this paper are endorsers? | Disable MathJax ( What is MathJax? ) We gratefully acknowledge support from our major funders , member institutions , and all contributors. About · Help · Contact · Subscribe · Copyright · Privacy · Accessibility · Operational Status (opens in new tab) Major funding support from

Related documents

• Learned Relay Representations for Forward-Thinking Discrete Diffusion Models#682689
• Physically-Relevant Information Learning in High-Dimensional Time-Derivatives Spaces#343499
• Gradient-Momentum Coupling: A Parameter-Space Proxy for Learning Progress#1105035
• Energy-Based Dynamical Models for Neurocomputation, Learning, and Optimization#785143
• Energy-Based Dynamical Models for Neurocomputation, Learning, and Optimization#783553