Source-linked AI summary
Online normalizer calculation for softmax
Maxim Milakov, Natalia Gimelshein
TL;DR
Softmax is widely used, while prior approaches generally still require the original normalization, motivating a faster implementation. The paper computes Softmax with fewer memory accesses, and benchmarks report up to 1.3x speedup for Softmax and up to 5x for fused Softmax+TopK. These gains are scope-dependent: TopK dominates as K grows, and CPU behavior was not experimentally evaluated.
Problem
Prior approaches generally still require the original Softmax, while targeted efforts to improve its performance were lacking.
Method
The paper computes the Softmax normalizer in a single pass over the input and develops a parallel version for concurrent hardware execution.
Results
Softmax reaches up to 1.3x faster performance, while fused Softmax+TopK reaches up to 5x improvement over separate traditional operations.
Takeaways & Limitations
The memory-access reduction also applies alongside several other Softmax optimization techniques, including Hierarchical Softmax, Importance Sampling, and SVD-Softmax.
Takeaways & Limitations
TopK fusion benefits are reported mainly for relatively small K and decline as K increases because partial TopK costs grow; CPU behavior was not tested.
Abstract
from arXiv · showhide
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.
1 Introduction
Softmax is a standard component for converting neural-network logits into probabilities, and prior acceleration methods generally still rely on computing Softmax. The paper targets this unaddressed performance problem by reducing Softmax memory accesses and benchmarking the result on hardware.
- 1 Introduction: Softmax transforms projection-layer logits into probability vectors for language modeling and other machine-learning applications.It is also used in multinomial logistic regression.
- 1 Introduction: Previous alternatives accelerate projection layers, reorganize Projection+Softmax computations, or approximate training probabilities.Examples include Differentiated Softmax, SVD-Softmax, Hierarchical Softmax, Importance Sampling, Noise Contrastive Estimation, and Blackout.
- 1 Introduction: Most prior approaches still compute the original Softmax, so they can benefit when that function runs faster.Noise Contrastive Estimation during training and Self-Normalized Softmax during inference are identified as exceptions.
- 1 Introduction: The paper addresses the lack of targeted optimization for original Softmax by computing it with fewer memory accesses and testing whether this improves real-hardware performance.The authors explicitly benchmark the relationship between reduced memory accesses and performance.
- 1 Introduction: The work is presented as a preprint and remains in progress.
2 Original softmax
Softmax converts an input vector into normalized output values, but naive and safe implementations require multiple passes and memory accesses. The safe version avoids exponent overflow or underflow by subtracting the input maximum before normalization.
- 2 Original softmax: Softmax is defined as y = Softmax(x), mapping input vector x to output vector y.The supplied definition introduces the function without further expanding its normalization expression.
- 2 Original softmax: The naive implementation scans the input twice, using two loads and one store per vector element.One pass calculates the normalization term, and the second computes output values.
- 2 Original softmax: The safe formulation prevents exponent overflow or underflow caused by the limited numeric range of real hardware.It does so by using exponentials of x_i−m_V rather than unshifted x_i.
- 2 Original softmax: Safe Softmax makes three passes—maximum, normalization, and final values—resulting in 4 memory accesses per vector element.Major deep-learning frameworks use this safe version.
3 Online normalizer calculation
Online Softmax computes the maximum and normalization term in one pass, reducing memory accesses while preserving numerical safety. Its associative formulation also supports parallel evaluation.
- 3 memory accesses per vector element replace Safe Softmax’s 4 by calculating the maximum and normalization term in one pass.The method adds two operations per element and is inspired by online variance calculation.
- Online Softmax maintains the running maximum m and normalizer d, adjusting d when a new maximum appears before adding the new input contribution.
- The algorithm is proved to compute the Softmax function and bounds d_j by 1 ≤ d_j ≤ j.32-bit storage supports processing up to 1.7 ∗ 10^37 elements; larger vectors require 64-bit storage.
- 3.1 Parallel online normalizer calculation: The online normalizer’s binary operation is associative and commutative, enabling parallel evaluation and implementation flexibility.Sequential left-to-right application is equivalent to the algorithm’s single-pass calculation.
4 Softmax and top-k fusion
Because TopK need not materialize every Softmax output, the paper fuses Softmax and TopK while retaining only the largest values and their indices. This reduces memory traffic substantially.
- Online Softmax normally uses three accesses per element, but TopK after Softmax can avoid computing all output values.
- TopK returns K indices referencing the largest input values together with those values.
- 1 memory access per input element enables fused Softmax+TopK, compared with 5 for separate Safe Softmax+TopK and 4 for separate Online Softmax+TopK.
5 Benchmarking
Benchmarks on an NVIDIA Tesla V100 evaluate Naive, Safe, and Online Softmax, plus Safe and Online Softmax fused with TopK, across large and small batches. Online Softmax improves performance by reducing memory accesses, reaching about 1.3x for Softmax and up to 5x for fused Softmax+TopK, while gains decline as K increases.
- Online algorithms: Online Softmax maintains running normalization values and TopK values and indices, enabling fused Softmax+TopK with one memory access per input element.Algorithm 4 updates the normalization state while inserting each input into a running TopK structure.
- Benchmark setup: The CUDA benchmark compares Naive, Safe, and Online Softmax across vector sizes for batches of 4,000 and 10 vectors.The large batch represents a saturated device, while the small batch represents online inference with too few vectors to fully occupy it.
- Softmax benchmarking: ∼1.3x faster at V = 4000: Online and Naive Softmax outperform Safe Softmax in the large-batch, bandwidth-limited regime.The improvement is close to the 1.33x reduction in memory accesses.
- Softmax benchmarking: ∼1.15x faster after V = 1000: Naive and Online Softmax outperform Safe Softmax for the small-batch case.GPU underutilization exposes latency, so the small-batch gain is lower than in the saturated large-batch case.
- Softmax and TopK benchmarking: Up to 5x faster at V = 25000: fused Online Softmax+TopK outperforms unfused Safe Softmax+TopK in the large-batch case.The 5x improvement corresponds to 2.5x from fusion and 2x from Online Softmax.
- Softmax and TopK benchmarking: 1.5x-2.5x faster: fused Online Softmax+TopK outperforms unfused Safe Softmax+TopK for small batches, despite latency-limited GPU utilization.Fusion alone provides substantial gains, and Online Softmax improves performance further.
- Softmax and TopK benchmarking: Performance gains decline as K increases, reaching 3.5x for K = 10, 2x for K = 15, and 1.4x for K = 30.For larger K, TopK increasingly dominates runtime, although fusing normalization into TopK remains beneficial.
6 Results
The proposed single-pass normalizer reduces Softmax memory accesses and yields measured performance gains on Tesla V100, including larger gains when fused with TopK. The improvements are also compatible with several other Softmax optimization techniques.
- 6 Results: 1.33x fewer memory accesses for Softmax alone results from calculating its normalizer in a single pass.The reduction applies to the Softmax function evaluation itself.
- 6 Results: 1.15x performance improvement occurs for V ≥1000 vector sizes, rising to 1.3x for large batches when V ≥4000 on Tesla V100.These benchmark results quantify the hardware benefit of the reduced memory traffic.
- 6 Results: Switching from Naive Softmax to Online Softmax improves numerical accuracy with no performance hit or only a negligible one.This comparison concerns Naive Softmax as the starting implementation.
- 6 Results: 1.5x–5x performance improvement is observed for fused Softmax+TopK on Tesla V100, with the 5x gain combining fusion and Online Softmax contributions.The reported 5x result comprises 2.5x from fusion and 2x from Online Softmax itself.
- 6 Results: These performance improvements are orthogonal to Hierarchical Softmax, Importance Sampling, and SVD-Softmax.The proposed improvements can therefore be applied alongside these listed techniques.
7 Discussion
Online Softmax reaches up to 1.3x speedup, while fusion with following TopK reaches up to 5x over running traditional Safe Softmax and TopK separately. The discussion identifies device dependence and remaining memory-bandwidth limits as important boundaries.
- 7 Discussion: Up to 1.3x speedup is reported on the latest-generation GPU relative to the version used by major deep-learning frameworks.The same approach enables fusion with following TopK, reported separately below.
- 7 Discussion: Up to 5x performance improvement is reported when Softmax and following TopK are fused instead of running traditional Safe Softmax and TopK separately.The fusion is enabled by Online Softmax.
- 7 Discussion: Experiments did not cover CPUs, where speedups could differ or slowdowns could occur; similar gains are only expected if vectorization is preserved.This is an explicit scope boundary and conditional expectation.
- 7 Discussion: The resulting Softmax and fused Softmax+TopK remain limited by memory bandwidth, and preceding-layer fusion could avoid memory round trips but is more challenging.Further improvement is described as possible through additional fusion.