英伟达提出全新Star Attention，10倍加速LLM推理！登顶Hugging Face论文榜
科技
2024-12-10 00:01
吉林
![]()
MLNLP社区是国内外知名的机器学习与自然语言处理社区，受众覆盖国内外NLP硕博生、高校老师以及企业研究人员。
社区的愿景是促进国内外自然语言处理，机器学习学术界、产业界和广大爱好者之间的交流和进步，特别是初学者同学们的进步。
当下的手机及AIPC中都会安装本地大模型，然而上下文长度增加，推理时的计算成本也会显著增长。最明显的一个后果就是，用户输入问题后需要等待很久才能看到结果。为此，已有多种优化方案提出，例如Flash Attention，而11月26日英伟达提出的Star Attention机制，可用于提升Transformer模型在处理长序列时的效率和准确性。值得一提的是，这篇文章受到了广泛的关注，登顶Hugging Face每日论文榜首。论文地址：https://arxiv.org/abs/2411.17116Star Attention如何降低推理成本
在了解Star Attention如何改进大模型推理前，让我们先看看当前大模型的推理过程涉及的两个步骤：1）prompt编码，即模型处理输入并在缓存中存储KV（键值）向量；2）token生成，即模型关注KV缓存并自回归生成新令牌，同时用新 的KV向量更新缓存。在许多长上下文任务中，输入由一个长上下文后跟一个短查询和一个短答案组成。当大模型的上下文变得越来越长之后，回答查询所需的信息通常局限在上下文的小部分内，意味着上下文只需关注附近的token，而查询token需要关注所有之前上下文涉及的内容。系统中所有设备被分组为多个主机（host），其中一个主机被标记 为「查询」主机。输入序列分为两个阶段处理。输入的上下文部分被分割成较小的块，并分配到各个主机。除了第一个块之外，所有块的前面都加上一个初始块，称为「锚点」块（anchor block）。每个主机处理其分配的块，并存储非锚点部分的KV缓存。输入查询被广播到所有主机，在每个主机中，它首先访问在第一阶段计算出的本地KV缓存。然后「查询」主机通过聚合所有主机的softmax归一化统计数据来计算全局注意力。这个过程对于每个生成的token都会重复。用一个不那么严谨的例子来概述上面的过程：想象一场烹饪比赛（上下文token），每个厨师（主机）负责准备一道菜的一部分（块）。为了确保味道一致，每个厨师除了准备自己的部分，还在前面加了一点「锚点」调料（锚点块）。每个厨师准备好自己的部分后，记住自己部分的口味（KV缓存）。阶段二的查询编码和token生成可视为：评委（查询token）来品尝菜肴，并决定下一道菜的口味（生成新token）。评委先品尝每个厨师的部分，看看哪个部分最符合他们的口味。最后，评委汇总所有厨师的意见，确定下一道菜的口味，并告诉厨师们。Star Attention的性能提升
Star Attention带来的性能提升，主要体现在以下两个方面：在多个长上下文基准测试上，Star Attention所加持的8B Llama3的推理速度显著提升，随着序列长度增加，加速比从1.1x提升到2.7x。而在参数量更大的Llama3.1-70B上，推理的加速比提升更为显著。与此同时，对比采用全局注意力的基准，Star Attention相对准确率的降低只在0～3%范围内。随着上下文长度的增加，star attention推理的准确性相比全局注意力几乎相同，但推理计算成本显著下降在更长的上下文尺度（128K）中，上下文编码过程中不同块的大小，也会影响推理的准确性和速度。块尺寸越大，Star Attention 的准确性越高。在 RULER 基准测试上，不同块大小对Star Attention准确性的影响，块大小范围从4K到32K，适用于序列长度为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显著减少了内存需求，使得在本地设备（如手机，笔记本中）用LLM处理更长的序列成为可能。实验发现，将块大小设置为总序列长度的约四分之一，可以在精度和速度之间取得最佳平衡。而用户也可以根据需求调整块大小，以在计算效率和精度之间进行权衡。结论
未来的研究，会尝试将Star Attention扩展到更长的序列（最长可达1M）和更大的模型，并希望能观察到甚至更的加速，同时保持相似水平的准确率。同时专注于优化「锚块」机制，并在更复杂的长上下文任务上提高性能，以增强Star Attention的可扩展性和稳健性。总的来看，对于想要开发部署本地大模型的厂商，Star Attention是一项不容错过的技术。使用Star Attention后，本地LLM能够更快地回复用户，还可在有限的内存中兼容更长的上下文序列，从而在RAG任务中阅读更长的文本。而对于云端大模型的提供商，Star Attention能够在几乎不影响用户体现的前提下，显著提升推理成本，实现「降本增效」，同时减少能源消费（碳足迹）。通过在多个主机间分配上下文处理，Star Attention使上下文长度能够随主机数量线性扩展。https://arxiv.org/abs/2411.17116![]()
扫描二维码添加小助手微信
即可申请加入自然语言处理/Pytorch等技术交流群关于我们
MLNLP 社区是由国内外机器学习与自然语言处理学者联合构建的民间学术社区，目前已经发展为国内外知名的机器学习与自然语言处理社区，旨在促进机器学习，自然语言处理学术界、产业界和广大爱好者之间的进步。社区可以为相关从业者的深造、就业及研究等方面提供开放交流平台。欢迎大家关注和加入我们。![]()
