Source-linked AI summary
Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformers
Jacob R. Stevens, Rangharajan Venkatesan, Steve Dai, Brucek Khailany, Anand Raghunathan
TL;DR
Transformer softmax operations are a significant runtime bottleneck, but existing accelerators have largely neglected them. Softermax combines base replacement, low-precision computation, online normalization, and accuracy-oriented fine-tuning; it achieves improved hardware efficiency with negligible accuracy impact.
Problem
Transformer attention makes softmax a significant runtime bottleneck, while existing deep-learning hardware has largely neglected this operation.
Method
Softermax combines base replacement, low-precision softmax operations, online normalization, and Softermax-aware fine-tuning implemented with dedicated accelerator hardware.
Results
2.35x greater energy efficiency at 0.90x the area is achieved with negligible accuracy impact; the worst accuracy drop is under 0.5%.
Takeaways & Limitations
Softermax provides an acceptable replacement for full-precision, numerically stable softmax in Transformer-based networks.
Abstract
from arXiv · showhide
Transformers have transformed the field of natural language processing. This performance is largely attributed to the use of stacked self-attention layers, each of which consists of matrix multiplies as well as softmax operations. As a result, unlike other neural networks, the softmax operation accounts for a significant fraction of the total run-time of Transformers. To address this, we propose Softermax, a hardware-friendly softmax design. Softermax consists of base replacement, low-precision softmax computations, and an online normalization calculation. We show Softermax results in 2.35x the energy efficiency at 0.90x the size of a comparable baseline, with negligible impact on network accuracy.
I. INTRODUCTION
Transformers’ growing importance and longer sequences make softmax a substantial runtime bottleneck, motivating Softermax, a hardware-friendly softmax design with hardware/software support and negligible accuracy impact.
- Transformers use self-attention alongside matrix multiplication, introducing softmax and dropout into each attention layer.
- Softmax contributes a large fraction of Transformer runtime, particularly at the longer sequence lengths of recent networks.
- Prior inference accelerators largely targeted matrix-multiply-dominated CNNs, MLPs, and LSTMs, leaving Transformer softmax comparatively neglected.
- Softermax combines base replacement, low-precision softmax computations, and online normalization.
- Softermax-aware fine-tuning reduces accuracy loss without additional training overhead, while the paper also details its inference-accelerator microarchitecture.
II. PRELIMINARIES AND RELATED WORK
This section introduces the computational characteristics and bottlenecks of Transformer-based networks before discussing efforts to address them.
- The section first describes computations and bottlenecks in Transformer-based networks, then reviews efforts addressing those bottlenecks.
A. Transformers
Transformer layers combine multi-headed attention with feed-forward processing, and softmax becomes an increasingly important bottleneck as sequence lengths grow.
- A Transformer layer contains a multi-headed attention block followed by a feed-forward block.
- Attention forms query, key, and value matrices, applies scaled softmax to query-key products, then multiplies the resulting attention matrix by the value matrix.
- Softmax and dropout are less common than matrix multiplication but become larger bottlenecks as Transformer sequence lengths increase.
- Existing DL hardware largely neglected softmax because it was minor in CNNs, MLPs, and LSTMs, whereas it is central to Transformer attention.
C. Related Work
Prior efficiency efforts mainly reduced Transformer model size, whereas Softermax comprehensively targets softmax computation through base replacement, reduced precision, and online normalization.
- Recent Transformer-efficiency methods commonly create smaller networks through knowledge distillation, inductive biases, or approximations rather than directly targeting softmax.
- Unlike prior efforts, Softermax addresses the explicit maximum-computation pass and replaces the natural exponential with a base-two calculation.
- Softermax comprehensively optimizes softmax through base replacement, reduced precision, and online normalization.
- The algorithm switches from e to 2, uses fixed-point low-precision exponentiation and division, avoids an explicit maximum pass through online normalization, and adds Softermax-aware fine-tuning.
A. Base Replacement
Softermax replaces the natural-exponential softmax with base-two, low-precision operations and hardware-friendly online normalization. Integer maxima make renormalization a shift, avoiding an additional explicit maximum pass.
- Base Replacement: Base replacement changes the exponential calculation from Euler’s number, e, to 2, reducing hardware overhead.The natural exponential otherwise requires base conversion in specialized hardware.
- Base Replacement: Low-precision exponentiation, accumulation, and division reduce the cost of softmax computation in custom hardware.Softermax-aware fine-tuning is used to minimize accuracy loss from these operations.
- Base Replacement: Online normalization calculates the maximum and denominator continuously, avoiding an additional explicit pass through the input vector.When a new maximum appears, the running sum is renormalized to preserve the correct accumulation.
- Base Replacement: Using an integer maximum guarantees integer differences between maxima, so renormalization by 2^(OldMax−NewMax) can be implemented with a shifter.This is the hardware co-design modification applied to online normalization.
IV. SOFTERMAX HARDWARE
The Unnormed Softmax unit computes local maxima, base-two powers, and accumulated denominators, using piecewise approximation and online reduction across vector slices.
- IV. SOFTERMAX HARDWARE: The Unnormed Softmax unit contains IntMax, Power of Two, and Reduction subunits.Together they determine local maxima, approximate powers of two, and accumulate the denominator.
- IV. SOFTERMAX HARDWARE: The IntMax unit applies a ceiling to each element in parallel before finding the maximum of a vector slice.This implements integer maximum rather than ordinary maximum.
- IV. SOFTERMAX HARDWARE: The Power of Two unit decomposes fixed-point inputs into integer and fractional parts, approximating the fractional part with a four-segment linear piece-wise function.The implementation uses four segments, requiring a shift left by two before lookup and interpolation.
- IV. SOFTERMAX HARDWARE: The Reduction unit sums unnormalized softmax outputs and renormalizes running sums when a newly observed local maximum exceeds the current maximum.For vectors spanning multiple slices, it reads the current maximum and running sum from buffers, applies a shifter, and combines the sums.
B. Normalization Unit
The Normalization Unit renormalizes the numerator and divides by the accumulated denominator to produce the final softmax result.
- B. Normalization Unit: The Normalization Unit renormalizes the numerator and performs the final division to obtain the result.Integer maxima make the renormalization difference integral, enabling a shifter; division uses a piecewise reciprocal unit and integer multiplier.
C. Accelerator Integration
Softermax is integrated into an existing accelerator, with hardware and software setups used to evaluate implementation overhead and accuracy.
- C. Accelerator Integration: The Unnormed Softmax Unit can integrate into tensor-processing hardware such as GPU tensor cores, TPUs, or dedicated DNN accelerators.Using MAGNet as a baseline, it can be integrated into the post-processing unit and sized to match MAC throughput.
- C. Accelerator Integration: The software evaluation modifies PyTorch Huggingface with percentile calibration, quantization-aware fine-tuning, and custom fixed-point softmax passes.The setup uses 8-bit weights and activations and straight-through-estimator-style backward passes.
- C. Accelerator Integration: Table I summarizes Softermax bitwidths using integer and fractional components, while Table II summarizes the experimental setup.These tables define the precision configuration and evaluation conditions.
- C. Accelerator Integration: The hardware evaluation uses high-level synthesis to integrate Softermax into MAGNet and compares it with a standard 16-bit floating-point softmax.The stated baseline is optimistic because current state-of-the-art accelerators use 32-bit precision.
VI. EVALUATION
The evaluation separates Softermax’s impact on accuracy from its hardware efficiency, testing whether hardware-friendly softmax changes preserve model accuracy and improve implementation efficiency.
- The evaluation first tests whether Softermax’s hardware-friendly softmax proposals negatively affect accuracy.
- The evaluation then measures whether Softermax produces more efficient hardware implementations.
- Accuracy preservation and hardware efficiency are evaluated as separate parts of the Softermax assessment.
A. Impact on Accuracy
Softermax preserves accuracy across evaluated Transformer tasks while reducing softmax hardware cost at both compute-unit and processing-element levels.
- Impact on Accuracy: Softermax causes negligible accuracy loss across BERT-Base and BERT-Large on SQuAD and GLUE.The worst accuracy drop is under 0.5%, while average accuracy increases 0.9% for BERT-Base and 0.7% for BERT-Large.
- Compute Unit Level Analysis: 4x smaller and 9.53x more energy efficient, the Unnormed Softmax unit outperforms the DesignWare-based baseline.
- Compute Unit Level Analysis: 1.54x smaller and 2.53x more energy efficient, the Normalization Unit also outperforms the baseline.
- PE-level Analysis: 1.11x more area efficient and 2.35x more energy efficient, Softermax improves efficiency when integrated into a 32-wide MAGNet PE.
- Sequence Length Sweep: Softermax scales better as sequence length increases, starting from a lower baseline with a shallower energy-consumption slope.The sweep evaluates both 32-wide and 16-wide PE configurations.
VII. CONCLUSION
The paper concludes that Softermax addresses Transformer softmax bottlenecks through software and hardware optimizations. It reports substantial area and energy improvements while maintaining accuracy relative to the quantized baseline.
- Softermax targets Transformer softmax bottlenecks with a set of software and hardware optimizations.
- 4x area and 9.53x energy improvements over conventional softmax units translate to 1.11x area and 2.35x energy improvements within a DNN accelerator.
- Softermax achieves these hardware improvements with no loss in accuracy compared to the quantized baseline.