首页 > 其他分享 >MemLong: 基于记忆增强检索的长文本LLM生成方法

MemLong: 基于记忆增强检索的长文本LLM生成方法

时间:2024-09-24 16:22:29浏览次数:13  
标签:检索 MemLong 模型 记忆 LLM 上下文 注意力


本文将介绍MemLong,这是一种创新的长文本语言模型生成方法。MemLong通过整合外部检索器来增强模型处理长上下文的能力,从而显著提升了大型语言模型(LLM)在长文本处理任务中的表现。

核心概念

MemLong的设计理念主要包括以下几点:

  1. 高效扩展LLM上下文窗口的轻量级方法。
  2. 利用不可训练的外部记忆库存储历史上下文和知识。
  3. 通过检索相关的块级键值(K-V)对来增强模型输入。
  4. 适用于各种仅解码器的预训练语言模型。
  5. 引入额外的记忆检索(ret-mem)组件和检索因果注意力模块。

MemLong的工作流程如图1所示:

MemLong: 基于记忆增强检索的长文本LLM生成方法_人工智能

图1:MemLong的记忆和检索过程示意图

MemLong的技术优势

1. 分布一致性

MemLong克服了传统方法中将信息存储在记忆中导致的分布偏移问题。通过精心设计的记忆机制,MemLong确保了缓存信息的分布与原始数据保持一致,从而提高了模型的稳定性和性能。

2. 训练效率

MemLong采用了创新的训练策略:

  • 冻结模型的底层参数,仅微调上层结构。
  • 显著降低了计算资源需求,例如:- 3B参数版本的MemLong在0.5B个token上的微调仅需8张3090 GPU运行8小时。

这种方法大幅提高了模型的训练效率,使得在有限资源条件下也能实现高效的模型优化。

3. 扩展的上下文窗口

MemLong在处理长文本方面表现出色:

  • 能够在单个3090 GPU上将上下文窗口扩展到80k个token。
  • 通过仅记忆单层的K-V对,有效控制了内存使用。

这一特性使MemLong在处理长文档、长对话等任务时具有显著优势,能够捕捉更长距离的语义依赖关系。

MemLong架构详解

概述

MemLong的核心架构如图2所示:

MemLong: 基于记忆增强检索的长文本LLM生成方法_大语言模型_02

图2:MemLong架构示例

工作流程主要包括以下步骤:

  1. 输入处理:将输入文本分割成块(ci),每个块对应原始文本ti。
  2. 底层处理:在冻结的底层中应用标准因果注意力机制。
  3. 记忆层操作:在最后一个底层(记忆层)后执行两个关键操作:- 检索:利用ti获取最相关的K-V对。- 缓存:存储获得的K-V对及其相关的块表示。
  4. 上层处理:在可训练的上层中,整合检索到的K-V对与当前输入上下文,并调整模型参数以优化检索过程。

这种设计允许模型有效地利用历史信息,同时保持计算效率和模型的适应性。

在接下来的部分中,我们将深入探讨MemLong的检索器设计、动态记忆管理策略、注意力机制重构以及推理过程等技术细节。

MemLong技术细节

检索器和动态记忆管理

检索过程

MemLong的检索过程是其核心功能之一,具体步骤如下:

  1. 对于每个查询块cq(等同于ci)及其对应的文本块tq(等同于ti),首先通过检索器R处理,得到表示嵌入rq:rq = R(tq), 其中 rq ∈ Rdret
  2. 利用得到的表示嵌入,对记忆M中的嵌入进行检索,获取k个最相关的块级索引。
  3. 计算检索表示rq与记忆M中存储的嵌入之间的余弦相似度。
  4. 选取相似度最高的k个索引:zq = TopK{Cos(rq)}, 其中 zq ∈ Rk
  5. 基于获得的索引,从记忆中检索相应的K-V对:˜zq ∈ Rk×τ×dmodel

这种方法允许模型高效地获取与当前输入最相关的历史信息。

记忆过程

