The Two-Pass Softmax Algorithm

01/13/2020 ∙ by Marat Dukhan, et al. ∙ Google 0

The softmax (also called softargmax) function is widely used in machine learning models to normalize real-valued scores into a probability distribution. To avoid floating-point overflow, the softmax function is conventionally implemented in three passes: the first pass to compute the normalization constant, and two other passes to compute outputs from normalized inputs. We analyze two variants of the Three-Pass algorithm and demonstrate that in a well-optimized implementation on HPC-class processors performance of all three passes is limited by memory bandwidth. We then present a novel algorithm for softmax computation in just two passes. The proposed Two-Pass algorithm avoids both numerical overflow and the extra normalization pass by employing an exotic representation for intermediate values, where each value is represented as a pair of floating-point numbers: one representing the "mantissa" and another representing the "exponent". Performance evaluation demonstrates that on out-of-cache inputs on an Intel Skylake-X processor the new Two-Pass algorithm outperforms the traditional Three-Pass algorithm by up to 28 proposed Two-Pass algorithm also outperforms the traditional Three-Pass algorithm on Intel Broadwell and AMD Zen 2 processors. To foster reproducibility, we released an open-source implementation of the new Two-Pass Softmax algorithm and other experiments in this paper as a part of XNNPACK library at



There are no comments yet.


page 6

page 8

This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.

1 Introduction

The softmax (also called softargmax) function is a smooth approximation to the argmax function. The softmax function

operates on a vector of real-valued scores

and normalizes them into a probability distribution

where and .

The softmax function is commonly used in machine learning to give a probabilistic interpretation to outputs of classification models [bridle1990probabilistic]

. Such machine learning models output a probability for every possible object class, and the number of classes in modern datasets can reach millions. For example, natural language processing models may predict probability distribution over each possible word in a vocabulary, and recommendation systems may model the probability distribution over users, products, web pages, or their interactions. Table 

1 summarized number of classes on several public classification datasets.

Dataset Class description Class Count
ImageNet [deng2009imagenet] Image category 21841
One Billion Word [chelba2013one] Unique Words 793471
Wikilinks [singh2012wikilinks] Wikipedia pages 2933659
DepCC [depcc] Web documents 364.80 million
Table 1: Characteristics of several public machine learning datasets.

Hierarchical Softmax [goodman2001classes] (HSM) and its modifications [EfficientSoftmaxApproximationGPUs] are the common techniques to scale classification models to large number of classes. HSM models jointly consider the softmax function and the matrix-matrix multiplication that produced its input, and replace them by a low-rank approximation. Thus, HSM methods improve performance by reducing the matrix-matrix multiplication cost, and do so by approximating the original machine learning model.

Unlike previous research, we focus on the softmax function in the context of inference using a pre-trained model. This situation commonly arises in machine learning frameworks, such as TensorFlow 


or PyTorch 

[PyTorch], when the training dataset or metaparameters needed to devise an accurate approximation to the model are not available. In this case, the model must be computed exactly according to the specification and unsafe approximations are not possible.

Our contributions in this paper are as follows:

  • We demonstrate that a well-optimized implementation of the softmax function can be memory-bound even in single-threaded execution. This result emphasizes the importance of eliminating memory operations for further improvements in the performance of the softmax function.

  • We introduce a novel algorithm for computing the softmax function. The new algorithm employs an exotic representation for intermediate values, where each value is represented as a pair of floating-point numbers: one representing the “mantissa” and another representing the “exponent”. Thanks to the special representation of intermediate results, the new algorithm needs only two passes over an input vector versus three passes for traditional algorithms.

  • We present and evaluate high-performance implementations of the new Two-Pass softmax algorithms for the x86-64 processors with AVX2 and AVX512F SIMD extensions. The experimental study confirms the speedups of up to on an Intel Skylake-X system.

