ARIA : 开源多模态原生 MoE 模型 - 文章 - 开发者社区

picture.image

点击下方卡片，关注“ 慢慢学AIGC ”

摘要

信息以多种形式存在。多模态原生 AI 模型对于整合现实世界中的信息并提供全面的理解至关重要。虽然已有专有的多模态原生模型，但它们缺乏开放性，给应用甚至适应带来了障碍。为弥补这一缺陷，我们引入了 ARIA，一个开放的多模态原生模型，在各种多模态、语言和编码任务中表现出一流的性能。 ARIA 是一个专家混合（MoE）模型，对每个视觉和文本 token 分别激活了 39 亿和 35 亿个参数。 它在多模态任务上超越了 Pixtral-12B 和 Llama3.2-11B，并在多项任务中与最好的专有模型具有竞争力。我们 从零开始按四阶段流水线对 ARIA 进行预训练，使模型逐步具备语言理解、多模态理解、长上下文窗口和指令跟随等强大能力 。我们开源了模型权重及相关代码库，便于 ARIA 在实际应用中的轻松适配。

论文：https://arxiv.org/pdf/2410.05993

代码: https://github.com/rhymes-ai/Aria
权重：https://huggingface.co/rhymes-ai/Aria
官网: https://rhymes.ai/

1、引言

在本报告中，我们介绍了 ARIA，这是首个开放的多模态原生专家混合 (MoE) 模型。多模态原生一词在以往文献中用于描述不同的模型能力，但尚未达成明确共识。我们在此给出一个可量化的定义：多模态原生模型指的是一个能够在多种输入模态（如文本、代码、图像、视频）中具备强大理解能力的单一模型，其性能可与相似容量的模态专用模型相媲美或超越。我们的定义与 GPT-4o 或 Gemini-1.5 等专有多模态模型的用户体验一致，用户无需区分来自不同模态的输入，模型能够无缝处理并整合多模态输入。

picture.image

虽然专有的多模态原生模型并不少见，但它们的训练方法大多未公开。因此，大多数开源模型是模态专用的，或在多模态上表现欠佳。我们的研究填补了这一空白，并提出了从头开发多模态原生模型的训练方法，包括以下关键方面：

模型架构 。我们模型的核心是一个 细粒度的专家混合解码器 ，相比于稠密解码器，专家专用化带来的更高效参数利用率使其具有更快的训练和推理速度。 ARIA MoE 每个文本 token 激活 35 亿参数，总共包含 249 亿参数 。 变长、不同大小和长宽比的视觉输入通过一个 4.38 亿参数的轻量视觉编码器编码为视觉 token 。ARIA 拥有 64k token 的长多模态上下文窗口。
数据。ARIA 在 6.4 万亿个语言 tokens 和 4000 亿个多模态 tokens 上进行了预训练。我们制定了一套严格的流程，从多种来源中筛选出高质量的数据。多模态预训练数据包含四大类：来自公共爬网的图文序列、合成图像描述、文档转录和问答对，以及合成视频描述和问答对。
训练流程 。我们设计了一个四阶段的训练流程，包括 语言预训练、多模态预训练、多模态长上下文预训练和多模态后训练 。每个阶段的设计都旨在逐步增强模型的特定能力，同时保留在早期阶段获得的能力。我们的流程高效利用数据和计算资源，以最大化模型性能。

picture.image

按照该训练方法，ARIA 展现了作为开放多模态原生模型的先进性能。与 Pixtral-12B [Mixtral, 2024] 和 Llama3.2-11B [Dubey 等人, 2024] 相比，ARIA 在各种多模态、语言和编码任务中表现更佳，并因激活参数数量更少而推理成本更低。此外，ARIA 在多项多模态任务中的表现也与 GPT-4o 和 Gemini-1.5 等专有模型相当。详细的基准测试结果见表 1。

picture.image

我们在 Apache 2.0 许可下发布 ARIA，可供学术和商业用途免费使用。为了便于采用，我们开源了一个训练框架，使得在各种数据源和格式上微调 ARIA 变得更加简便，最低只需一张 GPU。

2、模型细节

2.1 细粒度专家混合（MoE）模型

MoE 已成为比稠密模型更受欢迎的架构，用于构建计算高效的大型语言模型 [Fedus 等人, 2022; Jiang 等人, 2024; Dai 等人, 2024; Ludziejewski 等人, 2024]。 MoE 的核心思想是将 Transformer 中的每个前馈层 (FFN) 替换为一组专家，每个专家在结构上与 FFN 相同。每个输入 token 仅会被路由至每层中的一部分专家进行处理。专家激活的稀疏性确保了 MoE 层的计算效率。

由于多模态数据的多样性，我们假设专家专用化对于多模态 MoE 理解不同数据分布的输入至关重要。为此，我们采用大量细粒度的专家，具有比标准 FFN 更小的 FFN 隐藏维度，类似于 [Dai 等人, 2024]。特别地，ARIA 在每个 MoE 层中有 66 个专家，其中 2 个专家在所有输入之间共享以捕捉常识知识，而每个 token 通过路由模块激活另外 6 个专家。

