大家好,我是Olive,在MiniMax从事强化学习研究。今天我想简要介绍一下我们的新模型Minimax M3,以及我们如何从零开始进行原生多模态训练,并采用MSA(MiniMax稀疏注意力机制)进行优化。首先介绍一下我们的团队:我们是全球少数专注于多模态研究的实验室之一,此前已发布过LLM模型、语音模型、视频生成模型以及代码智能体模型。最近,我们发布了具有前沿性能的开源权重模型Minimax M3,它是首个结合以下三项前沿能力的开源模型。
第一项能力是我们所称的SODA(代码与智能体)功能。该模型能够执行多种想象得到的智能体任务,包括自动任务分解、在代码智能体中工作,并能够替代您当前90%的代码编写工作。第二项能力是模型拥有百万级别的上下文窗口,这极大地解锁了用户长交互场景的基础需求,例如处理完整的代码库。我们使用MSA进行训练,其预填充速度比全注意力机制快9倍,解码速度快15倍。
第三项能力是模型从零开始进行原生多模态训练,这意味着从起步阶段就将视觉理解与语言生成相结合。模型从一开始就能很好地掌握这两种能力,从而为构建更多应用(如计算机使用智能体)奠定基础。首先简要解释稀疏注意力的作用:我们的核心设计原则是更简单的架构总能带来更好的扩展性,这也是贯穿整个设计的理念。
稀疏注意力由两个分支组成:首先是索引分支,它执行轻量级的全注意力计算,使用四个查询头和一个键头。该分支通过KL损失进行训练,以匹配主分支的注意力分数。索引分支用于识别具有最高注意力分数的标记(tokens),随后由主分支对这些选定的标记进行关注。换句话说,稀疏分支并不关注所有标记,而是仅通过索引分支的轻量级计算选择子集进行关注。
接着,稀疏分支利用索引分支选定的块(blocks)来确定注意力计算所需的键和值。我想强调的是,我们的设计并非纯粹由算法驱动,而是采用了算法与基础设施协同设计的思路。以DSA为例,虽然它是一个强大的设计,但高度针对DeepSeek的特定架构(特别是MQA与大头维度的组合)进行了优化。
然而,当我们转向更常用的GQA架构时,直接应用相同的设计变得不那么直观,并引入了一些实际问题,尤其是在基础设施和实现层面。这促使我们做出设计选择:不再简单借用现有的稀疏注意力设计,而是在GQA约束和真实系统效率的限制下重新思考架构。
如果直接应用现有的稀疏注意力架构,会存在哪些问题?首先从算法角度看,GQA具有多个KV头,理想情况下应在不同KV组间提供更多样性。但如果直接重用相同的标记选择策略,多个KV头最终将关注完全相同的选择结果,这实质上抵消了GQA本应提供的多样性。
其次从基础设施角度看,该设计对GPU内存访问并不友好。GPU在读取连续数据时效率更高。在MQA设置中,KV布局可视为一个大头维度的单个KV头(例如1x512),这种结构相对连续且读取高效。然而在GQA中,相同的总维度被拆分到多个KV头(例如4x128),这不再是单次连续读取,而是需要在不同KV头间执行多次独立读取。
这种操作增加了内存访问开销并降低了计算与内存访问的比例。另一个问题在于top-k计算本身:为了选择最重要的标记,索引器需要在大量候选标记上计算top-k分数。然而,top-k存在固有的数据依赖性,难以在GPU上完全并行化。因此,标记级选择引入了显著的计算和合并开销。
那么我们在MSA中如何解决这些问题?在算法层面,我们移除了索引分支输出的聚合。在原设计中,所有索引头合并为单一的选择标记集,但这会消除GQA带来的多样性。相反,我们保持不同KV头间的选择差异。在稀疏注意力设计下使用GQA时,不同的KV头可以关注不同的标记集,这有助于保留GQA的多样性,而非强制所有头共享相同的稀疏注意力。
在基础设施层面,我们将标记级选择转变为块级检索。与标记级选择相比,块级检索显著减少了候选数量并降低了top-k计算及结果合并的开销。同时,块级检索对硬件更友好:它提高了内存访问效率并增加了计算与内存访问的比例。重要的是,我们证明该设计并未引入显著的性能下降。
最后,我们为训练和推理设计了高效的内核(kernels)。今天我不深入讨论内核细节,但这是使稀疏注意力设计变得实用且高效的系统级组件。凭借这一设计,我们能够非常高效地将M3模型的训练扩展至百万上下文窗口。
正如我们之前所讨论的,第三大支柱是原生多模态。我们的模型 MiniMax M3 是使用多模态数据从头开始训练的。这意味着什么?意味着我们的模型从第一步起就同时接受了文本和视觉理解能力的训练,这与当前开源权重模型的非常不同。我们的起点非常明确:在现阶段,独立的视觉语言模型的应用场景相对有限。
但将其与文本理解相结合,尤其是在训练过程中,具有挑战性。因此,关键问题和挑战在于:我们如何在不过度损害核心文本性能的前提下引入多模态能力?为了回答这个问题,我们在大规模上探索了几种训练策略。第一种也是最常见的方法是在预训练衰减阶段之后添加多模态训练,这本质上是一种持续预训练(CPT)风格的方法。
然而,在那时模型已经基本收敛。在此阶段添加多模态数据往往会干扰已习得的文本能力,导致性能下降。另一种方法是在预训练衰减阶段之前进入多模态训练。例如在模型已经见过大量文本标记之后。我们也尝试了这个方向,但发现结果对数据混合和训练超参数(如学习率)高度敏感。
在一个设置中,较低的学习率可能会减少性能下降,但这一结论不一定能推广到其他运行。更重要的是,在大规模预训练期间,这些参数不能仅为了多模态训练而随意更改。这引导我们采取了最终策略:从头开始的原生多模态训练。在这种设置下,多模态能力从一开始就被引入,而不是在文本模型形成其结构后再添加。
在相同的文本标记预算下,我们发现这种方法几乎不会损害文本能力。在这些实验中我们做了许多有趣的分析。其中之一是从注意力图中得出的一个有趣观察:你可以从这张图片中看到,对于 CPT 风格的多模态模型,文本标记往往很少关注视觉标记。在注意力图中,视觉标记区域大多是暗的。
这表明多模态能力并未完全融合。相比之下,原生多模态训练的注意力图看起来整合得更好。我们在这里没有看到文本标记和视觉标记之间的清晰分离。这表明模型从一开始就学会了更自然地融合视觉和文本信息。因此,总体而言,我们的训练策略不仅仅是增加视觉能力,而是在保留最强文本能力的同时将多模态整合到模型中。
当然,我们在多模态架构选择上也做了许多实验。我们的整体架构相对标准。图像或视频由 ViT 处理,产生视觉标记。这些视觉标记随后通过替换文本序列中原始的图像占位符插入到语言模型输入中。之后,视觉嵌入和文本嵌入一起被送入语言模型。
一个重要发现是:在 ViT 内部应用三维注意力对视觉理解非常有帮助。在标准的 ViT 设置中,注意力应用于每个单独的图像或帧内。对于视频输入而言,这意味着不同帧的视觉补丁在 ViT 阶段不会直接交互。但有了三维注意力,跨不同帧的补丁可以直接相互关注。
这赋予了模型更强的时序和跨帧视觉建模能力。我们发现这提高了视觉理解,且这种益处也转移到了某些图像基准测试中,表明它加强了整体视觉表示能力,而不仅仅是改善视频特定任务。一个实际的细节是:我们不能简单地将三维注意力扩展到无限数量的帧,因为那会引入额外的系统负载平衡挑战。
在我们的探索中,使用四帧在性能和效率之间提供了最佳权衡。因此,总体而言,凭借我们的架构设计、我们在各个阶段探索的训练策略以及我们对 ViT 的分析与探索,我们能够从头开始训练具备多模态理解能力的模型。今天的演示到此结束,谢谢。