用于长文本建模的记忆增强检索 MemLong

​一、论断写在前面

论文题目：MemLong: Memory-Augmented Retrieval for Long Text Modeling

论文链接：

LLMs在各个畛域的最新停顿取得了清楚的成功。但是，由于留意力机制的二次期间和空间复杂性以及生成环节中键值缓存的内存消耗始终参与，处置长高低文依然是LLMs的一个严重应战。

论文提出了MemLong，一种高效且轻量化的方法，用于裁减大型言语模型（LLMs）的高低文窗口。其外围现实是将过去的高低文和常识存储在一个无法训练的记忆库中，并进一步应用这些存储的嵌入来检索块级别的键值（K-V）对，以输入到模型中。MemLong实用于任何仅解码器的预训练言语模型，经过结合（1）一个额外的记忆检索组件用于记忆和检索，以及( 2 )一个检索因果留意力模块用于整合部分和记忆消息。MemLong的记忆和检索环节如图1(b)所示。在生成环节中，超出模型最大处置长度的文本被存储为高低文消息在记忆库中。随后，在长文档中生成最近的一个文本块时，论文经常使用检索器显式地检索过去的关系消息，经过索引对齐失掉额外的高低文消息。

MemLong提供了几个长处：（1）散布分歧性：与之前在消息存储到记忆中时阅历散布偏移的模型不同，MemLong确保缓存消息的散布坚持分歧。（2）训练高效：论文解冻模型的底层，只有微调高层，这大大缩小了计算老本。在论文的试验中，在0.5B个token上微调一个3B参数版本的MemLong仅需八块3090 GPU运转八小时。（3）宽泛的高低文窗口：由于只有记忆单层的K-V对，MemLong能够在单块3090 GPU上轻松裁减高低文窗口至80k个token。

少量试验标明，与其余上游的LLMs相比，MemLong在多个方面体现出优越的性能。在多个长高低白话语建模数据集上，MemLong优于OpenLLaMA和其余基于检索的模型。在检索增强的高低文学习义务中，MemLong比OpenLLaMA提高了多达10.2个百分点。

二、论文的便捷引见

2.1 论文的背景

由于传统的留意力机制的二次方期间和空间复杂性，裁减高低文长度具有应战性，这为触及长序列义务的运行带来了清楚的限度，例如长文档摘要和多轮对话。因此，LLMs理论被希冀具有常年间的上班才干（即长高低文LLMs），以有效应答这些厚道的场景。

图1：检索增强生成(RAG)和MemLong的记忆检索流程示用意。(a) 当检索消息的长度超越模型处置才干时，RAG甚至或者降低生成性能（黄色）。(b) 论文的方法应用外部检索器失掉历史消息，而后将其作为K-v对传递给模型，而不是以文本方式传递。

为了处置计算瓶颈疑问，曾经启动了少量致力。第一类上班并重于经过驳回稠密留意力操作来缩小个别留意力机制的计算量。虽然这些方法可以将计算复杂度降低到大概 O(n)，但理论会以模型容量为代价。因此，一些上班转向了记忆选用。这些方法，作为token级别的记忆选用，或者造成语义消息的截断。另一类最近的上班是检索增强言语建模。这些上班理论引入检索机制来增强模型处置长文本的才干。

但是，这些方法有几个缺陷。首先，由于训练环节中模型参数的变动，存储在记忆中的消息或者会阅历散布偏移。其次，这些方法理论须要从新训练，这在大型模型时代是不实际践的。最后，这些模型往往偏差于以就义预训练模型的原始才干为代价来处置长文本输入。为了处置先前钻研的局限性，论文提出了以下疑问：论文能否应用检索器的显式检索才干来近似模型外部的隐式检索环节？

2.2筹备常识

2.2.1 义务定义

言语模型旨在定义一系列token的概率散布，从而有效地预测给定言语中序列的或者性。给定一个序列 x_1, ..., x_n，规范方法建模其概率为 p ( x_1, ..., x_n )= sum pθ ( x_i | x< i )，其中 x< i , := x_1, ..., x_i-1 是 x_i 之前的token序列。

与规范言语建模目的不同，论文不只经常使用高低文启动下一个token预测，还应用外部检索失掉关系消息并在模型的高层启动常识融合。详细来说，给定一个由 l个token组成的序列，每个块的大小为 ν=l/τ，论文将序列划分为 τ个不堆叠的块，记为 C=( c_1, ..., cν )。相应地，其文本方式被划分为ν个文本块，记为 T=( t_1, ..., t_ν)。在每一步中，论文在高层对 c_i 启动因果言语建模，而在高层，论文对 t_i 启动细粒度可控检索以融合额外消息。成功此操作后，论文的言语建模目的变为

