Is Attention Better Than Matrix Decomposition

Is Attention Better Than Matrix Decomposition?

论文地址：

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

参考资料：

• https://zhuanlan.zhihu.com/p/369769485

• https://zhuanlan.zhihu.com/p/369855045

• https://zhuanlan.zhihu.com/p/370410446

整体思路以及计算方式

思路是利用矩阵分解代替Attention。

假设$\mathbf X\in \mathbb R^{d\times n}$可以分解为如下形式：

$$\mathbf {X}=\overline{\mathbf {X}}+\mathbf {E}=\mathbf {D}\mathbf C+\mathbf {E}$$

这些变量满足如下条件（$\mathbf E$表示噪声）：

$$\begin{aligned} \overline{\mathbf {X}}& \in \mathbb{R}^{d \times n} \\ \mathbf {E}& \in \mathbb{R}^{d \times n} \\ \mathbf D&\in \mathbb R^{d\times r}\\ \mathbf C&\in \mathbb R^{r\times n}\\ \end{aligned}$$

使用流程为通过某种方式计算$\mathbf D,\mathbf C$，最后输出$\mathbf D\mathbf C$。

作者给出了两种方式计算$\mathbf D,\mathbf C$：

• Soft VQ

  • for $k=1,\ldots, K$:

    • $\mathbf {C} \leftarrow \operatorname{Softmax}\left(\frac{1}{T} \operatorname{cosine}(\mathbf {D},\mathbf {X})\right)$

    • $\mathbf {D} \leftarrow \mathbf {X} \mathbf {C}^{\top} \operatorname{diag}\left(\mathbf {C} \mathbf{1}_{n}\right)^{-1}$

  • return $\overline {\mathbf X} =\mathbf D\mathbf C$

• NMF with MU

  • for $k=1,\ldots, K$:

    • ${\mathbf C}_{i j} \leftarrow {\mathbf C}_{i j} \frac{\left(\mathbf {D}^{\top}\mathbf X\right)_{i j}}{\left(\mathbf {D}^{\top} \mathbf {D C}\right)_{i j}}$

    • $\mathbf {D}_{i j} \leftarrow \mathbf {D}_{i j} \frac{\left(\mathbf {X} \mathbf C^{\top}\right)_{i j}}{\left(\mathbf {D C} \mathbf {C}^{\top}\right)_{i j}}$

  • return $\overline{\mathbf X} =\mathbf {DC}$

时间复杂度

Soft VQ时间复杂度：

• $\mathbf {C} \leftarrow \operatorname{Softmax}\left(\frac{1}{T} \operatorname{cosine}(\mathbf {D},\mathbf {X})\right)$，所以时间复杂度为$O(nrd)$

  • $\operatorname{cosine}(\mathbf {D},\mathbf {X})$需要计算$\mathbf D^{\top} \mathbf X$，即$r\times d,d\times n\to r\times n$，所以时间复杂度为$O(nrd)$

  • $\operatorname{Softmax}$：$r\times n \to r\times n$，所以时间复杂度为所以时间复杂度为$O(nr)$

  • 总时间复杂度为$O(nrd)$

• $\mathbf {D} \leftarrow \mathbf {X} \mathbf {C}^{\top} \operatorname{diag}\left(\mathbf {C} \mathbf{1}_{n}\right)^{-1}:$

  • $\operatorname{diag}\left(\mathbf {C} \mathbf{1}_{n}\right)^{-1}:r\times n,n\times 1 \to r\times 1 \to r\times r$，时间复杂度为$O(nr)$

  • $\mathbf {X} \mathbf {C}^{\top}:d\times n, n\times r \to d\times r$，时间复杂度为$O(nrd)$

  • $\mathbf {C}^{\top} \operatorname{diag}\left(\mathbf {C} \mathbf{1}_{n}\right)^{-1}: d\times r, r\times r \to r\times r$，时间复杂度为$O(dr^2)$

• 所以时间复杂度为$O(nrd+ dr^2)$

• 循环$K$次，时间复杂度为$O(K(nrd+ dr^2))$

• $\overline {\mathbf X} =\mathbf D\mathbf C: d\times r , r\times n \to d\times n$，时间复杂度为$O(nrd)$

• 总时间复杂度为$O((K+1)nrd + Kdr^2)$

NMF with MU时间复杂度：

• $\mathbf {C}_{i j} \leftarrow \mathbf {C}_{i j} \frac{\left(\mathbf {D}^{\top} \mathbf X\right)_{i j}}{\left(\mathbf {D}^{\top} \mathbf {D C}\right)_{i j}}$

  • $\mathbf {D}^{\top} \mathbf X: r\times d, d\times n\to r\times n$，时间复杂度为$O(nrd)$

  • $\mathbf {D}^{\top} \mathbf {D C}$

    • 先计算$\mathbf {DC}$，再计算$\mathbf {D}^{\top} \mathbf {D C}$

      • $\mathbf {DC}:d\times r, r\times n \to d\times n$，时间复杂度为$O(nrd)$

      • $\mathbf {D}^{\top} \mathbf {D C}: r\times d, d\times n \to r\times n$，时间复杂度为$O(nrd)$

    • 先计算$\mathbf {D}^{\top}\mathbf D$，再计算$\mathbf {D}^{\top} \mathbf {D C}$

      • $\mathbf {D}^{\top}\mathbf D: r\times d, d\times r \to r\times r$，时间复杂度为$O(dr^2)$

      • 再计算$\mathbf {D}^{\top} \mathbf {D C}:r\times r, r\times n\to r\times n$，时间复杂度为$O(nrd)$

    • 一般$r<n$，所以选择第二种算法，时间复杂度为$O(nrd+dr^2)$

  • 两次element wise乘法/除法：$r\times n \to r\times n \to r\times n$，时间复杂度为$O(nr)$

  • 循环$K$次，时间复杂度为$O(K(nrd+dr^2))$

  • $\overline {\mathbf X}=\mathbf {DC}: d\times r , r\times n \to d\times n$，时间复杂度为$O(nrd)$

  • 总时间复杂度为$O((K+1)nrd + Kdr^2)$

注意$r, K$一般不会很大（远小于$n$），所以该方法的时间复杂度关于序列长度大概能到线性。

训练以及loss

没有区别。

代码

• https://github.com/lucidrains/hamburger-pytorch

• https://github.com/Gsunshine/Enjoy-Hamburger

实验以及适用场景

目前只适用于Encoder结构，不适用于Decoder结构；实验主要是基于CV的，能达到和Attention相当的结果。

细节

在进行循环的时候不计算梯度，只有最后一次操作计算梯度。

简评

优点：

• 提供了一种新的理解Attention的视角；

• 方法实现比较简洁；

缺点：

• 没法直接应用到Decoder结构中，即无法训练lm；

总结

• 值得复现；