彻底改变语言模型:全新架构TTT超越Transformer,ML模型代替RNN隐藏状态
- 2024-07-25 11:48:00
- admin 原创
- 86
从 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 写在了个人网站介绍研究方向的第一行。
更多研究细节,可参考原论文。
发表评论
文章分类
联系我们
联系人: | 羊毛君 |
---|---|
Email: | soft850207@gmail.com |