Backbone | 谷歌提出LambdaNetworks：无需注意力让网络更快更强(文末获取论文源码)

2021-05-28 14:56:24 浏览数 (1)

本文提出新的层：Lambda Layers，用于替代自注意力！由其衍生出的LambdaResNets改进了分类模型的速度-准确性权衡。比EfficientNets快约4.5倍！同时在检测和分割任务中性能提升也比较明显，代码已开源！ 作者单位：谷歌大脑

1 简介

本文提出新的层：Lambda Layers，一种自注意力的替代方法，主要用于捕获输入和Content结构化的信息（例如一个像素被其他像素包围）之间的远程交互。Lambda Layer通过将Content转换成称为Lambda的线性函数，并将这些线性函数分别应用于每个输入以捕获此类交互。

与线性注意类似，Lambda层绕过了复杂的注意力图，但是相反，它们对内容和基于位置的交互都进行了建模，这使它们可以应用于大型结构化输入（例如图像）。由此产生的神经网络体系结构LambdaNetworks在ImageNet分类，COCO目标检测和COCO实例分割方面明显优于其卷积和注意力，同时具有更高的计算效率。

此外，作者还设计了LambdaResNets，是一个跨尺度的混合体系结构，可显著改善图像分类模型的速度-准确性权衡。LambdaResNets在ImageNet上具有出色的准确性，而其速度是流行的EfficientNets的3.2-4.4倍。在使用额外的130M伪标签图像进行训练时，LambdaResNets在相应的EfficientNet精度上可实现高达9.5倍的加速。

Lambda Layer主要贡献：

  • 1、提出了Content-Based和Position-Based的交互建模的Lambda层，而且很容易通过einsum操作和卷积核实现高效的操作；
  • 2、Lambda层提供了通Channel、Spatial和Linear Attention的统一视图；
  • 3、在ImageNet分类任务上Lambda层显著优于卷积和Attention，同时计算效率更高；
  • 4、Lambda层延迟低同时对于内存的需求也比较小；
  • 5、提出混合模型方法作为最大限度地权衡Lambdanworks速度-精度的方法；
  • 6、通过结合Lambda Layer设计了LambdaResNets，并且达到SOTA水平；
  • 7、在半监督学习中，使用额外的1.3亿个伪标签图像进行训练，实现了9.5倍的加速；
  • 8、使用Mask-RCNN对LambdaResNets进行COCO对象检测和实例分割的评估。LambdaResNet在2个任务的均有所提升。

2 解决的问题

  • 1、模型精度的提升长期依赖于数据集；
  • 2、Self-Attention虽然一定程度上解决了问题1，但是其巨大的内存需求阻碍了其在长序列和多维数据上的应用。

本文提出Lambda层主要是解决上诉2个问题，它以降低内存成本的方式对查询和同一组结构化Content元素之间的远程交互进行建模。Lambda层将每个Content转换为一个线性函数，称为Lambda函数，然后直接应用于相应的查询。Self-Attention定义了查询和Content元素之间的相似Kernel，而Lambda层则将Content信息总结为一个固定大小的线性函数(即矩阵)，从而绕过了对内存密集型注意映射的需求。

下图中说明了差异：

Global Content中包含3个查询块及其Local Content的示例

Self-Attention将每个查询块与其Content的Attention分布联系起来

Lambda层将每个Content转换为相应查询块的线性函数Lambda函数

本文所提的Lambda层是通用的，可以实现为全局、局部或Mask Content中Content-Based和Position-Based的交互建模。由此设计产生的神经网络(LambdaNetworks)计算效率高，以很小的内存成本模拟了长期依赖关系，因此可以应用于高分辨率图像等大型结构化输入。

3 远程交互建模

3.1 定义说明

符号

|n|

表示一个集合的基数，该集合的元素用

n

作索引；

x_n

表示

X

的第

n

行，

x_{ij}

表示

X

的第

i

行、第

j

列元素。在可能的情况下，采用self-attention来提高可读性，并突出差异。

Query和Context

这里使用向量

Q={(q_n,n)}

来表示Query，使用结构化集合

C={(c_m,m)}

来表示Context；其中每个query

