彻底改变语言模型：全新架构TTT超越Transformer，ML模型代替RNN隐藏状态

从 125M 到 1.3B 的大模型，性能都有提升。

难以置信，这件事终于发生了。

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

论文链接：https://arxiv.org/abs/2407.04620

该研究的作者来自斯坦福大学、加州大学伯克利分校、加州大学圣迭戈分校和 Meta。他们设计了一种新架构 TTT，用机器学习模型取代了 RNN 的隐藏状态。该模型通过输入 token 的实际梯度下降来压缩上下文。

该研究作者之一 Karan Dalal 表示，他相信这将根本性的改变语言模型方法。

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

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

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

与 Mamba 相比，TTT-Linear 的困惑度更低，FLOP 更少（左），对长上下文的利用更好（右）：

这不仅在理论上是线性的复杂度，而且实际运行时间也更快。

• 在论文上线后，作者公开了代码与 jax 以供人们训练和测试：https://github.com/test-time-training/ttt-lm-jax

• 还有 PyTorch 推理代码：https://github.com/test-time-training/ttt-lm-pytorch

01

方法介绍

长上下文的挑战是 RNN 层本质上所固有的：与自注意力机制不同&