# Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

arXiv: https://arxiv.org/abs/2605.22967 · HTML: https://mlsanity.com/paper/2605.22967/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 35 of its 45 references that are in the map.
- First posted on arXiv 21 May 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Chain of Thought Empowers Transformers to Solve Inherently Serial Problems](https://mlsanity.com/paper/2402.12875/) (2024)
- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling](https://mlsanity.com/paper/2409.02908/) (2024)
- [Hierarchical Reasoning Model](https://mlsanity.com/paper/2506.21734/) (2025)
- [Llama-Nemotron: Efficient Reasoning Models](https://mlsanity.com/paper/2505.00949/) (2025)
- [Fast-dLLM v2: Efficient Block-Diffusion LLM](https://mlsanity.com/paper/2509.26328/) (2025)
- [Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://mlsanity.com/paper/2505.24857/) (2025)
- [Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design](https://mlsanity.com/paper/2410.13643/) (2024)
- [Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?](https://mlsanity.com/paper/2410.08292/) (2024)
- [Distillation of Discrete Diffusion through Dimensional Correlations](https://mlsanity.com/paper/2410.08709/) (2024)
- [Continuously Augmented Discrete Diffusion model for Categorical Generative Modeling](https://mlsanity.com/paper/2510.01329/) (2025)
- [Latent Thought Models with Variational Bayes Inference-Time Computation](https://mlsanity.com/paper/2502.01567/) (2025)
- [Soft-Masked Diffusion Language Models](https://mlsanity.com/paper/2510.17206/) (2025)
- [Variational Autoencoding Discrete Diffusion with Enhanced Dimensional Correlations Modeling](https://mlsanity.com/paper/2505.17384/) (2025)
- [Planner Aware Path Learning in Diffusion Language Models Training](https://mlsanity.com/paper/2509.23405/) (2025)
- [Residual Context Diffusion Language Models](https://mlsanity.com/paper/2601.22954/) (2026)

## Built on by (most cited first, 1 of 1)

- [On Trajectory-Aware Training for Masked Diffusion Language Models](https://mlsanity.com/paper/2609.37974/) (2026)

## Abstract (arXiv)

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.