An Efficient Training Architecture for Nonlinear Softmax Function in Transformers

Transformers have achieved significant success in modern deep learning. However, their intensive and complicated computations challenge Transformers’ deployment on resource-constrained devices, especially the training process. Softmax, the crucial nonlinear function in Transformers, features complex exponential and division operations. Unfortunately, existing methods cannot effectively reduce the hardware complexity of Softmax while maintaining training accuracy. In this paper, we propose an efficient architecture for Softmax training. Specifically, we present a quantized Softmax training algorithm based on integer arithmetic with sufficient training accuracy. Additionally, we design a reconfigurable architecture to efficiently support various operations during different training phases, along with a staged parallel and pipelined (SPAP) dataflow to reduce latency and improve hardware efficiency. Experimental results show that our architecture achieves up to 1.0 Softmax GinS in throughput and 4.42 GinS/W in energy efficiency at 500MHz on the Xilinx ZCU102 FPGA, which significantly outperforms prior works.