GradMem：通过测试时梯度下降学习将上下文写入记忆 GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent
用测试时梯度下降优化记忆token，实现高效上下文压缩与复用
前置知识
KV-cache（键值缓存）
Transformer在推理时，每一层会为输入序列中的每个token计算Key和Value向量并缓存下来，这样在生成后续token时无需重新计算整个序列的注意力。KV-cache的大小与序列长度成正比，对于长上下文场景（如数万token的文档），KV-cache会占用大量GPU显存，成为部署瓶颈。
GradMem的核心动机就是用固定大小的记忆状态替代随上下文长度增长的KV-cache，理解KV-cache的局限性是理解本文必要性的前提。
测试时训练（Test-Time Training, TTT）
一种在推理阶段对模型进行自适应更新的技术范式。与传统推理不同，TTT在处理每个新样本时会执行若干步梯度更新，使模型能利用当前输入的结构信息进行自我调整。TTT-Layers是这一方向的代表工作，它在每一层维护可训练参数，在处理token时通过自监督目标在线更新。
GradMem可以看作TTT的一种变体，但它不更新模型权重，而是将梯度优化应用于独立的记忆状态，理解TTT的范式有助于把握GradMem的定位。
元学习与MAML
元学习（Meta-Learning）旨在让模型「学会学习」。MAML（Model-Agnostic Meta-Learning）是经典算法，其核心思想是找到一组好的模型初始化参数，使得在新任务上只需少量梯度步骤就能快速适应。MAML通过二阶梯度（对内循环优化过程求导）来更新初始化参数。
GradMem使用MAML风格的二阶元学习来训练记忆初始化M0和模型参数θ，使得测试时只需少量梯度步骤就能将上下文写入记忆。理解MAML的内外循环机制是理解GradMem训练过程的关键。
上下文压缩与记忆机制
在长上下文应用场景中，将长文档或对话历史压缩为紧凑的固定大小状态（记忆），使得后续查询可以基于该记忆而非完整上下文进行推理。这包括RMT（循环记忆Transformer）等分段处理方法，以及各种注意力压缩技术。目标是在有限的信息容量下尽可能保留任务相关的信息。
GradMem的定位是提供一种新的上下文压缩写入机制——用梯度下降而非前向传播来将上下文编码进记忆，理解这个领域的问题空间有助于评估其贡献。
自监督重构损失
一种不需要人工标注的训练信号。在GradMem中，WRITE阶段的损失函数定义为：给定记忆M和前缀token，模型需要重构（预测）上下文中的每个token。这个损失不需要任何下游任务标签，完全由模型自身的语言建模能力提供监督信号。最小化重构损失迫使记忆编码上下文中的关键信息。
GradMem的WRITE目标是任务无关的重构损失，这使得方法具有通用性——同一个重构目标可以服务多种下游任务，理解这一点对把握方法的泛化潜力很重要。
研究动机
当前大语言模型在处理长上下文（如文档、代码库、多轮对话）时，主流方案是保留完整的KV-cache来存储所有过去的激活值。这种方法虽然能保证信息无损，但KV-cache的显存占用与上下文长度成正比——对于数万token的上下文，KV-cache可能占用数十GB显存，严重限制了实际部署。更重要的是，KV-cache不具备可移植性：它绑定了特定的模型会话，无法轻松存储、传输或在不同查询间复用。在需要对同一上下文反复提问的场景中（如对同一文档的多轮问答），每次都必须重新加载和处理完整上下文，造成大量冗余计算。现有的前向写入方法（如RMT通过前向传播压缩上下文）虽然可以生成紧凑的记忆状态，但缺乏对写入质量的逐样本反馈——写入操作是「一次性的」，无法验证上下文是否被充分编码，也无法迭代修正写入错误。当上下文较长或包含高熵信息时，这种前向写入方法的容量会急剧下降。
本文的目标是GradMem的目标是提出一种基于梯度优化的上下文记忆写入机制，在固定的、小规模的记忆状态下（如8个记忆token），通过测试时的梯度下降步骤将上下文信息高效编码进记忆，使得后续的多个查询可以仅基于该记忆状态（而非完整上下文）进行准确推理。具体而言，GradMem希望实现：（1）在关联KV检索任务中，用8个记忆向量存储数十个键值对并保持高准确率；（2）梯度写入步骤应比前向写入步骤更有效地利用固定大小的记忆容量；（3）增加测试时的梯度优化步数应能可靠地提升记忆容量，形成可调的计算-精度权衡；（4）该方法应能迁移到自然语言任务和预训练语言模型上。
与已有工作不同的是，GradMem抓住了现有方法的一个被忽视的关键点：写入操作的质量不仅取决于「写入多少次」，更取决于「用什么方式写入」。前向写入方法（如RMT、ARMT）将记忆更新视为一个纯粹的前向计算过程，缺乏对写入质量的显式优化目标。而GradMem将记忆写入重新定义为一个优化问题——通过在重构损失上执行梯度下降，每一步梯度更新都能提供「当前记忆还缺少什么信息」的精确信号。这种损失驱动的写入机制自然地优先编码上下文中的高熵、不可预测部分，实现了一种结构化的信息压缩。此外，GradMem将所有测试时适应集中在一个单一的模型级记忆状态（而非像TTT-Layers那样在每一层维护可训练状态），使得记忆预算显式可控，且方法实现简洁。
核心方法
GradMem的核心思想可以用一个直觉类比来理解：标准的SGD训练可以看作通过梯度更新将数据「写入」模型参数；类比地，GradMem将当前上下文通过梯度更新「写入」一组独立的记忆token。与一次性前向编码（类似于「拍照」）不同，GradMem的梯度写入更像「反复描摹」——每一步梯度下降都会聚焦于模型当前预测不好的部分，逐步完善记忆。技术路线上，GradMem将每个任务样本分解为三部分：上下文C（要存储的信息）、查询Q（任务指令）、目标Y（要预测的输出）。系统分为WRITE和READ两个阶段：WRITE阶段执行K步梯度下降，优化记忆token以最小化上下文重构损失；READ阶段仅使用优化后的记忆和查询来预测目标，不接触原始上下文。记忆的初始化M0和模型参数θ通过二阶元学习（MAML风格）在训练集上学习，使得少量梯度步骤就能产生有用的记忆。
GradMem与已有方法最本质的区别在于：它将记忆写入定义为一个显式的、可迭代优化的问题，而非一个固定的前向计算。具体来说，现有方法（如RMT）通过前向传播将上下文映射到记忆状态——这个映射在推理时是固定的，无法根据当前上下文的难度自适应调整。GradMem则不同：它将记忆token视为可优化的「参数」，在推理时针对每个具体上下文执行梯度下降，使用模型级的自监督重构损失作为优化目标。这意味着：（1）每一步梯度更新都能获得「当前记忆还缺少什么信息」的精确反馈信号；（2）重构损失自然地集中在模型当前预测不好的token上，优先编码高熵、不可预测的信息；（3）通过增加梯度步骤数K，可以用额外的测试时计算换取更好的记忆质量——这是前向写入方法无法提供的灵活权衡。实验表明，即使是一步梯度更新也显著优于一步前向写入，而5步梯度更新可以将8个记忆向量的KV检索准确率从58.6%提升到99.1%（64个键值对），这是前向写入方法无论如何增加重复次数都无法达到的。
方法步骤详情
GradMem的完整流程分为训练阶段和推理阶段。训练阶段：（1）外循环——从训练集中采样一个batch，每个样本包含(C, Q, Y)；（2）内循环（WRITE）——对每个样本，从元学习到的初始化M0出发，执行K步梯度下降：第k步计算重构损失，即在记忆前缀条件下对上下文token的自回归交叉熵损失，然后更新记忆参数，其中alpha是内循环学习率；（3）READ——将优化后的记忆与查询Q拼接，计算下游任务损失；（4）外循环更新——通过WRITE优化过程反向传播任务损失（二阶梯度），更新模型参数θ和记忆初始化M0。推理阶段：（1）WRITE——对新上下文，从M0出发执行K步梯度下降（仅更新记忆，模型权重冻结）；（2）READ——使用优化后的记忆和查询进行预测，无需访问原始上下文。关键实现细节：WRITE和READ使用独立的预测头；记忆维度与模型输入嵌入维度相同；支持梯度裁剪等稳定化技术。在具体实现中，WRITE阶段的重构损失定义为自回归交叉熵：在当前记忆M和已看到的上下文前缀条件下，预测下一个token的概率取对数并求负和。每步梯度下降仅更新记忆M，模型参数保持冻结。READ阶段则将最终记忆与查询拼接作为模型输入，输出预测结果。
技术新颖性
GradMem的技术新颖性体现在三个层面。首先，在记忆写入机制上，它首次将模型级的显式重构损失应用于记忆token的测试时优化，区别于TTT-Layers在每一层使用激活重构、RMT使用前向传播传递记忆等已有方式。这种模型级的梯度写入提供了一个统一的、可控的写入接口。其次，在记忆架构上，GradMem将所有测试时适应集中在一个输入级的前缀记忆状态（m个d维向量），而非像TTT-Layers/Titans/Atlas那样在每一层维护可训练状态。这种设计使得记忆预算显式可控——只需指定m和d即可精确知道记忆大小，且方法实现无需修改模型内部架构。第三，在训练范式上，GradMem使用完整的二阶MAML元学习来训练记忆初始化和模型参数，使得内循环（WRITE）的少量梯度步骤能够可靠地产生有用记忆。这与一阶近似方法（如Reptile）不同——消融实验表明，去除二阶梯度信号会导致性能严重下降（8个KV对从99.7%降至12.9%）。
实验结果
GradMem在多个实验设置中展现了系统性的优势。在关联KV检索任务上，使用8个记忆向量的GradMem（K=1）在16个KV对时达到96.3%准确率，而相同记忆大小的前向写入方法RMT仅为45.5%——一步梯度更新的效果就远超一步前向写入。增加梯度步数进一步放大优势：GradMem（K=5）在32个KV对上达到99.9%（RMT为44.3%），在64个KV对上达到99.1%（RMT为19.3%），在96个KV对上达到88.4%（RMT为12.9%）。前向写入方法即使重复5次上下文处理，性能提升也很有限且不稳定（32个KV对从44.3%波动到37.0%），而梯度写入的额外步骤稳定提升。在与更大记忆状态的基线比较中，GradMem（8个记忆向量，K=5）在32个KV对上超过Mamba-130M（90.2%）和ARMT（54.9%），尽管后者的记忆状态分别大40倍和190倍。在NLP任务上，GradMem（GPT-2）在bAbI的QA1/QA4/QA5上匹配完整上下文Transformer的100%，在Short SQuAD上达到54.9%（K=5），超过RMT的42.6%。在语言建模上，GradMem的交叉熵为2.91-2.92，与RMT的2.91相当，且显著优于将上下文限制为128 token的GPT-2（3.20）。在更大模型上，Llama-3.2-1B使用1-2个记忆向量和1-2步WRITE即可在32个token内达到95-99%的重构准确率，实现了约16-32倍压缩。推理效率分析表明，当同一上下文被反复查询时（约64次以上），GradMem的总延迟开始低于完整上下文推理。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 关联KV检索（16对） | 精确匹配准确率 | 96.3%（K=1），99.6%（K=2），100.0%（K=5） | RMT前向写入45.5%（K=1） | K=1提升50.8个百分点，K=5提升54.5个百分点 |
| 关联KV检索（32对） | 精确匹配准确率 | 86.9%（K=1），98.3%（K=2），99.9%（K=5） | RMT前向写入44.3%（K=1），Mamba 90.2% | 比RMT提升42.6-55.6个百分点，K=5超过Mamba |
| 关联KV检索（64对） | 精确匹配准确率 | 58.6%（K=1），72.8%（K=2），99.1%（K=5） | RMT前向写入19.3%（K=1），ARMT 22.6% | 比RMT提升39.3-79.8个百分点 |
| 关联KV检索（96对） | 精确匹配准确率 | 32.6%（K=1），34.2%（K=2），88.4%（K=5） | RMT前向写入12.9%（K=1），ARMT 15.2% | 比RMT提升19.7-75.5个百分点 |
| bAbI QA3（三事实推理） | 精确匹配准确率 | 80.0%（K=1），79.3%（K=5） | RMT 87.9%，ARMT 92.3%，Mamba 96.7% | 低于RMT/ARMT/Mamba，但接近 |
| Short SQuAD | 精确匹配准确率 | 38.1%（K=1），54.9%（K=5） | RMT 42.6%，ARMT 39.0%，GPT-2全上下文64.2% | K=5超过RMT 12.3个百分点 |
| 语言建模（WikiText-103） | 交叉熵（越低越好） | 2.91-2.92 | RMT 2.91，GPT-2全上下文2.72 | 与RMT持平，优于限128上下文的GPT-2（3.20） |
| 文本重构（Llama-3.2-1B，K=2，1个记忆向量） | 最大序列长度@95%准确率 | 32 tokens | K=1时为8 tokens | K增加1步，容量提升4倍 |
局限与改进
GradMem存在几个明显的局限性。首先，训练成本显著高于前向写入方法：由于需要通过K步WRITE优化进行反向传播（二阶梯度），计算和显存开销都大幅增加。论文中提到，对于L=1024 token的序列，朴素实现的反向传播延迟约1000ms、显存占用约60GB，经过自定义优化后仍需约600ms和30GB。这限制了方法在大规模场景中的训练可行性。其次，推理时的WRITE阶段也有不可忽视的成本：执行K步梯度下降比一次前向传播慢得多。虽然论文展示了在多次查询场景下GradMem可以摊销WRITE成本，但对于单次查询场景，其效率不如直接处理完整上下文。第三，在bAbI的QA3任务上，GradMem的性能（80.0%）低于RMT（87.9%）和ARMT（92.3%），说明在需要多事实组合推理的场景中，梯度写入的优势并不总是成立。第四，方法目前仅在相对较小的模型（100M-3B参数）和较短的上下文（最多约256 token）上验证，对于真正长上下文（数千到数万token）的场景，8个记忆token的容量是否足够尚不清楚。第五，作者承认当前的重构损失可能不是最优的WRITE目标——它平等地对待所有token，但实际任务可能需要选择性地保留特定信息。
独立分析的弱点
从独立分析的角度看，GradMem有几个值得关注的弱点。第一，内循环学习率α的敏感性问题：虽然论文声称α在0.01到10范围内性能不强依赖，但不同任务仍需单独调整（KV检索需搜索，NLP任务默认0.4），且在大规模模型上的最优值可能不同，这增加了实际使用的调参负担。改进方向是学习一个自适应的学习率调度，或在元学习中同时学习最优学习率。第二，记忆token数量m的选择缺乏理论指导：论文中bAbI用8个token、SQuAD用32个token，但没有给出如何根据上下文复杂度自动确定m的建议。可以探索动态记忆分配——根据上下文的信息量自适应调整m。第三，WRITE阶段的梯度更新是「全局」的——每一步更新所有m个记忆token，缺乏对不同token的差异化处理。引入注意力机制或门控机制来选择性更新记忆token的子集，可能提高效率和容量。第四，当前的重构目标将所有token视为等价，但在实际任务中（如问答），某些token（如答案相关的关键信息）比其他token更重要。设计任务感知或信息论导向的WRITE目标（如最大化互信息）可能比简单的重构损失更有效。第五，方法在训练时需要完整展开WRITE内循环的计算图，这意味着K不能太大——论文中K最大为5，进一步增加K会导致训练不稳定。探索更高效的元学习算法（如iMAML、Reptile）是必要的。
未来方向
论文和基于成果可以延伸出多个研究方向。作者提出的方向包括：（1）探索更高效的元学习训练方法，如一阶/隐式梯度方法（iMAML、Reptile），以减少二阶梯度的计算开销；（2）设计超越重构的WRITE目标，学习自监督的编码变换以更好地保留任务相关信息；（3）将GradMem扩展到更长上下文和更大模型。基于论文成果可延伸的方向包括：（1）多级记忆架构——为不同类型的信息（事实、推理链、元数据）维护独立的记忆状态；（2）增量记忆更新——在上下文动态变化时（如流式输入），增量地更新记忆而非从头重写；（3）记忆共享与迁移——将一个上下文的记忆用于初始化相似上下文的记忆；（4）与检索增强生成（RAG）结合——用GradMem压缩文档为记忆token，作为RAG的紧凑索引；（5）多模态扩展——将图像、音频等非文本信息编码进统一的记忆空间。
复现评估
GradMem的代码已在GitHub开源（https://github.com/yurakuratov/gradmem），这为复现提供了基础。数据方面，论文使用的基准测试（关联KV检索、bAbI、SQuAD、WikiText-103）都是公开的标准数据集，无需额外获取。算力需求是复现的主要障碍：二阶元学习的训练需要较大显存，论文中提到使用A100 GPU进行所有实验，对于1024 token的序列，优化后的实现仍需约30GB显存。对于小规模的关联KV检索实验（4层、128维隐藏层），使用较小的GPU也可以复现。NLP任务的实验基于GPT-2（124M）和Pythia（160M），这些模型在消费级GPU上可以微调。论文的超参数在附录中有详细记录（内循环学习率、记忆token数量、训练初始化方式等），且提供了课程学习策略的描述，这些都有助于复现。总体而言，复现难度中等——小规模验证实验可行，但大规模训练实验需要专业GPU资源。
论文图表
