Elamur Research⌕
☰
home / papers / 2018 / Online normalizer calculation for softmax

Online normalizer calculation for softmax

Paper · 2018 · Maxim Milakov and Natalia Gimelshein

Summary

Softmax sits at the output of most language models. Many faster substitutes exist (hierarchical, sampling-based, SVD-based), but nearly all of them still run the exact softmax on some vector. This NVIDIA paper speeds up that exact function. The numerically safe variant used by TensorFlow, PyTorch, MXNET, Microsoft Cognitive Toolkit and Chainer makes three passes over the input: one for the maximum, one for the sum of exponentials, and one for the outputs. That costs four memory accesses per element.

The authors compute the maximum and the normalizer together in a single sweep. Each time the running maximum increases, the partial sum is rescaled by the exponential of the shift. They prove correctness by induction and show the accumulator stays bounded. They also recast the update as an associative, commutative merge, so it maps onto parallel reductions. A further variant fuses softmax with TopK for beam search and reads each input only once.

Because softmax is memory-bound, these savings translate into real speedups on a Tesla V100: up to 1.3x for softmax and up to 5x for fused Softmax+TopK. The method also combines with other softmax approximations.

Key points

  • Memory traffic: "Online Softmax" lowers accesses per element from 4 (safe) to 3, a 1.33x reduction. It keeps the safety of the max-subtraction form, and its normalizer satisfies 1 ≤ d_j ≤ j. In fp32 this covers vectors up to about 1.7·10^37 elements.
  • Softmax benchmarks: CUDA C with CUB v1.8.0, run on a Tesla V100 PCIe 16 GB in fp32.
    • Batch of 4,000 vectors: about 1.3x over safe softmax from V = 4000 onward.
    • Batch of 10 vectors: about 1.15x once V exceeds roughly 1000.
    • Below V ≈ 1000 all variants perform about the same, because caches absorb the extra traffic.
  • Fused Softmax+TopK (K = 5):
    • Large batch: 1.5x rising toward 5x near V = 25,000. The 5x splits into roughly 2.5x from fusion and 2x from the online normalizer.
    • Small, latency-bound batch: 1.5x–2.5x.
  • Effect of larger K: maintaining the partial top-K list gets expensive, so gains shrink to 3.5x (K = 10), 2x (K = 15) and 1.4x (K = 30).
  • Code: released at github.com/NVIDIA/online-softmax.

Interesting ideas

  • I1. Rescale instead of re-reading. Keep a running reference value (here the maximum) and correct the accumulated sum by a multiplicative factor whenever the reference changes. This removes a whole pass over the data. The same pattern could serve any streaming computation that needs a normalization constant, such as log-sum-exp over large vocabularies or scores read from disk.
  • I2. Make the per-chunk state mergeable. Storing (max, sum) pairs with an associative, commutative combine rule lets the computation be split across threads, blocks or machines and joined in any order. Designing reductions this way makes them easy to parallelize or shard.
  • I3. Fuse with a consumer that needs only a few outputs. If the next step (here TopK) uses only a handful of results, compute the normalizer alongside it and normalize just the survivors. The full output vector never gets written. This is useful wherever a probability layer feeds selection or ranking.
  • I4. Predict speedups from memory traffic, then verify. The authors estimated gains from counts of memory accesses, then checked them against measured throughput and profiler data. This is a cheap sanity model for optimizing any memory-bound kernel or data pipeline.
  • I5. Benchmark two regimes. Running both a saturating (large-batch) and a latency-bound (small-batch) workload shows where an optimization pays off and where it doesn't. This is a good evaluation habit for both training and online-inference settings.

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

🏢NvidiaSanta Clara · United States25 works · 5,058 citations🏢Nvidia (United Kingdom)Reading · United Kingdom18 works · 4,771 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksComputer Vision and Pattern Recognition · 2,869 worksComputational Theory and Mathematics · 252 worksHardware and Architecture · 81 works

💡 Concepts 1Key ideas and methods explained on our concept pages.Latent semantic indexing · 29 works

Cite

Cite as Maxim Milakov and Natalia Gimelshein. 2018. Online normalizer calculation for softmax. arXiv 2018 (NVIDIA).