其中 R( t_i) 表示检索 t_i 所在块的临近块。

2.2.2 模块与操作定义

如图 2 所示，Ret-Mem 模块由一个检索器（Retriever）和一个记忆组件（Memory）组成，用于消息替换。首先，论文将记忆组件定义为M，将检索器定义为 R，以及它们对应的操作M(.) 和 R(.)。此外，论文指定模型的维度为 d_model，检索器的维度作为 d_ret 。记忆模块包含两个部分：K-V 对和相应的表示嵌入。键和值的维度表示为R^d_model，而嵌入的维度表示为R^d_ret。必定强调的是，实践的检索环节触及表示块的嵌入，而不是 K-V 对。检索器实质上是一个预训练的密集嵌入器，具有杰出的表示才干。MemLong 经常使用它将每个块编码为表示嵌入。由于它为每个块生成一维表示向量，即使内存大小很大，内存占用也坚持最小。

图 2 ：MemLong 的一个示例：在较低层，模型坚持静态，对整个块 c_i口头因果言语建模，随后，c_i 以嵌入和 K-V 对的方式缓存。最后，高层被微调以协调检索偏好并整合检索到的内容。

2.3 MemLong

2.3.1 概述

如图 2 所示，每个步骤触及一个块 c_i 的输入，其中该块的原始文本为 t_i。在模型解冻的较低层，对整个 c_i 运行规范因果留意力。关于较低层的最后一层，论文称之为记忆层。在每次遍历记忆层后，口头两个关键操作。第一个操作是检索，由红线表示，其中 t_i 用于失掉最关系的 K-V 对。第二个操作，由蓝线表示，触及缓存取得的 K-V 对及其关系块表示。在模型的高层，检索到的 K-V 对与输入高低文整合，随后调整模型参数以校准检索参考。后续部分将讨论 MemLong 框架的各个方面及其细节，包含检索器和灵活内存治理 (Dynamic Memory Management)，留意力重构 ( Attention Reformulation )，以及经常使用 MemLong 启动推理 (Inference with MemLong)。

2.3.2 检索器与灵活内存治理

论文提供了一个关于检索环节和内存治理灵活的片面解释。

检索环节。鉴于论文的目的是用显式检索要代基于K-V对的传统kNN检索，论文旨在在可行的状况下预取所需消息，而后再缓存模型输入。详细来说，关于每个潜在的查问块c^q=c_i及其对应的文本块t^q=t_i，论文首先将其传递给检索器，而后取得一个表示嵌入r^q = R ( t^q )。随后，论文经常使用这个表示嵌入对\mathcal{M}中的嵌入启动检索，以失掉所需的k个块级索引。论文计算检索表示r^q与存储在内存M中的嵌入之间的余弦相似度。最后，论文失掉c^q的top-k索引z^q =Top K{Cos( r^q ) }。

记忆环节。记忆环节同步存储来自记忆层的K-V对以及先前计算的用于检索的表示嵌入，确保\mathrm{K-V}对的索引与其表示嵌入准确对应（见图2，右侧，蓝线）。关于每个或者的块记忆c^m=c_i及其对应的文本块t^m=t_i，论文将记忆环节分为两部分：第一部分详细说明如何缓存K-V对，第二部合成释如何存储相应的表示。首先，论文将c^m输入到MemLong中，并从记忆层失掉输入。论文的记忆操作十分高效，由于它仅触及存储检索所需的表示r^m=r^q，从而防止了冗余。在一切块对检索成功后，记忆操作——表示为M ( k, v ; r^m )——同步降级记忆，包含键值对及其对应的表示。

灵活记忆降级。当记忆溢出时，论文经常使用计数器智能降级记忆。在论文的试验中，论文保管最新10%的记忆内容，因其或者具无关系性，摈弃最旧的10%，因其或者已过期，并依据检索频率优先处置两边的80%，删除起码访问的条目，直到记忆经常使用量降至50%。这种选用性修剪平衡了时效性和关系性，保管了有价值的消息并删除了不太关系的数据。

图3：检索因果留意力示用意。部分因果留意力运行于最近的高低文，而经过检索方法取得的块级K-V对由于其历史性质而能够成功双向留意力，而不会出现消息暴露。

2.3.3留意力重构

