ICML2024高分论文!大模型计算效率暴涨至200%,来自中国AI公司
夕小瑶科技说 原创
作者 夕小瑶编辑部前段时间,KAN突然爆火,成为可以替代MLP的一种全新神经网络架构,200个参数顶30万参数;而且,GPT-4o的生成速度也是惊艳了一众大模型爱好者。
大家开始意识到——
大模型的计算效率很重要,提升大模型的tokens生成速度是很关键的一环。
而提升大模型的tokens生成速度,除了花钱升级GPU外,更长效的做法是改善Transformer模型架构的计算效率。
今天,笔者发现,终于有团队对Transformer计算最耗时的核心组件——多头注意力模块(MHA)下手了,将Transformer的计算性能提升了有2倍之高。
通俗的讲,如果这项工作未来能落地到大模型里面,那么大模型tokens生成速度翻倍式提升的一天就不远了。
这篇论文已经被今年的机器学习顶会ICML 2024录用,拿到了7分的高分,而且还开源了。
据透露,今年ICML 2024录用的paper平均得分在4.25-6.33之间
笔者扒了下,发现这个工作的背后是一家颇具影响力的国内公司——彩云科技,没错,就是打造爆火的“彩云小梦”产品的团队。
不急,先看看这篇论文,如何将Transformer模型计算效率暴涨100%的。
论文标题:
Improving Transformers with Dynamically Composable Multi-Head Attention
论文链接:
https://arxiv.org/abs/2405.08553
开源项目地址:
https://github.com/Caiyun-AI/DCFormer
Github上已开源这项工作的代码、模型和训练数据集。
我们知道,承载Transformer计算量的核心模块便是多头注意力(MHA)模块,位置(position=i)上的每一个注意力头(attention head)会与全部位置上的注意力头计算出一个注意力分布矩阵。在这个过程中,位置 i 上的各个注意力头计算出来的注意力分布矩阵是相互独立的。
忘了的小伙伴请自行扒拉Transformer论文
论文指出,这种多头独立计算的机制会带来两大问题:
低秩瓶颈(Low-rank Bottleneck):注意力矩阵的秩较低,模型的表达能力受限 头冗余(Head Redundancy):不同的注意力头可能会学习到相似的模式,导致冗余
因此,彩云科技提出了一种叫动态可组合多头注意力(DCMHA)的机制,DCMHA 通过一个核心的组合函数(Compose function),以输入依赖的方式转换注意力得分和权重矩阵,从而动态地组合注意力头,解决了传统MHA模块中存在的上述低秩瓶颈和头冗余问题。
值得强调的是,DCMHA旨在提高模型的表达能力,同时保持参数和计算效率,它可以作为任何Transformer架构中MHA模块的即插即用替代品,以获得相应的DCFormer模型。
论文通过实验表明,DCFormer在不同的架构和模型规模下,在语言建模方面显著优于Transformer,与计算量增加1.7倍至2倍的模型性能相匹配。例如,DCPythia-6.9B在预训练困惑度和下游任务评估方面优于开源的Pythia-12B。
DCMHA原理
DCMHA机制的核心是引入的Compose函数。这个Compose函数可以视为一个可学习的参数,它可以动态地组合不同头的QK矩阵和VO矩阵,内部通过一系列变换来分解和重构注意力向量。可以近似理解为:经过组合映射后,H个基础的注意力头可组合成多至H*H个注意力头。
你可以简单理解为,它能根据输入数据调整头之间的交互方式,一是打破头的独立性,二是可以根据输入数据动态组合,从而可以增强模型的表达能力。
利用矩阵分解实现高效的参数计算
尽管引入了动态组合,DCMHA的设计依旧注重参数和计算效率。通过矩阵分解DCMHA能够以较小的额外参数和计算开销实现动态组合,同时保持模型性能。
DCFormer可提高70%~100%的模型计算效率
还有很重要的一点是,DCMHA可以作为MHA的直接替代品应用于任何Transformer架构中,升级成DCFormer,实现计算效率的大幅提升,达到1.7倍-2倍的计算效率。
而且,实验结果表明在众多NLP下游任务和图像识别任务上的测评也验证了DCFormer的有效性。
1、DCFormer在不同参数规模下(405M到6.9B参数),对 Transformer 和 Transformer++ 模型的性能提升显著。
自2017年Transformer诞生至今,旋转位置编码RoPE和门控激活函数MLP被证明是最普世有效且广泛采用的改进,已融入到Transformer++架构,同时也是大名鼎鼎的Llama模型框架。
而DCFormer性能算力比的提升幅度超过这两项改进的提升幅度之和。
2、DCPythia-6.9B在多个下游任务中的表现优于Pythia-12B。
3、在ImageNet-1K数据集上的实验验证了DCMHA在非语言任务中也是有效性的。
从上图可以看出,在相同训练数据和算力下,一个被本文方法改进后69亿参数的模型,却拥有比120亿参数传统模型结构更好的效果。
换句话讲:相同的参数量下,使用DCFormer将具备更强的模型表达能力;用更少的参数量,拥有相同的模型表示效果。
DCFormer在不同的架构和模型规模下,在语言建模方面显著优于Transformer,与计算量增加1.7倍至2倍的模型性能相匹配。
距离大模型“光速”生成tokens不远了
笔者觉得这个工作还是蛮扎实的,如果能像RoPE一样在国内外的主流大模型落地,大模型“光速”生成tokens的一天并不遥远,而且从AI产业对电力能源的利用效率来说,也是一个很有意义的改善。
实话说,在如今这个“资本寒冬”,愿意为前瞻技术研究投入资金、人才支持的公司非常少了,能在ICML这个高含金量机器学习顶会上跑出来高分论文的团队,背后一定离不开公司层面的支持。
在写这篇文章的时候,笔者注意到,彩云科技团队也在进行大模型对齐和测评算法研究员、大模型推理优化、AIGC产品经理、后端工程师、前端工程师、SRE工程师等职位的招聘,这里附上简历投递二维码:
倘若能进入到发表ICML高分论文的团队参与AI方向的学术研究和产品落地,属实是一个非常珍贵的职业经历,感兴趣的小伙伴抓住机会。
笔者在搜彩云科技的时候,还无意间扒出来了意想不到的东西。
笔者发现,有一款服务500+家大客户的超大型B端产品——彩云天气竟然也是彩云科技旗下的。
没准,你手机里的、汽车车载系统里的天气APP背后走的很可能就是彩云天气API。
做过ToB业务的都知道,能获得100家大客户青睐的B端产品就已经具备相当的B端影响力了,而彩云天气不仅斩获了滴滴、小米、vivo、高德、360、小鹏汽车在内的500多家大客户,其甚至早在2014年就成为了中国气象局的战略合作伙伴,曾帮助多个部门和地区避免了自然灾害风险。
不夸张的讲,彩云天气已成为了国内事实上的气象服务基础设施。
这背后,无疑是彩云科技强悍的AI算法实力和强大的工程能力。
如果你对AI ToB产品觉得陌生,那彩云科技旗下的另一款爆款AI ToC产品请让我安利下,因为——
它真的太圈粉了!
作为文字工作者,笔者自ChatGPT爆火以来,玩遍了国内外几乎所有的文字创作类产品,但给笔者留下深刻印象&能持续用起来的产品不多,彩云小梦就是其中一款。
彩云小梦是一款网文辅助写作工具,也是一个 AI RPG 平台,用户可以在里面扮演各种角色,体验不同的人生。AI 写作助手具有文风独特、可以自动续写、支持自定义开头等特点和功能。
作为曾经的RPG游戏爱好者(玩过金庸群侠传、仙剑奇侠传、武林群侠传系列的小伙伴请举手🙋🏻♀️),笔者甚至用彩云小梦将金庸群侠传游戏剧情翻写过小说,因为彩云小梦AI生成的内容太有意思了,贴一段你们自己感受下:
在写网文这块,用过彩云小梦后就再也用不回ChatGPT了,体验差别悬殊。
但最让我停不下来的,倒不是写网文。而是你可以扮演网文里面的角色:
这种沉浸式角色的体验非常神奇,经常玩半天都停不下来:
彩云小梦的产品体验非常丝滑、稳定,技术出身的笔者,能深刻的在这份丝滑背后的工程能力和产品能力有多强大。
除了彩云天气和彩云小梦外,彩云科技旗下的彩云小译也是业界有口皆碑的同声传译软件,不仅有阅文集团、360和维基百科等广泛的客户群基础,其甚至给《三体》做过翻译,篇幅原因,这里就不展开讲了。
总之,通过进一步深挖彩云科技旗下的产品,笔者深感这是一家集强大的AI算法、工程和产品能力于一身的老牌科技公司,这种低调钻研技术、打磨产品、做扎实的价值创造的宝藏团队在国内属实稀缺。深得笔者喜爱。
最后贴下彩云科技的招聘信息,多个岗位正在火热招聘中,感兴趣的小伙伴抓住机会,招聘岗位详情请点击链接进一步了解:
http://colorfulclouds.com/jobs/
微信扫码关注该文公众号作者