Lambda Networks（★★★★★）

LambdaNetworks: Modeling Long-Range Interactions Without Attention

https://arxiv.org/abs/2102.08602

https://openreview.net/forum?id=xTJEN-ggl1b

https://hub.fastgit.org/d-li14/lambda.pytorch （这个是SENet作者的开源）

https://hub.fastgit.org/lucidrains/lambda-networks

https://hub.fastgit.org/leaderj1001/LambdaNetworks （感觉这个挺好）

17 February, 2021 ICLR2021 Google Brain

Irwan Bello

也是Attention Augmented Convolution和ResNetRS的一作，Stand Alone和GSA的作者之一

摘要：我们提出了lambda层——自我注意的另一个框架——用于捕捉输入和结构化上下文信息（例如，被其他像素包围的像素）之间的远程交互。Lambda层通过将能够获得的上下文元素用线性函数做变换，并将这种线性函数分别用于每个输入，以捕获这种交互。与线性注意力相似的是，lambda层没有做计算量大的注意力映射；与线性注意力不同的是，lambda层同时对基于内容和基于位置的交互进行建模，这使得lambda层可用于大型结构化输入，如图像。由lambda layer构成的网络lambda networks在ImageNet分类、COCO检测、COCO分割上，显著优于基于卷积和注意力的方法，同时计算得更快。另外，我们还设计了一系列不同大小的混合结构的LambdaResNets，大大提高图像分类的速度-精度的平衡。LambdaResNets在ImageNet上达到的精度，在现代机器学习加速器上比流行的EfficientNets快3.2-4.4倍。当使用额外的130M的伪标签数据进行训练时，LambdaResNets的速度是相应的EfficientNet的checkpoints的9.5倍。

1.引言

对数据中的长期依赖关系进行建模是机器学习中的一个核心问题。自注意力(Bahdanau et al.，2015；Vaswani et al.，2017)已经成为一种通用方法，但是自注意力需要占用很多的内存，这导致其很难用在长序列和多维数据上，例如图像。煮个栗子，将一个注意力头数为8的多头注意力，用在批量大小为128、分辨率为64x64的图像数据上，需要占用64GB的内存，这在实际中是没法用的。

线性注意力机制(Katharopoulos et al.，2020；Choromanski et al.，2020)为了解决这个问题提出一种补救办法，但是没有办法对数据的内部结构进行建模，例如像素之间的相对距离，以及一个图中节点之间的连边关系。

本文解决这两个问题。我们提出lambda层，其以较低的内存占用对一个query和一组结构化的上下文元素之间的远程交互进行建模。Lambda层将每个能获取的上下文元素变换为一个线性函数，称作lambda，然后将其直接应用在对应的query上。尽管自注意力在query和上下文元素之间定义了度量相似性的核，而lambda层则是将上下文信息汇总成一个固定大小的线性函数（即矩阵），从而绕过了占用很多内存的注意力映射的计算。它们的区别如图1所示。

图1：自注意力核lambda层的对比。左图：3个查询项和它们对应的在全局上下文内的局部上下文的示例。中间图：自注意力将每个查询项与其上下文元素的注意分布相关联。右图：lambda层将每个上下文元素变换为线性函数lambda，直接应用在对应的查询项上。

Lambda层是基于内容且基于位置的交互建模，上下文可以是全局的，也可以是局部的，也可以是带屏蔽的。由此产生的神经网络lambda networks，计算效率高，以较小的内存占用建立长程依赖模型，因此可以应用于大型结构化输入，如高分辨率图像。

我们评估了LambdaNetworks在计算机视觉任务中的应用，在这些任务中，使用自注意力会因为占用很多的内存而在应用上受到阻碍(Wang et al., 2018; Bello et al., 2019)，或者依赖于不切实际的实现(Ramachandran et al., 2019)，或者依赖于大量的数据(Dosovitskiy et al., 2020)。在我们的ImageNet分类、COCO检测、COCO实例分割的实验中，LambdaNetworks要显著优于卷积和注意力的同类网络，并且计算效率和速度也高于后者。

我们的贡献有：

• lambda层，这是一类层，对基于内容和基于位置的交互进行建模，不需要计算注意力映射。lambda层很容易用einsum操作和卷积核实现，这些操作在现代机器学习加速器上有着高效的实现。

• lambda层为通道注意力、空间注意力和线性注意力提供了一个统一的视角。我们的一些观察结果，例如使用多个查询项带来的计算好处，可以扩展到线性注意力上。

• lambda层在ImageNet分类上显著优于卷积和注意力的同类网络，并且计算更高效。例如，简单地用lambda层替换ResNet-50体系结构瓶颈块中的3x3卷积可产生+1.5%的top-1 精度提高，同时将参数减少40%。

• lambda层在延迟和内存占用上比很多的自注意力变体（局部注意力、轴向注意力）都有很大的计算优势。

• 我们研究了混合模型，发现混合模型可以最大化lambda netwoks的速度-精度平衡。

• 我们提出LambdaResNets，一种卷积和lambda混合的结构，基于Bello等人(2021)提出的训练和缩放策略。在ImageNet分类上速度是EfficientNet的4.4倍，且占用更少的内存。

• 在半监督学习设置下，使用额外的130M伪标签数据进行训练，LambdaResNets比EfficientNet noiseystudent checkpoints的速度提高了9.5倍（Xie et al.，2020）。

• 用MaskRCNN评估LambdaResNets在COCO检测和实例分割，在两个任务上都有一致的提升。

2 长程交互建模

本节正式定义查询项、上下文和交互。我们将keys作为捕获查询项和其上下文之间交互关系的依赖项，并表明，lambda层可作为能捕获长程交互的注意力机制的一种替用品。

符号我们用小写不加粗表示标量，用小写加粗表示向量，用大写加粗表示张量，例如 $n$ ， ${\bf{x} }$ ， ${\bf{X} }$ 。我们将 ${\bf{X} }$ 的第 $n$ 行表示为 ${\bf{x} }_{n}$ 。我们将 ${\bf{X} }$ 的元素表示为 $x_{ij}$ 。在可能的情况下，我们使用“自注意力”这个术语来简化表达、凸显差异。

查询项和上下文的定义 令 $Q =\left\{ ({\bf{ q }}_{n} , {n}) \right\}$ 和 $C=\left\{ ({\bf{c}}_{m} ,m) \right\}$ 表示向量的结构化集合。每个查询项query $({\bf{q}}_{n} ,n)$ 的特征是其内容 ${\bf{q}}_{n}\in R^{\vert k \vert }$ 与位置 $n$ 。类似的，每个上下文元素 $({\bf{c}}_{m} , {m})$ 的特征是其内容 ${\bf{c}}_{m}$ 与其在上下文中的位置 $m$ 。 $(n,m)$ 可指结构化元素之间的任何成对关系，例如图像中像素之间的相对距离、图中节点之间的连边。

交互的定义 我们考虑的是将一个查询项 $({\bf{q}}_{n} ,n)$ 在给定上下文 $C$ 及函数 $F:(({\bf{q}}_{n} ,{n}) ,C)\rightarrow {\bf{y}}_{n}$ ，映射到输出 ${\bf{y}}_{n}\in R^{\vert v \vert }$ 的一般问题。当处理结构化输入时，这样的函数可以是神经网络中的一个层。我们将交互 $({\bf{q}}_{n} , {\bf {c}}_{m} )$ 称为基于内容的交互， $({\bf{q}}_{n} ,(n,m))$ 称为基于位置的交互。我们注意到，尽管绝对位置信息有时是直接添加到查询项（或上下文元素）的内容上，我们认为这种类型的交互是基于内容的，因为它忽略了查询项和上下文元素位置之间的关系 $(n,m)$ 。使用绝对位置信息常被用于自然语言处理中，但很难用到图像上，因为在图像上，像素间的相对位置信息是关键的。

引入keys来捕获长程交互 在深度学习的背景下，我们优先考虑能够快速批量处理的线性运算，并使用点积操作作为我们的交互。这促使引入向量，这些向量可以通过点积操作与查询项交互，因此与查询项的维度相同。具体的是，基于内容的交互 $({\bf{q}}_{n} ,{\bf {c}}_{m} )$ 需要一个取决于 ${\bf{c}}_{m}$ 的 $k$ 维向量，一般地称作是键项 ${\bf {k}}_{m}$ 。相反，基于位置的交互 $({\bf{q}}_{n} ,(n,m))$ 需要一个相对位置嵌入 ${\bf {e}}_{mn}\in R^{\vert k \vert }$ 。由于查询项和键项的深度维度 $\vert k \vert$ 和上下文的空间维度 $\vert m\vert$ 不是在输出 ${\bf {y}}_{n} \in R^{\vert v \vert }$ 的维度空间里，这些维度需要作为层计算的一部分进行收缩。因此，捕获远程交互的每一层都可以根据它是先收缩查询深度维度还是先收缩上下文位置来描述。

注意力交互 收缩查询项的深度维度，首先是在查询项和上下文元素之间创建一个表示相似度的核（即注意力映射），这就是注意力操作。当上下文位置 $\vert m \vert$ 的数量增大，而输入和输出的维度 $\vert k \vert$ 和 $\vert v \vert$ 则不变，你可以认为这样计算注意力映射是很浪费计算资源的，因为层输出的向量的维度要小得多 $\vert v \vert <<\vert m \vert$ 。

Lambda交互 相反，将每个查询项以这样简单的方式映射到其输出会更高效： ${\bf {y}}_{n}=F(({\bf {q}}_{n} ,n) ,C)={\bf{\lambda }}(C，n)({\bf{q}}_{n})$ ，其中 $\lambda$ 是线性函数 ${\bf{\lambda }}(C，n):R^{\vert k \vert } \rightarrow R^{\vert v \vert }$ 。在这种情况下，上下文被聚合为一个固定规模的线性函数 ${\lambda}_{n}= {\bf{\lambda }}(C，n)$ 。每个 ${\lambda }_{n}$ 是一个小的（小规模？小计算量？）线性函数（这种机制让人想起函数式编程和lambda演算，后者激发了在本工作中起名为lambda），并且这个线性函数是独立于上下文（一旦计算）而存在，并在应用于其关联的查询项 ${\bf {q}}_{n}$ 后就会被丢弃。

3 Lambda层

3.1 Lambda层：将上下文变换为线性函数

lambda层将输入 ${\bf X}\in {R}^{{\vert n \vert } \times d_{in} }$ 和上下文 ${\bf C}\in {R}^{{\vert m \vert } \times d_{c} }$ 作为输入，生成线性函数lambdas，然后将其应用在查询项上，产生输出 ${\bf Y } \in {R}^{{\vert n \vert } \times d_{out} }$ 。在不丧失一般性的情况下，我们假设 $d_{in} =d_{c}=d_{out} =d$ 。与自注意力的情况类似，我们可能有 $C=X$ 。在本文的其余部分，我们将重点放在Lambda层的一个特定实例上，并表明它捕获了长程的基于内容和基于位置的交互，而不是具体化注意力映射。图2所示是lambda层的计算图。

图2 lambda层的计算图。位置的查询项的上下文信息被汇聚为一个lambda

{\lambda _{n} }\in {R}^{{\vert k\vert } \times \vert v \vert }

。动态地应用lambda，以分发上下文信息，以产生输出

{\bf{y}}_{n}={{\lambda }_{n}^T }{\bf{q}}_{n}

。这个过程捕获了基于内容的和基于位置的交互，而避免了计算注意力映射。

我们首先介绍将lambda层应用给单个的查询项 $({\bf{q}}_{n} ,n)$ 。

生成上下文lambda函数 我们希望生成一个线性函数 $R^{\vert k \vert} \rightarrow R^{\vert v \vert}$ ，即矩阵 ${\lambda _{n} }\in {R}^{{\vert k\vert } \times \vert v \vert }$ 。lambda层首先通过线性投影上下文来

计算键项矩阵 ${ \bf K }$ 和值项矩阵 ${\bf V}$ 。然后键项矩阵在上下文维度上通过一个softmax操作来归一化，以产生归一化的键项矩阵 $\overline{{ \bf K }}$ 。 ${\lambda _{n} }$ 矩阵即是，使用归一化的键项矩阵 $\overline{{\bf K }}$ 和位置嵌入 ${\bf{E} }_{n}$ 来聚合值项矩阵 ${\bf V}$ ：

其中，我们还定义了内容lambda ${\lambda }^c$ 和位置lambda ${\lambda }_{n}^p$ 。

• 所有位置 $n$ 的查询项的内容lambda ${\lambda }^c$ 是共享的，并且对上下文元素是排列不变的。它编码的是如何仅仅依赖于上下文内容来变换查询项 ${\bf {q}}_{n}$ 。

• 位置lambda ${\lambda }_{n}^p$ 通过位置嵌入 ${\bf {E}}_{ n }$ 是取决于查询项的位置 $n$ 。它编码的是如何根据内容元素 ${\bf{c}}_{m}$ 及其与查询项的相对位置 $(n,m)$ 来变换查询项 ${\bf {q}}_{n}$ 。

将lambda应用到查询项上 查询项 ${\bf {q}}_{n}\in {R}^{\vert k \vert }$ 是通过输入 ${\bf {x}}_{n}$ 经过一个可学习的线性投影来获得的，lambda层的输出是：

Lambda层的解释 矩阵 ${\lambda _{n} }\in {R}^{{\vert k\vert } \times \vert v \vert }$ 的每一行，可以看作是一个固定规模的 $\vert k \vert$ 个上下文特征。根据上下文的内容（基于内容的交互）和结构（基于位置的交互）对这些上下文特征进行聚合。应用lamdbda，然后基于查询项动态地分发这些上下文特征，以输出 ${\bf{y}}_{n}=\sum\nolimits_{k}q_{nk} {\lambda }_{nk}$ 。这个过程捕获了基于内容和基于位置的交互，但是没有计算注意力图。

归一化 你可以修改公式(1)和公式(2)，以把非线性操作和归一化操作也包括进去。我们的实验表明，在计算查询项和值项之后使用批量归一化，是能带来提升的。

3.2 多个查询项的计算方式，以减小复杂度

复杂度分析 对于一个批量，包含 $\vert b \vert$ 个样本，每个包含 $\vert n \vert$ 个输入（我能不能理解|b|为批量大小，|n|为像素数），应用lambda层所需的算术运算数和内存占用分别为 $\Theta (bnmkv)$ 和 $\Theta (knm+bnkv)$ 。由于相对位置嵌入 $e_{nm}$ ，在长序列上，我们的内存占用仍然是二次的。然而这个二次项并不随着批量大小增大而增大，也就是不像注意力操作中为每个样本产生注意力图一样。在实践中， $\vert k \vert$ 被设置为一个小值，例如 ${\vert k \vert } =16$ ，这样我们就可以处理大批量大小的大尺寸输入，这是注意力所无法办到的，见表4。此外，位置嵌入可以在lambda层之间共享，以保持其 $\Theta (knm)$ 内存占用不变，而注意力图的内存占用会随着层的数量而变化。注意力图通常需要存储起来进行反向传播。

多个查询项的lambda层减小空间复杂度和时间复杂度 回想一下，lambda层将输入 ${\bf{x}}_{n}\in R^{d}$ 映射我输出 ${\bf{y}}_{n}\in R^{d}$ 。对应到公式(2)上，这意味着 $\vert v \vert =d$ 。 $\vert v \vert$ 值过小可能会导致特征向量 $y_{n}$ 性能上的瓶颈，但是考虑到时间复杂度和空间复杂度为 $\Theta (bnmkv)$ 和 $\Theta (knm+bnkv)$ ，提高输出维度 $\vert v \vert$ 会导致很大的计算量。

我们提出将我们的lambda层的时间复杂度和空间复杂度与输出维度 $d$ 解耦。我们不使 $\vert v \vert =d$ ，我们创建 $\vert h \vert$ 个查询项 $\left\{ { \bf{q}}_{n}^h \right\}$ ，对每个查询项 ${\bf{q}}_{n}^h$ 应用相同的lambda ${\lambda _{n} }$ ，然后将各个输出拼接起来 ${{\bf {y}} _{n} }=concat(\lambda _{n} {\bf{q}}_{n}^1,...\lambda _{n} {\bf{q}}_{n}^{\vert h \vert })$ 。这样我们就有了 $\vert v \vert =d/\vert h \vert$ ，这样就可以成 $\vert h \vert$ 倍地减小复杂度。头数 $\vert h \vert$ 控制着lambdas ${\lambda _{n} }\in {R}^{{\vert k\vert } \times \vert d \vert /\vert h \vert }$ 与总体查询项 ${\bf {q}}_{n}\in {R}^{\vert hk \vert }$ 的相对大小。

我们将此操作称作为多查询项lambda层，并且在图3中展示了其einsum实现。lambda层对超参数 $\vert k \vert$ 和 $\vert n \vert$ 的选择是鲁棒的（见附录D.1），这样我们就可以很灵活地控制其复杂度。我们在大多数实验中设置 $\vert h \vert =4$ 。

图3 多查询项的lambda层的伪代码。在计算位置lambda时，位置嵌入可以满足各种条件，例如平移等变性（图中未显示）。

einsum运算表示任意维张量之间的广义压缩。它在数字上相当于广播其输入以共享其维度的并集，并对输出中未指定的所有维度按元素相乘再求和。

