使用RingAttention处理百万长度视频和语言的世界模型

摘要

2402.08268v1.pdf (arxiv.org)

当前的语言模型在理解不容易用语言描述的世界方面存在不足，并且在处理复杂、长篇的任务时也存在困难。视频序列提供了语言中不存在的有价值的时间信息，使其与语言联合建模变得具有吸引力。这样的模型可以发展出对人类文本知识和物理世界的理解，从而为人工智能提供更广泛的辅助人类的能力。然而，由于内存限制、计算复杂性和有限的数据集，从数百万个视频和语言序列令牌中学习带来了挑战。为了解决这些挑战，我们整理了一个包含各种视频和书籍的大型数据集，利用 RingAttention 技术对长序列进行可扩展训练，并逐步增加从 4K 到 1M 令牌的上下文大小。本文做出了以下贡献：（a）最大上下文大小神经网络：我们在长视频和语言序列上训练了具有最大上下文大小的变压器之一，为困难的检索任务和长视频理解设定了新的基准。（b）克服视觉语言训练挑战的解决方案，包括使用掩码序列打包混合不同序列长度、损失加权以平衡语言和视觉，以及为长序列聊天生成的模型生成的问答数据集。（c）使用 RingAttention、掩码序列打包和其他关键功能对数百万长度的多模态序列进行训练的高度优化的实现。（d）完全开源了一个能够处理超过 1M 令牌的长文本文档（LWM-Text、LWM-Text-Chat）和视频（LWM、LWM-Chat）的 70 亿参数模型家族。这项工作为在庞大的长视频和语言数据集上进行训练铺平了道路，以发展对人类知识和多模态世界的理解，以及更广泛的能力。

在这里插入图片描述

1 简介

目前对世界的建模方法大多局限于短语言序列或短图像和片段序列 [BMR+ 20, TLI+ 23, TMS+ 23, Ope23, TAB+ 23]。这导致模型缺乏对难以用文本或短片段表示的世界部分的理解，并且无法处理复杂的长篇语言和视觉任务。视频序列中的时间结构提供了有用的信息，这些信息在语言中缺失或在静态图像和短片段中不那么明显。长语言序列编码了短序列无法编码的信息，这对于各种应用（如长文档检索或编码）至关重要。长视频提供了丰富的上下文，这是短片段无法把握的，展示了场景如何连接、事件的发展和视频时间维度内动作的因果关系。通过联合建模长视频和书籍，该模型可以发展对多模态世界以及长文本和视频序列的理解，从而导致具有多模态理解的更先进的人工智能系统，能够在更广泛的任务中辅助人类。

为了从视频和语言序列中学习，我们需要训练一个能够处理每个序列中超过数百万个令牌的模型，并在一个非常大的数据集上进行训练。然而，由于内存成本高、计算复杂度高以及缺乏合适的数据集，对数百万个令牌进行建模是极其困难的。幸运的是，我们有了 RingAttention [LZA24]，这是一种可以任意扩展上下文大小而不进行近似或增加开销的技术，从而实现了对长序列的可扩展训练。我们从公共书籍和视频数据集中整理了一个大型的视频和语言数据集，其中包含了各种活动的视频和长篇书籍。考虑到高昂的计算成本，我们逐渐从较小的 4K 上下文增长到较大的 1M 上下文大小，以降低这一成本，并且这种方法在有效扩展上下文方面表现良好。此外，我们还确定了与视频和语言训练相关的挑战：我们发现对视频、图像和文本的混合训练对于实现最佳性能至关重要，因为图像代表更高的视觉质量，视频提供顺序信息，而文本保留语言理解。为了实现这一点，我们实施了一种高效的掩码序列打包，以有效地训练不同长度的序列，而不是标准的序列打包机制。此外，确定图像、视频和文本训练之间的正确平衡对于跨模态理解至关重要，我们提出了一种证明有效的比率。此外，为了解决缺乏长篇聊天数据集的问题，我们开发了一种模型生成的问答（QA）方法，通过使用短上下文模型从书籍中生成 QA 数据集。我们发现这对于长序列上的聊天能力至关重要。

