不受窗口长度限制的长文本生成全新思路:利用模型参数储存上文信息

2024-02-09 12:28 644 阅读 ID:1875
将门
将门

目前的长文本生成方面的研究主要集中在长度外推和扩充窗口长度上,其主要思想都是在模型的KV states中尽可能多且有效的储存上文的信息,并让模型在推理时候尽可能准确的用好这些信息。然而这种存储是有上限的,基于此本文研究者在新工作,With Greater Text Comes Greater Necessity: Inference-Time Training Helps Long Text Generation,提出了一种全新的思路,来支持无限长文本生成。主要思想简单有效到令人吃惊:将上文信息储存在模型参数中,而不是KV cache中,来降低对KV states的依赖。这些用来储存上文信息的参数被集中在一个临时的Lora模块(Temp-Lora)中,推理过程不断拿模型生成的token训练这个模块,来实现历史信息的存储。推理结束后,这个Temp-Lora模块可以被直接丢弃,从而防止对模型参数的永久化影响。

论文题目: 

With Greater Text Comes Greater Necessity: Inference-Time Training Helps Long Text Generation 

论文链接: 

https://arxiv.org/abs/2401.11504

目前的长文本生成方面的研究主要集中在长度外推和扩充窗口长度上,其主要思想都是在模型的KV states中尽可能多且有效的储存上文的信息,并让模型在推理时候尽可能准确的用好这些信息。然而这种存储是有上限的:

  1. 大部分模型的能力尚不足以完美利用所有KV states内的信息(如论文Lost in the Middle: How Language Models Use Long Contexts中的发现,模型的attention集中在上文的头部和尾部)。
  2. 更长的上文需要消耗更多的硬件资源(如,34B model不做任何量化,在单张A100上仅能容纳不到30K token)和推理时延。

我们最近做了一个有趣的新工作,With Greater Text Comes Greater Necessity: Inference-Time Training Helps Long Text Generation,提出了一种全新的思路,来支持无限长文本生成。主要思想简单有效到令人吃惊:将上文信息储存在模型参数中,而不是KV cache中,来降低对KV states的依赖。这些用来储存上文信息的参数被集中在一个临时的Lora模块(Temp-Lora)中,推理过程不断拿模型生成的token训练这个模块,来实现历史信息的存储。推理结束后,这个Temp-Lora模块可以被直接丢弃,从而防止对模型参数的永久化影响。

该框架的效果也相当惊人:在小说补全任务上,当历史文本超长时(>500K tokens),能将Llama2-7B的PPL降低13.2%。在小说翻译任务上,能将Yi-Chat-6B的PPL降低29.6%并提升6.6 BLEU score。更重要的是,在这个框架下,历史文本的增长不会带来额外的计算开销,即:无论context有多长,生成一个token所需要的浮点运算数(FLOPs)和推理延时(latency)都是基本恒定不变的。

                                                图一:Temp-Lora-Augmented Long Text Generation

一、方法

该方法的核心如图一所示,用原文的话说即:At the heart of this method lies the progressive training of a temporary Lora module, which is named Temp-Lora, on previously-generated text in an auto-regressive manner (自回归方式不断用模型生成的文本训练一个临时的Lora模块). 当生成结束之后,就实行“飞鸟尽,良弓藏;狡兔死,走狗烹”策略,直接把Temp-Lora模块扔进垃圾桶,以避免对模型主体参数的永久化影响。另外,当输入文本本身就特别长的时候,直接把输入文本也按照这样的方式切分成多个chunk并对Temp-Lora模块做一次预训练,能保证模型在推理开始阶段就更好的理解输入中包含的信息。详细算法如下图(注:该算法大部分都在处理边际情况,真正的核心都在图一。算法中唯一重要参数是Chunk size ,代表了隔多少token更新一次Temp-Lora模块,其他各参数定义请移步原文):

我们还提供了一个cache reuse策略来节省计算消耗,即更新了Temp-Lora之后不用新模型重新计算KV states,而是直接继续使用原来的cache中KV states的,直到不得不重新计算的时候(token数达到了窗口上限),再用新模型计算KV states。

二、实验

