---
title: "Online normalizer calculation for softmax — Milakov 2018"
description: "Softmax sits at the output of most language models. 18 references, cited by 103 papers here."
authors: ["Maxim Milakov", "Natalia Gimelshein"]
year: 2018
venue: "arXiv 2018 (NVIDIA)"
source: "https://arxiv.org/abs/1805.02867"
arxiv: "1805.02867"
licence: "arXiv — non-exclusive distribution license (arXiv OAI)"
licence_url: "http://arxiv.org/licenses/nonexclusive-distrib/1.0/"
summary_licence: "Elamur Cosmic License"
canonical: "https://research.elamur.ai/free/en/papers/milakov-2018-online-normalizer-calculation-softmax/"
---
[Elamur Research](https://md.research.elamur.ai/free/en/) · [this page in HTML](https://research.elamur.ai/free/en/papers/milakov-2018-online-normalizer-calculation-softmax/)

[home](https://md.research.elamur.ai/free/en/) / [papers](https://md.research.elamur.ai/free/en/papers/) / [2018](https://md.research.elamur.ai/free/en/papers/2018/page/4/) / Online normalizer calculation for softmax

# Online normalizer calculation for softmax

**Authors:** Maxim Milakov and [Natalia Gimelshein](https://md.research.elamur.ai/free/en/authors/natalia-gimelshein/) · **Year:** 2018 · **Venue:** arXiv 2018 (NVIDIA) · **Licence:** [arXiv — non-exclusive distribution license (arXiv OAI)](http://arxiv.org/licenses/nonexclusive-distrib/1.0/)

**Links:** [Publisher's page](https://arxiv.org/abs/1805.02867) · [PDF](https://arxiv.org/pdf/1805.02867) · [arXiv 1805.02867](https://arxiv.org/abs/1805.02867) · [OpenAlex](https://openalex.org/W2801655600)

**Citations:** 20 citations (OpenAlex, 2026-10-14) · [103](https://md.research.elamur.ai/free/en/papers/milakov-2018-online-normalizer-calculation-softmax/cited-by/) in our library, of the 149,439 papers whose bibliographies we read

## 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
- **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.
- **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.
- **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.
- **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.
- **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

- [Nvidia](https://md.research.elamur.ai/free/en/companies/nvidia/) — Company, Santa Clara, United States · 25 works · 5,058 citations
- [Nvidia (United Kingdom)](https://md.research.elamur.ai/free/en/companies/nvidia-united-kingdom/) — Company, Reading, United Kingdom · 18 works · 4,771 citations

📚 Shelves (4): [🖥️ Computer Science](https://md.research.elamur.ai/free/en/shelves/computer-science/) · 16,667 works · [👁️ Computer Vision and Pattern Recognition](https://md.research.elamur.ai/free/en/shelves/computer-science/computer-vision-and-pattern-recognition/) · 2,869 works · [🧮 Computational Theory and Mathematics](https://md.research.elamur.ai/free/en/shelves/computer-science/computational-theory-and-mathematics/) · 252 works · [🔌 Hardware and Architecture](https://md.research.elamur.ai/free/en/shelves/computer-science/hardware-and-architecture/) · 81 works

💡 Concepts (1): [💡 Latent semantic indexing](https://md.research.elamur.ai/free/en/concepts/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.

```bibtex
@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}
}
```

- [← previous A Large-Scale Test Set for the Evaluation of Context-Aware Pronoun Translation in Neural Machine Translation](https://md.research.elamur.ai/free/en/papers/muller-2018-large-scale-test-set-evaluation/)
- [next → Automatic Detection of Fake News](https://md.research.elamur.ai/free/en/papers/rosas-2018-automatic-detection-fake-news/)
- [↑ back to papers of 2018](https://md.research.elamur.ai/free/en/papers/2018/page/4/)

## 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](https://md.research.elamur.ai/free/en/papers/sutskever-2014-sequence-sequence-learning-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](https://scholar.google.com/scholar?q=G%2E+Hinton%2C+L%2E+Deng%2C+D%2E+Yu%2C+G%2E+E%2E+Dahl%2C+A%2E+r%2E+Mohamed%2C+N%2E+Jaitly%2C+A%2E+Senior%2C+V%2E+Vanhoucke%2C+P%2E+Nguyen%2C+T%2E+N%2E+Sainath%2C+and+B%2E+Kingsbury%2E+Deep+neural+networks+for+acoustic+modeling+in+speech+recognition%3A+)
3. Engel J. Polytomous logistic regression. Statistica Neerlandica, 42(4):233–252. [find](https://scholar.google.com/scholar?q=Engel+J%2E+Polytomous+logistic+regression%2E+Statistica+Neerlandica%2C+42%284%29%3A233%E2%80%93252%2E)
4. W. Chen, D. Grangier, and M. Auli. Strategies for Training Large Vocabulary Neural Language Models. ArXiv e-prints, December 2015, 1512.04906. [find](https://scholar.google.com/scholar?q=W%2E+Chen%2C+D%2E+Grangier%2C+and+M%2E+Auli%2E+Strategies+for+Training+Large+Vocabulary+Neural+Language+Models%2E+ArXiv+e%2Dprints%2C+December+2015%2C+1512%2E04906%2E)
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…](http://papers.nips.cc/paper/7130-svd-softmax-f…) [↗](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](http://dblp.uni-trier.de/db/conf/icassp/icassp2001.html\#Goodman01). [↗](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](https://scholar.google.com/scholar?q=T%2E+Mikolov%2C+S%2E+Kombrink%2C+L%2E+Burget%2C+J%2E+C%CC%8Cernock%C3%BD%2C+and+S%2E+Khudanpur%2E+Extensions+of+recurrent+neural+network+language+model%2E+In+2011+IEEE+International+Conference+on+Acoustics%2C+Speech+and+Signal+Process)
8. E. Grave, A. Joulin, M. Cissé, D. Grangier, and H. Jégou. [Efficient softmax approximation for GPUs](https://md.research.elamur.ai/free/en/papers/grave-2016-efficient-softmax-approximation-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](https://scholar.google.com/scholar?q=Yoshua+Bengio+and+Jean%2DS%C3%A9bastien+S%C3%A9n%C3%A9cal%2E+Quick+training+of+probabilistic+neural+nets+by+importance+sampling%2E+In+Proceedings+of+the+conference+on+Artificial+Intelligence+and+Statistics+%28AISTATS%29%2C+2003)
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](https://scholar.google.com/scholar?q=A%2E+Mnih+and+Y%2E+Whye+Teh%2E+A+Fast+and+Simple+Algorithm+for+Training+Neural+Probabilistic+Language+Models%2E+ArXiv+e%2Dprints%2C+June+2012%2C+1206%2E6426%2E)
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](https://scholar.google.com/scholar?q=S%2E+Ji%2C+S%2E+V%2E+N%2E+Vishwanathan%2C+N%2E+Satish%2C+M%2E+J%2E+Anderson%2C+and+P%2E+Dubey%2E+BlackOut%3A+Speeding+up+Recurrent+Neural+Network+Language+Models+With+Very+Large+Vocabularies%2E+ArXiv+e%2Dprints%2C+November+2015%2C+1511%2E)
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](https://md.research.elamur.ai/free/en/papers/devlin-2014-fast-robust-neural-network-joint/). 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](https://scholar.google.com/scholar?q=Mart%C3%ADn+Abadi%2C+Ashish+Agarwal%2C+Paul+Barham%2C+Eugene+Brevdo%2C+Zhifeng+Chen%2C+Craig+Citro%2C+Greg+S%2E+Corrado%2C+Andy+Davis%2C+Jeffrey+Dean%2C+Matthieu+Devin%2C+Sanjay+Ghemawat%2C+Ian+Goodfellow%2C+Andrew+Harp%2C+Geoffrey+I)
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](https://scholar.google.com/scholar?q=Adam+Paszke%2C+Sam+Gross%2C+Soumith+Chintala%2C+Gregory+Chanan%2C+Edward+Yang%2C+Zachary+DeVito%2C+Zeming+Lin%2C+Alban+Desmaison%2C+Luca+Antiga%2C+and+Adam+Lerer%2E+Automatic+differentiation+in+pytorch%2E+In+NIPS%2DW%2C+2017%2E)
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](https://md.research.elamur.ai/free/en/papers/chen-2015-mxnet-flexible-efficient-machine-learning/). 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…](http://doi.acm.org/10.1145/293…) — [doi:10.1145/2939672.2945397](https://doi.org/10.1145/2939672.2945397) [↗](http://doi.acm.org/10.1145/293…)
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](https://scholar.google.com/scholar?q=Seiya+Tokui%2C+Kenta+Oono%2C+Shohei+Hido%2C+and+Justin+Clayton%2E+Chainer%3A+a+next%2Dgeneration+open+source+framework+for+deep+learning%2E+In+Proceedings+of+Workshop+on+Machine+Learning+Systems+%28LearningSys%29+in+Th)
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](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) [↗](https://amstat.tandfonline.com/doi/abs/10.1080/00401706.1962.10490022)

*Updated 29 October 2026*