语言模型：全新架构TTT或将超越Transformer

9 个月前语言模型 123

一种全新的大语言模型（LLM）架构有望代替至今在 AI 领域如日中天的 Transformer，性能也比 Mamba 更好。2024年7月8日，有关 Test-Time Training（TTT）的论文成为了人工智能社区热议的话题。

斯坦福大学、加州大学伯克利分校、加州大学圣地亚哥分校和Meta的学者提出了一种全新架构，希望能用机器学习模型取代RNN的隐藏状态。这个架构通过对输入token进行梯度下降来压缩上下文，被称为“测试时间训练层（Test-Time-Training layers，简称TTT层）”。“共同一作”加州大学伯克利分校的Karen Dalal表示，我相信这将从根本上改变语言模型。

TTT-Linear 和 TTT-MLP 均能匹敌或击败Transformer和Mamba

在机器学习模型中，TTT 层直接取代 Attention，并通过表达性记忆解锁线性复杂性架构，使我们能够在上下文中训练具有数百万（有时是数十亿）个 token 的 LLM。

在 125M 到 1.3B 参数规模的大模型上进行了一系列对比发现，TTT-Linear 和 TTT-MLP 均能匹敌或击败最强大的 Transformers 和 Mamba 架构方法。

TTT 层作为一种新的信息压缩和模型记忆机制，可以简单地直接替代 Transformer 中的自注意力层。

过去这些年来，对大模型的研究和理解都绕不开“循环神经网络（RNN）”。RNN是一种深度学习模型，由许多相互连接的组件组成，经过训练后可以处理顺序数据输入并将其转换为特定的顺序数据输出，例如将文本从一种语言翻译成另一种语言。顺序数据是指单词、句子或时间序列数据之类的数据，其中的顺序分量根据复杂的语义和语法规则相互关联。

而“隐藏状态”是RNN模型中的一个关键概念。它可以看作是网络在每个时间步骤上的“记忆”，存储了之前时间步骤中的信息，并通过时间在不同步骤之间传递。隐藏状态可以捕捉到序列中的长期依赖性，从而使模型能够理解整个序列的上下文。

在传统的RNN中，隐藏状态的固定大小表达能力受限，也不好并行训练。例如，像Mamba这样的RNN层，会随着时间的推移压缩成一个固定大小的状态，它们虽然效率很高，但性能受限于其表达能力。

这个研究团队的对TTT层的想法来自于：与其让RNN隐藏状态被动地储存信息，不如让它主动学习。**设计的“TTT层”突破了“RNN层”在长上下文中性能受限的问题。**

他们在1.25亿~ 13亿个参数规模的大模型上进行一系列的对比后发现，TTT-Linear（线性模型）和TTT-MLP （MLP为多层感知器，基于前馈神经网络的深度学习模型）均能匹敌或击败最强大的Transformers和 Mamba架构方法。

隐藏状态时线性模型的TTT-Linear表现超过了Transformer和Mamba，用更少的算力达到更低的困惑度（下图左），也能更好利用长上下文（下图右）。此外，隐藏状态时MLP模型的TTT-MLP在32k长上下文时表现还要更好。