我们注意到，虽然这类似于多头注意力和多查询注意力（Shazeer, 2019提出一种多查询注意力以加速基于注意力的解码器），但动机是不同的。在注意力机制中，使用多个查询项是为了增加表示能力和复杂度。而在lambda层中，使用多个查询项会减小复杂度和表示能力（没考虑额外的查询项）。

将多查询计算形式扩展到线性注意力 最终，我们将我们的分析扩展到线性注意力上，线性注意力可以看做是只有内容交互的注意力层（附录C.3有详细讨论）。我们期望多查询项的计算形式也能为线性注意力机制到来计算上的好处。

3.3 使Lambda层具有平移等变性

使用相对位置嵌入 $e_{nm}$ 可以对上下文的结构做出明确的假设。特别是，在许多学习场景中，平移等变性（即输入偏移多少，输出就要偏移多少）是一种强归纳偏置。我们通过确保位置嵌入满足：对于任意的偏移 $t$ ，都有 $e_{nm} =e_{t(n)t(m)}$ ，以来获得位置交互作用中的平移等变性。在实践中，我们定义了相对位置嵌入 $R\in {R}^{{\vert r\vert } \times \vert k \vert }$ ，其中 $r$ 是所有对 $(n,m)$ 的可能相对位置，并将其再索引为 $E\in {R}^{{\vert n\vert } \times \vert m \vert \times \vert k \vert }$ ，这样就有 $e_{nm} =r_{r(n,m)}$ 。（我们建议读者看代码以了解细节。）

3.4 Lambda卷积：在局部上下文中长程交互建模

尽管长程交互建模有好处，但在许多任务中，局部模式仍然是一个很强的归纳偏置。使用全局上下文可能会产生噪声或计算量过大。因此，将位置交互的范围限制到查询项的位置 $n$ 周围的局部邻域可能是有用的，就像局部自注意力和卷积一样。这可以实现为将所需范围以外的上下文位置 $m$ 的相对嵌入置为0。但是这样做计算量仍是大的，因为计算仍在进行，它们只是被置零。

Lambda卷积 在上下文是多维网格状数据的情况下，我们可以使用常规卷积从局部上下文等价地计算位置lambda。

我们将这种操作称为lambda卷积。一个 $n$ 维的lambda卷积可以实现为一个带通道乘法器的 $n-d$ 维的深度卷积；也可以实现为 $(n+1)-d$ 的卷积，其将 $\pmb{V}$ 中的 $v$ 维视为额外的空间维。我们在附录B.1中介绍了这两种实现。:

由于计算现在被限制在局部范围内，lambda卷积的时间复杂度和内存占用相对于输入长度是线性的。Lambda卷积很容易附加扩张和步幅等其它功能（dilation and striding），并且在专用硬件上（其实现速度是被优化过的）也能享受硬件加速带来的好处（Nickolls & Dally, 2010; Jouppi et al., 2017）。这与局部自注意力（Parmar等人的ImageTransformer，2018；Ramachandran等人的StandAloneAttention，2019）需要对查询项和上下文局域块的有重叠地利用，因而导致延迟和内存占用（见表4）形成鲜明对比。

注释9：浮点运算量（时间复杂度）在专用硬件（如TPU或GPU）上不一定能很好的说明延迟特性。尽管lambda卷积的时间复杂度和空间复杂度是线性的，但是实际应用中它可能比全局lambda层慢，特别是当局部领域的范围比较大的时候。表4是个例子。

表4 lambda层达到了更高的ImageNet准确度，同时比自注意力各种变体方法速度更快、内存占用更少。内存占用情况是在全精度、批量大小为128、默认超参数的情况下的结果。存储lambda的内存占用与网络其余部分中激活函数的内存开销相当，因此可忽略。b:批量大小，h:头数/查询数，n:输入长度，m:上下文长度，k:查询项/键项深度，l:层数。

4. 相关工作

表2是以往的长程交互建模的方法，并将它们与本工作所提的多查询项lambda层进行了对比。我们在附录C中详细讨论了相关工作。

表2 捕获长程交互的可用方案。lambda层相比于相对注意力(Shaw等人，2018；Bello等人，2019)以较低的内存占用捕获基于内容和位置的交互。使用多查询lambda层将复杂度缩小|h|倍。此外，通过使用具有线性复杂度的lambda卷积，可以将基于位置的交互限制在局部范围内。b:批量大小，h:头数/查询项数，n:输入长度，m:上下文长度，r:局部范围大小，k:查询/键深度，d:维度输出。

通道注意力与线性注意力 lambda层的广义抽象概念，也即，将可用的上下文变换为线性函数并应用到查询项上，是非常通用的，因此包含了以往许多工作。与本工作相近是通道注意力机制与线性注意力机制：SENet，Performer，LinearTransformer(Transformers are RNNs)。这些注意力方法能捕获长程交互，无需计算注意力图，可以被视作只包含内容交互的lambda层特例。Lambda层将这些方法广义化，并且扩展为同时考虑基于内容和基于位置的交互，使它们能够作为高度结构化数据（如图像）上的独立层使用。我们没有试图像线性注意力那样逼近注意力核，而是专注于上下文lambda函数的高效计算方式的设计，并重新调整多查询项方法的用法(Shazeer，2019)，以进一步降低计算成本。 $\pmb{V}$

视觉领域中的自注意力 在NLP领域自注意力已经成为标配方法，而在视觉领域，自注意力的应用进展缓慢而稳定(Wang et al.，2018；Bello et al.，2019；Ramachandran et al.，2019；Carion et al.，2020)。在本项工作进展的同时，Dosovitskiy等人在ImageNet上用transformer达到88.6%的准确度，其是将图像分块并将块作为序列，其transformer是在300M大规模图像数据集上预训练的。

5 实验

我们在标准计算机视觉基准上评估lambda层：ImageNet分类（Deng et al.，2009）、COCO检测和实例分割（Lin et al.，2014）。视觉域非常适合展示lambda层的灵活性，因为（1）对于高分辨率图像，自注意力的内存占用会产生问题；（2）图像具有高度结构化，因此基于位置的交互至关重要。

LambdaResNets 我们通过替换ResNet结构瓶颈块中的3x3卷积来构建LambdaResNets。当替换所有这些3x3卷积时，在表中我们仅仅标示被测的层的名称来指代整个网络，例如conv+channel attention，或者lambda layer。我们用LambdaResNets表示表18（附录E.1）中的一系列混合结构。除非另有说明，所有lambda层都使用| k |=16，| h |=4，范围大小为| m |=23x23，实现如图3所示。更多的实验和细节可以在附录中找到。

5.1 Lambda层优于卷积层和注意层

我们首先考虑输入图像大小为224x224的标准ResNet-50结构。在表3中，我们将lambda层与(a)标准卷积，即基线ResNet-50；(b)通道注意力，即SE-ResNet；(c)多个自注意力变体进行比较。Lambda层在参数量上明显优于所有基线，并且比通道注意力有+0.8%的提升。

表3：用ResNet50架构比较ImageNet分类的lambda层和注意力机制。 lambda层在参数代价很小的情况下，明显优于各种注意力方案。所有模型都是在相似的设置中训练的（见附录E.2）。括号中是相比卷积基线的提升。超参数|u|的详情见附录B.4。†表示我们自己的实现。

5.2 Lambda层相比自注意力在计算上的优势

表4中，我们将lambda层与自注意力和当前吞吐量、内存复杂度和ImageNet准确度进行了比较。我们的研究结果突出了自注意力的弱点：由于记忆开销大，自注意力不能建模全局交互，轴向自注意力仍然内存占用大，局部自注意力太慢了。相比之下，lambda层可以捕获高分辨率图像上的全局交互，比局部自注意力准确率提高+1.0%，同时速度快了近3倍。此外，位置嵌入可以跨lambda层共享，这样可以进一步地降低内存需求。最后，lambda卷积的内存复杂度是线性的，这在大图像中（如在检测或分割中所看到的）是非常实用的。

我们还发现，在控制变量局部窗口大小时，lambda层的性能优于局部自注意力，在表中| m |=7x7时是78.1%比77.4%，这表明lambda层的好处不仅仅是提高了速度和可伸缩性。（注释11：基于内容的lambda层获取的仍是全局交互。）

注释10：局部自注意力的延迟量是Ramachandran et al.(2019)私下提供的，是基于查询块和重叠内存块的实现(Parmar et al., 2018)。专门设计的注意力核可以大大加速局部自注意力，这是个很有前景的研究方向。

5.3 混合算法提高了图像分类的速度-精度折平衡

