arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2407.04620cs.LGcs.AIcs.CL
收藏

学会(在测试时学习):具备强表达能力隐状态的RNN

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

Stanford University(斯坦福大学);UC San Diego(圣地亚哥大学);UC Berkeley(伯克利大学);Meta AI Core contributors(Meta AI 核心贡献者)

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

AI总结:

针对现有RNN长上下文性能受隐状态表达力限制的问题,提出隐状态为机器学习模型、更新基于自监督学习的TTT层,其长上下文表现优于Mamba,为相关研究提供了新方向。

摘要:

自注意力机制在长上下文场景下表现优异,但复杂度为二次方级。现有RNN层虽具备线性复杂度,但其长上下文性能受限于隐状态的表达能力。我们提出了一个实用框架,用于构建兼具线性复杂度与强表达隐状态的序列建模层。其核心思想是将隐状态本身设为一个机器学习模型,隐状态的更新规则则对应一步自监督学习过程。由于隐状态即使在处理测试序列时也通过训练完成更新,我们将这类层命名为Test-Time Training(测试时训练,TTT)层。我们设计了两种具体实现:TTT-Linear和TTT-MLP,二者的隐状态分别为线性模型和两层MLP。我们在1.25亿到13亿参数的规模下对所提实现进行评估,与性能强劲的Transformer以及现代RNN模型Mamba进行对比。与Transformer类似,TTT-Linear和TTT-MLP能够通过利用更多token持续降低困惑度,而Mamba在上下文长度达到16k后便无法再实现困惑度下降。TTT-MLP仍面临内存I/O方面的挑战,但在长上下文场景中展现出更大潜力,为未来研究指明了一个有前景的方向。

英文摘要:

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 states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. 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. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