本文经自动驾驶之心公众号授权转载，转载请联系出处。
在探索如何让AI更好地理解世界方面，最近的一项突破性研究引起了广泛关注。来自加州大学伯克利分校的研究团队发布了“Large World Model, LWM”，能够同时处理百万级长度的视频和语言序列，实现了对复杂场景的深入理解。这一研究无疑为未来AI的发展开启了新的篇章。
论文地址：World Model on Million-Length Video And Language With RingAttention
博客地址：Large World Models
huggingface: LargeWorldModel (Large World Model)
在传统方法中，AI模型往往只能处理较短的文本或视频片段，缺乏对长时间复杂场景的理解能力。然而，现实世界中的许多场景，如长篇书籍、电影或电视剧，都包含了丰富的信息，需要更长的上下文来进行深入理解。为了应对这一挑战，LWM团队采用了环形注意力(RingAttention)技术，成功扩展了模型的上下文窗口，使其能够处理长达100万个令牌(1M tokens)的序列。例如实现超过 1 小时的问答视频：
图1.长视频理解。LWM 可以回答有关超过 1 小时的 YouTube 视频的问题。
超过 1M 上下文的事实检索:
图 2. 针检索任务。LWM 在 1M 上下文窗口内实现了高精度，并且性能优于 GPT-4V 和 Gemini Pro。
图 3. 针检索任务。LWM 对于上下文窗口中不同的上下文大小和位置实现了高精度。
技术实现
为了训练和评估LWM，研究人员首先收集了一个包含各种视频和书籍的大型数据集。然后，他们逐步增加了训练的上下文长度，从4K tokens开始，逐步扩展到1M tokens。这一过程不仅有效降低了训练成本，还使模型能够逐步适应更长序列的学习。在训练过程中，研究人员还发现，混合不同长度的图像、视频和文本数据对于模型的多模态理解至关重要。具体包括：
模型训练分两个阶段：首先通过训练大型语言模型扩展上下文大小。然后进行视频和语言的联合训练。
Stage I: Learning Long-Context Language Models
扩展上下文：利用RingAttention技术，可以无近似地扩展上下文长度到数百万个token。同时，通过逐步增加训练序列长度，从32K tokens开始，逐步增加到1M tokens，以减少计算成本。此外，为了扩展位置编码以适应更长的序列，采用了简单的方法，即随上下文窗口大小增加而增加RoPE中的θ。
上下文扩展和视觉语言训练。使用 RingAttention 将书籍上的上下文大小从 4K 扩展到 1M，然后对长度为 32K 到 1M 的多种形式的视觉内容进行视觉语言训练。下面板显示了理解和响应有关复杂多模式世界的查询的交互功能。
训练步骤：首先从LLaMA-2 7B模型初始化，然后在5个阶段逐步增加上下文长度，分别是32K、128K、256K、512K和1M tokens。每个阶段都使用不同过滤版本的Books3数据集进行训练。随着上下文长度的增加，模型能够处理更多tokens。
任意对任意长序列预测。RingAttention 能够使用非常大的上下文窗口进行跨视频-文本、文本-视频、图像-文本、文本-图像、纯视频、纯图像和纯文本等多种格式的训练。请参阅LWM 论文了解关键功能，包括屏蔽序列打包和损失加权，它们可以实现有效的视频语言训练。
对话微调：为了学习长上下文的对话能力，构建了一个简单的问答数据集，将Books3数据集的文档分割成1000 token的块，然后利用短上下文语言模型为每个块生成一个问答对，最后将相邻的块连接起来构造一个长上下文的问答示例。在微调阶段，模型在UltraChat和自定义问答数据集上进行训练，比例为7:3。
语言评估结果：在单针检索任务中，1M上下文的模型可以在整个上下文中近乎完美地检索出随机分配给随机城市的数字。在多针检索任务中，模型在检索一个针时表现良好，在检索多个针时性能略有下降。在短上下文语言任务评估中，扩大上下文长度并没有降低性能。在对话评估中，增加对话交互能力可能会降低系统检索具体信息或“针”的精度。
Stage II: Learning Long-Context Vision-Language Models
架构修改：在第一阶段的基础上，对LWM和LWM-Chat进行修改，使其能够接受视觉输入。具体来说，使用预训练的VQGAN将256x256的输入图像转换为16x16的离散token，对视频进行逐帧的VQGAN编码并将编码连接起来。此外，引入了特殊的标记符号和来区分文本和视觉token，以及和来标记图像和视频帧的结束。
训练步骤：从LWM-Text-1M模型初始化，采用与第一阶段类似的逐步增加序列长度的训练方法，首先在1K tokens上训练，然后是8K tokens，最后是32K、128K和1M tokens。训练数据包括文本-图像对、文本-视频对以及下游任务的聊天数据，如文本-图像生成、图像理解、文本-视频生成和视频理解。在训练过程中，逐步增加下游任务的混合比例。
视觉-语言评估结果：在长视频理解方面，模型能够处理长达1小时的YouTube视频并准确回答问题，相较于现有模型具有明显优势。在图像理解和短视频理解方面，模型表现一般，但通过更严格的训练和更好的分词器，有潜力改进。在图像和视频生成方面，模型可以从文本生成图像和视频。Ablation研究表明，屏蔽序列填充对于图像理解等下游任务至关重要。
文本到图像。LWM 根据文本提示以自回归方式生成图像。
文本到视频。LWM 根据文本提示以自回归方式生成视频。
第二阶段通过逐步增加序列长度并在大量文本-图像和文本-视频数据上训练，成功扩展了第一阶段的语言模型，使其具备视觉理解能力。这一阶段的模型可以处理长达1M tokens的多模态序列，并在长视频理解、图像理解和生成等方面展现出强大的能力。
技术细节（Further Details）
训练计算资源：模型使用TPUv4-1024进行训练，相当于450个A100 GPU，使用FSDP进行数据并行，并通过RingAttention支持大上下文。
训练损失曲线：图10和图11展示了第一阶段语言模型和第二阶段视觉-语言模型的训练损失曲线。可以看出，随着训练进行，损失持续下降。
训练超参数：附录F提供了详细的训练超参数，包括参数量、初始化模型、序列长度、RoPE参数、每批tokens数、总tokens数、训练步骤数、学习率计划、学习率预热步数、最大学习率和最小学习率、计算资源等。
推断扩展：实现了RingAttention用于解码，支持对长达数百万tokens的序列进行推断，需使用至少v4-128 TPU，并进行32路tensor并行和4路序列并行。
量化：文档指出，模型使用单精度进行推断，通过量化等技术可以进一步提高扩展性。
一些例子
基于图像的对话。
图 6. 图像理解。LWM 可以回答有关图像的问题。
超过 1 小时的 YouTube 视频视频聊天。
图 7. 长视频聊天。
即使最先进的商业模型 GPT-4V 和 Gemini Pro 都失败了，LWM 仍能回答有关 1 小时长的 YouTube 视频的问题。每个示例的相关剪辑位于时间戳 9:56（顶部）和 6:49（底部）。