研究混合结构 尽管和自注意力相比占用内存要少得多，但是使用lambda层捕获全局上下文仍然是二次的时间复杂度（表2），当输入是高分辨时这样代价仍然是大的。另外你可以假设，特征在经过了一些操作处理之后有了语义信息，这是使用全局上下文是最能发挥作用的；而在前几层就使用全局信息则是浪费的。附录5.3中，我们研究了使用标准卷积来捕获局部上下文和lambda层来捕获全局上下文的混合设计。我们发现，这种卷积lambda混合模式，和其纯卷积对手相比，提高了表示能力，吞吐量的下降微乎其微。

LambdaResNets显著提高了ImageNet分类的速度-精度平衡。 基于我们对混合体系结构的研究和Bello等人（2021年）的扩展/训练策略，我们设计了一系列不同大小的混合LambdaResNets（见第E.1节）。图4显示了与TPUv3硬件上的速度-准确度的Pareto曲线， LambdaResNets和EfficientNets (Tan & Le, 2019)对比。为了控制变量，我们还比较了标准3x3卷积（表示为ResNet RS wo/SE）和具有通道注意力的3x3卷积（表示为ResNet RS w/SE）替换lambda层、保留宏观结构的网络。所有网络正则化方式、输入分辨率相同，都训练350个epoch。

在所有尺度下LambdaResNets的速度-准确度平衡都要优于基线。在相同的准确度下，LambdaResNets比EfficientNets快3.2-4.4倍，比ResNet RS快1.6-2.3倍，显著提高了图像分类的速度-准确度Pareto曲线（注释12）。我们最大的型号LambdaResNet-420训练图像大小为320，达到了84.9%的top-1 ImageNet准确度，比标准3x3卷积的相应网络结构高出0.9%，比带通道注意力的相应结构高出0.65%。

注释12：Ridnik et al.（2020）和Zhang et al.（2020）报告了在GPU上比EfficientNets快2倍的ImageNet分类高准确率网络。注释再注释：即TResNet和ResNeSt。

图4 LambdaResnet和EfficientNets之间的速度精度比较。当使用和EfficientNets相同的训练方式和正则化设置时，LambdaResnet比EfficientNets快3.2-4.4倍，比带SE块的ResNet-RS快1.6-2.3倍。LambdaResnet上的标注是深度（是指通道数？）和输入图像分辨率。我们最大的LambdaResNet，图像大小为320的LambdaResNet-420，达到了84.9%的最高准确度。

使用伪标签扩展到更大的数据集 我们使用来自JFT数据集的130M伪标签图像在半监督学习设置下训练LambdaResNets，就像训练EfficientNet-NoisyStudent checkpoints 一样（谢其哲2020）。表5比较了使用JFT数据集进行训练时具有相当准确度的一组典型模型的吞吐量和ImageNet分类准确度。在图像大小为288的情况下训练和评估的LambdaResNet-152实现了86.7%的top-1 ImageNet精度，同时在相同准确度下比EfficientNet NoiseStudent checkpoints具有更少的参数，速度是其9.5倍。

表5 根据额外数据训练的模型比较。ViT-L/16在JFT上预训练，并在ImageNet上以384x384的分辨率进行微调，而EfficientNet和LambdaResNet在ImageNet和JFT伪标签上共同训练。训练和推理吞吐量是在8个TPUv3核上运行的。

5.4 目标检测和实例分割结果

在表6中，我们评估了LambdaResNets作为Mask RCNN中COCO物体检测和实例分割任务的主干。使用lambda层可以在所有大小的物体上获得一致的增益，尤其是最难定位的小物体。这表明lambda层对于需要位置信息的更复杂的视觉任务也很有竞争力。

表6 1024x1024输入的基于Mask-RCNN结构的COCO物体检测和实例分割。小、中、大物体的平均精度是AP s/m/l。使用lambda层可以在所有大小物体上获得一致的增益，尤其是小物体。

6 讨论

Lambda层与注意力操作相比如何？ Lambda层与自注意力相比具有良好的伸缩性。使用自注意力的Vanilla Transformers的内存占用是 $\Theta (blhn^2 )$ ；而LambdaNetworks则是 $\Theta (lkn^2 )$ ，跨层共享位置嵌入时则是 $\Theta (kn^2 )$ 。这使得lambda层在更高分辨率和更大批量大小上的应用成为可能。此外，lambda卷积的实现比局部自注意力更简单更快。最后，ImageNet实验表明lambda层的性能优于自注意力，这表明lambda层的好处不仅仅是提高了速度和可伸缩性。

Lambda层与线性注意机制有何不同？ Lambda层广义化并扩展了线性注意力，以捕获基于位置的交互，这对于建模高度结构化的输入（如图像）至关重要（见附录D.1中的表10）。由于lambda层并不旨在逼近一个注意力核，lambda层允许更灵活的非线性和规范化，我们发现这也是有益的（见附录D.1中的表12）。最后，我们提出多查询项lambda层以减小复杂度，这个是与线性注意力中的多头注意力是不同的。附录C.3详细讨论了线性注意力。

如何在视觉域中最好地使用lambda层？ 与全局或局部注意力相比，lambda层的可伸缩性、速度和易实现性都有所提高，这使得它们成为视觉领域中可使用的有力候选对象。我们的消冗实验表明，当优化的目的是要取得速度-准确度的平衡时，lambda层在视觉结构的中、低分辨率阶段最为有利。设计一个完全由lambda层构成的网络也是可行的，其参数效率和flops效率会更高。在附录A中会讨论实用的建模建议。

Lambda层的推广 虽然这项工作侧重于静态图像任务，但我们注意到lambda层可以建模各种结构上的交互，如图、时间序列、空间晶格等。我们预计lambda层将有助于更多模式，包括多模式任务。我们在附录B.2中讨论了带屏蔽的上下文和自回归任务。

结论我们提出了一类新的层，称为lambda层，它提供了一个可伸缩的框架来捕获输入及其上下文之间的结构化交互。Lambda层将可用的上下文汇总为固定规模的线性函数，称为lambda，直接应用于其相关的查询项。由此产生的神经网络LambdaNetworks计算效率高，能够以较小的内存占用捕获长距离的依赖关系，使其能够应用于高分辨率图像等大型结构化输入。对计算机视觉任务的大量实验表明，它们比卷积网络和注意力网络用途更广泛、性能更优越。最值得注意的是，我们引入了LambdaResNets，这是一系列的混合结构，它达到了出色的ImageNet精度，比流行的EfficientNets实现了高达9.5倍的速度提升，显著提高了图像分类模型的速度-精度平衡。

致谢作者要感谢Barret Zoph和William Fedus的无休止的讨论、富有成效的建议和仔细的修改；Jonathon Shlens、Mike Mozer、Prajit Ramachandran、Ashish Vaswani、Quoc Le、Neil Housby、Jakob Uszkoreit、Margaret Li、Krzysztof Choromanski的许多有见地的评论；Hedvig Rausing的南极信息图表；Zolan Brinnes代表OST；Andrew Brock，Sheng Li代表协助分析效率网；Adam Kraft，Thang Luong和Hieu Pham代表协助进行半监督实验，Google Brain团队代表对论文进行有益的讨论

附录

A 实用的建模建议

如果你想让它在tpu/gpu上更快... 混合模型达到了更好的速度精度折衷。全局上下文可能会在计算上造成浪费，特别是在早期的高分辨率层中，在这些层中，特征缺少语义信息，可以用范围更小的lambda卷积（例如| m |=5x5或7x7）或标准3x3卷积来代替。此外，当从模型/训练设置开始时（when starting from a working model/training setup），使用混合结构可能需要更少的微调。

如果你希望使FLOPs最小（例如在移动端、嵌入式的应用）... 考虑具有反向瓶颈的混合结构，如第D.3.2节所述。在此基础上如果要进一步减小FLOPs，最好使用较小范围的lambda卷积（例如| m |=5x5或7x7）。

如果你遇到内存占用的问题... 通过跨层（尤其是具有最高分辨率的层）共享位置嵌入，可以减少内存占用。使用lambda卷积的内存占用效率更高。可以提高内存效率。减少查询深度| k |或增加头数| h |也可以减少内存占用。

如果你训练地不稳定... 我们发现将ResNet瓶颈块的最后一个batchnorm层中的γ参数初始化为0很重要（这是大多数代码库中的默认值）。沿着上下文的空间维度归一化（使用softmax）键项（keys）很重要。早期在同一残差块中连续使用2个lambda层的实验是不稳定的，这表明应避免连续使用2个lambda层。

你应该使用哪个lambda卷积实现？ 在我们使用TensorFlow1.x在TPUv3硬件上进行的实验中，我们发现 $n-d$ 的深度卷积与 $(n+1)-d$ 的卷积的实现的速度是相近的。这可能在不同的软件/硬件上会有差异。