在模型的可训练高层中，论文重构了留意力机制以融合常年记忆。如图3所示，与传统的Transformer解码层经常使用多头留意力不同，论文提出了一种检索因果留意力机制，将其裁减为联结留意力机制，并提出了一种常年记忆融合环节，使得每个token既能关注部分高低文，也能关注具有完整和延续语义的块级过去高低文。下一层的隐藏形态H^l计算如下：

为了防止训练初期检索留意力分数o_m的搅扰，论文驳回了多留意力机制，遵照LLaMA-adapter

最后，论文将V和V衔接起来失掉H^l：

2.3.4 经常使用MemLong启动推理

当MemLong接纳到超越长度的输入时，论文将其视为两个部分：前缀p和主体。论文将区分形容在推理阶段对长输入的编码和长输入的生成。当MemLong接纳到长输入时，它首先将前缀分红多个不堆叠的块，并从其内存层计算，这确保了介入留意力的token数量等于块大小，远小于输入长度。须要留意的是，每个块是互关系联的（例如，第t个块须要处置前t-1个块的）。

第二步是依据块级检索表示选用与主内容最关系的k个块，并失掉它们的关键和值表示。在此之后，关于高层检索层，检索的留意力窗口相当于k * τ，这也小于输入长度。最后，高效地口头长度受限的因果留意和检索留意。

2.4 试验

论文在须要内存中长高低文处置的各项义务上评价论文提出的MemLong模型：(a) 长高低白话语建模和检索增强言语建模；(b) 能够处置内存中少量演示示例的可裁减高低文学习。

2.4.1 成功细节

训练细节。论文经常使用OpenLLaMA-3B作为预训练的主干LLM，驳回旋转位置编码（rotation position coding）。由于配件限度，论文选用经常使用LoRA技术训练论文的模型。主干LLM具有L=26, H=32, d=100的架构。除非另有说明，论文经常使用第13层作为内存层，[14, 18, 22, 26]层作为检索增强层。检索增强顺应的训练仅在0.5B个token上迭代，序列长度为1024。Mem-Long的可训练参数来自14到26层。论文应用slimpajama数据集采样作为论文的训练语料库。

位置重映射。在生成环节中，M中检索到的块级别K-V有多个。由于每一步检索的不确定性，论文须要将位置嵌入重映射到检索到的块。与之前的上班（Tworkowski et al., 2024）相反，部分高低文（最多2048个token）接纳规范旋转位置编码，而内存键则被编码为在部分高低文窗口中具有位置0。

2.4.2 长高低白话语建模

论文首先在长高低白话语建模基准上评价MemLong，以评价其基本的言语建模才干。由于K-V缓存提供了清楚的背景和高低文消息，MemLong能够极速检索关系的K-V缓存并充沛应用它，从而在长高低文建模义务中增强模型的性能。

数据集。论文在四个宽泛的文本基准数据集上对论文的模型启动了评价：英语书籍PG-19和BookCorpus，维基百科文章Wikitext-103，以及数学论文Proof-Pile。试验结果标明，一切数据集上的困惑度都有清楚改善。论文的模型在从1024到32768个token的不同长度范围内启动了测试。经过应用外部检索器和内存，论文的模型在一切数据集上展现了清楚的性能优化，且内存开支最小。

设置。依照（Yen et al., 2024），论文计算每个序列最后2048个token的困惑度。此试验设置旨在验证不同检索器大小对模型全体性能的影响。为了成功高效的细粒度检索，论文经常使用faiss工具包在GPU上构建准确搜查索引，以存储文本块的表示嵌入并口头高效检索。关于MemLong，论文将token宰割并放入finetune-length = 1024的M中，用于进一步检索。

基线。在论文的试验中，论文驳回OpenLLaMA-3B模型作为基线。为了确保偏心比拟，论文经常使用相反的LoRA性能并微调了模型在相反数量的slimpajama数据集上的数据。此外，论文比拟了LongLLaMA-3B，该模型经常使用Focused Transformer（FoT）方法和5B token启动了微调。为了启动更片面的比拟，论文还测试了两个7B模型：LLaMA-2-7B和LongLoRA-7B-32K，以及两个位置编码模型：Yarn-7b-128k和Phi3-128k。

表1：不同高低文窗口裁减模型在PG19、Proof-pile、BookCorpus、Wikitext-103上的滑动窗口困惑度。一切试验都在一块3090 24GB GPU上启动。LongLLaMA-3B和MemLong-3Btoken为表示在没有Memory的状况下评价，LongLLaMA-3Btoken表示在有限Memory的状况下评价。论文还评价了MemLong在4K/32K Memory场景下的体现。"- / 6.95"表示模型在单GPU上造成内存无余（OOM）失误，而在双GPU上则发生相应结果。*

