新架构RNN反超Transformer：每个隐藏状态都是一个模型，一作：从根本上改变语言模型

495 0 0

文章摘要

【关键词】 TTT架构、Transformer挑战、测试时学习、隐藏状态优化、长上下文处理

一项新的架构TTT（Test-Time Training）向现有的Transformer模型发起了挑战，提出了一种在测试时也能学习的模型，旨在从根本上改变语言模型。
TTT的核心思想是用可学习的模型替换RNN中的隐藏状态，这样的模型在测试时仍可进行学习。
实验表明，TTT层拥有比RNN表达能力更强的隐藏状态，能够取代Transformer中昂贵的自注意力层。
TTT-Linear和TTT-MLP两种模型在实验中表现出色，使用较少的算力达到更低的困惑度，并且能更好地利用长上下文。

TTT模型的隐藏状态可以是任何模型，例如CNN或完整的Transformer，这为处理更长上下文提供了可能性。
TTT论文已经引起了学术界的关注，斯坦福博士生Andrew Gao认为这篇论文可能成为下一篇引领潮流的论文。
尽管如此，TTT是否能成为“Transformer杀手”还需要进一步的研究和实验来验证。

TTT模型的提出是基于对传统RNN和Transformer局限性的认识。
传统RNN的隐藏状态大小固定，表达能力受限，且难以并行训练。
而Transformer的自注意力机制随着上下文长度的增加，计算复杂度呈平方增长，导致资源消耗巨大。
TTT模型通过自监督学习方法在测试时对上下文进行压缩，使得隐藏状态模型能够记住产生大梯度的输入，从而获得更强的拟合和泛化能力。

TTT模型的实验主要在Mamba骨干上进行，因为现代RNN的骨干中包含时间卷积，对TTT也有帮助。
在Pile数据集的短上下文测试中，TTT-Linear和TTT-MLP在长上下文条件下表现优于Mamba和Transformer。
长上下文实验使用Pile的子集Books3，结果也显示TTT层的优势随着上下文长度的增长而扩大。

TTT模型的另一个优势是其时间复杂度保持线性，这得益于模型的两个嵌套学习循环设计。
外循环负责整体的语言建模，内循环通过自监督学习压缩上下文信息。
此外，TTT模型在速度上也具有优势，在A100上测试速度显示，TTT-Linear在预填充阶段比Mamba稍快，解码阶段速度相当。

尽管TTT模型展现出巨大潜力，但其共同一作Karan Dala表示，目前的研究只涉及了线性模型和小型MLP，对于更复杂的模型还有待进一步研究。
此外，隐藏状态模型的学习过程中，可以使用Adam优化器代替普通的梯度下降，以提高学习效率。
TTT模型的未来发展和应用，特别是在视频建模等领域，值得期待。