Reusing Softmax Hardware Unit for GELU Computation in Transformers

TL;DR

This work addresses the hardware inefficiency of separately computing GELU and softmax in transformers. It derives a mathematical transformation that expresses GELU as GELU$(z) = z\,\text{softmax}^2_1([k,-k])$ with $k = \sqrt{2/\pi}\,(z + 0.044715 z^3)$, enabling GELU to be computed using a softmax unit. A dual-mode softmax core is proposed to operate on either $N$ inputs (normal softmax) or $N/2$ two-element subvectors (GELU mode), along with overhead circuitry to compute $k$ and finalize GELU results; the unit can compute multiple GELUs in parallel. Experimental results show unchanged accuracy on GLUE benchmarks with the proposed approach and demonstrate hardware savings: around $6.1\%$ area and $11.9\%$ power on average for the combined GELU-softmax unit, with additional details on dual-mode overhead and throughput gains. Overall, the method enables efficient transformer acceleration by reusing softmax hardware, reducing area and power while maintaining NLP performance.

Abstract

Transformers have improved drastically the performance of natural language processing (NLP) and computer vision applications. The computation of transformers involves matrix multiplications and non-linear activation functions such as softmax and GELU (Gaussion Error Linear Unit) that are accelerated directly in hardware. Currently, function evaluation is done separately for each function and rarely allows for hardware reuse. To mitigate this problem, in this work, we map the computation of GELU to a softmax operator. In this way, the efficient hardware units designed already for softmax can be reused for computing GELU as well. Computation of GELU can enjoy the inherent vectorized nature of softmax and produce in parallel multiple GELU outcomes. Experimental results show that computing GELU via a pre-existing and incrementally modified softmax hardware unit (a) does not reduce the accuracy of representative NLP applications and (b) allows the reduction of the overall hardware area and power by 6.1% and 11.9%, respectively, on average.

Figures (4)

• Figure 1: An encoder-only transformer layer. For instance, BERT-base bert consists of twelve such layers.
• Figure 2: The dual-mode softmax hardware unit. Computation follows Eq. \ref{['e:softmax-log-domain']} that implements division in the logarithm domain. The extra logic required for supporting the dual-mode of operation is highlighted in blue.
• Figure 3: The combined GELU/softmax hardware unit. When in normal mode, softmax is driven by $N$ inputs $z_i$ and produces $N$ output probabilities. When in GELU mode, half of the inputs and outputs are used to compute $N/2$ GELU outcomes in parallel.
• Figure 4: The area and power of the combined GELU-softmax unit introduced in this work and the design that employs $N/2$ i-GELU ibert units and a state-of-the-art single-mode softmax unit edgebert. The dual-mode softmax unit used in the proposed design is an enhanced derivative of the same single-mode softmax unit.