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