本文的具体贡献如下：(a)我们在视频和文本序列上训练了迄今为止最大的上下文大小转换器之一，并在长视频理解（例如，参见图1）和长上下文事实检索（例如，参见图2）方面取得了迄今为止最好的报告结果。(b)我们发现了与视频和文本序列训练相关的一系列挑战，并为它们提出了解决方案：损失加权以平衡语言和视觉，掩码序列打包以有效地训练不同长度的序列，以及用于长序列聊天的模型生成的QA数据集。©一个高度优化的、开源的实现，具有RingAttention、掩码序列打包和其他关键功能，可用于数百万长度的多模态训练。(d)完全开源了一个能够处理长达1百万个令牌的长文本文档（LWM-Text、LWM-Text-Chat）和视频（LWM、LWM-Chat）的70亿参数模型家族。我们的工作为在长视频和语言的大规模数据集上进行训练铺平了道路，对于未来开发具有人类知识和多模态世界理解以及更广泛能力的AI系统非常有用。

2、概述

我们在一个非常大的上下文窗口（最多可达一百万个令牌）上训练了一个大型的自回归变换器模型，该模型基于Llama2 7B [TMS +23] 构建。为了实现这一目标，我们采用了多种策略：首先使用书籍将上下文扩展到1M（第3节），然后对长多模态序列进行联合训练，包括文本-图像、文本-视频数据和书籍（第4节）。

我们的训练阶段和数据集如图3所示，模型架构如图4所示。

3、阶段一：学习长上下文语言模型

本阶段的目标是首先开发LWM-Text和LWM-Text-Chat，这是一组通过训练逐渐增长序列长度的数据（使用RingAttention）并修改位置编码参数以适应更长序列长度的长上下文语言模型（参见第3.1节）。关于如何增长上下文大小的训练步骤详见第3.2节。在第3.3节中，我们将展示如何构建模型生成的QA数据，以实现长序列对话。

3.1 扩展上下文

在数百万个令牌的序列上学习长距离依赖关系需要（1）在这样的长文档上进行可扩展的训练，以及（2）稳定地扩展我们基础语言的上下文。

长文档上的可扩展训练。由于计算注意力权重的二次复杂性所施加的内存限制，对长文档进行训练变得非常昂贵。为了解决这些计算约束，我们使用了RingAttention [LZA24, LA23] 实现，该实现利用块级计算和序列并行性，理论上可以扩展到无限上下文，仅受可用设备数量的限制。我们进一步使用Pallas [BFH+ 18] 将RingAttention与FlashAttention [DFE+ 22, RS 21] 融合，以优化与XLA编译器相比的性能。一般来说，给定每个设备足够大的令牌数，RingAttention期间的通信成本完全与计算重叠，并不会增加任何额外的开销。

逐步增加上下文长度的训练。尽管我们的实现允许我们在数百万个令牌的长文档上进行训练，但由于注意力的二次计算复杂性仍然存在，因此这仍然是一个昂贵的操作，其中梯度步骤的时间大致与上下文大小成线性关系（给定每个批次的固定令牌数）。例如，当在1M令牌序列长度的7B模型上进行训练时，每个梯度步骤大约需要7分钟，这意味着在全天训练后只能进行总共200步。

因此，我们采用了一种受[JHY+ 23]启发的训练方法，其中我们的模型是在逐步增加的序列长度上进行训练的，从32K令牌开始，以2的幂次递增，直到1M令牌结束。从直觉上讲，这允许模型首先学习较短范围的依赖关系，然后再转向更长的序列，从而节省计算资源。通过这样做，与直接在最大目标序列长度上进行训练相比，我们能够训练更多数量级的令牌。逐步增加上下文大小的训练过程如图3所示。

位置推断用于长上下文

为了在长上下文中扩展位置嵌入，我们采用了一种简单且经过扩展的方法。这种方法在[RGG+ 23]中有所探索，其中RoPE的θ值[SAL+ 24]会随着上下文长度进行缩放。由于这种方法相对简单，只需要调整一个超参数，因此我们发现它是扩展具有长上下文的位置嵌入的稳定方法。我们随着上下文窗口大小一起扩展RoPE的θ值，这些值在表1中给出。

