
中国科学院自动化研究所、伦敦大学学院和香港科技大学(广州)共同承担了GTA项目的研发工作,并成功研发出一种新型的高效大模型框架,该框架显著提高了模型的性能和计算效率。其中,自动化所的孙罗洋博士生是该项目的参与者之一,他的研究方向集中在高效大模型的计算与优化领域。该项目的主要通讯作者包括香港科技大学(广州)的邓程博士、自动化所的张海峰教授以及伦敦大学学院的汪军教授。该成果为大模型的优化部署提供了创新解决方案。
GTA(Grouped-head latent Attention)技术惊艳亮相!该技术通过共享注意力矩阵及压缩潜在值表示,有效降低了计算量达62.5%,同时将KV缓存减少了70%,并且显著提升了prefill和decode的速度,实现了速度翻倍。在应对由海量数据形成的长序列任务时,亦或在计算资源极其匮乏的边缘设备上执行,GTA始终展现出其无与匹敌的效率与出色的性能,这无疑使其成为了大型语言模型优化领域的新标杆。
大型语言模型面临的效率困局
Transformer架构的崭新亮相显著加速了自然语言处理领域的迅猛进步,这不仅让大型语言模型在对话生成、文本摘要、机器翻译以及复杂推理等多个领域取得了卓越成就,而且展现了其令人瞩目的强大能力。然而,当模型参数的数量从数十亿急剧攀升至千亿级别,传统多头注意力机制(MHA)的固有缺陷逐渐显现,这已成为限制其广泛运用和持续进步的瓶颈。
计算冗余问题首当其冲。在多头注意力架构中,每个注意力头均相当于一个独立的“工作单元”,它们各自独立地处理查询(Query)、键(Key)和值(Value)向量的计算,从而引发了大量的重复运算。在应对持续时间较长的任务时,浮点运算的次数会急剧上升,呈平方增长趋势,这极大地降低了模型的处理速度,从而使原本就复杂的任务变得更加费时。
内存瓶颈问题同样不容忽视。每个注意力头均需保存其对应的键值对缓存,这一需求使得内存消耗随着序列长度和注意力头数量的增长而急剧上升。以处理长序列为例,KV 缓存的容量可能迅速超过数 GB,如此巨大的内存占用显著制约了大型模型在智能手机、物联网设备等边缘设备上的应用,使得其难以广泛进入普通家庭。
最终面临的是推理过程中的时间滞后问题。由于计算和存储资源的高要求,导致推理速度大幅降低,这便使得诸如语音助手即时响应、在线翻译无间断切换等对响应时间极为敏感的实时服务,难以给用户带来顺畅的使用体验。尽管研究人员在业界曾尝试运用多查询注意力(Multi-Query Attention,简称MQA)和分组查询注意力(Grouped-Query Attention,简称GQA)等技术来提升效率,然而这些方法通常需要在性能提升与资源消耗之间做出艰难的选择,难以实现两者之间的理想平衡。面对这一连串的艰巨考验,研究团队经过持续不断的奋斗,最终成功研发出Grouped-head latent Attention(GTA)这一创新技术。该技术以其革命性的设计理念,对注意力机制的效率上限进行了重新界定,为大型语言模型的进一步发展拓展了全新的路径。

GTA 的核心创新机制
GTA之所以取得显著成就,主要得益于其两项核心技术的重大突破。这两项技术巧妙地相互配合,即便在资源极为紧张的环境中,大型语言模型依然能够展现出前所未有的高效运行水平。
分组共享注意力矩阵机制
在MHA的传统架构里,每个注意力头均被视作一位“独行侠”,它们各自独立进行计算并管理自己的注意力得分。尽管这种分散的计算方式让模型具备了捕捉多种复杂依赖关系的能力,然而,它也带来了相当大的计算冗余问题。以一个拥有16个注意力模块的MHA为例,当每个模块单独对输入进行处理时,能够产出16个独立的注意力矩阵,从而导致整体的计算成本随着注意力模块数量的提升而呈现线性上升,效率相对较低。
与此形成鲜明对比,GTA引入了一种全新的“团队协作”模式。在这种机制下,注意力头被巧妙地划分为若干个逻辑小组,比如,每4个注意力头便构成一个小组,而这些小组的成员会共同使用一张统一的注意力矩阵。这种创新的共享设计方案确保,我们只需对注意力分数进行一次性计算,随后便能够高效地将其分配给组内所有注意力头,进而显著降低了浮点运算次数(FLOPs)。

