无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年,AI在线AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。如果您有优秀的工作想要分享,欢迎投稿或者联系报道。投稿邮箱:[email protected][email protected]李堉晖:北京大学智能学院硕士,受张弘扬老师和张超老师指导,研究方向为大模型加速和对齐,正在寻找25届工作机会魏芳芸:微软亚研院研究员,研究方向为具身智能、图像生成和AI agents张超:

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年,AI在线AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。如果您有优秀的工作想要分享,欢迎投稿或者联系报道。投稿邮箱:[email protected][email protected]

李堉晖:北京大学智能学院硕士,受张弘扬老师和张超老师指导,研究方向为大模型加速和对齐,正在寻找25届工作机会

魏芳芸:微软亚研院研究员,研究方向为具身智能、图像生成和AI agents

张超:北京大学智能学院研究员,研究方向为计算机视觉和机器学习

张弘扬:滑铁卢大学计算机学院、向量研究院助理教授,研究方向为LLM加速和AI安全

自回归解码已经成为了大语言模型(LLMs)的事实标准,大语言模型每次前向计算需要访问它全部的参数,但只能得到一个token,导致其生成昂贵且缓慢。

今日,一篇题为《EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees》的论文提出了动态草稿树投机采样,依据草稿模型的置信度动态调整草稿树的结构,最高可以将大语言模型的推理速度提高5倍,同时不改变大语言模型的输出分布,确保无损。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

论文链接:https://arxiv.org/pdf/2406.16858

项目链接:https://github.com/SafeAILab/EAGLE

Demo链接:https://huggingface.co/spaces/yuhuili/EAGLE-2

EAGLE-2在多轮对话数据集MT-bench上的加速效果(上图为贪婪生成,下图为采样生成):

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

使用EAGLE-2,2张RTX 3060($300)的推理速度可以超过A100($10000)。无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

背景

投机采样使用一个小的模型快速生成草稿,原始的大语言模型可以通过一次前向计算验证草稿的正确性,将正确的草稿作为输出,从而一次生成多个token,并确保无损。EAGLE是投机采样的一种改进。它在更有规律的特征层面而不是token层面进行自回归,同时输入采样结果(超前一个时间步的token)消除了不确定性,明显提升了草稿模型的准确率。

到目前为止,EAGLE在第三方测试Spec-Bench(https://github.com/hemingkx/Spec-Bench/blob/main/Leaderboard.md)中排名第一。

思路

EAGLE和Medusa等方法使用静态的草稿树,隐式地假设草稿token的接受率和上下文无关,下面是一个简单的例子

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

上文是“10+2”时,下一个token难以预测,EAGLE在这个位置添加两个候选token以增加草稿命中率,“10+2=”和“10+2+”有一个正确即可。当上文是“10+2=”时,下一个token明显是“1”,但是EAGLE使用静态的草稿结构,仍然添加两个候选“1”和“3”,“10+2=3”不可能通过大语言模型的检查,存在浪费。EAGLE-2旨在解决这一问题,如下图所示,当上文是“10+2=”时,EAGLE-2只增加一个候选token“1”,将节约出的token用于让草稿树更深,这样“10+2=12”通过大语言模型的检查,EAGLE-2可以一次生成更多的token。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

EAGLE-2的作者们在Alpaca数据集上进行了简单的测试,下图显示了不同位置的草稿token的接受率,左图中的P1-P6代表位置,与右图的横轴坐标对应。实验结果显示,在相同的位置上的草稿token的接受率也有较大的差异,这说明了使用动态草稿树可能取得比静态草稿树更好的效果。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

上述例子中,EAGLE-2根据预测草稿token的难易程度决定草稿树的结构,精确计算难易程度(接受率)需要原始大语言模型的计算结果,这违背了投机采样减少对原始大语言模型访问的初衷。幸运的是,EAGLE的草稿模型的置信度与接受率(难易程度)高度正相关。下图显示了草稿模型不同置信度区间的草稿token的平均接受率,红色虚线连接(0,0)和(1,1)。由此可见,草稿模型的置信度可以作为接受率的有效近似。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

方法

EAGLE-2包括两个阶段,扩展和重排,扩展阶段加深加大草稿树,重排阶段修剪草稿树,丢弃部分节点(token)。

为了保证无损,一个草稿token被接受的前提是它的祖先节点都被接受,所以EAGLE-2将一个节点的价值定义为它和它祖先的接受率的乘积,用置信度的乘积来近似。

在扩展阶段,EAGLE-2选择草稿树最后一层价值最高的m个节点(token)进行扩展。这些token被送入草稿模型,然后将草稿模型的输出作为子节点连接到输入节点,加深加大草稿树。在重排阶段,EAGLE-2按照价值对整棵草稿树进行重排序,保留前n个节点(token)。草稿token的置信度在0-1之间,两个节点价值相同时优先保留浅层节点,因此重排后保留的草稿树一定是连通的,保证了语义上的连贯性。重排后草稿树变小,降低了原始大语言模型验证的计算量。为了保证计算结果的正确性,还需要调整attention mask,确保每一个token只能看到它的祖先节点,不受其他分支的影响。下面是一个简单的例子。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

扩展(Expand)阶段的黄色框表示被选中进行扩展的节点,绿色框为以这些节点为输入时草稿模型的预测。重排(Rerank)阶段的蓝色框表示被保留的节点,之后它们被展平成一维作为原始大语言模型的输入。EAGLE-2根据树的结构调整attention mask,比如,”a”只能看到它的祖先“It”和“is”,看不到另一个分支的“has”。EAGLE-2也同时调整位置编码,确保和标准自回归解码的一致性。

实验

EAGLE-2在多轮对话、代码、数学推理、指令遵循、问答、总结六项任务上分别使用MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM、Natural Questions数据集进行了实验,与6种先进的投机采样方法(SpS、PLD、Medusa、Lookahead、Hydra、EAGLE)进行了比较。

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

无损加速最高5x,EAGLE-2让RTX 3060的生成速度超过A100

表格中的Speedup为加速比,τ 为平均接受长度,也就是原始大语言模型每次前向计算能生成的token数。EAGLE-2每次前向计算能生成大约4-5个token,而自回归解码每次生成1个token,因此EAGLE-2明显加速了大语言模型的生成,加速比为2.5x-5x。加速比和接受长度在代码生成任务(Humaneval数据集)上最高,这是因为代码中存在大量确定性的模板,草稿更容易命中。在所有任务和大语言模型上,EAGLE-2的加速比和平均接受长度都是最高的,明显优于其他方法。

应用

EAGLE-2也在工业界得到应用,集成至Intel/intel-extension-for-transformers等。

给TA打赏
共{{data.count}}人
人已打赏
工程

太酷了!iPhone、iPad、MacBook老旧设备组成异构集群,能跑Llama 3

2024-7-16 15:16:00

工程

ACL 2024 | 对25个开闭源模型数学评测,GPT-3.5-Turbo才勉强及格

2024-7-18 14:25:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
搜索