3.2 训练步骤

我们从LLaMA-2 7B模型[TMS {}^{+} 23]开始，并在五个阶段中逐步增加模型的有效上下文长度，这些阶段分别是：32K、128K、256K、512K和1M。对于每个阶段，我们在来自The Pile的Books 3数据集的不同过滤版本上进行训练[GBB+ 20]。表1详细说明了每个训练阶段的信息，如令牌数、总时间和Books3数据集的过滤约束。每个连续运行的初始化都来自前一个序列长度的运行。

3.3 长上下文学习的聊天微调

构建长上下文推理的QA数据

为了学习长上下文聊天能力，我们构建了一个简单的QA数据集。我们从Books 3数据集中将文档切分为固定大小的1000个令牌块，将每个块输入到我们的短上下文语言模型中，并提示它生成一个关于该段落的问答对。然后，给定一个如32K的上下文长度，我们通过将相邻的块连接在一起，并在序列末尾以聊天形式附加相关的QA对，来构建一个单个的32K令牌示例。

训练细节

对于聊天微调，我们以大约7:3的比例在UltraChat[DCX+ 23]和我们的自定义QA数据集上进行训练。我们发现将UltraChat数据预先打包到训练序列长度，并将其与我们的QA数据的示例分开至关重要，因为UltraChat数据通常具有更高的损失令牌比例（密集排列的短聊天序列），而我们的QA数据每个序列的损失令牌比例要低得多（<1%），因为给定的上下文中的长文档没有损失。表2显示了每个运行的进一步训练细节。请注意，对于每个聊天模型，并没有进行渐进式训练，而是从相同上下文长度的各自预训练模型进行初始化。

3.4 语言评估结果

3.4.1 单针检索

我们在流行的“大海捞针”任务[gka23]上进行了评估，更具体地说，是一个更容易评估的版本[AI23]，该任务从上下文中找到并检索分配给随机城市的随机数。图5显示了我们1M上下文模型的整个上下文中近乎完美的检索准确率。此外，图2显示与当前可用的最佳LLM相比，我们可以扩展到更大的上下文。附录A显示了我们的其他较短上下文长度模型的更多单针检索结果。

3.4.2 多针检索

我们还通过混合多根针并尝试检索其中的特定子集，来检查我们的模型在更复杂的针检索任务变体上的性能。图6显示了不同设置下的多针检索结果。我们的模型在从上下文中的多根针中检索单根针时表现出良好的泛化能力，当要求检索多于一根针时，性能略有下降。表3显示了我们的模型、Gemini Pro和GPT-4之间的多针比较，其中我们的模型在检索一根针时能够与GPT-4竞争或表现更好，或在检索多于一根针时性能略低。此外，我们的模型还能够在长达1M令牌的上下文中表现出色。然而，我们注意到在增加针检索任务的难度时，准确率会下降，这表明我们的模型在1M上下文利用方面仍有改进空间。我们相信，我们发布的模型将为未来开发更长上下文模型的工作提供基础，并鼓励制定更具挑战性的基准测试，其中包含需要更高层次合成的困难长距离任务，而不仅仅是纯粹的事实检索。

3.4.3 短上下文语言评估

表4比较了具有4K上下文的Llama2-7B模型与其从32K到1M的上下文扩展版本。评估涵盖了各种语言任务，结果表明，扩展上下文大小并不会损害短上下文任务的性能。实际上，结果表明，具有更大上下文容量的模型在这些任务上的表现同样出色，甚至可能更好。这一证据表明，上下文扩展没有产生负面影响，突显了模型适应不同任务要求而不损失较短上下文效率的能力。

3.4.4 聊天评估

我们还评估了我们的模型在MT-Bench [\mathrm{ZCS}^{+} 23 \mathrm{a}]上的对话能力。表5显示了每个模型的MT-Bench分数。表6说明了聊天和事实检索任务混合与MT-Bench分数和针检索准确度之间的关系。随着聊天比例的增加和事实检索的减少，MT-Bench分数提高，表明由MT-Bench测量的聊天性能更好。相反，针检索准确度降低，表明增加聊天交互能力可能会降低系统从输入上下文中检索特定信息或“针”的精确度。在不同的上下文大小中，我们发现支持更长输入序列的模型在MT-Bench分数上略有下降。我们假设这是因为我们在更长的序列训练上选择了较少的示例，并且可以通过简单地训练更多的数据来改进。此外，通过获取更接近UltraChat数据集聊天分布的高质量长上下文聊天数据，可能会解决这种权衡。

4 第二阶段：学习长上下文视觉语言模型

我们的第二阶段目标是在长视频和语言序列上有效地进行联合训练。在4.1节中，我们将介绍对LWM和LWM-Chat的架构修改，以融入视觉输入。4.2节将讨论在不同序列长度上的训练。评估结果将在4.3节中展示。在这一阶段，我们通过在各种长度的视觉语言数据上对先前开发的1M上下文语言模型进行微调，来增强其能力。所使用的数据集和训练过程中涉及的步骤如图3所示。

4.1 针对视觉的架构修改

模型架构如图4所示。我们使用了从aMUSEd [PBRvP24]预训练的VQGAN [ERO21]，它将256x256的输入图像标记化为16x16的离散标记。视频通过每帧应用VQGAN并进行代码串联来进行标记化。为了在生成时区分不同的模态以及知道何时切换，我们引入了机制来标记文本生成的结束/视觉生成的开始，反之亦然。为了定义视觉生成的结束，我们引入了新的标记和，它们分别表示帧结束（在每个非序列中最后一个视频帧的末尾）和视觉结束（在每个单个图像或在视频最后一帧的末尾）的边界。为了定义文本生成的结束，我们用和（作为文本）文本标记包裹视觉标记。模型是使用交织的视觉和文本标记串联进行训练的，并按自回归方式进行预测。
4.2 训练步骤

我们从LWM-Text-1M文本模型进行初始化，并在大量的组合文本-图像和文本-视频数据上执行类似的渐进式训练过程。不同的是，我们不再额外扩展RoPE θ，因为它已经支持高达1M的上下文。表7显示了每个训练阶段的详细信息，其中模型是从之前的较短序列长度阶段进行初始化的。对于每个阶段，我们训练以下数据：

LWM-1K：我们在一个大型文本-图像数据集上进行训练，该数据集由LAION-2Ben [SBV+ 22]和COYO-700M [BPK+ 22]的混合组成。数据集经过过滤，仅包含至少256分辨率的图像——总计约10亿个文本-图像对。在训练过程中，我们将文本-图像对进行串联，并随机交换模态的顺序以模拟文本-图像生成、无条件图像生成和图像标注。我们将文本-图像对打包成1K标记的序列。
LWM-8K：我们在WebVid10M [BNVZ21]和3M InternVid10M [WHL+ 23]的文本-视频数据集混合上进行训练。与先前的工作[HCS+ 22, HSG+ 22, VBK+ 22]类似，我们以50-50的比例联合训练图像和视频。我们将图像打包成8K标记的序列，并将30帧的视频以4FPS的速度进行处理。与图像训练类似，我们随机交换每个文本-视频对的模态顺序。
LWM-Chat-32K/128K/1M：对于最后3个阶段，我们在每个下游任务的组合聊天数据上进行训练：(1)文本-图像生成，(2)图像理解，(3)文本-视频生成，以及(4)视频理解。我们通过从预训练数据中随机抽样子集并增强聊天格式来构建文本-图像和文本-视频的聊天数据的简单版本。对于图像理解，我们使用了来自ShareGPT4V的图像聊天指导数据 $\left[\mathrm{CLD}^{+} 23\right]$ 。最后，对于视频理解聊天数据，我们结合了Valley-Instruct-73K [LZY {}^{+} 23 \mathrm{~b} ]和Video-ChatGPT-100K指导数据[MRKK23]。对于所有短上下文数据（图像生成、图像理解、视频生成），我们将序列打包到训练上下文长度。在打包过程中，我们发现掩盖注意力至关重要，这样每个文本-视觉对只关注自己，同时重新加权损失，使计算与非打包+填充训练方案中的训练相同。对于视频理解数据，如果视频太长，我们会均匀采样最大帧数以适应模型的训练上下文长度。在训练过程中，我们将每个批次的25%分配给4个下游任务中的每一个。