记忆过程是MemLong保持长期记忆的关键,其主要特点包括:

  1. 同步存储:记忆层的K-V对与用于检索的表示嵌入同步存储,确保索引的准确对应。
  2. 双重存储:- 缓存K-V对- 存储相应的表示
  3. 一致性保证:由于较低层在训练期间被冻结,输出K-V对的分布保持一致。
  4. 记忆更新:完成所有块对的检索后,执行记忆操作M(k, v; rm),同步更新记忆中的键值对及其对应表示。
动态记忆更新

为了有效管理长期记忆并防止记忆溢出,MemLong采用了动态记忆更新策略:

  1. 记忆分配:- 保留最新的10%记忆内容(可能相关)- 丢弃最旧的10%(可能已过时)- 根据检索频率优先考虑中间80%
  2. 内存控制:当记忆使用超过阈值时,删除最少访问的条目,直到使用率降至50%。
  3. 平衡机制:这种选择性修剪在保留有价值信息和删除不相关数据之间取得平衡。
  4. 效率保证:动态更新有效控制记忆大小,防止内存不足,同时保持检索效率。

注意力机制重构

MemLong对传统注意力机制进行了创新性的重构,以更好地融合长期记忆:

检索因果注意力
  1. 联合注意力机制:将标准注意力扩展为包含本地和长期记忆的联合机制。
  2. 长期记忆融合:使每个token能够同时关注局部上下文和具有完整连续语义的块级历史上下文。
  3. 双向注意力:通过检索方法获得的块级K-V对支持双向注意力,同时避免信息泄露。

图3展示了检索因果注意力的工作原理:

MemLong: 基于记忆增强检索的长文本LLM生成方法_人工智能_03

图3:检索因果注意力示意图

注意力计算

给定前一层的头部隐藏状态输出Hl−1 ∈ R|x|×dmodel和检索到的键值对˜zq = { ˜Ki, ˜ Vi} ωi=1 ∈ Rk×τ×dmodel,下一层的输出隐藏状态Hl的计算如下:

MemLong: 基于记忆增强检索的长文本LLM生成方法_LLM_04

为避免训练初期检索注意力分数Sm可能造成的干扰,MemLong采用多头注意力机制:

MemLong: 基于记忆增强检索的长文本LLM生成方法_大语言模型_05

最终,通过连接˜V和V得到Hl:

MemLong: 基于记忆增强检索的长文本LLM生成方法_RAG_06

这种注意力机制设计允许模型有效地整合局部和长期信息,提高了对长距离依赖的建模能力。

MemLong推理过程

MemLong在处理超长输入时采用了特殊的推理策略:

  1. 输入分割:将输入分为前缀和主体两部分。
  2. 前缀处理:- 将前缀分成多个非重叠的块- 通过记忆层处理每个块,确保参与注意力的token数量等于块大小- 注意块间的相互关联性(第t个块需要处理前t-1个块的信息)
  3. 主体处理:- 为主体选择k个最相关的块,基于块级检索表示- 获取这些块的键和值表示
  4. 上层检索:在上层检索层中,注意力窗口大小为k * τ,通常小于输入总长度。
  5. 高效执行:同时进行长度受限的因果注意力和检索注意力计算。

这种策略使MemLong能够高效处理远超其预训练长度的输入,同时保持计算效率和内存使用的平衡。

在下一部分中,我们将详细介绍MemLong的实验设置和结果,展示其在各种长文本任务中的优越性能。

MemLong实验评估与结论

实验设置

实现细节
训练配置
  • 基础模型:使用带有旋转位置编码的OpenLLaMA-3B作为预训练骨干网络。
  • 训练规模:在0.5B个token上进行检索增强适应训练,序列长度设置为1024。
位置编码策略

为了适应MemLong的特殊架构,采用了以下位置编码策略:

  1. 局部上下文(最多2048个token):使用标准的旋转位置编码。
  2. 记忆键:被编码为在局部上下文窗口中位置为0。

这种策略确保了模型能够有效区分和利用局部和长期信息。

评估数据集

为全面评估MemLong的性能,选择了四个广泛使用的长文本基准数据集:

  1. PG-19:英语图书数据集
  2. BookCorpus:另一个英语图书数据集
  3. Wikitext-103:维基百科文章集
  4. Proof-Pile:数学论文集

这些数据集涵盖了不同领域和文体的长文本,为模型的泛化能力提供了全面的测试。

评估指标

