【V-MoE】论文学习笔记

date

Jul 31, 2023

Last edited time

Jul 31, 2023 01:48 PM

status

Published

slug

v-moe

1 Abstract

背景：

MoE 在 NLP 领域表现出出色的可拓展性；

CV 的几乎所有 SOTA 网络都是 “dense”，i.e. 对于每个输入，每个参数都要参与计算。

MoE 是近年来快速发展的一种模型加速方法。

MoE 的思想：
研究者认为在面对不同类型的数据时，model 也可以选择更 “擅长” 的 Expert（特定的网络结构，MoE 中是 MLP）来处理。
在训练过程中，与其对每个输入都要训练整个模型，不如根据输入的不同，仅触发并更新该模型中的某一部分，以降低模型总体的训练代价，学习出更大参数量的模型；
在推理过程时，可以针对不同的下游任务，将大模型拆解得到一个体积更小，但是效果更好的子模型，以实现更加高效的推理。

本质：V-MoE 是一个 Vision Transformer（ViT）的稀疏版本。

优点：

V-MoE 有出色的可扩展性；

V-MoE 与最大的 dense 网络相当的性能；

图像识别领域，V-MoE 与最新的 SOTA 网络性能相近，但推理时只需要一半的计算开销。

Google Brain 提出的路由算法可以对对整个 batch 中的每个输入的子集进行优先级排序，从而实现每幅图像的自适应计算；这使得 V-MoE 可以更好地权衡性能和计算量；

V-MoE 具有 vision model 缩放的潜力：在 ImageNet 上训练一个 15B parameter model，精度达到 90.35%。

2 Introduction

conditional computaion：通过仅将参数的子集 (a subset of paranmeters) 应用于每个样本，增加 model 的容量，与保持训练和推理成本的大致恒定。

目的：Google Brain 探索了大规模 vision model 的条件计算，提出 V-MoE。用于图像分类任务。

创新点：

V-MoE 是 ViT 的一种稀疏变体。V-MoE 用稀疏的 MoE 层代替 ViT 中密集的前馈层的子集，其中每个图像 patch 被 “路由” 到 “Expert” 的子集。“Expert” 本质上是多层感知机（MLPs）。

由于路由的独特失效模式和不可微性，在深度稀疏模型中应用路由充满挑战。Google Brain 提出了一种有效的路由算法——批量优先路由算法（Batch Prioritized Routing Algorithm，BPR）：该算法重新利用模型稀疏性来跳过某些 patch 的计算，从而减少了对不提供信息的图像区域的计算。这使得 V-MoE 可以很好的权衡性能与计算量，其性能显著优于目前的密集网络。

主要贡献：

增大 vision model 的规模：V-MoE 是一个分布式、稀疏激活的 ViT。Google Brain 训练的 model 有多达 24 个 MoE 层，每层 32 个 Expert（MLPs），几乎达到 15B（15-bilion）参数。models 可以被稳定地训练，无缝地用于迁移，并只需 1000 个数据点上成功地 fine-tune。这个最大的 model 在经过 fine-tuned 后，在 ImageNet 上实现了 90.35% 的测试精度。

更强的性能：V-MoE 在上游任务（upstream：预训练）、小样本学习（Few-shot）和充分 finr-tune 指标上，显著的优于密集的 counterpart。在推断阶段，V-MoE 可以调整为：(i) 达到最大的密集 model 的性能，而只使用一半的计算量；(ii) 在相同的成本下显著优于最大的密集 model。

批量优先路由：我们提出了一种新的 BPR，允许 V-MoE 丢弃最没用的 patches。这样，我们对每个图像的计算量就减少了。而且，V-MoEs 可以在节省 20% FLOPs 的同时，达到和密集 models 一样的性能。

理论上的进步：提供了一些路由决策的可视化，揭示了有助于激励设计决策的模式和结论，并可能进一步提高对该领域的理解。

3 Methods

3.1 The Vision Mixture of Experts (V-MoE)

3.1.1 Conditional Computation with MoEs

condtional computation 旨在针对不同的输入，激活网络的不同子集。

MoE 是一种特定的实例化，其中不同模型 “Expert” 负责输入空间不同的区域。

其中，表示这一层的输入，表示 Expert 计算的函数，表示用来为 Expert 决定输入条件权重的路由函数。

但是按照这个定义，这仍然是一个密集的网络。然而，如果和是稀疏的，i.e. model 被限制为仅分配个非零权重，则不需要计算未使用的 Expert，这样就在训练和推理时降低了模型的计算量。

