LAMBDANETWORKS: MODELING LONG-RANGE INTERACTIONS WITHOUT ATTENTION

Abstract

We present lambda layers -an alternative framework to self-attention -for capturing long-range interactions between an input and structured contextual information (e.g. a pixel surrounded by other pixels). Lambda layers capture such interactions by transforming available contexts into linear functions, termed lambdas, and applying these linear functions to each input separately. Similar to linear attention, lambda layers bypass expensive attention maps, but in contrast, they model both content and position-based interactions which enables their application to large structured inputs such as images. The resulting neural network architectures, LambdaNetworks, significantly outperform their convolutional and attentional counterparts on ImageNet classification, COCO object detection and instance segmentation, while being more computationally efficient. Additionally, we design LambdaResNets, a family of hybrid architectures across different scales, that considerably improves the speed-accuracy tradeoff of image classification models. LambdaResNets reach excellent accuracies on ImageNet while being 3.2 -4.4x faster than the popular EfficientNets on modern machine learning accelerators. In large-scale semi-supervised training with an additinal 130M pseudo-labeled images, LambdaResNets achieve up to 86.7% ImageNet accuracy while being 9.5x faster than EfficientNet NoisyStudent and 9x faster than a Vision Transformer with comparable accuracies 1 .

1. INTRODUCTION

Modeling long-range dependencies in data is a central problem in machine learning. Selfattention (Bahdanau et al., 2015; Vaswani et al., 2017) has emerged as a popular approach to do so, but the costly memory requirement of self-attention hinders its application to long sequences and multidimensional data such as imagesfoot_1 . Linear (or efficient) attention mechanisms (Katharopoulos et al., 2020; Choromanski et al., 2020) offer a scalable remedy for high memory usage but fail to model internal data structure, such as relative distances between pixels or edge relations between nodes in a graph. This work addresses both issues. We propose lambda layers which model long-range interactions between a query and a structured set of context elements at a reduced memory cost. Lambda layers transform each available context into a linear function, termed a lambda, which is then directly applied to the corresponding query. Whereas self-attention defines a similarity kernel between the query and the context elements, a lambda layer instead summarizes contextual information into a fixed-size linear function (i.e. a matrix), thus bypassing the need for memory-intensive attention maps. This difference is illustrated in Figure 1 . Lambda layers are versatile and can be implemented to model both content-based and position-based interactions in global, local or masked contexts. The resulting neural networks, LambdaNetworks, are computationally efficient, model long-range dependencies at a small memory cost and can therefore be applied to large structured inputs such as high resolution images. We evaluate LambdaNetworks on computer vision tasks where works using self-attention are hindered by large memory costs (Wang et al., 2018; Bello et al., 2019) , suffer impractical implementations (Ramachandran et al., 2019) , or require vast amounts of data (Dosovitskiy et al., 2020) . In our experiments spanning ImageNet classification, COCO object detection and instance segmentation, LambdaNetworks significantly outperform their convolutional and attentional counterparts, while being more computationally efficient and faster than the latter. We summarize our contributions: G LO B A L C O N T E X T AT T E N T I O N M A P S L A M B DAS LO CA L C O N T E X TS Q U E R I E S • Lambda layers: a class of layers, that model content-based and position-based interactions without materializing attention maps. Lambda layers offer a unifying view of channel, spatial and linear attention (Appendix D.4). Some of our observations, such as the computational benefits of a multi-query formulation, extend to linear attention. Lambda layers are easily implemented with einsum operations and convolution kernels, operations with efficient implementations on modern machine learning accelerators. • Lambda layers significantly outperform their convolution and attention counterparts on the ImageNet classification task while being more computationally efficient. For example, simply replacing the 3x3 convolutions in the bottleneck blocks of the ResNet-50 architecture (He et al., 2016) with lambda layers yields a +1.5% top-1 ImageNet accuracy improvement while reducing parameters by 40% (Section 5.1). • Lambda layers achieve considerable computational benefits, both in latency and memory requirements, over multiple self-attention alternatives, including local and axial attention (Ramachandran et al., 2019; Wang et al., 2020a) . When used in a ResNet-50 architecture at image resolution 224, lambda layers reduce memory consumption by ∼200x compared to global attention (∼7x compared to axial attention) while being ∼3.7x faster than local attention (Section 5.2). • A study of hybrid convolution-lambda models as a means to maximize the speed-accuracy tradeoff (Section 5.3). Hybrid designs that first employ convolutions at the highest resolution and lambda layers in intermediate to low resolutions achieve the best speed-accuracy tradeoff. • LambdaResNets: a family of hybrids based on the training and scaling strategies recommended in Bello et al. (2021) . LambdaResNets achieve up to a 4.4x speed-up over EfficientNets on ImageNet, while being more memory-efficient. LambdaResNets can also be designed for parameter or flops efficiency. For example, a LambdaResNet with 42M parameters achieves 84.3% top-1 ImageNet accuracy at image resolution 320 (Section E.4). et al., 2017) . LambdaResNet backbones yield consistent gains across all metrics on both tasks (e.g. +1.8% mAP improvement for detecting small objects).



An updated version of this paper can be found on arXiv. For example, applying a single multi-head attention layer to a batch of 128 64x64 input images with 8 heads requires 64GB of memory, which is prohibitive in practice.



Figure 1: Comparison between self-attention and lambda layers. (Left) An example of 3 queries and their local contexts within a global context. (Middle) Self-attention associates each query with an attention distribution over context. (Right) The lambda layer transforms each context into a linear function lambda that is applied to the corresponding query.

In large-scale semi-supervised training with an additional 130M pseudo-labeled images, LambdaResNets achieve up to 86.7% top-1 ImageNet accuracy while being 9.5x faster than EfficientNet NoisyStudent (Xie et al., 2020) and 9x faster than a Vision Transformer(Dosovitskiy et al., 2020)  with comparable accuracies (Section 5.3).• An evaluation of LambdaResNets on COCO object detection and instance segmentation using Mask-RCNN (He

