🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

This paper introduces Learned Relay Representations (Relay), a method that enables Masked Diffusion Models to propagate latent information across denoising steps via a differentiable per-token channel, thereby improving performance on planning and coding tasks while significantly reducing inference latency.

Original authors: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Published 2026-05-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Hard Reset" Amnesia

Imagine you are trying to solve a complex puzzle, like a Sudoku or writing a piece of code. You are using a very smart assistant (an AI model) to help you.

Currently, this assistant works like a person with amnesia who forgets everything the moment they finish a single sentence.

The assistant looks at a blank page full of question marks.
It thinks hard, does some complex mental math, and decides to fill in one or a few letters.
The Hard Reset: As soon as it writes those letters down, it immediately throws away all the complex thoughts, calculations, and "feeling" it had about the rest of the page.
For the next step, it starts over from scratch, looking only at the letters it just wrote, completely forgetting the rich internal work it just did.

This is called the "Hard Reset" problem. The paper argues that this is inefficient. It's like a chef tasting a soup, writing down "needs salt," and then immediately forgetting the taste of the broth before adding the salt. They have to re-taste the whole soup from scratch every time they add an ingredient.

The Solution: The "Relay" Baton

The authors propose a new method called Relay.

Imagine a relay race. Instead of the runner stopping, forgetting the race, and starting over, they pass a baton to the next runner. This baton carries the momentum, the strategy, and the fatigue of the previous runner so the next one can pick up exactly where they left off.

In the AI world, the "baton" is a continuous stream of information (a hidden state) that the model carries forward.

Before Relay: The model computes, writes a token, and deletes the computation.
With Relay: The model computes, writes a token, and passes a "note" to itself for the next step. This note says, "Hey, I was thinking about this specific pattern, and I'm 80% sure about this next part."

How It Works: Learning to Pass the Baton

The paper introduces a few key tricks to make this work:

The "Forward-Thinking" Training:
Usually, AI models are trained to just get the current answer right. The authors trained the model to be "forward-thinking." They taught the model: "Don't just write the right letter today; write a note (the relay) that will help you write the right letter tomorrow."
Truncated Backpropagation (BPTT):
This is a fancy math term for "looking ahead to learn." Imagine you are practicing a dance routine. Instead of just practicing one move, you practice a sequence of three moves. If you stumble on the third move, you realize, "Oh, I messed up the first move because I didn't prepare for the third."
The paper uses a method called truncated BPTT to let the model look a few steps ahead during training. It learns to pass information forward specifically to make future steps easier.
The "Soft" Channel:
The information passed forward isn't just a word or a letter; it's a "soft" number (a continuous value). Think of it like a dimmer switch rather than an on/off light. It allows the model to carry subtle hints and probabilities that haven't been turned into a final decision yet.

The Results: Faster and Smarter

The authors tested this on two things:

Sudoku Puzzles: They treated Sudoku like a planning task. The "Relay" model solved puzzles with fewer attempts (fewer "forward passes") and made fewer mistakes than the standard model. It was better at keeping the "big picture" in mind while filling in the grid.
Coding (Writing Software): They took a state-of-the-art coding AI (Fast-dLLM v2) and added the Relay system.
- Better Accuracy: It wrote better code.
- Faster Speed: It needed to run its "thinking engine" 32% fewer times to get the same result.

Why This Matters

The paper claims that by stopping the "Hard Reset" and letting the model carry its own "thoughts" forward like a relay baton, we can make AI models:

Smarter: They can reason better over long sequences (like complex code or logic puzzles).
Faster: They don't have to re-compute everything from scratch every time, saving time and energy.

In short, the paper teaches AI models to stop forgetting their own work and start building on it, step by step, just like a human does when solving a hard problem.

Technical Summary: Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

1. Problem Statement: The "Hard Reset" in Masked Diffusion Models

Masked Diffusion Models (MDMs) generate discrete sequences through iterative denoising. Starting from a fully masked sequence, the model iteratively unmasks a fraction of positions. However, a critical inefficiency exists in standard MDM inference: after each forward pass, the model discards the rich internal continuous computations (hidden states) generated for masked positions. The next step begins solely from the newly committed discrete tokens.

The authors term this the hard reset problem. While MDMs perform multiple forward passes per generation, the inability to reuse intermediate continuous computation prevents the model from effectively amortizing reasoning across steps. This contrasts with recurrent computation, where unrolling a fixed-parameter model across steps has been shown to expand the function class the model can approximate, aiding performance on complex reasoning tasks. The central question addressed is how to enable MDMs to support recurrent computation that carries richer information across sequential unmasking steps without discarding intermediate latent states.

2. Methodology: Learned Relay Representations (Relay)

The proposed solution, Relay, introduces a continuous, differentiable latent state that is passed forward across decoding steps, allowing the model to be "forward-thinking."

Core Mechanism