picture.image

表 2 显示了 ARIA 详细的架构配置。

picture.image

ARIA 与之前的多模态 MoE 模型有显著不同，以往的模型要么设计特定模态的专家架构，要么依赖于从稠密模型中再利用 [Lin 等人, 2024b; Shen 等人, 2023; Lin 等人, 2024a]。我们的多模态原生 MoE 从头开始预训练，采用模态通用的专家。我们在第 4.2 节展示了多模态专家专用化在预训练后自然出现的现象。

2.2 视觉编码器

我们设计了一个 轻量级视觉编码器 ，将

视觉输入（如图像或视频帧）转换为与词嵌入相同特征维度的连续视觉 token

，从而使 MoE 能够无缝整合视觉和语言输入。

借鉴先前的研究 [Li 等人, 2023; Bai 等人, 2023; Laurençon 等人, 2024]，我们的 视觉编码器由一个视觉 Transformer (ViT) 和一个投影模块组成 。ViT 接收原始长宽比的图像，将其分割为不同长度的图像块序列 [Lee 等人, 2023; Dehghani 等人, 2023]，以保留图像的固有信息结构。我们将图像尺寸分为三类：
(1) 中等分辨率图像 ，长边调整为 490 像素；
(2) 高分辨率图像 ，长边调整为 980 像素；
(3) 超高分辨率图像 ，动态分解为多个高分辨率图像，类似于 Liu 等人 [2024] 的方法。

我们用 SigLIP-SO400M 模型 [Zhai 等人, 2023] 初始化 ViT 权重，并在我们的多模态数据上继续对 ViT 进行预训练。

我们的 投影模块将 ViT 的图像嵌入序列转换为视觉 token 序列 。它 由一个交叉注意力层和一个前馈层 (FFN) 组成 。交叉注意力层使用一组可训练向量作为 query，将图像嵌入作为 key。中等分辨率图像使用 128 个 query 处理，高分辨率图像则使用额外的 128 个 query（共 256 个 query）。交叉注意力层的输出传入 FFN 层，最终输出供 MoE 解码器进一步处理的视觉 tokens 。

2.3 基础设施

ARIA 使用经过大幅修改的 Megatron 框架 [Shoeybi 等人, 2019] 进行训练。我们 放弃了流水线并行 ，而是 采用了专家并行 [Lepikhin 等人, 2020] 和 ZeRO-1 数据并行 [Rajbhandari 等人, 2020] 的组合来优化性能。得益于精心设计的并行方法和较小的模型规模， ARIA 可以在不使用张量并行的情况下高效训练 ，从而显著减少通信开销并提升训练效率。

为防止路由崩溃并确保专家激活的均衡性，我们实现了 负载均衡损失 。我们发现，由于专家数量众多，以往研究中的专家级负载均衡损失 [Fedus 等人, 2022; Dai 等人, 2024] 对于我们的 MoE 模型而言限制过多。因此，我们将 负载均衡放宽至专家组 ，每组包含 8 个细粒度专家。此外，我们采用了 z-loss [Zoph 等人, 2022] 来稳定训练过程。

3、训练

在本节中，我们详细描述了四阶段训练流程。 每个阶段中，模型的目标是学习新能力，同时保持之前习得的能力 。我们在 每个阶段进行评估 ，以确保在数据和计算资源方面实现高效的目标。

3.1 语言预训练

第一个阶段使用大量筛选过的 语言数据 对 MoE 解码器 进行预训练，数据被转换为离散文本 token，并通过下一 token 预测损失来训练，从而使 MoE 学习有关世界的通用知识。 上下文窗口长度设为 8K 个 tokens 。

语言数据。我们的语言预训练数据总计包含 6.4 万亿个 tokens ，来自多种数据源，内容截至 2024 年 5 月。我们在不同粒度上对数据进行去重，并通过规则和模型驱动的质量过滤相结合的方式严格筛选数据。为了增强模型的上下文学习能力，我们采用 数据聚类 并在训练过程中将相似数据打包在同一序列中，类似于 Shi 等人 [2023] 的方法。然而，原方法在处理数万亿 token 时不够具备可扩展性，且容易生成大量长尾结构。我们改用 最小生成树算法进行语言数据聚类 ，显著提高了性能。

3.2 多模态预训练

第二个阶段使用 语言和多模态数据的混合 对 MoE 解码器和视觉编码器 进行预训练，采用相同的下一 token 预测损失。该阶段旨在赋予模型广泛的多模态理解能力，同时保持或提高其语言理解能力。为此， 语言数据包含了 1 万亿个高质量 tokens，涵盖代码、推理和知识等主题 。 多模态数据包含了来自多种来源的 4000 亿个 tokens ，分为以下四大类：