The proposed improvements to the softmax implementation are orthogonal to matrix-matrix multiplication optimizations, and can be combined with sparsification [SCNN, HolisticSparseCNN], low-rank decomposition [Denton2014], low-precision arithmetic [GEMMLOWP, FBGEMM, Tulloch2017], or hardware acceleration [EIE, TPU] for the matrix-matrix multiplication that produce softmax input.

2 Related work

The softmax function is widely used in the modern machine learning models and algorithms. In particular, the softmax function gained wide popularity in Natural Language Processing as a building block for language models, which predict a word or n-gram out of a large vocabulary 

[LMLimits]. As the softmax function plays an important role in the machine learning models the several approximations have been proposed [HardwareSoftmax, HighSpeedLowComplexitySoftmax, SoftmaxVLSI, SVDSoftmax, SigSoftmax, EfficientHardwareSoftmax, EfficientSoftmaxApproximationGPUs] in order to make the calculations efficient and fast. We refer the reader to Jozefowicz et al. [EfficientSoftmaxApproximationGPUs] for a good overview of the recent approaches to speed up the softmax calculations.

Although a lot work has been done to address the computation complexity of the softmax function on a large inputs, the proposed solutions either 1) require special hardware (e.g FPGA) or 2) require to use an approximation to the softmax function during model training. In the latter case the softmax approximation can not be directly applied to a pre-trained machine learning model, which dramatically limits its use in existing machine learning frameworks e.g. TensorFlow [TensorFlow] and PyTorch[PyTorch]. Therefore, we developed a novel algorithm that improves performance of the original softmax function on widely available hardware, can be used with pre-trained machine learning models, and can be implemented in any machine learning framework.

3 The Three-Pass Algorithm

Direct calculation of the softmax function according to the formula is conjugate with numerical issues. Single-precision function overflows111Produce floating-point infinity because result is too large to be represented as a finite single-precision floating-point number for and underflows222Produce 0 because result is too small to be distinguishable from zero in the single-precision floating-point format for , and, in turn, cause NaN333Not a Number: a special floating-point value defined by IEEE 754 standard indicating invalid result outputs in the naïve implementations of the softmax function. In practice, the parts of the machine learning models that produce input to the softmax function are rarely bounded, and thus implementation can’t assume that the input would fall into such narrow range.

In order to overcome numerical instability issues machine learning frameworks adapt a workaround by utilizing the equivalence [goodfellow2016deep]:

which holds for any value. In particular, if , then:

  • No inputs to function exceed zero

  • There is at least one zero input to the function, and thus the denominator of the fraction is non-zero.

Together, these properties result in good numerical behavior of the computation and lead to Algorithm 1.

function SoftmaxThreePassRecompute(, )
      Pass 1: read X
      Pass 2: read X
     for all  do
          Pass 3: read X, write Y
     end for
end function
Algorithm 1 The Three-Pass algorithm with re-computation of exponential function

Both of the Pass 2 and the Pass 3 in Algorithm 1 compute with the same and values. This observation hints a potential optimization: if computing function is expensive, we could save the computed values to avoid recomputing them in the Pass 3. Such modification is presented in Algorithm 2.

function SoftmaxThreePassReload(, )
      Pass 1: read X
     for all  do
          Pass 2: read X, write Y
     end for
     for all  do
          Pass 3: read Y, write Y
     end for
end function
Algorithm 2 The Three-Pass algorithm with re-loading of exponential computations

Algorithm 2 computes values only once, but this reduction in the number of computations comes at a cost: the second pass of Algorithm 2 does both a read and a write for each element, unlike Algorithm 1 where the second pass does only reads.

4 The Two-Pass Algorithm

The Three-Pass Algorithms 1 and 2 avoid numerical issues by normalizing inputs relative to their maximum value, but then require an additional memory pass to find the maximum value. In this section we suggest that it is possible to get the numerical stability without the extra memory pass, and present a practical Two-Pass algorithm for softmax computation.

