完全改变语言模型：TTT超越Transformer的新体系结构，ML模型取代RNN隐藏状态-人工智能

首页 2024-07-10 23:24:32

从 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

方法介绍

长上下文的挑战是 RNN 层本质上是固有的：不同于自我注意机制，RNN 层必须将上下文压缩成固定大小的隐藏状态，更新规则需要发现数千甚至数百万 token 底层结构与关系。

研究小组首先观察到，自我监督学习可以将大量的训练集压缩成 LLM 等模型的权重，而 LLM 模型通常表现出对其训练数据之间语义关系的深刻理解。

受此观察的启发，研究团队设计了一种新的序列建模层，其中隐藏状态是一种模型，更新规则是自我监督和学习的一步。由于更新测试序列中隐藏状态的过程相当于测试时的训练模型，研究团队称这一新层为测试时的训练（Test-Time Training，TTT）层。

研究小组介绍了两个简单的例子：TTT-Linear 和 TTT-MLP，隐藏状态分别是线性模型和两层 MLP。TTT 层可以集成到任何网络架构中，并进行端到端优化，类似于 RNN 层和自注意力。

为了让 TTT 这项研究采用了一些技巧来提高层次效率 TTT 层：

首先，类似于在常规训练期间采用小批量序列 gradient step 本研究旨在获得更好的并行性 TTT 小批量使用期间 token。

其次，每一项研究都是这样 TTT 为了更好地利用现代化，小批量操作开发了双重形式 GPU 和 TPU。双重形式的输出等同于简单的实现，但训练速度更快 5 倍以上。如图所示 3 所示，TTT-Linear 在 8k 上下文中比 Transformer 更快，与 Mamba 相当。

研究团队认为，所有序列建模层都可以被视为将历史上下文存储到隐藏状态，如图所示 4 所示。

例如，RNN 层（如 LSTM、RWKV 和 Mamba 层)将上下文压缩成跨时间固定状态。这种压缩有两个后果：一方面，输入标记 x_t 映射到输出 token z_t 因为每一个，都是高效的 token 更新规则和输出规则需要恒定的时间。另一方面，RNN 层在上下文中的性能受其隐藏状态的限制 s_t 的表现力。

自我注意力也可以从上述角度来看，但它的隐藏状态(通常称为 Key-Value 缓存)是一个随行 t 线性增长列表。它的更新规则只是当前的更新规则 KV 元组（tuple）添加到列表中，并扫描输出规则 t 以前的所有元组形成注意力矩阵。隐藏状态清楚地存储了所有的历史上下文，不需要压缩，这使得自我注意力在长上下文中比较 RNN 层更有表现力。然而，扫描线性增长的隐藏状态所需的时间也是线性增长。研究人员需要更好的压缩启发，以保持长上下文的高效性和表现力。具体来说，需要数千万或数百万 token 压缩到一个隐藏的状态，以有效地捕捉它们的底层结构和关系。这听起来有点难，但实际上很多人都很熟悉这种启发式。

骨干架构。将任何 RNN 将层集成到更大架构中最简单的方法是直接替换 Transformer 自我注意力在这里被称为骨干。然而，现有的 RNN（如 Mamba 和 Griffin 等等)都用了 Transformer 骨干层不同。最值得注意的是，它们的骨干层在 RNN 层之前包含时间卷积，这可能有助于收集跨时间的局部信息。在对 Mamba 经过主干网的试验，研究人员发现它也可以改进 TTT 因此，将层的混乱纳入建议方法，详见图片 16。

实验结果

研究者将在实验中进行 TTT-Linear 、 TTT-MLP 与 Transformer、Mamba 这两条基线进行了比较。

短文本

从图 11 以下结论可以从中得出：

2k 上下文，TTT-Linear (M)、Mamba 和 Transformer 性能相当，因为大部分线条重叠。在 FLOP 当预算较大时，TTT-MLP (M) 性能稍差。尽管 TTT-MLP 在各种模型的大小下进行比较 TTT-Linear 有更好的困惑，但是 FLOPs 额外的成本抵消了这一优势。
8k 上下文，TTT-Linear (M) 和 TTT-MLP (M) 表现明显优于 Mamba，这与 2k 上下文的观察结果完全不同。即使是使用 Transformer 主干网络的 TTT-MLP (T) 在 1.3B 左右也比 Mamba 略胜一筹。一个显著的现象是，随着上下文长度的增加，TTT 层相对于 Mamba 层层的优势也在扩大。
达到上下文的长度 8k，Transformer 每个模型尺寸下的混乱度还是不错的，但是因为 FLOPs 成本的原因，已经没有竞争力了。

