长上下文训练方法的探究

为什么要写下长上下文训练方法的探究这篇blog呢，是因为博主研究生的研究方向是长视频（小时级别长度）的时序定位，正在通过阅读不同方向的任务来获得灵感。偶然间发现目前火热的大语言模型的方向出现了“大海捞针“（Needle in a Haystack，NIAH）任务，与长视频时序定位这个任务在概念上非常接近，希望能从该任务中找到一些模型设计或者模型训练方法上的灵感。

既然如此，就先介绍一下NIAH这个任务，并延续探究在这个任务上大模型处理长上下文的方法。

Needle in a Haystack

任务描述

在2022年底chatgpt横空出世之后，大语言模型变得异常火热。然而，当时的gpt的上下文长度仅有4k tokens，这限制了其处理大文档的效果，在其后短时间内推出的各种模型的上下文长度也都不超过128k，大部分在8k以内。同时，虽然模型号称的长度越来越长，但模型的性能却不一定跟上，比如说和模型一开始约定好一个格式，随着推理的进行，模型却忘掉了要遵守的格式；或者说模型忘掉了之前推理中提到的一个信息。基于这一痛点，gkamradt等人提出了大海捞针任务（LLMTest_NeedleInAHaystack: Doing simple retrieval from LLM models at various context lengths to measure accuracy），旨在测试大模型对于长文本中事实的检索效果。
具体来说分为以下几步：

1.Place a random fact or statement (the ‘needle’) in the middle of a long context window (the ‘haystack’) 在长上下文窗口（“草堆”）中间放置一个随机的事实或陈述（“针”）

2.Ask the model to retrieve this statement 要求模型检索此声明

3.Iterate over various document depths (where the needle is placed) and context lengths to measure performance 迭代各种文档深度（放置针的位置）和上下文长度以测量性能

（补充LLMTest_NeedleInAHaystack相关材料，这篇文章里的内容是怎么做的）

任务发展

大海捞针任务后来还有了一些改进，如RULER和MMNeedle

RULER（RULER: What’s the Real Context Size of Your Long-Context Language Models?），增强了大海捞针的检索任务（Retrieval），还添加了多跳跟踪（Multi-hop Tracing）、信息聚合（Aggregation）、带干扰的问答（QA）三种任务；

具体的bench的实操细节详看论文，这张图有几个有意思的点可以关注一下：

1、图4中的第一张和第二张图对比显示，总体而言，较大的上下文大小会带来更好的性能，但对于长序列，排名可能不一致。例如，使用1M上下文大小（LWM-1M）训练的模型比长度为256K、使用512K的模型差，这可能是由于在RoPE中调整新基频的训练不足。此外，我们观察到，当模型需要外推到看不见的长度时（例如，给定256K的输入，LMW-128K），性能会突然下降，并且在最大训练上下文大小内，输入长度在对数尺度上几乎呈线性下降。

2、为了消除模型大小的影响，我们评估了Yi-34B-200k、Yi-9B-200k和Yi-6B-200k，它们都使用相同的数据混合训练到相同的上下文长度。图4（右中）显示，在RULER上，34B模型在4K长度下的性能和相对退化方面明显优于6B模型，这表明缩放模型大小有利于更好的长上下文建模。

3、我们评估了两种具有非Transformer架构的模型的有效上下文长度：RWKV-v5（Peng等人，2023）和Mamba-2.8B-limpj（Gu和Dao，2023）。我们发现，当将上下文大小扩展到8K时，这两个模型都表现出明显的退化，并且在4K长度之前，它们的性能都远远低于Transformer基线Llama2-7B，超过4K长度后，Llama2的长度外推性能较差（图4右）。

多模态大海捞针MMNeedle（Multimodal Needle in a Haystack: Benchmarking Long-Context Capability of Multimodal Large Language Models），专门设计用于评估 MLLMs 的长上下文能力。除了多图像输入外，我们还采用图像拼接进一步增加输入上下文长度，并开发了一种协议来自动生成子图像级别检索的标签。具体而言就是设置一组图像（M，N，N），其中M是拼接图像的序列号，N*N是这个序列中的子图数量。将N*N个子图拼接成一个大图，反复操作得到多个大图的图像序列形成最终的Haystack。给予模型某张子图的caption输入，从Haystack中寻找到和这个caption相对应的图像序列，称之为Multimodal Needle。为了确保评估的科学性和可重复性，其根据模型上下文最长支持长度设置了图像拼接的最大分辨率2048和子图分辨率的下限256。

本质上，MMNeedle 通过压力测试 MLLMs 在基于文本指令和图像内容描述的情况下，在一组图像（干草堆）中定位目标子图像（针）的能力来评估它们。

同期复旦和上海ai lab也提出了MM-NIAH（Needle In A Multimodal Haystack）

长上下文训练方案

长上下文模型的训练方法

从上述的一些长上下文能力测试的benchmark中可以发现，大语言模型对于长上下文中的定位能力不足。针对这个问题，有不少研究者也在探寻解决的方法。文本类的大模型处理长上下文的方法很新颖，且已经推广到应用级别，比如Kimi，GLM，DeepSeek，Gemini等，一般都需要查看对应大模型研究团队自己开源的技术报告：

GLMLong 技术报告

ChatGLM官方在24年7月给出了GLMLong的训练方法（GLM Long: Scaling Pre-trained Model Contexts to Millions），写成了技术报告的形式发表在Medium平台上。其中ChatGLM-9B-Chat-1M模型在四项长上下文检测的benchmark（LongBench-Chat、InfiniteBench、RULER、NIAH）中均取得同量级模型参数下最好的结果。GLM的训练过程指出长上下文的能力不是一蹴而就的，是需要经历多个阶段才能逐步激活和保持模型的长文本能力。下面这张图给出了GLM逐步获得长上下文能力的训练流程：

