Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

• URL: http://arxiv.org/abs/2605.22967v1
• Date: Thu, 21 May 2026 18:53:22 GMT
• Title: Learned Relay Representations for Forward-Thinking Discrete Diffusion Models
• Abstract summary: We propose Learned Relay Representations (Relay) to avoid a hard reset between denoising rounds.<n>Relay introduces a differentiable per-token channel that passes information between forward passes and is trained via truncated backpropagation through time.<n>We show that Relay can be scaled to state-of-the-art Diffusion Language Models (DLMs) and is seamlessly compatible with techniques like block diffusion and KV caching.
• Score: 34.17541648016911
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Subliminal Clocks: Latent Time Modelling in Diffusion Language Models [19.412917551000742]
  Diffusion Language Models (DLMs) have emerged as a promising alternative to autoregressive models.<n>We show that DLMs do in fact encode a latent representation related to the diffusion timestep within their residual streams.<n>We further demonstrate that steering the model along a low-dimensional subspace associated with the inferred timestep allows us to systematically modulate its notion of denoising progress.
  arXiv  Detail & Related papers  (2026-07-02T06:45:42Z)
• MemDLM: Memory-Enhanced DLM Training [56.40248490616793]
  Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models.<n>They suffer from a notable train-inference mismatch.<n>We propose MemDLM (Memory-Enhanced DLM), which embeds a simulated denoising process into training.
  arXiv  Detail & Related papers  (2026-03-23T17:39:56Z)
• MetaState: Persistent Working Memory for Discrete Diffusion Language Models [13.775010785119761]
  We introduce textbfMetaState, a lightweight recurrent augmentation that equips a frozen dLLM backbone with a persistent, fixed-size working memory.<n>On LLaDA-8B and Dream-7B, textbfMetaState introduces negligible trainable parameters while keeping the backbone frozen, and it consistently improves accuracy over frozen baselines.
  arXiv  Detail & Related papers  (2026-03-02T00:16:35Z)
• Causal Autoregressive Diffusion Language Model [70.7353007255797]
  CARD reformulates the diffusion process within a strictly causal attention mask, enabling dense, per-token supervision in a single forward pass.<n>Our results demonstrate that CARD achieves ARM-level data efficiency while unlocking the latency benefits of parallel generation.
  arXiv  Detail & Related papers  (2026-01-29T17:38:29Z)
• Streaming-dLLM: Accelerating Diffusion LLMs via Suffix Pruning and Dynamic Decoding [36.74241893088594]
  Diffusion Large Language Models (dLLMs) offer a compelling paradigm for natural language generation.<n>Recent works have accelerated inference via KV cache reuse or decoding, but overlook the intrinsic inefficiencies within the block-wise diffusion process.<n>We propose Streaming-dLLM, a training-free framework that streamlines inference across both spatial and temporal dimensions.
  arXiv  Detail & Related papers  (2026-01-25T17:36:04Z)
• WeMMU: Enhanced Bridging of Vision-Language Models and Diffusion Models via Noisy Query Tokens [69.97021957331326]
  We propose Noisy Query Tokens, which learn a distributed representation space between the VLM and Diffusion Model via end-to-end optimization.<n>We also introduce a VAE branch with linear projection to recover fine-grained image details.
  arXiv  Detail & Related papers  (2025-12-02T09:02:20Z)
• Denoising Task Routing for Diffusion Models [19.373733104929325]
  Diffusion models generate highly realistic images by learning a multi-step denoising process. Despite the inherent connection between diffusion models and multi-task learning (MTL), there remains an unexplored area in designing neural architectures. We present Denoising Task Routing (DTR), a simple add-on strategy for existing diffusion model architectures to establish distinct information pathways.
  arXiv  Detail & Related papers  (2023-10-11T02:23:18Z)
• Can SAM Boost Video Super-Resolution? [78.29033914169025]
  We propose a simple yet effective module -- SAM-guidEd refinEment Module (SEEM) This light-weight plug-in module is specifically designed to leverage the attention mechanism for the generation of semantic-aware feature. We apply our SEEM to two representative methods, EDVR and BasicVSR, resulting in consistently improved performance with minimal implementation effort.
  arXiv  Detail & Related papers  (2023-05-11T02:02:53Z)
• Diffusion Recommender Model [85.9640416600725]
  We propose a novel Diffusion Recommender Model (named DiffRec) to learn the generative process in a denoising manner.<n>To retain personalized information in user interactions, DiffRec reduces the added noises and avoids corrupting users' interactions into pure noises like in image synthesis.
  arXiv  Detail & Related papers  (2023-04-11T04:31:00Z)
• DiffTAD: Temporal Action Detection with Proposal Denoising Diffusion [137.8749239614528]
  We propose a new formulation of temporal action detection (TAD) with denoising diffusion, DiffTAD. Taking as input random temporal proposals, it can yield action proposals accurately given an untrimmed long video.
  arXiv  Detail & Related papers  (2023-03-27T00:40:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.