The immediate reason for the numerical instability of a naïve softmax implementation is the saturation of function for inputs outside of the narrow range of . Therefore, we have to look inside the function for a solution.

The function can be implemented in infinite number of ways, but practical implementations [AccurateTables, XScaleFP, IA64ElementaryFunctions, dukhan2013methods]

in IEEE floating-point arithmetic follow the traditional structure of elementary function implementation 

[muller2010handbook], and include three steps:

  1. Range reduction, where the problem of approximating on the infinite input domain is reduced to a problem of approximating on a small finite range. For , a natural range reduction derives from the equivalence

    which decompose approximating on into approximating on and multiplying the result by where is an integer.

  2. Approximation of the function on the reduced range, i.e. on for . This step is achieved through a polynomial or rational approximation, often in combination with table look-ups.

  3. Reconstruction of the final value from the approximation on the reduced range. For , the reconstruction step consists of multiplication of by , and can be achieved at low cost by manipulating the exponent field of a binary floating-point number . It is this step where the underflow and overflow situations arise: and thus always fits into a single-precision floating-point number, but can exceed the range of the exponent field, causing underflow or overflow.

The key idea that enables the Two-Pass Softmax algorithm is to remove the reconstruction step, and instead keep the result of as a pair of floating-point values , where . Mathematically, , but in general this expression can not be evaluated in floating-point arithmetic without overflowing or underflowing the result. The representation of as a floating-point number is important: although is by design always an integer, it can have a very large magnitude, and fall outside of the range of standard integer formats. Therefore, the result must be represented as two, rather than one, floating-point numbers. Using multiple floating-point numbers to represent the real-valued result has a long history in double-double, triple-double, quad-double [hida2001algorithms] representations and Error-Free Transformations [graillat2007error]. However, these representations use multiple floating-point numbers to improve precision of floating-point arithmetic, whereas we suggest to use two floating-point numbers to extend its dynamic range.

function SoftmaxTwoPass(, )
     for all  do
          Pass 1: read X
     end for
     for all  do
          Pass 2: read X, write Y
     end for
end function
Algorithm 3 The Two-Pass softmax algorithm. ExtExp denotes an exponential function that produce a pair of floating-point values.

Algorithm 3 presents the softmax computation in just two passes by implementing addition for representation. The reduction pass keeps track of the running maximum value among all elements, and accumulates the scaled values to the running sum. It avoids the floating-point overflow by scaling values by the difference between the corresponding values and maximum of values. As this difference is never positive, values are never scaled up, which ensures the absence of the floating-point overflow.

5 Theoretical Analysis

Algorithm Memory reads Memory writes Bandwidth cost
Three-Pass (Recompute) 1 3N 1N 4N
Three-Pass (Reload) 2 3N 2N 5N
Two-Pass 3 2N 1N 3N
Table 2: Theoretical analysis of memory complexity and bandwidth costs of the three softmax algorithms.

While the number of memory passes in the presented softmax algorithms is evident from the names, the number of actual memory operations is more nuanced. Every pass of the Three-Pass algorithm with Recomputing reads the input array, while the last pass also writes the output array. The Three-Pass algorithm with Reloading just reads the input array in the first pass, reads the input array and writes the output array in the second pass, and reads-modifies-writes the output array in the last pass. The Two-Pass algorithm reads the input array in both passes, and also writes the output array in the second pass. Thus, the memory bandwidth requirements of the Two-Pass algorithm are similar to just the last two passes of the Three-Pass algorithm with Recomputing.

Table 2 summarize the number of memory reads, memory writes, and the memory bandwidth cost for the three algorithms on arrays of elements. Per Table 2, the Two-Pass algorithm has a memory bandwidth advantage of over the Three-Pass algorithm with Recomputing and over the Three-Pass algorithm with Reloading. In practice, we should treat these numbers as upper bounds, because higher computational complexity of the Two-Pass algorithm cuts into gains from bandwidth reduction.

