开云网址提高大模型计算效率!彩云科技推出大模型 DCFormer效率是 Transformer 两倍
2024-06-23 次
开云网址提高大模型计算效率!彩云科技推出大模型 DCFormer效率是 Transformer 两倍最近KAN突然爆火,让人们开始意识到大模型的计算效率至关重要,而提高大型模型生成tokens的速度同样至关重要。相对于不断加码GPU,改善Transformer模型架构的计算效率是一种更为长远有效的方法。近期,彩云科技的研究团队对Transformer计算最耗时的核心组件——多头注意力模块(MHA)进行了改进,推出动态组合多头注意力改进Transformer(DCFormer),将Transformer的计算性能提升有2倍之高。该论文已被ICML 2024接收为oral(今年oral的接收率仅为1.5%)。
集智俱乐部邀请论文一作、彩云科技首席科学家肖达老师在中深入解读了DCFormer的主要工作,介绍这项工作背后的动机和研究历程,并探讨模型结构研究的未来发展方向,本文是对此次读书会的整理。DCFomer的论文作者均来自彩云科技NLP算法组,核心成员出自集智俱乐部。彩云科技旗下产品包括彩云天气、彩云小梦、彩云小译等,致力于通过人工智能产品让生活更美好。团队目前正在招聘多个算法岗位,期待志同道合的朋友加入!详情见文末。
在大模型时代,无论是处理文本、图像、语音还是视频任务,大多以Transformer作为基础模型。前段时间一篇非常受欢迎的论文,名为KAN,它之所以受到广泛关注,是因为人们认为KAN可能会对Transformer进行改进,提升模型的效率。我们的工作同样旨在改进Transformer。 (参看《》)
我把Transformer引领的新一代模型范式比作搭积木的游戏,并将其分为三个派别:大力飞砖派、系统交互派和结构创新派。
首先是“大力飞砖派”,也就是我们常说的“大力出奇迹”的 Scaling Law。如果我们把Transformer block看作积木块,那么使用更多的积木块和数据就能得到更好的效果。这种方法的最大优势在于确定性强,因此我们称之为“靠谱青年”。
另一派——“系统交互派”认为,不应将模型视为孤立的点,而应将其置于系统中,让其进行更丰富的交互,例如通过强化学习训练,或提供思维链提示等。这种思路更接近人类的学习方式,因此我们称之为“文艺青年”。
第三个派别——“结构创新派”则更为直接,他们认为,既然前两派都把Transformer的积木块当作黑盒来使用,为什么不拆开看看里面究竟是什么呢?他们主张深入研究Transformer的工作原理,探讨如何改进甚至替换它。这一派的做法确实引发了一些质疑。
许多人会问:“真的有必要改进Transformer吗?这个结构不是已经足够好了吗?”这些疑问并非无的放矢,因为Transformer已经诞生七年了,无数实践似乎已经证明它的强大和通用性。并且有论文显示对Transformer的改进大多并不像原始Transformer那样普适有效和可扩展。
正是在这样的背景下,我们决定进行这项看似有些中二的工作——改进Transformer的结构。我们设计了一种名为动态组合多头注意力(DCMHA) 的新架构,用以替换Transformer中的多头注意力 (MHA) 机制,从而得到一个具有更高计算效率的模型。这意味着在达到相同效果的情况下,我们可以节省1.7到2倍的算力,或者在使用相同算力的情况下获得更好的效果。简而言之,这个新模型让我们能够以更少的算力训练出更优秀的模型。
本次分享围绕三个方面展开:第一部分是改进的动机,即我们是如何产生这个改进想法的;第二部分是DCFormer本身的设计和实现细节;第三部分是对DCFormer研究历程的回顾,包括我们的经验教训和对未来的展望,探讨结构改进和创新的可能方向。
我们都知道Transformer由两大模块组成,一块叫做多头注意力 (MHA) ,另一块叫做全连接层 (MLP) 。前一段比较很火的KAN架构本质上是一项关于MLP的工作,而DCFormer则更加关注多头注意力模块 (MHA) 。我们认为,多头注意力的重要性在于:Transformers中跨token的信息流动都是在注意力模块中完成的,而MLP模块中则没有这种信息流动。因此,很显然,所有关于上下文的信息的整合处理以及算法的实现都要依靠Attention注意力。
我们的第一步工作是研究注意力模块里每个注意力头是如何工作的。在这里,我们需要引入一个回路 (Transformer circuits) 的概念。每个注意力头有四个权重矩阵:Wq、Wk、Wo、Wv,当它们两两进行配对时,便会形成了两个信息回路——QK和OV。QK回路决定了从当前位置关注上下文中的哪些位置 (如下图紫色部分所示) 。OV回路决定了从关注到的位置取回什么信息/属性写入当前位置的残差流,进而预测下一个token (如下图黄色部分所示) 。
• OV/变换回路:决定从关注到的位置取回什么信息/属性写入当前位置的残差流,进而预测下一个token
图中展示了一个经典案例——我们可以称之为“归纳头” (induction head) 。对于致力于模型可解释性的研究者来说,理解这种模型通过内部回路完成任务的概念至关重要。
现在,让我们通过一些具体的例子来探讨这种组合方式是如何工作的,以及它可能面临的问题。以中文句子“小明有个钢笔 ... 小明有个 - 钢笔”为例,模型会使用一个类似induction head的机制来确定如何从前面的内容中提取信息。模型会识别到“小明”作为主语,并与宾语“钢笔”有关联。通过识别前文相同的主语,模型用一个QK回路能够找到相关的信息——这里是“钢笔”,并用一个OV回路直接复制了“钢笔”这个信息到残差流用于输出预测。
接下来,我们再看两个例子:“苹果是一种 - 水果。模型使用QK来关注前面的词汇“苹果”,并通过OV回路,来提取“苹果”的类别信息,即“水果”。
另一个例子是:“钢笔用来 - 写字。”这里,模型使用QK来识别“钢笔”,并通过OV来提取关于“钢笔”的用途信息,即它是用来写字的。这样,模型就能够利用这些信息来进行更准确的预测。
我们注意到,QK回路的查找 (从哪拿) 和OV回路取属性 (拿什么) 本来是独立的两件事,理应可以分别指定并按需自由组合,MHA硬把它们放到一个注意力头的QKOV里“捆绑销售”,限制了灵活性和表达能力。假设有个模型存在注意力头A、B、C其QK和OV回路能够完成上面的例子,那换成“小明有一支钢笔...小明有一种-文具” (head A 关注相同主语的宾语的QK + head B 取类别属性的OV) 或“小红有一支钢笔...小红能用它-写字” (head A 关注相同主语的宾语的 QK + head C 取用途属性的OV) ,需要交叉组合现有注意力头的QK和OV回路,模型就可能“转不过弯儿”了。这就像是在肯德基点餐,如果只有固定的套餐组合,而顾客想要的组合并不在菜单上,那就无法满足需求。
实际上,即使是这样简单的例子,模型的表现也并不总是令人满意,尤其是对于较小的模型,比如6B以下的模型,它们的表现会出乎意料地差。这促使我们去深入分析原因,并寻求改进的方法。
那么,我们应该如何将QK和OV进行动态组合呢?我们可以在多头注意力的几个计算步骤中进行组合。比如,将Wq、Wk、Wo、Wv等权重组合,或对输出进行组合。而我们在DCFomer中采用的其实是将它生成的注意力分数和注意力权重矩阵做组合。
接下来,我们通过之前的例子来说明在注意力机制的层面上,组合不同头 (Head) 的信息优势。假设我们有一句话,其中每一对注意力的起始词和终点词 (例如“pen”和“of”) 都可以计算出一组注意力分数,它们是通过权重矩阵Wq和Wk相乘得到的。假设有八个头 (H=8) ,那么就有八组Wq和Wk的组合。这些组合计算出的8个注意力分数被称为注意力向量,它们决定了从“pen”到“of”应该传递哪些信息。例如,假设第六个头 (Head 6) 的QK回路特别关注了“pen”,在正常情况下,第六个头的数值会较大,模型就会倾向于使用第六个头的OV回路来传递信息。
但如果第六个头的OV并不获取我们需要的属性信息,比如“kindOf” (类别) ,那么我们就不能简单地使用它。为了解决这个问题,我们可以调整这些注意力分数,假设第四个头 (Head 4) 的OV包含所需的“kindOf”信息。我们可以通过将Head 6的注意力分数增加给Head 4,让它在计算中起到更大的作用,从而实现在Head 6的QK搭起的通路上传递Head 4的OV获取的属性信息。
通过这样“偷梁换柱”,我们就实现了不同头之间的信息组合,让模型能够更灵活地处理信息。通过这种方式,我们可以操作注意力向量,使模型能够更准确地捕捉和传递关键信息。
那么,如何将一个注意力向量A:ij转换成我们想要的另一个向量A:ij呢?一种简单的方法是通过矩阵乘法。当我们将一个向量与一个特定的矩阵相乘时,它就可以转换成另一个向量。这个矩阵的性质决定了转换的具体方式。
例如,如果矩阵设计成特定的形式,它就可以实现特定的变换,比如将原本属于第四个头的信息移动到第六个头,或者反之。此外,矩阵还可以调整各个头的权重,比如增强或减弱某个头的影响力。比如,如果我们想要减弱第46个位置的权重,我们可以在矩阵的对角线上相应位置设置较小的值,这样在乘法操作后,该位置的权重就会降低。
我们将这些用于调整和组合的矩阵称为“Composition Map” (组合映射) 。这些映射不是针对单个注意力起点-终点对,而是全局性的。
如果这些组合映射对于所有起点-终点对都是相同的,我们称之为“静态”的组合映射,其效果等同于扩展了头的QK (Query-Key) 和OV (Output Value) 的维度,比如如果有8个头,就可以将每个头的维度扩展到原来的8倍,从而增强模型的表达能力。
在DCMHA中,我们引入了一个“compose”的关键操作,是将一个注意力向量转换成另一个不同的注意力向量,从而实现更灵活的注意力分配。
具体来说,compose操作会接收一个原始的注意力向量,并输出一个新的、经过转换的注意力向量。在多头注意力机制中,我们会在两个关键的位置插入compose操作:一个在softmax归一化步骤之前,另一个在softmax之后。通过这种方式,DCMHA能够动态地调整每个头的贡献,创造出更加精确和有效的注意力模式。提高了模型的表达能力,而且增强了其对复杂数据关系的捕捉能力,从而在各种任务中都能取得更好的性能。
Noam Shazeer等人的Talking-heads attention工作提出过类似的注意力头组合思想 (其本质是上文的静态组合映射) ,与他们的工作不同,我们的工作则更加关注动态组合的概念。所谓动态,意味着根据不同的输入,在每个位置生成不同的composition map,使模型能够动态调整不同头之间的组合方式。例如,给定前文“小明有只钢笔”,对“小明有一种 - 文具”,模型需要将找相同主语的QK与取属性的OV组合,而对“小明拿它 - 写字”,则需要与取用途的OV组合。
动态组合的优势在于它能够更灵活地适应不同的语言模式,但这也带来了计算上的挑战。因为模型需要为每个token生成一个动态的组合映射,这在计算上可能会非常昂贵。特别是当模型规模较大时,需要处理的参数和计算量会显著增加。
为了解决这个问题,我们在DCMHA中采用composition map分解的方法。这种方法将compose操作分解为几个步骤,每个步骤都有其特定的作用。例如,我们可以通过低秩分解来限制组合的复杂性,从而减少计算开销。低秩意味着每个特定位置需要的组合数量是有限的,这样可以有效地控制计算资源的使用。此外,我们还引入了门控机制,它通过调整composition map对角线上的值来增强或减弱特定头的影响。这种方法允许模型在不同情况下灵活调整每个头的贡献。
最终,通过将这些不同步骤的结果结合起来,我们可以得到一个综合的输出。如果我们从整体张量的角度来理解这个过程,实际上是对组合张量进行了行列分解和低秩对角分解,这些技术在其他领域也有应用,但我们将其创新性地应用于组合映射上。
这种方法引入了一些额外的参数,但其计算复杂度相对较小。增加的计算量 (flops) 也不多,并且随着模型规模的增大,这种复杂度相对减少,即模型越大,增加的计算复杂度越少。
接下来,让我们看看主要的测试结果——性能算力比。性能算力比衡量的是单位算力投入能带来多少性能提升。实验结果显示,在相同的训练数据 (300B tokens) 下,一个改进后的69亿参数模型,其效果甚至超过了120亿参数的模型。例如,如果GPT-4o能够应用这项技术,推理一次128k上文的成本就可能从4元降低到2元。此外,在下游任务的评估中,随着模型规模的增大,我们的改进效果更加显著,如6.9B参数模型与2.8B参数模型的比较。在保持超参数和训练数据一致的情况下,无论是损失值 (Loss) 还是语言模型的困惑度 (PPL) ,DCPythia-6.9B的表现都优于Pythia-6.5B和Pythia-12B。
尽管理论上的复杂度很小,但实际训练中我们确实发现了一些额外开销。这主要是因为compose操作对注意力矩阵进行了一系列的变换,虽然每个变换的计算量不大,但由于Attention矩阵本身很大,导致了较大的I/O需求。针对这个问题,我们在实践层面进行了一些优化。优化结果显示,无论是在训练还是推理中,相比于Transformer模型开云网址·(中国)官方网站,这种额外开销是可以接受的。例如,在13B或6.9B模型上,尽管推理速度慢了5.5%,但性能却超过了13B模型。如果与13B模型相比,实际上速度是变快的。
从这个角度来看,随着算力的提升,模型规模的增大带来了更大的收益,额外开销实际上随着算力的增加而降低。
此外,我们还进行了消融实验,以评估DCMHA中五个分支的重要性。实验结果表明,动态组合比静态组合更重要,尤其是跨头的组合,比单个头的自增强或自抑制 (即门控) 更为关键,这符合我们的直觉;关注侧的信息 (query-wise) 比被关注侧的信息 (key-wise) 更重要。即使只使用一侧的信息,也能达到很好的效果。基于这些观察,我们可以进一步采用一些方法,比如使用更多的局部注意力层,只用关注侧的信息 (QW) ,这不仅能保持类似的效果,还能进一步提高速度。
我们还进行了另一项实验,测试了DCFormer是否能够解决需要注意力头组合的任务。我们构建了一个测试集,包含了各种类似于“小明有支钢笔”这样的测试用例,并进行了不同的QK和OV变换。结果表明,在6.9B参数的模型上,DCFormer在这种需要头组合的任务上的表现确实显著优于普通模型。
接下来,我想谈谈我们工作中存在的一些局限性。首先,关于实现方面,我们的模型主要是使用JAX在TPU上进行训练的。虽然我们也提供了在GPU上使用JAX和PyTorch进行训练的代码,但我们必须承认,这部分代码并没有经过深入优化,因此在GPU上的性能可能会差于TPU上的性能,而推理部分是在GPU上使用PyTorch来执行。
其次,关于现有模型的适配问题。我们希望能够将DCFormer技术应用到现有的模型上,而不是从头开始训练一个新模型,因为从头训练的开销是相当大的。我们尝试过将现有的模型,比如Llama2或Llama3,直接应用DCFomer技术,并通过一些微调来提升性能。但实验结果表明,这种简单的直接应用是不够的。原因在于,我们的DCFormer依赖于不同头之间的组合,而现有的模型在训练时并没有考虑到这种组合关系。例如,一个头的QK (Query-Key) 需要与另一个头的OV (Output Value) 进行组合,这种组合的前提是它们必须在同一层中。但是,在现有的Llama模型中,需要组合的QK和OV不一定位于同一层,这就导致了无法有效实现组合。这不仅是一个技术挑战,而且在分布式训练环境中,跨层组合的实现难度和开销都是相当大的。因此,目前我们的一个局限性在于,虽然跨层组合在技术上是可行的,但要实现它并达到良好的性能,还存在一定的难度和成本。
关于这篇论文的工作,以上介绍的就是我们研究的主要内容。现在,我想简单地回顾一下我们进行这项研究的初衷。我们最初的出发点是提高模型的可解释性。具体来说,我们设计了一系列任务,并开发了自己的分析方法和工具来探究模型的工作机制:模型是如何执行任务的,以及它为何能或不能完成这些任务。我们的尝试相当于将模型逆向工程从“黑盒”变为“灰盒”,使我们能够对其进行调试并寻找改进方案。
费曼说:“如果你不能创造一样东西,你就不能真正理解它”;我觉得反过来:如果你不能理解一样东西,你也无法改进或创新它。然而,这并非唯一的途径。许多杰出的研究人员,比如之前提到的Noam Shazeer,他们的论文往往不会详细解释设计结构的动机,我们也无法确切知道他们的灵感来源。可能在更高层次的逻辑中,存在着其他的思路和方法,比如直觉的引导。当实践经验和视野积累到一定程度时,直觉可能会自然而然地引领你走向正确的道路。
然而,在实际研究过程中,我们最深的体会是迭代速度的重要性。即使你的想法在理论上是可行的,但在实践中往往需要不断地调整和实验,快速地完成闭环。我们总共进行了500多次大大小小的实验,才验证了我们方法的有效性。这意味着,如果你能加快迭代速度,将原本一年内完成的实验压缩到半年,那么你就有可能在半年内取得显著的成果。
此外,面向底层硬件的优化也是至关重要的。在模型优化的后期,我们发现性能提升的瓶颈往往不在于PyTorch或JAX等框架层面,而在于TPU、GPU等底层硬件。最终,我们可能需要编写CUDA Kernel等底层代码,这似乎是不可避免的。最近许多架构创新的工作,如Flash Attention、Mamba、GLA等,它们的实现都依赖于定制化的Kernel,否则模型效率将受到影响。因此,我曾跟团队半开玩笑说,当前架构创新的尽头是CUDA编程。如果你的想法仅限于PyTorch、JAX层面,那么你的发挥空间可能受到限制。而如果你掌握了CUDA、Pallas等底层技术,你将有更多的空间去做出更出色的工作。
我坚信底层结构的创新不仅可行,而且是一个充满潜力的领域。自从Mamba模型出现后,它激发了一系列令人兴奋的研究工作,这一领域也逐渐受到更广泛的关注。在过去,许多人认为Transformer架构已经达到了底层结构的极限,无需进一步改进。然而,Mamba、DCFormer等工作的出现打破了这一刻板印象,为我们打开了新的探索空间。
特别值得一提的是,我们对混合架构非常感兴趣,尤其是将Attention机制与RNN或状态空间模型 (SSM) 结合起来。这两种架构具有不同的内在偏好 (Inductive Bias) ,它们的结合能够形成强大的混合模型。例如,简单地将一层Attention与一层RNN结合,就能显著提升模型的性能。无论是已发表的研究还是我们的实验,都证实了这种混合方法的有效性。此外,这种方法与人脑的工作方式更为贴近。人脑在维持当前状态的同时,能够通过远程Attention机制进行记忆检索,而混合架构正是这种能力的体现。
此外,信息流动性也是我们关注的焦点。无论是Mamba还是DCFormer,它们都通过输入相关的方式增强了模型内部的信息流动。尽管Transformer模型本身也能动态形成信息流通路径,但我们认为这些路径的动态性还有待提高。因此,探索如何以更高效的方式促进信息流动通路的动态形成,是一个充满希望的研究方向。
肖达,北京邮电大学讲师,彩云科技首席科学家,集智俱乐部核心成员。长期从事深度学习、大模型的算法研究及其在天气预报、机器翻译、机器辅助创作、系统和代码安全等领域的应用研究,相关工作发表在 ICLR、ICML、TDSC等顶级会议和期刊。目前研究兴趣包括大模型机制可解释性、高效模型架构设计。
DCFomer的论文作者均来自彩云NLP算法组,核心成员出自集智俱乐部。团队继承了集智俱乐部自由讨论和乐于奉献的精神,鼓励开放开云网址·(中国)官方网站、真诚、平等的沟通,秉承着精益求精的精神,致力于“让田间地头的农民大哥”和“字宙里的航天员”一样,都能通过人工智能产品让生活更美好开云网址·(中国)官方网站。
参与⼤模型对齐方向的研发,涉及数据构造,SFT及RL训练等。持续提升自研基座模型在业务场景下的能力
985/211高校硕士及以上学历或优秀本科生,计算机、自动化、人工智能、机器学习、数学或统计学等相关专业,两年及以上NLP相关经验
对Transformers模型、LLM预训练/微调、RLHF/RLAIF等基本原理有深入理解和通过实践掌握的实现细节
对现有模型结构进行基础性改进和创新,设计更高效的模型结构,结合优化算法、数据混合策略等维度的改进、大幅度提高基础模型预训练的效率
985/211高校硕士及以上学历或优秀本科生,计算机、自动化、人工智能、机器学习、数学或统计学等相关专业,两年及以上大模型相关经验
2022年11月30日,一个现象级应用程序诞生于互联网,这就是OpenAI开发的ChatGPT。从问答到写程序,从提取摘要到论文写作,ChatGPT展现出了多样化的通用智能。于是,微软、谷歌、百度、阿里、讯飞,互联网大佬们纷纷摩拳擦掌准备入场……但是,请先冷静一下…… 现在 all in 大语言模型是否真的合适?要知道,ChatGPT的背后其实就是深度学习+大数据+大模型,而这些要素早在5年前的AlphaGo时期就已经开始火热了。5年前没有抓住机遇,现在又凭什么可以搭上大语 言模型这趟列车呢?
集智俱乐部特别组织,由北师大教授、集智俱乐部创始人张江老师联合肖达、李嫣然、崔鹏、侯月源、钟翰廷、卢燚等多位老师共同发起,旨在系统性地梳理ChatGPT技术,并发现其弱点与短板。读书会已完结,现在报名可加入社群并解锁回放视频权限。