在前两个训练阶段（LWM-1K和LWM-8K）中，我们另外将批次的16%混合为来自OpenLLaMA的纯文本数据[GL23]，因为我们发现，在视觉数据上训练时保留语言功能是有益的。

4.3 视觉语言评估结果

4.3.1 长期视频理解

尽管视觉语言模型 $\left[\mathrm{LZY}^{+} 23 \mathrm{a}, \mathrm{Ope}^{23}, \mathrm{TAB}^{+} 23\right]$ 可以处理长视频，但由于上下文长度有限，通常需要对视频帧进行大量时间下采样。例如，Video-LLaVA$ [LZY { }^{+} 23 \mathrm{a} ] $只能从视频中均匀采样8帧，无论原始视频有多长。因此，模型可能会丢失更多重要的细粒度时间信息，这对于准确回答有关视频的任何问题都很重要。相比之下，我们的模型是在1M个token的长序列上进行训练的，因此可以同时处理数千帧视频，以在短时间间隔内检索细粒度信息。图1显示了我们的模型正确回答有关长达1小时、由500多个独立片段组成的YouTube合辑的问题的示例。另一方面，我们的基准方法由于帧数有限，通常难以回答问题。更多的结果如图7和附录B所示。

尽管我们证明了我们的模型可以在复杂的长格式视频上进行问答，但我们注意到，在所有1M个token中，仍有改进的空间以更好地利用上下文，因为我们的模型生成的答案可能并不总是准确的，而且模型在处理需要更高层次视频理解的复杂问题时仍然会遇到困难。我们希望我们的模型能够帮助未来的工作开发出改进的基础模型，以及长视频理解的基准。

4.3.2 图像理解和短视频理解

表8和表9显示了图像理解和短视频理解的常见基准测试结果。图17显示了图像理解的定性示例。我们的模型在基线中表现平均，并且低于最先进的模型。我们假设这可能是由于文本-图像和文本-视频的对齐训练有限，而基线可以利用已经经过更广泛、基于大规模CLIP训练的视觉骨干网络。相比之下，我们的模型使用VQGAN令牌，需要从头开始学习文本-图像对齐，并且由于VQGAN重建图像中文本的能力不够忠实，通常会在OCR任务上遇到困难。然而，我们相信我们的模型将是未来基于VQ的视觉语言模型的一个有前途的方向，可以通过更严格的训练和学习更好的分词器来表现得更好。附录B显示了更多的定性图像理解和附录B显示了更多的定性视频理解示例。

4.3.3 图像和视频生成

除了能够进行图像/视频标题以及问答之外，我们的模型还可以从文本中生成图像和视频。图8显示了这种能力的一些示例。我们采用与先前工作类似的方式 $\left[\mathrm{YXK}^{+} 22, \mathrm{GPA}^{+}{ }^{2} 2\right]$ 对自回归采样的logits进行无分类器指导[HS22]。对于无条件分支，我们使用初始化每个序列。附录E和D显示了更多的图像和视频生成示例。

4.3.4 掩码序列打包消融实验

正如4.2节所提到的，正确地掩码注意力并重新加权损失对于下游任务的某些方面至关重要，特别是图像理解。表10显示了我们的模型在是否进行打包校正的情况下的比较。简单地打包会导致图像理解任务的准确性大幅下降。我们假设简单的打包会降低性能，因为它会降低较短的文本令牌答案的权重，这对于良好的图像理解基准性能是一个重要的方面。

5 更多细节

MFU（模型浮点运算利用率）. 我们使用TPUv4-1024进行模型训练，这大约相当于450个A100 GPU，使用FSDP [Fac23] 和RingAttention进行大上下文训练，批处理大小为8M。图9显示了每个训练阶段的模型浮点运算利用率（MFU）。蓝色条形图表示语言训练，橙色条形图表示视觉语言训练。我们的训练即使在非常大的上下文大小下也能实现良好的MFU。