如果你的任务不需要基于位置的交互呢？ lambda层的计算开销主要是基于位置的交互。如果您的任务不需要它们，您可以尝试只包含内容的lambda层或任何其他线性注意力机制。我们建议使用多查询项公式（与通常的multi-head相反）并缩放模型的其他维度。

B 其它变体

B.1 Lambda卷积的实现代码

图5 多查询项lambda层和1d的lambda卷积的伪代码。一个n-d维的lambda卷积可等效地实现为一个常规的(n+1)-d卷积或具有通道乘法器的n-d深度卷积。使用einsum实现计算位置lambda时，可以使嵌入满足各种条件（例如平移等变性和掩蔽）。

B.2 从带屏蔽的上下文生成lambda

文字介绍略。图6是伪代码。

图6 从带屏蔽的上下文生成lambda 的伪代码

B.3 多头与多查询项的lambda层的区别

在本节中，我们使用多查询项的方法，这与自注意力中常用的多头方法不同。图7所示是一个多头lambda层的实现。表7比较了多头和多查询lambda层的复杂度。与多头相比，使用多查询可将计算量缩小| h |倍（每个lambda的查询项数）。我们在早期的实验中还发现，多查询lambda在速度和精度上有更好的折衷。此外，多头lambda层不像lambda卷积那样可用简单的局部实现。

图7 多头lambda层的实现的伪代码。我们建议不用这个，而是使用多查询项的方式。

表7 多头和多查询lambda层之间的复杂度比较。与标准的多头方法相比，使用多查询将复杂性降低了一个因子|h|（每个lambda的查询数）。

B. 4 增加额外维度的表现力

我们简要地试验了一种变体，它的目的是使计算获得lambda的成本增加，而应用lambda（到查询项上，笔者注）的成本不变。这是通过在键、位置嵌入和值中引入一个额外的维度来实现的，称为内部深度，相应超参数记为| u |。每个键（或位置嵌入）现在是一个| k |×| u |矩阵，而不是一个| k |维向量。类似地，每个值现在都是一个| v |×| u |矩阵，而不是一个| v |维向量。lambda是通过对上下文位置和内部深度位置| u |求和获得的，并且具有与默认情况类似的| k |×| v |形状。图8是实现的伪代码，表8是复杂度。实验（见附录D.1）表明，这种变体可以提高精度，但我们发现，在现代机器学习加速器上控制速度时，使用| u |=1（即默认情况）是最佳的。

图8 具有内部深度|u|的多查询lambda层的伪代码实现。lambda是通过上下文位置和内部深度维度求和得到的。这个变体分配更多的计算来生成lambda，而应用lambda的成本不变。相应的n-d lambda卷积可以用(n+1)-d常规卷积等价实现。

表8 具有内部深度|u|的多查询lambda层的复杂度。

C 其他相关工作

在这一节中，我们将回顾注意力操作以及改进其可扩展性的相关工作。我们讨论了lambda层和通道、空间或线性注意力机制之间的联系，并展示了如何将它们转换为lambda层的不太灵活的特例。最后，我们简要回顾了视觉领域中的自注意力，并讨论了与专家模型的联系。

C.1 Softmax注意力

Softmax注意力 Softmax注意力为每个查询 $\pmb{q}_{n}$ 生成一个分布，即 $\pmb{a}_{n} =softmax({\pmb{K}}{\pmb{q}_{n}})\in R^{\vert m\vert }$ ，其中键 $\pmb{K}$ 是从上下文 $\pmb{C}$ 获得的。然后使用注意分布 $\pmb{a}_{n}$ 形成从上下文获得的值 $\pmb{V}$ 的线性组合，即 $\pmb{y}_{n} =\pmb{V}^T{\pmb{a}_{n}}=\sum\nolimits_{m}a_{nm} \pmb{v}_{m}\in R^{\vert v\vert }$ 。当我们采取值项的加权求和时，我们将查询项 $\pmb{q}_{n}$ 转换为输出项 $\pmb{y}_{n}$ ，并丢弃其注意力分布。此操作捕获基于内容的交互，但不捕获基于位置的交互。

相对注意力 为了建立基于位置的交互模型，相对注意力(Shaw et al.，2018)引入了一个可学习的位置嵌入 $| m |$ 的矩阵 $\pmb{E}_{n}\in R^{{\vert m\vert}\times {\vert k\vert}}$ ，然后计算注意力分布为 $\pmb{a}_{n} =softmax(({\pmb{K}+{\pmb{E}_{n}}}){\pmb{q}_{n}})\in R^{\vert m\vert }$ 。注意力分布现在还取决于查询项 $n$ 与上下文元素 $m$ 的相对位置，因此，相对注意力捕获基于内容和基于位置的交互。

C.2 稀疏注意力

相对注意力难以应用到高分辨率图像或长序列上，因为存储注意力映射需要二次 $\Theta (\vert bnm \vert )$ 的内存占用。因此，许多最近的工作建议将特定的模式强加给注意力映射，以减小上下文元素数量，从而减小内存占用。这些方法包括局部注意模式（Dai et al.，2019；Parmar et al.，2018；Ramachandran et al.，2019）、轴向注意力模式（Ho et al.，2019；Wang et al.，2020a）、静态稀疏注意力模式（Child et al.，Beltagy et al.，2020）或动态稀疏注意模式（Kitaev et al.，2020）。参见Tay等人（2020）的综述。它们的实现可能相当复杂，有时需要底层核实现来获得计算上的好处，或者可能依赖于对输入形状的特定假设（例如，轴向注意）。

相比之下，lambda层使用简单的einsum和卷积对于全局和局部上下文都很容易实现，并且可以捕获密集的内容和基于位置的交互，而不需要对输入形状进行任何假设。

C.3 线性注意力：联系与区别

另一种降低注意力机制计算要求的方法是，使用时间复杂度和空间复杂度为线性的操作来近似注意力操作，这些被称为线性注意力或高效注意力。线性注意力机制可以追溯到de Brebisson&Vincent´（2016）；Britz et al.（2017），后来由Chen et al.（2018）；Shen et al.（2018）引入视觉领域。最近，随着很多工作修改了在序列任务上很流行的transformer（Katharopoulos et al.，2020；Wang et al.，2020b；Choromanski et al.，2020），线性注意力重新流行起来。

基于核因子分解的线性注意 线性注意力通常是通过将注意力重新解释为一个相似核，并利用一些特征函数 $\phi$ 来做低秩核分解：

$Attention(\pmb{Q},\pmb{K},\pmb{V})=softmax({\pmb{Q}}{\pmb{K}}^T ){\pmb{V}}\sim{\phi( {\pmb{Q}})}({\phi( {\pmb{K}^T })}{\pmb{V}})$ $(3)$

计算 ${\phi( {\pmb{K}^T })}{\pmb{V}}\in R^{{\vert k\vert}\times {\vert v\vert}}$ 就不用去计算 ${\phi( {\pmb{Q}})}{\phi( {\pmb{K^T }})}$ ，因而对于输入长度 $| n |$ 具有线性复杂度。

特征函数 $\phi$ 有多种选择。例如，Katharopoulos et al.(2020)使用 $\phi (x)=elu(x)+1$ ，而Choromanski et al.（2020）使用正值正交随机特征来近似原始softmax注意核。在视觉领域，Chen et al.(2018)和Shen et al.(2018)（2018）都使用 $\phi (x)=softmax(x)$ 。这是为了保证（尽管没有显式计算）注意映射 $\phi (Q)\phi (K^T )$ 的行和为1，就像在常规注意力中一样。

接下来讨论lambda层和线性注意力机制之间的主要区别。

1) Lambda层扩展了线性注意力，也考虑了基于位置的交互。 对于单个查询 $\pmb{q}_{n}$ ，公式(3)中的核逼近可重写为：

$\pmb{{y}_{n} }= ({\phi( {\pmb{K}^T })}{\pmb{V}})^T{\phi( {\pmb{{q}_{n} } })}$ $(4)$

它类似于公式(1)中内容交互的输出 $\pmb{y}^c_{n} =(\pmb{\lambda }^c)^T\pmb{q}_{n} =({\overline{\pmb{K}} }^T {\pmb{V}})^T\pmb{q}_{n}$

。Lambda层扩展线性注意机制，也考虑基于位置的交互：

$\pmb{y}_{n} =(\pmb{\lambda }^T_{n})\pmb{q}_{n} =(\pmb{\lambda }^c+\pmb{\lambda }^p_{n})^T\pmb{q}_{n} =(({\overline{\pmb{K}}+\pmb{E}_{n}) }^T {\pmb{V}})^T\pmb{q}_{n}$ $(5)$

