image.png

前言

这个作者太能写了,三十页!。。。

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

相关概念

  • 在本文的语境中,为了与self-attention对比和便于阅读,沿用了许多类似的概念,包括query-key-value。但是不同的是,这里实际上更加强调的是query和context的概念(但是key和value可以由context转化得到)。
    • 对于query和context的定义中,不光包含自身的content(LambdaNetworks: Modeling long-range Interactions without Attention - 图2),还包含着对应的position(LambdaNetworks: Modeling long-range Interactions without Attention - 图3),即,二者之间的交互既有内容的交互,亦有位置的交互。
    • 对于query的位置索引使用字母n,而context则使用m,二者亦可表示各自的长度,所以LambdaNetworks: Modeling long-range Interactions without Attention - 图4可以用来索引任意结构化元素之间的成对的位置关系,例如像素之间的相对距离或者是图节点之间的边。
  • 不论是什么的注意力的计算过程,最终的目的就是将输入的query(LambdaNetworks: Modeling long-range Interactions without Attention - 图5)映射到输出上LambdaNetworks: Modeling long-range Interactions without Attention - 图6,文本也不例外。

    • 本文构造的函数关系可以简单表示为:LambdaNetworks: Modeling long-range Interactions without Attention - 图7,即在给定context(LambdaNetworks: Modeling long-range Interactions without Attention - 图8)时,将包含了query的content和position信息整合后得到输出。这和神经网络中的一个层的效果类似,将输入变换后得到输出。
    • 在这个变换过程中,将LambdaNetworks: Modeling long-range Interactions without Attention - 图9交互的形式成为content-based,而将LambdaNetworks: Modeling long-range Interactions without Attention - 图10交互的形式称为position-based。即前者仅考虑context的content,后者则考虑query和context之间的相对位置关系。
    • 说起位置信息,之前的self-attention方法中往往会引入位置编码。他们有时会将绝对位置信息直接添加到query(或context元素)中,但在作者看来,这种类型的交互是实际上仍然是基于内容的,因为它忽略了查询的位置和上下文元素位置之间的关系LambdaNetworks: Modeling long-range Interactions without Attention - 图11 。这里其实非常像相对位置编码。虽然本文在此处也提到了视觉领域中,绝对位置编码[Attention is all you need]效果有限,而像素之间的相对位置信息(相对位置编码)[Attention augmented convolutional networks, Self-attention with relative position representations]却很重要。本文虽然也是相对位置编码,但是使用方式和之前的并不相同。之前的思路仍然局限于Self-Attention的框架中。 :::info 这里在具体形式上进行比较下:(这里softmax中都忽略了放缩因子,另外其中的一些关于形状和复杂度计算包含一定的我的主管推测,不一定完全贴切)
      image.png
  • Softmax Attention,即最原始的self-attention形式,更多强调内容上的交互,缺少位置交互

    • LambdaNetworks: Modeling long-range Interactions without Attention - 图13
      • LambdaNetworks: Modeling long-range Interactions without Attention - 图14
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图15
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图16
          • 因为需要保存多头attention map作为中间结果
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图17
  • Relative Attention[Self-attention with relative position representations],进一步补充位置交互
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图18
      • LambdaNetworks: Modeling long-range Interactions without Attention - 图19
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图20
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图21
          • 因为需要保存多头attention map和LambdaNetworks: Modeling long-range Interactions without Attention - 图22作为中间结果
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图23
  • Linear Attention,尝试降低计算需求,在线性空间和时间复杂度上近似attention操作(后面的相关工作中提供了一些介绍)
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图24
      • 这里先计算后两项可以实现关于输入长度LambdaNetworks: Modeling long-range Interactions without Attention - 图25的线性复杂度,一般有LambdaNetworks: Modeling long-range Interactions without Attention - 图26或者LambdaNetworks: Modeling long-range Interactions without Attention - 图27
      • LambdaNetworks: Modeling long-range Interactions without Attention - 图28
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图29
          • 这里由于m和n在复杂度式子上是相加的关系,所以是可以省略掉m的,因为一般m和n相差不大
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图30
          • 只需要保存后两项的结果
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图31
  • Linear Attention近似relative attention时
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图32
      • LambdaNetworks: Modeling long-range Interactions without Attention - 图33
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图34
          • 这里由于相对编码的存在导致运算次数增加
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图35
          • 这里由于相对编码的存在导致存储空间增加
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图36
  • 本文的Lambda Layer不再刻意去拟合一个原始的attention kernel(即softmax部分)
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图37
      • LambdaNetworks: Modeling long-range Interactions without Attention - 图38
        • 对于mutli-query,则h个query共享相同的LambdaNetworks: Modeling long-range Interactions without Attention - 图39
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图40
        • LambdaNetworks: Modeling long-range Interactions without Attention - 图41
        • 如果基于卷积实现的lambda convolution,,设LambdaNetworks: Modeling long-range Interactions without Attention - 图42为局部区域元素数量,则这里的位置嵌入会进一步简化LambdaNetworks: Modeling long-range Interactions without Attention - 图43
    • LambdaNetworks: Modeling long-range Interactions without Attention - 图44


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

  • 为了便于query和context的交互,需要引入一个中间的依赖于context的量,这里称之为keyLambdaNetworks: Modeling long-range Interactions without Attention - 图48。而基于位置的交互则需要引入相对位置嵌入LambdaNetworks: Modeling long-range Interactions without Attention - 图49,二者都是k维的矢量。在这里k可以看作是query和key的通道数(亦或是多头注意力中的单个头的通道数)。由于LambdaNetworks: Modeling long-range Interactions without Attention - 图50和context空间维度LambdaNetworks: Modeling long-range Interactions without Attention - 图51并没有体现在输出LambdaNetworks: Modeling long-range Interactions without Attention - 图52中,所以这些尺寸需要进行收缩,这些操作都会作为层计算的一部分。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元素的数量变得更大,输入和输出的维度LambdaNetworks: Modeling long-range Interactions without Attention - 图53维持不变,当给定层输出是一个具有相对较小维度的向量(LambdaNetworks: Modeling long-range Interactions without Attention - 图54)时,可能会假设计算注意力图会变得浪费。所以在基于lambda的交互中,对于某些线性函数LambdaNetworks: Modeling long-range Interactions without Attention - 图55来说,简单的按照LambdaNetworks: Modeling long-range Interactions without Attention - 图56映射query到输出可能会更有效。这种场景下,context被集成到一个固定大小的线性函数LambdaNetworks: Modeling long-range Interactions without Attention - 图57。每个LambdaNetworks: Modeling long-range Interactions without Attention - 图58可以看做一个独立于context(仅计算一次)而存在的小的线性函数,在被应用到与之相关的LambdaNetworks: Modeling long-range Interactions without Attention - 图59上后,就会被舍弃。
  • 这里的lambda的概念实际上来自于函数式编程中的λ-calculus(lambda演算)。

    Lambda Layers

    image.png

    lambda layers

  • 层的计算涉及到如下两个输入,但是实际上亦可使用同一个,即类似于Self-Attention的情况

    • 输入LambdaNetworks: Modeling long-range Interactions without Attention - 图61,用来生成query
    • contextLambdaNetworks: Modeling long-range Interactions without Attention - 图62,用来生成线性函数lambdas,作用于queries
  • 输出LambdaNetworks: Modeling long-range Interactions without Attention - 图63
  • lambda层主要包含两部分计算

image.png

  • 生成基于上下文的lambda函数。其中包含了两部分,content lambdaLambdaNetworks: Modeling long-range Interactions without Attention - 图65和position lambdaLambdaNetworks: Modeling long-range Interactions without Attention - 图66。这里的key LambdaNetworks: Modeling long-range Interactions without Attention - 图67和value LambdaNetworks: Modeling long-range Interactions without Attention - 图68通过对context LambdaNetworks: Modeling long-range Interactions without Attention - 图69线性投影获得,然后,被softmax操作归一化处理的key LambdaNetworks: Modeling long-range Interactions without Attention - 图70和对应于LambdaNetworks: Modeling long-range Interactions without Attention - 图71的相对位置嵌入LambdaNetworks: Modeling long-range Interactions without Attention - 图72一同来集成value LambdaNetworks: Modeling long-range Interactions without Attention - 图73,从而可以获得最终的lambda函数。
    • 可以观察到,这里的content lambda实际上在所有的n是共享的,与n无关,这对于context element的置换是不变的,编码了如何基于context content来单独地转换query。
    • position lambda则是通过位置嵌入从而依赖于query的位置n,这编码了如何基于context元素LambdaNetworks: Modeling long-range Interactions without Attention - 图74(有LambdaNetworks: Modeling long-range Interactions without Attention - 图75来提供context的信息)和他们的相对位置关系LambdaNetworks: Modeling long-range Interactions without Attention - 图76来转换query。

image.png

  • 将lambda应用于对应的query上。这里的queryLambdaNetworks: Modeling long-range Interactions without Attention - 图78通过LambdaNetworks: Modeling long-range Interactions without Attention - 图79投影所得。
    • 整个lambda层,通过在lambda函数中基于context’s content(content-based interactions)和structure(position-based interactions)来整合上下文特征,并基于query动态的分配这些特征产生最终的输出。不再需要产生attention maps了。
    • 在实验中验证了,在两部分(生成lambda和应用lambda)计算中加入非线性和归一化操作的有效性。

      multi-query形式

      即面对一个batch包含的LambdaNetworks: Modeling long-range Interactions without Attention - 图80个样本,每个样本内包含LambdaNetworks: Modeling long-range Interactions without Attention - 图81个输入的情况,即对应有bn个query元素。
      此时的计算规则从代码来看最清楚:
      image.png
      multi-query的实现
      image.png
      对应于原始的multi-head的实现
      这里相对于self-attention中的multi-head,提出了一个对应的mutli-query,也就是说对query进行分组,每h个query共享一个lambdaLambdaNetworks: Modeling long-range Interactions without Attention - 图84,而对应的有LambdaNetworks: Modeling long-range Interactions without Attention - 图85
      需要注意与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).

image.png

lambda convolution

尽管远距离交互有很多好处,但在许多任务中,局部性仍然是一个强烈的归纳偏置。使用全局上下文可能会导致噪声或计算量过大。 因此,将位置交互作用的范围限制为query位置周围的局部邻域可能是有用的。这可以通过将上下文位置m在所需范围之外的相对位置嵌入置零来完成。但是,对于较大的LambdaNetworks: Modeling long-range Interactions without Attention - 图87值,此策略仍然消耗巨大,因为仍会进行计算,它们只是被清零。在将上下文排列在多维网格中的情况下,我们可以使用常规卷积等效地从局部上下文中计算位置lambda。我们称此操作为lambda convolution。这可以使用具有通道数整数倍的n-d深度分离卷积或将V中的v维视为额外空间维度的(n+1)-d卷积来实现n维lambda卷积。
由于现在将计算限制在局部范围内,因此lambda卷积获得了相对于输入长度的线性时间和存储复杂度。lambda卷积易于与其他卷积特性配合使用,例如扩张和跨步,并在专用硬件加速器上享有优化的实现方式。这与局部self-attention的实现形成鲜明对比,后者需要实现重叠query和context block的特征patch,从而增加了内存消耗和延迟。
image.png
image.png

intra-depth

文章简要地尝试了一个变体,该变体能够增加计算lambda的成本,同时保持应用它们的成本不变。这通过在key,位置嵌入和value中引入一个附加维度(称为具有相应超参数LambdaNetworks: Modeling long-range Interactions without Attention - 图90的内部深度)来实现的。
现在,每个键(或位置嵌入)都是LambdaNetworks: Modeling long-range Interactions without Attention - 图91矩阵,而不是LambdaNetworks: Modeling long-range Interactions without Attention - 图92维向量。类似地,每个值现在是LambdaNetworks: Modeling long-range Interactions without Attention - 图93矩阵,而不是LambdaNetworks: Modeling long-range Interactions without Attention - 图94维向量。
lambdas通过对上下文位置和intra-depth positionLambdaNetworks: Modeling long-range Interactions without Attention - 图95求和而得,并具有类似于默认情况的LambdaNetworks: Modeling long-range Interactions without Attention - 图96的形状。
实验表明,此变体可提高准确性,但我们发现在现代机器学习加速器上控制速度时,使用LambdaNetworks: Modeling long-range Interactions without Attention - 图97(即默认情况)是最佳选择。
image.png
image.png
image.png

相关工作

Linear Attention

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

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

所以本文主要是针对视觉任务,希望实现有效模拟query和结构化上下文元素之间的长距离交互,同时有着更低的存储需求。 :::info 可以从科学空间的几篇高质量博文了解下linear attention:

读者也许想不到,制约Attention性能的关键因素,其实是定义里边的Softmax!事实上,简单地推导一下就可以得到这个结论。QK⊤这一步我们得到一个n×n的矩阵,就是这一步决定了Attention的复杂度是O(n^2);如果没有Softmax,那么就是三个矩阵连乘QK⊤V,而矩阵乘法是满足结合率的,所以我们可以先算K⊤V,得到一个d×d的矩阵,然后再用Q左乘它,由于d≪n,所以这样算大致的复杂度只是O(n)(就是Q左乘那一步占主导)。
也就是说,去掉Softmax的Attention的复杂度可以降到最理想的线性级别O(n)!这显然就是我们的终极追求:Linear Attention,复杂度为线性级别的Attention。
https://spaces.ac.cn/archives/7546#%E6%91%98%E6%8E%89Softmax ::: Lambda层通用化并扩展了线性注意公式,并进一步捕获基于位置的交互作用,这对于建模高度结构化的输入(例如图像)至关重要。
image.png
由于其目的不是近似注意内核,因此lambda层允许更灵活的非线性和规范化,实验也发现这是有益的。
image.png
最后,与线性注意工作中通常使用的多头(或单头)公式相比,提出了multi-query lambda layer作为降低复杂度的一种方法。

Self Attention

image.png
二者的差异

  • 自注意力定义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 LambdaNetworks: Modeling long-range Interactions without Attention - 图104, where LambdaNetworks: Modeling long-range Interactions without Attention - 图105 is the excitation weight for channel LambdaNetworks: Modeling long-range Interactions without Attention - 图106 in the query LambdaNetworks: Modeling long-range Interactions without Attention - 图107. (即将每一个通道看做了一个单独的query)
    • This can be viewed as using a diagonal lambda which is shared across query positions LambdaNetworks: Modeling long-range Interactions without Attention - 图108.
    • 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 LambdaNetworks: Modeling long-range Interactions without Attention - 图109 where LambdaNetworks: Modeling long-range Interactions without Attention - 图110 is the attention weight for position LambdaNetworks: Modeling long-range Interactions without Attention - 图111 in the input query.
    • This can be viewed as using (position-dependent) scalar lambdas LambdaNetworks: Modeling long-range Interactions without Attention - 图112, where LambdaNetworks: Modeling long-range Interactions without Attention - 图113 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