ASTRA: Reconfigurable Training Architecture Design for Nonlinear Softmax and Activation Functions in Transformers

The efficient training of Transformer-based neural networks on resource-constrained personal devices is attracting continuous attention due to domain adaptions and privacy concerns. However, Transformers’ intensive and complicated computations, especially the crucial nonlinear Softmax and activation (Act) functions, pose challenges for training deployment on edge. This brief proposes an efficient training architecture for both Softmax and Act functions. Specifically, we present a quantized nonlinear training algorithm based on fully integer (int) arithmetic with sufficient training accuracy. Then, we develop a reconfigurable hardware architecture to efficiently support various operations during the training of these nonlinear functions. Furthermore, a staged parallel and pipelined (SPAP) dataflow is presented 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 500 MHz on Xilinx ZCU102 FPGA, outperforming prior works.