3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

高效解码n -token序列,CLLMs Jacobi解码框架。传统上,大型语言模型(LLMs)被认为是顺序解码器,逐个解码每一个token。来自上海交通大学、加利福尼亚大学的研究团队展示了预训练的LLMs可以轻松地被教导成为高效的并行解码器,并介绍了一种新的并行解码器族,称为一致性大语言模型(CLLMs),能够通过在每一个判断步骤中高效地解码一个n -token序列来降低判断延迟。在此篇论文中,研究表明:「模仿人类在头脑中形成完整句子后逐字表达的认知进程,可以通过简单地微调预训练的LLMs来有效地学习。」具体而言,C

高效解码n -token序列,CLLMs+Jacobi解码框架。

传统上,大型语言模型(LLMs)被认为是顺序解码器,逐个解码每一个token。

来自上海交通大学、加利福尼亚大学的研究团队展示了预训练的LLMs可以轻松地被教导成为高效的并行解码器,并介绍了一种新的并行解码器族,称为一致性大语言模型(CLLMs),能够通过在每一个判断步骤中高效地解码一个n -token序列来降低判断延迟。

在此篇论文中,研究表明:「模仿人类在头脑中形成完整句子后逐字表达的认知进程,可以通过简单地微调预训练的LLMs来有效地学习。」

具体而言,CLLMs通过将任何随机初始化的n -token序列映射到尽可能少的步骤中,产生与自返回(AR)解码相同结果,来进行并行解码的训练。

实验结果表明,利用该研究团队所提出的法子获得的CLLMs非常有效,在生成速度上显示出该法子获得了2.4倍至3.4倍的改进,与其他快速判断技术如Medusa2和Eagle相媲美甚至更好,且在判断时不需要额外的内存利润来容纳辅助模型组件。

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

论文名称:《CLLMs:Consistency Large Language Models》

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

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

                               图1:在GSM8K上,利用Jacobi解码时,CLLM-ABEL-7B-001是baseline ABEL-7B-001大约3倍速度的演示。

Jacobi解码

大型语言模型(LLMs)正在改变人类生活的面貌,从编程到提供法律和健康建议。

然而,在判断进程中,LLMs利用自返回解码逐token生成响应,如图1所示,这导致了较长响应的高延迟。利用自返回解码,通常需要进行架构修改、辅助组件或初稿模型等,以通过一次生成多个token来加快判断速度。

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

                                 图2:传统的自返回(AR)解码示意图:一次生成一个token。

Jacobi解码源自Jacobi和Gauss-Seidel定点迭代求解非线性方程的法子,经证明与利用贪婪解码的自返回生成完全相同。

Jacobi解码将顺序生成进程重新构造为一个包含n个变量的n个非线性方程组,并基于Jacobi迭代可以并行求解。

每一个迭代步骤可能会预测出多个精确的token(所谓的“精确”是指在贪婪采样策略下与自返回解码结果对齐),从而潜在地加快自返回解码。

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

                               图3:Jacobi解码示意图:将n -token序列馈送到LLM中,并进行迭代直至收敛。

具体来说,Jacobi解码法子首先从输入提示中随机猜测序列的下一个token(以下简称为n -token序列,除非另有说明)。

然后,将n -token序列连同提示一起馈送到LLM中,以进行迭代更新。这个进程会持续进行,直到n -token的序列稳定下来,不再发生变化,达到一个固定点。

值得注意的是,Jacobi解码不需要比自返回(AR)解码更多的对LLM的查询。最终,n -token的序列会收敛到在贪婪策略下由AR解码生成的输出。从最初的随机猜测到最终的AR生成结果的这一进程被称为「Jacobi轨迹」。

Jacobi解码迭代进程和Jacobi轨迹的一个实例在图2中进行了说明。

Jacobi解码的局限:

然而,在实践中,普通的Jacobi解码对LLMs的加快效果仅有微弱的提升,例如,平均加快比只有1.05倍。这是因为当LLM在先前的token中存在错误时,很难产生精确的token。

因此,大多数Jacobi迭代只能为n -token的序列获得一个校正,导致如图3左侧所示的较长轨迹。

前瞻解码和推测解码法子试图缓解Jacobi解码和传统的自返回解码的低效问题,但在判断时会产生额外的内存利润。

而CLLMs则不需要这些额外的内存利润。

一致性大型语言模型(CLLMs)

初步Jacobi解码:

给定一个prompt x和一个预训练的LLM p(·|x),通常研究者会利用标准的自返回(AR)解码法子在贪婪策略下获得模型的响应,即:3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

Jacobi解码重新构建了LLM判断进程,将其视为解决非线性方程组的进程,以将解码进程转化为可并行计算的形式。考虑到: 

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

研究者可以将上述方程重写为一个非线性方程组:

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

需要注意的是:

该进程在某个k值处退出,使得:3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

然后,定义3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了作为固定点,并且将3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了作为Jacobi轨迹。

为了解决该问题,该研究团队提出调整预训练的LLMs,使它们能够一致地将Jacobi轨迹 J 上的任意点 y 映射到固定点 y*

令人惊讶的是,他们发现这样的宗旨类似于一致性模型的宗旨——一种扩散模型的主要加快法子。

在该团队提出的法子中,利用从宗旨模型收集的Jacobi轨迹来训练模型,并利用一种受益函数,该函数鼓励在Jacobi迭代进程中实现单步收敛。

对于每一个要调整为CLLM的宗旨模型 p ,训练包括两个部分:

(1)Jacobi轨迹准备:

对于每一个提示,作者按顺序对每一个token截断进行Jacobi解码,直到整个响应序列 l 被生成为止,这相当于所有连续固定点的串联。

 沿轨迹生成的每一个序列都会被算作一个数据条目。

此处需要注意的是,对于包含N(N ≫ n)个token的 I 的长响应,这种截断避免了对长输入的慢速模型评估。

(2)利用一致性和AR受益进行训练:

作者联合优化两个受益来调整CLLMs,一致性受益确保一次预测多个token,而AR受益则防止CLLM偏离宗旨LLM,以保持生成质量。3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

                                图4: one-step收敛一致性训练的示意图:将宗旨LLM调整为在Jacobi轨迹上的任何状态作为输入时始终预测固定点。

一致性和AR受益:

(1) 一致性受益

p 表示宗旨LLM。

使3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了表示为被初始化为 p 的参数 θ 的 CLLM。

对于prompt x 和相应的Jacobi轨迹 J ,令 y y* 分别表示轨迹上的随机状态和固定点。

可以通过最小化以下受益来促使CLLM在输入为 y* 时输出 y ,称为全局一致性(GC)受益

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了在此公式里,3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

作者大量利用符号来表示从数据集中均匀抽样。

D(·||·) 表示两个分布之间的距离,选择则在GKD法子中进行了讨论,在本文中主要利用前向KL。

或者,根据一致性模型中的公式,利用局部一致性(LC)受益。

其中相邻状态:3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了在Jacobi轨迹 中,被驱使产生相同的输出:

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

(2)AR受益:

为了避免偏离宗旨LLM的分布,作者结合了基于宗旨LLM p 的生成 l 的传统AR受益:

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

通过将两种受益结合在一起,利用一些权重 ω ,训练CLLM的总受益为:

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

实验

结果:

总的来说,该实验涵盖了三个特定领域的任务:

(1)Spider(文本到SQL)

(2)Human-Eval(Python代码完成)和GSM8k(数学)

(3)更广泛的开放域会话挑战MT-bench。

所报告的实验利用微调过的编码器LLM、Deepseek-coder-7B-instruct、LLaMA-2-7B或ABEL-7B-001作为宗旨模型,具体利用则取决于任务。