6 Experimental Evaluation

6.1 Platform

We evaluate the performance of the three softmax algorithms on the Intel Xeon W-2135 processor based on Skylake-X microarchitecture and with the characteristics listed in Table 3. To improve performance stability we disabled dynamic frequency scaling in the processor for the duration of our experiments.

Characteristic Value
Microarchitecture Skylake-X
Number of cores 6
Number of hyperthreads 12
Base frequency 3.7 GHz
L1 cache size (per core) 32 KB
L2 cache size (per core) 1 MB
L3 cache size (shared by all cores) 8.25 MB
AVX2 / AVX512 FMA throughput 2 / cycle
AVX2 / AVX512 FMA latency 4 cycles
Table 3: Characteristics of the Intel Xeon W-2135 processor used for experimental evaluation of the softmax algorithms.

Additionally, in Sec. 6.8 we replicate a subset of experiments on a Broadwell-based Intel Xeon E5-2696 v4 processor and on an AMD Ryzen 9 3900X processor with Zen 2 microarchitecture.

6.2 Protocol

We use Google Benchmark framework to estimate sustained performance of the softmax implementations. We set the minimum run-time for each measurement to 5 seconds, repeat each measurement 25 times and record the median of the 25 runs. In each benchmark run we simulate the cache state during neural network inference: output vector is evicted from the cache before each iteration, but input tensor stays in cache as long as it fits.

6.3 Implementation

We developed highly optimized implementations of the Three-Pass Algorithms 1 and 2, and the Two-Pass Algorithm 3 in C. For all three algorithms, we did two templated implementations targeting AVX2 and AVX512 instruction sets. We expressed high-level optimization parameters, such as unroll factor for the loops and the number of accumulator variables in reduction functions, as meta-parameters of the templated implementations, and employed auto-tuning to discover their optimal values.

An efficient implementation of vectorized function is a key component of all softmax variants. For our implementation, we adapted the throughput-optimized methods of Dukhan and Vuduc [dukhan2013methods] to single-precision floating-point evaluation. Algorithm 4 presents the resulting table-free, branch-free, and division-free algorithm.

function Exp()
      Cody-Waite range reduction
      Polynomial approximation
end function
Algorithm 4 Calculation of in the Three-pass softmax algorithms

Algorithm 4 follows the traditional structure of elementary function implementation [muller2010handbook], described in Sec. 4. It starts with a range reduction to reduce approximation on the infinite input domain to approximation on a small and finite range. The calculation of the reduced argument in Algorithm 4 uses Cody-Waite range reduction [cody1980software]: is represented as a sum of two single-precision constants, and , to improve the accuracy of this step. Range reduction results in a reduced argument in the range and a reduced integer argument . Next, is approximated on with a degree-5 polynomial. The polynomial coefficients are produced by the algorithm of Brisebarre and Chevillard [brisebarre2007efficient] as implemented in Sollya software [chevillard2010sollya]. Following [dukhan2013methods], we evaluate the approximation polynomial with Horner scheme using Fused Multiply-Add instructions to minimize the number of floating-point instructions and maximize the throughput. In the last stage the Algorithm 4 reconstructs the final output value of the function by multiplying polynomial approximation result by . In AVX2 implementation, we do this multiplication by directly manipulating floating-point exponent to construct a scale number :

and compute the final step as . This reconstruction trick has two buit-in assumptions: the argument to the is negative444Always the case for the evaluation in the Three-Pass softmax algorithms., and subnormal floating-point numbers can be flushed to zero without significant accuracy impact. The reconstruction step in the AVX512 implementation leverage the new VSCALEFPS instruction [cornea2015intel] which computes as a single hardware operation.

The resulting implementation has maximum error under 2 ULP, validated through exhaustive evaluation on all valid inputs. This accuracy is comparable to other vectorized elementary function implementations, e.g. SIMD functions in GNU LibM library guarantee maximum error under 4 ULP.