上图结果显示了将军 TTT 层从 Mamba 切换到主干网络 Transformer 主要网络的影响。研究人员假设，当序列建模层的隐藏状态表现较低时，Mamba 时序卷积在主干网络中更有帮助。线性模型的表现力低于线性模型 MLP，因此，从卷积中获益更多。

长文：书籍

研究人员使用长上下文来评估他们的能力 Pile 流行子集 Books3，以 2 倍的增量对 1k 到 32k 上下文长进行实验。这里的训练方法和 Pile 相同，并且 TTT 所有层的实验都是在一次训练操作中完成的。从图表开始 12 在结果子集中，他们得到了以下观察结果：

在 Books 的 2k 上下文中，Pile 2k 所有的观察结果仍然成立，但是 Mamba 现在的表现略好于 TTT-Linear(而且它们的线条在 Pile 2k 大致重叠)。

在 32k 上下文中，TTT-Linear (M) 和 TTT-MLP (M) 所有的表现都比 Mamba，类似于 Pile 8k 观察结果。甚至使用 Transformer 主干的 TTT-MLP (T) 在 32k 上下文的表现也略好于 Mamba。

TTT-MLP (T) 在 1.3B 仅略差于规模 TTT-MLP (M)。如上所述，由于缺乏清晰的线性拟合，经验缩放定律很难获得。然而，TTT-MLP (T) 强劲的趋势表明，Transformer 超出我们的评价范围，主干可能更适合更大的模型和更长的上下文。

时钟时间

LLM 训练和推理可以分解为前、后和生成。推理过程中的提示处理(也称为预填充)和前向操作在训练过程中是相同的，但后向操作不需要存储中间激活值。

由于前向(在训练和推理过程中)和后向可以并行处理，这里采用了双重形式。生成新的 token(也叫解码)本质上是顺序的，所以这里采用了原始形式。

研究人员提到，由于资源的限制，本文的实验使用 JAX 编写，并在 TPU 上运行。在 v5e-256 TPU pod 上，Transformer 上下文为基线 2k 每次迭代训练都需要 0.30 秒，而 TTT-Linear 每次迭代都需要 0.27 秒，没有任何系统优化，速度更快 10%。鉴于 Mamba（用 PyTorch、Triton 和 CUDA 实现)只能在那里 GPU 为了进行公平的比较，研究人员对本文的方法进行了初步的系统优化，使其能够进行比较 GPU 上运行。

图 15 每个模型的前向内核显示在左侧，批量尺寸为 16 时间延迟。所有模型都是 1.3B（Mamba 为 1.4B）。值得注意的是，这里的 Transformer 基线要比 Mamba 论文要快得多，因为它在这里使用 vLLM ，而不是 HuggingFace Transformer 。

此外，研究人员还为生成编写了另一篇文章 GPU 内核，并在图中 15 右侧为批量尺寸 512 测试基准速度。另一个常用的挂钟时间（wall-clock time）指标是吞吐量（throughput），它考虑了使用更大批量尺寸的潜在好处。上述所有观察结果和方法之间的排序对吞吐量仍然有效。

主要作者

在 TTT 论文作者之一在研究提交后，UCSD 助理教授 Xiaolong Wang 表示祝贺。他说，TTT 研究持续了一年半，但测试时间训练（TTT）这个想法从诞生到现在已经五年了。虽然最初的想法与现在的结果完全不同。

TTT 来自斯坦福的三位主要作者，UC Berkeley 和 UCSD。

其中 Yu Sun 他是斯坦福大学的博士后，博士毕业于斯坦福大学 UC Berkeley EECS，长期以来的研究方向一直是 TTT。

Xinhao Li 是 UCSD 他本科毕业于电子科技大学攻读博士学位。

Karan Dalal 是 UC Berkeley 在读博士的时候，他在高中的时候和别人一起创办了一个名字 Otto 兽医远程医疗初创公司。

以上三个人，都把 test-time training 本文介绍了个人网站研究方向的第一行。

更多研究细节，请参考原论文。

以上是完全改变语言模型的详细内容：新架构TTT超越Transformer，ML模型取代RNN隐藏状态，请关注其他相关文章！

阅读全文