当前位置: 华文世界 > 科技

从根本上改变语言模型!全新架构TTT超越Transformer和Mamba,长上下文中性能受限问题或被彻底解决,12人研究团队中一半是华人

2024-07-10科技

从2017年Google Brain团队推出Transformer架构逐步取代长短期记忆(LSTM)等「循环神经网络(RNN)模型」成为首选模型,到后来首个线性时间序列架构Mamba推出又对Transformer架构构成挑战,大语言模型底层架构的迭代正在迅速改变人们对于AI的认知和理解。

美东时间周一(7月8日),一种全新的大语言模型(LLM)架构有望代替至今在AI领域如日中天的Transformer,性能也比Mamba更好。

图片来源:arXiv

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

但对于该论文,也有人提出质疑,认为 只有30亿~70亿参数的可用演示模型才足以了解其实用性。

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

过去这些年来,对大模型的研究和理解都绕不开「循环神经网络(下称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长上下文时表现还要更好。

图片来源:arXiv

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

图片来源:arXiv

12人团队半数为华人,排首位的「共同一作」为斯坦福博士后

本篇论文一共有12名作者,有一半(6人)成员为华人。 其中,排名第一的「一作」Yu Sun博士毕业于加州大学伯克利分校电气工程与计算机科学系,目前是斯坦福大学的博士后,研究重点便是TTT的算法架构。排在第二位的 「一作」 Xinhao Li为加州大学圣地亚哥分校研二学生,其研究集中在深度学习和计算机视觉,「三作」之一的Xiaolong Wang为其导师。

Yu Sun 图片来源:个人主页

论文称,Yu Sun于2022年11月便开始和Xinhao Li做这个项目,2023年6月份开始进入全职工作状态。Yu Sun提出了项目的概念框架,设计了小批量的TTT层和「双重形式」,在他人的帮助下撰写了论文,并领导了整个团队的日常运作和实验。

该研究「一作」之一 、加州大学伯克利分校信息工程学系学生Karen Dalal在X上表示, 他相信这将从根本上改变语言模型。 他称,「我们设计了一个新的架构,用机器学习模型取代了RNN的隐藏状态。该模型通过输入标记的实际梯度下降来压缩上下文。我们将我们的方法称为‘测试时间训练层’。 TTT层通过表达性记忆解锁了线性复杂性架构,使我们能够在上下文中用数百万(有朝一日甚至数十亿)个token来训练大语言模型 。」

图片来源:X

Karen Dalal还称,「TTT-Linear已经比最快的SSM(注:指‘状态空间模型’)更快,并且在大小和上下文方面具有很强的可扩展性。 这个架构内的探索空间是巨大的,我们的论文只是迈出了一小步。

图片来源:X

该论文的「三作」之一、加州大学伯克利分校博士后,现加州大学圣地亚哥分校电气与计算机工程助理教授Xiaolong Wang则在X上激动地称:「我真不敢相信这终于发生了。」

「TTT层理念是我们已经研究了5年的架构……今天的TTT和我刚开始做博士后研究的时候已经完全不同了,它已经是一个网络层,用机器学习模型取代了RNN的隐藏状态。我们的TTT层并不是使用特定的向量来表达记忆,而是维护一个小型神经网络来压缩输入标记…… 这种架构目前应用于语言建模,但想象一下将其应用于视频。未来,在长视频建模时,我们可以对帧进行密集采样,而不是以1 FPS的速度采样,这些密集的帧对Transformer架构来说是负担,但对TTT层来说却是福音。因为它们本质上只是在TTT内训练更好网络的‘时间增强’。

图片来源:X

【每日经济新闻】记者注意到,在Xiaolong Wang的推文下方,也有不少质疑者,例如,一个名为John Bollenbacher的用户称,「这是一个有希望的迹象,但我认为我需要看到2.5个以上的示例才能认真对待。而一个30亿~70亿参数的可用演示模型才可以让我们了解它的实用性。」

对此,Xiaolong Wang回复称,「 (我们正)朝着那个方向努力。我希望通过这个版本我们可以得到更多的计算……就目前而言,学术实验室仍然很难训练70亿参数的模型。

图片来源:X

每日经济新闻