Elamur Research⌕
☰
home / papers / 2021 / BASE Layers: Simplifying Training of Large, Sparse Models

BASE Layers: Simplifying Training of Large, Sparse Models

Paper · 2021 · Mike Lewis, Shruti Bhosale, Tim Dettmers et al.

Summary

Sparse expert models increase capacity cheaply by sending each token to one of many expert sub-networks. In practice they are hard to train because routing tends to become uneven. Earlier systems such as Sparsely Gated MoE and Switch counter this with auxiliary balancing losses, extra hyperparameters and capacity limits that drop tokens. Lewis et al. propose BASE layers, which treat routing during training as a linear assignment problem. Token–expert affinity scores are maximised under the constraint that every expert gets exactly the same number of tokens. Each expert's output is added to the token through a residual connection scaled by a sigmoid gate, so experts can specialise without extra loss terms. At inference, each token simply goes to its best-scoring expert. In language-model experiments with up to 110B parameters, BASE matches or beats dense data/model-parallel training and earlier sparse methods at equal compute, and it trains faster because workers communicate less. Its value is that it works as a drop-in conditional-compute layer with no new hyperparameters. Code is in fairseq.

Key points

  • The balanced assignment is computed with the auction algorithm (Bertsekas, 1992), chosen because it parallelises on GPUs more easily than the Hungarian algorithm. The problem is sharded so that each worker solves it for its own batch, and tokens are randomly shuffled across workers beforehand.

  • Training throughput on 128 GPUs, in tokens per second:

    ModelTokens/s
    Data Parallel600k
    BASE545k
    BASE ×2475k
    Switch469k
    Sparsely Gated MoE292k
    Model Parallel ×2224k

  • A 43B-parameter BASE model trains at roughly 90% of the speed of a 1.5B data-parallel baseline. About 1.5B is the largest model that fits in pure data parallelism in the authors' setup.

  • A single large BASE layer is already effective, and adding more BASE layers improves results further.

  • Experts tend to specialise on shallow, local cues such as the preceding token. Examples include experts for numbers, auxiliary verbs, possessives and specific noun or adjective groups.

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

🎓University of WashingtonSeattle · United States402 works · 212,762 citations🏢MetaMenlo Park · United States214 works · 241,492 citations🏢Meta (United Kingdom)London · United Kingdom54 works · 19,348 citations

 Fields, topics and concepts

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

Cite

Cite as Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal and Luke Zettlemoyer. 2021. BASE Layers: Simplifying Training of Large, Sparse Models. arXiv:2103.16716.

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

@misc{lewis2021base,
  title = {{BASE Layers: Simplifying Training of Large, Sparse Models}},
  author = {Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal and Luke Zettlemoyer},
  year = {2021},
  doi = {10.48550/arxiv.2103.16716},
  eprint = {2103.16716},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2103.16716}
}

References

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

  1. Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
  2. Bapna, A., Arivazhagan, N., and Firat, O. Simple, scalable adaptation for neural machine translation. arXiv preprint arXiv:1909.08478, 2019.
  3. Bertsekas, D. P. Auction algorithms for network flow problems: A tutorial introduction. Computational optimization and applications, 1(1):7–66, 1992. find
  4. Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
  5. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
  6. Lample, G., Sablayrolles, A., Ranzato, M., Denoyer, L., and Jégou, H. Large memory layers with product keys. arXiv preprint arXiv:1907.05242, 2019. find
  7. Conneau, A., Khandelwal, K., Goyal, N., Chaudhary, V., Wenzek, G., Guzmán, F., Grave, E., Ott, M., Zettlemoyer, L., and Stoyanov, V. Unsupervised cross-lingual representation learning at scale. arXiv preprint arXiv:1911.02116, 2019.
  8. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. find
  9. Correia, G. M., Niculae, V., and Martins, A. F. Adaptively sparse transformers. arXiv preprint arXiv:1909.00015, 2019.
  10. Dettmers, T. and Zettlemoyer, L. Sparse networks from scratch: Faster training without losing performance. arXiv preprint arXiv:1907.04840, 2019. find
  11. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  12. Evci, U., Gale, T., Menick, J., Castro, P. S., and Elsen, E. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pp. 2943– 2952. PMLR, 2020. find
  13. Fan, A., Bhosale, S., Schwenk, H., Ma, Z., El-Kishky, A., Goyal, S., Baines, M., Celebi, O., Wenzek, G., Chaudhary, V., et al. Beyond english-centric multilingual machine translation. arXiv preprint arXiv:2010.11125, 2020. find
  14. Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O., Stoyanov, V., and Zettlemoyer, L. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. arXiv preprint arXiv:1910.13461, 2019.
  15. Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019.
  16. Mostafa, H. and Wang, X. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In International Conference on Machine Learning, pp. 4646–4655. PMLR, 2019. find
  17. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  18. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683, 2019.
  19. Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021. find
  20. Roy, A., Saffar, M., Vaswani, A., and Grangier, D. Efficient content-based sparse attention with routing transformers. arXiv preprint arXiv:2003.05997, 2020.
  21. Houlsby, N., Giurgiu, A., Jastrzebski, S., Morrone, B., De Laroussilhe, Q., Gesmundo, A., Attariyan, M., and Gelly, S. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning, pp. 2790–2799. PMLR, 2019. find
  22. Sennrich, R., Haddow, B., and Birch, A. Neural machine translation of rare words with subword units. arXiv preprint arXiv:1508.07909, 2015.
  23. Khandelwal, U., Levy, O., Jurafsky, D., Zettlemoyer, L., and Lewis, M. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00172, 2019. find
  24. Khandelwal, U., Fan, A., Jurafsky, D., Zettlemoyer, L., and Lewis, M. Nearest neighbor machine translation. arXiv preprint arXiv:2010.00710, 2020. find
  25. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. find
  26. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  27. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  28. Strubell, E., Ganesh, A., and McCallum, A. Energy and policy considerations for deep learning in nlp. arXiv preprint arXiv:1906.02243, 2019.
  29. Kuhn, H. W. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83– 97, 1955. find
  30. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. arXiv preprint arXiv:1706.03762, 2017.
  31. Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524– 10533. PMLR, 2020. def balanced_assignment(scores, max_iterations=100): num_worke… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/lewis-2021-base-layers-simplifying-training-large/