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

https://papers.cool/arxiv/2407.04620

Authors: Yu Sun ; Xinhao Li ; Karan Dalal ; Jiarui Xu ; Arjun Vikram ; Genghan Zhang ; Yann Dubois ; Xinlei Chen ; Xiaolong Wang ; Sanmi Koyejo ; Tatsunori Hashimoto ; Carlos Guestrin

Summary: Self-attention performs well in long context but has quadratic complexity. Existing RNN layers have linear complexity, but their performance in long context is limited by the expressive power of their hidden state. We propose a new class of sequence modeling layers with linear complexity and an expressive hidden state. The key idea is to make the hidden state a machine learning model itself, and the update rule a step of self-supervised learning. Since the hidden state is updated by training even on test sequences, our layers are called Test-Time Training (TTT) layers. We consider two instantiations: TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B parameters, comparing with a strong Transformer and Mamba, a modern RNN. Both TTT-Linear and TTT-MLP match or exceed the baselines. Similar to Transformer, they can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. With preliminary systems optimization, TTT-Linear is already faster than Transformer at 8k context and matches Mamba in wall-clock time. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research.

Q: 这篇论文试图解决什么问题？

A: 这篇论文提出了一种新的序列建模层——Test-Time Training (TTT) 层，旨在解决现有循环神经网络（RNN）在处理长文本序列时所面临的性能限制问题。具体来说，论文指出现有的RNN层虽然具有线性复杂度的优势，但在长文本序列的表达能力上受限于其隐藏状态的固定大小。与之相对的，自注意力（Self-attention）机制虽然在长文本序列上表现更好，但其计算复杂度是二次方的，这限制了其在更大规模序列上的可行性。

为了解决这个问题，论文提出了TTT层，其核心思想是将隐藏状态设计为一个机器学习模型本身，并通过自监督学习的更新规则来更新这个状态。这样，即使在测试时序上，也可以通过训练来更新隐藏状态，从而提高模型对长文本序列的理解和表达能力。TTT层的设计使得它在保持线性复杂度的同时，能够更好地捕捉和利用长文本序列中的信息，提高了RNN在长文本序列建模中的性能。

Q: 有哪些相关研究？

A: 论文中提到了多个与TTT层相关的研究领域和具体工作，主要包括以下几个方面：

1. 现代RNNs（如Mamba、RWKV、xLSTM、GLA等）：这些模型都是Structured State-Space Models，它们尝试通过不同的方式改进RNN的隐藏状态表示，以提高对长序列的建模能力。

2. 自注意力机制（Self-attention）：自注意力是Transformer模型的核心，它通过Key-Value存储历史上下文信息，虽然表达能力强，但计算复杂度较高。

3. 测试时学习（Test-Time Training, TTT）：这是一个在测试时对模型进行训练以适应特定输入的概念，先前的研究已经展示了TTT在异常值检测和视频流处理中的有效性。

4. 快速权重（Fast Weights）：这是一种在特定数据上更新模型参数的概念，与TTT有关联，但在更新规则上通常不依赖于显式的优化问题。

5. 学习学习（Learning to Learn）：也称为元学习或双层优化，指的是模型能够学习如何通过优化过程来学习新任务。

6. 非参数学习器（Nonparametric Learners）：如最近邻、支持向量机（SVMs）和核岭回归等，这些方法不依赖于参数化模型，而是直接使用训练数据。

7. Transformer模型：作为当前NLP领域的主流模型，Transformer通过自注意力机制实现了优秀的性能，但其二次方的计算复杂度限制了其在更长序列上的应用。

8. Mamba模型：一种现代RNN，使用Structured State-Space Models来提高对长序列的建模能力。

9. RWKV模型：另一种现代RNN，具有矩阵值状态和动态递归。

10. xLSTM和GLA：这些模型通过改进LSTM和引入门控机制来提高性能。

11. 其他RNN变体：如LSTM、GRU等，它们在序列建模中有着广泛的应用。

论文通过将这些相关研究与提出的TTT层进行比较，展示了TTT层在长序列建模中的潜力和优势。

Q: 论文如何解决这个问题？

A: 论文通过提出一种新的序列建模层——Test-Time Training (TTT) 层来解决现有RNN在长文本序列处理中的性能限制问题。具体解决方案包括以下几个关键点：