在上述公式中，计算位置（或内容）lambda具有相同的时间复杂度 $\Theta (bmkv)$ 。由于位置lambda不在查询位置 $n$ 之间共享，因此对于所有 $| n |$ 个查询重复此开销，从而导致总时间复杂度 $\Theta (bnmkv)$ 。与线性注意力机制不同，lambda层因为考虑基于位置的交互，相对于输入长度具有二次时间复杂度（在全局上下文中）。

2) Lambda层不一定试图逼近注意力核。 虽然注意力核的近似理论上是有动机的，但我们认为它们可能是不必要的限制。例如，公式3中的核近似要求 $\pmb{Q}$ 和 $\pmb{K}$ 上的特征函数 $\phi$ 相同，并且排除了使用更灵活的非线性和归一化方案。相比之下，lambda层并不试图逼近注意力核。这简化了其设计，并允许更灵活的非线性和归一化方案，这在我们的消冗研究中很有用（见附录D.1中的表12）。考虑到位置嵌入是独立于健项的，这就使得lambda卷积的很简单而有效的局部实现方法成为可能。相对注意力核的近似，要求归一化带健项的位置嵌入，也就是 $\phi (\pmb{K}+\pmb{E}_{n} )$ ，而不是 $\phi (\pmb{K})+\pmb{E}_{n}$ ，这不能在带卷积的局部上下文情况下实现。

3) lambda广义化揭示了多查询方式的计算优势。 最后，本文提出将 ${\overline{\pmb{K}} }^T {\pmb{V}}$ 和 $\pmb{E}^T_{n} {\pmb{V}}$ 矩阵抽象为直接应用于查询的线性函数（内容和位置lambda）。lambda抽象揭示了多查询公式（相对于传统的多头注意公式）作为一种降低计算成本的方法的优点。

C.4 通道注意力和空间注意力作为lambda层

我们证明lambda抽象概括了通道和空间注意力机制，两者都可以看作lambda层的具体实例。这一观察结果与我们的实验一致，实验表明lambda层在计算效率上优于通道和空间注意力。

通道注意力 通道注意力，例如SENet(Hu et al.，2018c；b)、FiLM layer(Perez et al.，2017)，通过整合整个特征图的信号，根据跨通道相互作用重新校准特征。特别地，SE操作可以写成 $y_{nk}＝w_kq_{nk}$ ，其中 $w_k$ 是查询 $\pmb{q}_{n}$ 中的通道 $k$ 的激励权重(excitation weight)。这可以看作是使用一个对角线lambda，它在查询位置之间是共享的： $\pmb{\lambda }_{n} =diag(w_{1} \cdots w_{\vert k\vert })$ 。通道注意力机制作为卷积的补充已被证明是有用的，但不能作为一个独立的层，因为其丢弃空间信息。

空间注意力 相反，空间注意机制根据所有通道聚集的信号对每个位置重新加权（Xu等人，2015；Park等人，2018；Woo等人，2018）。这些机制可以写成 $y_{nk}＝w_nq_{nk}$ ，其中 $w_n$ 是输入查询 $\pmb{Q}$ 中位置 $n$ 的注意力权重。这可以看作是使用（位置相关的）标量lambdas $\pmb{\lambda }_{n} =w_{n} \pmb{1 }$ ，其中 $\pmb{1 }$ 是单位矩阵。空间注意力机制作为卷积的补充也被证明是有用的，但不能作为一个独立的层，因为其丢弃通道信息。

C.5 视觉领域的自注意力

自注意已经被用于视觉领域的无数任务中。其中包括图像分类（Bello et al.，2019；Ramachandran et al.，2019；Cordonnier et al.，2019；Zhao et al.，2020；Wu et al.，2020；Dosovitskiy et al.，2020）；目标检测和目标中心任务（Wang et al.，2018；Hu et al.，2018a；Carion et al.，2020；Locatello et al.，2020）；视频任务（Sun et al.，2019；Liao et al.，2019）；自回归/对抗生成建模（Parmar et al.，2018；Zhang et al.，2019；Brock et al.，2019；Chen et al.，2020a）和多模态文本视觉任务（Chen et al.，2020b；Lu et al.，2019；Li et al.，2019；Radford et al.，2021）

在视觉中首次使用自注意力可以追溯到非局部块（ non-local block，Wang等人，2018），该块在ConvNet的低分辨率阶段添加了一个单头全局自注意力残差，用于远程相关性建模。非局部块已被证明可用于补充卷积，但不能用作独立层，因为它不模拟基于位置的交互。

在低分辨率，用全局相对注意力取代卷积。 Bello等人（2019）引入了一种二维相对注意力机制，该机制被证明是卷积的替代品，单独用的时候具有竞争力，而将注意力特征与卷积特征拼接时会产生更强大的结果。采用二维相对位置嵌入的全局多头自注意力机制代替ResNet结构瓶颈块中的空间卷积。由于全局注意力占用内存大，该操作只能在低分辨率特征图上使用，所提出的结构是conv-transformer混合结构。

Srinivas等人（Bottleneck transformers for visual recognition，2021年）最近利用现代训练和缩放技术重新研究了类似的混合设计。其建议在ResNet架构的最后阶段使用1的步长来提高性能，而不是拼接卷积特征映射。

在高分辨率，用局部/轴向相对注意力取代卷积。 全局注意力的巨大内存占用很快被多项研究所解决，这些研究提出了局部注意力（Ramachandran et al.，2019；Hu et al.，2019）和轴向注意力（Ho et al.，2019；Wang et al.，2020a；Shen et al.，2020）等注意上下文的大小（见C.2节）。这些方法能够将注意力用在更高的分辨率的特征图上，并且促成了完全由注意力机制构成的网络，，但由于使用专门的注意力模式，速度可能较慢。

缩放胜过归纳偏置(Scaling trumps inductive bias)

同时，ViT（Dosovitskiy et al.，2020）建议只关注图像块（而不是单个像素），解决占用内存大的问题。虽然图像块不能保留准确的位置信息和平移等变性，但通过对大规模数据集（如300M图像）进行预训练，可以恢复归纳偏置的损失。最显著的是，当对ImageNet数据集进行微调时，ViT达到接近最先进的准确度，相比卷积对手（Kolesnikov et al.，2020；Xie et al.，2020）只需要更少的训练。这一结果重新激发了人们对在视觉领域中使用自注意力的兴趣，已经有很多后续工作是建立在ViT的基础上（Touvron等人，2021）。（注释14：很多后续工作声称在小的数据集上胜过了ViT，但是小规模数据集原本就不是ViT的目的。）尽管取得了令人印象深刻的图像分类结果，但基于图像块的方法是否能够扩展到更大的图像，能否应用到需要准确位置信息的任务中（如检测），仍然令人担忧。

我们强调的是，用图像块而非像素来减小内存占用，与使用特定的操作来减小内存占用，两者是正交的，并且我们预期lambda层（或线性注意）可以成功地作为像素块方案的补充。

C.6 与超网络和专家模型的联系

LambdaNetworks生成自己的计算，即lambda，即 $\pmb{y}_{n} =\pmb{\lambda }_{n} \pmb{q}_{n}$ 。因此，它们也可以被视为超网络的一个扩展（Ha等人，2016），基于上下文信息动态生成计算。

最后，LambdaNetworks与稀疏激活的专家模型有一些联系（Shazeer et al.，2017；Fedus et al.，Switch transformers，2021）。稀疏激活的专家模型根据输入的查询从权重库中选择计算（即lambda），而LambdaNetworks则根据上下文信息（包括输入查询）生成计算。

D 附加实验

D.1 消冗研究

我们进行了多次消冗，验证了位置相互作用、长程相互作用和灵活的归一化方案的重要性。除非另有说明，本节中的所有实验结果都是ImageNet分类准确度，网络结构都是用lambda层替换ResNet-50中的空间卷积的LambdaNetwork。

不同的查询深度、头数和内部深度。 表9显示了查询深度 $| k |$ 、头数 $| h |$ 和内部深度 $| u |$ 对性能的影响（内部深度 $| u |$ 的表示见附录B.4）。我们的实验表明，lambda层在较宽的超参数范围内优于卷积基线和注意力基线，证明了该方法的鲁棒性。

表9 ResNet50体系结构中使用lambda层时ImageNet分类任务的消冗。所有配置都以较低的参数代价超过卷积基线。正如预期的那样，我们通过增加查询深度

| k |

或内部深度

| u |

获得了额外的改进。头数最好设置为中间值，如

| h |=4

。头太多会降低值项的深度

| v |=d/| h |

。头太少会导致导致查询太少，这两种情况都会影响性能。

内容交互 vs. 位置交互 表10显示了ImageNet分类任务中基于内容和基于位置的交互的相对重要性。我们发现，基于位置的交互对于达到高精度至关重要，而基于内容的交互只会比基于位置的交互带来微小的改进。

