主题 28：什么是 Mixture-of-Mamba？

社区文章发布于 2025 年 2 月 20 日

🔳 我们讨论如何使用 Transformer 混合概念和模态感知稀疏性，使 Mamba 选择性状态空间模型 (SSM) 能够处理多模态数据

在 Turing Post，我们特别热衷于探索与 Transformer 等流行方法不同的 LLM 架构。其中之一是 Mamba 选择性状态空间模型 (SSM)，我们曾在最初的 AI 101 剧集中介绍过它。由于其高效处理长序列、高速度和减少内存使用，它是 Transformer 的主要竞争对手之一。AI 最吸引人的地方在于观察不同的架构如何升级以适应新兴趋势。例如，Mamba 并非处理多模态数据的有效选项，而这正是 Mixture-of-Mamba (MoM) 的用武之地。它通过使用专家混合 (MoE) 概念来增强 SSM 的多模态任务处理能力，从而扩展了 Transformer 的优势。MoM 的主要特点——模态感知稀疏性——将 Mamba 核心转变为满足多模态需求的新型强大架构。让我们探讨 MoM 如何改变 Mamba 以及这个迷人的复杂系统如何运作。

📨 点击关注！如果你想直接在收件箱中接收我们的文章，请在此订阅

在本期节目中，我们将讨论

Mixture-of-Mamba (MoM)：有什么想法？
MoM 如何运作？
MoM 有多好？
MoM 的优势
并非没有局限性
结论：Mixture-of-Mamba 为何脱颖而出？
奖励：深入学习资源

Mixture-of-Mamba：有什么想法？

Mamba 是最强大的选择性状态空间模型 (SSM) 之一。SSM 本质上是一种 AI 模型，可以高效地处理数据序列，例如句子或视频。它们被认为是 Transformer 的有力替代品，因为 Transformer 虽然强大但计算成本高昂。Mamba 尤其高效，并且相对于 Transformer 具有以下优势：

高效处理长序列： Mamba 实现的序列长度缩放是线性的，而 Transformer 则是二次方缩放。
更快的推理： 由于其线性时间处理，Mamba 的推理速度可以比 Transformer 快五倍。
减少内存使用： 它避免了 Transformer 中注意力机制对内存的巨大需求。
可并行训练： 通过将 SSM 表示为卷积，Mamba 能够实现类似于卷积神经网络 (CNN) 的并行训练，从而缩短训练时间。

然而，存在一个大问题——Mamba 无法很好地利用不同类型的数据，并且将所有输入数据（如文本、图像或语音）都以相同的方式处理。这限制了 Mamba 在多模态任务中的有效性。

问题来了：我们如何才能将 Mamba 的优势扩展到多模态数据，使其成为一个更强大的架构？

来自斯坦福大学、卡内基梅隆大学和 Meta 的 FAIR 的研究人员找到了解决方案。他们转向了专家混合 (MoE) 的想法，该技术允许模型仅使用其结构的一部分来处理特定输入。特别是，他们受到 Transformer 混合 (MoT) 的启发，该技术根据输入类型选择性地激活不同的处理组件。因此，他们在此基础上构建了新的 SSM 架构——Mixture-of-Mamba (MoM)，这使得模型更能“感知”不同数据类型，同时保持计算效率。让我们来探讨 MoM 如何使 Mamba 成为多模态的。

MoM 如何运作？

“MoM 通过 Mamba 块的模态特定参数化引入模态感知稀疏性。”（原论文）——让我们一步步分解这句话。

MoM 将模态感知稀疏性整合到 Mamba 的核心。这意味着 MoM 不再对所有数据类型应用相同的参数，而是动态地为每种输入类型（文本、图像或语音）选择正确的处理方法。其内部运作方式如下。

MoM 使用 Mixture-of-Mamba 块构建，这些块对不同类型的输入数据应用单独的处理规则，同时在有意义的地方共享通用组件。

图片来源：原论文

模态感知稀疏机制像一个动态路由系统

MoM 使用模态掩码来区分不同类型的输入 token（文本、图像或语音）。
它根据模态激活正确的权重集。这个过程称为模态特定参数化。
MoM 同时处理不同模态的 token，确保高效的训练和推理。

