Learning to (Learn at Test Time): RNNs with Expressive Hidden States

自注意力机制在处理长文本时表现良好，但其复杂度为二次方。现有的循环神经网络层复杂度为线性，但其在长文本上的表现受到隐藏状态表达能力的限制。我们提出了一种新的序列建模层，具有线性复杂度和表达能力强的隐藏状态。关键思想是将隐藏状态本身作为一个机器学习模型，并将更新规则作为自监督学习的一步。由于隐藏状态通过对测试序列进行训练而更新，因此我们的层被称为测试时间训练（TTT）层。我们考虑了两种实例化：TTT-Linear和TTT-MLP，其隐藏状态分别为线性模型和两层MLP。我们在125M到1.3B参数的范围内进行了评估，与强大的Transformer和现代RNN Mamba进行了比较。TTT-Linear和TTT-MLP都与基准线相匹配或超过。类似于Transformer，它们可以通过对更多标记的调节来不断降低困惑度，而Mamba在16k上下文之后无法做到。通过初步的系统优化，TTT-Linear已经在8k上下文时比Transformer更快，并与Mamba在墙钟时间上匹配。TTT-MLP仍面临着内存I/O方面的挑战，但在长上下文方面显示出更大的潜力，为未来的研究指明了一个有前途的方向。

提出一种新的序列建模层，解决了传统RNN层表达能力不足和Self-Attention层计算复杂度高的问题。

将隐藏状态作为一个机器学习模型，并通过自监督学习的方式进行更新，从而实现线性复杂度和表达能力强的序列建模。

所提出的TTT-Linear和TTT-MLP序列建模层在125M到1.3B参数范围内的实验中，表现优于或与Transformer和Mamba等强基线相当。TTT-Linear已经在8k上下文的情况下比Transformer更快，并且与Mamba的墙钟时间相当。TTT-MLP在内存I/O方面仍面临挑战，但在长上下文方面具有更大的潜力。

最近的相关研究包括《Attention is All You Need》和《Mamba: Integrating Batch and Momenta Adaptation for Online Optimization of Sequence Objectives》等。