Elamur Research⌕
☰
home / papers / 2021 / Scaling Vision with Sparse Mixture of Experts

Scaling Vision with Sparse Mixture of Experts

Paper · 2021 · Carlos Riquelme, Joan Puigcerver, Basil Mustafa et al.

Summary

Large vision networks are usually dense: each parameter touches every input, so growing capacity drives up compute. This paper brings sparse Mixture of Experts, already successful in NLP, to image recognition. The authors build V-MoE on the Vision Transformer (ViT). Some of its feedforward MLP blocks are swapped for banks of expert MLPs, and a learned router sends each image patch, not the whole image, to only a few experts. Models are pre-trained on JFT-300M and then transferred through linear few-shot evaluation and full fine-tuning. The paper also proposes Batch Prioritized Routing. This scheme ranks tokens across a batch so that low-value patches can be dropped, which allows compute to vary per image and to be reduced at test time without retraining. V-MoE equals or beats dense ViT models at roughly half the inference cost. The largest version, with about 15B parameters, is the biggest vision model reported so far. These results show that conditional computation is a workable way to scale vision models efficiently.

Key points

  • Largest model: ~15B parameters, up to 24 MoE layers with 32 experts each; fine-tuned ImageNet accuracy reaches 90.35%.
  • Default setup: k = 2 experts selected per token, E = 32 total experts, and a capacity ratio C = 1.05 during upstream training. MoE layers sit on alternate blocks (Every-2) or only on the final blocks (Last-n).
  • Routing applies top-k after the softmax, which makes k = 1 trainable, and adds small Gaussian noise. Fixed expert buffers and auxiliary losses keep the load balanced across experts.
  • Pre-training data: JFT-300M, with ~305M images across 18,291 classes. Models still transfer well when fine-tuning uses only about 1,000 examples.
  • Batch Prioritized Routing matches dense model quality while using 20% fewer training FLOPs. The values of k and C can be changed after training to trade accuracy against cost.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢GoogleMountain View · United States492 works · 440,463 citations🏢AmazonSeattle · United States128 works · 41,700 citations

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksComputer Vision and Pattern Recognition · 3,486 works

Cite

Cite as Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, Andre Susano Pinto, Daniel Keysers and Neil Houlsby. 2021. Scaling Vision with Sparse Mixture of Experts. arXiv:2106.05974.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2106.05974 · arXiv 2106.05974

@misc{riquelme2021scaling,
  title = {{Scaling Vision with Sparse Mixture of Experts}},
  author = {Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, Andre Susano Pinto, Daniel Keysers and Neil Houlsby},
  year = {2021},
  doi = {10.48550/arxiv.2106.05974},
  eprint = {2106.05974},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2106.05974}
}

References

