150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?

近日,土豪Google AI继之前发布的20亿参数的ViT-G/14模型之后,又发布了参数量为150亿的视觉模型V-MoE,这可以称得上迄今为止最大的视觉模型,其在ImageNet上的Top-1 accuracy达到90.35%,这个也超过之前Google提出的Meta Pseudo-Labelling (MPL)(注意这里ViT-G/14模型的训练成本只有MPL的70%左右),但是略低于ViT-G/14模型。

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第1张图片

MPL模型是基于半监督和JFT-300M,而ViT-G/14模型是基于ViT和JFT-3B来实现更好的效果。这次Google提出的V-MoE核心设计是Sparse Mixture of Experts,这其实是一个稀疏模型(sparse model),这个也是借鉴了NLP领域已有的成果,即实现视觉领域的稀疏模型V-MoE,其模型结构如下图所示:

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第2张图片

可以看到V-MoE也是在ViT基础上构建的,都包含个同质的ViT模块,这里最大的变动是将transformer模块里面的MLP替换成一些列的MLPs,每个MLP放在一个单独的device上,只负责处理固定数量的tokens(或者说是image patchs),每个MLP可以看成一个expert,替换后的layer这里称为MoE layer。在处理上,每个图像的每个token在每个MoE层只有个expert来负责处理它,图中的,这样其实每个token只不过经过一个MLP的处理。对一个图像来说,虽然V-MoE模型很大,但是只用了V-MoE模型的很小一部分来处理,所以说这是一个稀疏模型。具体地,每个MoE层可以用以下公式表示:

这里的是输入token,特征维度为,而就是第个expert,即一个2层MLP,这里的就是计算各个experts的权重系数(是expert的总数)。所以MoE这样看起来是一种模型集成,但是实际上这里要限制每个token的experts数量为,只需要使只有个权重值非零,其它值全为零即可。这样就是变成了一个选址函数(routing function),具体实现为:

这里的操作指的是只选择前个较大的,其它的权重全置为0,论文中的或者,而对于权重为0的expert其实我们也不需要计算,所以虽然模型参数量很大,但是计算开销并不会太大,当时只增加了Router带来的开销。

对于稀疏模型来说,一个挑战是模型最终有效的可能其中的一小部分expert,或者说大部分tokens只选择了其中几个experts,这会导致稀疏模型近似等价于一个dense模型,稀疏模型的大容量没有利用好,而且计算上也不均衡。那么解决方案就是限制每个expert在训练过程中的处理能力:所能处理的tokens的最大量,这里称为buffer capacity,记为。假定训练过程的图像总数为,每个图像的patchs总数为,那么可以计算容量比(capacity ratio):

上图中给的实例。当时,这意味着所有experts的容量是用不完的,这个时候会出现前面说的不平衡问题,论文中额外增加了一个辅助loss(load balancing losses)来避免这种问题,在上游训练或者预训练过程中(upstream training),有必要选择大于1的容量以更好地迁移到下游任务中(需要finetune),论文中设定的是。当时,这意味着所有experts的容量是不够的,那么就有一部分token无法被处理,或者说一些token被drop掉了,其实MLP本身就带有残差结构,所以只相当于少了一个MLP的处理,另外时也会减少每个token都不被处理的概率。论文中在下游任务finetune时采用,这个时候有一个问题,就是选择留下哪些tokens来处理,传统的做法是序列地处理每个图像的tokens直到容量用尽,由于图像的tokens顺序是固定的,这对后面的tokens就不太好。论文中提出的方案是Batch Prioritized Routing (BPR),就是先确定各个tokens的重要性,再按这个顺序来处理tokens:

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第3张图片

下图是图像的tokens在不同的下在第一个MoE层被drop的情况:

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第4张图片

设定较小的那么模型的训练成本就可以降低,比如下图基于BPR的V-MoE-L/32 (k=2)在时在性能上就超过原始的Dense ViT-L/32,而此时的训练所需要的FLOPS,前者只有后者的不到90%。

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第5张图片

同时inference time也可以减少,如V-MoE-L/16只需要dense VIT-H/14一半的推理时间,但性能前者还稍好于后者:

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第6张图片

150亿!谷歌发布史上最大视觉模型V-MoE,却最有希望减少碳排放?_第7张图片

image.png-129.2kB

虽然V-MoE参数量更大,但是由于是稀疏的,其实训练和推理的成本反而会更好一点:

This has interesting connotations for recent work in NLP using sparse models; recent analysis shows model sparsity is the most promising way to reduce model CO2 emissions [46] and that 90% of the footprint stems from inference costs.

你可能感兴趣的:(机器学习系列,计算机视觉,机器学习,人工智能,深度学习,神经网络)