LambdaNetworks: Modeling long-range Interactions without Attention

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 images2 . 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 . 1An updated version of this paper can be found on arXiv . 2For 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 . 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 : • 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 ) . • 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 et al. , 2017 ) . LambdaResNet backbones yield consistent gains across all metrics on both tasks ( e.g . +1.8 % mAP improvement for detecting small objects ) . 2 MODELING LONG-RANGE INTERACTIONS . In this section , we formally define queries , contexts and interactions . Starting from first principles , we motivate keys and relative position embeddings as a requirement for capturing structured interactions between queries and their contexts . We then show that lambda layers arise as an alternative to attention mechanisms for capturing long-range interactions . Notation . We denote scalars , vectors and tensors using lower-case , bold lower-case and bold upper-case letters , e.g. , n , x and X . We denote |n| the cardinality of a set whose elements are indexed by n. We denote xn the n-th row of X . We denote xij the |ij| elements of X . When possible , we adopt the terminology of self-attention to ease readability and highlight differences . 2.1 MOTIVATING QUERIES , KEYS , POSITION EMBEDDINGS AND VALUES . Defining queries and contexts . Let Q = { ( qn , n ) } and C = { ( cm , m ) } denote structured collections of vectors , respectively referred to as the queries and the context . Each query ( qn , n ) is characterized by its content qn ∈ R|k| and position n. Similarly , each context element ( cm , m ) is characterized by its content cm and its position m in the context . The ( n , m ) pair may refer to any pairwise relation between structured elements , e.g . relative distances between pixels or edges between nodes in a graph . Defining interactions . We consider the general problem of mapping a query ( qn , n ) to an output vector yn ∈ R|v| given the context C with a function F : ( ( qn , n ) , C ) 7→ yn . Such a function may act as a layer in a neural network when processing structured inputs . We refer to ( qn , cm ) interactions as content-based and ( qn , ( n , m ) ) interactions as position-based . We note that while absolute positional information is sometimes directly added to the query ( or context element ) content3 , we consider this type of interaction to be content-based as it ignores the relation ( n , m ) between the query and context element positions . Introducing keys and relative position embeddings to capture long-range interactions . In the context of deep learning , we prioritize fast batched linear operations and use dot-product operations as our interactions . This motivates introducing vectors that can interact with the queries via a dotproduct operation and therefore have the same dimension as the queries . In particular , content-based interactions ( qn , cm ) require a |k|-dimensional vector that depends on cm , commonly referred to as the key km . Conversely , position-based interactions ( qn , ( n , m ) ) require a relative position embedding enm ∈ R|k| ( Shaw et al. , 2018 ) . As the query/key depth |k| and context spatial dimension |m| are not in the output yn ∈ R|v| , these dimensions need to be contracted as part of the layer computations . Therefore Every layer capturing long-range interactions can be characterized based on whether it contracts ( 1 ) the query depth or ( 2 ) the context positions first . 2.2 ATTENTION VS LAMBDA LAYERS .. ( 1 ) Attention layers . Contracting the query depth first creates a similarity kernel ( the attention map ) between the query and context elements and is known as the attention operation . As the number of context positions |m| grows larger and the input and output dimensions |k| and |v| remain fixed , one may hypothesize that computing attention maps become wasteful , given that the layer output is a vector of comparatively small dimension |v| |m| . ( 2 ) Lambda layers . Instead , it may be more efficient to simply map each query to its output as yn = F ( ( qn , n ) , C ) = λ ( C , n ) ( qn ) for some linear function λ ( C , n ) : R|k| → R|v| . In this 3This approach is often used in natural language processing tasks ( Vaswani et al. , 2017 ) but has had limited success in the visual domain where relative position information between pixels is crucial ( Bello et al. , 2019 ) . scenario , the context is aggregated into a fixed-size linear function λn = λ ( C , n ) . Each λn acts as a small linear function4 that exists independently of the context ( once computed ) and is discarded after being applied to its associated query qn . 3 LAMBDA LAYERS . 3.1 LAMBDA LAYER : TRANSFORMING CONTEXTS INTO LINEAR FUNCTIONS .. A lambda layer takes the inputs X ∈ R|n|×din and the context C ∈ R|m|×dc as input and generates linear function lambdas that are then applied to the queries , yielding outputs Y ∈ R|n|×dout . Without loss of generality , we assume din = dc = dout = d. As is the case with self -attention , we we may have C = X . In the rest of this paper , we focus on a specific instance of a lambda layer and show that it captures long-range content and position-based interactions without materializing attention maps . Figure 2 presents the computational graph of the lambda layer . We first describe the lambda layer when applied to a single query ( qn , n ) . Generating the contextual lambda function . We wish to generate a linear function R|k| → R|v| , i.e . a matrix λn ∈ R|k|×|v| . The lambda layer first computes keys K and values V by linearly projecting the context , and keys are normalized across context positions via a softmax operation yielding normalized keys K̄ . The λn matrix is obtained by using the normalized keys K̄ and position embeddings En to aggregate the values V as λn = ∑ m ( k̄m + enm ) v T m = K̄ T V︸ ︷︷ ︸ content lambda + ETnV︸ ︷︷ ︸ position lambda ∈ R|k|×|v| ( 1 ) where we also define the content lambda λc and position lambda λpn . • The content lambda λc is shared across all query positions n and is invariant to permutation of the context elements . It encodes how to transform the query qn solely based on the context content . • The position lambda λpn depends on the query position n via the position embeddingEn . It encodes how to transform the query qn based on the context elements cm and their relative positions to the query ( n , m ) . Applying lambda to its query . The query qn ∈ R|k| is obtained from the input xn via a learned linear projection and the output of the lambda layer is obtained as yn = λ T nqn = ( λ c + λpn ) Tqn ∈ R|v| . ( 2 ) 4This mechanism is reminiscent of functional programming and λ-calculus which motivates the lambda terminology . Interpretation of lambda layers . The columns of the λn ∈ R|k|×|v| matrix can be viewed as a fixed-size set of |k| contextual features . These contextual features are aggregated based on the context ’ s content ( content-based interactions ) and structure ( position-based interactions ) . Applying the lambda then dynamically distributes these contextual features based on the query to produce the output as yn = ∑ k qnkλnk . This process captures content and position-based interactions without producing attention maps and can be viewed as an efficient relative attention mechanism . Normalization . One may modify Equations 1 and 2 to include non-linearities or normalization operations . Our experiments indicate that applying batch normalization ( Ioffe & Szegedy , 2015 ) after computing the queries and the values is helpful .

Summary: This paper presents an efficient method to model long-range interaction. The proposed lambda layer removes the nonlinearity of the original attention operation and makes the matrix multiplication independent of the context, hence skipping expensive computation and storage of large attention maps. Two kinds of lambda functions in lambda layer, i.e., content lambda and position lambda,  allows the model to capture both dense content and long-range interaction.  In addition, the lambda layer can be further extended to working with local context and to being more efficient by docomposing a query into multiple short ones. Its effectivess has been demonstrated on extensive experiments on different backbone network architectures and tasks. Its speed-accuracy tradeoff perform very favorably against SOTA methods.