训练损失曲线. 图10和图11分别显示了语言模型和视觉语言模型每个训练阶段的训练损失曲线。

训练超参数. 参见附录F

推理扩展. 我们还扩展了我们的推理代码，通过实现RingAttention来支持百万长度的序列解码。对于这种长序列的推理，至少需要v4-128，使用32个张量并行性和4个序列并行性（环形维度）的TPU网格分片。我们在纯单精度下进行推理，可以通过可扩展性技术（如量化）进行进一步的改进。

6 相关工作

我们的工作与一些旨在扩展语言模型上下文窗口以允许更多令牌的研究相关，如 $ \left[\mathrm{CWCT23}, \mathrm{TSP}^{+} 23, \mathrm{LYZ}^{+} 23\right. , inter alia], $ 等，这些研究通常通过使用新颖的外推方法来扩展预训练的位置编码，并在更长的上下文数据上对模型进行微调。我们的模型采用了一种简单的方法，即在训练过程中逐步增加RoPE位置编码中的θ，并发现这种方法是有效的。也有研究在不建模成对交互的架构上进行，如稀疏注意力和滑动窗口 [CGRS19, BPC20]；我们的工作利用RingAttention [LZA24] 和BPT [LA23] 来在长序列中建模精确的成对交互以获得最佳性能。通过跳过因果掩码计算的负载平衡也可以进一步提高训练性能$\left[\mathrm{BNQ}^{+} 23\right] $ 。

我们的工作还与指令调优的相关工作相关 $\left[\mathrm{TGZ}^{+} 23, \mathrm{CLL}^{+} 23, \mathrm{GGL}^{+} 23\right. , inter alia]$ 等。这些研究重点是通过使用对话数据对模型进行微调，以增强其在语言任务上的能力。我们的方法旨在推进模型对复杂、长序列的视频和语言的理解。为了实现这一点，我们通过在书籍和长视频上进行训练，并在模型生成的QA数据上进行微调来学习长序列上的聊天能力，从而扩展了模型的上下文大小。

我们的工作还与结合视觉和语言的研究相关 $}^{+} 23 \mathrm{a}, \mathrm{AGG}^{+} 23 , \mathrm{ZHZ}^{+} 23, \mathrm{JXC}^{+} 23 , inter alia]$ 等。这些研究通常使用CLIP $\left[\mathrm{RKH}^{+} 21\right]$ 或BLIP [LLXH22] 将视觉信息编码为嵌入，然后输入到语言模型中。它们有可能利用CLIP的跨模态理解来编码图像中的文本信息。然而，它们只能被训练成根据视觉输入预测文本，而不能反过来，这限制了它们从多种视觉和语言信息格式中学习的能力。在它们预测视觉令牌的情况下 $\left[\mathrm{JXC}^{+} 23\right]$ ，由于CLIP嵌入的相对有损性质，通常需要一个更强的扩散解码器。另一方面，我们的工作是自回归的“令牌输入，令牌输出”，这使我们能够灵活地建模多种形式的图像-文本、文本-图像、文本-视频、视频-文本以及纯视频、图像或文本格式。

7 结论

在本文中，我们解决了通过结合语言和视频来更好地学习模型理解世界的挑战。我们利用RingAttention在大量的长视频和书籍数据集上进行可扩展的训练，并逐步增加序列长度，从 $\mathrm{~K}$ 到 $\mathrm{M}$ 令牌，以保持计算的可管理性。我们开发了掩码序列打包和损失加权，以有效地在多样化的视频、图像和书籍数据集上进行训练。最后，我们证明了LWM具有高效的 $\mathrm{M}$ 上下文大小，这是迄今为止最大的，使其能够成功处理涉及长视频和语言序列的复杂任务。我们开源了我们针对数百万长度序列训练的RingAttention、掩码序列打包和其他关键功能的优化实现，以及一个能够处理超过 $\mathrm{M}$ 多模态令牌的70亿参数模型。我们希望这项工作能够为推进具有可靠推理和对世界的扎实理解的AI模型以及更广泛的能力铺平道路。