注释15：这种观察受到与本工作同期进的ViT的挑战（Dosovitskiy等人，2020年），他们表明，在大规模数据集（例如300M图像）上进行预训练时，基于内容的交互对于图像分类是足够的。

表10：内容和位置交互的贡献。正如预期的那样，位置交互作用对于完成图像分类任务至关重要

窗口大小的重要性 LambdaNetworks的内存占用很小，即使在相对较高的分辨率下，也可以考虑全局上下文。表11给出了在LambdaNetwork中改变窗口大小时的flops计数和top-1 ImageNet准确度。我们发现使用更大范围的好处，在| m |=15x15附近性能处于高位平台，这验证了与ResNet架构中使用的通常3x3空间卷积相比，长距离交互的重要性。我们在图像分辨率较大的实验中，选择| m |=23x23作为默认值。

表11：改变位置lambda的作用域大小对ImageNet分类任务的影响。我们用lambda层（输入图像大小为224x224）替换ResNet-50的最后两个阶段中的3x3空间卷积。Flops随着作用域的大小而显著增加，但是我们强调，在使用einsum实现时，较大的作用域不会转换为较慢的延迟（参见图3）。

归一化 表12描述了lambda层设计中的归一化操作。我们发现，对健项的归一化是至关重要的，softmax以外的其它归一化函数也可考虑。对对查询项和值项使用批量归一化也很有用。

Lambda层中归一化方案的影响。键的归一化是沿着上下文空间维度

m

的方向。查询项的归一化是验证查询项的深度

k

的方向。

D.2 混合模型的研究

在本节中，我们研究使用标准卷积来捕获局部上下文和lambda层来捕获全局上下文的混合设计。

注释16：我们也可以使用lambda卷积来捕获局部上下文。

Lambda层在哪里最有用？ 表13显示了混合LambdaNetwork架构的吞吐量和精度，控制变量是ResNet-50架构中卷积和lambda层位置的函数。在考虑速度-准确度权衡时，我们观察到lambda层在最后两个阶段（通常称为c4和c5）最有帮助。我们将在ResNet的最后两个阶段用lambda层替换3x3卷积的体系结构称为LambdaResNet-C4。

表13：混合型实现了更好的速度-准确度权衡。宏观结构是ResNet50，输入尺寸是224x224，L是lambda，C是卷积。与标准3x3卷积相比，第5阶段C5中的Lambda层几乎没有使速度降低。第4阶段C4中的Lambda层比标准的3x3卷积相对较慢，但产生显著的精度增益。

进一步推高速度-准确度Pareto高台 在表14中，我们进一步研究了c4阶段lambda层的数量对吞吐量和精度的影响。我们的结果表明，lambda层的最大优点可以通过(a)在c4阶段用lambda层替换一些3x3卷积和（b）在c5阶段替换所有3x3卷积来获得。由此产生的混合LambdaResNets体系结构增加了表示能力，而与普通ResNet体系结构相比，吞吐量几乎可以忽略不计。表18给出了我们的LambdaResNets系列的详细块配置和lambda层的放置。

表14 Lambdaresnet c4阶段lambda层数的影响。lambda层能带来的大部分好处可以通过在c4中有几个lambda层来获得。这样的混合设计最大限度地提高了速度-准确度的权衡。LambdaResNet-C4结构在C4和C5中使用lambda层。LambdaResNet块配置见表18。在ImageNet分类任务中，模型被训练了350个时期。

比较混合lambda和注意力模型。 在上述混合设计中，lambda层的内存节省与注意力模型相比不太显著，因为操作发生在较低的分辨率下。因此，在考虑混合设计时，很自然会问lambda层是否仍比自注意力有好处。我们以我们最大的混合模型为例（见表18）。LambdaResNet-420在320x320输入上训练，在c4中使用8个lambda层，每个TPU-v3核可以容纳32个样本。这为lambda层增加了38.4MB的成本（如果共享位置嵌入，则为4.8MB），而使用注意层则需要0.625GB。这一增长在实践中可能并不显著，仔细衡量混合注意力变体17将是一件有趣的事情。我们指出，来自表4的实验表明lambda层的好处不仅仅是提高了可伸缩性，而且强调对于需要较大输入的任务（如检测），内存节省更为明显。

注释17：我们将在这个论文的未来版本中对这些架构进行基准测试

D.3 计算效率结果

D.3.1 与大型EfficientNet的计算效率比较

在表15和表16中，我们展示了LambdaNetworks的参数和flops效率。我们发现LambdaResNet-C4在ResNet体系结构的最后两个阶段取代了3x3卷积，在这两个阶段产生的参数代价最高，LambdaResNet-C4的参数效率和flops效率要好于EfficientNets。这些结果具有重要意义，因为EfficientNets是由神经架构搜索（Zoph&Le，2017）专门设计的，以使用高度计算效率的深度卷积（Tan&Le，2019）最小化计算成本。

表15:LambdaResNet-C4和EfficientNet-B6之间的参数效率比较。尽管使用较小的图像大小，LambdaResNet-C4的参数效率更高。增大图像大小可能会提高精度，同时保持参数的数量不变。模型被训练了350个epoch。

表16:LambdaResNet-C4和EfficientNet-B6之间的Flops效率比较。我们使用较小的局部作用域(|m|=7x7)来减少lambda层中的flops。模型被训练了350个epoch。

D.3.2 资源受限场景中的lambda层

最后，我们使用MobileNetv2架构简要研究了资源受限场景中的lambda层（Sandler et al.，2018）。MobileNets（Howard et al.，2017；Sandler et al.，2018；Howard et al.，2019）采用轻型反向瓶颈块，其包括以下序列：1）用于扩展通道数量的pointwise卷积，2）用于空间混合的深度卷积，以及3）用于通道混合的最终pointwise卷积。深度卷积的使用（与常规卷积相反）减少了参数和flops，使得反向瓶颈特别适合于嵌入式应用程序。

轻量化lambda块 我们构造一个轻量级lambda块，如下所示。我们用一个小范围的lambda卷积来代替反向瓶颈中的深度卷积，这个lambda卷积的大小是|m|=5x5，查询深度是|k|=32，头数是| h |=4。我们还将第一个逐点卷积更改为输出相同数量的通道（而不是增加通道数量），以进一步减少计算量。

在MobileNetv2中添加lambda层 我们希望评估lambda层是否可以提高MobileNet体系结构的flops-准确度权衡（或参数-准确度平衡）。我们实验了一种简单的策略，即用我们提出的轻量级lambda块替换几个反向瓶颈，这样得到的体系结构具有与它们的基线相似的计算需求。在MobileNet-v2体系结构中，用轻量级lambda块替换第10个和第16个反向瓶颈块的简单过程将参数和flops减少了∼10%，同时将ImageNet的精度提高了0.6%。这表明lambda层可能非常适合用于资源受限的场景，例如嵌入式视觉应用（Howard et al.，2017；Sandler et al.，2018；Howard et al.，2019）。

表17:Lambda层在资源受限的情况下提高了ImageNet的准确度。在MobileNet-v2结构中，用轻量级lambda块替换第10个和第16个反向瓶颈块可将参数和flops减少10%，同时将ImageNet ac精确度提高0.6%。

E 实验细节

E.1 网络结构的细节

Lambda层实现详细信息 除非另有规定，否则所有lambda层都使用查询深度| k |=16、| h |=4头和内部深度| u |=1。位置lambda是使用大小为| m |=23x23的局部上下文生成的，内容lambda是使用图3中描述的einsum实现使用全局上下文生成的。局部的位置lambda可以与lambda卷积替代( interchangeably)实现，或者通过使用全局einsum、在局部上下文之外掩蔽位置嵌入来实现（图5）。后者可以更快，但由于 $\Theta (knm)$ 项，具有更高的flops和内存占用（见表2）。在我们的实验中，我们只对输入长度| n |>852或帧内深度| u |>1使用卷积版的实现。当帧内深度增加到| u |>1时，我们切换到卷积实现，并将作用域大小减小到| m |=7x7以减少flops。

使用单位正态分布 $N(0,1)$ 随机初始化位置嵌入。对于lambda层中的线性投影，我们使用fan-in初始化。用 $N(0,\vert d \vert ^{-1/2} )$ 分布随机初始化计算 $\pmb{K}$ 和 $\pmb{V}$ 的投影。用 $N(0,\vert kd \vert ^{-1/2} )$ 分布随机初始化计算 $\pmb{Q}$ 的投影（这类似于缩放点积注意力机制，有一点不同就是缩放被吸收在投影中）。我们对 $\pmb{Q}$ 和 $\pmb{V}$ 进行批量归一化，并通过softmax操作对 $\pmb{K}$ 进行归一化。

