Linformer: Self-Attention with Linear Complexity

URL

https://arxiv.org/pdf/2006.04768.pdf

TL;DR

本方法—— Linformer 使用矩阵的低秩来降低原始 Transformer 的 Multi-HEAD Attention 计算的时空复杂度
不同 Transformer 结构的复杂度

Algorithm

原始 `Transformer` 使用的 `Multi-HEAD Attention`

$head_i = Attention(QW_i^Q,KW_i^K,VW_i^V)=softmax[\frac{QW_i^Q(KW_i^K)^T}{\sqrt d_k}]VW_i^V$

其中： $K,Q,V \in \mathbb R^{n\times d_m} \ \ \ W_i^Q,W_i^K\in \mathbb R^{d_m\times d_k}$

所以： $softmax[\frac{QW_i^Q(KW_i^K)^T}{\sqrt d_k}] \in \mathbb {R} ^{n\times n}$ ，n 表示序列长度，所以原始 Transformer 使用的 Multi-HEAD Attention 的时空复杂度为 $O(n^2)$

`Linformer` 对 `Multi-HEAD Attention` 的修改

将 $KW_i^K, VW_i^V\in \mathbb R^{n\times d_k}$ 投影到 $E_iKW_i^K, F_iVW_i^V\in \mathbb R^{k\times d_k}$ ，其中 k 是一个常数，时空复杂度变成了 $O(n)$ ，其中，E、F 都是可学习的投影矩阵， $E,F \in \mathbb R^{k\times n}$
$\bar{head_i} = Attention(QW_i^Q,E_iKW_i^K, F_iVW_i^V)=softmax[\frac{QW_i^Q(E_iKW_i^K)^T}{\sqrt d_k}]F_iVW_i^V$
投影矩阵 E、F 可共享参数，分为：
- Headwise sharing： $E_i=E,\ \ F_i=F, \ \ for\ each\ layer$
- Key-value sharing： $E_i=E = F_i, \ \ for\ each\ layer$
- Layerwise sharing： $E, F, \ \ layer\ sharing$

理论依据与结果

特征值的长尾分布

效果（与 BERT-base 对比）

Thoughts

文中提到不使用奇异值分解来得到低秩矩阵的原因是：奇异值分解会引入额外的计算量，并且无法共享参数
代码被打包为了 linformer 的 pip 包，可以在 torch 框架下直接使用

URL