(q_n,n)

以为Context

q_nin R^{|k|}

和对应位置

n

为特征；类似地，每个Context元素

(c_m,m)

以其Context

c_m

和在Context的位置

m

为特征。

(n,m)

对可以指结构化元素之间的任何成对关系，例如图中像素之间的相对距离或节点之间边的相对距离。

interactions

这里使用映射函数

F:((q_n,n), C)mapsto y_n

和给定的Context C 将query

(q_n,n)

映射输出向量

y_n

。当处理结构化输入时，此函数可作为神经网络中的一层。参考

(q_n,c_m)

Context-Based的交互和

(q_n,(n,m))

Position-Based的交互注意到，虽然绝对位置信息有时直接添加到query(或context元素)content，但作者认为这种类型的交互是Position-Based，因为它忽略了关系

(n,m)

查询和Context元素之间的位置。

3.2 引入Keys来捕捉long-range interactions

在深度学习的背景下，优先考虑Fast batched linear操作和使用dot-product操作作为交互。这促使引入向量，这些向量可以通过dot-product操作与query交互，因此具有与query相同的维数。

特别是，Content-Based交互

(q_n,c_m)

需要依赖于

c_m

的

|k|

维向量，通常称为关键

k_m

。相反，Position-Based交互

(q_n,(n,m))

需要一个相对位置嵌入

e_{nm}in R^{|k|}

。由于 query/key 深度

|k|

和context空间维度

|m|

不在输出

y_n

中，所以这些维度需要压缩作为层计算的一部分。

因此，每个捕获远程交互的层都可以根据它是先缩小query深度还是先缩小context位置来描述。

3.3 Attentional interactions

收缩查询深度首先在query元素和Context元素之间创建一个相似核(Attention Map)，称为Attention操作。由于Context位置

|m|

的数量越来越大，而输入和输出维度

|k|

和

|v|

保持不变，假设层输出是一个相对小维度的向量

|v|<<|m|

，则可以假设计算Attention Map是没必要的。

3.5 Lambda interactions

与Attentional interactions相反，对于一些线性函数

lambda(C,n): R^{|k|}rightarrow R^{|v|}

来说，简单地将每个query映射为

y_n=F((q_n,n),C)=lambda (C,n)(qn)

或许是更有效的。在这种情况下，Context被聚合成一个固定大小的线性函数

lambda n=lambda(C,n)

。每个

lambda_n

作为一个小的线性函数，独立于Context存在，并在应用到其关联的query

q_n

后被丢弃。

4. Lambda Layer

4.1 Context转换为线性函数

Lambda Layer将输入

X

和Context C作为输入，并生成线性函数lambdas，然后应用于Query产生输出

y

。不失一般性假设

d_{in}=d_c=d_{out}=d

。和self-attention一样可能存在

C=X

。

下图给出了lambda层的计算过程：

这里大致描述应用到single query的lambda layer：

1、生成与Context相关的Lambda函数：

这里希望生成一个

R^{|k|}

到

R^{|v|}

的线性函数，即一个矩阵

lambda_n

。Lambda层首先通过线性预测Context和keys（通过Softmax规范化操作得到）在Context的位置，计算

V

值和

K

键。用标准化键

overline K

和位置嵌入

En

将

V

集合得到

lambda_n

矩阵如下：

式中，定义content lambda

lambda^c

和position lambda

lambda^p_n

为：

  • content lambda
lambda^c

：content lambda在所有query位置

n

上共享，并且对context元素的排列是不变的。它对如何仅基于context内容转换query

q_n

进行编码;

  • position lambda
lambda^p_n

：position lambda通过位置嵌入

E_n

依赖于query位置

n

。它编码如何根据context元素

c_m

及其相对位置转换query

q_n

到query

(n,m)

。

2、将Lambda函数用于query：

query

q_nin R^{|k|}

是通过学习线性投影从输入

x_n

得到的，Lambda层的输出为:

3、Lambda层的解释：

lambda_nin R^{|k|×|v|}

矩阵的列可以看作是一个固定大小

|k|

的context特征集。这些context特征基于context的content(content-based interactions)和结构(position-based interactions)进行聚合。然后应用lambda根据query动态分布这些context特征，以产生

