Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformers

• Transformers use self-attention alongside matrix multiplication, introducing softmax and dropout into each attention layer.Source: §I · p. 1
• Softmax contributes a large fraction of Transformer runtime, particularly at the longer sequence lengths of recent networks.Source: §I · p. 1, Figure 1 · p. 1
• Prior inference accelerators largely targeted matrix-multiply-dominated CNNs, MLPs, and LSTMs, leaving Transformer softmax comparatively neglected.Source: §I · p. 1
• Softermax combines base replacement, low-precision softmax computations, and online normalization.Source: §I · p. 1
• Softermax-aware fine-tuning reduces accuracy loss without additional training overhead, while the paper also details its inference-accelerator microarchitecture.Source: §I · p. 1

• The section first describes computations and bottlenecks in Transformer-based networks, then reviews efforts addressing those bottlenecks.Source: §II · p. 2

• A Transformer layer contains a multi-headed attention block followed by a feed-forward block.Source: A. Transformers · p. 2, Figure 2 · p. 2
• Attention forms query, key, and value matrices, applies scaled softmax to query-key products, then multiplies the resulting attention matrix by the value matrix.Source: A. Transformers · p. 2
• Softmax and dropout are less common than matrix multiplication but become larger bottlenecks as Transformer sequence lengths increase.Source: A. Transformers · p. 2
• Existing DL hardware largely neglected softmax because it was minor in CNNs, MLPs, and LSTMs, whereas it is central to Transformer attention.Source: B. Softmax Bottlenecks · p. 2

• Recent Transformer-efficiency methods commonly create smaller networks through knowledge distillation, inductive biases, or approximations rather than directly targeting softmax.Source: C. Related Work · p. 2
• Unlike prior efforts, Softermax addresses the explicit maximum-computation pass and replaces the natural exponential with a base-two calculation.Source: C. Related Work · p. 3
• Softermax comprehensively optimizes softmax through base replacement, reduced precision, and online normalization.Source: C. Related Work · p. 3
• 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.Source: §III · p. 3, Figure 3 · p. 3

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.Source: A. Base Replacement · p. 3, C. Hardware-friendly Online… · p. 4
• 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.Source: B. Low-precision Softmax Operations · p. 3
• 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.Source: C. Hardware-friendly Online… · pp. 3–4
• 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.Source: C. Hardware-friendly Online… · p. 4

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.Source: A. Unnormed Softmax Unit · p. 4
• 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.Source: A. Unnormed Softmax Unit · p. 4
• 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.Source: A. Unnormed Softmax Unit · p. 4
• 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.Source: A. Unnormed Softmax Unit · pp. 4–5

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.Source: B. Normalization Unit · p. 5

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.Source: C. Accelerator Integration · p. 5
• 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.Source: §V · p. 5
• 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.Source: Table I · p. 5, Table II · p. 5
• 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.Source: §V · p. 5

• The evaluation first tests whether Softermax’s hardware-friendly softmax proposals negatively affect accuracy.Source: §VI · p. 5
• The evaluation then measures whether Softermax produces more efficient hardware implementations.Source: §VI · p. 5
• Accuracy preservation and hardware efficiency are evaluated as separate parts of the Softermax assessment.Source: §VI · p. 5

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.Source: A. Impact on Accuracy · p. 5, Table III · p. 6

Compute Unit Level Analysis

• 4x smaller and 9.53x more energy efficient, the Unnormed Softmax unit outperforms the DesignWare-based baseline.Source: B. Impact on Hardware Efficiency · pp. 5–6
• 1.54x smaller and 2.53x more energy efficient, the Normalization Unit also outperforms the baseline.Source: B. Impact on Hardware Efficiency · p. 6

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.Source: B. Impact on Hardware Efficiency · p. 6, Table IV · p. 6

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.Source: B. Impact on Hardware Efficiency · p. 6, Figure 5 · p. 6

• Softermax targets Transformer softmax bottlenecks with a set of software and hardware optimizations.Source: §VII · p. 6
• 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.Source: §VII · p. 6
• Softermax achieves these hardware improvements with no loss in accuracy compared to the quantized baseline.Source: §VII · p. 6