大世界模型 Large World Model(LWM):开源界Gemini Pro

开源 0

摘要:2024年2月,UC Berkeley开源了大世界模型(LWM),支持1M token(与Gemini1.5持平)、1h视频问答、及视频图片生成,相当于开源版Gemini1.5pro。

目录

一、前言

二、模型架构

三、核心技术

四、训练过程

五、效果与性能

六、验证


一、前言

目前,在paperswithcode网站研究趋势榜单中排名第一。

大世界模型(LWM)是一种通用的多模态自回归模型。它使用RingAttention在各种长视频和书籍的大型数据集上进行训练,可以执行语言,图像和视频理解和生成。

大世界模型(LWM)具备的能力:

  1. LWM可以与图像聊天。
  2. LWM可以跨1M上下文检索事实,具有高准确性。
  3. LWM可以在1小时的YouTube视频上回答问题。
  4. LWM可以从文本生成视频和图像。

图像问答是大部分商业和开源VLM(Vison-Language Model)模型的通用能力,如GPT-4V、Qwen-VL。1M的上下文窗口目前与Gemini1.5持平,超越了大部分VLM,也为视频类理解和生成能力的融入铺平道路。同时,能做到在长视频上的问答、并具备图像和视频的生成,可以说UC伯克利的这项工作直接与VLM统一架构标杆的商用产品谷歌Gemini对标,并且开源。

二、模型架构

LLaMA和RingAttention。LWM采用具有高达1M tokens上下文序列的自回归transformer模型架构,基于Llama2-7B和RingAttention。

VQGAN和BPE编码器。图片和视频帧采用VQGAN编码为视觉tokens(对于图片,将256x256输入图像转换16x16=256个tokens。对于视频,通过每帧应用VQGAN并将代码连接在一起),与经过BPE编码后的文本tokens结合,统一送入LWM中,自回归方式进行token预测支持理解和生成任务。

Any-to-Any多模态任务训练。该模型基本上是使用多种模态以任意对任意(any-to-any)的方式进行训练的。输入和输出tokens的顺序反映了不同的训练数据格式,包括图像-文本、文本-图像、视频问答、文本-视频和纯文本问答等格式。为了区分图像和文本标记,并进行解码,我们用特殊的分隔符<vision>和</vision>包围视频和图像标记。我们还采用<eof>和<eov>视觉标记,以标记图像和视频中的中间帧和最后帧的结束。

CFG图像视频生成。LWM也可以从文本生成图像和视频。LWM在logits上使用CFG(classifier-free guidance)进行自回归采样,CFG在文生图扩散模型如SD、DALLE2、Imagen中广泛使用,原始噪声预测被有、无条件噪声预测的线性组合代替,可进一步提升生成质量。在LWM中对于无条件分支,我们用<bos><vision>初始化每个序列。

三、核心技术

环注意力机制:RingAttention

为了增强LWM的长文本处理能力,开发者应用了RingAttention机制。RingAttention是该团队去年提出的一种窗口扩增方式,《Ring Attention with Blockwise Transformers for Near-Infinite Context》论文入选了ICLR 2024。

RingAttention运用了“分而治之”的思想,将长文本分成多个块,用多个计算设备做序列并行处理,然后再进行叠加,理论上允许模型扩展到无限长的上下文。由于transformer自注意力架构中,内存与输入长度呈二次方关系。RingAttention思路是先把SA(Self-Attention)和FFN(前向反馈网络)计算分块Block到不同设备上,借鉴百度提出的分布式训练技术——环形模式(RingAllReduce),在不同设备上分布式计算transformer块,解决上下文的扩展瓶颈。

在LWM中,RingAttention还与FlashAttention结合使用,并通过Pallas框架进行优化,从而提高性能。

RingAttention详细解析

为了能够拓展Transformer可支持的序列长度。作者提出以分块方式执行自注意力和前馈网络计算跨多个主机设备分布序列维度,从而实现并发计算和通信,由于该方法将环中主机设备之间的键值块通信与块计算重叠,因此将其命名:环注意(Ring Attention)。

具体地,该方法在主机设备之间构建注意力计算块的外循环,每个主机设备具有一个查询块,并通过键值块遍历主机设备环,以逐块的方式进行注意力和前馈网络计算。当计算注意力时,每个主机将键值块发送到下一个主机,同时从前一个主机接收键值块。这里作者使用与原始 Transformer 相同的模型架构,但重新组织了计算。具体如下图所示:

