英伟达提出全新 Star Attention：10 倍加速大语言模型推理，登顶 Hugging Face 论文榜

如今，大型模型拥有不断增长的上下文长度，这也导致了推理成本上升。英伟达最新推出的 Star Attention 提出了一种方法，旨在减少推理计算量，从而在不降低精度的情况下助力边缘计算。

当前，手机和AIPC上都安装了本地大型模型，但随着上下文长度增加，推理成本也在显著增长。用户往往需要等待很长时间才能获取结果。

为了解决这个问题，各种优化方案相继提出，包括 Flash Attention。而在11月26日，英伟达提出了Star Attention机制，旨在提升Transformer模型在处理长序列时的效率和准确性。

这篇文章受到了广泛关注，登上了Hugging Face每日论文榜首。

论文地址：https://arxiv.org/abs/2411.17116

Star Attention如何优化推理成本

在探讨Star Attention如何改进大型模型推理之前，让我们先了解当前大型模型推理涉及的两个步骤：

1）prompt编码，即模型处理输入并在缓存中存储KV（键值）向量；

2）token生成，即模型关注KV缓存并自回归生成新令牌，同时用新的KV向量更新缓存。

许多长上下文任务中，输入由一个长上下文后跟一个短查询和一个短答案组成。随着大型模型上下文变得越来越长，回答查询所需的信息通常仅限于上下文的一小部分，这意味着上下文只需关注附近的标记，而查询标记则需要关注所有之前上下文的内容。

Star Attention下的两阶段推理

系统中所有设备被分组为多个主机（host），其中一个主机被标记为「查询」主机。输入序列分为两个阶段处理。

阶段一：上下文编码

输入的上下文部分被划分为较小的块，并分配到各主机。除第一个块外，所有块的前面都加上一个初始块，称为「锚点」块（anchor block）。每个主机处理其分配的块，并存储非锚点部分的KV缓存。

阶段二：查询编码和token生成

输入查询被广播到所有主机，每个主机首先访问第一阶段计算出的本地KV缓存。然后，「查询」主机通过聚合所有主机的softmax归一化数据来计算全局注意力。这一过程会针对每个生成的标记重复进行。

用一个不那么严谨的例子概述上述过程：想象一场烹饪比赛（上下文标记），每个厨师（主机）负责准备一部分菜肴（块）。

为了保持口味一致，除了准备自己的部分外，每位厨师在自己部分之前添加了一点「锚点」调味料（锚点块）。厨师准备好自己的部分后，会记住自己部分的口味（KV缓存）。

阶段二的查询编码和token生成可以类比为：评委（查询标记）品尝菜肴，并决定下一道菜的口味（生成新标记）。评委首先品尝每位厨师的部分，看看哪个最符合他们的口味。

最终，评委汇总所有厨师的意见，确定下一道菜的口味，并告知厨师们。

Star Attention的性能优化

Star Attention的性能提升主要体现在以下两个方面：

1）高达11倍的加速

在多个长上下文基准测试中，Star Attention带来的8B Llama3的推理速度显著提升，随着序列长度的增加，加速比从1.1倍提升至2.7倍。

而在参数更多的Llama3.1-70B上，推理的加速比提升更为明显。

此外，与采用全局注意力的基准相比，Star Attention相对准确率仅降低了0～3%。

Star Attention推理在处理较长上下文时，与全局注意力相比准确性几乎一致，但推理计算成本显著减少。

在更长的上下文（128K）中，不同块大小在上下文编码阶段会影响推理的准确性和速度。块尺寸越大，Star Attention的准确性也更高。

在RULER基准测试中，不同块大小对Star Attention准确性的影响，适用于序列长度为128K的Llama-3.1-8B instruct模型。

RULER用于评估时包含了13个任务，分为4个领域：大海捞针（检索）、多跳追踪、聚合和问答。

在不同任务中，全局注意力和Star Attention在准确性方面的差异对比

当上下文长度增加至1048K时，Star Attention仍然保持了原始基准90%的推理准确性，推理加速比达到了10.8×～16.9×。

在更大的Llama3.1-70B模型中，Star Attention实现了更大的加速比，并同时保持相似水平的准确率下降。

由于其运行机制不牵扯具体模型，Star Attention能够无缝集成到大多数基于Transformer的LLMs中，而无需进行额外的模型微调。

通过减少推理计算成本，Star Attention显著降低了内存需求，使得在本地设备（如手机、笔记本）上处理更长的序列成为可能。

实验证明，将块大小设置为总序列长度的约四分之一，可以在精度和速度之间取得最佳平衡。用户也可以根据需求调整块大小，以在计算效率和精度之间进行权衡。

结论

未来的研究将尝试将Star Attention扩展到更长的序列（最长可达1M）和更大的模型，并希望观察到更多的加速效果，同时保持相似水平的准确率。同时，重点将放在优化「锚块」机制上，以提升Star Attention的可扩展性和稳健性，以提高在更复杂的长上下文任务中的性能。

总的来说，对于希望开发部署本地大型模型的厂商而言，Star Attention是一项不可忽视的技术。使用Star Attention后，本地LLM能够更快地为用户提供结果，同时在有限内存情况下与更长的上下文序列兼容，在RAG任务中可以处理更长的文本。

对于云端大型模型的供应商来说，Star Attention能够在几乎不影响用户体验的前提下显著提升推理成本，实现「降本增效」，同时减少能源消耗（碳足迹）。

通过在多个主机之间分配上下文处理，Star Attention使得上下文长度能够与主机数量线性扩展。

参考资料：

• https://arxiv.org/abs/2411.17116

本文来自微信公众号：新智元（ID：AI_era）

广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考。