首先使用大量的文本训练数据（约10T个token）进行预训练，获得一个基础模型。然后利用少部分长文本数据（约10B+10B=20B个token）进行进一步预训练，使模型获得处理长文本的能力（原文的话是：This step is crucial for stimulating the model’s ability to handle long texts.）。考虑到从基础模型8k的上下文长度直接扩展到1M的跨度过大，采取分阶段预训练：第一阶段扩展到128k的上下文长度，第二阶段扩展到1M的上下文长度。

1、数据处理阶段

为了保持模型general processing的能力，同时又具有处理长文本的能力，在两阶段的训练过程中，需要对预训练数据进行混合采样技术。具体来说，在第一阶段扩展到128k长度中，原始预训练数据包含4B个token；在原始预训练数据中对于长度超过8K的数据进行上采样，大约3B个token；在原始预训练数据中对于长度超过32K的数据进行上采样，大约3B个token，总共约10B个token的数据。在处理这两种类型的上采样数据时，努力确保每个长度区间内的token总数保持一致。例如，28K至36K区间和60K至68K区间内的标记总数应大致相同。在这一平衡原则下，60K至68K区间内的数据项数量约为28K至36K区间内数量的一半。————————>这里有个疑惑，这里的上采样是怎么做的？和后文提及到的情况一样吗在第二阶段扩展到1M的训练时，训练数据非常稀少且分布极不均匀，主要来自书籍。为了解决这个问题，引入人工长文本数据，通过预训练语料库中的数据和公开可用的数据进行拼接生成，为了保证拼接的合理性，设计了一种两阶段的聚类策略：首先对文档数据进行抽样，并使用语言模型对其进行分类，然后在每个分类下使用embedding model获得数据的向量表示。通过这种方法，确保了8K-1M长度内token数量的平衡，生成所需的人工长文本数据。最终通过两阶段的策略，预训练的数据分布情况如下：原始预训练数据3.5B个token，上采样的长度超过8k的数据1.5B个token，上采样的长度超过32k的数据1.5B个token，上采样的长度超过128k的数据1.5B个token，人工生成的长文本数据2B个token。

2、训练阶段

在持续预训练阶段，采用了与预训练阶段相拥的方法，但是选择了一种更简洁的位置编码方法，调整扩展RoPE Base（因为有研究表明，虽然不同的RoPE变体在无需训练直接外推的场景下差异显著，但是在经过充分预训练之后性能差异较小）。

与Llama 3的预训练阶段类似，我们采用了带有注意力分离的打包训练策略。以第一阶段128K的上下文窗口为例，每个训练样本可能包含多个不同长度的独立文本，在注意力计算过程中它们互不干扰（通过基于Flash Attention的Varlen实现）。这种注意力分离对于激活模型处理长文本的能力至关重要。与直接应用128K的全注意力相比，它有效地避免了建立许多无效的长距离依赖关系。

Qwen简单提了一下Qwen/Qwen2.5-72B-Instruct · Hugging Face — Qwen/Qwen2.5-72B-Instruct · Hugging Face

Jamba [2408.12570] Jamba-1.5: Hybrid Transformer-Mamba Models at Scale

Deepseek [2412.19437] DeepSeek-V3 Technical Report4.3部分

Prolong [2410.02660] How to Train Long-Context Language Models (Effectively)

MegaBeam aws-prototyping/MegaBeam-Mistral-7B-512k · Hugging Face似乎也没有讲长上下文

InternVideo2.5 视频级别的超长上下文定位能力

长上下文模型的技术现状

由于工程技术、模型架构设计和算法设计的进步，最近引入了许多长上下文语言模型。

FlashAttention（Dao等人，2022；Dao，2023）和RingAttention（Liu等人，2023）显著减少了处理长上下文所需的内存占用。

采用各种稀疏注意力机制（Child等人，2019；Jaszczur等人，2021），如转移稀疏注意力（Chen等人，2024）、扩展注意力（Ding等人，2023）和注意力汇（Han等人，2023；Xiao等人，2024b），以实现有效的上下文缩放。

提出了新的位置嵌入方法来改进Transformers中的长度外推（Vaswani等人，2017），包括ALiBi（Press等人，2022）、xPOS（Sun等人，2023b）和RoPE（Su等人，2023）变体（Chen等人，2023；Xiong等人，2023，Peng等人，2024，Liu等人，2024b，Ding等人，2024，Zhu等人，2024）。（该方法可以看一下本博主写的另一篇blog

另一项研究侧重于减少上下文大小。这可以通过使用递归机制缓存先前的上下文（Zhang等人，2024a；Bulatov等人，2023；Martins等人，2022；Wu等人，2022）、从上下文中检索相关信息（Xu等人，2024a；Mohtashami&Jaggi，2023；Wang等人，2024；Tworkowski等人，2024a；Xiao等人，2024B）或通过压缩保留显著信息（Jiang等人，2023）来实现。

最后，还提出了诸如Mamba（Gu&Dao，2023）和RWKV（Peng等人，2023）之类的新型架构（Gu等人，2022；Fu等人，2023a；Poli等人，2023；Fu等，2023b；Sun等人，2023a；Beck等人，2024；Sun等人），以有效地处理长上下文输入。(（该方法可以看一下本博主写的另一篇blog)

参考文献和链接

[1] C.-P. Hsieh等, 《RULER: What’s the Real Context Size of Your Long-Context Language Models?》, 2024年8月6日, arXiv: arXiv:2404.06654. doi: 10.48550/arXiv.2404.06654.
[2] W. Wang等, 《Needle In A Multimodal Haystack》, 2024年10月9日, arXiv: arXiv:2406.07230. doi: 10.48550/arXiv.2406.07230.