新架构,再次向Transformer发起挑战!
核心思想:将RNN中的隐藏状态换成可学习的模型。
甚至在测试时都可以学习,所以该方法称为TTT(Test-Time Training)。
共同一作UC伯克利的Karen Dalal表示:我相信这将从根本上改变语言模型。
一个TTT层拥有比RNN表达能力更强的隐藏状态,可以直接取代Transformer中昂贵的自注意力层。
在实验中,隐藏状态是线性模型的TTT-Linear表现超过了Transformer和Mamba,用更少的算力达到更低的困惑度(左),也能更好利用长上下文(右)。
此外,隐藏状态是MLP模型的TTT-MLP在32k长上下文时表现还要更好。
Karen Dalel还指出,理论上可学习的隐藏状态可以是任意模型,对于更长上下文来说,可以是CNN、甚至可以是完整的Transformer来套娃。
目前刚刚出炉的TTT论文已经在学术界引起关注和讨论,斯坦福博士生Andrew Gao认为,这篇论文或许能成为下一篇Attention is all you need。
另外有人表示,众多新架构能否真正击败Transformer,还要看能不能扩展到更大规模。
Karen Dalel透露,马上就会推出7B模型。
用机器学习模型来压缩上下文
传统RNN,隐藏状态固定大小表达能力受限,也不好并行训练。
Transformer强大,但自注意力机制随上下文长度呈平方复杂度,非常昂贵。
最近一系列基于RNN的架构创新中:
RWKV,用线性注意力结合RNN和Transformer的优点,在训练时可以并行计算。
Mamba,赋予模型选择性记住或遗忘信息的能力来压缩上下文,同时设计了面向硬件的高效并行算法。
它们的表现在短上下文时追上甚至超越了Transformer,但在32k超长上下文以上,Trasformer依旧称霸。
TTT团队的想法来自于:与其让隐藏状态被动地储存信息,不如让它主动学习。
就像Transformer模型作为一个整体在压缩互联网数据到参数中一样,可学习的隐藏状态模型也在少量参数上不断缩上下文信息。
这种“隐藏状态模型”随着时间的推移仍然具有固定的大小(固定的模型参数),但表达能力更强了。
论文的联合指导UCSD助理教授王小龙认为:
Transformer显式地储存所有输入token,如果你认为个神经网络是压缩信息的好方法,那么压缩这些token也将是有意义的。
如此一来,整个框架的时间复杂度还是线性的,
至此,序列建模被拆解为两个嵌套的学习循环,外循环负责整体的语言建模,内循环通过自监督学习压缩上下文信息。
外循环的参数变成了内循环的超参数,也就是元学习的一个变种了。
标准的元学习是训练一个适应不同任务的模型,而TTT是让模型去适应每一个测试样本。单个样本虽然信息量小,但用来训练隐藏状态模型也绰绰有余。
特别的,在内循环是一个线性模型时,相当于线性注意力。当内循环是一个Nadaraya-Watson estimator时,TTT等价于自注意力。
在测试时学习
在TTT层里,使用自监督学习方法将上下文压缩到隐藏状态。
上下文就是未标记的数据集,隐藏状态不再是一个固定的向量,可以是线性模型、小型神经网络或任何机器学习模型,更新规则采用了在自监督损失上的一步梯度下降。
这样一来,隐藏状态模型可以记住产生大梯度的输入,并且可以获得比选择性遗忘机制更强的拟合和泛化能力,并且在测试时仍然为每个输入序列训练不同的参数。
到目前为止,朴素的TTT层已经有效了,但还无法并行化。
团队提出的解决方案为mini-batch梯度下降,把一个batch内的梯度计算并行化。
再通过Dual form方法,只在mini-batch结束时计算权重以及输出token,避免冗余计算。在JAX版实现中快了5倍以上。
TTT能否成为“Transformer杀手”?
理论上都走的通了,那么TTT在实验中表现到底如何?
最简单干净的测试方法,应该是直接替换掉Transformer中的自注意力层。
但是在研究过程中,团队发现Mamba等现代RNN的骨干中在RNN层之前还包含时间卷积,对TTT也有帮助。
所以实验中TTT-Linear和TTT-MLP主要应用到Mamba骨干上,其他训练细节也严格遵照Mamba论文中的设置。
最终在Pile数据集短上下文测试中:
-
2k上下文时,TTT-Linear、Mamba和Transform具有相当的性能,TTT-MLP的表现略差。
-
8k上下文时,TTT-Linear和TTT-MLP都优于Mamba和Transformer,应用在Transformer骨干的TTT-MLP(T)在1.3B参数左右也略好与Mamba。
总的来说,随着上下文长度的增长,TTT层相对于Mamba的优势也会扩大。
另外团队猜测,线性模型比MLP表达能力差,因此从Mamba骨干的卷积中受益更多。
长上下文实验使用Pile的子集Books3:
-
32k上下文,TTT-Linear和TTT-MLP的表现都优于曼巴,类似于Pile 8k的观察。即使是带有Transformer骨干的TTT-MLP(T)表现也略好于曼巴。
-
1.3B参数尺度上,TTT-MLP(T)仅比TTT-MLP(M)稍差,Transformer骨干可能更适合论文评估范围之外的更大模型和更长的上下文。
在A100上测试速度,TTT-Linear在预填充阶段比Mamba稍快,解码阶段几乎与Mamba速度相同。TTT-MLP相比Transformer整体上也有线性复杂度的优势。
共同一作Karan Dala表示:我一直被问到的一个问题是,我们是否相信TTT就是“Transformer杀手”,我仍然认为我们需要继续努力。
隐藏状态可以是任意模型,但目前的研究只涉及了线性模型和小型MLP,更复杂的还有待研究。
隐藏状态模型的学习可以用Adam代替普通的梯度下降等等。
如何学习大模型 AI ?
由于新岗位的生产效率,要优于被取代岗位的生产效率,所以实际上整个社会的生产效率是提升的。
但是具体到个人,只能说是:
“最先掌握AI的人,将会比较晚掌握AI的人有竞争优势”。
这句话,放在计算机、互联网、移动互联网的开局时期,都是一样的道理。
我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。
我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。
第一阶段(10天):初阶应用
该阶段让大家对大模型 AI有一个最前沿的认识,对大模型 AI 的理解超过 95% 的人,可以在相关讨论时发表高级、不跟风、又接地气的见解,别人只会和 AI 聊天,而你能调教 AI,并能用代码将大模型和业务衔接。
- 大模型 AI 能干什么?
- 大模型是怎样获得「智能」的?
- 用好 AI 的核心心法
- 大模型应用业务架构
- 大模型应用技术架构
- 代码示例:向 GPT-3.5 灌入新知识
- 提示工程的意义和核心思想
- Prompt 典型构成
- 指令调优方法论
- 思维链和思维树
- Prompt 攻击和防范
- …
第二阶段(30天):高阶应用
该阶段我们正式进入大模型 AI 进阶实战学习,学会构造私有知识库,扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架,抓住最新的技术进展,适合 Python 和 JavaScript 程序员。
- 为什么要做 RAG
- 搭建一个简单的 ChatPDF
- 检索的基础概念
- 什么是向量表示(Embeddings)
- 向量数据库与向量检索
- 基于向量检索的 RAG
- 搭建 RAG 系统的扩展知识
- 混合检索与 RAG-Fusion 简介
- 向量模型本地部署
- …
第三阶段(30天):模型训练
恭喜你,如果学到这里,你基本可以找到一份大模型 AI相关的工作,自己也能训练 GPT 了!通过微调,训练自己的垂直大模型,能独立训练开源多模态大模型,掌握更多技术方案。
到此为止,大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗?
- 为什么要做 RAG
- 什么是模型
- 什么是模型训练
- 求解器 & 损失函数简介
- 小实验2:手写一个简单的神经网络并训练它
- 什么是训练/预训练/微调/轻量化微调
- Transformer结构简介
- 轻量化微调
- 实验数据集的构建
- …
第四阶段(20天):商业闭环
对全球大模型从性能、吞吐量、成本等方面有一定的认知,可以在云端和本地等多种环境下部署大模型,找到适合自己的项目/创业方向,做一名被 AI 武装的产品经理。
- 硬件选型
- 带你了解全球大模型
- 使用国产大模型服务
- 搭建 OpenAI 代理
- 热身:基于阿里云 PAI 部署 Stable Diffusion
- 在本地计算机运行大模型
- 大模型的私有化部署
- 基于 vLLM 部署大模型
- 案例:如何优雅地在阿里云私有部署开源大模型
- 部署一套开源 LLM 项目
- 内容安全
- 互联网信息服务算法备案
- …
学习是一个过程,只要学习就会有挑战。天道酬勤,你越努力,就会成为越优秀的自己。
如果你能在15天内完成所有的任务,那你堪称天才。然而,如果你能完成 60-70% 的内容,你就已经开始具备成为一名大模型 AI 的正确特征了。