arXiv:2603.13875cs.CLcs.LG2026-03被引 3

用测试时优化把长文本压缩进记忆,高效存取。

GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent

  • 通过测试时梯度下降优化记忆令牌,动态写入上下文。
  • 在键值检索任务中,相同内存下优于传统前向写法。
  • 适合需要长上下文记忆的模型部署场景。

许多大语言模型应用需处理长上下文。传统Transformer依赖每层存储大量键值缓存(KV-cache),导致内存开销大。理想方案是压缩记忆:仅读取一次上下文,将其编码为紧凑状态,后续多次查询均基于该状态。本文研究上下文移除场景,即推理时无法访问原始上下文。提出GradMem,通过每样本测试时优化将上下文写入记忆。给定上下文后,GradMem在保持模型权重冻结的前提下,对少量前缀记忆令牌执行若干步梯度下降,并显式优化模型级自监督上下文重建损失,实现带有迭代纠错的损失驱动写入,区别于单向前向方法。在关联键值检索任务中,GradMem在相同内存容量下超越前向写法;额外梯度步骤比重复前向写更有效提升容量。进一步实验表明,该方法可迁移至自然语言任务:使用预训练语言模型,在bAbI和SQuAD变体上取得有竞争力的表现,仅依赖记忆中的编码信息。

原文摘要 · Abstract (English)

Many large language model applications require conditioning on long contexts. Transformers typically support this by storing a large per-layer KV-cache of past activations, which incurs substantial memory overhead. A desirable alternative is compressive memory: read a context once, store it in a compact state, and answer many queries from that state. We study this in a context removal setting, where the model must generate an answer without access to the original context at inference time. We introduce GradMem, which writes context into memory via per-sample test-time optimization. Given a context, GradMem performs a few steps of gradient descent on a small set of prefix memory tokens while keeping model weights frozen. GradMem explicitly optimizes a model-level self-supervised context reconstruction loss, resulting in a loss-driven write operation with iterative error correction, unlike forward-only methods. On associative key--value retrieval, GradMem outperforms forward-only memory writers with the same memory size, and additional gradient steps scale capacity much more effectively than repeated forward writes. We further show that GradMem transfers beyond synthetic benchmarks: with pretrained language models, it attains competitive results on natural language tasks including bAbI and SQuAD variants, relying only on information encoded in memory.

记忆机制测试时优化上下文压缩

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