闻乐 发自 凹非寺
量子位 | 公众号 QbitAI
提高大模型记忆这块儿,美国大模型开源王者——英伟达也出招了。
联合Astera研究所、斯坦福大学、UC伯克利、加州大学圣地亚哥分校等机构推出了TTT-E2E方法。
在128K超长文本上处理速度比全注意力模型快2.7倍,处理2M上下文时提速达35倍,性能还不打折。
这项技术与前几天大火的DeepSeek条件记忆模块有所不同。
DeepSeek的Engram模块依赖的是“按需查表”的静态学习路径,而英伟达走的是动态学习的路子,关键在于上下文压缩
通过实时学习将关键内容压缩到自身权重中,让模型在测试阶段依然保持学习状态。
这样既避免了额外缓存的负担,又能精准捕捉长文本中的核心逻辑。
给模型装上记忆压缩包
TTT-E2E并没有依赖复杂特殊架构,反而是基于带滑动窗口注意力的标准Transformer,容易部署。
这个方法的核心思路是将长文本建模从架构设计问题转化为「持续学习」任务
在测试阶段,模型会基于当前读取的上下文进行下一个词预测。
每读取一段文本,就通过梯度下降更新自身参数,通过这种方式持续训练自身,把读到的文本信息动态压缩到权重中,这样就不用额外存储冗余数据。
在训练阶段,团队通过元学习为模型做初始化准备,让模型天生适应「测试时学习」的模式。
把每个训练序列都模拟成测试序列,先在内循环中对其进行测试时训练,再在外循环中优化模型的初始参数,确保初始状态就能快速适配测试时的学习需求,实现了训练与测试的端到端对齐优化。
为了平衡效率与稳定性,TTT-E2E还设计了三项关键优化。
一是采用「迷你批处理+滑动窗口」的组合策略。将测试时的训练数据分成多个迷你批,配合8K大小的滑动窗口注意力,既解决了单token梯度更新易爆炸的问题,又保证模型能记住批内上下文,提升计算并行度;
二是精准更新策略。只更新模型的MLP层(冻结嵌入层、归一化层和注意力层),并且只更新最后1/4的网络块,在减少计算成本的同时避免参数更新混乱;
三是双MLP设计。在需更新的网络块中加入一个静态MLP层,专门存储预训练知识,另一个动态MLP层负责吸收新上下文,来防止模型学新忘旧。
从实验数据来看,TTT-E2E的表现很亮眼。
在3B参数模型的测试中,TTT-E2E在128K上下文长度下的测试损失与全注意力Transformer持平甚至更优,而Mamba 2、Gated DeltaNet等同类模型在长文本场景下性能均出现明显下滑;
在延迟上,它的推理延迟不随上下文长度增加而变化,与RNN类似,在H100显卡上处理128K文本时,速度比全注意力模型快2.7倍。
在解码长序列任务中,经Qwen-8B模型评估,TTT-E2E生成的文本质量稳定,损失值持续低于传统模型。
通过实验结果也可以看出,该方法的推理延迟与上下文长度无关,始终保持恒定,这也意味着无论处理8K还是128K文本,用户都能获得一致的快速响应体验。
不过,TTT-E2E也存在一些小局限。
在大海捞针这类需要精准回忆细节的任务中,它的表现远不如全注意力模型。
这是因为它的核心是压缩记忆,会过滤掉看似无关的细节,而全注意力模型能近乎无损地召回所有信息。
另一方面,训练阶段的元学习需要计算梯度的梯度,目前实现比标准预训练要慢。
目前,TTT-E2E的代码和相关论文已完全开源。
这项研究的项目总负责人是斯坦福的博士后研究员Yu Sun,他同时是该研究的核心贡献者。
他研究的总体目标是让人工智能系统能够像人类一样持续学习。自2019年以来,他就在开发“测试时训练”的概念框架,TTT-E2E项目的早期构想就是他提出的。
论文地址:https://arxiv.org/abs/2512.23675
代码地址:https://github.com/test-time-training/e2e
参考链接:https://x.com/karansdalal/status/2010774529120092481
热门跟贴