开云网址·(中国)官方网站

　　开云网址提高大模型计算效率！彩云科技推出大模型 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技术，并发现其弱点与短板。读书会已完结，现在报名可加入社群并解锁回放视频权限。