对于内循环,每个设备计算其各自的块式注意力和前馈操作。在内循环期间,每个设备将用于块式计算的键值块的副本发送到环中的下一个设备,同时从前一个设备接收键值块。由于块计算比块传输需要更长的时间,与标准Transformer相比,此过程不会增加开销。具体如下图所示:

使用RingAttention,每个设备的内存需求与块大小成线性关系,而与原始输入序列长度无关,这也消除了各个设备施加的内存限制。其中,b为批大小,h为隐藏维数,n为头数,s为序列长度,c为块大小。


 

四、训练过程

训练分两阶段:

  • 阶段一:渐进式的纯文本训练
  • 阶段二:多模态训练

第一阶段的目标是建立一个能够处理长文本序列的语言模型,以理解复杂的文档和长文本内容。为实现这一目的,研究人员采取了渐进式的训练方式,使用总计33B Token、由图书内容组成的Books3数据集,从32k开始训练,逐步将窗口扩增至1M。

第二阶段则是将视觉信息(如图像和视频)整合到模型中,以提高对多模态数据的理解能力。在此阶段,研究人员对第一阶段的LWM-Text模型进行了架构修改,以支持视觉输入。他们使用VQGAN将图像和视频帧转换为token,并与文本结合进行训练。

这一阶段同样采用循序渐进的训练方法, LWM首先在文本-图像数据集上进行训练,然后扩展到文本-视频数据集,且视频帧数逐步增多。

理解/生成任务混合训练。在训练过程中,模型还会随机交换文本和视觉数据的顺序,以学习文本-图像生成、图像理解、文本-视频生成和视频理解等多种任务。

五、效果与性能

这里主要关注多模态相关能力。

1、论文中样例测试

视频理解问答:1小时时长视频问答,在自己制作的视频和问题上准确回答,而GPT-4V、GeminiPro、Video-LLaVA等标杆产品均回答错误。

穿霸王龙服饰的人骑的什么车?LWM正确回答motorcycle。

图像视频生成:

2、论文中benchmark结果分析

图像通用问答(VQA):在VQAv2测评集上,LWM 55.8,CogAgent 83.7,LLaVA-1.6(Vicuna-7B版本)81.8,差距很大,这个成绩根本进不了VLM第一梯队。

富文本图像问答(Text-rich VQA):论文仅给出TextVQA任务测试结果,LWM 18.8,开源标杆CogAgent 76.1,差距巨大。DocVQA、ChartQA等任务没有进行测试。一方面,训练中并没有针对此类任务数据专项提升,另一方面,与CogAgent比,LWM图像输入分辨率偏低,没有ViT及CA等规模化视觉处理机制,也造成在复杂的富文本图像理解任务上天然没有优势。

整体看,LWM当前版本在图像和视频理解问答上的表现并不出众、不在第一梯队。论文中对此进行解释。“我们的模型在基线中表现平均,表现逊于SOTA模型。我们假设这可能是由于有限的文本图像和文本视频对齐训练,而基线利用了经过更广泛、大规模数据训练的CLIP视觉骨干。相比之下,我们的模型使用VQGAN tokens,需要从头开始学习文本-图像对齐,并且由于VQGAN重建图像的文本能力较差,通常难以完成OCR任务。然而,我们相信,我们的模型将是未来基于VQ的视觉语言模型架构的一个很有前途的方向,并且可以通过更严格的训练和学习更好的tokenizers来提升表现。”

六、验证

LWM官方未提供演示系统,需自己搭建(据官方issue:32KText版约1xA100 80G,512KText版约8xA100 80G,1M多模态版至少8个A100 80G且,基于JAX),后续进行本地实际部署验证,并给出多模态能力的实测结果。

附录

  • 《WORLD MODEL ON MILLION-LENGTH VIDEO AND LANGUAGE WITH RINGATTENTION》,https://arxiv.org/pdf/2402.08268v1.pdf
  •  《Ring Attention with Blockwise Transformers for Near-Infinite Context》,https://browse.arxiv.org/pdf/2310.01889.pdf
  • LWM代码:GitHub - LargeWorldModel/LWM

也许您对下面的内容还感兴趣: