Vision Transformer之LambdaNetworks

原始文档： https://www.yuque.com/lart/papers/lhnua2

前言

这个作者太能写了，三十页！。。。

ICLR 2021的一篇文章。提出了一种Self-Attention的备选方案，用来捕获输入和结构化上下文信息之间的长距离交互，例如像素和周围的其他的像素，同时进一步降低attention计算的存储需求。
提出了一种lambda layer，这可以通过将需要关注的上下文内容转化为针对当前位置的的一种线性函数/变换（这被称为lambdas）以独立的作用于每个输入。而且这里的lambda layer用途广泛，可以实现以在全局，局部或蒙版上下文中对基于内容的交互和基于位置的交互进行建模。
生成的神经网络LambdaNetworks计算效率高，可以以较小的内存成本对远距离依赖性进行建模，因此可以应用于大型结构化输入（例如高分辨率图像）。

---

• 基于内容的交互(content-based interaction)，考虑上下文的内容，但忽略query位置和上下文(context)之间的关系(例如，两个像素之间的相对距离)。
• 基于位置的交互(position-based interaction)，考虑query位置和上下文(context)位置之间的关系。

相关概念

• 在本文的语境中，为了与self-attention对比和便于阅读，沿用了许多类似的概念，包括query-key-value。但是不同的是，这里实际上更加强调的是query和context的概念（但是key和value可以由context转化得到）。
  • 对于query和context的定义中，不光包含自身的content($q_n\&c_m$)，还包含着对应的position($n\&m$)，即，二者之间的交互既有内容的交互，亦有位置的交互。
  • 对于query的位置索引使用字母n，而context则使用m，二者亦可表示各自的长度，所以$(n,m)$可以用来索引任意结构化元素之间的成对的位置关系，例如像素之间的相对距离或者是图节点之间的边。
• 不论是什么的注意力的计算过程，最终的目的就是将输入的query($q_n$)映射到输出上$y_n$，文本也不例外。
  • 本文构造的函数关系可以简单表示为：$F:((q_n,n),\mathcal{C})\mapsto y_n\in {\mathbb{R}}^{|v|}$，即在给定context( C = { ( c m , m ) } \mathcal{C} = \{ (c_m, m) \} C={(cm​,m)})时，将包含了query的content和position信息整合后得到输出。这和神经网络中的一个层的效果类似，将输入变换后得到输出。
  • 在这个变换过程中，将$(q_n,c_m)$交互的形式成为content-based，而将$(q_n,(n,m))$交互的形式称为position-based。即前者仅考虑context的content，后者则考虑query和context之间的相对位置关系。
  • 说起位置信息，之前的self-attention方法中往往会引入位置编码。他们有时会将绝对位置信息直接添加到query（或context元素）中，但在作者看来，这种类型的交互是实际上仍然是基于内容的，因为它忽略了查询的位置和上下文元素位置之间的关系$(n,m)$ 。这里其实非常像相对位置编码。虽然本文在此处也提到了视觉领域中，绝对位置编码[Attention is all you need]效果有限，而像素之间的相对位置信息（相对位置编码）[Attention augmented convolutional networks, Self-attention with relative position representations]却很重要。本文虽然也是相对位置编码，但是使用方式和之前的并不相同。之前的思路仍然局限于Self-Attention的框架中。

---

这里在具体形式上进行比较下：（这里softmax中都忽略了放缩因子，另外其中的一些关于形状和复杂度计算包含一定的我的主观推测，不一定完全贴切）

• Softmax Attention，即最原始的self-attention形式，更多强调内容上的交互，缺少位置交互
  • $Y=softmax(Q{K}^{\top })V$
    • $Q\in {\mathbb{R}}^{b\times n\times h\times k},K\in {\mathbb{R}}^{b\times m\times h\times k},V\in {\mathbb{R}}^{b\times m\times h\times d/h}\to$
      • ${\Theta }_T(bnmhk+bhnmd/h)\to {\Theta }_T(bnm(hk+d))$
      • ${\Theta }_S(bhnm)$
        • 因为需要保存多头attention map作为中间结果
  • $y_n=V^{\top }softmax(K{q}_n)$
• Relative Attention[Self-attention with relative position representations]，进一步补充位置交互
  • $Y=softmax(Q{K}^{\top }+Q{E}^{\top })V$
    • $Q\in {\mathbb{R}}^{b\times n\times h\times k},K\in {\mathbb{R}}^{b\times m\times h\times k},E\in {\mathbb{R}}^{n\times m\times h\times k},V\in {\mathbb{R}}^{b\times m\times h\times d/h}\to$
      • ${\Theta }_T(bnmhk+bhnmd/h)\to {\Theta }_T(bnm(hk+d))$
      • ${\Theta }_S(bhnm)$
        • 因为需要保存多头attention map和$Q{E}^{\top }$作为中间结果
  • $y_n=V^{\top }softmax((K+E_n)q_n)$
• Linear Attention，尝试降低计算需求，在线性空间和时间复杂度上近似attention操作（后面的相关工作中提供了一些介绍）
  • $Y=\varphi (Q)(\varphi (K{)}^{\top }V)\sim softmax(Q{K}^{\top })V$
    • 这里先计算后两项可以实现关于输入长度$n$的线性复杂度，一般有$\varphi (x)=elu(x)+1$或者$\varphi (x)=softmax(x)$
    • $Q\in {\mathbb{R}}^{b\times n\times k},K\in {\mathbb{R}}^{b\times m\times k},V\in {\mathbb{R}}^{b\times m\times d}\to$
      • ${\Theta }_T(bnkd+bkmd)\to {\Theta }_T(bnkd)$
        • 这里由于m和n在复杂度式子上是相加的关系，所以是可以省略掉m的，因为一般m和n相差不大
      • ${\Theta }_S(bkd)$
        • 只需要保存后两项的结果
  • $y_n=V^{\top }\varphi (K)\varphi (q_n)\sim V^{\top }softmax(K{q}_n)$
• Linear Attention近似relative attention时
  • $Y=\varphi (Q)(\varphi (K+E{)}^{\top }V)\sim softmax(Q(K+E{)}^{\top })V$
    • $Q\in {\mathbb{R}}^{b\times n\times k},K\in {\mathbb{R}}^{b\times m\times k},E\in {\mathbb{R}}^{n\times m\times k},V\in {\mathbb{R}}^{b\times m\times d}\to$
      • ${\Theta }_T(bmkd+bnmkd+bnkd)\to {\Theta }_T(bnmkd)$
        • 这里由于相对编码的存在导致运算次数增加
      • ${\Theta }_S(bkd+bnkd)\to {\Theta }_S(bnkd)$
        • 这里由于相对编码的存在导致存储空间增加
  • $y_n=V^{\top }\varphi (K+E_n)\varphi (q_n)\sim V^{\top }softmax((K+E_n)q_n)$
• 本文的Lambda Layer不再刻意去拟合一个原始的attention kernel（即softmax部分）
  • $Y=Q(\bar{K}+E{)}^{\top }V$
    • $Q\in {\mathbb{R}}^{b\times h\times n\times k},K\in {\mathbb{R}}^{b\times m\times k},E\in {\mathbb{R}}^{n\times m\times k},V\in {\mathbb{R}}^{b\times m\times d/h}\to$
      • 对于mutli-query，则h个query共享相同的${\lambda }_n\in {\mathbb{R}}^{b\times k\times d/h}$
      • ${\Theta }_T(bmkd/h+bnmkd/h+bnhkd/h+bnhkd/h)\to {\Theta }_T(bnmkd/h)$
      • ${\Theta }_S(\underset{E}{\underbrace{nmk}}+\underset{otheritems}{\underbrace{\cdots }}+\underset{{\lambda }^c={\bar{K}}^{\top }V}{\underbrace{bkd/h}}+\underset{{\lambda }^p=E^{\top }V}{\underbrace{bnkd/h}}+\underset{Q{\lambda }^c}{\underbrace{bnhd/h}}+\underset{Q{\lambda }^p}{\underbrace{bnhd/h}})\to \Theta (nmk+bndk/h)$
      • 如果基于卷积实现的lambda convolution，设$r$为局部区域元素数量，则这里的位置嵌入会进一步简化$\Theta (kr+bnkd/h)$
  • $y_n=V^T(\bar{K}+E_n)q_n$

虽然attention kernel的近似是理论上的动机，但文章认为它们可能是不必要的限制。
例如，Linear Attention中的核逼近在$Q\&K$上都需要相同的特征函数，因此无法使用更灵活的非线性和归一化方案。
相反，lambda layer 不会尝试近似attention kernel，这简化了它们的设计，并允许使用更加灵活的非线性和标准化方案（消融试验中有所体现）。
考虑到与键无关的位置嵌入，尤其可以通过设计的lambda convolution来实现简单有效的针对局部上下文的实现，即由于近似相对注意力内核将需要标准化键和位置嵌入（即$\varphi (K+E_n)$而不是$\varphi (K)+E_n$），这在基于卷积的局部上下文情况下无法实现。

---

• 为了便于query和context的交互，需要引入一个中间的依赖于context的量，这里称之为key$k_m$。而基于位置的交互则需要引入相对位置嵌入$e_{nm}$，二者都是k维的矢量。在这里k可以看作是query和key的通道数（亦或是多头注意力中的单个头的通道数）。由于$|k|$和context空间维度$|m|$并没有体现在输出$y_n\in {\mathbb{R}}^{|v|}$中，所以这些尺寸需要进行收缩，这些操作都会作为层计算的一部分。Every layer capturing long-range interactions can therefore be characterized based on whether it contracts the query depth or the context positions first.（也因此，每个捕获远程交互的层都可以根据它是先缩小查询深度还是先缩小上下文位置来描述，这说的有点绕）。
• 传统的基于attention的交互形式，压缩query的深度需要先在query和context元素之间创建一个相似核（即注意力图）。由于context元素的数量变得更大，输入和输出的维度$|k|\&|v|$维持不变，当给定层输出是一个具有相对较小维度的向量（$|v|\ll |m|$）时，可能会假设计算注意力图会变得浪费。所以在基于lambda的交互中，对于某些线性函数$\lambda (\mathcal{C},n):{\mathbb{R}}^{|k|}\mapsto {\mathbb{R}}^{|v|}$来说，简单的按照$y_n=F()$映射query到输出可能会更有效。这种场景下，context被集成到一个固定大小的线性函数${\lambda }_n=\lambda (\mathcal{C},n)$。每个${\lambda }_n$可以看做一个独立于context（仅计算一次）而存在的小的线性函数，在被应用到与之相关的$q_n$上后，就会被舍弃。
• 这里的lambda的概念实际上来自于函数式编程中的λ-calculus（lambda演算）。

Lambda Layers

lambda layers

• 层的计算涉及到如下两个输入，但是实际上亦可使用同一个，即类似于Self-Attention的情况
  • 输入$X\in {\mathbb{R}}^{|n|\times d_{in}}$，用来生成query
  • context$C\in {\mathbb{R}}^{|m|\times d_c}$，用来生成线性函数lambdas，作用于queries
• 输出$Y\in {\mathbb{R}}^{|n|\times d_{out}}$
• lambda层主要包含两部分计算
  • 生成基于上下文的lambda函数。其中包含了两部分，content lambda${\lambda }^c$和position lambda${\lambda }_n^p$。这里的key $K$和value $V$通过对context $C$线性投影获得，然后，被softmax操作归一化处理的key $\bar{K}$和对应于$q_n$的相对位置嵌入$E_n$一同来集成value $V$，从而可以获得最终的lambda函数。
    • 可以观察到，这里的 content lambda 实际上在所有的n是共享的，与n无关，这对于context element的置换是不变的，编码了如何基于context content来单独地转换query。
    • 而 position lambda 则是通过位置嵌入从而依赖于query的位置n，这编码了如何基于context元素$c_m$（有$V$来提供context的信息）和他们的相对位置关系$(n,m)$来转换query。
  • 将lambda应用于对应的query上。这里的query$q_n\in {\mathbb{R}}^{|k|}$通过$x_n$投影所得。
• 整个lambda层，通过在lambda函数中基于context’s content(content-based interactions)和structure(position-based interactions)来整合上下文特征，并基于query动态的分配这些特征产生最终的输出。不再需要产生attention maps了。
• 在实验中验证了，在两部分（生成lambda和应用lambda）计算中加入非线性和归一化操作的有效性。

multi-query形式

即面对一个batch包含的$|b|$个样本，每个样本内包含$|n|$个输入的情况，即对应有bn个query元素。
此时的计算规则从代码来看最清楚：

multi-query的实现

对应于原始的multi-head的实现

这里相对于self-attention中的multi-head，提出了一个对应的mutli-query，也就是说对query进行分组，每h个query共享一个lambda${\lambda }_n\in {\mathbb{R}}^{|k|\times |h|/|d|}$，而对应的有$q_n\in {\mathbb{R}}^{|hk|}$。

需要注意与self-attention中的multi-head的差异：

• Using multiple queries in the attention operation increases representational power and complexity.
• In contrast, using multiple queries in the lambda layer decreases complexity and representational power (ignoring the additional queries).

lambda convolution

尽管远距离交互有很多好处，但在许多任务中，局部性仍然是一个强烈的归纳偏置。使用全局上下文可能会导致噪声或计算量过大。 因此，将位置交互作用的范围限制为query位置周围的局部邻域可能是有用的。这可以通过将上下文位置m在所需范围之外的相对位置嵌入置零来完成。但是，对于较大的$|m|$值，此策略仍然消耗巨大，因为仍会进行计算，它们只是被清零。在将上下文排列在多维网格中的情况下，我们可以使用常规卷积等效地从局部上下文中计算位置lambda。我们称此操作为lambda convolution。这可以使用具有通道数整数倍的n-d深度分离卷积或将V中的v维视为额外空间维度的**(n+1)-d卷积**来实现n维lambda卷积。

由于现在将计算限制在局部范围内，因此lambda卷积获得了相对于输入长度的线性时间和存储复杂度。lambda卷积易于与其他卷积特性配合使用，例如扩张和跨步，并在专用硬件加速器上享有优化的实现方式。这与局部self-attention的实现形成鲜明对比，后者需要实现重叠query和context block的特征patch，从而增加了内存消耗和延迟。

intra-depth

文章简要地尝试了一个变体，该变体能够增加计算lambda的成本，同时保持应用它们的成本不变。这通过在key，位置嵌入和value中引入一个附加维度（称为具有相应超参数$|u|$的内部深度）来实现的。
现在，每个键（或位置嵌入）都是$|k|\times |u|$矩阵，而不是$|k|$维向量。类似地，每个值现在是$|v|\times |u|$矩阵，而不是$|v|$维向量。
lambdas通过对上下文位置和intra-depth position$|u|$求和而得，并具有类似于默认情况的$|k|\times |v|$的形状。
实验表明，此变体可提高准确性，但我们发现在现代机器学习加速器上控制速度时，使用$|u|=1$（即默认情况）是最佳选择。

相关工作

Linear Attention

本文工作和linear attention比较像类似，提出的lambda layers也是在尝试绕过计算attention map，但是不同的是：

• 提出的lambda layers在模拟基于上下文和位置的交互，这也确保了应用在大的结构化输入（如图像）的能力。
• 而现有的Linear attention虽然提供了一种针对高昂存储需求的可扩展的补救措施，但未能对内部数据结构建模，例如像素之间的相对距离或图中节点之间的边缘关系。

所以本文主要是针对视觉任务，希望实现有效模拟query和结构化上下文元素之间的长距离交互，同时有着更低的存储需求。

---

可以从科学空间的几篇高质量博文了解下linear attention：

• 苏剑林. (Jul. 04, 2020). 《线性Attention的探索：Attention必须有个Softmax吗？ 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/7546
• 苏剑林. (Dec. 01, 2020). 《Performer：用随机投影将Attention的复杂度线性化 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/7921
• 苏剑林. (Feb. 16, 2021). 《Nyströmformer：基于矩阵分解的线性化Attention方案 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/8180

读者也许想不到，制约Attention性能的关键因素，其实是定义里边的Softmax！事实上，简单地推导一下就可以得到这个结论。$Q{K}^{\top }$这一步我们得到一个n×n的矩阵，就是这一步决定了Attention的复杂度是$O(n^2)$；如果没有Softmax，那么就是三个矩阵连乘$Q{K}^{\top }V$，而矩阵乘法是满足结合率的，所以我们可以先算$K^{\top }V$，得到一个$d\times d$的矩阵，然后再用Q左乘它，由于$d\ll n$，所以这样算大致的复杂度只是$O(n)$（就是Q左乘那一步占主导）。
也就是说，去掉Softmax的Attention的复杂度可以降到最理想的线性级别$O(n)$！这显然就是我们的终极追求：Linear Attention，复杂度为线性级别的Attention。
https://spaces.ac.cn/archives/7546#%E6%91%98%E6%8E%89Softmax

---

Lambda层通用化并扩展了线性注意公式，并进一步捕获基于位置的交互作用，这对于建模高度结构化的输入（例如图像）至关重要。

由于其目的不是近似注意内核，因此lambda层允许更灵活的非线性和规范化，实验也发现这是有益的。

最后，与线性注意工作中通常使用的多头（或单头）公式相比，提出了multi-query lambda layer作为降低复杂度的一种方法。

Self Attention

二者的差异

• 自注意力定义query和上下文元素之间的相似核
• lambda layer则是通过先归纳总结上下文信息到一个固定尺寸的线性函数（变换矩阵）中，因此可以绕过memory-intensive attention maps

Channel and Spatial Attention

关于这部分的解释确实很有意思，将attention各种线性类型都给统一到了lambda layer的形式下。
lambda抽象概括了通道和空间注意机制，这两种机制都可以被视为lambda层的具体实例。这个观察结果与实验是一致的，实验表明lambda层优于通道和空间注意力，同时计算效率更高。

• Channel attention mechanisms, such as Squeeze-and-Excitation (SE) and FiLM layers, recalibrate features via cross-channel interactions by aggregating signals from the entire feature map.
  • In particular, the SE operation can be written as $y_{nk}=w_n{q}_{nk}$, where $w_k$ is the excitation weight for channel $k$ in the query $q_n$. （即将每一个通道看做了一个单独的query）
  • This can be viewed as using a diagonal lambda which is shared across query positions ${\lambda }_n=diag(w_1,\cdots ,w_{|k|})$.
  • Channel attention mechanisms have proven useful to complement convolutions but cannot be used as a stand-alone layer as they discard spatial information. 只能用来修正卷积层的特征，而不能作为单独的层结构。
• Spatial attention mechanisms, reweight each position based on signals aggregated from all channels.
  • These mechanisms can be written as $y_{nk}=w_n{q}_{nk}$ where $w_n$ is the attention weight for position $n$ in the input query.
  • This can be viewed as using (position-dependent) scalar lambdas ${\lambda }_n=w_n\mathbb{I}$, where $\mathbb{I}$ is the identity matrix.
  • Spatial attention has also proven helpful to complement convolutions but cannot be used as a stand-alone layer as it discards channel information.

链接

• ICLR open review页面：https://openreview.net/forum?id=xTJEN-ggl1b
• 论文：https://openreview.net/pdf?id=xTJEN-ggl1b