训练和评估则都在NVIDIA A100 40GB服务器上进行。

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

图5:CLLM在不同下游任务上的加快效果。结果显示:「CLLM比预训练模型快得多,并且与Medusa相比实现了可比较的加快,但在判断时没有额外的利润。」

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

图6:CLLM与其他基准在特定领域任务(Spider、CSN-Python、GSM8k)以及MT-bench上的对比示意图。CLLM在与Medusa2的比较中实现了类似或甚至更好的加快效果,同时不引入额外的判断利润(根据FLOPS和内存消耗判断)。

专业领域:

从图5中,可以看到,与其他基准(包括原始宗旨模型、Medusa2和猜测解码)相比,CLLMs实现了最显著的加快。

开放域会话挑战(MT-bench):

利用ShareGPT数据集从LLaMA2-7B训练的CLLM与前瞻解码(lookahead decoding)结合利用时,可以实现与Medusa2大致相同的加快,并在MT-bench上获得可比较的分数。

然而,CLLM具有更高的适应性和内存效率,因为它不需要对宗旨模型的原始架构进行修改,也不需要辅助组件。

训练利润:

CLLMs的微调利润是适度的。

例如,对于LLaMA-7B,只需传递大约1M token就可以在Spider数据集上实现它的3.4倍的加快。在数据集size较大的情况下(例如对于CodeSearchNet-Python),只需要利用数据集的10%来生成训练CLLMs的Jacobi轨迹,从而获得大约2.5倍的加快。

可以通过以下方式估算token的总数:

N = 平均每一个prompt的轨迹量 × 平均轨迹长度 × Prompt数量。

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

3倍生成速度还降内存利润,超越Medusa2的高效解码框架终于来了

图7:Spider上宗旨LLM和CLLM之间的Jacobi轨迹对比。Jacobi轨迹上的每一个点都是颜色编码的序列:与AR结果匹配的精确标记为蓝色,不准确的标记为红色。CLLM表现出增强的效率,收敛到固定点的速度比宗旨LLM快2倍。CLLM的这种增强效率可以归因于一致性受益,它促进了每一个给定前缀的n-token 序列的结构的学习。

图6左侧显示,宗旨LLM通常在一个迭代中只生成一个精确的token。相比之下,在CLLMs中,作者发现了快速推进现象,即在单个Jacobi迭代中连续多个token被精确预测。

此外,在宗旨LLM中,提前精确生成的token(例如图7左侧索引6和7处的「country」和「H」),往往在随后的迭代中被不准确地替换。

另一方面,CLLMs表现出了预测精确token的能力,即使在先前有错误token的情况下,也确保token保持不变。

作者将这样的token称为「固定token」。这两种现象共同促成了CLLMs在Jacobi解码中的快速收敛,从而实现了相当大的生成速度提升。

研究团队还观察到,通过训练,CLLMs获得了一个关键的语言概念——搭配:「一系列词或术语,它们的共同出现频率高于随机机会所预期的。」

语言不仅由孤立的单词组成,而且还严重依赖于特定的词对。搭配的示例在自然语言和编程语言中都很丰富。

它们包括:

动词+介词组合(例如「talk to」,「remind … of …」)

动词+名词结构(例如「make a decision」,「catch a cold」)

许多领域特定的句法结构(例如「SELECT … FROM …」,「if … else」用于编程)。

一致性生成宗旨使CLLMs能够从Jacobi轨迹的任何点判断出这样的结构,促进CLLMs掌握大量的搭配,并因此同时预测多个词以最小化迭代步骤。

参考链接:

https://hao-ai-lab.github.io/blogs/cllm/

给TA打赏
共{{data.count}}人
人已打赏
应用

AI在用 | AI模仿器《换你来当爹》上线,「马斯克」10岁就被我练废

2024-5-10 17:58:00

应用

苹果启动AI云效劳器打算,芯片直接用M2 Ultra

2024-5-11 11:04:00

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