1. 隐藏状态作为模型：TTT层的核心思想是将隐藏状态设计为一个机器学习模型本身，这意味着隐藏状态可以包含模型的参数，并且可以通过学习过程进行更新。

2. 自监督学习的更新规则：TTT层的更新规则是自监督学习的一步，这意味着在测试时，即使没有标签，也可以通过训练来更新隐藏状态。这个过程称为测试时训练（TTT）。

3. 线性复杂度：尽管TTT层在更新隐藏状态时进行了训练，但其整体复杂度仍然是线性的，这使得它在处理长序列时比自注意力机制更高效。

4. 实例化：论文提出了两种TTT层的实例化方法，分别是TTT-Linear和TTT-MLP，它们的隐藏状态分别是线性模型和两层MLP（多层感知机）。

5. 硬件效率：为了提高TTT层在实际硬件上的性能，论文提出了两种实用创新：mini-batch TTT和对内部操作的双重形式（dual form），以更好地利用现代GPU和TPU。

6. 端到端优化：TTT层可以集成到任何网络架构中，并与RNN层和自注意力一样进行端到端优化。

7. 自监督任务学习：TTT层的自监督任务是可学习的，这意味着可以通过外循环来优化任务本身，从而更好地适应语言建模的最终目标。

8. 实验验证：论文通过在不同规模的参数（从125M到1.3B）上的实验，验证了TTT层与强Transformer基线和现代RNN Mamba相比的性能。

通过这些方法，TTT层能够在保持线性复杂度的同时，提高对长文本序列的表达能力，从而解决了现有RNN在长文本序列处理中的性能限制问题。

Q: 论文做了哪些实验？

A: 论文中进行了一系列的实验来评估提出的TTT层与现有的Transformer和Mamba这两种流行的模型的性能对比。实验主要分为以下几个方面：

1. 短上下文（Short Context）实验：在2k和8k的上下文长度上，使用Pile数据集进行实验。这部分实验的目的是评估TTT层在处理较短序列时的性能。

2. 长上下文（Long Context）实验：在1k到32k的上下文长度上，使用Books3数据集进行实验。这部分实验的目的是评估TTT层在处理较长序列时的性能和可扩展性。

3. 模型大小对比：实验涵盖了不同大小的模型，包括125M、350M、760M和1.3B参数的配置，以评估模型规模对性能的影响。

4. 性能和计算复杂度（FLOPs）的对比：评估不同模型在相同计算复杂度下的困惑度（Perplexity），以及在不同FLOPs预算下的性能表现。

5. 时间效率（Wall-clock Time）的对比：评估不同模型在实际硬件上运行时的前向传播和生成（解码）的速度。

6. Transformer微调（Finetuning）实验：在长上下文设置中，比较从头开始训练的Transformer与在长上下文中微调的Transformer的性能差异。

7. 不同背景架构的对比：评估TTT层在Transformer背景架构和Mamba背景架构中的性能差异。

8. 不同TTT层实例化（TTT-Linear和TTT-MLP）的对比：评估两种不同隐藏状态（线性模型和两层MLP）的TTT层的性能。

9. 系统优化：评估mini-batch TTT和双重形式（dual form）对提高TTT层在硬件上运行效率的影响。

10. 理论等价性验证：通过实验验证TTT层在特定条件下与线性注意力（Linear Attention）和自注意力（Self-attention）的等价性。

这些实验结果用于全面评估TTT层的性能，并与现有的Transformer和Mamba模型进行比较。论文还讨论了实验中发现的一些开放性问题，并鼓励社区参与解决这些问题。

Q: 有什么可以进一步探索的点？

A: 论文中提出了一些可以进一步探索的点，主要包括以下几个方向：

1. 外循环参数化（Outer-loop Parameterization）：论文中提到，可以探索更多方式来参数化多视图重建任务或更一般的自监督任务家族。这意味着寻找新的方法来设计自监督任务，可能会提高TTT层的性能。

2. 系统优化（Systems Optimization）：尽管论文中已经进行了一些系统优化，但仍有改进空间。例如，可以通过时间上的流水线并行性来处理更长的序列，或者进一步优化现有的GPU和TPU实现。