为了实现模态感知，MoM 修改了 Mamba 模型中的关键投影层

输入投影层： 将原始数据转换为其模态特定的初始表示。
中间处理层： 应用模态特定的变换来细化模型内部的数据表示。
输出投影层： 将处理后的数据转换为其最终输出格式。

通过这样做，模型为每种数据类型应用最有效的处理方法，使训练更快、更高效。

虽然模型中的投影层是分离的，但其某些部分仍然是共享的。它们是什么？

状态转换和卷积保持共享，因为它们不依赖于数据类型。它们也用于保持计算效率。

状态转换跟踪并更新随时间变化的序列信息，像记忆单元一样帮助捕获所有模态的长期依赖关系。
一维卷积层捕获序列中的局部模式，帮助模型处理不同模态（包括文本、图像、语音和视频）的时间相关信息。

这种设计使 MoM 既能专门处理不同模态，又能通过共享关键操作实现计算效率。现在让我们来看看 MoM 的性能结果。

MoM 有多好？

为了展示 MoM 模型的优势，研究人员在三种不同的 AI 训练设置中测试了 MoM

Transfusion – 它融合了文本和连续图像数据。
Chameleon – 一个混合文本和离散图像数据的设置。
三模态框架 – 一个更复杂的设置，除了文本和图像之外，还包括语音作为第三种模态，并将所有这些都视为离散 token。

结果令人印象深刻。

Transfusion 设置
- 图像性能
  - 在 1.4B 规模下，MoM 降低了 2.2% 的损失，同时仅使用了 34.76% 的 FLOPs（每秒浮点运算次数，衡量计算性能）。
  - 较小模型（760M、163M）显示出类似趋势，损失降低 2.37%，FLOPs 减少 60%。
- 文本性能： MoM 提高了验证损失并更好地泛化，同时使用的 FLOPs 更少。它更快地达到准确性，提高了训练效率。
- 整体效率： MoM 降低了跨任务的训练损失，同时将 FLOPs 削减了高达 86.11%。

图片来源：原论文

Chameleon 设置： 图像和文本都被视为离散 token。
- 图像性能： MoM 将图像损失降低了高达 3.46%，仅使用了 25.9% - 42.5% 的 FLOPs。
- 文本性能： 在 1.5B 规模下，MoM 将文本损失降低了 3.01%，使用了 65.4% 的 FLOPs。

图片来源：原论文

三模态设置： 图像、文本和语音
- 语音性能
  - 443M 规模：训练损失降低 7.14%，使用 19.2% 的 FLOPs。
  - 1.5B 规模：损失降低 5.75%，仅使用 24.8% 的 FLOPs。
- 总体指标
  - 训练损失提高高达 3.67%，FLOPs 削减高达 56.2%。
  - MoM 在所有三种模态中持续提高性能。

图片来源：原论文

在 Mixture-of-Mamba 的结果中，MoM 在使用更少 FLOPs 的同时始终获得更高的准确性，使其成为比 Mamba 和 Transformer 更高效的模型。接下来，我们将汇总 MoM 的所有优势，以精确阐明它如何改进 Mamba 选择性 SSM。

MoM 的优势

在 MoM 中，仅激活相关参数，从而优化了计算、训练速度和成本。
模态感知稀疏性优势： 它使 MoM 能够更有效地专注于文本、图像和语音处理，以适应其独特结构的方式处理模态。在模型的每个部分使用模态感知稀疏性比选择性应用它能带来更好的结果。
可扩展且灵活： 它适用于不同的训练策略，如基于扩散的图像学习和基于 token 的处理。
MoM 在三种多模态设置中始终优于传统密集模型： 文本 + 连续图像；文本 + 离散图像；以及文本 + 图像 + 语音。
MoM 提高了计算效率，将成本降低了高达 65%，同时保持或提高了性能。
MoM 以更少的训练步骤达到相同的准确性，使其比 Mamba Dense 和 Flex-Attention Transformer 训练更快。
与基线模型相比，实现了显著的损失减少，从而更好地泛化到未见过的数据。
它降低了能源消耗，使 AI 更易于访问且更环保。
MoM 可以与 MoE 技术结合，为多模态 AI 的进一步效率改进开辟了可能性。