结果。结果如表1所示。论文驳回困惑度（PPL）作为言语模型的评价目的。较低的PPL表示更强的言语建模才干。与两个齐全微调的模型相比，OpenLLaMA-3B和LLaMA-2-7B，论文的模型在测试长度在其预训练限度内（OpenLLaMA-3B为2048，LLaMA-2-7B为4096）时，在多个数据集上体现出相当的性能。

但是，一旦测试长度超越这些预训练限度，论文的模型即使在微调长度1024和预训练长度2048之后，仍能继续降低困惑度，展现了其优越的泛化才干。

相比之下，OpenLLaMA-3B和LLaMA-2-7B模型无法泛化到超出其预训练长度的输入，并且由于留意力机制的二次复杂性，体现出清楚参与的内存开支。论文还与LongLoRA启动了比拟。虽然LongLoRA中提出的Shifted Sparse Attention清楚缩小了内存经常使用，但它也削弱了模型在短文本上的性能。

相比之下，LongLLaMA由于其内存经常使用有限增长，在测试长度变得过长时也会遇到OOM疑问。位置编码模型具有弱小的泛化才干。但是，这种方法的性能只能保障长距离生成性能不降低。与这些方法相比，MemLong应用外部检索器处置更长的输入token，并成功了更好的困惑度改良。同时，由于高存储效率，MemLong可以有效控制GPU的经常使用，防止OOM疑问。

2.4.3 高低文学习

传统的高低文学习（ICL）将大批的非参数化示例与查问一起输入模型。但是，这些方法理论受限于模型的输入长度。在本试验中，由于MemLong可以将示例以参数化方式存储在其记忆中，论文重要钻研MemLong能否能有效应用其记忆中存储的常识以增强其突现才干。结果如表2所示。

与仅依赖非参数化常识的OpenLLaMA相比，在相反数量的高低文示例下，MemLong可以应用其记忆中存储的额外示例。随着记忆中示例数量的参与，性能进一步优化或坚持稳固。在与LongLLaMA的对比剖析中，论文观察到在相反的保管内存示例条件下，论文的模型在大少数数据集上体现优于LongLLaMA。值得留意的是，与LongLLaMA相比，论文的模型在训练参数（2亿对比0.3亿）和微调数据量（0.5亿对比5亿）方面清楚缩小。这凸显了论文模型在应用外部检索器启动消息失掉方面的效率，展现了在资源大幅缩小的状况下，能够更有效地综合和应用常识的才干。

2.5 Ablation Study

2.5.1 训练设置

在训练阶段，论文讨论了不同检索层对模型的影响，并调查了论文的方法能否能充沛处置MemTrm中讨论的散布偏移疑问。如前所述，论文的方法为散布偏移提供了一种低老本的处置打算。如图4所示，棕色线（图片顶部的线条；训练方法相似于MemTrm，微调模型的一切参数，并且在内存层之后的一切层都介入检索）在性能和拟合速度方面清楚劣于论文一切其余方法（即使是设置最不正当的方法）。论文将在稍后剖析推理阶段的性能。

2.5.2 推理性能

图 4：训练阶段PPL的变动水平。y轴的目的为PPL。论文重要关注训练参数和检索层。

Q1：记忆长度能否影响模型的性能？如图5所示，论文对同一模型在不同记忆大小下的性能启动了审核，结果标明记忆容量与模型效率之间存在清楚的关系性。趋向标明，记忆大小的参与会逐渐优化性能。此外，在记忆大小为65536时，模型才干阅历了一个清楚的飞跃。这标明，虽然裁减记忆提供了实质性的好处，但其有效性存在一个实践的下限，这或者遭到数据散布纤细差异的影响。

Q2：论文须要引入多少层额外的记忆消息？如图4所示（粉色线）和表3（RPL+TH）中展现的，当检索层的数量设置为[13,17,21,25]时，模型体现最佳。依据阅历以为，假设在模型的一切高层都引入检索消息，会造成模型对部分高低文的留意力降低。因此，以适当的距离选用检索层实践上可以增强模型的才干。

表3：不同的检索层会影响MemLong的性能。token为的MemLong表示在没有Memory的状况下启动评价。一切经常使用Memory的方法的大小设置为32768。RA表示跨一切高层检索；TA表示不解冻参数的训练；RP表示跨较少高层检索，RPL表示跨更少高层检索。

图 5：在不同记忆大小下评价不同数据集。在每个子图中，除记忆大小外，一切参数均相反。

本文转载自​​，作者：​​