ResNet 我们使用ResNet-v1实现，并将瓶颈块的最后一个批量归一化(Ioffe&Szegedy，2015)中的 $\gamma$ 参数初始化为0。Squeeze-and-Excitation层采用4的挤压比。与ResNet RS(Bello et al.，2021)类似，我们使用ResNet-D(He et al.，2018)，并用3x3的步幅卷积替换网络前面的最大池化层。我们的块分配和缩放策略（即选定分辨率作为模型深度的函数）也严格遵循ResNet RS的缩放建议（Bello等人，2021）。

LambdaResNets 通过用我们提出的lambda层替换ResNet RS结构的瓶颈块的空间3x3卷积来构造LambdaResNets，除了stem保持不变。我们在lambda层之后使用核为3x3、步长为2的平均池化来代替步幅卷积。Lambda层均匀地分布在c4阶段，c5中的所有瓶颈都使用Lambda层。表18给出了混合LambdaResNets的详细块配置和lambda层的位置。在采用lambda层而不是标准3x3卷积的瓶颈块中，我们不使用Squeeze-and-Excitation。

表18 Pareto曲线中LambdaResNets模型块配置和lambda层的放置位置。LambDaresnes使用He等人（2016）、Bello等人（2021）的块分配

E.2 训练细节

ImageNet训练设置。 我们考虑为ImageNet分类任务提供两种训练设置，训练90个epoch和训练350个epoch。训练90个epoch使用标准的预处理，并可以和经典作品公平比较。训练350个epoch使用改进的数据增强和正则化，并且更接近现代作品中使用的训练方法，具有最先进的精确度。

表19：详细的LambdaResNets结果。延迟是指使用bfloat16激活的8个TPU-v3核上1024个批大小的每个训练步骤的时间。

使用vanilla ResNet的监督学习训练90epoch 训练90epoch的设置，我们使用vanilla ResNet与以前的作品公平比较。我们使用了官方实现中的默认超参数，没有进行额外的调优。使用动量为0.9的SGD，通过反向传播对所有网络进行端到端90个epoch的训练。批量大小 $B$ 为4096，分布在32个TPUv3芯上（Jouppi等人，2017年），权重衰减设置为1e-4。学习率在5个epoch内从0线性扩展到0.1B/256，然后使用cosine衰减（Loshchilov&Hutter，2017）。我们在可训练参数上使用衰减为0.9999的批量归一化和权重为0.9999的指数移动平均，标签平滑为0.1。输入图像大小设置为224x224。我们使用标准的训练数据扩充（随机裁剪和50%概率的水平翻转）。

与表3相比，大多数工作都使用类似的训练设置，并用他们提出的方法替换ResNet体系结构中的3x3空间卷积。我们注意到，Ramachandran等人（2019年）训练的时间更长（130个epoch而不是90个epoch），但没有使用标签平滑，这可能会混淆我们的比较。

监督学习350个epoch的设置 ImageNet上的更高精度通常是通过更长时间的训练以及增强和正则化来获得的（Lee等人，2020；Tan&Le，2019）。与Bello等人（2021年）类似，权重衰变减少到4e-5，我们采用了RandAugment（Cubuk等人，2019年）和2层，dropout（Srivastava等人，2014年）和随机深度（Huang等人，2016年）。具体超参数见表20。批量大小B为4096或2048，分布在32或64个TPUv3核上，具体取决于内存限制。

表20 训练LambdaResNet的超参，训练350个epoch，使用RandAugment，dropout，随机深度

我们使用一个包含大约为2%ImageNet训练集（1024个中取20个）的held-out验证集来优化我们的模型。我们对最大模型的验证集执行提前停止，从分辨率为288x288的LambdaResNet-350开始，对小模型我们只报告最终准确度。

伪标签半监督学习。 我们的半监督学习按照谢其哲的NoisyStudent的设置。我们使用相同的数据集，有130M张图像、经过类别过滤和平衡的JFT数据集，使用EfficientNet-L2模型（88.4%的ImageNet准确度）生成的伪标签。超参数与监督的ImageNet 350 epochs实验相同。

延迟测量。 图4报告了使用混合精度训练（bfloat16激活）在8个TPUv3核上处理1024幅图像的训练延迟（即每个训练步骤的时间）。训练延迟最初是在8个TPUv3核上测量的，从1024个（即每个核128个）的总批量大小开始，将批量大小除以2，直到它适合内存。然后我们在图4中报告了标准化的延迟。例如，如果用512（而不是1024）的批大小来测量延迟，我们通过将测量的延迟乘以2来规范化报告的延迟。表4、表13和表14使用全精度（即float32激活）报告了8个TPUv3核的推断吞吐量。ViT的延迟（Dosovitskiy等人，2020年）由作者私下提供。

FLOPS计数。 在使用图3中的einsum实现计算位置lambda时，我们没有把屏蔽掉的位置的FLOPs计数置零。FLOPs计数高度依赖于作用域大小，默认情况下作用域大小相当大（| m |=23x23）。在表11中，我们表明可以显著减小作用域的大小，从而以最小的性能降低来实现最小化FLOPs。

COCO物体检测 在Mask-RCNN结构中，我们采用了改进的ImageNet训练设置中的结构作为主干。所有模型都在1024x1024个图像上进行训练，从零开始进行130k步，批量大小为256，分布在128个TPUv3核上，并进行同步批量归一化。我们在训练过程中使用[0.1，2.0]的多尺度抖动。在训练的90、95和97.5%的步骤中，将学习率从0到0.32预热1000个步骤，再除以10。权重衰减设置为4e-5。

MobileNet训练设置 所有MobileNet结构都在Imagenet上接受350个epoch的训练，采用224x224分辨率，标准预处理。我们使用与Howard等人（2019）相同的超参数。更具体地说，我们使用的RMSProp动量为0.9，批量大小为4096，分为32个TPUv3核。学习率线性上升到0.1，然后每3个周期乘以0.99。我们使用重量衰减1e-5和下降概率为0.2的下降。

表19：详细的LambdaResNets结果。延迟是指使用bfloat16激活的8个TPU-v3核上1024个批量大小的每一个训练步的时间

注：在翻译的时候可能有个错误，原文好一律用LambdaResNets are 3.2 - 4.4x faster than EfficientNets，笔者前面可能翻译成xxx速度是xxx的xx倍，后面可能翻译成xxx速度比xxx快xx倍。

最后编辑于：2021.07.05 17:38:11

人面猴
序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
沈念sama阅读 199,064评论 5赞 466
死咒
序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
沈念sama阅读 83,606评论 2赞 376
救了他两次的神仙让他今天三更去死
文/潘晓璐我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
开封第一讲书人阅读 146,011评论 0赞 328
道士缉凶录：失踪的卖姜人
文/不坏的土叔我叫张陵，是天一观的道长。经常有香客问我，道长，这世上最难降的妖魔是什么？我笑而不...
开封第一讲书人阅读 53,550评论 1赞 269
港岛之恋（遗憾婚礼）
正文为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
茶点故事阅读 62,465评论 5赞 359
恶毒庶女顶嫁案：这布局不是一般人想出来的
文/花漫我一把揭开白布。她就那样静静地躺着，像睡着了一般。火红的嫁衣衬着肌肤如雪。梳的纹丝不乱的头发上，一...
开封第一讲书人阅读 47,919评论 1赞 275
城市分裂传说
那天，我揣着相机与录音，去河边找鬼。笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。我是一名探鬼主播，决...
沈念sama阅读 37,428评论 3赞 390
双鸳鸯连环套：你想象不到人心有多黑
文/苍兰香墨我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
开封第一讲书人阅读 36,075评论 0赞 254
万荣杀人案实录
序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
沈念sama阅读 40,208评论 1赞 294
护林员之死
正文独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋以下内容为张勋视角年9月15日...
茶点故事阅读 35,185评论 2赞 317
白月光启示录
正文我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
茶点故事阅读 37,191评论 1赞 328
活死人
序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
沈念sama阅读 32,914评论 3赞 316
日本核电站爆炸内幕
正文年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
茶点故事阅读 38,482评论 3赞 302
男人毒药：我在死后第九天来索命
文/蒙蒙一、第九天我趴在偏房一处隐蔽的房顶上张望。院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
开封第一讲书人阅读 29,585评论 0赞 19
一桩弑父案，背后竟有这般阴谋
文/苍兰香墨我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。一阵脚步声响...
开封第一讲书人阅读 30,825评论 1赞 255
情欲美人皮
我被黑心中介骗来泰国打工，没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。一个月前我还...
沈念sama阅读 42,194评论 2赞 344
代替公主和亲
正文我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
茶点故事阅读 41,703评论 2赞 339

Lambda Networks（★★★★★）

推荐阅读更多精彩内容