arXiv:2603.01331cs.CLcs.AI2026-03被引 1

给离散扩散语言模型加了个持续记忆,让推理更连贯。

MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models

  • 用可循环的记忆模块,让每步推理结果能保留并更新。
  • 在数学和代码任务上平均提升4.5个百分点,参数增加仅0.6%。
  • 适合需要长链推理的场景,如数学证明或复杂代码生成。

离散扩散语言模型(dLLMs)通过逐步去噪掩码序列生成文本。然而,标准dLLMs仅依赖当前硬掩码序列进行每一步去噪,而中间的连续表示在采样和重掩码后被丢弃。我们称此瓶颈为‘信息孤岛’问题:连续信息被隔离在单个去噪步骤中,无法跨轨迹传播。该问题在需要长期推理状态维持的场景中尤为严重。为此,我们提出轻量级递归增强模块MetaState,为冻结的dLLM主干赋予持久、固定大小的工作记忆。MetaState包含三个共享时间条件器的模块:跨注意力的Mixer将主干激活读入记忆槽,GRU风格的Updater在步骤间整合信息,以及跨注意力的Injector将更新后的记忆写回主干。我们通过专用的K步展开训练管道学习多步动态。MetaState仅引入约0.6%的可训练参数,保持主干冻结,且在数学推理与代码生成基准上持续提升性能,所有评估中平均提升4.5个百分点。

原文摘要 · Abstract (English)

Discrete diffusion language models (dLLMs) generate text by iteratively denoising a masked sequence. However, standard dLLMs condition each denoising step solely on the current hard-masked sequence, while intermediate continuous representations are discarded after sampling and remasking. We term this bottleneck the \textbf{Information Island} issue: continuous information remains isolated within individual denoising steps and fails to propagate across the trajectory. This bottleneck is especially harmful for reasoning, which requires intermediate reasoning state to be preserved and updated across many denoising steps. To address this limitation, we introduce \textbf{MetaState}, a lightweight recurrent augmentation that equips a frozen dLLM backbone with persistent, fixed-size working memory. MetaState comprises three modules with a shared time conditioner: a cross-attention \textbf{Mixer} that reads backbone activations into memory slots, a GRU-style \textbf{Updater} that integrates information across steps, and a cross-attention \textbf{Injector} that writes the updated memory back into the backbone. We train these modules with a dedicated $K$-step unrolling pipeline to learn multi-step dynamics. MetaState adds only ${\sim}0.6\%$ trainable parameters while keeping the backbone frozen, and consistently improves reasoning performance over frozen baselines on mathematical reasoning and code generation benchmarks, with an average gain of 4.5 percentage points across all evaluations.

扩散模型推理增强工作记忆

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。