学习在测试时学习：具有表征性隐藏状态的 RNN
机器学习
2025-09-03 v4 人工智能
计算与语言
摘要
自注意力在长序列上下文中表现良好，但具有二次复杂度。现有的RNN层具有线性复杂度，但其在长序列上下文中的性能受限于隐藏状态的表达能力。我们提出了用于实例化具有线性复杂度和表征性隐藏状态的序列建模层的实用框架。关键思想是将隐藏状态本身充当机器学习模型，而更新规则则是自监督学习的一步。由于隐藏状态是通过在测试序列上训练进行更新的，我们称这些层为测试时训练（TTT）层。我们考虑了两种实例化：TTT-Linear和TTT-MLP，其隐藏状态分别为线性模型和两层MLP。我们在1.25亿至13亿参数的规模上评估了这些实例化，与强大的Transformer和Mamba——一种现代RNN进行比较。类似于Transformer，TTT-Linear和TTT-MLP可以通过条件化更多标记继续降低困惑度，而Mamba在16k上下文后无法继续降低困惑度。TTT-MLP仍面临内存I/O方面的挑战，但在长序列上下文方面显示出更大的潜力，指向未来研究的有前景的方向。
引用
@article{arxiv.2407.04620,
title = {Learning to (Learn at Test Time): RNNs with Expressive Hidden States},
author = {Yu Sun and Xinhao Li and Karan Dalal and Jiarui Xu and Arjun Vikram and Genghan Zhang and Yann Dubois and Xinlei Chen and Xiaolong Wang and Sanmi Koyejo and Tatsunori Hashimoto and Carlos Guestrin},
journal= {arXiv preprint arXiv:2407.04620},
year = {2025}
}
备注
The current version contains updates on related work and limitations. All experiments were completed in the first version