H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequences

论文地址：

• https://arxiv.org/abs/2107.11906

参考资料：

• https://www.mis.mpg.de/scicomp/Fulltext/WS_HMatrices.pdf

整体思路以及计算方式

利用层次化的方式计算Attention（本质上任然是稀疏的方法），核心思路是只计算如下位置的Attention：

$$\tilde{\mathbf A}^{(0)} \propto\left[\begin{array}{c|c|c|c|c|c|c|c} 2 & 2 & & & & & & \\ \hline 2 & 2 & 2 & & & & & \\ \hline & 2 & 2 & 2 & & & & \\ \hline & & 2 & 2 & 2 & & & \\ \hline & & & 2 & 2 & 2 & & \\ \hline & & & & 2 & 2 & 2 & \\ \hline & & & & & 2 & 2 & 2 \\ \hline & & & & & & 2 & 2 \end{array}\right],\tilde{\mathbf A}^{(1)} \propto\left[\begin{array}{l|l|l|l} & 2 & & \\ \hline 2 & & 2 & \\ \hline & 2 & & 2 \\ \hline & & 2 & \end{array}\right],\tilde{\mathbf A}^{(2)} \propto\left[\begin{array}{l|l} & 2 \\ \hline 2 & \end{array}\right]$$

计算公式为：

$$\begin{aligned} \mathbf Y =\mathbf A\mathbf V=\mathbf Y^{(0)}+\mathbf P^{(0)}\left(\tilde{\mathbf Y}^{(1)}+\mathbf P^{(1)} \tilde{\mathbf Y}^{(2)}\right) \\ \mathbf Y^{(0)}=\mathbf A^{(0)} \mathbf V^{(0)}, \tilde{\mathbf Y}^{(l)}=\tilde{\mathbf A}^{(l)} \tilde{\mathbf V}^{(l)}, l=1,2 \end{aligned}$$

其中$\mathbf P^{(i)}$为预先计算好的矩阵，$\mathbf A^{(i)}$的计算方式如下：

$$\begin{aligned} \tilde{\mathbf A}^{(i)} &=\exp(\tilde{\mathbf S}^{(i)})= \exp \left(\tilde{\mathbf Q}^{(i)}{\tilde{\mathbf K}^{(i)}}^\top\right)\\ \tilde{\mathbf Q}_{j}^{(l+1)} &=\frac{1}{2}\left(\tilde{\mathbf Q}_{2 j}^{(l)}+\tilde{\mathbf Q}_{2 j+1}^{(l)}\right) \\ \tilde{\mathbf K}_{j}^{(l+1)} &=\frac{1}{2}\left(\tilde{\mathbf K}_{2 j}^{(l)}+\tilde{\mathbf K}_{2 j+1}^{(l)}\right) \\ \tilde{\mathbf V}_{j}^{(l+1)} &=\left(\tilde{\mathbf V}_{2 j}^{(l)}+\tilde{\mathbf V}_{2 j+1}^{(l)}\right)\\ \tilde{\mathbf Q}^{(0)}&=\mathbf Q, \tilde{\mathbf K}^{(0)}=\mathbf K, \tilde{\mathbf V}^{(0)}=\mathbf V \end{aligned}$$

时间复杂度

$O(knd)$。

训练以及loss

不变。

代码

• https://github.com/lucidrains/h-transformer-1d

实验以及适用场景

适用于所有场景。

细节

具体实现需要细读代码。

简评

很新颖的思路，可以趁此机会学习Hierarchical Matrix。