Online normalizer calculation for softmax

By Maxim Milakov, Natalia Gimelshein. 2018

201 citations recorded by the source corpus.

The Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothesize that this reduction in memory accesses should improve Softmax performance on actual hardware. The benchmarks confirm this hypothesis: Softmax accelerates by up to 1.3x and Softmax+TopK combined and fused by up to 5x.

Subfields: Artificial Intelligence, Computational Theory and Mathematics, Hardware and Architecture, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Theory and Mathematics

In Hardware and Architecture

In Optimization & Training

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “Conceptually, both F LASH A TTENTION and Rabe and Staats [69] operate on blocks of the attention matrix using the well-established technique of tiling (or softmax scaling) [53, 62].”
Ring Attention with Blockwise Transformers for Near-Infinite Context (2023, methodology): “5 [33] with context length 16K, GPT-4 [30] with context length 32k, MosaicML’s MPT [26] with context length 65k, and Anthropic’s Claude [1] with context length 100k. Preprint.”
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023, methodology): “As the softmax couples entire rows or blocks of row, online softmax [11, 13] can split the attention computation into blocks, and rescale the output of each block to ﬁnally get the right result (with no approximation).”