英伟达开源128K记忆压缩方案，免额外缓存提速2.7倍

闻乐 发自 凹非寺
量子位 | 公众号 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项目的早期构想就是由他提出的。