图文交错的网页数据 。我们从 Common Crawl 提取并过滤网页。首先删除图像或文本质量较低的网页，然后去重图像，删除图文总体 CLIP 得分较低的网页 [Radford 等人, 2021]。此外，如果某句与图像的 CLIP 得分较高且位于图像之前，则将图像移至该句之前。最终我们获得了 1900 亿个图文交错 tokens 。
合成图像描述 。直接从网页图像中提取的替代文本通常简短、不够描述性且含有噪声。先前研究表明， 大规模的合成数据可以改进多模态预训练 [Li 等人, 2022]。因此，我们使用一个 小型模型将替代文本重写为更长、更具描述性的图像描述 ，为 LAION-400M 数据集中的 3 亿图像生成了合成描述 ，总计 700 亿多模态 tokens 。
文档转录和问答 。为了提高模型对文字稠密图像的理解能力，我们使用公开 OCR 方法 将文档图像转录为文本，并使用纯文本、图表 JSON 或表格/公式 LaTeX 代码渲染图像。此外，我们使用语言模型生成合成的问答对，以提升模型的文本转录和理解能力。最终，我们的 多模态文档数据包含 1020 亿个 tokens 。
视频描述和问答 。我们收集了 440 万个长度不一的视频 ，来源多样。我们训练模型生成视频的帧级稠密描述，然后利用语言模型基于这些稠密描述生成问答对和视频摘要。 视频数据总计包含 350 亿个 tokens 。在多模态预训练中，我们选择长度在 8K 以下的样本。

3.3 多模态长上下文预训练

在该阶段，我们在 长序列数据 上进行预训练，将模型的 上下文窗口扩展至 64K 个 tokens 。语言长序列数据从预训练数据源中选取，多模态长序列数据则包括 长视频、长文档以及由短多模态数据构建的合成长序列 。特别地，我们将一系列独立图像串联为输入，并将其图像描述串联为目标。本阶段消耗了 120 亿语言 tokens 和 210 亿多模态 tokens ，其中 33 亿 tokens 中有 69% 为长序列。我们将 RoPE 的基础频率超参数从 100K 增加到 5M。

经过此阶段后，模型在上下文窗口长达 64K 的情况下完美解决了大海捞针任务 [Kamradt, 2023]，并在长视频理解和长多模态文档理解任务上显示出显著的性能提升。

3.4 多模态后训练

最终的后训练阶段将学习率逐步收敛，以确保模型的最终收敛。学习的重点在于 提升模型的问答和指令跟随能力 ，使用高质量的开源数据集与人工标注数据集的混合，涵盖多模态、代码、数学和推理等领域。本阶段共处理了 200 亿个 tokens 。

4、评估与分析

4.1 基准测试结果

在表 1 中，我们将 ARIA 与同规模的领先开源模型及专有模型进行了广泛基准测试比较。

picture.image

在表 3 和表 4 中，我们分别检查了其长上下文多模态理解能力和指令跟随能力。

picture.image

基于评估结果，我们总结出以下关键观察：

ARIA 是同类最佳的开源多模态原生模型 ，在多模态、语言和代码任务上显著优于 Pixtral-12B 和 Llama3.2-11B。
在 多模态任务 （包括 文档理解、图表阅读、场景文本识别和视频理解 ）上， ARIA 与闭源模型相比仍具有竞争力 。
ARIA 在长上下文多模态理解方面表现出色 。现实中的多模态数据本质上十分复杂，通常包含长序列的图文交错输入，例如视频字幕或多页文档等。 ARIA 在理解此类数据方面表现出色，显著优于开源模型，例如 Qwen2-VL-7B [Bai et al., 2023] 和 LLaVA-OneVision-72B [Li et al., 2024]。与专有模型相比，ARIA 在长视频理解上超越了 GPT-4o mini [Wu et al., 2024]，在长文档理解方面优于 Gemini-1.5-Flash，使其成为以高计算效率和时间效率处理长多模态数据的首选。
ARIA 拥有强大的指令跟随能力 ，在多模态和纯语言基准测试上均优于其他开源模型。具体的示例请见第 4.3 节。

4.2 专家模式的模态专门化

我们分析了 ARIA MoE 在多模态预训练阶段后的各层专家激活行为 。为分析，我们选取来自三个领域的多模态数据：自然图像、视频和 PDF 格式图像。 对于每个专家，我们首先计算它对视觉 token 和文本 token 的激活率，分别记为 Rv 和 Rt 。Rv 表示激活该专家的视觉 token 数量与该层所有专家处理的视觉 token 总数之比，而 Rt 表示文本 token 的同类比率。然后，我们计算 Rv/Rt 比值，该比值代表该专家的视觉专门化程度。 视觉专门化程度越高，说明该专家比文本 token 更频繁地被视觉 token 激活 。

图 1 展示了专家视觉专门化值的可视化（上限为 50）。

picture.image