实验数据充分表明,这种巧妙的设计能够有效降低总体的计算负担,对于处理超长序列的任务,实现了显著的推理速度提升。这就像一位技艺高超的厨师,事先准备齐全所有食材,然后分派给不同的助手进行精细的操作,不仅大大节约了宝贵的时间,而且还保证了最终成品的高品质和一致性。
压缩潜在值表示技术
MHA架构存在一个显著的难题,那就是其KV缓存对内存的消耗极大。这是因为每个注意力头的值向量都必须被完整保存,这使得模型的内存需求随着输入序列的长度和注意力头数量的增加而急剧上升,从而成为部署大型模型的重大障碍。GTA通过其独特的“压缩+解码”设计,巧妙地解决了这一难题。
这项技术首先对所有的注意力头值向量进行了高效的压缩,将其转换成了一个低维度的潜在表示(即Latent Representation),这一过程显著降低了所需的存储空间。紧随其后,模型借助一个轻量级且运作高效的 WaLU(加权加性线性单元)非线性解码器,能够依据每个注意力头特有的需求,从该紧凑的潜在表示中灵活地、个性化地构建出所需的完整值向量。

此创新举措不仅大幅降低了珍贵的内存消耗,而且巧妙地维护了每个注意力模块的独特表达能力,有效防止了信息流失。实验数据令人振奋,GTA 的 KV 缓存量实现了高达 70% 的缩减,这一重大突破为大型语言模型在性能受限的边缘设备上的广泛应用奠定了基础,使其能够更加广泛地服务于各种应用场景。
实验验证:GTA 的卓越性能与效率
研究团队经过一系列细致周密的实验,对Grouped-head Latent Attention(GTA)在多种模型尺寸、输入序列长度以及不同硬件平台上的表现和效能进行了详尽的测试与分析。实验数据有力地证实,GTA 在显著提高计算效能和内存使用效率方面表现出色,不仅维持了,甚至在某些核心性能指标上超过了目前流行的注意力机制模型,充分体现了其卓越的实用价值以及巨大的应用潜力。
模型有效性验证
为确保实验数据的公正性与精确度,研究团队在实验方案中实施了严谨的变量控制措施:将所有与注意力无关的模型参数(诸如隐藏层的大小、多层感知机MLP的规模等)设定为恒定值,以此保证模型参数数量的任何变动仅由注意力机制的独立创新设计所引起。
160M 参数模型表现
在针对参数量达到160M规模的模型进行测试时,无论是输入序列长度为2048个token,还是4096个token,GTA模型始终表现出显著的超群性能。具体来看,运用GTA2配置的模型在2048个token序列长度上,实现了比传统MHA、GQA和MLA更低的评估损失,并且展现出更出色的Wikitext困惑度(PPL)表现。另外,采用GTA1配置的模型在多个下游任务中获得了更高的平均准确率,这充分证明了它在实际应用中的高效性。特别需要指出的是,GTA 在内存使用效率上展现出卓越的表现,其每一层所需的键值对缓存容量仅是 MHA 的 12.5%(具体数据为 192 维度与 MHA 的 1536 维度相比),这种大幅度的减少充分彰显了 GTA 在内存优化领域的强大实力。具体的实验数据详尽地展示在下面的表格里。

500M 参数模型表现
当模型参数量增至500M时,GTA在性能方面依旧处于领先位置。在2048个token序列长度的测试中,GTA不仅达到了更低的评估损失,而且在下游任务中的平均准确率也有所提升。与此同时,GTA在Wikitext困惑度方面与MHA、GQA等主流模型持平,甚至在某些方面表现更为出色。GTA 在内存使用上展现出其独特的优势,其 KV 缓存容量仅为 MHA 的 12.5%,具体来说,是 320 维度与 MHA 的 2560 维度相比。即便是在使用更小的缓存(比如 192 维度,仅占 MHA 的 7.5%)时,GTA 依然能保持与 MHA 相当的性能,这充分证明了它在内存效率和性能之间达到了完美的平衡。在应对涉及4096个token的序列处理任务时,GTA不仅成功达到了MHA在评估损失方面的水平,而且在Wikitext的困惑度以及下游任务的平均准确率方面展现了更为卓越的成果。具体的数据详情已详尽地整理并展示在下方的表格之中。

1B 参数语言模型扩展性
为了深入检验GTA在大型模型中出色的扩展性能与稳定性,研究团队特地培育了参数量达到10亿级别的GTA-1B和GQA-1B模型。下方的图表直观地描绘了GTA-1B与GQA-1B在经过50,000次训练迭代后的损失曲线与梯度范数曲线,观察结果显示,这两个模型均呈现出了令人满意的稳定收敛态势。

GTA-1B 在设计上选择了较小的缓存规模,然而其轨迹损失与 GQA-1B 极为相似,这一现象显著地印证了 GTA 内存高效架构的优越性,即在降低资源消耗的同时,并未对模型的学习能力造成影响。在一系列严格的评估测试中,GTA-1B(包含经过SFT优化的版本)与GQA-1B的表现相当,甚至更胜一筹,特别是在平均准确率方面实现了显著进步。这一成果充分证明了GTA即便在资源有限的环境中,也能通过微调技术,在多种复杂任务中实现有效泛化,并维持其强劲的竞争力。具体的实验数据已详尽地展示在下方的表格中。