我们在两个数据集上做了实验:PG19:一个包含了许多1919年的小说的数据集 和 Guofeng:WMT 2023中的小说翻译数据集,包含了200本中国网络小说以及它们的对应译文。

我们将PG19数据集按照上文长度的不同分成了不同的子集,从0-100K 到 500K+。理论上讲,上文越长,长文本理解能力就越重要,而实验结果也证明了这一点:在500K+子集上,Llama2-7B-4K的PPL从4.61降低到了4.00 (13.2%),而其他模型也有10%左右的PPL下降;与之对应,在0-100K子集上,PPL的降低幅度只有大约3%。平均来看,在整个PG19数据集上,Temp-Lora方法都能未各模型带来大约5%的PPL下降。

在小说翻译上,Temp-Lora方法取得了更显著的效果:直接让Yi-6B-Chat的PPL从5.67降低到了3.99,BLEU上涨了6.6,COMET上涨了6.1。作者也分析了为何Temp-Lora在小说翻译任务上更显著的效果:小说补全是一个没有标准答案的任务,即使是一个正在读某本小说的人类读者,也很难“猜”到作者接下来会安排怎样的剧情发展,因此,Temp-Lora只能小幅度的增加真实文本的生成概率。对于小说翻译来说,为了保持一致性,大多数名词实体的翻译都是唯一的。通过有效地在 Temp-Lora 模块中存储这些翻译映射,可以轻易地实现大幅度的翻译效果提升。

我们还在Llama2-7B-32K上分析了Temp-Lora跟计算效率之间的关系,如上图,并给出结论:相比于将长文本硬塞进context window造成的计算开销,训练Temp-Lora的开销相对来说是很低的。结论:在最“经济适用”的配置下,将模型的window size限制为3K,并用2K的chunk size,能在降低3.8% PPL的同时降低70.5%的浮点运算量(FLOPs)和51.5%的推理时延。而对"土豪"来说,直接不做任何改变,使用Temp-Lora,能以19.3%的额外计算量和19.3%的额外推理延时为代价获得5%的PPL下降。

三、结论和应用

我一(xiang)拍(le)脑(ban)袋(tian),闹弄了个中二的漫威体口号来总结:文本越长,越需要Temp-Lora。这是一个跟长度外推、窗口扩充等正交的方法。另外,顺便也跟大家分享一些在实际场景应用Temp-Lora的经验:

  1. 我们与微软的工作 “Fine-tuning or retrieval? comparing knowledge injection in llms ”有着相似的结论:Temp-Lora不能取代RAG。或者说,Retrieval的结果还是需要放在prompt里面才能有最佳效果。这是因为简单的1-2 steps的训练不足以让模型准确记住一些数字之类的factual information。
  2. 在角色扮演(CharacterAI类产品),尤其是扮演已有角色时候,Temp-Lora方法有神奇的效果。服务提供者可以在用户上传的角色背景资料上训练一个Temp-Lora模块,来为模型注入该角色的背景知识。当用户删除该角色的时候,Temp-Lora模块也随之删除。
  3. 在对话式助手(如ChatGPT)中,可以用Temp-Lora模块储存用户与模型的对话历史。一般来说,用了Temp-Lora模块以后,模型就不再需要扩窗口了,在应用时候保持与pre-train阶段一样的窗口长度即可。当用户删除某个对话session时,Temp-Lora模块也随之删除
  4. 在游戏NPC & Generative Agents领域(如斯坦福小镇),Temp-Lora模块可以用来记录一个npc过去所经历的所有事件,来取代现在常用的memory-summarization-reflection框架。一般来说可以有1 npc/1 Lora 和1 npc + 1 user/ 1 Lora两种模式,取决于应用形态(千人千面 or 千人一面)
  5. 在人机交互式翻译(如腾讯Transmart)中,简单的用同一个项目中已经翻译好的文本训练Temp-Lora模块,能保证机器翻译跟译员的翻译习惯(如特定名词的翻译)高度一致

本文转载自:

https://zhuanlan.zhihu.com/p/679713147

Illustration From IconScout By Delesign Graphics

-The End-

免责声明:作者保留权利,不代表本站立场。如想了解更多和作者有关的信息可以查看页面右侧作者信息卡片。
反馈
to-top--btn