Implementation of the ExtExp in the Two-Pass softmax algorithm is similar to Algorithm 4 with the reconstruction step removed. Thus, implementations of both the Three-Pass and the Two-Pass algorithms use exactly the same range reduction and approximating polynomials to compute the exponential function.

6.4 The Three-Pass Algorithms and Bandwidth Saturation

Figure 1: Performance comparison of the Softmax algorithms 1 and 2 in the AVX512 implementations on the Skylake-X system. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

Fig. 1 presents the performance of the Three-Pass softmax Algorithm 1 with recomputing of exponentials and the Three-Pass softmax Algorithm 2 with reloading of computed exponentials in the AVX512 implementations. Reloading of exponential computations delivers speedup when the data is small enough to fit into private L1 and L2 caches, but turns into slowdown when operating on L3 cache, and eventually levels off at advantage when working set exceeds last-level cache.

Figure 2: Performance comparison of the Softmax algorithms 1 and 2 in the AVX2 implementations on the Skylake-X system. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

The AVX2 implementation of the same Three-Pass softmax Algorithm 1 and Algorithm 2 is illustrated in Fig. 2 and demonstrates similar trends. As the working set increases, the speedup from reloading of exponential computations goes down, and eventually levels off at for large arrays.

The small difference between recomputing and reloading of exponential computations on Fig. 1 and Fig. 2 suggests that despite the expensive exponential function, softmax computation might be memory-bound for large arrays. To directly test this hypothesis, we decompose the Algorithms 1 and 2 into individual memory passes, and compare measured bandwidth to STREAM benchmarks [STREAM].

Figure 3: Measured memory bandwidth on the Skylake-X system in the three passes of the Softmax algorithms 1 and 2, and in the STREAM benchmark. Both the softmax implementations and the STREAM benchmark use AVX512 instructions.
Figure 4: Measured memory bandwidth on the Skylake-X system in the three passes of the Softmax algorithms 1 and 2, and in the STREAM benchmark. Both the softmax implementations and the STREAM benchmark use AVX2 instructions.

Fig. 3 and 4 illustrate the memory bandwidth in each pass of the Three-Pass softmax Algorithms 1 and 2, as well as Copy and Scale STREAM benchmarks [STREAM]. As recommended in STREAM documentation, we set the array size to four times the size of last-level cache (8650752 single-precision elements for Softmax, and 4325376 double-precision elements for STREAM). The first softmax pass (max-reduction) is the same in both versions of the Three-Pass algorithm, and thus presented only once. This pass reads one input array, and doesn’t have a direct equivalent in STREAM, but achieves similar bandwidth to STREAM Copy and Scale benchmarks, which both read one array and write one array. The second pass in Algorithm 1 reads one array, computes exponentials on the inputs, and accumulates them. It achieves of STREAM Copy bandwidth in AVX512 version and in AVX2 version. The second pass Algorithm 2 is similar, but additionally stores computed exponents into the output array. Although it achieves higher bandwidth than the second pass in Algorithm 1, it takes substantially longer to complete; the higher bandwidth is due to doubling the number of transferred bytes with a less than proportional increase in run time. The third pass of Algorithm 1 reads one array, computes exponentials on the inputs, scales them, and writes results to another array. This pass does the same number of memory operations as STREAM Scale benchmark, but substantially more computational operations. Yet, our auto-tuned implementations exceed the performance of STREAM Scale benchmark in both the AVX512 and the AVX2 versions. The third pass of the Algorithm 2 is an in-place variant of STREAM Scale benchmark. The processor clearly favors in-place operation: it is faster than STREAM Scale with AVX512, and faster with AVX2.

