4. Lambda Layer

4.1 Context转换为线性函数

Lambda Layer将输入和Context C作为输入，并生成线性函数lambdas，然后应用于Query产生输出。不失一般性假设。和self-attention一样可能存在。

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

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

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

这里希望生成一个到的线性函数，即一个矩阵。Lambda层首先通过线性预测Context和keys（通过Softmax规范化操作得到）在Context的位置，计算值和键。用标准化键和位置嵌入将集合得到矩阵如下：

式中，定义content lambda 和position lambda 为：

• content lambda ：content lambda在所有query位置上共享，并且对context元素的排列是不变的。它对如何仅基于context内容转换query 进行编码;
• position lambda ：position lambda通过位置嵌入依赖于query位置。它编码如何根据context元素及其相对位置转换query 到query 。

2、将Lambda函数用于query：

query 是通过学习线性投影从输入得到的，Lambda层的输出为:

3、Lambda层的解释：

矩阵的列可以看作是一个固定大小的context特征集。这些context特征基于context的content(content-based interactions)和结构(position-based interactions)进行聚合。然后应用lambda根据query动态分布这些context特征，以产生的输出。这个过程捕获内容和基于位置的互动，而不产生attention maps。

4、Normalization：

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

4.2 减少复杂性的Multi-Query

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

本文从输出维解耦lambda层的时间和空间复杂性。不是强加，而是创建 query ，对每个query 应用相同的，并将输出cat为。现在有，这将复杂度降低了的1倍。head的数目控制了的大小相对于query 的总大小。

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

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

本文通过确保嵌入的位置满足对任何translation  获得translation equivariance位置的相互作用。在实践中,定义了一个张量的相对位置嵌入, index r为可能的相对位置对。

4.4 Lambda卷积

尽管远距离相互作用有诸多好处，但在许多任务中，局部性仍然是一种强烈的感应偏向。使用全局context可能会被证明是noisy或computationally excessive。因此，就像local self-attention和卷积一样将位置交互的范围限制在query位置周围的一个局部邻域可能是有用的。这可以通过将context位置在所需范围之外的相对嵌入置零来实现。然而，这种策略对于的大值仍然是复杂的，因为计算仍然会发生——它们只是被置零。

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

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

# 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有θ内存footprint，而LambdaNetworks有θ内存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