TTT大语言模型架构发布,成功撼动了Transformer与Mamba模型

智能也得细细瞧 2024-07-12 19:08:51

上期图文,我们刚,觉得 Mamba 模型的发布可以有效改善 transformer 模型长序列输入复杂度的问题,应该能够成为下一个大语言模型的基石。谁想Mamba2 还没有发布几天,这边最新的大语言模型TTT模型发布,成功撼动了 transformer 与 Mamba 模型。

transformer 模型的自注意力机制在长上下文中表现良好,但其计算复杂度随着输入序列的长度平方的增长。现有的RNN循环神经网络模型具有线性复杂度,但RNN循环神经网络模型在长上下文中的表现受到隐藏状态表达能力的限制。UCSD, UC伯克利,斯坦福与 meta 科研人员提出了一种新的序列建模层。模型不仅具有线性复杂度,且使用了表达能力强的隐藏状态。关键思想是将隐藏状态本身设为一个机器学习模型,并将更新规则设为自监督学习的一步。由于隐藏状态在测试序列上也通过训练进行更新,因此称这些层为测试时训练(Test-Time Training, TTT)层。

TTT模型考虑了两种实例:TTT-Linear和TTT-MLP,其中隐藏状态分别是线性模型和两层MLP。研究人员在125M到1.3B参数规模下评估了TTT模型,并与Transformer和Mamba进行比较。TTT-Linear和TTT-MLP都匹配或超过了基线模型。与Transformer类似,它们可以通过更多的tokens来持续降低困惑度,而Mamba在16k上下文后则不能持续降低困惑度。经过初步的系统优化,TTT-Linear在8k上下文中已经比Transformer更快,并在wall-clock 时间上匹配Mamba的性能。TTT-MLP在内存I/O方面仍面临挑战,但在长上下文中显示了更大的潜力性能。从下图可以看出,transformer 模型随着输入长度的增加,每个 token 的前向传递世界也是直线上升,而 Mamba 与TTT模型几乎不会变化,transformer 模型真的要被比下去了吗?

自注意力机制隐藏状态(通常称为Key-Value (KV)缓存)是一个随着时间 t 线性增长的列表。它的更新规则只是将当前的KV元组追加到计算列表中,输出规则则扫描直到时间 t 的所有元组以形成注意力矩阵。隐藏状态明确地存储了所有的历史上下文而不进行压缩,这使得自注意力机制在长上下文中比RNN层更具表现力。然而,扫描这个线性增长的隐藏状态也需要线性增长的时间。

为了在长上下文中保持高效和表现力,大语言模型需要更好的压缩方法。一个通用的序列建模层可以表示为一个根据更新规则转换的隐藏状态。所有的序列建模层都可以看作是下图中三个组件(初始状态、隐藏状态和输出状态)不同实例的表现。自注意力机制的隐藏状态随着上下文的增长而增长,因此每个token的成本也会增加。而简单的RNN和TTT层则将增长的上下文压缩成一个固定大小的隐藏状态,因此它们每个token的成本保持不变。

TTT模型使用自监督的方式更新隐藏状态,其TTT模型直接使用TTT层代替了 transformer 模型的注意力机制层。使用RNN的线性复杂度方法建立大语言模型,但是基于RNN循环神经网络模型,其无法实现并行化运算。

为了实现并行化计算,TTT模型使用Mini-batch梯度下降,使用 b代表一个批量大小。论文使用了 b等于 16。 从下图可以看出,G1, G2,..., Gb没有任何前后时间步骤的关系,可以直接使用并行化计算,这样可以同时计算 b个梯度,加速模型的并行操作。

https://arxiv.org/pdf/2407.04620Learning to (Learn at Test Time): RNNs with Expressive Hidden States

人工智能领域苦 transformer 模型久矣,但是 transformer 模型是所有模型的基础,虽然最新的模型Retnet,Mamba,TTT模型都有效改进了 transformer 模型在长输入序列的缺点,但transformer 模型可是 2017 年发布的模型,7 年前的模型被成功应用在各种大语言模型上面,其模型框架更是被各种机器学习任务魔改,包含计算机视觉任务,语音识别与语音合成等,其 transformer 模型功不可没。

而 transformer 模型更是大模型的基础,学习了 transformer 模型的核心知识,再去掌握其他模型更是轻而易举了。目前包含 ChatGPT,Gemini,llama 等模型都是基于 transformer 的注意力机制,而其他模型是否可以成功应用到各个任务上,需要更多参与者进入尝试了。

0 阅读:0

智能也得细细瞧

简介:感谢大家的关注