MemDLM: Memory-Enhanced DLM Training

• URL: http://arxiv.org/abs/2603.22241v1
• Date: Mon, 23 Mar 2026 17:39:56 GMT
• Title: MemDLM: Memory-Enhanced DLM Training
• Abstract summary: 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.
• Score: 56.40248490616793
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, they suffer from a notable train-inference mismatch: DLMs are trained with a static, single-step masked prediction objective, but deployed through a multi-step progressive denoising trajectory. We propose MemDLM (Memory-Enhanced DLM), which narrows this gap by embedding a simulated denoising process into training via Bi-level Optimization. An inner loop updates a set of fast weights, forming a Parametric Memory that captures the local trajectory experience of each sample, while an outer loop updates the base model conditioned on this memory. By offloading memorization pressure from token representations to parameters, MemDLM yields faster convergence and lower training loss. Moreover, the inner loop can be re-enabled at inference time as an adaptation step, yielding additional gains on long-context understanding. We find that, when activated at inference time, this Parametric Memory acts as an emergent in-weight retrieval mechanism, helping MemDLM further reduce token-level attention bottlenecks on challenging Needle-in-a-Haystack retrieval tasks. Code: https://github.com/JarvisPei/MemDLM.

Related papers

• Dynamic-dLLM: Dynamic Cache-Budget and Adaptive Parallel Decoding for Training-Free Acceleration of Diffusion LLM [35.77320928871856]
  We propose Dynamic-dLLM, a training-free framework that enhances dLLM inference efficiency through two components.<n>Experiments on models like LLaDA-8B-Instruct, LLaDA-1.5, and Dream-v0-7B-Instruct demonstrate that Dynamic-dLLM significantly improves inference speed.
  arXiv  Detail & Related papers  (2026-05-27T02:47:50Z)
• Learned Relay Representations for Forward-Thinking Discrete Diffusion Models [34.17541648016911]
  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.
  arXiv  Detail & Related papers  (2026-05-21T18:53:22Z)
• Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models [8.430043498727617]
  Memory-Efficient Looped Transformer (MELT) is a novel architecture that decouples reasoning depth from memory consumption.<n>We show that MELT achieves constant-memory iterative reasoning without sacrificing LoopLM performance.
  arXiv  Detail & Related papers  (2026-05-08T13:25:27Z)
• MemSifter: Offloading LLM Memory Retrieval via Outcome-Driven Proxy Reasoning [78.46301394559903]
  Large Language Models (LLMs) are increasingly used for long-duration tasks.<n>Current methods face a trade-off between cost and accuracy.<n>MemSifter is a novel framework that offloads the memory retrieval process to a small-scale proxy model.
  arXiv  Detail & Related papers  (2026-03-03T02:57:38Z)
• 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)
• NextMem: Towards Latent Factual Memory for LLM-based Agents [58.35585202907478]
  NextMem is a latent factual memory framework that utilizes an autoregressive autoencoder to efficiently construct latent memory.<n>Extensive experiments demonstrate that NextMem achieves superior performance.
  arXiv  Detail & Related papers  (2026-02-26T14:35:27Z)
• Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference [16.71963410333802]
  Large language models (LLMs) have demonstrated remarkable success across diverse artificial intelligence tasks.<n>MoC substantially reduces activation memory during pre-training.<n>MoC delivers significant memory savings and throughput gains while maintaining competitive model performance.
  arXiv  Detail & Related papers  (2025-11-12T13:30:57Z)
• Saber: An Efficient Sampling with Adaptive Acceleration and Backtracking Enhanced Remasking for Diffusion Language Model [98.35868970993232]
  Diffusion language models (DLMs) are emerging as a powerful and promising alternative to the dominant autoregressive paradigm.<n>We introduce efficient Sampling with Adaptive acceleration and Backtracking Enhanced Remasking (i.e., Saber) to achieve better inference speed and output quality in code generation.
  arXiv  Detail & Related papers  (2025-10-20T23:38:12Z)
• DAF: An Efficient End-to-End Dynamic Activation Framework for on-Device DNN Training [41.09085549544767]
  We introduce a Dynamic Activation Framework (DAF) that enables scalable and efficient on-device training through system-level optimizations.<n>DAF achieves both memory- and time-efficient dynamic quantization training by addressing key system bottlenecks.<n> Evaluations on various deep learning models across embedded and mobile platforms demonstrate up to a $22.9times$ reduction in memory usage and a $3.2times$ speedup.
  arXiv  Detail & Related papers  (2025-07-09T08:59:30Z)
• Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training [78.93900796545523]
  Mini-Sequence Transformer (MsT) is a methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. integrated with the huggingface library, MsT successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.
  arXiv  Detail & Related papers  (2024-07-22T01:52:30Z)
• MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training [24.066283519769968]
  Large Language Models (LLMs) have been trained using extended context lengths to foster more creative applications.<n>We propose MEMO, a novel framework for fine-grained activation memory management.<n>MeMO achieves an average of 1.97x and 1.80x MFU compared to Megatron-LM and DeepSpeed.
  arXiv  Detail & Related papers  (2024-07-16T18:59:49Z)
• SHERL: Synthesizing High Accuracy and Efficient Memory for Resource-Limited Transfer Learning [63.93193829913252]
  We propose an innovative METL strategy called SHERL for resource-limited scenarios. In the early route, intermediate outputs are consolidated via an anti-redundancy operation. In the late route, utilizing minimal late pre-trained layers could alleviate the peak demand on memory overhead.
  arXiv  Detail & Related papers  (2024-07-10T10:22:35Z)
• Dynamic Sparse No Training: Training-Free Fine-tuning for Sparse LLMs [67.38165028487242]
  We introduce Dynamic Sparse No Training (DSnoT), a training-free fine-tuning approach to fine-tune large language models (LLMs) Inspired by the Dynamic Sparse Training, DSnoT minimizes the reconstruction error between the dense and sparse LLMs. Our paper offers fresh insights into how to fine-tune sparse LLMs in an efficient training-free manner and open new venues to scale the great potential of sparsity to LLMs.
  arXiv  Detail & Related papers  (2023-10-13T07:38: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.