y_n=sum_kq_{nk}lambda_{nk}

的输出。这个过程捕获内容和基于位置的互动，而不产生attention maps。

4、Normalization：

作者实验表明，在计算query和value后应用batch normalization是有帮助的。

4.2 减少复杂性的Multi-Query

由于输出维度比较大可能会带来比较大的计算复杂度，因此作者还设计了Multi-query Lambda Layer以减少复杂度，进而降低推理时间。

本文从输出维

d

解耦lambda层的时间和空间复杂性。不是强加

|v|=d

，而是创建

|h|

query

{q_n^h}

，对每个query

q_n^h

应用相同的

lambda_n

，并将输出cat为

y_n=concat(lambda_nq_n^1,···,lambda_nq_n^{|h|})

。现在有

|v|=d/|h|

，这将复杂度降低了

|h|

的1倍。head的数目

|h|

控制了

lambda_nin R^{|k|times |d|/h}

的大小相对于query

q_nin R^{|hk|}

的总大小。

代码语言：javascript复制
def lambda layer(queries, keys, embeddings, values):
    """Multi−query lambda layer."""
    # b: batch, n: input length, m: context length,
    # k: query/key depth, v: value depth,
    # h: number of heads, d: output dimension.
    content lambda = einsum(softmax(keys), values, 'bmk,bmv−>bkv')
    position lambdas = einsum(embeddings, values, 'nmk,bmv−>bnkv')
    content output = einsum(queries, content lambda, 'bhnk,bkv−>bnhv')
    position output = einsum(queries, position lambdas, 'bhnk,bnkv−>bnhv')
    output = reshape(content output   position output, [b, n, d])
    return output

虽然这类似于multi-head或multi-query attention formulation，但动机是不同的。在attention操作中使用多个query增加了表征能力和复杂性。相反，在lambda层中使用多个query降低了复杂性和表示能力(忽略额外的query)。

最后，作者指出，可以将Multi-Lambda Layer扩展到linear attention，可以将其视为只包含content的lambda层。

4.3 使Lambda层translation equivariance

使用相对位置嵌入

e_{nm}

可以对context的结构做出明确的假设。特别地，在许多学习场景中，translation equivariance(即移动输入导致输出等效移动的特性)是一种强烈的归纳bias。

本文通过确保嵌入的位置满足对任何translation

t
e_{nm}=e_{t(n)t(m)}

获得translation equivariance位置的相互作用。在实践中,定义了一个张量的相对位置嵌入

Rin R^{|r|times |k|}

, index r为可能的相对位置

(n,m)

对。

4.4 Lambda卷积

尽管远距离相互作用有诸多好处，但在许多任务中，局部性仍然是一种强烈的感应偏向。使用全局context可能会被证明是noisy或computationally excessive。因此，就像local self-attention和卷积一样将位置交互的范围限制在query位置

n

周围的一个局部邻域可能是有用的。这可以通过将context位置

m

在所需范围之外的相对嵌入置零来实现。然而，这种策略对于

|m|

的大值仍然是复杂的，因为计算仍然会发生——它们只是被置零。

在context排列在multidimensional grid的情况下，可以通过使用正则卷积从local contexts等效地计算位置lambda。作者把这个运算称为lambda卷积。n维的lambda卷积可以使用n-d与channel乘法器的深度卷积或

(n 1)-d

卷积来实现，将

v

维中的

v

维视为额外的空间维。

由于计算现在被限制在局部范围内，lambda卷积的时间和内存复杂度与输入长度成线性关系。lambda卷积很容易与其他功能一起使用，比如dilation和striding，并在专门的硬件加速器上实现了优化。这与local self-attention的实现形成了鲜明的对比，后者需要具体化重叠查询和context块的特性补丁，增加了内存消耗和延迟。