70 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. A. Abbas and Y. Andreopoulos. Biased mixtures of experts: Enabling computer vision inference under data transfer limitations. IEEE Trans. Image Processing, 29:7656–7667, 2020. find
  2. K. Ahmed, M. H. Baig, and L. Torresani. Network of experts for large-scale image categorization. In ECCV, 2016. find
  3. A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lučić, and C. Schmid. ViViT: A video vision transformer. arXiv preprint arXiv:2103.15691, 2021. find
  4. E. Bengio, P.-L. Bacon, J. Pineau, and D. Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06297, 2015. — [arXiv:1511.06297](https://arxiv.org/abs/1511.06297) find
  5. Y. Bengio. Deep learning of representations: Looking forward. In International Conference on Statistical Language and Speech Processing, pages 1–37, 2013. find
  6. Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013. find
  7. L. Breiman, J. Friedman, C. J. Stone, and R. A. Olshen. Classification and regression trees. CRC press, 1984. find
  8. A. Brock, S. De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization. arXiv preprint arXiv:2102.06171, 2021. — [arXiv:2102.06171](https://arxiv.org/abs/2102.06171) find
  9. T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
  10. K. Chen, L. Xu, and H. Chi. Improved learning algorithms for mixture of experts in multiclass classification. Neural networks, 12(9):1229–1252, 1999. find
  11. G. Cheng, J. Han, and X. Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEEE, 105(10):1865–1883, Oct 2017. find
  12. K. Cho and Y. Bengio. Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning. arXiv preprint arXiv:1406.7362, 2014. find
  13. M. Cimpoi, S. Maji, I. Kokkinos, S. Mohamed, and A. Vedaldi. Describing textures in the wild. In Computer Vision and Pattern Recognition (CVPR), 2014. find
  14. E. D. Cubuk, B. Zoph, J. Shlens, and Q. Le. Randaugment: Practical automated data augmentation with a reduced search space. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, 2020. find
  15. A. Davis and I. Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint arXiv:1312.4461, 2013. — [arXiv:1312.4461](https://arxiv.org/abs/1312.4461) find
  16. J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, 2009. find
  17. L. Denoyer and P. Gallinari. Deep sequential neural network. arXiv preprint arXiv:1410.0510, 2014. — [arXiv:1410.0510](https://arxiv.org/abs/1410.0510) find
  18. J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
  19. G. S. Dhillon, P. Chaudhari, A. Ravichandran, and S. Soatto. A baseline for few-shot image classification. In ICLR, 2020. find
  20. A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
  21. D. Eigen, M. Ranzato, and I. Sutskever. Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314, 2013. — [arXiv:1312.4314](https://arxiv.org/abs/1312.4314) find
  22. W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021. find
  23. D. M. Gavrila and S. Munder. Multi-cue pedestrian detection and tracking from a moving vehicle. International journal of computer vision, 73(1):41–59, 2007. find
  24. A. Geiger, P. Lenz, and R. Urtasun. Are we ready for autonomous driving? The KITTI vision benchmark suite. In Computer Vision and Pattern Recognition (CVPR), 2012. find
  25. S. Gross, M. Ranzato, and A. Szlam. Hard mixtures of experts for large scale weakly supervised vision. In CVPR, 2017. find
  26. J. V. Hansen. Combining predictors: comparison of five meta machine learning methods. Information Sciences, 119(1-2):91–105, 1999. find
  27. T. Hastie, R. Tibshirani, and J. Friedman. The elements of statistical learning: data mining, inference, and prediction. Springer, 2017. find
  28. P. Helber, B. Bischke, A. Dengel, and D. Borth. EuroSAT: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 12(7):2217–2226, 2019. find
  29. D. Hendrycks and K. Gimpel. arXiv:1606.08415, 2016. Gaussian error linear units (gelus). arXiv preprint find
  30. Y. H. Hu, S. Palreddy, and W. J. Tompkins. A patient-adaptable ECG beat classifier using a mixture of experts approach. IEEE Trans. Biomedical Engineering, 44(9):891–900, 1997. find
  31. R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79–87, 1991. find
  32. Z. Jiang, Q. Hou, L. Yuan, D. Zhou, X. Jin, A. Wang, and J. Feng. Token labeling: Training a 85.5% top-1 accuracy vision transformer with 56m parameters on imagenet. arXiv preprint arXiv:2104.10858, 2021. — [arXiv:2104.10858](https://arxiv.org/abs/2104.10858) find
  33. J. Johnson, B. Hariharan, L. van der Maaten, F.-F. Li, C. Lawrence Zitnick, and R. Girshick. CLEVR: A diagnostic dataset for compositional language and elementary visual reasoning. In Computer Vision and Pattern Recognition (CVPR), 2017.
  34. M. I. Jordan and R. A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural computation, 6(2):181–214, 1994. find
  35. Kaggle and EyePacs. Kaggle diabetic retinopathy detection, 2015. find
  36. A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung, S. Gelly, and N. Houlsby. Big transfer (BiT): General visual representation learning. In ECCV, 2020. find
  37. A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. find
  38. Y. LeCun, F. J. Huang, and L. Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In Computer Vision and Pattern Recognition (CVPR), 2004. find
  39. D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021. find
  40. M. Lewis, S. Bhosale, T. Dettmers, N. Goyal, and L. Zettlemoyer. Base layers: Simplifying training of large, sparse models. arXiv preprint arXiv:2103.16716, 2021.
  41. F.-F. Li, R. Fergus, and P. Perona. Learning generative visual models from few training examples: An incremental Bayesian approach tested on 101 object categories. Computer Vision and Pattern Recognition Workshop, 2004. find
  42. L. Matthey, I. Higgins, D. Hassabis, and A. Lerchner. dSprites: Disentanglement testing sprites dataset, 2017. find
  43. Y. Netzer, T. Wang, A. Coates, A. Bissacco, and A. Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. find
  44. M.-E. Nilsback and A. Zisserman. Automated flower classification over a large number of classes. In Sixth Indian Conf. on Computer Vision, Graphics & Image Processing, 2008. find
  45. O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. Jawahar. Cats and dogs. In CVPR, 2012. find
  46. D. Patterson, J. Gonzalez, Q. Le, C. Liang, L.-M. Munguia, D. Rothchild, D. So, M. Texier, and J. Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2104.10350, 2021.
  47. S. Pavlitskaya, C. Hubschneider, M. Weber, R. Moritz, F. Huger, P. Schlicht, and M. Zollner. Using mixture of expert models to gain insights into semantic segmentation. In CVPR Workshops, 2020. find
  48. M. E. Peters, M. Neumann, M. Iyyer, M. Gardner, C. Clark, K. Lee, and L. Zettlemoyer. Deep contextualized word representations. In NAACL, 2018.
  49. H. Pham, Z. Dai, Q. Xie, M.-T. Luong, and Q. V. Le. Meta pseudo labels. arXiv preprint arXiv:2003.10580, 2020. find
  50. C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683, 2019.
  51. P. Ramachandran and Q. V. Le. Diversity and depth in per-example routing models. In ICLR, 2018. find
  52. C. Rosenbaum, I. Cases, M. Riemer, and T. Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint arXiv:1904.12774, 2019. — [arXiv:1904.12774](https://arxiv.org/abs/1904.12774) find
  53. C. Rosenbaum, T. Klinger, and M. Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. arXiv preprint arXiv:1711.01239, 2017. — [arXiv:1711.01239](https://arxiv.org/abs/1711.01239) find
  54. N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017. find
  55. C. Sminchisescu, A. Kanaujia, Z. Li, and D. Metaxas. Learning to reconstruct 3D human motion from Bayesian mixtures of experts. A probabilistic discriminative approach. Dept. Comput. Sci., Univ. Toronto, Tech. Rep. CSRG-502, 2004. find
  56. E. Strubell, A. Ganesh, and A. McCallum. Energy and policy considerations for deep learning in NLP. arXiv preprint arXiv:1906.02243, 2019.
  57. C. Sun, A. Shrivastava, S. Singh, and A. Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In ICCV, 2017.
  58. J. Tani and S. Nolfi. Learning to perceive the world as articulated: an approach for hierarchical learning in sensory-motor systems. Neural Networks, 12(7-8):1131–1141, 1999. find
  59. H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012.12877, 2020.
  60. H. Touvron, M. Cord, A. Sablayrolles, G. Synnaeve, and H. Jégou. Going deeper with image transformers. arXiv preprint arXiv:2103.17239, 2021. find
  61. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In NeurIPS, 2017.
  62. B. S. Veeling, J. Linmans, J. Winkens, T. Cohen, and M. Welling. Rotation equivariant CNNs for digital pathology. In Medical Image Computing and Computer Assisted Intervention (MICCAI), 2018. find
  63. X. Wang, F. Yu, L. Dunlap, Y.-A. Ma, R. Wang, A. Mirhoseini, T. Darrell, and J. E. Gonzalez. Deep mixture of experts via shallow embedding. In Uncertainty in Artificial Intelligence, 2020. find
  64. B. Yang, G. Bender, Q. V. Le, and J. Ngiam. Condconv: Conditionally parameterized convolutions for efficient inference. arXiv preprint arXiv:1904.04971, 2019. find
  65. L. Yuan, Y. Chen, T. Wang, W. Yu, Y. Shi, F. E. Tay, J. Feng, and S. Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. arXiv preprint arXiv:2101.11986, 2021. find
  66. S. E. Yuksel, J. N. Wilson, and P. D. Gader. Twenty years of mixture of experts. IEEE transactions on neural networks and learning systems, 23(8):1177–1193, 2012. find
  67. A. J. Zeevi, R. Meir, and R. J. Adler. Time series prediction using mixtures of experts. In NeurIPS, 1997. find
  68. X. Zhai, A. Kolesnikov, N. Houlsby, and L. Beyer. Scaling vision transformers, 2021. find
  69. X. Zhai, J. Puigcerver, A. Kolesnikov, P. Ruyssen, C. Riquelme, M. Lucic, J. Djolonga, A. S. Pinto, M. Neumann, A. Dosovitskiy, L. Beyer, O. Bachem, M. Tschannen, M. Michalski, O. Bousquet, S. Gelly, and N. Houlsby. A large-scale study of representation learning with the visual task adaptation bench… — [arXiv:1910.04867](https://arxiv.org/abs/1910.04867) find
  70. X. Zhai, J. Puigcerver, A. Kolesnikov, P. Ruyssen, C. Riquelme, M. Lucic, J. Djolonga, A. S. Pinto, M. Neumann, A. Dosovitskiy, et al. A large-scale study of representation learning with the visual task adaptation benchmark. arXiv preprint arXiv:1910.04867, 2019. *Table 3: Comparison of routing func… — [arXiv:1910.04867](https://arxiv.org/abs/1910.04867) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/riquelme-2021-scaling-vision-sparse-mixture-experts/