Two-Pass Softmax Algorithm

• Published in: 2020 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW) pp. 386 - 395
• Main Authors: Dukhan, Marat, Ablavatski, Artsiom
• Format: Conference Proceeding
• Language: English
• Published: IEEE 01.05.2020
• Subjects: Approximation algorithms; Bandwidth; Computational modeling; Machine learning; Machine learning algorithms; Probability distribution; Program processors; SIMD; softargmax; softmax
• DOI: 10.1109/IPDPSW50202.2020.00074

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% in AVX512 implementation, and by up to 18% in AVX2 implementation. The proposed Two-Pass algorithm also outperforms the traditional Three-Pass algorithm on Intel Broadwell and AMD Zen 2 processors.