To summarize, passes 1 and 3 of the Algorithm with Recomputing 1 demonstrate similar memory performance to STREAM benchmark, and pass 2 in AVX512 implementation is not far behind. Passes 1 and 2 of the Algorithm 2 with Reloading similarly perform close to STREAM bandwidth, and pass 3 is significantly faster than STREAM Scale benchmark. These results confirm that performance of Three-Pass softmax algorithms is limited by achievable memory bandwidth, and suggest that softmax computation can be further accelerated only through reducing the number of memory operations.

6.5 The Two-Pass Algorithm

Figure 5: Performance comparison of the Algorithms 1, 2, and 3 in the AVX512 implementations. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

On Fig. 5 we present the performance of the Two-Pass softmax algorithm in comparison with the two versions of the Three-Pass algorithm in AVX512 implementations. On out-of-cache working sets the proposed Two-Pass softmax algorithm outperforms Three-Pass algorithms by .

Figure 6: Performance comparison of the Algorithms 1, 2, and 3 in the AVX2 implementations. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

Fig. 6 similarly compares performance of the Two-Pass and the Three-Pass algorithms in the AVX2 implementations. Here, the Two-Pass algorithm outperforms Three-Pass algorithm with Reloading of exponential computations by on out-of-cache workloads. The smaller speedups, compared to AVX512 implementation, are explained by relatively higher cost of recomputing exponentials in AVX2 compared to AVX512. If we compare to the Three-Pass Algorithm 1, which similarly recomputed exponentials, the Two-Pass algorithm wins by .

Figure 7: Absolute runtime of the passes in the Algorithms 1, 2, and 3 in both the AVX2 and the AVX512 implementations. The algorithms were evaluated on arrays of 8,650,752 single-precision elements on the Skylake-X system.

On Fig.7 we decompose the absolute run-time for the three algorithms and two SIMD instruction sets into individual memory passes and offers insight into the origin of performance improvements with the Two-Pass algorithm. The two passes of the Two-Pass softmax algorithm have similar, but slightly higher absolute run-time to the last two passes of the Three-Pass softmax algorithm with recomputation of re-computation of exponential function, which share the same memory access pattern. The slightly higher run-time in the passes of the Two-Pass algorithm can be explained by larger number of operations needed for accumulation on the representation compared to just accumulating scalar floating-point values.

6.6 Multi-Threaded Performance

The benchmarks in Sec. 6.4 and Sec. 6.5 presented performance in single-threaded softmax computation, and demonstrated that on HPC-class systems softmax saturates memory bandwidth even when running on a single core. Utilizing multiple cores increases available computational resources faster than achievable memory bandwidth, and therefore increases the advantage of the bandwidth-saving Two-Pass softmax algorithm. To quantify this advantage, we fix the size of the array at 4 times the last-level cache size [STREAM], and vary the number of threads from 1 to 6 (number of cores in the system) to 12 (number of logical processors, including hyperthreads, in the system).

Figure 8: Weak scaling of the softmax algorithms in the AVX512 implementations on the Skylake-X system.

Fig. 8 illustrates weak multi-core scaling of the AVX512 implementations. As the number of threads grows, the advantage of the Two-Pass over Three-Pass algorithms remains unchanged at . Interestingly, the reloading variant of the Three-Pass algorithm scales worse than the recomputing variant, and the recomputing Algorithm 1 outperforms the reloading Algorithm 2 when at least 3 cores are being utilized.

Figure 9: Weak scaling of the softmax algorithms in the AVX2 implementations on the Skylake-X system.

Fig. 9 similarly illustrates weak multi-core scaling of the AVX2 implementations. The advantage of the Two-Pass over Three-Pass algorithms grows from on a single core to when utilizing all 6 cores to when also using hyperthreads.

6.7 Comparison with Intel DNNL

The results in Sec. 6.5-6.6 demonstrate that on out-of-cache inputs the Two-Pass softmax algorithm outperforms the Three-Pass softmax algorithms in our implementation, but leaves out the question of whether our implementations are competitive with state-of-the-art. To demonstrate the absolute effectiveness of the Two-Pass algorithm, we compared our implementations of the three softmax algorithm to the softmax primitive in Intel Deep Neural Network Library (DNNL) version 1.1.1.

