L AMBDA N ETWORKS : M ODELING L ONG -R ANGE I NTERACTIONS W ITHOUT A TTENTION
Anonymous authors
Paper under double-blind review
A BSTRACT
We present a framework for capturing long-range interactions between an input and structured contextual information (e g. a pixel surrounded by other pixels).
Our method, called the lambda layer, captures such interactions by transforming available contexts into linear functions, termed lambdas, and applying these linear functions to each input separately.
Lambda layers may be implemented to model content and position-based interactions in global, local or masked contexts.
As they bypass the need for expensive attention maps, lambda layers can routinely be applied to inputs of length in the thousands, enabling their applications to long sequences or high-resolution images.
The resulting neural network architectures,
LambdaNetworks , are computationally efﬁcient and simple to implement using direct calls to operations available in modern neural network libraries.
Experiments on ImageNet classiﬁcation and COCO object detection and instance segmentation demonstrate that LambdaNetworks signiﬁcantly outperform their convolutional and attentional counterparts while being more computationally efﬁcient.
Finally, we introduce LambdaResNets, a family of architectures that considerably improve the speed-accuracy tradeoff of image classiﬁcation models.
LambdaResNets reach state-of-the-art accuracies on ImageNet while being ∼ 4.5x faster than the popular
EfﬁcientNets on modern machine learning accelerators.
I NTRODUCTION
Modeling long-range interactions is of central importance in machine learning.
Attention (Bahdanau et al., 2015; Vaswani et al., 2017) has emerged as the paradigm of choice for capturing long-range interactions.
However, the quadratic memory footprint of self-attention has hindered its applicability to long sequences or multidimensional inputs such as images which typically contain tens of thousands of pixels.
For example, applying a single multi-head attention layer to a batch of 256 of 64 x 64 input images with 8 heads requires 32 GB of memory, which is prohibitive in practice.
This work presents a class of layers, termed lambda layers, which provide a general framework for capturing long-range interactions between an input and a structured set of context elements.
Lambda layers transform available contexts into individual linear functions, termed lambdas , that are directly applied to each input separately.
We motivate lambda layers as a natural alternative to attention mechanisms.
Whereas attention deﬁnes a similarity kernel between the input and context elements, lambda layers summarize contextual information into a ﬁxed-size linear function, thus bypassing the need for memory-expensive attention maps.
This contrast is illustrated in Figure 1.
We demonstrate the versatility of lambda layers and show that they may be implemented to capture content-based and position-based interactions in global , local or masked contexts.
The resulting neural networks, LambdaNetworks , are computationally efﬁcient, model long-range dependencies at a small memory cost and can therefore be routinely applied to large structured inputs such as high resolution images.
We evaluate LambdaNetworks on computer vision tasks where self-attention has shown promise (Bello et al., 2019; Ramachandran et al., 2019) but has suffered from large memory costs and impractical implementations.
Controlled experiments on ImageNet classiﬁcation and
COCO object detection and instance segmentation indicate that LambdaNetworks signiﬁcantly outperform their convolutional and attentional counterparts while being more computationally efﬁcient and much faster than the latter.
Finally, we introduce LambdaResNets, a family of hybrid LambdaNetworks across different scales, which considerably improve the speed-accuracy tradeoff of im-
[Figure]
age classiﬁcation models.
In particular, LambdaResNets reach state-of-the-art ImageNet accuracies while being 4.5x faster than EfﬁcientNets.
M ODELING LONG - RANGE INTERACTIONS
In this section, we formally deﬁne queries, contexts and interactions.
We motivate keys as a requirement for capturing interactions between queries and their contexts and 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 x n the n -th row of X and { x n } the collection of its | n | rows.
We denote x ij the | ij | elements of X .
When possible, we adopt the terminology of self-attention to ease readability and highlight differences.
Deﬁning queries, contexts and interactions.
Let Q = { ( q n , n ) } and C = { ( c m , m ) } denote structured collections of vectors, respectively referred to as the queries and the context .
Each query ( q n , n ) is characterized by its content q n ∈ R | k | and position n .
Similarly, each context element ( c m , m ) is characterized by its content c m and its position m in the context.
The ( n, m ) pair may refer to any type of pairwise relation between structured elements.
For example, it could refer to the 2D relative distance between pixels arranged in a two-dimensional grid or to edge relations between nodes in a graph.
We consider the general problem of mapping a query ( q n , n ) to an output vector y n ∈ R | v | given the context C with a function F : (( q n , n ) , C ) (cid:55)→ y n .
Such a function may act as a layer in a neural network when processing structured inputs.
We refer to ( q n , c m ) interactions as content-based and ( q n , ( n, m )) interactions as position-based .
Additionally, we say that F captures global interactions when the output y n depends on all ( q n , c m ) (or ( q n , ( n, m )) )
interactions and local when only a restricted smaller context around n is considered.
Finally, these interactions are deﬁned as dense if they include all | m | elements in the context and sparse otherwise.
Introducing keys to capture long-range interactions.
In the context of deep learning, we prioritize fast batched linear operations and choose our interactions to be captured by dot-product operations.
This motivates introducing vectors that can interact with the queries via a dot-product operation and therefore have the same dimension as the queries.
In particular, content-based interactions ( q n , c m ) require a | k | -dimensional vector that depends on c m , commonly referred to as the key k m .
Conversely, position-based interactions ( q n , ( n, m )) require a positional embedding e nm ∈ R | k | , sometimes called a relative key (Shaw et al., 2018).
As the query/key depth | k | and context spatial dimension | m | are not in the output y n ∈ R | v | , these dimensions need to be contracted as part of the layer computations.
Every layer capturing long-range interactions can therefore be characterized based on whether it contracts the query depth or the context positions ﬁrst.
[Table]
Attentional interactions.
Contracting the query depth ﬁrst creates a similarity kernel (the attention map) between the query and context elements and is known as the attention operation.
This mechanism can be viewed as addressing a differentiable memory which motivates the query, key, value terminology.
As the number of context positions | m | grows larger and the input and output dimensions | k | and | v | remain ﬁxed, one may hypothesize that computing attention maps become wasteful, given that the layer output is a vector of comparatively small dimension | v | (cid:28) | m | .
Lambda interactions.
Instead, it may be more efﬁcient to simply map each query to its output via a linear function as y n = F (( q n , n ) , C )
=
λ ( C , n )( q n ) for some linear function λ ( C , n ) .
In this scenario, the context is aggregated into a ﬁxed-size linear function λ n =
λ ( C , n ) .
Each λ n acts as a small linear function that exist independently of the context (once computed) and is discarded after being applied to its associated query q n .
This mechanism is reminiscent of functional programming and λ -calculus which motivates the lambda terminology.
L AMBDA LAYERS
A lambda layer takes the inputs X ∈ R | n |× d in and the context C ∈ R | m |× d c as input and generates linear function lambdas that are then applied to the queries, yielding outputs Y ∈ R | n |× d out .
Note that we may have C = X , as is the case for self-attention.
Without loss of generality, we assume d in = d c = d out = d .
In the rest of this paper, we focus on a speciﬁc instance of a lambda layer and show that it enables dense long-range content and position-based interactions without materializing attention maps.
3. T HE LAMBDA LAYER : T RANSFORMING CONTEXTS INTO LINEAR FUNCTIONS
We ﬁrst describe our lambda layer in the context of a single query ( q n , n ) .
As we wish to generate a linear function lambda R | k | → R | v | , we interchangeably refer to R | k |×| v | matrices as functions.
Hyperparameters, parameters and other quantities of interest of our lambda layer are presented in
Table 1.
Generating the contextual lambda function.
Our lambda layer ﬁrst computes keys and values by linearly projecting the context, and keys are normalized across context positions via a softmax operation yielding normalized keys ¯ K .
The λ n function is obtained by summing contributions from the context as
[Equation] [Equation]
The content lambda λ c = (cid:80) m ¯ K m V Tm is invariant to permutation of the context elements, shared across all query positions n and encodes how to transform the query q n solely based on the context content.
In contrast, the position lambda λ pn = (cid:80) m E nm V
Tm encodes how to transform the query content q n based on the content c m and positions ( n, m ) , enabling modeling structured inputs such as images.
Applying lambda to its query.
The input x n is then transformed into a query q n = W Q x n and the output of the lambda layer is obtained as
[Equation]
Lambda interpretation.
The columns of the λ n ∈ R | k |×| v | matrix can be viewed as a ﬁxed-size set of | k | | v | -dimensional contextual features.
These contextual features are aggregated from the context’s content (content-based interactions) and structure (position-based interactions).
Applying the lambda linear function dynamically distributes these contextual features to produce the output as y n = (cid:80) k
q nk λ nk .
This process captures dense content and position-based long-range interactions without producing attention maps.
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.
3. L AMBDA LAYERS WITH STRUCTURED CONTEXTS
This section presents how to adapt our lambda layer to structured contexts, such as relative and local contexts.
We discuss masked contexts and their applications in the Appendix A.
Translation equivariance Translation equivariance is a strong inductive bias in many learning scenarios.
The content-based interactions are permutation equivariant and hence already translation equivariant.
We obtain translation-equivariance in position interactions by ensuring that the position embeddings satisfy E nm =
E t ( n ) t ( m ) for any translation t .
In practice, we deﬁne a tensor of relative position embeddings R ∈ R | k |×| r |×| u | , where r indexes the possible relative positions for all ( n, m ) pairs, and reindex it into E ∈ R | k |×| n |×| m |×| u | such that E nm =
R r ( n,m ) .
Lambda convolution Despite the beneﬁts of long-range interactions, locality remains a strong inductive bias in many tasks.
Using global contexts may prove noisy or excessive from a computational standpoint.
It may therefore be useful to restrict the scope of position interactions to a local neighborhood around the query position n as is the case for local self-attention and convolutions.
This can be done by zeroing out the position embeddings for context positions m outside of the desired scope.
However, this strategy remains costly for large values of | m | since the computations still occur (they are only being zeroed out).
In the case where the context is arranged on a multidimensional grid, we can generate positional lambdas from local contexts by using a regular convolution that treats the v dimension in V as an extra spatial dimension .
For example, let’s assume we want to generate positional lambdas with local scope size | r | on 1d sequences.
The relative position embedding tensor R ∈ R | r |×| u |×| k | can be reshaped to ¯ R ∈ R | r |× 1 ×| u |×| k | and used as the kernel of a 2d convolution to compute the desired position lambda as
[Equation]
We term this operation the lambda convolution .
As the computations are now restricted to a local scope, the lambda convolution obtains linear time and memory complexities with respect to the input length .
The lambda convolution is readily usable with additional functionalities such as dilation and striding and enjoys highly optimized implementations on specialized hardware accelerators (Nickolls & Dally, 2010; Jouppi et al., 2017).
This is in stark contrast to implementations of local self-attention (Parmar et al., 2018;
Ramachandran et al., 2019) which require materializing feature patches of overlapping query and memory blocks, increasing memory consumption and latency (see Table 4).
[Table]
3. R EDUCING COMPLEXITY WITH MULTIQUERY LAMBDAS .
Complexity analysis.
For a batch of | b | elements, each containing | n | inputs, the number of arithmetic operations and memory footprint required to apply our lambda layer are respectively Θ( bnmkv ) and Θ( knm + bnkv ) .
We still have a quadratic memory footprint with respect to the input length due to the E nm parameters that capture position-based interactions.
However this quadratic term does not scale with the batch size as is the case with the attention operation which produces per-example attention maps.
In practice, the hyperparameter | k | is set to a small value (such as | k | =16) and we can process large batches of large inputs in cases where attention cannot (see Table 4).
Multiquery lambdas reduce complexity.
Recall that the lambdas map queries q n ∈ R k to outputs y n ∈ R d .
As presented in Equation 2, this implies that | v | =
d .
Small values of | v | may therefore act as a bottleneck on the feature vector y n but larger output dimensions | v | can incur an excessively large computational cost given our Θ( bnmkv ) and Θ( bnkv + knm ) time and space complexities.
We propose to decouple the time and space complexities of our lambda layer from the output dimension d .
Rather than imposing | v |
=d, we create | h | queries { q hn } , apply the same lambda function λ n to each query q hn , and concatenate the outputs as y n = concat ( λ n q 1 n , · · · , λ n q | h | n ) .
We refer to this operation as a multiquery lambda layer as each lambda is applied to | h | queries.
This can also be interpreted as constraining the lambda to a smaller block matrix with | h | equal repeated blocks.
We now have d = | hv | and our time and space complexities become Θ( bnmkd/h ) and Θ( knm + bnkd/h ) .
We note that while this resembles the multihead or multiquery (Shazeer, 2019) attention formulation, the motivation is different.
Using multiple queries in the attention operation increases representational power and complexity.
In our case, using multiquery lambdas reduces complexity and representational power.
Table 2 compares time and space complexities of the multiquery lambda layer and the multihead attention operation.
Batched multiquery implementation with einsum.
The multiquery lambda layer can be applied on a batch of inputs using einsum 1 as:
[Equation]
and a reshaping operation Y bnhv → Y bnd .
In the special case | u | = 1 , we work with the squeezed tensors and the indice u can be removed from the einsum equations.
Local positional lambdas may instead be obtained via the lambda convolution as in Equation 3.
[Table] [Table]
R ELATED W ORK
While it has not been explicitly stated [Table] the abstraction of transforming available contexts into linear functions that are applied to queries is quite general and therefore encompasses many previous works.
Closest to our work are channel and linear attention mechanisms which can be cast as less ﬂexible speciﬁc instances of content-only lambda interactions.
Lambda layers formalize and extend such approaches to consider both content-based and position-based interactions, which enables their use as a stand-alone layer on highly structured inputs such as images.
Rather than attempting to closely approximate attention maps as is the case in linear attention formulations, the lambda abstraction shifts the focus to the design of efﬁcient contextual lambda functions.
This enables more ﬂexible non-linearity and normalization schemes and leads to multiquery lambdas.
Controlled experiments demonstrate that lambda layers signiﬁcantly outperform regular and linear attention alternatives while being more computationally efﬁcient.
We discuss related work in details in the
Appendix B
[Table]
E XPERIMENTS
In subsequent experiments [Table] we test LambdaNetworks on standard large-scale high resolution computer vision benchmarks: ImageNet image classiﬁcation task (Deng et al., 2009), COCO object detection and instance segmentation (Lin et al., 2014).
The visual domain is well-suited to showcase the ﬂexibility of lambda layers since i) the memory footprint of self-attention becomes problematic for high-resolution imagery and ii) images are highly structured, making position-based interactions
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
crucial.
We construct LambdaResNets by replacing the 3x3 convolutions in the ResNet architecture (He et al., 2016).
Unless speciﬁed otherwise, all lambda layers use | k | =16, | h | =4 and | u |
= with a scope size of | m | =23x23 and are implemented using the einsum implementation (Equation 4).
All experiments are implemented with Tensorﬂow and code will be open-sourced upon publication.
Experimental details can be found in the Appendix C.
LambdaNetworks outperform convolutions and attentional counterparts.
In Table 3, we perform controlled experiments to compare LambdaNetworks against a) the baseline ResNet50, b) channel attention ( Squeeze-and-Excitation or SE ) and c) prior works that use self-attention to complement or replace the 3x3 convolutions in the ResNet50.
The lambda layer strongly outperforms these approaches at a fraction of the parameter cost and notably obtains a +0.8% improvement over channel attention.
The lambda layer also compares favorably against the linear attention mechanisms used in Shen et al.
; Chen et al.
In Table 4, we compare lambda layers against self-attention and present their throughputs, memory complexities (speciﬁcally the nm term) and ImageNet accuracies.
Our results highlight the weaknesses of self-attention: self-attention cannot model global interactions due to large memory costs, axial self-attention is still memory expensive and local self-attention is prohibitively slow.
In contrast, the lambda layer can capture global interactions on high-resolution images and obtains a +1 . 0% improvement over local self-attention while being almost 3x faster.
Additionally, positional embeddings can be shared across lambda layers to further reduce memory requirements, at [Table] a minimal degradation cost.
Finally, the lambda convolution has linear memory complexity, which becomes practical for very large images as seen in detection or segmentation.
Model ablations We perform several ablations and validate the importance of positional interactions, long-range interactions and ﬂexible normalization schemes in the Appendix E. Tablepresents the impact of the query depth |
k | , number of heads | h | and intra depth | u | on performance.
Our experiments indicate that the lambda layer outperforms convolutional and attentional baselines for a wide range of hyperparameters, demonstrating the robustness of the method.
The lambda layer outperforms local self-attention when controlling for the scope size (78.1% vs 77.4% for | m | =7x7), suggesting that the beneﬁts of the lambda layer go beyond improved speed and scalability.
LambdaResNets signiﬁcantly improve the speed-accuracy tradeoff of ImageNet classiﬁcation.
In the Appendix D, we study and motivate hybrid LambdaNetwork architectures as a mean to maximize the speed-accuracy tradeoff of LambdaNetworks.
The resulting hybrid LambdaResNets architectures have increased representational power at a negligible decrease in throughput compared to their vanilla ResNet counterparts.
We construct hybrid LambdaResNets across various model scales by jointly scaling the depth from 50 to 420 layers and the image size from 224 to 320.
Figure 2 presents the speed-accuracy curve of LambdaResNets compared to ResNets with or without channel attention and the popular EfﬁcientNets
(Tan & Le, 2019).
LambdaResNets outperform the baselines across all depth and image scales with the largest LambdaResNet reaching a state-of-theart accuracy of 84.8.
Most remarkably, LambdaResNets are ∼ 4.5x faster than EfﬁcientNets when controlling for the accuracy and signiﬁcantly improve the speed-accuracy Pareto curve of image classiﬁcation.
Computational efﬁciency.
In Table 5 and Table 6, we ﬁnd that it is also possible to construct LambdaResNets to improve upon the parameter and ﬂops efﬁciency of large EfﬁcientNets.
These results are signiﬁcant because EfﬁcientNets were speciﬁcally designed by neural architecture search (Zoph & Le, 2017) to minimize computational costs using highly computationally efﬁcient depthwise convolutions.
In Appendix F, we show that lambda layers improve the computationalaccuracy tradeoff of mobilenet architectures (Howard et al., 2017).
These results suggest that lambda layers may be well suited for use in resource constrained scenarios such as embedded vision applications.
Object Detection and Instance Segmentation
Lastly, we evaluate LambdaResNets on the COCO object detection and instance segmentation tasks using a Mask-RCNN architecture (He et al., 2017).
Using lambda layers yields consistent gains at all IoU thresholds and all object scales, especially the harder to locate small objects.
This indicates that lambda layers are also competitive for more complex visual tasks that require localization information.
D ISCUSSION
How do lambda layers compare to the attention operation?
Similarly to attention, lambda layers can be implemented to model interactions between inputs and global, local or masked contexts.
However, lambda layers scale favorably compared to self-attention.
Vanilla Transformers using self-attention have Θ( blhn 2 ) spatial complexity, whereas LambdaNetworks have Θ( lkn 2 ) spatial complexity (or Θ( kn 2 ) when sharing positional embeddings across layers).
This enables the use of lambda layers at higher-resolution and on large batch sizes.
Additionally, the lambda convolution (the local version of the lambda layer) enjoys a simpler and faster implementation than local selfattention.
Finally, our ImageNet experiments show that lambda layers outperforms self-attention, demonstrating that the beneﬁts of lambda layers go beyond improved speed and scalability.
In contrast to the attention operation, lambda layers may not capture ﬁne-grained interactions or enable pointing to speciﬁc elements in the context such as in PointerNetworks (Vinyals et al., 2015).
How are lambda layers different than linear attention mechanisms?
Linear attention mechanisms typically attempt to approximate the attention operation by leveraging a low-rank factorization of the attention maps.
Such approaches do not consider position-based interactions and may be unnecessarily constrained in trying to closely approximate an attention kernel.
In contrast, lambda layers can model position-based interactions which are crucial to model highly structured inputs such as images (see Table 10 in Appendix E).
As the aim is not to approximate an attention kernel, lambda layers allow for more ﬂexible non-linearities and normalizations which we also ﬁnd beneﬁcial (see Table 12 in Appendix E).
Finally, we propose multiquery lambdas as a means to reduce complexity compared to the multihead (or single head) formulation typically used in linear attention works.
See Appendix B for a detailed discussion of linear attention.
How to best use lambda layers in the visual domain?
The improved scalability, speed and ease of implementation of lambda layers compared to global or local attention makes them a strong candidate for use in the visual domain.
Our ablations demonstrate that lambda layers are most beneﬁcial in the intermediate and low-resolution stages of vision architectures.
In particular, hybrid LambdaResNets which employ a few lambda layers in these stages show signiﬁcant performance improvements at a negligible increase in latency on modern machine learning accelerators.
Generality of lambda layers.
While this work focuses on static 2D image recognition, we note that lambda layers can be instantiated to model interactions on structures as diverse as graphs, time series, spatial lattices, etc.
We anticipate that lambda layers will be helpful in more modalities.
We discuss masked contexts and auto-regressive tasks in the Appendix A.
Conclusion.
We propose a new class of layers, termed lambda layers, which provide a general and scalable framework for capturing structured long-range interactions between inputs and their contexts.
Lambda layers summarize available contexts into ﬁxed-size linear functions, lambdas, that are directly applied to their associated queries.
The resulting neural networks, LambdaNetworks, are simple to implement, computationally efﬁcient and capture long-range dependencies at a small memory cost, enabling their application to large structured inputs such as high-resolution images.
Extensive experiments on computer vision tasks showcase their versatility and superiority over convolutional and attentional networks.
Most notably, we introduce LambdaResNets which reach state-of-the-art ImageNet accuracies and signiﬁcantly improve the speed-accuracy tradeoff of image classiﬁcation models
A G ENERATING LAMBDAS FROM MASKED CONTEXTS
In some applications such as denoising tasks or auto-regressive training, it may be useful to restrict interactions to a sub-context C n ⊂ C when generating λ n for query position n .
For example, for parallel auto-regressive training, it is necessary to mask the future by ensuring that the output y n only depends on past context positions m < n .
Self-attention achieves this by zeroing out the irrelevant attention weights a nm (cid:48)
=
∀ m (cid:48) / ∈ C n , thus guaranteeing that y n = (cid:80) m a nm v m only depends on C n .
Similarly, we can block interactions between queries and masked context positions when generating lambdas by applying a mask before summing the contributions of context positions.
Using the einsum notation, general masking can be implemented as
[Equation]
where p nm = 1[ m ∈ C n ] and ∗ is a broadcasted element-wise multiplication.
One can also normalize the keys by only considering the elements in their contexts.
Computing masked lambdas still does not require to materialize per-example attention maps and the complexities are the same as for global lambdas case.
B A DDITIONAL RELATED WORK
In this section, we review the attention operation and related works on improving its scalability.
We discuss connections between lambda layers and channel, spatial or linear attention mechanisms and show how they can be cast as less ﬂexible speciﬁc instances of lambda layers. This section assumes | u | =1 to favor simplicity and highlight differences between lambda layers and attention mechanisms. B. S OFTMAX ATTENTION Softmax attention Softmax-attention produces a distribution over the context for each query q n as a n = softmax ( Kq n ) ∈ R | m | where the keys K are obtained from the context C .
The attention distribution a n is then used to form a linear combination of values V obtained from the context as y n = V T a n = (cid:80) m a nm v m ∈ R | v | .
As we take a weighted sum of the values, we transform the query q n into the output y n and discard its attention distribution a n .
This operation captures content-based interactions, but not position-based interactions. Relative attention In order to model position-based interactions, relative attention (Shaw et al., 2018) introduces a learned matrix of | m | positional embeddings E n ∈ R | m |×| k | and computes the attention distribution as a n = softmax (( K + E n ) q n ) ∈ R | m | . The attention distribution now also depends on the query position n relative to positions of context elements m .
Relative attention therefore captures both content-based and position-based interactions.
B. A TTENTION WITH SPARSE PATTERNS
A signiﬁcant challenge in applying (relative) attention to large inputs comes from the quadratic Θ( | bnm | ) memory footprint required to store attention maps.
Many recent works therefore propose to impose speciﬁc patterns to the attention maps as a means to reduce the context size | m | and therefore the memory footprint of the attention operation.
These approaches include local attention patterns (Dai et al., 2019; Parmar et al., 2018;
Ramachandran et al., 2019), axial attention patterns (Ho et al., 2019; Wang et al., 2020a), static sparse attention patterns (Child et al. ; Beltagy et al., 2020) or dynamic sparse attention patterns (Kitaev et al., 2020). See Tay et al. (2020) for a review.
Their implementations can be rather complex, sometimes require low-level kernel implementations to get computational beneﬁts or may rely on speciﬁc assumptions on the shape of the inputs (e.g., axial attention).
In contrast, lambda layers are simple to implement for both global and local contexts using simple einsum and convolution primitives and capture dense content and position-based interactions with no assumptions on the input shape.
B. L INEAR ATTENTION
Another approach to reduce computational requirements of attention mechanisms consists in approximating the attention operation in linear space and time complexity, which is referred to as linear (or efﬁcient) attention.
Linear attention mechanisms date back to de Br´ebisson & Vincent (2016); Britz et al. and were later introduced in the visual domain by Chen et al.
Shen et al.
They are recently enjoying a resurgence of popularity with many works modifying the popular Transformer architecture for sequential processing applications (Katharopoulos et al., 2020;
Wang et al., 2020b; Choromanski et al., 2020).
Linear attention via kernel factorization Linear attention is typically obtained by reinterpreting attention as a similarity kernel and leveraging a low-rank kernel factorization as
[Equation] 6) for some feature function φ .
Computing φ ( K T ) V ∈ R | k |×| v | ﬁrst bypasses the need to materialize the attention maps φ ( Q ) φ ( K T ) and the operation therefer has linear complexity with respect to the input length | n | . Multiple choices for the feature function φ have been proposed. For example, Katharopoulos et al.
(2020) use φ ( x ) = elu ( x ) + 1 , while Choromanski et al. (2020) use positive orthogonal random features to approximate the original softmax attention kernel. (2018) and Shen et al. (2018) use φ ( x ) = softmax ( x ) .
This choice is made to guarantee that the rows of the (non-materialized) attention maps φ ( Q ) φ ( K ) T sum to 1 as is the case in the regular attention operation.
B. L AMBDA LAYERS VS LINEAR ATTENTION
We now discuss the main differences between lambda layers and linear attention mechanisms.
1) Lambda layers extend linear attention to also consider position-based interactions.
The kernel approximation from Equation 6 can be rewritten for a single query q n as
[Equation]
which resembles the output of the content lambda y
cn =
( λ c
[Equation]
n from Equationsand 2. Lambda layers extend linear attention mechanisms to also consider position-based interactions as [Equation] assuming | u | =1.
In the above equation, computing the position (or content) lambda has Θ( | bmkv | ) time complexity.
As the position lambdas are not shared across query positions n , this cost is repeated for all | n | queries, leading to a total time complexity Θ( | bnmkv | ) .
Unlike linear attention mechanisms, lambda layers have quadratic time complexity with respect to the input length (in the global context case) because they consider position-based interactions.
2) Lambda layers do not necessarily attempt to approximate an attention kernel.
While approximations of the attention kernel are theoretically motivated, we argue that they may be unnecessarily restrictive.
For example, the kernel approximation in Equation 6 requires the same feature function φ on both Q and K and precludes the use of more ﬂexible non-linearities and normalization schemes.
In contrast, lambda layers do not attempt to approximate an attention kernel.
This simpliﬁes their design and allows for more ﬂexible non-linearity and normalization schemes, which we ﬁnd useful in our ablations (See Table 12 in Appendix E).
Considering the position embeddings independently of the keys notably enables a simple and efﬁcient local implementation with the lambda convolution.
Approximating the relative attention kernel would require normalizing the position embeddings with the keys (i.e., φ ( K + E n ) instead of φ ( K )+ E n ), which cannot be implemented in the local context
case
with a convolution. The lambda abstraction reveals the computational beneﬁts of using a multiquery formulation. Finally, this work proposes to abstract the ¯ K T V and E Tn V matrices as linear functions (the content and position lambdas) that are directly applied to the queries.
The lambda abstraction reveals the beneﬁts of multiquery lambdas (as opposed to the traditional multi-head attention formulation) as a means to reduce computational costs.
B. C HANNEL AND SPATIAL ATTENTION .
Finally, we show that the lambda abstraction generalizes channel and spatial attention mechanisms, both of which can be viewed as speciﬁc instances of lambda layers. This observation is consistent with our experiments which demonstrate that lambda layers outperform both channel and spatial attention while being more computationally efﬁcient. Channel attention Channel attention mechanisms, such as Squeeze-and-Excitation (SE) and FiLM layers, recalibrate features via cross-channel interactions by aggregating signals from the entire feature map (Hu et al., 2018b;a; Perez et al., 2017).
In particular, the SE operation can be written as y nk = w k q nk where w k is the excitation weight for channel k in the query q n .
This can be viewed as using a diagonal lambda which is shared across query positions λ n = diag ( w 1 · · · w | k | ) .
Channel attention mechanisms have proven useful to complement convolutions but cannot be used as a stand-alone layer as they discard spatial information.
Spatial attention Conversely, spatial attention mechanisms, reweigh each position based on signals aggregated from all channels (Xu et al., 2015; Park et al., 2018;
Woo et al., 2018).
These mechanisms can be written as y nk = w n q nk where w n is the attention weight for position n in the input query Q .
This can be viewed as using (position-dependent) scalar lambdas λ n = w n I where I is the identity matrix.
Spatial attention has also proven helpful to complement convolutions but also cannot be used as a stand-alone layer as it discards channel information.
B. H YPERNETWORKS
LambdaNetworks can alternatively be viewed as an extension of HyperNetworks (Ha et al., 2016) that dynamically generate their computations based on the inputs contexts.
C E XPERIMENTAL DETAILS
ResNets.
We use the ResNet-v1 implementation and initialize the γ parameter in the batch normalization (Ioffe & Szegedy, 2015) layer at the end of the bottleneck blocks to 0 .
Squeeze-andExcitation layers employ a squeeze ratio of 4.
Lambda layer implementation details Unless speciﬁed otherwise, all lambda layers use query depth | k | =16, | h
|
=
4 heads and intra-depth | u | =1.
The position lambdas are generated with local contexts of size | m | =23x23 and the content lambdas with the global context as described in Equation 4.
When the intra-depth is increased to | u | > 1, we reduce the scope of size | m | =7x7 and switch to the convolution implementation to reduce ﬂops.
The projections to compute Q and V are followed by batch normalization and K is normalized via a softmax operation.
Positional embeddings are initialized at random using the unit normal distribution.
Local positional lambdas can be implemented interchangeably with the lambda convolution or by using the global einsum implementation from
Equation 4 and masking the position embeddings outside of the local contexts.
The latter can be faster but has a higher memory footprint and FLOPS due to the Θ( knm ) term (see Equation 4).
In our experiments, we use the convolution implementation only for input length | n | > 85 2 or intra-depth | u | > 1 .
LambdaResNets.
We construct our LambdaResNets by replacing the spatial (3x3) convolutions in ResNet architectures by our proposed lambda layer, with the exception of the stem which is left unchanged.
We apply 3x3 average-pooling with stride 2 after the lambda layers to downsample in place of the strided convolution.
The number of residual blocks per stage for the deeper ResNets are [4, 29, 53, 4] for ResNet-270, [4, 36, 72, 4] for ResNet-350, and [4, 44, 87, 4] for ResNet-420.
When working with hybrid LambdaNetworks, we use a single lambda layer in c4 for LambdaResNet50, 3 lambda layers for LambdaResNet101, 6 lambda layers for LambdaResNet-152/200/270/350 and 8 lambda layers for LambdaResNet-420.
Lambda layers are uniformly spaced in the c4 stage for hybrid architectures.
ImageNet training setups.
We consider two training setups for the ImageNet classiﬁcation task.
The 90 epochs training setup trains models for 90 epochs using standard preprocessing and allows for fair comparisons with classic works.
The 350 epochs training setup trains models for 350 epochs using improved data augmentation and regularization and is closer to training methodologies used in modern works with state-of-the-art accuracies.
ImageNet 90 epochs training setup.
We use the vanilla ResNet for fair comparison with prior works.
We used the default hyperparameters as found in ofﬁcial implementations without doing additional tuning.
All networks are trained end-to-end for 90 epochs via backpropagation using
SGD with momentum 0.9.
The batch size B is 4096 distributed across 32 TPUv3 cores (Jouppi et al., 2017) and the weight decay is set to 1e-4.
The learning rate is scaled linearly from 0 to 0.1B/256 for 5 epochs and then decayed using the cosine schedule (Loshchilov & Hutter, 2017).
We use batch normalization with decay 0.9999 and exponential moving average with weight 0.
over trainable parameters and a label smoothing of 0.1.
The input image size is set to 224x224.
We use standard training data augmentation (random crops and horizontal ﬂip with 50% probability).
Most papers compared against in Table 3 use a similar training setup and also replace the 3x3 spatial convolutions in ResNet architectures by their proposed methods.
This allows for a fair comparison.
ImageNet 350 epochs training setup.
Higher accuracies on ImageNet are commonly obtained by training longer with increased augmentation and regularization (Lee et al., 2020; Tan & Le, 2019).
In the 350 epochs training setup, we replace the baseline architecture with the ResNet-D (He et al., 2018) and use squeeze-and-excitation in the residual blocks that do not employ lambda layers for the hybrid LambdaResNets.
We additionally replace the max pooling layer in the stem by a strided 3x convolution.
Networks are trained for 350 epochs with a batch size B of 4096 or 2048 distributed across 32 or 64 TPUv3 cores, depending on memory constraints.
We employ RandAugment (Cubuk et al., 2019) with a magnitude of 15 as the data augmentation strategy.
We use a smaller weight decay of 4e-5 and dropout with a drop probability of 0.3.
All architectures deeper than ResNetare trained with stochastic depth with a drop probability of 0.2.
Tuning Each training setup uses a constant set of hyperparameters across model scales.
The improved 350 epoch training setup was found by tuning the baseline architectures to identify a robust training setup across different scales.
While individual accuracies may be improved with further tuning, we favor simplicity and use the same training hyperparameters for all experiments.
We do not perform early stopping and simply report the ﬁnal accuracies.
Throughputs.
Figure 2 reports the latency to process a batch of 4096 images on 32 TPUv3 cores using mixed precision training (ı.e bﬂoat16 activations).
Table 4, Table 8 and Table 9 report inference throughput on 8 TPUv3 cores using ﬂoat32 precision.
FLOPS count.
We do not count zeroed out ﬂops when computing positional lambdas with the einsum implementation from Equation 4.
Flops count is highly dependent on the scope size which is rather large by default ( | m | =23x23).
In Table 11, we show that it is possible to signiﬁcantly reduce the scope size and therefore FLOPS at a minimal degradation in performance.
Computational efﬁciency.
In these experiments, we replace the last two stages of the ResNet architecture (where the convolutions are the most computationally expensive) with lambda layers.
The parameter-efﬁcient LambdaResNets in Table 5 employ an image size of 320.
For ﬂops efﬁciency, we additionally reduce the lambda scope size to | m | =7x7 and set the image size to 256.
COCO object detection.
We employ the architecture from the improved ImageNet training setup as the backbone in the Mask-RCNN architecture.
All models are trained on 1024x1024 images from scratch for 130k steps with a batch size of 256 distributed across 128 TPUv3 cores with synchronized batch normalization.
We apply multi-scale jitter of [0.1, 2.0] during training.
The learning rate is warmed up for 1000 steps from 0 to 0.32 and divided by 10 at steps 90, 95 and 97.5% of training.
The weight decay is set to 4e-5.
D H YBRID L AMBDA N ETWORKS
In Table 8 and Table [Table] we study the throughput and accuracy of hybrid LambdaNetwork architectures.
We ﬁnd that lambda layers are most helpful in the last two stages of the ResNet architecture (commonly referred to as c4 and c5 )
when considering the speed-accuracy tradeoff (see Table 8).
In particular, lambda layers in the c5 stage incur almost no speed decrease compared to 3x3 convolutions.
Lambda layers in the c4 stage are relatively slower than convolutions but are crucial to reach high accuracies.
In Table
[Table]
we test how the speed and ﬁnal accuracy is impacted by the number of lambda layers in the c4 stage.
Our results reveal that most beneﬁts from lambda layers can be obtained by 1) replacing a few 3x3 convolutions with lambda layers in the second last stage (commonly referred to as c4 ) of the ResNet architecture and 2) replacing all 3x3 convolutions in the last stage ( c5 ).
The resulting hybrid LambdaResNets architectures have increased representational power at a virtually negligible decrease in throughput compared to their vanilla ResNet counterparts [Table] [Table] [Table]
[Table]
[Table]
E A BLATIONS
Content vs position interactions Table 10 presents the relative importance of content-based and position-based interactions on the ImageNet classiﬁcation task.
As expected, position-based interactions are necessary to reach high accuracies, while content-based interactions only bring marginal improvements over position-based interactions. Importance of scope size The small memory footprint of LambdaNetworks enables considering global contexts, even in the early high resolution layers of the networks. Table 11 presents ﬂops counts and top-1 ImageNet accuracies when varying scope sizes for LambdaResNet50 on 224x inputs. We ﬁnd beneﬁts from using larger scopes, with a plateau around | m | =15x15, which validates the importance of long-range interactions.
We choose | m | =23x23 as the default to account for experiments that use larger image sizes.
Normalization Table 12 ablates normalization operations in the design of the lambda layer.
We ﬁnd that normalizing the keys is crucial for performance and that other normalization functions besides the softmax can be considered.
Additionally, applying batch normalization to the queries and values is also helpful
[Table]
[Table]
[Table]
Varying query depth and number of heads. In Table [Table] we study the impact of query depth | k | , number of heads | h | and intra-depth | u | on the accuracy.
Our experiments indicate that LambdaNetworks outperform the convolutional and attentional baselines for a wide range of | k | and | h | hyperparameters. As expected, we get additional improvements by increasing the query depth | k | or intra-depth | u | . The number of heads is best set at intermediate values such as | h | =4. A large number of heads | h | excessively decreases the value depth | v | = d/ | h | , while a small number of heads translates to too few queries, both of which hurt performance.
F L AMBDA LAYERS IN A RESOURCE CONSTRAINED SCENARIO
In this section, we study lambda layers in a resource-constrained scenario using the MobileNetv2 and
MobileNetv3 architectures (Sandler et al., 2018; Howard et al., 2019).
These architectures employ lightweight inverted bottleneck blocks which consist of the following sequence: 1) a pointwise convolution for expanding the number of channels, 2) a depthwise convolution for spatial mixing and 3) a ﬁnal pointwise convolution for channel mixing.
The use of a depthwise convolution (as opposed to a regular convolution) reduces parameters and ﬂops, making inverted bottlenecks particularly well-suited for embedded applications.
Lightweight lambda layers.
We obtain a lightweight lambda block as follows.
We replace the depthwise convolution in the inverted bottleneck with a lambda layer.
The lightweight lambda layer uses query depth | k | =32, | h | =4 and a smaller scope size | m | =5x5 to reduce computational demand.
Finally, we change the ﬁrst pointwise convolution to output the same number of channels (instead of increasing the number of channels) to reduce computations.
MobileNet experiments We wish to assess whether lambda layers can improve the ﬂops-accuracy (or parameter-accuracy) tradeoff of mobilenet architectures.
We therefore experiment with hybrid
LambdaNetworks and replace only a few inverted bottlenecks so that the resulting architectures have similar computational demands as their baselines (see Table 14).
For MobileNetv2, this simple strategy reduces parameters and ﬂops by ∼ 10% while improving ImageNet accuracy by 0.6%.
The improvement is smaller for the MobileNetv3 architecture.
However we note that the MobileNetv architecture is a local optimum obtained via a combination of automated architecture search and human design assuming only access to inverted bottleneck blocks.
We hypothesize that adding lambda layers in the search space might lead to stronger improvements. [Table]
Mobilenet training setup.
All mobilenet architectures are trained for 350 epochs on Imagenet with standard preprocessing at 224x224 resolution.
We use the same hyperparameters as Howard et al.
More speciﬁcally, we use RMSProp with 0.9 momentum and a batch size of 4096 split {+across 
+}
TPUv3 cores.
The learning rate is warmed up linearly to 0.1 and then multiplied by 0.
every 3 epochs.
We use a weight decay 1e-5 and dropout with drop probability of 0.