Scaling Vision with Sparse Mixture of Experts
Carlos Riquelme Google Brain、Joan Puigcerver * Google Brain、Basil Mustafa * Google Brain、Maxim Neumann Google Brain Rodolphe Jenatton Google Brain、André Susano Pinto Google Brain、Daniel Keysers Google Brain、Neil Houlsby Google Brain These authors contributed equally. Correspondence to { rikel, jpuigcerver, basilm
Abstract
Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are “dense”, that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains $90.35\%$ on ImageNet.
中文速览
稀疏激活的混合专家网络（Mixture of Experts, MoE）在自然语言处理领域已被证明具有出色的扩展性，但视觉领域的高性能模型几乎清一色是"密集"网络，即每张图片都要经过所有参数的计算。本文将这一稀疏化思路引入视觉领域，提出 Vision MoE（V-MoE）——在视觉 Transformer（ViT）的基础上，用稀疏 MoE 层替换部分前馈层，让每个图像块（patch）只被路由到少数几个"专家"MLP 处理。研究团队进一步提出"批量优先路由"（Batch Prioritized Routing）算法，让模型在推理时自动跳过信息量少的图像区域，从而节省计算量。实验结果显示，V-MoE 在保持与最优密集模型相当精度的同时，推理计算量最多可减少一半；团队还成功训练了一个约 150 亿参数的 V-MoE 模型，在 ImageNet 上达到 90.35% 的准确率，创下视觉模型规模的新纪录。这项工作表明，稀疏条件计算同样适用于视觉大模型，为在有限算力下扩展视觉系统提供了切实可行的路径。
原文 arXiv:2106.05974;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2106.05974v1