# When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute

论文地址：

* [https://arxiv.org/abs/2102.12459](https://arxiv.org/abs/2102.12459)

## 整体思路以及计算方式

将SRU中的全连接层替换成MHA，增加模型表达性，首先回顾SRU：

$$
\mathbf{U}^{\top}=\left(\begin{array}{l} \mathbf{W} \\ \mathbf{W}^{\prime} \\ \mathbf{W}^{\prime \prime} \end{array}\right) \mathbf{X}^{\top}
$$

其中$$\mathbf{U} \in \mathbb{R}^{L \times 3 \times d}, \mathbf{X} \in \mathbb{R}^{L \times d}.$$

然后利用递推式计算：

$$
\begin{aligned} \mathbf{f}[t] & =\sigma(\mathbf{U}[t, 0]+\mathbf{v} \odot \mathbf{c}[t-1]+\mathbf{b}) \\ \mathbf{r}[t] & =\sigma\left(\mathbf{U}[t, 1]+\mathbf{v}^{\prime} \odot \mathbf{c}[t-1]+\mathbf{b}^{\prime}\right) \\ \mathbf{c}[t] & =\mathbf{f}[t] \odot \mathbf{c}[t-1]+(1-\mathbf{f}[t]) \odot \mathbf{U}[t, 2] \\ \mathbf{h}[t] & =\mathbf{r}[t] \odot \mathbf{c}[t]+(1-\mathbf{r}[t]) \odot \mathbf{x}[t] . \end{aligned}
$$

这里的改进是，将$$\mathbf U$$部分替换为MHA：

$$
\begin{aligned} \mathbf{U}^{\top}&=\mathbf{W}^o\mathrm{Layernorm}(\mathbf{Q}+\alpha \cdot \mathbf{A}) \\ \mathbf{A}^{\top}&=\operatorname{softmax}\left(\frac{\mathbf{Q}^{\top} \mathbf{K}}{\sqrt{d^{\prime}}}\right) \mathbf{V}^{\top} \\ \mathbf{Q}& =\mathbf{W}^q \mathbf{X}^{\top} \\ \mathbf{K}& =\mathbf{W}^k \mathbf{Q} \\ \mathbf{V}& =\mathbf{W}^v \mathbf{Q} \end{aligned}
$$

## 代码

* [https://github.com/asappresearch/sru](https://github.com/asappresearch/sru)

## 细节

作者测试了每几层增加Attention的效果，最后的结论是，只要增加一层Attention，就能比纯SRU的效果好很多。

## 简评

不错的工作，最后的结论也有启发意义。