代码语言：javascript复制
# b: batch, n: input length, m: context length, r: scope size,
# k: query/key depth, v: value depth, h: number of heads, d: output dimension.
def compute position lambdas(embeddings, values, impl=’einsum’):
    if impl == ’einsum’: # embeddings shape: [n, m, k]
        position lambdas = einsum(embeddings, values, ’nmk,bmv−>bnkv’)
    else: # embeddings shape: [r, k]
        if impl == ’conv’:
            embeddings = reshape(embeddings, [r, 1, 1, k])
            values = reshape(values, [b, n, v, 1])
            position lambdas = conv2d(values, embeddings)
        elif impl == ’depthwise conv’:
            # Reshape and tile embeddings to [r, v, k] shape
            embeddings = reshape(embeddings, [r, 1, k])
            embeddings = tile(embeddings, [1, v, 1])
            position lambdas = depthwise conv1d(values, embeddings)
        # Transpose from shape [b, n, v, k] to shape [b, n, k, v]
        position lambdas = transpose(position lambdas, [0, 1, 3, 2])
    return position lambdas
    
    
def lambda layer(queries, keys, embeddings, values, impl=’einsum’):
    """Multi−query lambda layer."""
    content lambda = einsum(softmax(keys), values, ’bmk,bmv−>bkv’)
    position lambdas = compute position lambdas(embeddings, values, impl=impl)
    content output = einsum(queries, content lambda, ’bhnk,bkv−>bnhv’)
    position output = einsum(queries, position lambdas, ’bhnk,bnkv−>bnhv’)
    output = reshape(content output   position output, [b, n, d])
    return output

5 问题讨论

1、lambda层与attention操作相比如何？

Lambda层规模有利比较self-attention。使用self-attention的Vanilla Transformers有

θ(blhn^2)

内存footprint，而LambdaNetworks有

θ(lkn^2)

内存footprint。这使得lambda层能够以更高的分辨率和更大的批处理规模使用。

此外，lambda卷积的实现比它的local self-attention对等物更简单、更快。最后，ImageNet实验表明lambda层优于self-attention，证明了lambda层的好处不仅仅是提高了速度和可伸缩性。

2、lambda层与线性注意力机制有何不同？

Lambda层推广和扩展了线性注意力公式，以捕获基于位置的交互，这对于建模高度结构化的输入(如图像)至关重要的；

由于目标不是近似一个attention kernel，lambda层可以通过使用非线性和规范化进一步提升性能，

3、在视觉领域如何最好地使用lambda层？

与global或local attention相比，lambda层改进了可伸缩性、速度和实现的简易性，这使得它们成为可视化领域中使用的一个强有力的候选对象。消融实验表明，当优化速度-精度权衡时，lambda层在视觉架构的中分辨率和低分辨率阶段最有利。也可以设计完全依赖lambda层的架构，这样可以更有效地进行参数化处理。作者在附录A中讨给出了使用的意见。

4、lambda层的泛化性如何？

虽然这项工作主要集中在静态图像任务上，但作者注意到lambda层可以被实例化来建模各种结构上的交互，如图形、时间序列、空间格等。lambda层将在更多的模式中有帮助，包括多模态任务。作者在附录中对masked contexts和auto-regressive进行了讨论和复现。

6 实验

6.1 ImageNet分类

该实验主要是针对基于ResNet50的架构改进设计和实验，通过下表可以看出在ResNet50的基础上提升还是比较明显的：

通过下表可以看出Lambda层可以捕捉在高分辨率图像上的全局交互，并获得1.0%的提升，且相较于local self-attention速度提升接近3倍。此外，位置嵌入可以跨lambda层共享，以最小的退化成本进一步减少内存需求。最后，lambda卷积具有线性的内存复杂度，这对于检测或分割中看到的非常大的图像是实用的。

通过下表可以看出，基于EfficientNet的改进在不损失精度的情况下，可以将训练速度提升9倍之多，推理速度提升6倍之多。

6.2 COCO目标检测实验

在下表中，作者在COCO目标检测和实例分割任务评估了LambdaResNets作为Mask-RCNN的Backbone的性能。使用lambda层可以在所有目标大小上产生一致的增益，特别是那些最难定位的小对象。这表明lambda层对于需要local信息的更复杂的视觉任务也具有竞争力。

7 参考

[1].LAMBDANETWORKS: MODELING LONG-RANGE INTERACTIONS WITHOUT ATTENTION

serverless linux 批量计算

0 人点赞