SOLE: Hardware-Software Co-design of Softmax and LayerNorm for Efficient Transformer Inference
This work addresses real-time inference speed and efficiency limitations in Transformers for NLP and CV applications, offering a domain-specific optimization.
The paper tackles the inefficiency of Softmax and LayerNorm in Transformer inference by proposing SOLE, a hardware-software co-design that achieves 3.04x to 3.86x energy-efficiency improvements and 2.82x to 3.32x area-efficiency improvements over prior state-of-the-art custom hardware without retraining.
Transformers have shown remarkable performance in both natural language processing (NLP) and computer vision (CV) tasks. However, their real-time inference speed and efficiency are limited due to the inefficiency in Softmax and Layer Normalization (LayerNorm). Previous works based on function approximation suffer from inefficient implementation as they place emphasis on computation while disregarding memory overhead concerns. Moreover, such methods rely on retraining to compensate for approximation error which can be costly and inconvenient. In this paper, we present SOLE, a hardware-software co-design for Softmax and LayerNorm which is composed of E2Softmax and AILayerNorm. E2Softmax utilizes log2 quantization of exponent function and log-based division to approximate Softmax while AILayerNorm adopts low-precision statistic calculation. Compared with state-of-the-art designs, we achieve both low-precision calculation and low bit-width storage on Softmax and LayerNorm. Experiments show that SOLE maintains inference accuracy without retraining while offering orders of magnitude speedup and energy savings over GPU, achieving 3.04x, 3.86x energy-efficiency improvements and 2.82x, 3.32x area-efficiency improvements over prior state-of-the-art custom hardware for Softmax and LayerNorm, respectively.