主要使用困惑度(Perplexity, PPL)作为评估指标。对每个序列的最后2048个token计算PPL,较低的PPL值表示更强的语言建模能力。

基线模型

选择OpenLLaMA-3B作为主要基线模型,同时与LLaMA-2-7B等其他模型进行了对比。

实验结果与分析

长上下文语言建模

图7展示了MemLong与基线模型在不同数据集和输入长度下的性能比较:

MemLong: 基于记忆增强检索的长文本LLM生成方法_大语言模型_07

图7:不同模型在各数据集上的困惑度(PPL)比较

关键发现:

  1. 预训练长度内的性能:在不超过预训练长度(OpenLLaMA-3B为2048,LLaMA-2-7B为4096)的测试中,MemLong与完全微调的模型表现相当。
  2. 超长输入处理能力:- MemLong在超过预训练和微调长度的输入上仍能持续降低困惑度,展现出优秀的泛化能力。- 相比之下,OpenLLaMA-3B和LLaMA-2-7B在超出预训练长度的输入上表现明显下降。
  3. 计算效率:- MemLong通过外部检索器有效处理长输入,同时保持较低的内存开销。- 传统模型在处理超长输入时,由于注意力机制的二次复杂度,内存消耗显著增加。
  4. 一致性改进:MemLong在各个数据集上都实现了困惑度的降低,表明其方法的有效性和稳定性。
上下文学习能力

表1展示了MemLong在5个NLU任务上的少样本学习(4-shot和20-shot)性能:

MemLong: 基于记忆增强检索的长文本LLM生成方法_人工智能_08

表1:不同模型在NLU任务上的少样本学习准确率(%)

关键观察:

  1. 记忆利用:与仅依赖非参数知识的OpenLLaMA相比,MemLong能有效利用存储在记忆中的额外示例,提高性能。
  2. 示例数量影响:随着记忆中示例数量的增加,MemLong的性能进一步提升或保持稳定,显示出良好的可扩展性。
  3. 与LongLLaMA对比:在相同的内存演示条件下,MemLong在大多数数据集上优于LongLLaMA,证明了其记忆机制的优越性。
  4. 任务适应性:MemLong在不同类型的NLU任务中均表现出色,表明其具有广泛的应用潜力。

局限性讨论

尽管MemLong展现出优秀的性能,研究团队也认识到以下局限性:

  1. 模型规模限制:当前研究主要集中在OpenLLaMA-3B上,未来需要在更大规模的模型上进行验证。
  2. 信息稳定性:观察到单层K-V对提供的额外语义信息可能存在不稳定性,这一问题需要在未来的研究中进一步探讨和解决。
  3. 领域特异性:虽然在多个数据集上进行了测试,但可能还需要在更多特定领域的任务中评估MemLong的表现。
  4. 计算复杂度:虽然MemLong相对高效,但在处理极长文本时的计算需求仍需进一步优化。

总结

MemLong作为一种创新的长文本处理方法,在以下方面展现了显著优势:

  1. 上下文扩展:成功将模型的有效上下文窗口从2k扩展到80k个token,大幅提升了长文本处理能力。
  2. 性能提升:在长距离文本建模和理解任务中,MemLong展现出明显的竞争优势,相比全上下文模型,在某些任务中实现了高达10.4个百分点的性能改进。
  3. 效率平衡:通过创新的记忆和检索机制,MemLong在保持高性能的同时,有效控制了计算复杂度和内存使用。
  4. 广泛适用性:在多种NLU任务和不同类型的长文本数据集上均表现出色,显示出良好的泛化能力。

MemLong的成功为长文本语言模型的发展开辟了新的方向,其核心思想和技术可能对未来的大规模语言模型设计产生深远影响。未来的研究方向可能包括进一步优化记忆机制、探索更大规模模型的应用,以及在更多领域特定任务中的表现验证。


标签:检索,MemLong,模型,记忆,LLM,上下文,注意力
From: https://blog.51cto.com/deephub/12100844