综合分析,GTA-1B 在基础模型阶段以及经过微调之后,均展现出了与 GQA-1B 相当的优异性能。此外,它的 KV 缓存容量仅为 GQA-1B 的三分之一,而自注意力计算的开销更是大幅降低至 GQA-1B 的37.5%。这些引人注目的数据显著突显了内存与计算高效架构在大型语言模型应用扩展领域所蕴含的巨大潜能,预示着未来人工智能的发展趋势将更加侧重于可持续性和资源的高效利用。
效率评估
理论效率分析


通过 LLM-Viewer 进行经验基准测试
为了实现理论上的优势向具体可衡量的性能转变,该研究团队借助LLM-Viewer框架,在装备有NVIDIA H100 80GB GPU的高性能计算设备上,对GTA-1B与GQA-1B模型进行了全方位的基准测试。图示直观地对比了两种模型在多种配置下的预填充与解码所需时间。观察结果清晰表明,GTA-1B 在处理计算密集型的预填充环节以及I/O密集型的解码环节均持续领先于GQA-1B,这充分彰显了其在延迟性能方面的显著优势以及更高的执行效率。

实际部署表现
为了更准确地衡量GTA-1B在实际应用领域的表现,研究人员运用了transformers库,对包括服务器级别的NVIDIA H100、NVIDIA A800,以及消费级产品RTX 3060,还有边缘设备如Apple M2和BCM2712在内的多种异构硬件平台,开展了细致的推理测试。

图中清晰呈现了GTA-1B与GQA-1B在多种配置条件下的预填充及解码耗时差异。GTA-1B(以蓝色实线表示)在各个测试环境中均保持了相对于GQA-1B(以橙色虚线表示)更短的预填充时间,特别是在面对2k token等较长的输入序列时,两者间的性能差异更加明显,这充分说明了GTA-1B在处理长文本方面的显著优势。在解码环节,GTA-1B展现出了非凡的能力,尤其是在扩展生成内容长度方面,其表现优势在各类硬件平台上均保持稳定,这充分彰显了其设计的高可靠性。

图中详细揭示了开启缓存卸载功能后的性能变化。在 NVIDIA H100 硬件上,GTA-1B 在处理较长的输入序列时,依旧显现出其预填充的领先优势,同时,在解码速度方面,相较于 GQA-1B,GTA-1B 实现了更为显著的提升。在各个平台上的这一持续发展态势,显著彰显了GTA-1B在I/O密集型环境中的卓越性能,此类环境中,缓存卸载要求GPU与CPU内存间进行频繁的数据交换,然而GTA-1B即便在这样的复杂环境中,依旧展现出优异的表现。
总体来看,GTA-1B在各个硬件平台上,无论是在预填充还是解码所需的时间上,都全面胜过了GQA-1B。此外,在处理较长的输入序列时,GTA-1B也展现出明显的性能优势。它不仅在标准的推理配置下表现出色,即便是在开启缓存卸载的I/O密集型环境下,其表现同样卓越。这充分展示了GTA-1B在不同硬件配置和批处理规模下的强大多功能性。这种卓越的适应性让GTA-1B成为服务器端与消费端部署的理想选择,它显著减少了计算难度和内存要求,从而大幅提高了大型语言模型中注意力机制的整体运作效率。
技术局限与未来方向
尽管 Grouped-head latent Attention(GTA)在效率与性能上实现了显著的进展,然而,作为一项新兴技术,它仍面临若干关键的技术难题,这些难题值得我们持续关注并深入挖掘。在模型压缩的环节,非线性解码器可能会产生细微的近似错误,对此,未来的研究需在架构设计及训练策略方面进行进一步的优化,以保障模型输出的精确度。此外,目前针对 GTA 的研究和验证主要聚焦于自然语言处理领域,而在计算机视觉以及多模态任务中的应用潜力和性能表现,尚需开展更为全面和深入的探究与测试。
尽管面临诸多限制,研究团队已为GTA的进一步发展设定了明确且充满潜力的后续研究路线。他们打算不断优化非线性解码器的架构,旨在确保解码效率的同时,还能进一步降低信息损耗,从而提高模型性能的极限。研究团队怀揣壮志,拟将GTA技术扩展至更大规模模型之中,旨在检验其在超大规模场景中的扩展能力与效率优势,进而助力大型语言模型迈向更广泛的应用天地。
本文采摘于网络,不代表本站立场,转载联系作者并注明出处:https://mjgaz.cn/fenxiang/277481.html