3. 更长的上下文和更大的模型（Longer Context and Larger Models）：由于资源限制，论文中的实验没有在百万或十亿级别的上下文长度上进行训练。在更长的上下文中，TTT层的优势可能会更加明显，同时可能需要更大的模型来处理这些长序列。

4. 更雄心勃勃的f实例化（More Ambitious Instantiations of f）：随着上下文长度的增加，可能需要更大的f来处理视频任务和具身代理，其中上下文长度可以轻松扩展到数百万或数十亿。

5. 多级学习到学习（Multi-level Learning to Learn）：如果f本身是一个自注意力层，那么可以将其视为现有内循环内部的另一个嵌套学习问题。这样，可以构建多个嵌套学习问题的层次结构。

6. 人类学习与TTT的联系（Connection between Human Learning and TTT）：论文提出，与人类学习过程相比，TTT可能是一种更自然的建模方式，因为人类学习通常不涉及独立同分布（i.i.d.）的实例或训练-测试分割。进一步探索这种联系可能会提供对人类学习和机器学习之间差异的更深入理解。

7. 自监督任务的改进（Improvement of Self-Supervised Tasks）：论文中提到，自监督任务是TTT层成功的关键，可以探索更复杂或更高级的任务设计，以提高模型的学习能力。

8. 硬件效率的进一步提升（Further Improvement of Hardware Efficiency）：尽管已经采用了mini-batch TTT和双重形式来提高硬件效率，但仍有可能通过其他技术进一步提高效率，特别是在大规模并行处理和内存优化方面。

这些方向为未来的研究提供了丰富的探索空间，有助于推动TTT层及相关技术的发展。

Q: 总结一下论文的主要内容

A: 这篇论文的主要内容可以概括为以下几个要点：

1. 问题识别：论文指出现有的循环神经网络（RNN）在处理长文本序列时受限于其隐藏状态的固定大小，而自注意力机制虽然在长文本上表现更好，但计算复杂度较高。

2. TTT层的提出：为了解决上述问题，论文提出了一种新的序列建模层——Test-Time Training (TTT) 层。TTT层的核心思想是将隐藏状态设计为一个机器学习模型，并通过自监督学习的更新规则在测试时更新这个状态。

3. TTT层的特性：

   • 线性复杂度：TTT层保持了RNN的线性复杂度优势。

   • 表达性：通过自监督学习，TTT层能够更好地捕捉长文本序列中的信息。

4. 实例化方法：论文提出了两种TTT层的实例化方法——TTT-Linear和TTT-MLP，分别使用线性模型和两层MLP作为隐藏状态。

5. 系统优化：为了提高TTT层在实际硬件上的性能，论文提出了mini-batch TTT和双重形式（dual form）的技术，以更好地利用现代GPU和TPU。

6. 自监督任务学习：TTT层的自监督任务是可学习的，可以通过外循环优化，以适应语言建模的最终目标。

7. 实验评估：论文通过在不同规模的参数配置下进行实验，比较了TTT层与Transformer和Mamba模型的性能，证明了TTT层在长文本序列建模中的潜力。

8. 未来研究方向：论文讨论了多个未来研究的方向，包括外循环参数化、系统优化、更长上下文和更大模型的探索、更复杂的隐藏状态模型、多级学习到学习等。

9. 理论与实践的结合：论文不仅提出了理论上的创新，还通过实验验证了TTT层的有效性，并探讨了其与人类学习过程的联系。

10. 代码和数据集的公开：论文提供了代码和使用的数据集，以便研究社区可以复现实验结果并在此基础上进行进一步的研究。

总的来说，这篇论文提出了一种新的序列建模方法，通过在测试时对模型进行训练来提高其在长文本序列上的表达能力，同时保持了计算效率，为未来的研究提供了新的方向和工具。

Q: 想要进一步了解论文

A: 以上只是了解一篇论文的几个基本FAQ。如果你还想与Kimi进一步讨论该论文，请点击 这里  为你跳转Kimi AI网页版，并启动一个与该论文相关的新会话。

转自: https://papers.cool/arxiv/search?highlight=1&query=Learning+to+Learn+at+Test+Time+RNNs+with+Expressive+Hidden+States