Press "Enter" to skip to content

ImageNet准确率超过90%!谷歌大脑开源V-MoE,用稀疏条件计算来训练目前最大的视觉模型!(NeurIPS…

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

  写在前面

 

稀疏的专家混合网络 (Sparsely-gated Mixture of Experts networks (MoEs)) 在自然语言处理中显示出出色的可扩展性。然而,在计算机视觉中,几乎所有SOTA网络都是 “密集的”,也就是说,对于每个输入,每个参数都要参与计算 。

 

在本文中,作者提出了一种视觉MoE (V-MoE),这是视觉Transformer的稀疏版本,它是可扩展的,并且可以与最大的密集网络性能相近。当应用于图像识别时,V-MoE与最新网络的性能相近,同时在推理时只需要一半的计算开销。

 

此外,作者提出了路由算法的扩展,该算法可以对整个batch中的每个输入的子集进行优先级排序,从而实现每幅图像的自适应计算。这使得V-MOE能够在测试时能够更好的权衡性能和计算量。最后,作者展示了V-MoE在视觉模型缩放方面的潜力,并在ImageNet上训练了一个15B参数的模型,达到了准确率90.35% 。

 

  1. 论文和代码地址

 

 

Scaling Vision with Sparse Mixture of Experts

 

论文地址:https://arxiv.org/abs/2106.05974

 

代码地址:https://github.com/google-research/vmoe

 

  2. Motivation

 

深度学习的一系列研究表明,增加网络容量和数据集大小通常会提高性能。在计算机视觉中,在大型数据集上预训练的大型模型通常会达到SOTA的水平。这种方法在自然语言处理(NLP)中取得了更大的成功,在自然语言处理中,大型的预训练模型无处不在,并且在许多任务上表现得非常好。

 

Text Transformers是迄今为止最大的模型,有些模型的参数超过100B。然而,这种模型的训练和推理成本都是非常昂贵的。这是因为这些深度网络通常是 “ 密集的 ”—每个样本都使用每个参数进行处理,因此具有很高的计算成本。

 

相比之下, 条件计算旨在通过仅将参数子集应用于每个样本来增加模型容量,同时保持训练和推理成本大致恒定 。在NLP中, 稀疏混合专家(MOE) 越来越受欢迎,能够用更少的资源进行训练和推理,同时解锁万亿个参数的模型。

 

在这项工作中,作者探索了大规模视觉模型的条件计算,提出了 视觉MoE (V-MoE) ,这是最近的ViT结构的稀疏变体,用于图像分类。V-MoE用稀疏MoE层代替ViT中的密集前馈层的子集,其中每个图像patch被 “路由” 到 “专家” 的子集。

 

由于路由模式的独特性和不可微性,在深度稀疏模型中进行路由是具有挑战性的。作者探索了各种设计选择,并为V-MoE的预训练和迁移学习提供了一个有效的方法,且性能显着优于其密集的网络。 作者进一步证明了V-MoE模型是非常灵活的,通过调节输入或模型权重的稀疏水平,可以在推理期间对已经训练好的模型进行性能和计算开销的平衡。

 

使用V-MoE,作者可以扩展到15B参数的模型大小,这是迄今为止最大的视觉模型,该模型与最先进的密集模型的性能相匹配,同时只需要更少的时间来训练。V-MOE可以与VIT的成本相匹配,同时实现更好的性能。为了控制性能和计算的权衡,作者提出了 批量优先路由算法(Batch Prioritized Routing) ,该算法重新利用模型稀疏性来跳过某些patch的计算,从而减少了对无信息图像区域的计算。

 

  3. 方法

 

3.1. The Vision Mixture of Experts

 

3.1.1. Conditional Computation with MoEs

 

条件计算旨在针对不同的输入,激活网络的不同子集 。专家混合模型是一种特定的实例化,其中不同的模型 “专家” 负责输入空间的不同区域 。MoE的计算可以表示为,其中是输入的数据,是专家计算的函数,是 “路由” 函数,用来决定每个专家的权重。和都由神经网络参数化。

 

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

 

3.1.2. MoEs for Vision

 

作者在视觉Transformer (ViT) 的背景下探索稀疏性在视觉中的应用。事实证明,ViT在迁移学习设置中可以很好地扩展,能够在更少的预训练计算下达到比CNN具有更高的准确性。ViT将图像处理为patch序列,输入图像首先被分为大小相等的patch,然后被线性投影为token。添加positional embeddings后,token由Transformer处理,该Transformer主要由交替的自注意和MLP层组成。

 

 

MLP具有两层FC结构和GeLU非线性函数:。对于视觉MoE,作者用MoE层替换其中的一个子集,其中每个专家都是MLP。如上图所示。所有专家具有相同的结构,但权重不同。

 

3.1.3. Routing

 

对于V-MOE中的每个MoE层,作者使用路由函数,其中是除了最大k个值的元素,将向量的其他元素设置为零的操作;是在上的采样结果。

 

实验中,作者使用k=1或k=2。在视觉Transformer的上下文中,是网络某一层上的图像token的表示。因此,V-MOE路由patch表示,而不是整个图像。

 

作者在专家权重的Softmax之后用。这允许模型能够在k=1的情况下进行训练,并且在k>1的情况下执行得更好。最后,作者在激活值中添加了少量标准差为的噪声,从而进一步提高了性能。

 

3.1.4. Expert’s Buffer Capacity

 

在训练期间,稀疏模型可能只训练一小部分专家 。这种模式会导致两个问题。首先, 统计效率低下 :在模型专注于单个专家的学习时,该模型并不比密集模型更强大。其次, 计算效率低下 : 向专家分配数据的不平衡可能会导致硬件利用率低下。

 

为了消除不平衡并简化实现,作者将每个专家的缓冲容量(即每个专家处理的token数量)固定,并使用辅助损失函数训练模型,以保持负载平衡。

 

作者将专家的缓冲容量 ()定义在batch处理中的图像数量 (),每个图像的token数量 (),每个token的选定专家数量 (),专家总数 (和容量比率 ()之上,即:Ž。

 

如果路由器向给定专家分配了多于个token,则仅处理其中的个token。剩余的token并不完全 “丢掉”,因为它们的信息由剩余的连接保留。另外,如果,则有多位专家会处理每个token,token永远不会完全丢弃。如果为专家分配的token少于个token时,则其缓冲区的其余部分就用零填充。

 

作者使用容量比率来调整专家的能力。当C>1时,将添加一个空闲容量,以考虑潜在的路由不平衡。当新数据可能来自与上游训练中非常不同的分布时,这通常有助于微调。在C <1的情况下,路由器会被迫忽略一些分配。

 

3.2. Skipping Tokens with Batch Prioritized Routing

 

为了提高路由效率,作者提出了一种新的路由算法(BPR),允许模型对重要token进行优先排序。通过同时减少每个专家的容量,丢弃最没用的token。直观地,并非每个patch对于给定图像进行分类都同样重要的,例如,可以删除大多数背景patch,使模型仅关注具有相关实体的patch。

 

Vanilla Routing

 

路由函数逐行应用于一个batch的输入。一个batch处理包含N个图像,每个图像由P个token组成;的每一行对应于图像的特定token的D维表示。相应地,表示第t个token和第i个专家的路由权重。在所有路由算法中,对于i <j,每个TOP-i分配都比任何TOP-j分配具有更高的优先级。路由器首先尝试分配所有第i个专家选项,然后再分配第j个选项。普通路由算法如下所示:

 

 

Batch Prioritized Routing (BPR)

 

为了关注“最重要”的token,作者提出了计算每个token的优先级分数,并在分配之前对进行相应排序。根据token的最大路由权重对token进行排序,即。前k个权重的总和用同样的方式计算,即。

 

 

作者将路由器输出重用为分配优先级的代理。上图可视化了容量越来越小的Batch Prioritized Routing (BPR)算法的token优先级结果。由于batch中所有图像上的所有token彼此竞争,因此不同的图像可以接收不同的计算量。BPR的算法如下所示:

 

 

  4.实验

 

4.1. Upstream results

 

 

从上表可以看出,本文的模型在各个数据集上都表现不错,并且在ImageNet上基于fine-tuning设置达到了90.35%的准确率。

 

 

上图显示了不同V-MOE和ViT变体的总训练计算量和时间。

 

4.2. Linear few-shot results

 

 

上图显示,在5-shot ImageNet下的计算量和训练时间的结果。

 

 

上表展示了在1,5,10-shot下,本文方法的性能优于其他的方法。

 

4.3. Full fine-tuning results

 

 

上图展示了在ImageNet上进行fine-tuning的结果。

 

  5. 总结

 

在本文中, 作者使用稀疏条件计算来训练一些迄今为止最大的视觉模型 ,显示了在表示学习和迁移学习方面的显着改进。除了V-MoE之外,作者还提出了 批量优先路由(Batch Prioritized Routing ,BPR) ,这利用了模型的稀疏性来引入输入的稀疏性,可以在不进一步调整模型的情况下完成此操作,从而允许重新使用具有稀疏条件计算的训练模型。

 

作者认为,这只是视觉条件计算的开始;工作的扩展可以包括扩大专家数量,减少对数据的依赖以及改善稀疏模型产生特征的迁移。

 

▊  作者简介

 

研究领域:FightingCV公众号运营者,研究方向为多模态内容理解,专注于解决视觉模态和语言模态相结合的任务,促进Vision-Language模型的实地应用。

 

知乎/公众号:FightingCV

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注