Intel DNNL implements the Three-Pass softmax Algorithm 2 with reloading of computed exponentials. It includes implementations for SSE4.1, AVX, and AVX512 instruction sets, and automatically dispatch to AVX512 implementation on the Skylake-X processor. Unlike our implementations, Intel DNNL generates implementation in runtime using Just-in-Time (JIT) technology. JIT code generation potentially allows adaptation of implementation to parameters of a particular softmax operator (e.g. number of channels).

Figure 10: Performance comparison of our implementation of Algorithms 1, 2, and 3, with the softmax implementation in Intel DNNL library. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

Fig. 10 presents the comparison between two implementations (Ours and DNNL) of the Three-Pass algorithm with reloading of exponentials, and the Two-Pass softmax algorithm in our implementation. For the Three-Pass algorithm with reloading, our implementation ourperforms Intel DNNL for the majority of problem sizes. Its advantage is particularly high – over 2X – when data fits into L1, diminish to within L2, and levels off at for out-of-cache problem sizes. As the implementation in Intel DNNL is less efficient than ours, our Two-Pass softmax algorithm outperforms DNNL softmax primitive on all problem sizes: it is faster on out-of-cache problem sizes, and up to when input fits into L1 cache.

6.8 Validation on Alternative Processors

The results in Sec. 6.4-6.7 were all collected on the Xeon W-2135 processor with the Intel Skylake-X microarchitecture, which prompts a question as to whether the advantage of the Two-Pass softmax algorithm is restricted to a specific type of processor. To demonstrate that the Two-Pass algorithm generalize to other types of processors, we replicated results of Sec. 6.5 on Xeon E5-2696 v4 processor with Intel Broadwell microarchitecture and Ryzen 9 3900X with AMD Zen 2 microarchitecture. Both of these processors support the AVX2, but not the AVX512 instruction set, and have different cache hierarchy parameters than the Intel Skylake-X system.

Figure 11: Performance comparison of Algorithms 1, 2, and 3 on an Intel Broadwell-based system. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

Fig. 11 presents performance of the three softmax algorithms on the Intel Broadwell system. Although the Two-Pass softmax algorithm demonstrates inferior performance on problems which fit into L2 cache, it gets competitive with the Three-Pass softmax algorithms on L3-sizes problems, and outperforms them by on out-of-cache problems.

Figure 12: Performance comparison of Algorithms 1, 2, and 3 on a Ryzen 9 3900X system. Gray dotted lines denote boundaries of level-1, level-2, and level-3 caches.

Fig. 12 shows a similar picture on AMD Zen 2 microarchitecture. Here, the Three-Pass algorithms deliver superior performance as long as data fits into L3 cache, but loose to the Two-Pass algorithm when the data exceeds cache size.

7 Conclusion

We presented a novel Two-Pass algorithm for softmax computation and demonstrated that the new Two-Pass algorithm is up to 28% faster than the traditional Three-Pass algorithm on large input vectors. The algorithm, however, offers no advantage over reloading variant of the Three-Pass algorithm when the data fits into the processor’s cache.

This study focused on performance on a CPU, but the algorithm has great potential for GPU and hardware AI accelerators. These platforms further shift the balance between compute and memory performance towards expensive memory and cheap floating-point operations, and would favor reduced memory intensity of the presented Two-Pass softmax algorithm.

To foster reproducibility, we released an open-source implementation of the new Two-Pass Softmax algorithm and other experiments in this paper as a part of XNNPACK library at


We thank Matthias Grundmann, Sergey Ioffe, Juhyun Lee, Karthik Raveendran, and Richard Vuduc for helpful feedback and discussion, and Vijay Thakkar for sharing access to the Ryzen 9 3900X system.