文献阅读：Linformer: Self-Attention with Linear Complexity

最新推荐文章于 2025-03-31 17:13:10 发布

Espresso Macchiato

最新推荐文章于 2025-03-31 17:13:10 发布

阅读量5.6k

点赞数 5

分类专栏：文献阅读文章标签： LInformer Transformer优化长句计算 attention优化效率优化

本文链接：https://blog.csdn.net/codename_cys/article/details/124233661

版权

文献阅读专栏收录该内容

55 篇文章

订阅专栏

文献阅读：Linformer: Self-Attention with Linear Complexity

文献链接：https://arxiv.org/pdf/2006.04768.pdf

1. 问题描述

这篇文章同样是我在阅读Transformer Quality in Linear Time这篇文章时想到的一个工作，所以就来考个古，把这篇文章也翻出来整理一下，算是给自己做个笔记了。

Linformer这篇工作是20年facebook提出来的一个工作，目的是优化transformer的计算量，尤其是针对长句计算希望可以减少计算的复杂度，同时尽可能地保持模型的效果。

众所周知，self-attention layer的计算复杂度是和句长成正比关系的，这就导致transformer在针对长句的计算上面尤其耗时，几乎不具备可用性。

当然，针对这方面已经有了不少优化，比如从算子算力方面的半精度模型优化，或者另辟蹊径用蒸馏方式来直接用小模型替代大模型。

而对于模型本身的优化，也同样有sparse transformer以及reformer这样的工作在前。Linformer算是另一种针对模型的attention结构本身进行优化的一种方式。

他的核心思路就是直接先将attention投影到一个低阶矩阵，从而在计算attention时避免掉 $O(n^2)$ 的计算量，而是转换成 $O(k\cdot n)$ 的计算量（ $k$ 是一个事先确定的常数），从而使得模型存在对长句的编码能力。

我们摘取文中对当时所有方法以及对应的复杂度总结表格如下：

在这里插入图片描述

2. 核心方法

1. vanilla attention layer

在介绍Linformer的核心结构之前，我们首先回顾一下基础的transformer的attention layer的结构。

为了简化问题，我们这里暂时不考虑多头的情况。

我们可以直接写出attention层的表达公式如下：

$softmax(\frac{QW_Q \cdot (KW_K)^T}{\sqrt{d}})VW_V$

针对self-attention的情况，我们假设句长为 $n$ ，embedding维度为 $d$ ，则 $\in \mathbb{R}^{n\times d}$ ， $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ 。

因此，当 $\gg d$ 时，上述的attention层的算法复杂度就是 $O(n^2)$ 。

2. attention优化

Linformer的核心方法其实还是比较trivial的，本质上就是将attention投影到一个低阶矩阵当中，从而规避掉 $O(n^2)$ 的attention计算。

具体而言，我们在attention计算中增加两个矩阵E和F，使得n阶的attention降维到固定的k维矩阵：

$softmax(\frac{QW_Q \cdot (E \cdot KW_K)^T}{\sqrt{d}})(F \cdot VW_V)$

其中， $\in \mathbb{R}^{k \times n}$ ， $k$ 是一个常数。

由此，我们就可以将 $O(n^2)$ 的计算复杂度降维到 $O (k n)$ ，其中 $k$ 不过是一个常数而已。

更进一步的，为了更进一步的缩小参数量，文中还尝试了在不同的层之间share相同的 $E, F$ ，甚至干脆令 $E = F$ 。

令人惊讶的是，这样同样可以得到一个还过得去的效果，简直震惊。

3. 分析 & 证明

有了上面这个看似粗暴的手法，下面，我们来看看上述方法是否在逻辑上合理。

1. self-attention是低阶的

首先，上述attention投影的基础在于一个现象，即：

self-attention的权重矩阵是一个低阶矩阵。

也就是说，如果我们跑去计算一下权重矩阵的本征值，我们可以发现，得到的大部分本征值都是0或者接近于0的。

文中给出了一个图表来对这个现象进行了展示：

在这里插入图片描述

可以看到，尾部有相当一部分的本征值事实上是接近于0的。

更数学化的，文中还给出了一个定理以及其对应的推导，不过这里就不过多展开了，就只把相应的定理摘录如下：

Theorem 1. (self-attention is low rank)
For any $\in \mathbb{R}^{n \times d}$ , and $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ , for any column vector $\in \mathbb{R}^n$ of matrix $VW_V$ , there exists a low-rank matrix $\tilde{P} \in \mathbb{R}^{n\times n}$ such that
$Pr(||\tilde{P}w^T - Pw^T|| < \epsilon||Pw^T||) > 1 - o(1)$
and $rank(\tilde{P}) = \Theta(log(n))$ .
where the context mapping matrix $P$ is defined as:
$softmax(\frac{QW_Q \cdot (KW_K)^T}{\sqrt{d}})$

2. linear self-attention效果与vanilla self-attention相仿

有了上述定理，那么，我们总可以找到一个低阶的attention权重矩阵 $\tilde{P}$ 来替换掉原始的 $P$ ，但是，我们如何来寻找这个低阶矩阵呢？

或者说，我们在上一个小节当中已经给出的那个暴力的降维方案，是否真的可以达到相仿的效果呢？

同样的，文中依然给出了一个数学证明，不过这里同样还是不具体展开了，只是摘录文中的定理如下：

Theorem 2. (Linear self-attention)
For any $\in \mathbb{R}^{n \times d}$ and $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ ,
if $min{\Theta(9d log(d))/ \epsilon^2, 5 \Theta(log(n)/\epsilon^2)}$ , then there exists matrices $\in \mathbb{R}^{n \times k}$ such that,
for any row vector $w$ of matrix $QW_Q(KW_K)^T/\sqrt{d}$ , we have:
$Pr(||softmax(wE^T)FVW_V - softmax(w)VW_V|| \leq \epsilon||softmax(w)|| \cdot ||VW_V||) > 1 - o(1)$