Skip to content
DeepMind研究成本大起底,一篇ICML论文烧掉1290万美元
DeepMind最近被ICML 2024接收的一篇论文,完完全全暴露了他们背靠谷歌的「豪横」。一篇文章预估了这项研究所需的算力和成本,大概是Llama 3预训练的15%,耗费资金可达12.9M美元。
发一篇顶会论文,需要多少实验预算?最近,DeepMind发表了一项研究,对LLM扩大规模时各种算法和架构细节,比如参数和优化器的选择,进行了广泛的实证调查。这篇论文已被ICML 2024接收。
论文地址:https://arxiv.org/abs/2407.05872
63页的论文涵盖了数以万计的模型,备选方案包括3种优化器、4种参数化方案、几种对齐假设、十多个学习率,以及最高达26.8B的14种参数规模。
需要进行实验的4种参数化方案仅仅听到这些数字,就不难知道,这项研究必定涉及海量的模型运行实验。而有一位忠实读者,为了测试自己对论文内容的理解,统计了其中进行的所有实验,并估算出了复现论文的成本。将所需算力全部加在一起,林林总总,居然达到了惊人的1290万美元。考验基本功的时刻到了,假如你是研究团队的leader,根据实验计划对所需算力和成本进行预估是一项必不可少的技能。那就让我们跟着这篇博客文章盘一遍,这一千多万美元,究竟烧在哪里。
Transformer架构信息
论文附录C提供了关于模型算法和架构的各种细节设置,比如使用decoder-only架构、层归一化、GeLU激活函数、无dropout、T5分词器、批大小为256、用FSDP并行等等。
实验模型的参数规模统计通过架构方面的信息,我们可以大致估算出训练中每个token所需的FLOPS,记为M。由于论文没有描述到任何GQA/MQA机制,所以就假设Rkv=1,此外还有lseq=512,Dhead=128,L=8(深度),V=32101(分词器词汇量)。模型总参数量可以表示为:因此,就可以得到M的计算公式:默认情况下,每次实验处理的token数(tokens per experiment, TPE)为5k(训练步数)×256(批大小)×512(lseq),约为6.5536e9。
def M(d: int, L=8, l_seq=512, V=32101) -> int:
return 6*d * (L*(12*d + l_seq) + V)
TPE = 50000 * 256 * 512
对齐实验
假设对齐实验中,直接使用了后面的学习率扫描得出的最优结果,并没有单独进行学习率扫描,因此这一步的成本计算比较简单:
def alignment() -> int:
return 4 * TPE * sum(M(d) for d in [1024,2048,4096])
# >>> f'{alignment():.3E}'
# '3.733E+20'
# >>> cost_of_run(alignment())[0]
# 888.81395400704
如果H100每运行1小时的花费以3美元计算,对齐实验的成本大致为888美元。
学习率
子问题:最佳评估损失(eval loss)实验
论文的表E1记录了6种模型规模下,所有可能的优化器×参数化方案×模型大小×实验设置的组合,分别进行基础学习率扫描,以获得最佳评估损失。总共包括如下几个实验变量:- 模型维度D∈3072,4096,6144,8192,12288,16384- 4种参数化方案- 3种优化器,其中SGD仅有5个实验设置,Adam和Adam+Param Scaling有7个实验设置假设这里的实验都是单独进行,没有从其他地方复制结果,因此如果全部运行一遍,有成本上限预估:
H = [1,2,4,6,8,12,16,20,24,32,48,64,96,128]
D = [h * 128 for h in H]
def table_e1() -> int:
sets_x_optims = 5 + 7 + 7
return 4 * sets_x_optims * TPE * sum(M(d) for d in D[-6:])
# >>> f'{table_e1():.3E}';cost_of_run(table_e1())
# '1.634E+23'
# (388955.9991064986, 16206.499962770775)
这部分的成本就接近40万美元,虽然仍属于可接受范围内,但对于大多数学术预算来说,已经算是非常昂贵了。表E1给出了最佳评估损失,但没有描述LR的扫描策略,每张图上的点数也不尽相同。由于没有得到论文作者的答复,我们也无法确定具体机制,因此假设每个最佳评估损失都经过了15次实验(目测发现,每条线的点数约为10~15)。
β参数
根据论文4.2节内容,学习率还涉及到两个超参数的选择:β和γ。如果仅有β参数,则被称为「LR+default」设置:这部分包括3×优化器,4×参数化,加上全局和单层(GlobalLR、Perlayer-fullalign)分别进行实验,以及未知的LR扫描数量:
def beta_only() -> int:
return 3*4*2*PpL * TPE * sum(M(d) for d in D)
# 7.988E+23 (1902022.3291813303, 79250.93038255542)
从公式就可以看出,成本和下文的epsilon实验类似,都是200万美元。
γ参数
相比β参数的实验,这部分有两个细节差异。首先,除了GlobalLR、Perlayer-fullalign两种设置外,还需要加上Perlayer-noalign设置。其次,仅针对d=1024=b,进行3D超参数搜索(γ_1,γ_h,γ_L+1),因此有额外的800次运行。两者结合后的计算公式为:这部分的预估成本与Adam的epsilon热力图实验接近,约为320万美元。
def gamma_expts() -> int:
return 36*TPE * (800*M(1024) + PpL*sum(M(d) for d in D))
# gamma_expts 1.354E+24 (3224397.534237257, 134349.8972598857)
Adam优化器的Epsilon参数
论文4.3节所述的Epsilon参数实验是计算量的大头。根据上面的推断,每次找到最佳评估损失时都尝试过15个不同的学习率(points per line),那么图6所示的epsilon参数变化图耗费的计算量为:计算结果透露出一种简洁的昂贵,也就是200万美元的账单而已。
PpL = 15 # unprincipled estimate
def eps_variants() -> int:
return 4 * 6 * PpL * TPE * sum(M(d) for d in D)
'''
>>> f'{eps_variants():.3E}';cost_of_run(eps_variants())
'7.988E+23'
(1902022.3291813303, 79250.93038255542)
'''
除了图6左侧的折线图,还有附录F热力图的结果。假设每个方块值都是经过13次学习率扫描后得到的结果,这部分计算量则为:结果发现,仅仅要得到这8张热力图,成本就是320万美元。而且,由于我们将LR扫描数量建模为常数13,这个数字可能低于实际成本。
def eps_heatmaps() -> int:
# eps-type * eps-val * parameterizations * LR range * ...
return 2 * 6 * 4 * 13 * TPE * sum(M(d) for d in D[-6:])
'''
>>> f'{eps_heatmaps():.3E}';cost_of_run(eps_heatmaps())
'1.341E+24'
(3193533.466348094, 133063.89443117057)
'''
权重衰减
权重衰减实验(附录G)比较好理解,对4×参数化方案以及所有参数进行一次基本的LR扫描:比epsilon实验便宜不少,也就是湾区工程师一年的工资——31.7万美元。
def weight_decay() -> int:
return 4 * PpL * TPE * sum(M(d) for d in D)
'''
>>> f'{weight_decay():.3E}'; cost_of_run(weight_decay())
'1.331E+23'
(317003.7215302217, 13208.488397092571)
'''
Adafactor优化器
这部分实验在附录C3中有详细描述,是为了检验Adafactor和Adam+parameter scaling是否有相似的宽度缩放机制。共有2×4张图,其中每个优化器收集11个数据点,因此计算公式为:账单上再加18.8万美元。
def adafactor() -> int:
return 2*2*4*PpL*TPE*sum(M(d) for d in D[:11])
'''
>>> f'{adafactor():.3E}'; cost_of_run(adafactor())
'7.918E+22'
(188532.80765144504, 7855.533652143543)
'''
计算最优化
论文尝试改变注意力头H的数量,希望找到计算最优化的设置,但其中涉及步长和数据集的改变,因此这部分不使用公式描述,计算代码如下:
def P(d: int, L=8, V=32101) -> int:
return 2 * d * (6*L*d + V)
def compute_optimal():
indices_50k = (14, 14, 12)
return 4*PpL*sum([
TPE * sum(sum( M(d) for d in D[:i] ) for i in indices_50k),
20 * sum(P(d)*M(d) for d in D[:11]) *3,
])
# compute_optim 7.518E+23 (1790104.1799513847, 74587.67416464102)
总结
将以上各部分实验的算力和成本汇总在一起:
alignment 3.733E+20 (888.81395400704, 37.033914750293334)
table_e1 1.634E+23 (388955.9991064986, 16206.499962770775)
eps_variants 7.988E+23 (1902022.3291813303, 79250.93038255542)
eps_heatmaps 1.341E+24 (3193533.466348094, 133063.89443117057)
beta_only 7.988E+23 (1902022.3291813303, 79250.93038255542)
gamma_expts 1.354E+24 (3224397.534237257, 134349.8972598857)
weight_decay 1.331E+23 (317003.7215302217, 13208.488397092571)
adafactor 7.918E+22 (188532.80765144504, 7855.533652143543)
compute_optim 7.518E+23 (1790104.1799513847, 74587.67416464102)
结果发现,整篇论文的运算量为5.42e24 FLOPS。这个数字仅仅是Llama 3训练计算量的15%,如果在10万卡H100集群上运行,只需要2天时间即可完成所有实验。
total_flops=5.421E+24
rental price: US$12.9M
h100 node months required: 746.9595590938408
(sanity check) D=[128, 256, 512, 768, 1024, 1536, 2048, 2560, 3072, 4096, 6144, 8192, 12288, 16384]
(sanity check) model sizes: ['0.00979B', '0.0227B', '0.058B', '0.106B', '0.166B', '0.325B', '0.534B', '0.794B', '1.1B', '1.87B', '4.02B', '6.97B', '15.3B', '26.8B']
(sanity check) M/6P: ['63.4%', '68.5%', '75.3%', '79.7%', '82.8%', '86.8%', '89.3%', '91.0%', '92.2%', '93.9%', '95.7%', '96.7%', '97.7%', '98.3%']
然而,如果不从LLM预训练的标准来衡量,仅把DeepMind的这篇论文看做一篇学术研究,这个计算量就显得相当奢侈了。如果实验室仅有10张H100,就根本不可能进行这个量级的研究。有100张H100的大型实验室,或许能用几年时间跑完以上所有实验。
参考资料:
https://152334h.github.io/blog/scaling-exponents/
https://news.ycombinator.com/item?id=41107721https://arxiv.org/abs/2407.05872
贾扬清:大模型尺寸正在重走CNN的老路;马斯克:在特斯拉也是这样
Transformer大模型尺寸变化,正在重走CNN的老路!
看到大家都被LLaMA 3.1吸引了注意力,贾扬清发出如此感慨。
拿大模型尺寸的发展,和CNN的发展作对比,就能发现一个明显的趋势和现象:
在ImageNet时代,研究人员和技术从业者见证了参数规模的快速增长,然后又开始转向更小、更高效的模型。
听起来,是不是和GPT哐哐往上卷模型参数,业界普遍认同Scaling Law,然后出现GPT-4o mini、苹果DCLM-7B、谷歌Gemma 2B如出一辙?
贾扬清笑称,“这是前大模型时代的事儿,很多人可能都不咋记得了:)”。
而且,贾扬清不是唯一一个感知到这一点的人,AI大神卡帕西也这么觉得:
大模型尺寸的竞争正在加剧……但是卷的方向反着来了!
模型必须先追求“更大”,然后才能追求“更小”,因为我们需要这个过程,帮咱把训练数据重构成理想的、合成的格式。
他甚至拍着胸脯打赌,表示我们一定能看到又好、又能可靠地思考的模型。
而且是参数规模很小很小的那种。
连马斯克都在卡帕西的评论区连连称是:
以上,大概可以称之为“大佬所见略同”。
展开说说
贾扬清的感慨,要从只在最强王座上短暂待了一天的LLaMA 3.1说起。
那是首次实现“最强开源模型=最强模型”,不出意外,万众瞩目。
However,贾扬清在这个时候提出了一个观点:
“但我认为, 行业会因小型垂直模型而真正蓬勃发展。”
至于啥是小型垂直模型,贾扬清也说得很清楚,比如以Patrouns AI的Iynx(该公司的幻觉检测模型,在幻觉任务上超过GPT-4o)为代表的那些很棒的中小模型。
贾扬清表示,就个人喜好而言,他本人是非常喜欢千亿参数模型的。
但现实情况里,他观察留意到,7B-70B参数规模之间的大模型,大家用起来更顺手:
它们更容易托管,不需要巨大的流量即可盈利;
只要提出明确的问题,就能得到质量还不错的输出——与和之前的一些看法相反。
与此同时,他听说OpenAI最新的、速度很快的模型也开始变得比“最先进的”大模型尺寸更小。
“如果我的理解是正确的,那么这绝对表明了行业趋势。”贾扬清直接表明了自己的观点,“即在现实世界中,使用适用的、具有成本效益、且仍然强大的模型。”
于是乎,贾扬清简单梳理了CNN的发展历程。
首先,是CNN的崛起时代。
以AlexNet(2012)为起点,开启了大约三年的模型规模增长时期。
2014年出现的VGGNet就是一个性能和规模都非常强大的模型。
其次,是缩小规模时期。
2015年,GoogleNet把模型大小从“GB”缩小到了“MB”级别,即缩小了100倍;但模型性能并没有因此骤减,反而保持了不错的性能。
遵循类似趋势的还有2015年面世的SqueezeNet模型等。
然后的一段时间,发展重点在追求平衡。
后续研究,如ResNet(2015)、ResNeXT(2016)等,都保持了一个适中的模型规模。
值得注意的是,模型规模的控制并没有带来计算量的减少——其实,大伙儿都愿意投入更多的计算资源,寻求一种“同等参数但更高效”的状态。
紧接着就是CNN在端侧起舞的一段时期。
举个例子,MobileNet是谷歌在2017年推出的一项有趣的工作。
有趣就有趣在它占用的资源超级少,但是性能却非常优异。
就在上周,还有人跟贾扬清提到:“Wow~我们现在还在用MobileNet,因为它可以在设备上运行,而且在出色的特征嵌入泛化(Feature Embedding Generality)。”
最后,贾扬清借用了来源于Ghimire等人的《A Survey on Efficient Convolutional Neural Networks and Hardware Acceleration》里的一张图:
并再一次发出自己的疑问:
大模型尺寸,会遵循与CNN时代相同的趋势来发展吗?
网友怎么看?
其实GPT-4o mini这样走在大模型发展道路上“不大反小”的例子不在少数。
当上述几位表达出这样的观点后,立马有人点头如捣蒜,还拿出了一些别的类似例子,证明他们看到了相同的趋势。
有人立马跟上:
我这儿有个新的正面例子!Gemma-2就是把27B参数大小的模型知识蒸馏成更小的版本。
还有网友表示,开发更大的模型,意味着能给后续几代更小、更垂直的模型的训练“上强度”。
这个迭代过程最终会产生所谓的“完美训练集”。
这样一来,较小的大模型在特定领域,能与现在参数巨大的大模型一样聪明,甚至更聪明。
一言以蔽之,模型必须先变大,然后才能变小。
大多数讨论此观点的人,还是对这个趋势比较认同,有人直言“这是一件好事,比‘我的模型比你的模型大’参数竞赛更实用和有用。”
但是,当然了!
翻遍网络评论区,也有人发出不同的声音。
比如下面这位朋友就在贾扬清推文底下留言:
Mistral Large(背后公司Mistral AI)、LLaMA 3.1(背后公司Meta)和OpenAI,持有最强竞争力模型的公司,目前可能都正在训练更大的模型。
我没发现有“更小型号模型搞定技术突破”的趋势哟。
面对这个问题,贾扬清倒也及时回复了。
他是这么说的:“没错!我说大模型尺寸可能在走CNN的老路,绝对不意味着号召大家停止训练更大的模型。”
他进一步解释道,这么说的本意是,随着技术(包括CNN和大模型)落地实践越来越广,大家已经开始越来越关注性价比更高的模型了。”
所以,或许更高效的小·大模型,能够重新定义AI的“智能”,挑战“越大越好”的假设。
你赞同这个观点不?
参考链接:
[1]https://x.com/jiayq/status/1818703217263624385
[2]https://x.com/fun000001/status/1818791560697594310
[3]https://www.patronus.ai/
[4]https://twitter.com/karpathy/status/1814038096218083497
Transformer作者回流谷歌,Character.AI创始团队被「收购」,只要人不要公司
AI 初创者的归宿还是大厂?
一觉醒来,生成式 AI 的「吃鸡大赛」再次缩圈了。
初创公司 Character.AI 周五宣布已与谷歌签署协议,谷歌将获得 Character.AI 的大型语言模型(LLM)技术的非独家许可。
谷歌还宣布重新雇佣 Noam Shazeer 和 Daniel De Freitas。其中,Noam Shazeer 是 Character.AI 的创始人、CEO,也是 Transformer 论文作者之一,他曾在谷歌任首席软件工程师。而 Daniel De Freitas 是 Character.AI 的总裁,曾在谷歌担任高级软件工程师。
Daniel de Freitas(左)和 Noam Shazeer。图源:https://www.bizjournals.com/sanjose/inno/stories/news/2023/03/24/q-a-interview-with-characterai-founders.html
2021 年,Noam Shazeer 和 Daniel De Freitas 因对谷歌这家搜索巨头的官僚主义感到失望而离开谷歌,并在 2022 年创办了 Character.AI。而现在,他们又将与约 30 人的研究团队一起回到 Google DeepMind 工作。
谷歌发言人在一封电子邮件中表示:「我们特别高兴地欢迎 Noam 回来,他是机器学习领域的杰出研究员。」
Character.AI 剩余的大约 140 名员工将留下来,面临着下一步抉择。Character.AI 官方发布了一份公开信,内容如下:
2022 年,我们创立了 Character.AI,旨在为全球用户带来个性化的超级智能。在过去的两年里,我们在这一目标上取得了巨大进展。我们构建了越来越智能的模型,推出了与虚拟角色对话的沉浸式新功能,并迅速发展到服务数百万用户,成为他们日常生活的一部分。
当 Noam 和 Daniel 创办 Character.AI 时,我们实现个性化超级智能的目标需要全栈式方法。我们必须对模型进行预训练和后训练,以保证用户在 Character.AI 上能够获得独特的体验,并构建一个能够使全球用户共同使用的平台。然而,在过去的两年里,技术环境发生了变化 —— 现在有更多的预训练模型可用。鉴于这些变化,我们认为联合利用第三方大型语言模型(LLM)和我们自己的模型将具有优势。这使我们能够投入更多资源用于后训练和为不断增长的用户群体创造新的产品体验。
我们很高兴地宣布,我们已与谷歌达成协议,这将使我们能够加速进步。根据该协议,Character.AI 将为谷歌提供现有 LLM 技术的非独家许可。这项协议将为 Character.AI 提供更多资金,以继续增长,并专注于为全球用户构建个性化 AI 产品。
Noam、Daniel 和我们研究团队的部分成员也将加入谷歌。Character.AI 的大多数才华横溢的团队成员将继续留在公司,继续构建 Character.AI 产品并服务于我们不断增长的用户群。
Character.AI 的总法律顾问 Dominic Perella 已担任临时首席执行官一职。Perella 之前是 Snap Inc. 的长期高管,自 2023 年中期以来一直是 Character.AI 核心领导团队的一员。这些变动将立即生效。
在我们进入下一个增长阶段时,我们将继续投资于我们的后训练能力,灵活使用我们自己的或外部可用的大语言模型。我们对 Character.AI 的未来充满期待,并致力于通过创新型产品来服务我们的用户。
我们对 Noam、Daniel 和其他团队使 Character.AI 从梦想化为现实表示无比感激。我们期待在他们已有贡献的基础上,Character.AI 在下一个增长阶段继续航行。
虽然在技术层面上,公司的股份没有易手,但谷歌会以 25 亿美元的估值向 Character.AI 的投资者支付其股权价值。
据消息人士透露,Character.AI 的员工也将根据其已归属的股份按该估值获得现金,并且随着其现有股票转移归属,他们还将获得偿付。
Character.AI 此前从包括 Andreessen Horowitz 在内的投资者处筹集了 1.93 亿美元的风险资本,其最后一次已知估值为 10 亿美元。该公司也曾在谈判中表示希望从谷歌筹集数亿美元资金。
这一协议类似于微软、亚马逊等公司在过去几个月与初创公司达成的协议。这些协议正受到监管机构的审查。科技巨头正投入数十亿美元来增强其 AI 基础设施,并从初创公司中招聘最优秀的研究人员。
今年 3 月,微软支付 6.5 亿美元引入 AI 初创公司 Inflection 的联合创始人及数十名员工。前 Inflection 首席执行官 Mustafa Suleyman,已成为微软执行副总裁和新成立的微软 AI 组织的首席执行官。
前 Inflection 首席执行官 Mustafa Suleyman。
与之相似,6 月,亚马逊则从另一家 AI 初创公司 Adept 中招聘了多名联合创始人和员工。
这场从 AI 初创公司招募人才以扩展羽翼,为业务再开发赋能的战略部署已初见端倪。大型科技企业对 AI 初创公司的「蚕食」可能才刚刚开始。
参考链接:
https://www.theverge.com/2024/8/2/24212348/google-hires-character-ai-noam-shazeer
https://www.reuters.com/technology/artificial-intelligence/google-hires-characterai-cofounders-licenses-its-models-information-reports-2024-08-02/
https://blog.character.ai/our-next-phase-of-growth/
阿里「轨迹可控版Sora」,告别「抽卡」,让视频生成更符合物理规律
你规定路线,Tora 来生成相应轨迹的视频。
目前,扩散模型能够生成多样化且高质量的图像或视频。此前,视频扩散模型采用 U-Net 架构 ,主要侧重于合成有限时长(通常约为两秒)的视频,并且分辨率和纵横比受到固定限制。
Sora 的出现打破了这一限制,其采用 Diffusion Transformer(DiT)架构,不仅擅长制作 10 到 60 秒的高质量视频,而且还因其生成不同分辨率、各种纵横比、且遵守实际物理定律的能力而脱颖而出。
可以说 Sora 是 DiT 架构最有利的证明,然而,基于 Transformer 的扩散模型在有效生成可控动作视频方面还未被充分探索。
针对这一问题,来自阿里的研究者提出了 Tora,这是第一个面向轨迹的 DiT 架构,它将文本、视觉和轨迹条件同时集成在一起以生成视频。
Tora 的设计与 DiT 的可扩展性无缝契合,允许精确控制具有不同持续时间、宽高比和分辨率的视频内容。大量实验证明,Tora 在实现高运动保真度方面表现出色,同时还能细致模拟物理世界的运动。
论文地址:https://arxiv.org/pdf/2407.21705
论文主页:https://ali-videoai.github.io/tora_video/
论文标题:Tora: Trajectory-oriented Diffusion Transformer for Video Generation
一艘老式的木制帆船沿着规定好的路线在迷雾笼罩的河流上平稳地滑行,周围是茂密的绿色森林。
一条鲫鱼优雅地游过火星的红色岩石表面,鱼的轨迹向左,火星的轨迹向右。
热气球沿着不同的轨迹升入夜空,一个沿着规定的斜线,另一个沿着有弯度的轨迹。
两只可爱的小猫并排走在宁静的金色沙滩上。
气泡沿着轨迹轻轻地漂浮在盛开的野花中。
枫叶在清澈的湖面上颤动,映照着秋天的森林。
山间的瀑布倾泻而下,主题、背景的运动都可以按照不同的路线运动。
在 Tora 与其他方法的比较中,可以看出 Tora 生成的视频流畅度更高,更遵循轨迹,且物体不会存在变形的问题,保真度更好。
方法介绍
Tora 采用 OpenSora 作为其 DiT 架构的基础模型,包含一个轨迹提取器 (TE,Trajectory Extractor)、时空 DiT(Spatial-Temporal DiT )和一个运动引导融合器 (MGF,Motion-guidance Fuser) 。TE 使用 3D 视频压缩网络将任意轨迹编码为分层时空运动 patch。MGF 将运动 patch 集成到 DiT 块中,以生成遵循轨迹的一致视频。图 3 概述了 Tora 的工作流程。
时空 DiT(ST-DiT)
ST-DiT 架构包含两种不同的块类型:空间 DiT 块 (S-DiT-B) 和时间 DiT 块 (T-DiT-B),它们交替排列。S-DiT-B 包含两个注意力层,每个层按顺序执行空间自注意力 (SSA) 和交叉注意力,后面跟着一个逐点前馈层,用于连接相邻的 T-DiT-B 块。T-DiT-B 仅通过用时间自注意力 (TSA) 替换 SSA 来修改此架构,从而保持架构一致性。在每个块中,输入在经过规范化后,通过跳跃连接连接回块的输出。通过利用处理可变长度序列的能力,去噪 ST-DiT 可以处理可变持续时间的视频。
轨迹提取器
轨迹已被证明是一种更加用户友好的方法来控制生成视频的运动。然而,DiT 模型采用视频自编码器和 patch 化过程将视频转换为视频 patch。在这里,每个 patch 都是跨多个帧导出,因此直接采用帧间偏移是不合适的。为了解决这个问题,本文提出的 TE 将轨迹转换为运动 patch,运动 patch 与视频 patch 位于相同的潜在空间。
运动引导融合器
为了将基于 DiT 的视频生成与轨迹结合起来,本文探索了三种融合架构变体,将运动 patch 注入每个 ST-DiT 块。这些设计如图 4 所示。
实验结果
在实现细节上,研究者基于 OpenSora v1.2 权重来训练 Tora。训练视频的分辨率由 144p 到 720p 不等。为了平衡训练 FLOP 以及每次迭代不同分辨率和帧数所需的内存,研究者相应地将批大小从 1 调整到 25。
至于训练基础设施,研究者使用了 4 块英伟达 A100 和 Adam 优化器,学习率为 2 × 10^−5。
研究者将 Tora 与流行的运动指导视频生成方法进行了比较。评估中使用了三种设置,分别为 16、64 和 128 帧,所有设置都是 512×512 的分辨率。
结果如下表 1 所示,在 U-Net 方法常用的 16 帧设置下,MotionCtrl 和 DragNUWA 能够更好地与所提供的轨迹实现对齐,但仍弱于 Tora。随着帧数增加,U-Net 方法在某些帧中出现明显偏差,并且错位误差传播会导致后续序列中出现变形、运动模糊或物体消失。
相比之下,得益于集成了 Transformer 的缩放能力,Tora 对帧数变化表现出很高的稳健性。Tora 产生的运动更加流畅,且更符合物理世界。对于 128 帧测试设置下的评估,Tora 的轨迹精度达到其他方法的 3 到 5 倍,展现出了卓越的运动控制能力。
在下图 5 中,研究者对不同分辨率和持续时长的轨迹误差进行分析。结果显示,不同于 U-Net 随时间推移出现明显的轨迹误差,Tora 的轨迹误差随时间推移出现渐进增加。这与 DiT 模型中视频质量随时间增加而下降相一致。Tora 在更长的时间下保持了有效的轨迹控制。
下图 6 展示了 Tora 与主流运动控制方法的比较分析,在包含两人共同运动的场景中,所有方法都能生成相对准确的运动轨迹。不过,Tora 的视觉质量更好,这要归功于更长序列帧的使用,有助于实现更平滑的运动轨迹和更逼真的背景渲染。
可以看到,在 Tora 生成的自行车场景中,人的双腿表现出逼真的踩踏动作,而 DragNUWA 的双腿几乎水平漂浮,违反了物理真实性。此外,DragNUWA 和 MotionCtrl 在视频结尾处都出现了严重的运动模糊。
在另一个生成灯笼的场景中,DragNUWA 随着所提供轨迹的持续升降出现了严重的变形。MotionCtrl 的轨迹虽然相对准确,但生成的视频与两个灯笼的描述不相符。Tora 不仅严格地遵循了轨迹,而且最大程度地减少了物体变形,确保了更高保真度的动作表示。
小技巧大功效,「仅阅读两次提示」让循环语言模型超越Transformer++
在当前 AI 领域,大语言模型采用的主流架构是 Transformer。不过,随着 RWKV、Mamba 等架构的陆续问世,出现了一个很明显的趋势:在语言建模困惑度方面与 Transformer 较量的循环大语言模型正在快速进入人们的视线。
令人兴奋的是,这些架构在推理期间使用了恒定量的内存。不过,受制于有限的内存,循环语言模型(LM)无法记忆并使用长上下文中的所有信息,这导致了上下文学习(in-context learning,ICL)质量的不佳。因此,获得高效大语言模型的关键挑战在于选择存储或者丢弃哪些信息。
在最近的论文《Just read twice: closing the recall gap for recurrent language models》中,来自斯坦福大学、布法罗大学的研究者通过简单观察发现,数据在推理期间涌入循环语言模型的排序极大地影响了在有限内存中预测存储哪些信息的难度。
我们假设根据文档 D(比如伽利略・伽利莱的详细维基百科)来提问:伽利略是什么时候搬到的佛罗伦萨?这时,如果提示遵循了 [Q, D] 的排序,则模型只需要记住文档 D 中的一个事实即可。相反,如果提示遵循了 [D, Q] 的排序,则模型需要记住所有事实。如下图 1(左)所示。
因此,本文首先从理论上形式化了数据排序如何影响内存需求,然后提出两种方法来减轻对数据排序的依赖,分别是 Just-read-twice(JRT)提示策略和 JRT 循环架构。本文主要分为以下几个部分展开:
理解数据排序的作用。研究者得出的第一个洞见是:记忆问题的 hardness 要降低到与设置剥离(set disjointness,SD)相同,这是通信复杂度理论中持续数十年的最典型问题。SD 要求一种流算法(比如循环模型)来决定上下文中提供的输入集是否剥离:
理论分析和实验结果表明,第一个集 | A | 掌控了求解 SD 所需的内存。因果模型需要存储 A 中的所有元素以与 B 中的元素进行比较。这表明了,使用上下文中的「正确数据排序」(如将最小 min (|A|, |B|) 的集放在首位)将有助于内存受限的模型。更进一步,观察到上下文非因果逻辑的模型可在空间最小的 (|A|, |B|) 中求解 SD,而无需考虑数据排序。
其次是利用「正确的」排序。本文提出了一种非常简单的 JRT-Prompt 策略,在模型生成答案之前在上下文中将信息重复多次(如上图 1 右所示)。在第二以及更多轮次中,语言模型在决定存储哪些信息时要以完整的上下文为条件,从而有效避免了将数据排序「归正」的问题。
结果表明,JRT-Prompt 在 16 个已有循环语言模型和 6 项 ICL 任务上,实现了平均 11.0 ± 1.3 百分点的提升,而吞吐量是 FlashAttention-2(长度 32k、批大小 16)的 11.9 倍。JRT-Prompt 虽然增加了上下文长度,但渐进来看仍然比注意力更加地计算和内存高效。
超越因果模型。本文提出了 JRT-RNN,它的灵感来源于简单的 Prefix-LM 编码器解码器架构设计。大多数的上下文学习输入包含两部分内容,分别是输入的提示(上下文、指令)和作为输出的模型生成文本。在 Prefix-LM 架构中,LM 并没有遵循因果逻辑地处理提示区域,而对输出进行了因果解码,其中在因果区域仅使用了标准的下一个 token 预测损失,以及非因果区域上的损失。
不过遗憾的是,此前 Prefix-LM 模型的训练方法取得的成功有限,并使用了低效的 Transformer 主干。因此本文通过一些简单的改变来提高质量和效率,包括改进训练损失并使用称之为「Prefix Linear Attention,PLA」 的线性注意力公式。研究者发现,使用他们的 IO 感知实现,JRT-RNN 在 360m 和 1.3b 参数设置下,分别可以提供 13.7 和 6.9 百分点的平均质量改进,吞吐量是 FA2 的 19.2 倍。
论文地址:https://arxiv.org/pdf/2407.05483
项目主页:https://github.com/HazyResearch/prefix-linear-attention
JRT-Prompt 方法概览
上下文学习任务以 (C, Q, Y) 作为输入,其中 C 为一些上下文来源(如文档或代码存储库),Q 为给定上下文时对模型的一些问题或请求,Y 为答案。对于使用自回归 LM A 的标准上下文学习,研究者输入 C 和 Q,并根据正确的完成情况 Y 来评估生成的输出 Yˆ = A (C, Q)。
JRT-Prompt 是一种极其简单的方法,在提示模型输出答案之前会在上下文中重复提示中的信息(如问题和文档),例如下图 1 右的 Yˆ = A (C, Q, C, Q)。因此,在上下文第二次出现时,模型根据完整的上下文来决定存储哪些信息。
此外,JRT-Prompt 可以与现成的 LLM 一起使用。研究者在零样本提示下,在一系列记忆密集型上下文任务上评估了以下 LM:
Based 预训练 LM,参数规模为 1.3B,在 Pile 的 10 − 50B 个 token 上进行训练;
Mamba 预训练的 LM,参数规模为 130M、370M、1.4B 和 2.8B,在 Pile 的 300B 个 token 上进行训练;
Gated Linear Attention 预训练的 LM,参数规模为 1.3B 和 2.7B,在 SlimPajama 数据集的 100B 个 token 上进行训练;
Mamba-2 预训练的 LM,参数规模为 130M、370M、1.3B 和 2.7B,在 Pile 的 300B 个 token 上进行训练。
结果如下表 1 所示,通过增加状态(state)大小,研究者发现 JRT-Prompt 方法在各个模型和任务上平均带来了 11.0 ± 1.3 百分点的性能提升,利用该方法的 Based 模型平均优于利用标准提示的 Transformer 模型。
他们还发现,JRT-Prompt 可以使 Transformer 模型受益,并且该方法在一些任务上(附录 2)比少样本学习更加有效。值得注意的是,Springer 等人在论文《Repetition improves language model embeddings》中提出使用自回归 Transformer 模型来重复上下文以实现生成嵌入的目的,本文的研究结果也类似。研究者专注于亚二次架构和上下文学习任务。
JRT-Prompt 虽然由于重复而增加了上下文长度,但是其使用的亚二次循环架构仍比使用二次 Transformer 模型更高效。研究者发现,在序列长度 N = 32768、批大小为 16 时,使用 JRT-Prompt(序列长度 2N)在英伟达 H100 上提供的吞吐量是 FlashAttention-2(序列长度 N)的 11.9 倍。
JRT-RNN:编码器 – 解码器循环架构
JRT-RNN 的灵感来自于 Prefix-LMs,但侧重于扩展质量 – 效率权衡空间的帕累托边界(Pareto frontier)。为了提高质量,JRT-RNN 在编码器端使用了单独的 k_e 和 v_e 映射,在解码器端使用了 k_d 和 v_d 映射。虽然 Prefix LM 模型对编码器和解码器区域使用了共享映射权重,但研究者发现使用两组映射可以提高质量。
为了提高效率,JRT-RNN 为编码器使用了非因果线性注意力,而为解码器使用标准因果线性注意力。研究者称为 Prefix Linear Attention(PLA)(图 1 右),公式如下:
JRT-RNN 训练目标。Prefix LMs 通常不计算非因果区域的损失,而 JRT-RNN 将下一个 token 预测与掩码语言建模(MLM)目标进行了结合。并且对于添加的 MLM 目标,研究者用一个 [MASK] token 替换了来自编码器区域 {u_1, …, u_M} 的比例为 P 的 tokens,并在预测原始 token 时测量了交叉熵损失。
损失如下:
实验结果
在实验中,研究者评估了 JRT-RNN 在以下三个指标上的质量和效率:
上下文学习质量
整体语言建模
生成
上下文学习质量
如下表 2 所示,研究者发现,JRT-RNN 在参数为 360M(30B tokens)时比仅解码器的基线(Based)平均高出 13.7 个百分点,在参数为 1.3B(50B tokens)时平均高出 6.9 个百分点。
同时,JRT-RNN 在参数为 360M 和 1.3B 时与 Transformer++ 的差距分别缩小到了 0.5 个百分点和 1.9 个百分点之内。
在下表 3 中,研究者比较了当 prefill 长度 l 小于编码器长度 M 时,JRT-RNN 与同类推理策略的表现。
整体自然语言理解
根据以往研究,研究者进一步将困惑度分为了两组:联想记忆「AR slice」包括了被称为「AR hits」的 tokens,它们需要模型按照顺序执行记忆以正确地预测下一个 token;而「Other slice」包含剩余的 tokens(如记忆的知识)。
对于记忆频率,JRT-RNN 在「AR slice」表现出色。对于训练期间不常见的二元组(即不太可能在模型参数中被记住的),JRT-RNN 的困惑度相对于 Based 和 Mamba 这两个强大的因果循环基线有所改善。
对于记忆距离,在「AR slice」中,JRT-RNN 与仅解码器基线之间的差距随着上下文中重复二元组的增加而扩大。这也进一步证明了 JRT-RNN 可以帮助完成更长的上下文记忆任务。
非记忆频率。对于训练期间很少见到的二元组的非记忆「Other slice」,JRT-RNN 的困惑度比仅解码器的 LM 更差。这是意料之中的结果,因为 JRT-RNN 计算了仅解码器 LM 的 65% tokens 的损失。
我们预计这一差距会随着规模和训练时间的延长而缩小(随着二元语法频率的增加而增加)(图 3,左上角)。
生成吞吐量
生成可以分解为提示「prefill 处理」和解码「下一个 token 预测」两步。相较于标准的仅解码器循环模型,JRT-RNN 不会修改解码步骤,因此讨论重点在 prefill 阶段。
使用 Simran Arora 等人论文《Simple linear attention language models balance the recall-throughput tradeof》中提出的 Based CUDAn 内核,JRT-Prompt 在处理 prefill 时吞吐量分别是 FlashAttention-2 和 FLA Triton 内核的 11.9 和 13.7 倍,如下表 5 所示。
当研究者将批大小增加到 64 时,JRT-Prompt 吞吐量分别是 FlashAttention-2 和 FLA Triton 内核的 6.1 倍和 7.2 倍。
接下来他们扩展了 Based 内核以支持 JRT-RNN,并且证明了当将序列长度增加到 32768 时,吞吐量分别是 FlashAttention-2 和 FLA 的 19.2 倍和 22.0 倍。当将批大小增加到 64 时,JRT-RNN 分别又提供了 9.7 倍和 11.5 倍的吞吐量提升。JRT-RNN 所需的时间是 Based prefill 的 1.24 倍,比 JRT-Prompt 更加高效。
更多技术细节和实验结果请参阅原论文。
准确率达60.8%,浙大基于Transformer的化学逆合成预测模型,登Nature子刊
逆合成是药物发现和有机合成中的一项关键任务,AI 越来越多地用于加快这一过程。
现有 AI 方法性能不尽人意,多样性有限。在实践中,化学反应通常会引起局部分子变化,反应物和产物之间存在很大重叠。
受此启发,浙江大学侯廷军团队提出将单步逆合成预测重新定义为分子串编辑任务,迭代细化目标分子串以生成前体化合物。并提出了基于编辑的逆合成模型 EditRetro,该模型可以实现高质量和多样化的预测。
大量实验表明,模型在标准基准数据集 USPTO-50 K 上取得了出色的性能,top-1 准确率达到 60.8%。
结果表明,EditRetro 表现出良好的泛化能力和稳健性,凸显了其在 AI 驱动的化学合成规划领域的潜力。
相关研究以「Retrosynthesis prediction with an iterative string editing model」为题,于 7 月 30 日发布在《Nature Communications》上。
论文链接:https://www.nature.com/articles/s41467-024-50617-1
分子合成路径设计是有机合成的一项重要任务,对生物医学、制药和材料工业等各个领域都具有重要意义。
逆合成分析是开发合成路线最广泛使用的方法。它包括使用已建立的反应将分子迭代分解为更简单、更易于合成的前体。
近年来,AI 驱动的逆合成促进了对更复杂分子的探索,大大减少了设计合成实验所需的时间和精力。单步逆合成预测是逆合成规划的重要组成部分,目前已有几种基于深度学习的方法,且效果优异。这些方法大致可分为三类:基于模板的方法、无模板的方法和半基于模板的方法。
在此,研究人员专注于无模板逆合成预测。提出将问题重新定义为分子字符串编辑任务,并提出基于编辑的逆合成模型 EditRetro,可以实现高质量和多样化的预测。
图示:所提出的基于分子串的逆合成的 EditRetro 方法的示意图。(来源:论文)
该研究的核心概念是,通过使用 Levenshtein 操作的迭代编辑过程生成反应物字符串。该方法从基于编辑的序列生成模型的最新进展中汲取灵感。具体来说,采用了 EDITOR 中的操作,EDITOR 是一种基于编辑的 Transformer,专为神经机器翻译而设计。
EditRetro 模型包含三种编辑操作,即序列重新定位、占位符插入和标记插入,以生成反应物字符串。它由一个 Transformer 模型实现,该模型由一个编码器和三个解码器组成,两者都由堆叠的 Transformer 块组成。
重新定位解码器:重新定位操作包括基本的 token 编辑操作,例如保留、删除和重新排序。它可以与识别反应中心的过程进行比较,包括重新排序和删除原子或基团以获得合成子。
占位符解码器:占位符插入策略(分类器)预测要在相邻 token 之间插入的占位符数量。它在确定反应物的结构方面起着至关重要的作用,类似于识别从序列重新定位阶段获得的中间合成子中添加原子或基团的位置。
Token 解码器:token 插入策略(分类器),负责为每个占位符生成候选 token。这对于确定可用于合成目标产品的实际反应物至关重要。该过程可以看作是合成子完成的类似过程,结合占位符插入操作。
EditRetro 模型通过其非自回归解码器提高了生成效率。尽管结合了额外的解码器来迭代预测编辑操作,但 EditRetro 在每个解码器内并行执行编辑操作(即非自回归生成)。
当给定一个目标分子时,编码器将其字符串作为输入并生成相应的隐藏表示,然后将其用作解码器交叉注意模块的输入。类似地,解码器也在第一次迭代时将产品字符串作为输入。在每次解码迭代期间,三个解码器依次执行。
研究人员在公共基准数据集 USPTO-50K 和 USPTO-FULL 上评估了所提方法。大量实验结果表明,该方法在预测准确度方面优于其他基线,包括最先进的基于序列的方法 R-SMILES 和基于图编辑的方法 Graph2Edits。
EditRetro 在基准逆合成数据集 USPTO-50K 上进行的大量实验表明,EditRetro 取得了优越的性能,top-1 精确匹配准确率达到 60.8%。
此外,在更大的 USPTO-FULL 数据集上,其中 top-1 精确匹配准确率达到 52.2%,证明了其在更多样化和更具挑战性的化学反应中是有效的。
EditRetro 在 RoundTrip 和 MaxFrag 准确率方面也表现出优于基线方法的性能。这证明了 EditRetro 能够有效地学习化学规则。
此外,EditRetro 通过精心设计的推理模块提供多样化的预测。该模块结合了重新定位采样和序列增强,有助于生成多样化和变化的预测。重新定位采样对重新定位动作的预测进行采样,从而能够识别不同的反应位点。序列增强从不同的产品变体到反应物生成不同的编辑途径,从而提高了预测的准确性和多样性。这两种策略共同作用,提高了预测的准确性和多样性。
进一步的实验验证了 EditRetro 在一些更复杂的反应中的优越性,包括手性、开环和成环反应。结果证实了 EditRetro 在这些具有挑战性的场景中的优越性,证明了它能够处理不同类型的化学转化。
特别是,EditRetro 在四个多步骤逆合成规划场景中的成功应用证明了其实用性。
为了评估 EditRetro 在合成规划中的实用性,通过连续的逆合成预测设计完整的化学途径。研究人员选择了四种具有重要药用价值的目标化合物进行评估:非布司他、奥希替尼、GPX4 的变构激活剂和 DDR1 激酶抑制剂 INS015_037。
图示:EditRetro 的多步逆合成预测。(来源:论文)
所有四个示例都产生了与文献中报道的途径非常一致的逆合成途径,大多数预测排名在前两位。在考虑的 16 个单独步骤中,有 10 个步骤的预测准确率为 1。这些结果证明了 EditRetro 在实际逆合成预测中的实际潜力。
通过提供有价值的见解并促进高效合成路线的设计,该方法有望在逆合成规划领域得到实际应用。
八问八答搞懂Transformer内部运作原理
七年前,论文《Attention is all you need》提出了 transformer 架构,颠覆了整个深度学习领域。
如今,各家大模型都以 transformer 架构为基础,但 transformer 内部运作原理,仍是一个未解之谜。
去年,transformer 论文作者之一 Llion Jones 宣布创立人工智能公司 Sakana AI。近期,Sakana AI 发表了一篇题为《Transformer Layers as Painters》的论文,探究了预训练 transformer 中的信息流,并针对仅解码器和仅编码器冻结 transformer 模型进行了一系列实验。请注意,该研究没有对预训练模型进行任何类型的微调。
论文地址:https://arxiv.org/pdf/2407.09298v1
该研究认为 transformer 的内部机制(特别是中间层)可以类比画家作画流水线来理解。
作画流水线通常是将画布(输入)传递给一系列画家。有些画家擅长画鸟类,而另一些画家则擅长画轮子。每个画家从其下一级画家那里收到画布,然后其决定是否给画作添加一些笔画,或者只是将其传递给其上一级画家(使用剩余连接)。
这个类比并不是一个严格的理论,而是一个思考 transformer 层的工具。受这个类比的启发,该研究测试验证了一些假设:
各层是否都在使用相同的表征空间?
所有层都是必要的吗?
中间层都执行相同的功能吗?
层的顺序重要吗?
这些层可以并行运行吗?
对于某些任务来说,顺序是否比其他因素更重要?
循环有助于层并行吗?
哪些变体对模型性能影响最小?
该研究对预训练 LLM 进行了一系列实验,其中包括试验标准 transformer 执行策略的变化,并在仅解码器 (Llama) 和仅编码器 (BERT) 模型的各种基准上测量这些变化对模型性能的影响。
各层是否都在使用相同的表征空间?
为了回答不同层是否使用相同的表征空间,作者测试了 Transformer 在跳过特定层或切换相邻层的顺序时是否具有稳健性。例如,在 Llama2-7B 中,第 6 层通常期望接收第 5 层的输出。如果给第 6 层以第 4 层的输出,它是否会出现「灾难性」的行为?
在图 2 中,我们可以看到,除了第一层和最后几层之外,Llama2-7B 的各层对跳层或切换层都相当稳健。
该实验表明,中间层共享一个表征空间,且与「外围层」(第一层和最后几层)拥有不同的表征空间。为了进一步验证这一假设,作者效仿之前的研究,测量了基准中模型(Llama2-7B、Llama2-13B 和 BERT-Large)不同层的隐藏状态激活之间的平均余弦相似度。图 3 显示了所有中间层之间的一致性。
这表明该模型可能具有「开始」、「中间」和「结束」层的三个不同的表征空间。回答问题 1:是的,中间层似乎共享一个共同的表征空间。
所有层都是必要的吗?
为了进一步测试中间层的重定向空间是否真正共享(除了具有接近的余弦相似度之外),该研究尝试了「跳过层」,即将第 N 层的输出直接发送到第 N + M 层(其中 M > 1)的输入中,从而「跳过」M − 1 层,如图 1a 所示。该实验是为了看看第 N + M 层是否可以理解第 N 层的激活,尽管它仅根据从第 N + M − 1 层发来的输入进行训练。图 4 显示,Llama2-7B 和 BERT-Large 在许多基准测试上性能均出现适度下降。回答问题 2,是否所有层都是必要的:
不,至少可以删除一些中间层而不会发生灾难性故障。
中间层都执行相同的功能吗?
如果中间层都共享一个共同的表征空间,这是否意味着除此之外的中间层是多余的呢?为了测试这一点,研究者们重新运行了前一子节中的「跳过」实验,他们将中间层的权重替换为中心层的权重,有效地在被替换的每一层上循环 T – 2N + 1 次,其中 T 是总层数(Llama2-7B 为 32 层,BERT-Large 为 24 层)。
如图 5 所示,可以观察到,随着被替换层数的增加,模型在基准测试的得分迅速下降。从后文的图 11 看来,这种替换层的做法比研究者们尝试的其他方法都更糟糕。因此,研究者得出结论:中间层执行的是不同的功能,让中间层之间共享权重并不可行。
层的顺序重要吗?
之前的实验表明,中间层共享一个表示空间,但在该空间中负责不同的功能。下一个需要解决的问题是,这些功能的顺序有何意义。为了解决这个问题,研究者们设计了两组实验。首先,以与训练时相反的顺序来运行中间层。具体来说,取第 T – N 层的输出,将其输入到第 T – N – 1 层,然后将这一层的输出输入到第 T – N – 2 层,依此类推,一直到第 N 层,再将这一层的输出发送到后面的 T – N 层。在第二组实验中,研究者采用随机顺序运行中间层,并在 10 个种子值上取平均值。
图 6 和图 7 分别显示了反向和以随机顺序运行中间层的结果,模型在所有基础测试集中都显示出了逐渐下降的趋势。这也表明虽然层的顺序对模型来说有一定的重要性,但即使改变了顺序,这些层仍然能够发挥作用。
更有趣的是,随机打乱层的顺序比完全反过来效果更好。这可能是因为,随机打乱的顺序在某些方面保留了层之间的一些原有关系(即层 i 在层 j 之后,其中 i > j),而完全反过来则完全打破了这些关系。
这些层可以并行运行吗?
为了验证层本身存在比执行的顺序更重要,研究者们设计了一个实验,并行运行中间层,将它们的平均结果发送给最终的 N 层。
如图 8 所示,模型在所有基准测试中的表现均呈现了一种平缓下降趋势,然而,这种趋势并不适用于 GSM8K 中的数学应用题。
实验结果显示,大部分情况下这种方法都是有效的,只是一些复杂的数学题处理得不太好。这种并行处理方法相比直接跳过一些层,效果更好,但不如按反向顺序运行层的效果出色。基于此,研究者得出结论:并行运行层在一般情况下是可行的,但对于需要顺序逻辑理解的数学问题,这种方法可能不太适用。
对于某些任务来说,顺序是否比其他因素更重要?
对于大多数经过「改造」的模型,在面对抽象推理(ARC)或数学推理(GSM8K)基准测试时,它们往往显示出最陡峭的下降趋势。这一现象可能源于逐步推理任务对于模型层级顺序的敏感度远高于那些主要依赖语义理解的常识性任务。与那些仅通过理解语义便能完成的任务不同,推理任务要求模型同时把握结构与含义。这种观察与模型在单次处理过程中可能进行一定程度的顺序依赖性推理的假设相吻合。
研究者使用了一个比喻来说明:如果画一幅由许多不同元素组成的拼贴画,那么画的顺序可能不那么重要;但如果是要画一幅精确的建筑场景,那么每一笔的顺序就变得非常重要了。据此,研究者得出了结论:数学和推理任务对模型层的顺序具有更高的依赖性,而对于那些主要依赖语义理解的任务,顺序的影响则相对较小。
循环有助于层之间并行吗?
沿用上一节中画画的的比喻,当画家在画一幅画时,不是一开始就画所有东西,而是先画一部分,比如车身,然后再根据这部分来添加其他的东西,比如车轮。在 AI 模型中,层就是所谓的画家,处理信息就是在画画,如果先得到了正确的信息,也就先画出了所谓的车身,那么它们就能更好地完成自己的工作,为画作添加车轮。
对于 transformer 而言,当给予适当的输入时,层可能只在前向传播中做出贡献,并非通过残差连接「传递」输入。如果情况确实如此,那么迭代上一个实验中的并行层应该比单次执行并行层更能提高模型的性能。基于此,研究者通过将并行层的平均输出反馈到同一层中进行固定次数的迭代来测试这一点。
图 9 展示了将并行层循环 3 次的结果。循环并行 3 次的结果显著优于单次迭代(并行层)。起始层 N 设定为 15(针对 Llama2-7B 模型)或 11(针对 BERT 模型)时,即处于每种情况的极左端点,仅有单一的层级受到影响。在这种特定情况下,三次循环并行的效果等同于单纯地将中间层重复三次。与此同时,对于这一点上的并行层而言,其性能与完整模型无异。
研究者们还针对不同的迭代次数重复了相同的实验。图 10 展示了 Llama2-7B 的性能随并行化层数 M 和迭代次数的变化情况。每个 M 的最高性能迭代次数用红框标出。除了 M=29 和 M=31(几乎并行化所有层)外,最佳迭代次数大致与并行化层数成线性比例。因此,研究者得出的结论是:最佳迭代次数与并行化层数成正比。
如何调整层,对模型性能的影响最小?
最后,在图 11 中,研究者们将所有实验中对 Transformer 的「改造」进行了比较,在一个图表上显示了所有基准测试的中位数或平均性 。
中间重复 —— 用相同数量的中间层副本替换中间层 —— 表现最差, 很快就降到了随机基线的性能。相反,循环并行和随机层顺序的影响最小。因此,研究者得出的结论是:重复单一层的影响最严重。随机化层顺序和循环并行的影响最小。
这些实验整体上显示出平缓的性能下降,但研究者仍然不清楚为什么这些层在大多数扰动下还能保持一定的稳健性,这个问题还需在未来的研究中进一步探讨。
更多细节请参见原论文。
参考链接:https://arxiv.org/pdf/2407.09298v1