中国科学院自动化研究所、伦敦大学学院与香港科技大学（广州）共同承担了GTA项目的研发工作，他们共同研发出一种高效的大模型框架，这一框架显著提高了模型的性能和计算效率。其中，自动化所的孙罗洋博士生是该项目的参与者，他的研究方向专注于大模型的高效计算与优化。该项目的主要通讯作者包括香港科技大学（广州）的邓程博士、自动化所的张海峰教授以及伦敦大学学院的汪军教授。该成果为大模型的优化部署提供了创新解决方案。
GTA（Grouped-head latent Attention）技术惊艳亮相！该机制通过共享注意力矩阵及压缩潜在值表示，有效降低了计算量达62.5%，同时将KV缓存需求减少了70%，并且显著提升了prefill和decode的速度，实现了2倍的增长。在应对由海量数据形成的长序列任务时，亦或在计算资源极其匮乏的边缘设备上执行，GTA均呈现出令人瞩目的效率和非凡的表现，毫无疑义，它将确立为大型语言模型优化领域的新典范。
大型语言模型面临的效率困局
近年来，Transformer 架构的崭露头角极大地促进了自然语言处理领域的迅猛进步，大型语言模型在对话生成、文本摘要、机器翻译以及复杂推理等多个前沿领域接连取得显著成就，其卓越能力令人叹为观止。然而，当模型参数的规模从数十亿急剧膨胀至千亿级别，传统多头注意力机制（Multi-Head Attention，MHA）本身存在的缺陷也逐渐显现，这已成为限制其广泛运用和继续进步的障碍。
计算冗余问题首当其冲。在多头注意力架构中，每个注意力头均扮演着独立的“工作单元”角色，它们各自独立地处理查询（Query）、键（Key）和值（Value）向量的计算，从而引发了大量的重复计算现象。在应对那些涉及长序列的任务时，浮点运算的次数（FLOPs）会急剧上升，呈平方级增长，这极大地降低了模型的处理速度，导致原本就复杂的任务变得更加耗时。
内存瓶颈问题同样不容忽视。每个注意力头均需保留其对应的键值对（KV）缓存，这一需求使得内存消耗随着序列长度及注意力头数量的增长而急剧上升。以处理长序列为例，KV缓存的容量很容易超过数GB，如此巨大的内存占用严重制约了大型模型在智能手机、物联网设备等边缘设备上的应用，使得其难以广泛普及于千家万户。
最后需关注推理过程中的延迟现象。这种延迟主要源于计算和内存的高需求，进而显著降低了推理速度。因此，对于对延迟极为敏感的实时应用，如语音助手即时响应和在线翻译的无缝切换，往往难以保证用户享受到流畅的使用体验。尽管研究人员在业界曾探索过运用多查询注意力（MQA）和分组查询注意力（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 维度的 GTA 与 1536 维度的 MHA 相比，这一大幅度的降低充分彰显了 GTA 在内存优化方面的强大实力。具体实验数据详尽地展示在下方的表格里。
500M 参数模型表现
当模型参数量达到500M时，GTA依旧在性能方面保持领先。在2048个token的序列长度测试中，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在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 技术拓展至更大规模的模型之中，旨在检验其在超大规模场景中的可扩展性与效率上的优势，进而助力大型语言模型迈向更为广泛的应用领域。
