用于长上下文语言建模的门控可微工作记忆
计算与语言
2026-01-21 v1
摘要
长上下文对 Transformer 提出了挑战：注意力分数在数千个 token 中被稀释，关键信息通常在中间丢失，且模型在推理时难以适应新模式。最近关于测试时适应的工作通过维护一种工作记忆（即在当前上下文上更新的瞬态参数）来解决这一问题，但现有方法依赖统一的写入策略，在低效用区域浪费计算资源，并在语义异构的上下文中产生高梯度方差。在本工作中，我们将测试时适应重新构建为预算受限的记忆巩固问题，侧重于在有限计算量下应将上下文的哪些部分巩固到工作记忆中。我们提出了 Gdwm（门控可微工作记忆，Gated Differentiable Working Memory），一种引入写入控制器以门控巩固过程的框架。该控制器估计上下文效用（Contextual Utility，一种对长程上下文依赖的信息论度量），并据此分配梯度步数，同时保持全局覆盖。在 ZeroSCROLLS 和 LongBench v2 上的实验表明，Gdwm 以比统一基线少 4 的梯度步数实现了相当或更优的性能，为测试时适应确立了新的效率-性能 Pareto 前沿。
引用
@article{arxiv.2601.12906,
title = {Gated Differentiable Working Memory for Long-Context Language Modeling},
author = {Lingrui Mei and Shenghua Liu and Yiwei Wang and Yuyao Ge and Baolong Bi and Jiayu Yao and Jun Wan and Ziling Yin and Jiafeng Guo and Xueqi Cheng},
journal= {arXiv preprint arXiv:2601.12906},
year = {2026}
}