论文标题：ASTRA: Reconfigurable Training Architecture Design for Nonlinear Softmax and Activation Functions in Transformers
论文作者：Haikuo Shao，Zhongfeng Wang
关键词：Activation function, algorithm, architecture, Softmax, training, Transformer

ASTRA提出了一种在资源受限设备上优化Transformer模型的Softmax函数与GELU函数的方法，并考虑了反向传播阶段。在Xilinx ZCU102上（500 MHz），论文提出的Softmax单元实现了每秒处理1G个输入的吞吐量（1.0 GinS）以及4.42 Gins/W的能效。

Motivation

Transformer模型在深度学习领域取得了巨大的成功，然而其庞大而又复杂的计算为其在资源受限设备上的部署带来了挑战，特别是训练阶段。

现有面向训练阶段的加速器仅考虑了提升线性层的效率，而忽略了非线性的Softmax函数与激活函数。论文在一个配备了16$\times$16大小脉动阵列、并行度为16的Softmax单元以及激活单元的FPGA加速器上进行了延迟分析：当序列长度从128增长到4096时，Softmax函数所占用的时间比例从3.2%增长到了37.7%。

Challenges and Solutions

Challenges

不同于CNN通常只在最后的分类器中需要Softmax的计算，Transformer的每个注意力层都包含Softmax操作，因此，对其进行激进的估计计算可能会导致误差传播，极大影响模型的准确率。
之前的非线性函数加速器通常只考虑推理，而训练阶段由于训练样本的长度变化、注意力掩码机制，导致固定并行度的硬件单元效率较低。
大多数研究都将Softmax和其他激活函数分开优化，而忽略了其间潜在的数学一致性。

Solutions

论文提出了一种硬件友好并且可实现量化训练的算法，Softmax与指数型激活函数可采用int类型进行计算，并可以在多种任务中实现与全精度相当的训练准确率。
论文设计了一个高效的架构来支持Softmax在前向传播（FP）与反向传播（BP）阶段的计算，该架构可配置，可用作指数激活函数的计算。一种分阶段并行化与流水化（SPAP）的数据流被进一步提出用来提升训练的吞吐率与训练效率，能够自适应训练阶段序列长度动态与计算量的要求。

Methods

Hardware-Firendly Training Algorithm

Integer-Only Arithmetic for Softmax Training

首先采用log-sum-exp方法避免softmax中的除法：
$$
y_i=\frac{e^{x_i-x_{\max}}}{\sum_{j=1}^se^{x_j-x_{\max}}}=e^{x_i-x_{\max}-\ln\left(\sum_je^{x_j-x_{\max}}\right)}
$$
令$G=\sum^s_{j=1}e^{x_j-x_{max}}> 0$，$\ln G=\ln 2\cdot \log_2 G$，将$G$以$G=2^n\times (1.m)$表示，并采用对数变换$\log_a(M\cdot N)=\log_a M + \log_b N$，与泰勒近似$\log(1+m)\approx m, m\in[0,1)$将$\log_2 G$转变为
$$
\log_2(2^n\times1.m)\approx n+m
$$

然后令$z_i = x_{max}-x_i$，通过对数变换将$y_i$的计算转换为以2为底的指数函数，即
$$
y_i=2^{(-z_i-\ln G)\cdot\log_2e}
$$
再由$\ln 2\times \log_2 e=1$，进一步转换为：
$$
y_i=2^{-z_i\cdot\log_2e-(n+m)}
$$
将指数分为整数$u_i$与小数$v_i$两部分，即：
$$
y_i = 2^{(-u_i)+(-v_i)}=2^{-v_i} \gg u_i
$$
用PWL方法估计$2^{-v_i}$部分，从而将指数函数转换为分段线性函数与移位。

Quantized Acts Based on Integer Exponential

对于指数型激活函数，可以转换为$\sigma(x)$函数与其他函数的复合，如$\text{GELU}(x)=x\cdot \sigma(2t)$，$\text{SiLU}(x)=x\cdot \sigma(x)$，$Tanh(x)=2\sigma(2x)-1$，因此论文对$\sigma(x)$函数进行优化。为了避免在x<0时计算$e^{-x}$溢出，$\sigma(x)$作$\sigma(x)=\frac{e^x}{1+e^x}$变换。
$$
\sigma(x)=\left{
\begin{array}{ll}
e^{0-\ln\left(1+e^{-x}\right)}=e^{0-\ln(1+q)}, & \text{if } x \geqslant 0, \
e^{x-\ln\left(1+e^{x}\right)}=e^{x-\ln(1+q)}, & \text{if } x < 0.
\end{array}
\right.
$$
其中，$q=e^{-|x|}$，并采用以下估计$\ln(1+q)=\ln2\cdot \log_2(1+q)\approx \ln2 \cdot (1.34q-0.34q^2)$。
进一步复用将底换为2的方法执行该指数函数。

Hardware and Dataflow Design

可重构的Softmax单元如下图所示，其中蓝色表示BP阶段的数据流。
前向传播步骤：

找到全局最大值，并将一整个向量的数据保存到Buffer A中；
然后计算$x_i-x_{max}$并乘以$c=(1.011)2$（即$\log_2e$），得到$-z_ic$，保存该向量到Buffer B中，同时将$-z_i$送入指数单元EU中，然后对指数函数的结果进行累加得到$G=\sum^s{j=1}e^{-z_j}$，$G$随后送入对数单元NLU得到$n+m$。
最后计算$-z_ic-(n+m)$，并送入EU得到最后的输出$y_i$

反向传播步骤：
论文在执行反向传播时对公式进行了重组，使得其可以复用计算单元，具体来说：
原本的梯度计算公式：
$$
\frac{\partial y_{i}}{\partial x_{j}}=
\begin{cases}
y_{i} - y_{i} y_{i}, & \text{if } i = j \
0 - y_{i} y_{j}, & \text{if } i \neq j
\end{cases}
$$
令$\mathrm{d}\boldsymbol{y}$为损失函数关于$\boldsymbol{y}$的梯度，则：
$$
\mathrm{d} \boldsymbol{x} = \mathrm{d} \boldsymbol{y} \frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}} = \mathrm{d} \boldsymbol{y} \times \left( \operatorname{diag}(\boldsymbol{y}) - \boldsymbol{y}^{T} \times \boldsymbol{y} \right)
$$

论文将其转换为：
$$
\begin{aligned}
\mathrm{d} x_{i} &= \mathrm{d} y_{1}\left(-y_{1} y_{i}\right) + \cdots + \mathrm{d} y_{i}\left(y_{i} - y_{i} y_{i}\right) + \cdots + \mathrm{d} y_{s}\left(-y_{s} y_{i}\right) \
&= y_{i}\left(\mathrm{d} y_{i} - \left(y_{1} \mathrm{d} y_{1} + \cdots + y_{i} \mathrm{d} y_{i} + \cdots + y_{s} \mathrm{d} y_{s}\right)\right) \
&= y_{i}\left(\mathrm{d} y_{i} - \boldsymbol{y} \times \mathrm{d} \boldsymbol{y}^{T}\right)
\end{aligned}
$$