Learning to (Learn at Test Time): RNNs with Expressive Hidden States：测试时训练层突破RNN长上下文瓶颈
2024年7月，斯坦福大学等机构提出测试时训练（TTT）层，将隐藏状态本身视为一个机器学习模型，并在测试序列上通过自监督学习更新。TTT-Linear和TTT-MLP两种实例化在125M至1.3B参数规模下，与Transformer和Mamba对比，TTT层在超过16k上下文后仍能持续降低困惑度，而Mamba无法做到。TTT-MLP在长上下文上展现更大潜力，但面临内存I/O挑战。
发展脉络
- 首次出现Learning to (Learn at Test Time): RNNs with Expressive Hidden States：测试时训练层突破RNN长上下文瓶颈arXiv cs.AI
- 当前判断TTT层为构建高效长上下文模型提供了新范式，可能影响下一代语言模型和序列处理系统的架构设计。对于需要处理长文档、视频、基因组序列等场景的AI产品，TTT层有望在降低计算成本的同时提升性能。Agent Pulse · 分析
该工作提出了一种新的序列建模层TTT，其核心创新在于让隐藏状态成为一个可学习的模型，并在测试时通过自监督学习更新，从而在保持线性复杂度的同时获得类似Transformer的长期依赖建模能力。这改变了RNN在长上下文任务中表现不佳的局面，为构建高效且强大的序列模型提供了新方向。
TTT层将隐藏状态定义为一个小型机器学习模型（如线性模型或两层MLP），更新规则为一步自监督学习（如重构输入token）。这使得隐藏状态能够根据测试序列动态调整，表达能力远超传统RNN的固定状态。实验在125M至1.3B参数规模下进行，与Transformer和Mamba对比，TTT层在16k上下文后仍能持续降低困惑度，而Mamba饱和。TTT-MLP在长上下文上表现更优，但训练和推理速度受限于内存I/O。该方法在保持线性复杂度的同时，实现了接近Transformer的长上下文性能。
TTT层为构建高效长上下文模型提供了新范式，可能影响下一代语言模型和序列处理系统的架构设计。对于需要处理长文档、视频、基因组序列等场景的AI产品，TTT层有望在降低计算成本的同时提升性能。
建议AI基础设施团队评估TTT层在长文档理解、代码生成、多轮对话等场景的替换潜力。可考虑在现有RNN或Transformer模型中集成TTT层进行小规模实验，验证其在特定业务数据上的收益。
关注TTT-MLP的内存I/O优化进展，以及更大规模（如7B+）的预训练实验。若能在效率上取得突破，TTT层可能成为Transformer和Mamba的有力竞争者，推动长上下文AI应用的普及。