Augmented State Trajectory: At each inference step $k$ , the model maintains an augmented state $s_k = (x_{t_k}, h_k)$ , where $x_{t_k}$ is the current partially masked sequence and $h_k$ is a learned relay state (continuous latent vector).
Architecture: The model consists of a backbone $f_\theta$ $f_{θ}$ , token embeddings $Emb_\theta$ $E m b_{θ}$ , an unembedding head $UnEmb_\theta$ $U n E m b_{θ}$ , and a specific relay module $R_\theta$ $R_{θ}$ .
- The next hidden state is computed as: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ .
- The relay state $h_k$ is projected and added to the embedded tokens before entering the backbone.
- The output logits $\ell_k$ are derived from $h_{k+1}$ via $UnEmb_\theta$ .
Training via Truncated BPTT: Unlike standard MDMs trained with independent steps, Relay is trained using Truncated Backpropagation Through Time (BPTT).
- The model performs a rollout of $K$ steps (e.g., $K=2$ ) starting from a fully masked sequence.
- During the rollout, positions are unmasked based on a policy (e.g., confidence thresholds), and the committed tokens are "teacher-forced" to ground truth values to prevent error accumulation during the rollout.
- The loss is the sum of cross-entropy losses over the $K$ steps.
- Gradients are backpropagated through the relay path ( $h_k \to h_{k+1}$ ), allowing the model to learn not just to predict the current step, but to encode information in $h_k$ that is maximally useful for future steps.

Inference

At inference time, the process remains largely unchanged in terms of the unmasking schedule or sampling policy. The only addition is that the model carries the relay state $h_k$ forward alongside the committed tokens. The relay state serves as a differentiable memory channel for information not yet committed to discrete tokens.

3. Key Contributions

Relay Framework: A general method for incorporating recurrent computation into MDMs by training the model to pass a learned latent relay forward across decoding steps via truncated BPTT. It is architecture-agnostic and compatible with existing MDM acceleration techniques like block diffusion and KV caching.
Scalability and Adaptation: Validation of Relay at the Large Language Model (LLM) scale through full-parameter adaptation of Fast-dLLM v2 (1.5B). The method successfully adapts a state-of-the-art Discrete Language Model (DLM) with minimal additional training FLOPs.
Performance-Latency Improvements: Empirical demonstration that Relay advances the performance-latency Pareto frontier, outperforming standard supervised fine-tuning (SFT) on coding tasks while significantly reducing inference latency.
Design Space Analysis: Extensive ablation studies on a challenging Sudoku-based planning task that isolate the contributions of the rollout procedure, the relay channel, and the BPTT training objective.

4. Experimental Results

Sudoku Planning Task

Using the Sudoku-Extreme dataset, the authors compared four training objectives:

MLM: Standard uniform masked diffusion (single step, no relay).
Rollout: Multi-step rollout with teacher forcing but no relay channel ( $R_\theta \equiv 0$ ).
Relay (sg): Includes the relay path but uses stop-gradients (no BPTT through the relay).
Relay: Full method with $K=2$ BPTT through the relay.

Findings:

Accuracy-NFE Frontier: Relay achieved the best trade-off between accuracy and Number of Forward Evaluations (NFE).
Impact of Components: Moving from MLM to Rollout improved performance. Adding the relay channel (Relay sg) provided a significant jump. Finally, enabling BPTT (Relay) yielded the best results, demonstrating that the model learns to keep the partial board self-consistent under more aggressive unmasking.
Efficiency: At a fixed confidence threshold, Relay committed more cells per forward pass while maintaining legality (no row/column/box violations), reaching the same accuracy in fewer total forward passes compared to baselines.

Pretrained Model Adaptation (Fast-dLLM v2)

The authors adapted Fast-dLLM v2 (1.5B) using Relay on a mixture of code and math datasets.

Coding Performance: On HumanEval and MBPP benchmarks, Relay outperformed standard SFT.
- HumanEval: Relay achieved 42.1% accuracy (vs. 38.4% for SFT) with 88.3 NFE (vs. 130.7 for SFT).
- MBPP: Relay achieved 46.6% accuracy (vs. 43.9% for SFT) with 78.8 NFE (vs. 84.8 for SFT).
Latency: Relay reduced inference latency by up to 32% compared to standard SFT while improving accuracy.
Memory Overhead: Profiling on an A100 GPU showed that while Relay increases live activation memory during the second forward pass, the peak memory usage remains nearly identical to vanilla SFT (approx. 20.1 GiB vs. 21.2 GiB) because the bottleneck is the vocabulary projection backward pass, not the relay state storage.

5. Significance and Claims

The paper claims that Relay fundamentally changes the nature of iterative non-autoregressive generation by allowing MDMs to accumulate intermediate continuous computation.

Forward-Thinking: The method enables models to explicitly optimize intermediate representations for future unmasking decisions, rather than treating each step as a fresh prediction problem.
Pareto Frontier Advancement: By reducing the number of forward passes required to reach a target accuracy, Relay advances the performance-latency frontier, making discrete diffusion models more efficient for complex reasoning and generation tasks.
Compatibility: The approach is presented as a drop-in enhancement that works with existing state-of-the-art DLMs and acceleration techniques (e.g., block diffusion, KV caching) without requiring architectural overhauls.

The authors conclude that Relay represents a meaningful step toward non-greedy, forward-thinking generation, offering a continuous substrate for latent reasoning that can be further explored for interpretability and application in other modalities (e.g., images, molecules).