相关文章

  • 自学大模型LLM,看这本书就够了!帮你快速入门大模型!(含PDF)
    一、Transformer模型介绍Transformer模型是一种基于自注意力机制的深度神经网络模型,在自然语言处理(NLP)领域取得了显著的成果,并在机器翻译、文本摘要、问答系统等多个任务中表现出色。1.基本原理Transformer模型主要由自注意力机制和前馈神经网络两部分组成。(1)自注意力机制:功能:自注......
  • EI Compendex、Scopus检索-国际会议论文征稿:第三届传感器技术与控制国际研讨会(ISSTC 2
    收录率高丨往届均已EI检索 | 往届会后3个月EI检索|IEEE出版第三届传感器技术与控制国际研讨会(ISSTC2024)20243rd InternationalSymposiumonSensorTechnologyandControl2024年10月25-27日,珠海大会官网:http://www.isstc.net/【论文投稿】主办单位联合主办......
  • 【图像检索】基于纹理(LBP)和形状特征的图像检索,matlab实现
       博主简介:matlab图像代码项目合作(扣扣:3249726188)~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~       本次案例是基于纹理(LBP)和形状特征(hu特征)的图像检索,用matlab实现。 一、案例背景和算法介绍       这次博客文......
  • 2025秋招LLM大模型多模态面试题(八)- langchain完整面试题
    目录什么是LangChainLangChain包含哪些核心模块模型输入/输出(ModelI/O)组件管理数据处理链式组合记忆与上下文管理外部集成一些核心概念什么是LangChainAgent?什么是LangChainmodel?LangChain包含哪些特点?LangChain如何使用?LangChain如何调用......
  • LLM正当时,火山引擎数智平台如何赋能AIGC增长?
    更多技术交流、求职机会,欢迎关注字节跳动数据平台微信公众号,回复【1】加入官方交流群本文将从推广获客、产品优化、用户运营三个角度,分享火山引擎数智平台VeDI旗下营销套件如何多方面为LLM领域的相关业务提供深度支持,帮助大模型应用公司在应用推广环节实现健康快速增长。前言当......
  • 2024最新LLM大模型学习:从零开始精通AI大模型,全面详细的入门指南,非常高详细收藏我这一
    一.初聊大模型1.为什么要学习大模型?在学习大模型之前,你不必担心自己缺乏相关知识或认为这太难。我坚信,只要你有学习的意愿并付出努力,你就能够掌握大模型,并能够用它们完成许多有意义的事情。在这个快速变化的时代,虽然新技术和概念不断涌现,但希望你能静下心来,踏实地学习。一......
  • LLM大模型: Denoising Diffusion Probabilistic Models 原理解析与核心代码
      根据文本生成图片是AI的核心应用之一,2020年后主流的生成方式都是基于DenoisingDiffusionProbabilisticModels原理的,逐渐替代了之前使用GAN的方式生成图片!那么DDPM为啥能取代GAN了?其优势在哪?或者说GAN的劣势在哪?  1、CLIP模型都知道吧?text和image都通过各自的enco......
  • OpenAI o1模型揭秘:通过LLMs学习推理能力
    OpenAI推出了o1,这是一种通过强化学习训练的大型语言模型,专门用于进行复杂的推理任务。o1在回答问题之前会“思考”,能够在响应用户之前生成一条长的内部思维链。在编程竞赛问题(Codeforces)中,OpenAIo1的排名在89%分位,位列美国数学奥林匹克预选赛(AIME)前500名学生之列,并且在物理、生......
  • OpenAI o1模型揭秘:通过LLMs学习推理能力
    OpenAI推出了o1,这是一种通过强化学习训练的大型语言模型,专门用于进行复杂的推理任务。o1在回答问题之前会“思考”,能够在响应用户之前生成一条长的内部思维链。在编程竞赛问题(Codeforces)中,OpenAIo1的排名在89%分位,位列美国数学奥林匹克预选赛(AIME)前500名学生之列,并且在物理、生......
  • EI, Scopus检索 | 计算机领域国际会议论文征稿:第三届信息学,网络与计算技术国际学术会
    河南大学主办,SPIE独立出版,连续2届稳定EI检索! 往届录用稿件已100%EI检索--见刊后1个月检索-高录用-稳检索!第三届信息学,网络与计算技术国际学术会议(ICINC2024)20243rd InternationalConferenceonInformatics,NetworkingandComputing大会时间:2024年10月25-27日......