The paper itself, not this page. BibTeX from the library's record of the work. arXiv 1805.02867

@misc{milakov2018online,
  title = {{Online normalizer calculation for softmax}},
  author = {Maxim Milakov and Natalia Gimelshein},
  year = {2018},
  eprint = {1805.02867},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/1805.02867}
}

References

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

  1. I. Sutskever, O. Vinyals, and Q. V. Le. Sequence to Sequence Learning with Neural Networks. ArXiv e-prints, September 2014, 1409.3215.
  2. G. Hinton, L. Deng, D. Yu, G. E. Dahl, A. r. Mohamed, N. Jaitly, A. Senior, V. Vanhoucke, P. Nguyen, T. N. Sainath, and B. Kingsbury. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IEEE Signal Processing Magazine, 29(6):82–97, Nov 2012. IS… find
  3. Engel J. Polytomous logistic regression. Statistica Neerlandica, 42(4):233–252. find
  4. W. Chen, D. Grangier, and M. Auli. Strategies for Training Large Vocabulary Neural Language Models. ArXiv e-prints, December 2015, 1512.04906. find
  5. Kyuhong Shim, Minjae Lee, Iksoo Choi, Yoonho Boo, and Wonyong Sung. Svd-softmax: Fast softmax approximation on large vocabulary neural networks. In Advances in Neural Information Processing Systems 30, pages 5463–5473. Curran Associates, Inc., 2017. URL http://papers.nips.cc/paper/7130-svd-softmax-f… ↗
  6. Joshua Goodman. Classes for fast maximum entropy training. ICASSP, pages 561–564. IEEE, 2001. ISBN 0-7803-7041-4. http://dblp.uni-trier.de/db/conf/icassp/icassp2001.html\#Goodman01. ↗
  7. T. Mikolov, S. Kombrink, L. Burget, J. Černocký, and S. Khudanpur. Extensions of recurrent neural network language model. In 2011 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5528–5531, May 2011. In URL find
  8. E. Grave, A. Joulin, M. Cissé, D. Grangier, and H. Jégou. Efficient softmax approximation for GPUs. ArXiv e-prints, September 2016, 1609.04309.
  9. Yoshua Bengio and Jean-Sébastien Sénécal. Quick training of probabilistic neural nets by importance sampling. In Proceedings of the conference on Artificial Intelligence and Statistics (AISTATS), 2003. find
  10. A. Mnih and Y. Whye Teh. A Fast and Simple Algorithm for Training Neural Probabilistic Language Models. ArXiv e-prints, June 2012, 1206.6426. find
  11. S. Ji, S. V. N. Vishwanathan, N. Satish, M. J. Anderson, and P. Dubey. BlackOut: Speeding up Recurrent Neural Network Language Models With Very Large Vocabularies. ArXiv e-prints, November 2015, 1511.06909. find
  12. Jacob Devlin, Rabih Zbib, Zhongqiang Huang, Thomas Lamar, Richard Schwartz, and John Makhoul. Fast and robust neural network joint models for statistical machine translation. In Proceedings of ACL2014, pages 1370–1380, 2014.
  13. Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenber… find
  14. Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W, 2017. find
  15. Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems. In Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Twenty-…
  16. Frank Seide and Amit Agarwal. Cntk: Microsoft’s open-source deep-learning toolkit. In Proceedings of the 22Nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD ’16, pages 2135–2135, New York, NY, USA, 2016. ACM. ISBN 978-1-4503-4232-2. URL http://doi.acm.org/10.1145/293… — [doi:10.1145/2939672.2945397](https://doi.org/10.1145/2939672.2945397) ↗
  17. Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: a next-generation open source framework for deep learning. In Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Twenty-ninth Annual Conference on Neural Information Processing Systems (NIPS), 2015. URL http://l… find
  18. B. P. Welford. Note on a method for calculating corrected sums of squares and products. Technometrics, 4(3):419–420, 1962. URL https://amstat.tandfonline.com/doi/abs/10.1080/00401706.1962.10490022. — [doi:10.1080/00401706.1962.10490022](https://doi.org/10.1080/00401706.1962.10490022) ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/milakov-2018-online-normalizer-calculation-softmax/