MetaState: 持久工作记忆增强离散扩散语言模型的推理能力
MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models
AI总结：
MetaState通过引入轻量级循环增强模块，为冻结的离散扩散语言模型提供持久固定大小的工作记忆，提升推理性能，平均提升4.5%。
摘要：
离散扩散语言模型（dLLMs）通过迭代去噪掩码序列生成文本。然而，标准dLLMs仅将每个去噪步骤的当前硬掩码序列作为条件，中间连续表示在采样和重新掩码后被丢弃。我们称此瓶颈为信息孤岛问题：连续信息在个体去噪步骤中孤立，无法跨轨迹传播。为解决此限制，我们引入MetaState，一个轻量级循环增强模块，为冻结的dLLM主干提供持久、固定大小的工作记忆。MetaState包含三个模块，共享时间条件器：一个交叉注意力Mixer读取主干激活到内存槽，一个GRU风格的Updater整合跨步骤信息，一个交叉注意力Injector将更新的内存写回主干。我们使用专用K-step展开管道训练这些模块以学习多步动态。MetaState仅增加约0.6%的可训练参数，同时保持主干冻结，并在数学推理和代码生成基准上持续优于冻结基线，平均提升4.5%。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
