“WORLD MODEL ON MILLION-LENGTH VIDEO AND LANGUAGE WITH RINGATTENTION”
为了应对这些挑战，UC伯克利的研究者整理了一个包含各种视频和书籍的大型数据集，并且提出了大世界模型（Large World Model，LWM），利用RingAttention 技术对长序列进行可扩展训练，将上下文大小从 4K 增加到 1M Token。
项目主页：https://largeworldmodel.github.io/
论文地址：https://arxiv.org/pdf/2402.08268.pdf
Github地址：https://github.com/LargeWorldModel/LWM
摘要
当前的语言模型在理解难以用文字描述的世界方面存在局限性，并且在处理复杂的长篇任务时遇到困难。视频序列提供了语言和静态图像所缺乏的有价值的时间信息，使其成为与语言联合建模的有吸引力的选择。这样的模型可以发展对人类文本知识和物理世界的理解，从而为辅助人类的广泛AI能力提供支持。然而，从数百万个标记的视频和语言序列中学习面临着内存限制、计算复杂性和有限数据集的挑战。为了解决这些问题，我们策划了一个包含多样化视频和书籍的大型数据集，利用RingAttention技术在长序列上进行可扩展训练，并逐渐增加上下文大小从4K到1M标记。
本文的贡献如下：
- a）最大上下文大小的神经网络：我们训练了一个具有最大上下文大小的transformer模型，对困难的检索任务和长视频理解设置了新的基准。
- b）解决克服视觉语言训练挑战的方法，包括使用掩码序列打包来混合不同的序列长度，损失加权来平衡语言和视觉，以及模型生成的QA数据集用于长序列对话。
- c）高度优化的实现，包括RingAttention、掩码序列打包和其他关键特性，用于在数百万长度的多模态序列上进行训练。
- d）完全开源的一系列7B参数模型，能够处理超过1M标记的长文本文档(LWM-Text、LWM-Text-Chat)和视频(LWM、LWM-Chat)。
这项工作为在大规模视频和语言数据集上进行训练，以发展对人类知识和多模态世界的理解以及更广泛的能力铺平了道路。
概述
本文介绍了一个使用大型自回归变换器模型进行训练的方法，上下文窗口可以达到一百万个标记。为了实现这个目标，使用了几种策略，包括使用书籍扩展上下文，以及在长的多模态序列上进行联合训练。
阶段1：训练长上下文语言模型
LWM-Text和LWM-Text-Chat是一组长上下文语言模型，通过RingAttention和修改位置编码参数来训练逐渐增加序列长度的数据。
扩展上下文
学习长序列的远程依赖关系需要可扩展的训练和稳定的扩展基础语言上下文。通过使用RingAttention和FlashAttention的组合来解决训练长文档的计算限制，并通过Pallas进行性能优化。采用渐进式训练方法，从较短的序列长度开始逐渐增加上下文长度，以节省计算资源。采用简单的位置外推方法来扩展长上下文的位置嵌入，通过调整单个超参数来实现稳定的扩展。
训练阶段
使用LLaMA-2 7B [TMS + 23]模型进行文本生成的训练。通过逐步增加模型的有效上下文长度，分为5个阶段进行训练，分别为32K、128K、256K、512K和1M。每个阶段都使用来自The Pile [GBB + 20]的Books3数据集的不同过滤版本进行训练。每次运行都是从前一个序列长度的运行中初始化的。
长上下文学习的对话微调
本文介绍了如何构建一个用于学习长文本推理的简单QA数据集。作者将Books3数据集中的文本分成1000个标记的块，然后将每个块输入到短上下文语言模型中，提示其生成一个关于该段落的问题-答案对。然后，作者将相邻的块连接在一起，构建一个32K标记的示例，并在末尾以聊天形式附加相关的QA对。作者还介绍了聊天微调的训练细节，包括使用UltraChat和自定义QA数据集进行训练，以及如何处理不同数据集之间的差异。
语言评估结果
Single-Needle检索
我们评估了模型在Needle In A Haystack任务上的表现，结果显示该模型在1M上下文模型中表现出了几乎完美的检索准确性，并且可以比当前最好的LLMs扩展到更大的上下文。
Multi-Needle检索
结果显示，当从多个Needle中检索单个Needle时，我们的模型表现良好，但在检索多个Needle时稍有下降。
短上下文语言评估
Llama2-7B模型与其上下文扩展版本进行了比较分析。结果表明，扩展上下文大小不会影响短上下文任务的性能。事实上，结果表明具有更大上下文容量的模型在这些任务上表现得同样好，甚至更好。这些证据表明上下文扩展没有负面影响，突出了模型在不损失短上下文效率的情况下适应不同任务要求的能力。
对话评估
结果表明，随着聊天比例的增加和事实检索比例的减少，MT-Bench得分提高，表明聊天性能更好。然而，针对检索任务的准确性会下降，这表明在增加聊天交互能力时，可能会降低系统从输入上下文中检索特定信息的精度。此外，模型支持更长输入序列的情况下，MT-Bench得分略有下降，这可能是因为训练数据较少，可以通过增加数据量来改善。最后，可以通过获取更接近UltraChat数据集聊天分布的高质量长上下文聊天数据来解决这种权衡。
阶段2：训练长上下文视觉-语言模型
第二阶段旨在有效地联合训练长视频和语言序列。我们将在LWM和LWM-Chat中引入架构修改，以包含视觉输入。我们讨论了在不同序列长度上的训练。在这个阶段，我们通过在不同长度的视觉语言数据上微调先前开发的1M上下文语言模型来增强其能力。
针对视觉的架构修改
使用预训练的VQGAN将256×256的图像分词为16×16的离散标记，并将视频逐帧分词，然后将代码连接在一起。为了区分生成的模态以及知道何时切换，引入了标记文本生成/视觉生成结束的机制。为了定义视觉生成的结束，引入了新的标记和，分别表示帧的结束和视觉的结束。为了定义文本生成的结束，将视觉标记用和包装起来。模型通过交替连接视觉和文本标记进行训练，并进行自回归预测。
训练阶段
从我们的LWM-Text-1M文本模型初始化，并在大量的文本-图像和文本-视频组合数据上执行类似的渐进式训练过程。对于每个阶段，我们在以下数据上进行训练:
- LWM-1K：我们在由laon – 2ben [SBV + 22]和COYO-700M [BPK + 22]混合而成的大型文本-图像数据集上进行训练。在训练过程中，连接文本-图像对，并随机交换模态的顺序，以对文本-图像生成、无条件图像生成和图像描述进行建模。
- LWM-8K：我们在WebVid10M [BNVZ21]和3M InternVid10M [WHL + 23]示例的文本-视频数据集混合上进行训练。与图像训练类似，随机交换每个文本-视频对的模态顺序。
- WM-Chat-32K/128K/1M：在最后的3个阶段，我们对每个下游任务的聊天数据的组合混合进行训练：(1)文本-图像生成，(2)图像理解，(3)文本-视频生成，以及(4)视频理解。
在训练的前两个阶段(LWM-1K和LWM-8K)，我们还将批数据的16%混合为来自OpenLLaMA [GL23]的纯文本数据，因为我们发现在对视觉数据进行训练时保留语言能力是有益的。
视觉-语言评估结果
长视频理解
模型可以同时处理数千帧视频，提取细粒度的时间信息，从而更准确地回答关于视频的问题。与现有的模型相比，该模型可以处理更长的视频序列，但仍有改进的空间。
图像理解和短视频理解
结果表明，该模型在基准测试中表现一般，低于SOTA模型。作者认为这可能是由于训练数据不足，需要更严格的训练和更好的分词器。此外，该模型使用的VQGAN令牌需要从头开始学习文本-图像对齐，因此在OCR任务中表现不佳。
图像和视频生成
该模型不仅可以执行图像/视频字幕和QA，还可以从文本生成图像和视频。
掩码填充消融分析
本文指出，在某些下游任务中，正确地掩码注意力并重新加权损失对于图像理解等方面至关重要。通过比较我们的模型在有和没有打包修正的情况下的表现，发现朴素的打包会导致准确性大幅下降，这可能是因为它会降低文本标记答案的权重，而这对于良好的图像理解基准性能是重要的。
总结
本文介绍了一种结合语言和视频的学习模型，利用RingAttention在大规模数据集上进行训练，并逐步增加序列长度，最终实现了1M标记的高效处理。作者还开源了优化实现的RingAttention和其他关键特性，以及一个可处理超过1M多模态标记的7B参数模型。未来的研究可以探索更好的视频标记化、更多的模态和更好的视频数据集。
评论0