作者 | 北方的郎 编辑 | 汽车人
原文链接:https://zhuanlan.zhihu.com/p/683154572
点击下方卡片,关注“自动驾驶之心”公众号
戳我-> 领取自动驾驶近15个方向学习路线
本文只做学术分享,如有侵权,联系删文
在探索如何让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 小时的问答视频:

超过 1M 上下文的事实检索:


技术实现
为了训练和评估LWM,研究人员首先收集了一个包含各种视频和书籍的大型数据集。然后,他们逐步增加了训练的上下文长度,从4K tokens开始,逐步扩展到1M tokens。这一过程不仅有效降低了训练成本,还使模型能够逐步适应更长序列的学习。在训练过程中,研究人员还发现,混合不同长度的图像、视频和文本数据对于模型的多模态理解至关重要。具体包括:
模型训练分两个阶段:首先通过训练大型语言模型扩展上下文大小。然后进行视频和语言的联合训练。
Stage I: Learning Long-Context Language Models
扩展上下文:利用RingAttention技术,可以无近似地扩展上下文长度到数百万个token。同时,通过逐步增加训练序列长度,从32K tokens开始,逐步增加到1M tokens,以减少计算成本。此外,为了扩展位置编码以适应更长的序列,采用了简单的方法,即随上下文窗口大小增加而增加RoPE中的θ。

训练步骤:首先从LLaMA-2 7B模型初始化,然后在5个阶段逐步增加上下文长度,分别是32K、128K、256K、512K和1M tokens。每个阶段都使用不同过滤版本的Books3数据集进行训练。随着上下文长度的增加,模型能够处理更多tokens。

对话微调:为了学习长上下文的对话能力,构建了一个简单的问答数据集,将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研究表明,屏蔽序列填充对于图像理解等下游任务至关重要。




第二阶段通过逐步增加序列长度并在大量文本-图像和文本-视频数据上训练,成功扩展了第一阶段的语言模型,使其具备视觉理解能力。这一阶段的模型可以处理长达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路序列并行。

量化:文档指出,模型使用单精度进行推断,通过量化等技术可以进一步提高扩展性。
一些例子
基于图像的对话。

超过 1 小时的 YouTube 视频视频聊天。


投稿作者为『自动驾驶之心知识星球』特邀嘉宾,欢迎加入交流!
① 全网独家视频课程
BEV感知、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、大模型与自动驾驶、Nerf、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频(扫码即可学习)
② 国内首个自动驾驶学习社区
国内最大最专业,近2700人的交流社区,已得到大多数自动驾驶公司的认可!涉及30+自动驾驶技术栈学习路线,从0到一带你入门自动驾驶感知(2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪)、自动驾驶定位建图(SLAM、高精地图、局部在线地图)、自动驾驶规划控制/轨迹预测等领域技术方案、大模型、端到端等,更有行业动态和岗位发布!欢迎扫描下方二维码,加入自动驾驶之心知识星球,这是一个真正有干货的地方,与领域大佬交流入门、学习、工作、跳槽上的各类难题,日常分享论文+代码+视频

③【自动驾驶之心】技术交流群
自动驾驶之心是首个自动驾驶开发者社区,聚焦2D/3D目标检测、语义分割、车道线检测、目标跟踪、BEV感知、多模态感知、Occupancy、多传感器融合、transformer、大模型、在线地图、点云处理、端到端自动驾驶、SLAM与高精地图、深度估计、轨迹预测、NeRF、Gaussian Splatting、规划控制、模型部署落地、cuda加速、自动驾驶仿真测试、产品经理、硬件配置、AI求职交流等方向。扫码添加汽车人助理微信邀请入群,备注:学校/公司+方向+昵称(快速入群方式)
④【自动驾驶之心】平台矩阵,欢迎联系我们!