3.1.2 MoEs for Vision

ViT：将图像处理为 patch 序列，输入图像首先被分为大小相等的 patch，然后被线性投影为 token。添加 positional embeddings 后，tokens 由 Transformer 处理，该 Transformer 主要由交替的 self-attention 层和 MLP 层组成。MLP 由 2 个 FC 层 (fully-connected layers) 和 GeLU 非线性函数组成：

V-MoE 的模型结构和 MoE Transformer 的结构类似，如下图所示。研究者将 ViT 中的 feedforward 层（FFN，即图中的 MLP）替换为多个 FFN 的混合集成（其中每一个 FFN 为一个 Expert）。之后，对于每一个模块的输入，ViT 通过学习路由单元（Router）来选择性地激活若干个 Expert，并为每一个 Expert 赋予其权重。

V-MoE 其中的每一个 Expert 都是一个 MLP。如上图所示，Expert 具有相同的结构：，但有不同的权重。

3.1.3 Routing

在 V-MoE 中，每张图片中不同的片段（即不同的 Token）将被 Router 分配到不同的几个 FFN 中。一般而言，Router 会为每个 Token 选择 Top-K 个 FFN（K 一般取 1 或 2），而在每一个 ViT 模块中大概会有 E 个 FFN（在本论文中 E 取值为 32）。通过这种方式，V-MoE 模型的大小可以得到提升，而且在模型规模提升的情况下，每一个 Token 对应的计算量都不会发生太大变动。

对于 V-MOE 中的每个 MoE 层，上述路由函数（Router）定义为：

其中，：将输入向量中除了最大的个元素外元素都设置为。实践中，一般取或。表示网络某一层上的图像 token。因此，V-MoE 只路由 patch representations，而非整幅图像。

不同于 MoE 中先再 softmax ⁡ softmax，V-MoE 中先 softmax ⁡ softmax 再。这样才使得时也可以训练，否则路由梯度会处处为 0。最后，通过添加少量符合正态分布标准差的噪声，进一步提高性能。

3.1.4 Expert’s Buffer Capacity

Routing 带来的问题：

在训练期间，稀疏模型可能只训练一小部分 Expert 。这种模式会导致两个问题。

统计效率低：model 专注于单一 Expert 的学习，该 model 不会比密集的 model 更强大。

计算效率低：向 Expert 分配数据的不平衡可能会导致硬件利用率低。

解决方案：

将每个 Expert 的缓冲容量（Buffer Capacity，即每个 Expert 处理的 tokens 数量）固定，并使用辅助损失函数训练 model，以保持负载均衡。

Expert 的缓冲能力：

其中，是 batch 重要处理的图像数量，是每个图像的 token 数量，是每个 token 的选定 Expert 数量，是总的 Expert 数量，是容量比率。

如果路由器向 Expert 分配了多于个 tokens，则仅处理其中的个 tokens。剩余的 tokens 并不完全 “丢掉”，因为它们的信息由剩余的连接保留。另外，如果，则有会多位 Expert 处理每个 token，tokens 永远不会完全丢弃。

如果路由器为 Expert 分配少于个 tokens，则其缓冲区的其余部分就用零填充。

Google Brain 使用容量比率来调整 Expert 的能力。当时，将添加一个空闲容量，以考虑潜在的路由不平衡。当新数据可能来自与上游训练非常不同的分布时，这种设置非常有助于 fine-tuning。在的情况下，路由器会被迫忽略一些分配。在 3.2 3.2 3.2 中提供了一个新算法，利用 来丢弃最无用的 tokens，并在推理阶段节省算力。

3.2 Skipping Tokens with Batch Prioritized Routing（BPR）

我们前面无数次提到过：为了提高路由效率，Google Brain 提出了一种新的路由算法（BPR）。

基本思想：允许模型对重要 token 进行优先排序。通过同时减少每个专家的容量，丢弃最没用的 token。直观地，并非每个 patch 对于给定图像进行分类都同样重要的，例如，可以删除大多数背景 patch，使模型仅关注具有相关实体的 patch。

3.2.1 Vanilla Routing

路由函数逐行应用于一个 batch 的输入。一个 batch 处理包含个图像，每个图像由个 token 组成；的每一行对应于图像的特定 token 的维表示。相应地，表示第个 token 和第个 Expert 的路由权重。在所有路由算法中，对于，每个分配都比任何分配具有更高的优先级。路由器首先尝试分配所有第个 Expert，然后再分配第个 Exeprt。普通路由算法如下所示：