LambdaNetworks: Modeling long-range Interactions without Attention
Authors: Irwan Bello
Venue: ICLR 2021 Spotlight (2021)
Keywords: deep learning
neural networks
attention
transformer
vision
image classification

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 additional 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.

Published as a conference paper at ICLR 2021
LAMBDANETWORKS: MODELING LONG-RANGE INTERACTIONS WITHOUT ATTENTION
Irwan Bello
Google Research, Brain team ibello@goog com
ABSTRACT
We present lambda layers an alternative framework to self-attention - for cap- turing long-range interactions between an input and structured contextual infor- mation (e.g. a pixel surrounded by other pixels). Lambda layers capture such interactions by transforming available contexts into linear functions, termed lamb- das, 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 appli- cation large structured inputs such 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. Addition- ally, we design Less a family of hybrid architectures across different scales, that considerably improves the speed-accural tradeoff of image classi- fication models. LambdaResNets reach excellent accuracies on ImageNet while being 3.2 4.4x faster than the popular EfficientNets on modern machine learn- ing accelerators. In large-scale semi-supervised training with an additinal 130M pseudo-labeled images, LesbdaResNets achieve up to 86.7 6.77% ImageNet accuracy while being 9.5x faster than EfficientNet NoisyStudent and 9x faster than a Vision Transformer with comparable accuracies!
INTRODUCTION
Modeling long-range dependencies data a central problem machine learning. Self- attention (Bahdanau al., 2015; Vaswani et al., 2017) has emerged popular approach do so, but the costly memory requirement of self-attention hinders its application to long sequences and multidimensional data such images². 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 graph.
This work addresses both issues. We propose lambda layers which model long-range interactions between a query and structured set of context elements at reduced memory cost. Lambda layers transform each available context into linear function, termed lambda, which 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.c. matrix), thus bypassing the need for memory- -nntesive attention maps. This difference 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, LambdaNenworks, computationally efficient, model range dependencies at small memory cost and can there- fore be applied to large structured inputs such as high resolution images.
updated version of this paper can found on arXiv. For example, applying single multi-head attention layer to batch of 128 64x64 input images with 8 heads requires 64GB of memory, which is prohibitive practice.

---

Published as a conference paper at ICLR 2021
GLOBAL CONTEXT
QUERIES
LOCAL CONTEXTS
ATTENTION MAPS
LAMBDAS
Figure 1: Comparison between self-attention and lambda layers. (Leff) An example of queries and their local contexts within a global context. (Middle) Self-attention associates each query with an attention distribution over its context. (Right) The lambda layer transforms each context into linear function lambda that is applied the corresponding query.
We evaluate LambdaNetworks on computer vision tasks where works using self-attention are hin- dered by large memory costs (Wang et al., 2018; Bello et al., 2019), suffer impractical implementa- tions (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, Letwoosetworrs 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- content-based and position-based interac- tions without materializing attention maps. Lambda layers offer a unifying view of chan- nel, 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 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 the bottleneck blocks of the ResNet-50 architec- ture (He et al., 2016) with lambda layers yields a +1.5% top-1 ImageNet accuracy improve- ment while reducing parameters by 40% (Section 5.
Lambda layers achieve considerable computational benefits, both in latency and mem- ory requirements, over multiple self-attention alternatives, including local and axial at- tention (Ramachandran al., 2019; Wang et al., 2020a). When used in ResNet-50 chitecture 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-lante lambda models as means to maximize the speed- accuracy tradeoff (Section 5.3). Hybrid designs that first employ convolutions the highest resolu- tion and lambda layers in intermediate to low resolutions achieve the best speed-accurage tradeoff.
LestdaResNets: family of hybrids based on the training and scaling strategies rec- ommended in Bello (2021). LambdaResNets achieve up to a 4.4x speed-up over EfficientNets ImageNet, while being more memory-efficiess designed for parameter or flops efficiency. For example, a LembdaResNet with 42M pa- rameters achieves 84.3% top-1 ImageNet accuracy image resolution 320 (Section E.4).
LambdaResNets can also
In large-scale semi-supervised training with an additional 130M pseudo-labeled images, Nest 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 Vision Trans- former (Dosovitskiy et al., 2020) with comparable accuracies (Section 5.3).
An evaluation of LestdaResNets using Mask-RCNN (He et al., 2017). LembdaResNet backbones yield consistent gains across all metrics both tasks (e.g. +1.8% mAP improvement for detecting small objects).
COCO object detection and instance segmentation
2

---

Published as a conference paper at ICLR 2021
A content-based interaction considers the content the context but ignores the relation between the query position and the context (c.g relative distance between two pixels). A position-based interaction considers the relation between the query position and the context position.
Table 1: Definition of content-based vs position- interactions.
2 MODELING LONG-RANCE 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 inter- actions between queries and their contexts. We then show that lambda layers arise 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, & and X. We denote |n| the cardinality of a set whose elements are indexed by We denote nn the n-th row of X. We denote Tıy the ıijl 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 = {(gn: n)} and C = {(mm, m)} denote structured col- lections of vectors, respectively referred to as the queries and the context. Each query (gn, n) is characterized by its content 4n E R|| 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 graph.
Defining interactions. We consider the general problem of mapping a query (qn, n) to an output vector yn E glel given the context C with function F ((gm, n), + yn. Such function may act as a layer in a neural network when processing structured inputs.
We refer to (qn Cm) interactions content-based and (gn: (n, m)) interactions as position-based We note that while absolute positional information is sometimes directly added to the query (or context element) content', we consider this type of interaction to be content-band baseet as 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 interactions. This motivates introducing vectors that can interact with the queries via dot- product operation and therefore have the same dimension as the queries. In particular, content-based interactions (qn, Cm) require a [kll-dimeniional vector that depends on Cm. commonly referred to as the key km. Conversely, position- interactions (gn, (n, m)) require a relative position em- bedding e.m E Rlkk (Shaw et al., 2018). As the query/key depth [k| and context spatial dimension |m| are not in the output yn E IR "l, 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 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 contextelemmnts and is known the attention operation. As the number of context positions m grows larger and the input and output dimensions [k] and |e| remain fixed, one may hypothesize that computing attention maps become wasteful, given that the layer output a vector of comparatively small dimension |ol << m/]
(2) Lambda layers. Instead, it may be more efficient to simply map each query to its output as yn = F((gn. n),C) C) = X(C, (g)))) for some linear function X(C. n) Rlkl + R! In this
TThis 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).
3

---

Published as a conference paper at ICLR 2021
E1
Name
Descrittin
AP
query, value depth
E
X E
lnguts
AP
C exami
context
A
X
XWO RR
K CWK cRIm/xIR
Reys
V CWV E 2lmix o(K) softman K. mism)
valaes
normalized keys
A
B, Rlmix|RI
relative pusition
embeddings
X R
conteor lambda
X E V
position
Inpots
A x +
lambdas.
Generating lambdes from the context
Applying
queries
Figure 2: Computational graph of the lambda layer. Contextual information for query position summarized into lambda X, E Riiiill Applying the lambda dynamically distributes textual features to produce the output as yv = X In- This process captures content-based and position-based interactions without producing attention maps.
scenario, the context is aggregated into fixed-size linear function A,, X),, Each A,, acts as a small linear function" that exists independently of the context (once computed) and is discarded after being applied its associated query Qn-
LAMBDA LAYERS
3.1 LAMBDA LAYER: TRANSFORMING CONTEXTS INTO LINEAR FUNCTIONS.
lambda layer takes the inputs X E Ii/xxdd and the context C' E R mmim/xde as input and gener- ates linear function lambdas that are then applied to the queries, yielding outputs Y E IR 1 m||xdoue Without loss of generality,
assume dnn de dout d. As is the with self-attention, we may have C = X. In the rest of this paper, we focus on a specific instance of lambda layer and show that captures long-range content and position-basee interactions without materializing attention maps. Figure 2 presents the computational graph of the lambda layer.
We first describe the lambda layer when applied single query (gn, nt)
Generating the contextual lambda function. We wish generate linear function [R Rlal, iee. matrix An R[k]x|ol, 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 An matrix is obtained by using the normalized keys K and position embeddings E, to aggregate the values V as
A, Σ (km + enm) m
ETV
ER|MX|
(1)
content lambda
position lamhda
where we also define the content lambda X and position lambda X;
The content lambda X is shared across all query positions n and is invariant to permutation of the context elements. It encodes how to transform the query solely based on the context content.
The position lambda X depends the query position via the position embedding E,, It encodes how to transform the query 4n based on the context elements Cm and their relative positions to the query (n, m).
Applying lambda its query. The query qn glk is obtained from the input wn via learned linear projection and the output of the lambda layer is obtained as
Un XT qm (X
XP)T q. E 1o
"This mechanism is reminiscent of functional programming and X-calculus which motivates the lambda terminology.
4

---

Published as a conference paper at ICLR 2021
Interpretation of lambda layers. The columns of the X, E | matrix can be viewed as a fixed-size set of |kj contextual features. These contextual features are aggregated based on the context's content (content-basee interactions) and structure (position-basets interactions). Applying the lambda then dynamically distributes these contextual features based on the query to produce the output as yn >n Oak-Ant- 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 and 2 to include non-lineen lineaiiasittes normalization operations. Our experiments indicate that applying batch normalization (loffe & Szegedy, 2015) after computing the queries and the values is helpful.
3.2 A MULTI-QUERY FORMULATION TO REDUCE COMPLEXITY.
Complexity analysis. For a batch of |b] examples, each containing |n| inputs, the number of arithmetic operations and memory footprint required to apply our lambda layer are respectively ((bmmke) and @(knm + bnkv). We still have a quadratic memory footprint with respect to the in- put length due to the pmm relative position embeddings. 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 small value (such as [k]==66 and we can process large batches of large inputs in cases where attention cannot (see Table 4). Additionally, position embeddings can be shared across lambda layers to keep their (krm) memory footprint constant whereas the memory footprint of attention maps scales with the number of layers".
def lambda (queries, keys, embeddinge, valuea):
" -Muutiiquerry lambda layer. "ue
b: batch, n: input length, m: context length,
x: query/key depth, v: value depth, h: number of heads, d: output dimension. content lambda eineum (keys), values, 'bak, bav->>bkr) position.lambdas
lambdas einsum (embeddings, values, 'nnk, bawbbbvvbbbaaaae >bnk'') .ouppt einsum (nrsum(qquriaa (queries, content lambda,, 'bhnk, >bnhv')
position.output = einsum (queries, position.lambdas, 'bhnk, bnkv mank-->ban >bnh') output reshape (content output + position output, [b, n, d]) return output
Figure 3: Pseudo-code for the multi-query lambda layer. The position embeddings can be made to satisfy various conditions, such as translation equivariance, when computing positional lambdas (not shown). The lambda layer can be adapted to other tasks/modalities by adjusting the choice of embeddings (Section A.2).
Multi-query lambda layers reduce time and space complexities. Recall that the lambda layer maps inputs Fr I to outputs Yn E Rd. As presented in Equation 2, this implies that ju|=d. Small values of [o] may therefore act as bottleneck on the feature vector ye but larger output dimensions |u| can incur excessively large computational cost given our (bbremko) and O(krm + bnke) time and space complexities.
We propose to decouple the time and space complexities of our lambda layer from the output di- mension d. Rather than imposing |u|=d, we create |4] queries { apply the same lambda A to each query qh and concatenate the outputs yn Concat(A..gl....
An (l). We now have
|u|=d/| hb). which reduces complexity by a factor of [/]. The number of heads [h| controls the size of the lambdas An E IR x|x|d//h| relative to the total size of the queries qn E [kk
We refer to this operation as multi-query lambda layer and present an implementation using einsum in Figure 3. The lambda layer is robust to [k] and /h| hyperparameter choices (see Ap- pendix E.1), which enables flexibility in controlling its complexity. We use //|=4 in most experi-
Attention maps typically need be stored for propagatonnt (Kitaev et al., 2020).
'The einsum operation denotes general contractions between tensors of arbitrary dimensions. It is numeri- cally equivalent to broadcasting its inputs to share the union of their dimensions, multiplying element- and summing across all dimensions specified in the output.
5

---

Published as a conference paper at ICLR 2021
Head configuration
Time
Space complexity
Operation
Interactions
complexity
Attention
multi-head multi-head multi-heas
content-only content position
(bnm + d)) e (bmm + d))
(bhnm)
Relative attention
(hhmm)
Linear attention
content-only
((brkd)
(
Lambda layer
multi-query content position content & position
+ (bmmkd/h) (bnrkd/hh)
( bokd/h h) (kr + bnkd/ h))
Lambda convolution
multi-query query
Table 2: Alternatives for capturing long-range interactions. The lambda layer captures content and position-based interactions at reduced memory cost compared to relative attention (Shaw et al., 2018: Bello al., 2019). Using a multi-query lambda layer reduces complexities by factor of [h] Additionally, position-based interactions can be restricted to local scope by using the lambda convolution which has linear complexity. b: batch size, h; number of heads/queries, n: input length, context length, local scope size, k: query/key depth, d: dimension output.
ments. We note that while this resembles the multi head or multi-query (Shazeer, 2019)? attention formulation, the motivation is different. Using multiple queries in the attention operation increases representational power and complexity. In contrast, using multiple queries in the lambda layer de- creases complexity and representational power (ignoring the additional queries).
Extending the multi-query formulation to linear attention. Finally, we point that our analysis extends to linear attention which can be viewed as a content-only lambda layer (see Appendix D.3 for detailed discussion). We anticipate that the multi-query formulation can also bring computational benefits to linear attention mechanisms.
3.3 MAKING LAMBDA LAYERS TRANSLATION EOCTVARIANT.
Using relative position embeddings e. enables making explicit assumptions about the structure of the context. particular, translation equivariance (i.c. the property that shifting the inputs results in equivalent shift of the outputs) is strong inductive bias in many learning scenarios. We obtain translation equivariance in position interactions by ensuring that the position embeddings satisfy C.m C((n)((m) for any translation t. In practice, define tensor of relative position embeddings E |||[[|]], where r indexes the possible relative positions for all (n, m) pairs, and reindex* it into E |R|x|x|||x||| such that Exm = Tr(n.m)
3.4 LAMBDA CONVOLUTION: LOCAL CONTEXTS THE GRID.
Despite the benefits of long-range interactions, locality remains a strong inductive bias in many tasks. Using global contexts may prove noisy or computationally excessive. It may therefore be useful to restrict the scope of position interactions to 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 relative embeddings for context positions 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.
Lambda convolution In the case where the context is arranged in multidimensional grid, can equivalently compute positional lambdas from local contexts by using regular convolution. We term this operation the lambda convolution. A n-dimensional lambda convolution be im- plemented using an n-d depthwise convolution with channel multiplier or (n+1) convolution that treats the dimension in V extra spatial dimension. We present both implementations in Appendix C.I.
As the computations are now restricted local scope, the lambda convolution obtains linear time and memory complexities with respect to the input length". The lambda convolution is readily
7 (Shazeer, 2019) proposes a multi-query formulation to speed-up attention-based decoding. We refer the reader the code for more details.
'FLOPs (time complexity) is not necessarily good proxy for latency on TPUs/GPUs. Eventhough the lambda convolution has linear time/space complexities, it can be slower than than the global lambda layer in practice, especially for large convolution scope sizes. See Table for example.

---

Published as a conference paper at ICLR 2021
usable with additional functionalities such as dilation and striding and enjoys optimized implemen- tations on specialized hardware accelerators (Nickolls & Dally, 2010; Jouppi et al., 2017). This is in stark contrast to implementations of local self-attention that require materializing feature patches of overlapping query and context blocks (Parmar et al., 2018; Ramachandran et al., 2019), increasing memory consumption and latency (see Table 4).
4 RELATED WORK
Table 2 reviews alternatives for capturing long-range interactions and contrasts them with the pro- posed multi-query lambda layer. We discuss related works in details in the Appendix D.
Channel and linear attention The lambda abstraction, i.e. transforming available contexts into linear functions that are applied to queries, is quite general and therefore encompasses many pre- vious works. Closest to our work are channel and linear attention mechanisms (Hu et al., 2018c; Katharopoulos et al., 2020; Choromanski et al., 2020). Such mechanisms also capture long-range in- teractions without materializing attention maps and can be viewed as specific instances of content- only lambda layer. Lambda layers formalize and extend such approaches to consider both content- based and position-based interactions, enabling their use as a stand-alone layer on highly structured data such images. Rather than attempting to closely approximate an attention kernel as is the case with linear attention, we focus on the efficient design of contextual lambda functions and repurpose multi-query formulation (Shazeer, 2019) to further reduce computational costs.
Self-attention in the visual domain In contrast to natural language processing tasks where it is now the de-facto standard, self-attention has enjoyed steady but slower adoption in the visual domain (Wang et al., 2018; Bello et al., 2019; Ramachandran et al., 2019; Carion et al., 2020). Concurrently to this work, Dosovitskiy et al. (2020) achieve a strong 88.6% accuracy on ImageNet by pre-training Transformer on sequences of image patches large-scale dataset of 300M images.
EXPERIMENTS
In subsequent experiments, we evaluate lambda layers on standard computer vision benchmarks: ImageNet classification (Deng al., 2009), COCO object detection and instance segmentation (Lin et al., 2014). The visual domain is well-suited to showcase the flexibility of lambda layers since (1) the memory footprint of self-attention becomes problematic for high-resolution imagery and (2) images are highly structured, making position- interactions crucial.
LesbdaResNets We construct Ness by replacing the 3x3 convolutions in the bottle- neck blocks of the ResNet architecture (He et al., 2016). When replacing all such convolutions, we simply denote the name of the layer being tested (e.g. conv channel attention or lambda layer). We denote Ness the family of hybrid architectures described in Table 19 (Appendix F.2). Unless specified otherwise, all lambda layers use | |/|=4 with scope size of mm/=23x23 and implemented as in Figure 3. Additional experiments and details can be found in the Appendix.
5.1 LAMBDA LAYERS OUTPERFORM CONVOLUTIONS AND ATTENTION LAYERS
We first consider the standard ResNet-50 architecture with input image size 224x224. In Table 3, we compare the lambda layer against (a) the standard convolution (i.e. the baseline ResNet-50) (b) channel attention squeeze-and-exciiation and (c) multiple self-attention variants. The lambda layer strongly outperforms all basclines at a fraction of the parameter cost and notably obtains a +0.8% improvement over channel attention.
5.2
COMPUTATIONAL BENEFITS OF LAMBDA LAYERS OVER SELF-ATTENTION
In Table 4, we compare lambda layers against self-attention and present throughputs, memory com- plexities 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 lo- cal self-ateention while being almost 3x !o Additionally, positional embeddings can be shared
10 Latencies for local self-attention were provided privately by Ramachandran et al. (2019) based on an implementation that relies on query blocks and overlapping memory blocks (Parmar et al., 2018). Specialized attention kemels may greatly speed up local self-attention, making it promising avenue for future research.
7

---

Published as a conference paper at ICLR 2021
Layer
Params (M)
top-1
Conv (He et al., 2016)!
25.6
76.9
Conv +
channel attention (Hu et al., 2018c)
28.1
77.6 (+0.7)
Conv +
linear attention (Chen al., 2018)
33.0
77.0
Conv + linear attention (Shen et al., 2018)
77.3 (+1.2)
Conv +
relative self-attention (Bello et al.. 2019)
25.8
77.7 (+1.3)
Local relative self-attention (Ramachandran et al., 2019)
18.0
77.4 (40.5)
Local relative self-attention (Hu et al., 2019)
23.3
77.3 (+1.0)
Local relative self-attention (Zhao al., 2020)
20.5
78.2 (+1.3)
Lambda layer
15.0
78.4 (+1.5)
Lambda layer ([u|=4)
16.0
78.9 (42.0)
Table Comparison of the lambda layer and attention mechanisms on ImageNet classification with ResNet50 architecture. The lambda layer strongly outperforms attention alternatives at fraction of the parameter cost. All models are trained in mostly similar setups (see Appendix F.3) and we include the reported improvements compared to the convolution baseline in parentheses. See Appendix C.4 for description of the |u| hyperparameter. Our implementation.
across lambda layers to further reduce memory requirements, at a minimal degradation cost. Finally, the lambda convolution has linear memory complexity, which becomes practical for very large im- ages as seen in detection or segmentation. We also find that the lambda layer outperforms local self-attention when controlling for the scope size" (78.1 1..1% % vs 77.4% for |m m=77), suggesting that the benefits of the lambda layer go beyond improved speed and scalability.
Layer
Space Complexity
Memory (GB)
Throughput
top-1
Global self-attention
O(blhn²) )
120
OOM
0OM
Axial self-attention
Obllnn/n)
4.8
960 ex/s
77.5
Local self-attention (7x7)
(blhnm))
440 ex/s
77.4
Lambda layer
(Iknn))
1.9
1160ex/s
78.4
Lambda layer ([k|=8)
(
0.95
1640 ex/s
77.9
Lambda layer (shared embeddings)
(kkr?
0.63
1210 ex/s
78.0
Lambda convolution (7x7)
+ (Ikum)
1100
78.1
Table 4: The lambda layer reaches higher ImageNet accuracies while being faster and more memory-efficien than self-attention alternatives. Memory reported assuming full precision for batch of 128 inputs using default hyperparameters. The memory cost for storing the lambdas matches the memory cost activations the rest of the network and is therefore ignored. b: batch size, h; number of heads/queries, input length, m: context length, A: query/key depth, & number of layers.
5.3 HYBRIDS IMPROVE THE
TRADEOFF OF IMAGE CLASSIFICATION
Studying hybrid architectures. In spite of the memory savings compared to self-attention, cap- turing global contexts with the lambda layer still incurs quadratic time complexity (Table 2). which remains costly at high resolution. In Appendix 5.3, we study hybrid designs that use standard convolutions to capture local contexts and lambda layers to capture global contexts. We find that such convolution-lamble hybrids have increased representational power at negligible decrease in throughput compared to their purely convolutional counterparts.
LambdaResNee significantly improve the speed-accuracy tradeoff of ImageNet classification. We design family of hybrids based on our study of hybrid architectures and the scaling/training strategies from Bello al. (2021) (Section F.2). Figure presents the speed-accuracy Pareto curve of LambdaResNetss compared to EfficientNets (Tan & Le, 2019) on TPUv3 hardware. In order to isolate the benefits of lambda layers, we additionally compare against the same architectures when replacing lambda layers by (1) standard 3x3 convolutions (denoted ResNet-RS wo/ SE) and (2) 3x3
Note that the content-based lamhda still captures global interactions.
8

---

Published as a conference paper at ICLR 2021
convolutions with squeeze-and gqueeze-and-excitatie (denoted ResNet-RS w/ SE). All architectures trained for 350 epochs using the same regularization methods and evaluated at the same resolution they trained at. LambdaResNets outperform the baselines across all scales on the speed- -eed-acc trade-off.
Speed-Accuracy
87
speed-up
84
Top-3 Imagettet accuntt
83
85.0
(EA,
02-
(156,
1270. 228
84.0
81
83.5
80
82.5
Lambdialtestet
81
Efficienthic:
82.0
ResNet-RS wy SE
78
81.5
83
ResNet
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Time per training step (5) 1024 images.
Figure 4: Speed-accuracy comparison between Ness and EfficientNets with match- ing training and regularization setups. LambdaResNett (annotated with (depth, image size)) are 3.2- 4.4x faster than EfficientNets and 1.6-2.3x faster than ResNet-RS with squeeze-nn--exciaation thus significantly improving the speed-accuracy Pareto curve image classification. LesedaResNett 420 (image size 320), reaches a strong 84.9% top-1 accuracy, 0.9% over the corresponding architec- ture with standard 3x3 convolutions and 0.65% over the corresponding architecture with squeeze- and-excitation.
Scaling to larger datasets with pseudo-labels We train Nesss in semi-supervised learning setting using 130M pseudo-labeled images from the JFT dataset, as done for training the SoeiientNet-NoisyStuddee
checkpoints (Xie al., 2020). Table 5 compares the throughputs and ImageNet accuracies of representative set of models with similar accuracies when trained using the JFT dataset. LambdaResNNt-1552 trained and evaluated at image size 288, achieves strong 86. 7% top-1 ImageNet accuracy while being more crrramtter-fffcienn and 9.5x faster than the EfficientNet- NoisyStudent checkpoint with the same accuracy.
Architecture
Params (M)
Train (ex/s)
Infer (ex/s)
ImageNet top-1
LambdaResNet-152
51
1620
6100
86.7
EfficientNet-B7
66
170 (9.5x)
980 (6.2x)
86.7
VIT-L/16
307
180 (9.0x)
640 (9.5x)
87.1
Table 5: Comparison of models trained extra data. ViT-I/16 is pre-trained JFT and fine- tuned on ImageNet at resolution 384x384, while EfficientNet and LambdaResNet are co-trained on ImageNet and JFT pseudo-labels. Training and inference throughput is shown for 8 TPUv3 cores.
CONCLUSION
We propose new class of layers, termed lambda layers, which provide scalable framework for capturing structured interactions between inputs and their contexts. Lambda layers summarize avail- able contexts into fixed-size linear functions, termed lambdas, that are directly applied their as- sociated queries. The resulting neural networks, LambdaNetworks, are computationally efficient and capture long-range dependencies at a small memory cost, enabling their application to large structured inputs such as high-resolution images. Extensive experiments showcase their versatility and superiority over convolutional and attentional networks. We introduce NesbdaResNets, family of hybrid LambdaNetworks which reach excellent ImageNet accuracies and achieve up to 9.5x speed-ups over the popular EfficientNets and Vision Transformers, signifi- cantly improving the speed-accuracy tradeoff of image classification models.
computer vision tasks

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
The author would like to thank Barret Zoph William Fedus endless discussions, fruitful suggestions and careful revisions; Jonathon Shlens, Mike Mozer, Prajit Ramachandran, Ashish Vaswani, Quoe Lc, Neil Housby, Jakob Uszkoreit, Margaret Li, Krzysztof Choromanski for many insightful comments; Hedvig Rausing for the antarctic infographics: Zolan Brinnes for the OST; Andrew Brock. Sheng Li for assistance with profiling EfficientNets: Adam Kraft, Thang Luong and Hieu Pham assistance with the semi-supervised experiments and the Google Brain team for useful discussions the paper.
REFERENCES
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin lonescu. weights to attend to the recent past. 2016.
Using fast
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations, 2015.
Irwan Bello, Hieu Pham, Quoe V. Le, Mohammad Norouzi, and Samy Bengio. Neural combinatorial optimization with reinforcement learning. 2016. URL 09940.
Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoe V. Le. Attention augmented convolutional networks. CoRR, abs/1904.09925,, 2019. URL 1904. 09925.
org/bbs/
Irwan Bello, William Fedus, Xianzhi Du, Ekin D. Cubuk, Aravind Srinivas, Tsung- Lin, Jonathon Shlens, and Barret Zoph. Revisiting resnets: Improved training methodologies and scaling rules. 2021.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. 2020.
Denny Britz, Melody Y. Guan, and Minh- Thang Luong. Efficient attentioo using a fixed-size mem- ory representation. CoRR, 0ts/1777.00010 2017. URL http://aaax/ org/ abspg//bs//7// 1777. 00110.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. 2020.
Mark Chen, Alec Radford, Rewon Child, Jeff Wu, Heewoo Jun, Prafulla Dhariwal, David Luan, and Ilya Sutskever. Generative pretraining from pixels. 2020a. URL httph:ttt//// htRLL:tts:///ppaa //ppnai...p..../////// com/ bllog/mmage-gei
Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and Jingjing Liu. Uniter: Universal image-text representation learning. 2020b.
Yumpeng Chen, Yannis Kalantidis, Jianshu Li, Shuicheng Yan, and Jiashi Feng. A2-nets: Double attention networks. CoRR, Abs/1810.111799 2018. URL http: http///axvvv../ org/ abs/181C. 11579.
Rewon Child, Scott Gray, Alec Radford, and Sutskever Ilya. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1904./944//0// 10509.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin. Lukasz Kaiser. David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. 2020.
Jean-Baph Ctan-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self- wwwg/bbss......oo.oo.... abs 1911 03584.
attention and convolutional layers. 2019. URL /ttt:://aree org
10

---

Published as a conference paper at ICLR 2021
Ekin D. Cubuk, Barret Zoph. Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with a reduced search space. 2019.
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc Lc, and Ruslan Salakhutdinov. Transfomee-xll: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. Association for Compu- tational Linguisties, 2019. doi: 10. 10.18653/v1/P19- 1285. URL https: wwwww//wwwwwbbw..........o//....
acwweb. org/
anthology/P11 Ptt.llogy/P11128888
Alexandre de Brébisson and Pascal Vincent. A cheap linear attention mechanism with fast lookups and fixed-size representations. 2016.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition. IEEE, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlshy. An image is worth 16x 16 words: Transformers for image recognition scale. 2020.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. 2021.
David Ha, Andrew M. Dai, and Quoc Le. Hypernetworks. CoRR, 0tss//609.09006 2016. URL http:/ www.ii.wgggoa........oo.o. hatgg/abs//e 609.6099000 091006
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In 2017 IEEE Inter- national Conference on Computer Vision (ICCV), pp. 2980-2988, 2017.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. 2018.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi- mensional transformers. arXiv preprint arXiv: 12180 2019.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chich Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, and Adam Hartwig. Searching for mobilenetv3. 2019.
Andrew Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint (eXXiv:/704.044 2017.
Han Hu, Jiayuan Gu, Zheng Zhang, Jifeng Dai, and Yichen Wei. Relation networks for object detection. 2018a.
Han Hu, Zheng Zhang, Zhenda Xie, and Stephen Lin. Local relation networks for image recognition. arXiv preprint /rXiv:9904. 1491, 2019.
Jie Hu, Li Shen, Samuel Albanie, Gang Sun, and Andrea Vedaldi. Gather-excite: Exploiting feature context in convolutional neural networks. In Advances in Neural Information Processing Systems, bb.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-excitan networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018c.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Weinherger. Deep networks with stochas- tic depth. 2016.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Learning Representations, 2015.

---

Published as a conference paper at ICLR 2021
Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Ba- jwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le. Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of a tensor processing unit. SIGARCH Comput. Archit. News, 45(2):1-12, June 2017. ISSN 0163-5964. doi: 10.1145/ 1140659.3080244 URL http: /tw.oooooooooo//oooooooo acm org /10. 1141/ 31 4594..66.... 3080246. 6068.300024.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Flcuret. Transformers are mas: Fast autoregressive transformers with linear attention. 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint wrXiv:2001.044511 2020.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning, 2020.
Jungkyu Lee, Taeryun Won, Tae Kwan Lee, Hyemin Lee, Geonmo Gu, and Kiho Hong. Compound- ing the performance improvements of assembled techniques in a convolutional neural network, 2020.
Liunian Harold Li. Mark Yatskar, Da Yin, Cho-Jui Hsieh, and Kai-Wei Chang. Visualbert: A simple and performant baseline for vision and language. 2019.
Xingyu Liao, Lingxiao He, Zhouwang Yang, and Chi Zhang. Video-based person re-identification via 3d convolutional networks and non-local attention. 2019.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pictro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference Computer Vision, pp. 740-755. Springer, 2014.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. 2020.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In Inter- national Conference on Learning Representations, 2017.
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolin- guistic representations for iisionandd-lnguuge tasks. 2019.
John Nickolls and William J Dally. The gpu computing era. IEEE micro, 30(2):56-69, 2010.
Jongchan Park, Sanghyun Woo, Joon- Young Lee, and In So Kweon. Bam: bottleneck attention module. In British Machine Vision Conference, 2018.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, 2018.
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with a general conditioning layer. CoRR, 1b9.07799,7771, 2017.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. 2021. URL https:/ //spnaai http///lllll//////
12

---

Published as a conference paper at ICLR 2021
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention vision models. CoRR, abs/1906.05909, 2019. URL http:
/aaxiv.crg/abs/aa......a.aaa. htt.:rggga.aa/a//aa 05909.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang- Chigh Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4510-4520, 2018.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. arXiv preprint DeXiv:8803.02115 2018.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expeeee layer. CoRR, 701170011663333 2017. URL http: aariiv www.riiw.r......eeeeeee.eeeoeeee abs 1701 06538.
Zhuoran Shen, Mingyuan Zhang, Shuai Yi, Junjie Yan, and Haiyu Zhao. Efficient attention: Self- attention with linear complexities. CoRR, abs/1812.01243, 2018. URL
abs/1812, 01243.
Zhuoran Shen, Irwan Bello, Raviteja Vemulapalli, Xuhui Jia, and Ching-Hui Chen. Global self- attention networks for image recognition, 2020.
Aravind Srinivas, Tsung-Yr Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck transformers for visual recognition. 2021.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 19(6699999- 2014. URL http: www.mml.rgagappeeeeeeeeeem
srivastaval4a ha..sttaaala.hhm......
Chen Sun, Austin Myers, Carl Vondrick, Kevin Murphy, and Cordelia Schmid. Videobert: A joint model for video and language representation learning. 2019. URL http://arxiv.or abs/1904. 01766.
Mingxing Tan and Quoe V. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. CoRR, htbs/1905.1194 2019. URL http: arxiv. /ttg/as/////eeeooeeee/./eeeaaaee 1905. 14646.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2020.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, pp. 5998-6008, 2017.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly, Pointer networks. In Advances in Neural Information Processing Systems, 2015.
Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, and Liang-Chich Chen. Axial-deeplab: Stand-alone axial-attention for panoptic segmentation. 2020a.
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Seff-attention with linear complexity. 2020b.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7794- 7803, 2018.
Sanghyun Woo, Jongchan Park, Joon- Young Lee, and In So Kweon. Cbam: Convolutional block attention module. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
13

---

Published as a conference paper at ICLR 2021
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision. 2020.
Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V. Le. Self-training with noisy student improves imagenet classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. 2021.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. Proceedings of Machine Learning Research, pp. 2048 2057. PMLR, 2015.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. 2019.
Hengshuang Zhao, Jiaya Jia, and Vladlen Koltun. Exploring self-attention for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), June 2020.
Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning. In Interna- tional Conference on Learning Representations, 2017.
14

---

Published as a conference paper at ICLR 2021
CONTENTS
Introduction
2
Modeling Long-Range Interactions
2.1
Motivating queries, keys, position embeddings and values
2.2
Attention vs lambda layers.
3
Lambda Layers
3.1
Lambda layer: transforming contexts into linear functions.
3.2
A multi-query formulation to reduce complexity.
3.3
Making lambda layers translation equivariant.
6
3.4
Lambda convolution: local contexts the grid.
6
Related Work
5
Experiments
5.1
Lambda layers outperform convolutions and attention layers
5.2
Computational benefits of lambda layers over self-attention
5.3
Hybrids improve the speed-accuracy tradeoff of image classification
6
Conclusion
Discussion
16
A.1
General discussion
16
A.2 Extending lambda layers other modalities
16
Practical Modeling Recommendations
17
C Additional Variants
18
C.1 Complete code with lambda convolution
18
C.2 Generating lambdas from masked contexts
18
C.3 Multi-head vs multi-query lambda layers
18
C.4 Adding expressivity with an extra dimension
19
D Additional Related Work
20
D.1 Softmax attention
20
D.2
Sparse attention
................................
21
D.3 Linear attention: connections and differences
................................
21
D.4
Casting channel and spatial attention lambda layers. Self-Attention in the visual domain
................................
22
D.5
22
D.6
HyperNetworks, expert models and context- weights
23
Additional Experiments
24
E.1
Ablation study
24
E.2
Hybrid models study
25
E.3
Object detection and instance segmentation results
26
E,4
Parameter and FLOPs efficiency results
27
Experimental Details
29
F.1
Detailed LambdaResNets results
29
F22
Architectural details
29
F3
Training details
30
15

---

Published as a conference paper at ICLR 2021
A DISCUSSION
A.1 GENERAL DISCUSSION
How do lambda layers compare to the attention operation? Lambda layers scale favorably compared to self-attention. Vanilla Transformers using self-attention have e ((blhn²) memory foot- print, whereas LambdaNetworks have @(Ikn²) memory footprint (or O(k)²) ² when sharing posi- tional embeddings across layers). This enables the use of lambda layers at higher-resolution and on larger batch sizes. Additionally, the lambda convolution enjoys a simpler and faster implementation than its local self-attention counterpart. Finally, ImageNet experiments show that lambda layers outperforms self-attention, demonstrating that the benefits of lambda layers go beyond improved speed and scalability.
How are lambda layers different than linear attention mechanisms? Lambda layers general- ize and extend linear attention formulations to capture position-basee -astiion--bsdd interactions, which crucial for modeling highly structured inputs such as images (see Table 9 in Appendix E.1). As the aim not to approximate an attention kernel, lambda layers allow for more flexible non-linearities and normalizations which we also find beneficial (see Table 11 in Appendix E.I). Finally, we pro- pose multi-query lambda layers as means to reduce complexity compared to the multi-head (or single-head head) formulation typically used in linear attention works. Appendix D.3 presents 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 strong candidate for use in the visual domain. Our ablations demonstrate that lambda layers are most beneficial in the intermediate and low-resolution stages of vision architectures when optimizing for the speed-accuracy tradeoff. lambda layers which can be more parameter and flops efficient. We discuss practical modeling recommendations in Appendix B.
also possible to design architectures that rely exclusively on
A.2 EXTENDING LAMBDA LAYERS TO OTHER MODALITIES
While this work focuses on static image recognition, we note that lambda layers may be instantiated to model structured interactions on structures as diverse as graphs, time series, spatial lattices, etc. We anticipate that lambda layers will be helpful in more modalities, including multimodal tasks. We discuss masked contexts and auto-regressive tasks in the Appendix C.2.
Lambda layers can instantiated on other tasks simply by adapting the choice of structural/posi- tion embeddings to the task of interest and following the pseudo-code presented in Figure 3. The choice of embeddings dictates the memory costs the lambda layer. The assumption underlying the (knm) space complexity of the lambda layer (Section 3.2) is that all examples in the batch share the same structure, i.e. relative position embeddings have shape k x n x m. This assumption does not hold when the data structure is different across examples (e.g. graphs with variable edge relations between nodes), in which case embeddings have shape x le x n x m. In such cases, the lambda layer has (bkrmm) space complexity, similar self-attention.
16

---

Published as a conference paper at ICLR 2021
B PRACTICAL MODELING RECOMMENDATIONS
I want to make faster on TPUs/GPUS.. Hybrid models reach better speed-accuracy tradeoff. Global contexts be computationally wasteful, especially in the early high resolution layers where features lack semantic information, and can be replaced by lambda convolutions with smaller scopes (e.g. |m|=5x5 7x7) the standard 3x3 convolution. Additionally, using a hybrid can require less tuning when starting from working model/training setup.
I want to make minimize FLOPS (e.g. embedded applications).. Consider a hybrid with inverted bottlenecks, as done in Section E,4.2. To further reduce FLOPS, prefer lambda convolutions with smaller scopes (e.g. m =5x5 or 7x7).
I encounter memory issues... Memory footprint can be reduced by sharing position embeddings across layers (especially layers with the highest resolution). Using the lambda convolution is more memory efficient. Reducing the query depth |k or increasing the number of heads also decreases memory consumption.
I'm experiencing instability... We found it important to initialize the y parameter in the last batchnorm layer of the ResNet's bottleneck blocks to (this is the default in most codebases). Nor- malizing the keys (i.e. with the softmax) along the context's s length important. Early experiments which employed lambda layers sequentially in the same residual block were unstable, suggesting that using lambda layers in sequence should be avoided.
Which implementation of the lambda convolution should use? In our experiments using Ten- sorflow 1.x on TPUv3 hardware, we found both the n-d depthwise and (n+I)-d convolution imple- mentations have similar speed. We point out that this can vary across software/hardware stacks.
What if my task doesn't require position-based interactions? Computational costs in the lambda layer are dominated by position-based interactions. If your task doesn't 't require them, you try the content-only lambda layer any other linear attention mechanism. We recommend us- ing the multi-query formulation (as opposed to the usual multi-head) and scaling other dimensions of the model.
17

---

Published as a conference paper at ICLR 2021
C ADDITIONAL VARIANTS
C.1 COMPLETE CODE WITH LAMBDA CONVOLUTION
b: batch, input length, context length, scope
query/key depth, value depth, h: number of heads, d: output dimension. def wtw.ppio..aaaaaaaa....eooo.oaeee (embedinns,,,,, values, impl='einsun));
if impl 'einsum': embeddinge shape: [n, m, k]
position lambdas einsum (enumeddemddd, values, bmb->) else: embeddings shape: [r, k] impl 'conv':
bev
[r, 1, 1, k)) values reshape(valuee) [b, n, v, 1])
embeddings reshape
position_lambdas
convld(values, embeddinge)
elif impl == 'depthvise.con'::
# Reshape and tile embeddings to [r, v, x] shape
embeddings reshape (enbeddinge, [r, 1, k])
embeddings tile (embedddid,, [1, v, 13)
position_lambdan
aepthuise.conviddvalll
embeddings)
Transpose from shape [b, n, v, x] to shape [b, k, v] position lambdas (positonn lambdas, [0, 1, 3, 2])
return position lambdas
def (qyeeu keys, embeddings, values, impl='einsun)
""Multi-query lambda layer. .*"
lambda einaum (sof (keys), values, bmkkbbrv-bbkv')
position_lambdas compute compute_pasition positiion_ambdas (embeddgnn,,, values, impl-impl) content_output einsum (queries, content_lanbda, "bhnk, bkv (ww->bnhv'') position.output einsum (queries, seittionllabddas "bhnk, babv->b output (content return output
+ position.outpur, [b, n, d])
Figure 5: Pseudo-code for the multi-query lambda layer and the 1d lambda convolution. A n-d lambda convolution can equivalently be implemented via regular (n+1)-d convolution or a n- d depthwise convolution with channel multiplier. The embeddings can be made satisfy various conditions (e.g. translation equivariance and masking) when computing positional lambdas with the einsum implementation.
C.2 GENERATING LAMBDAS FROM MASKED CONTEXTS
In some applications, such denoising tasks or auto-regressive training, it is necessary to restrict interactions a sub-context C, when generating X, for query position n. For example, parallel auto-regressive training requires masking the future to ensure that the output yn only depends on past context positions m n. Self-attention achieves this by zeroing out the irrelevant attention weights aam 0 Vm Cn, thus guaranteeing that yn Σm arum Um only depends C,,-
Similarly, one can block interactions between queries and masked context positions when generating lambdas by applying mask before summing the contributions of context positions. As long as the mask is shared across all elements in the batch, computing masked lambdas does not require materializing per-example attention maps and the complexities the same for global context case. See Figure 6 for an implementation.
C.3 MULTI-HEAD MULTI-QUERY LAMBDA LAYERS
In this section, we motivate using a multi-query formulation opposed to the usual multi-head for- mulation used self-attention. Figure presents the implementation of a multi-head lambda layer. Table 6 compares complexities for multi-head and multi-query lambda layers. Using multi-query formulation reduces computations by factor |h| (the number queries per lambda) compared to the multi-head formulation. We also found in early experimentation that multi-query lambdas yield accuracy trade-off. Additionally, the multi-head lambda layer does not enjoy a simple local implementation the lambda convolution.
a better speed- speed-accuracy
18

---

Published as a conference paper at ICLR 2021
masked lambda layer (queries, normalized keys, embeddings, values, mask): "Maakedd multi-query lambda layer.
Args:
queries; a tensor with shape [b, h, n, k] normalized keys: a tensor with shape [b, m, k]
embeddinge: a tensor with shape [k, n, m]
values: tensor with shape [b, vl.
mask: tensor of 0 and is with shape [n, m]
We show the general case but cumulative sum may be faster for masking the future. Note that each query nou also has its oun content_lambee since every query interacts with different context.
Keys should be normalized by only considering the elements their contexts. content.mu einsum (normalizee..eey.. values, 'bmk, b b content_labbaae einaum (content,mm) mu, mask, 'bmkv, nm->bukv')
embeddings einsum (embeddings, mask, nm >knm') apply mask to enbeddings position_lambdas einsum (eneeuddeemmdding values, 'kmm, bmv bmy->>bkv')
content output einsum (queries, content lambda, boky bnkkv>>bkk>> >bbbh')
position output einsum (queries, position lambdas, "bhnk, bnkv->bahv output (content return output
+ iosition.outpuu [b, n, d])
Figure 6: Pseudo-code for masked multi-query lambda layer.
def multihead lambda (queries, keys, embeddings, values, "oppummminsmmm)
**"Multi heaulti--edddd lambda layer.
lambea einsum (softmax (keys), values, bbhnk, bhmv >bhkv') positionlammdaa lambdas einsum(embedding values, 'hnnk, blav->>bnhkv")
content output einsum (queries, content lambda, "bhnk, bhkv->bnhv") position_output einsun(queries, lambaas, "bhnk, bakv->bahy output reeshppeconttettoouppe return output
position.output, [b, n, d])
Figure 7: Pseudo-code for the multi-head lambda layer. This is only shown an example recommend the multi-query lambda laayer instead.
Operation
Time complexity
Space complexity
Multi-hcar head lambda layer Multi-query lambda layer
(bbmmkd) @bmmkkd/h)
@kmm + bnkd) (hknm + bnkdd h)
Table 6: Complexity comparison between multi-head and multi-query lambda layer. Using a multi-query formulation reduces complexity by factor [/] (the number of queries per lambda) compared to the standard multi-head formulation.
C.4 ADDING EXPRESSIVITY WITH AN EXTRA DIMENSION
We briefly experiment with variant that enables increasing the cost of computing the lambdas while keeping the cost of applying them constant. This is achieved by introducing an additional dimension, termed the intra-depth with corresponding hyperparameter |ul, in keys, position embeddings and values. Each key (or positional embedding) is now |k| x |u] matrix instead of |A]-ddmmnnsiona vector. Similarly, each value is now |u] x Iu] matrix instead of [w] del-dimmnsional vector. The lambdas are obtained via summing over context positions and the intra-depth position |u| and have |k| [o] shape similar to the default case. See Figure for an implementation and Table for the complexities. Experiments (see Appendix E.1) demonstrate that this variant results in accuracy improvements but we find that using |w/=1 (i.e. the default case) is optimal when controlling for speed on modern machine learning accelerators.
19

---

Published as a conference paper at ICLR 2021
compute position lambdas (embeddings, values, impl='einsu")), "o"Compue position 1ambdas with intra-depth u. "".
if impl 'conv':
values: [b, n, V, u] shape
embeddings: [r, 1, u, k] shape
position lambdas conv2d (values, embeddinge)
Transpose from shape [b, n, v, k] to shape [b, k, v] position_lamddaa trampose@peatiionllambbe elif impl == 'einsum':
[o, 1, 3, 2])
embeddinge: [k, m, u] shape
position_lambdas einsum (embeddinge, values, *knmu, bavu->bbkv') return position_lambdas
def Lamtca.layyerquurrie,, keys, embeddings, values, impl=/einsum)):
"Muulti-query lambda layer with intra-depth u. ""n
content.lambda einsum tmx (keys), values, 'bmku, bavu >bkv') position_aamdaas
popppuueppsiition.. (enbbeonn (embeddings, values, lambda_conv)
content output einsum (queries, content lambda, 'bhnk, Biikkkbbbbbbbbrre >>bhh')) position output einsum (queries, position lambdas, "bhnk, bakv->
output reshape return output
+ opsition.output. [b, n, d])
Figure 8: Pseudo-code for the multi-query lambda layer with intra-depth lu]. Lambdas obtained by reducing over the context positions and the intra-depth dimension. This variant allocates more computation for generating the lambdas while keeping the cost of applying them constant. The equivalent n-d lambda convolution can be implemented with a regular (n+1)+s ) convolution.
Operation
Time complexity
Space complexity
Lambda layer (ju] >1)
O(bnmkud/h)
(knmu + bnku)
Table 7: Complexity for multi-query lambda layer with intra-depth lu].
D ADDITIONAL 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 cast as less flexible specific instances of lambda layers. We conclude with a brief review of self-attention in the visual domain and discuss connections with expert models.
D.1 SOFTMAX ATTENTION
Softmax attention Softmax-attention produces a distribution over the context for each query q. an softmax K qm. E IR where the keys K obtained from the context C. The attention distribution a, is then used to form a linear combination of values V obtained from the context as yn V a. Σm Onm Um. E [R|0| As take weighted sum of the v2 we transform the query into the output and discard its attention distribution a. This operation captures content-based interactions, but not position- based interactions.
Relative attention In order to model position-base based interactions, relative attention (Shaw et al., 2018) introduces a learned matrix of |m positional embeddings E,, E R|m|x[k| and computes the attention distribution as a, = softmax [K + E, En)q,)) E R!ml The attention distribution now also depends the query position relative to positions of context elements m. Relative attention therefore captures both content-based and position-based interactions.
"ommeimess the attention operation is instead used to point specific context elements (Vinyals et al., 2015; Bello al., 2016), which not supported by lambda layers.
20

---

is the case in the regular Published as a conference paper at ICLR 2021
D.2 SPARSE ATTENTION
A significant challenge in applying (relative) attention to large inputs comes from the quadratic O[Imml] memory footprint required store attention maps. Many recent works therefore propose to impose specific patterns to the attention maps as a means to reduce the context size |m| and consequently 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; Shen et al., 2020)
static sparse attention patterns (Child et al.; Beltagy et al., 2020)
dynamic spurse attention patterns (Kitaev et al., 2020)
See Tay et al. (2020) for a review. Their implementations can rather complex, sometimes require low-level kernel implementations get computational benefits may rely specific assumptions on the shape of the inputs (c.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.
D.3 LINEAR ATTENTION: CONNECTIONS AND DIFFERENCES
Another approach to reduce computational requirements attention mechanisms consists in ap- proximating the attention operation in linear space and time complexity, which is referred to as linear (or efficient) attention. Linear attention mechanisms date back to de Brébisson & Vincent (2016); Britz et al. (2017) and were later introduced in the visual domain by Chen et al. (2018); Shen et al. (2018). They are recently enjoying a resurgence of popularity with many works mod- ifying the popular Transformer architecture for sequential processing applications (Katharopoulos et al., 2020; Wang et al., 2020b; Choromanski et al., 2020; Xiong et al., 2021).
Linear attention via kernel factorization Linear attention is typically obtained by reinterpreting attention as similarity kernel and leveraging low-rank kernel factorization
Attention Q, V) softmax QK")V ((Q)
for some feature function di. Computing di (((K) E R kk el first bypasses the need to materialize the attention maps ((Q)((KT) and the operation therefore has linear complexity with respect to the input length [n).
Multiple choices for the feature function e 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. In the visual domain, both Chen al. (2018) and Shen er al. (2018) use (dxr) = softmax (x). This choice is made to guarantee that the rows of the (not-materialized) attention maps oi Q rob K""" sum attention operation.
We 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 3 can rewritten for single query qn as
yn = (6(K)) V) ((q.)
(4)
which resembles the output of the content lambda y. (At)T qm = (RTV)T qn from Equation 1. Lambda layers extend linear attention mechanisms to also consider baased interactions
yn X qm (X*+ x ((K E.,) VT. In
(5)
In the above equation, computing the position (or content) lambda has O(bmkv)) time complexity. As the position lambdas are not shared across query positions n, this cost is repeated for all |n] queries, leading to total time complexity O (bmmkkv). Unlike linear attention mechanisms, lambda layers have quadratic time complexity with respect to the input length (in the global context case) because they consider baaed interactions.
21

---

Published as a conference paper at ICLR 2021
2) Lambda layers do not necessarily attempt to approximate attention kernel. While ap- proximations of the attention kernel are theoretically motivated, we argue that they may be unnec- essarily restrictive. For example, the kernel approximation in Equation 3 requires the same feature function & on both Q and K and precludes the use of more flexible non-linearities and normaliza- tion schemes. In contrast, lambda layers do not attempt approximate simplifies their design and allows for more flexible non-linearin and normalization schemes, which we find useful in our ablations (See Table 11 in Appendix E. Considering the position embed- dings independently of the keys notably enables simple and efficient local implementation with the lambda convolution. Approximating the relative attention kernel would require normalizing the position embeddings with the keys (i.e., of (KK + E,,) instead of of ()) + E,,, which cannot be implemented in the local context case with a convolution.
attention kernel. This
3) The lambda abstraction reveals the computational benefits of the multi-query formulation. Finally, this work proposes to abstract the k and E V matrices as linear functions (the content and position lambdas) that are directly applied to the queries. The lambda abstraction reveals the benefits of as a means to reduce computational costs.
formulation (as opposed to the traditional multi-head attention formulation)
D.4 CASTING CHANNEL AND SPATIAL ATTENTION AS LAMBDA LAYERS.
We show that the lambda abstraction generalizes channel and spatial attention mechanisms, both of which can be viewed as specific 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 efficient.
Channel attention Channel attention mechanisms, such as gqueeze-and-Excitatiin (SE) (Hu al., 2018c;b) and FiLM layers (Perez et al., 2017), recalibrate features via cross-channel inter- actions by aggregating signals from the entire feature map. In particular, the SE operation can be written Ynk wkenk where Wi is the excitation weight for channel ie in the query 9n- This can be viewed as using diugonal lambda which is shared across query positions X = diag(w) ... wpp/ Channel attention mechanisms have proven useful to complement convolutions but cannot be used as stand-alone layer they discard spatial information.
Spatial attention Conversely, spatial attention mechanisms, reweigh each position based on sig- nals aggregated from all channels (Xu et al., 2015; Park et al., 2018: Woo al., 2018). These mechanisms can be written as Vnk = Wnynkk where wn, is the attention weight for position n in the input query Q. This can be viewed as using (position-ddeeedent) scalar lambdas X,, = w,,, I. where II is the identity matrix. Spatial attention has also proven helpful to complement convolutions but cannot be used as a stand-alone layer as it discards channel information.
D.5 SELF-ATTENTION IN THE VISUAL DOMAIN
Self-attention has been used in a myriad of tasks the visual domain. These include image classi- fication (Bello et al., 2019; Ramachandran et al., 2019; Cordonnier et al., 2019; Zhao et al., 2020; Wu et al., 2020; Dosovitskiy et al., 2020); object detection and object-centric tasks (Wang et al., 2018; Hu et al., 2018a; Carion et al., 2020; Locatello et al., 2020); video tasks (Sun et al., 2019; generative modeling (Parmar et al., 2018; Zhang et al., 2019; Brock al., 2019; Chen et al., 2020a) and multi-modal text-vision tasks (Chen al., 2020b; Lu et al., 2019; 90 Li et al., 2019; Radford et al., 2021)
Liao et al., 2019); autoregressive/adversss
The first of self-attention in vision dates back to the non-local block (Wang al., 2018), which added a single-head global self-attention residual in the low resolution stages of a ConvNet for long- range dependency modeling. The non-local block has proven useful complement convolutions but cannot be used as stand-alone layer as it does not model position-baa based interactions.
Global relative attention replaces convolutions at low resolution. Bello et al. (2019) introduced a 2d relative attention mechanism that proved competitive as a replacement to convolutions but gives even stronger results when used to concatenate convolutional features with self-attention features. The spatial convolutions in the bottleneck block of the ResNet architecture were replaced with
22

---

Published as a conference paper at ICLR 2021
global multi-head self-attention mechanism with 2d relative position embeddings. Due to the large memory constraints of global attention, this operation was restricted to low resolution feature maps and the proposed architecture was con-trransforme hybrid.
A similar hybrid design has recently been revisited by Srinivas et al. (2021) using modern training and scaling techniques. Srinivas al. (2021), rather than concatenating convolutional feature maps, propose to use a stride of in the last stage of the ResNet architecture for improved performance.
Local/axial relative attention replaces convolutions high resolution. The large memory foot- print of global attention was quickly solved by multiple works which proposed to limit the size of the attention contexts such as local attention (Ramachandran et al., 2019; Hu al., 2019) and avial attention (Ho et al., 2019; Wang et al., 2020a; Shen et al., 2020) (See Section D.2). Such approaches enable using attention at higher resolution and facilitate fully-attentional models but can be slow due to the use of specialized attention patterns.
Scaling trumps inductive bias Concurrently to this work, ViT (Dosovitskiy et al., 2020) propose to simply apply attention on pixel patches (as opposed to individual pixels) remedy large memory requirements. While patch-based attention does not maintain accurate positional informa- tion translation equivariance, the loss of inductive bias is recovered by pre-training on large-scale datasets (e.g. 300M images). Most remarkably, VIT achieves close to state-of-the accuracy when fine-tuned on the ImageNet dataset, while requiring less training compute that convolutional alter- natives (Kolesnikov et al., 2020: Xie et al., 2020). This result has reinvigorated interest in using self-attention in the visual domain with multiple follow-up works already building upon this ap- proach (Touvron al., 2021)² In spite of the impressive image classification results, concerns remain as to whether the patch-based approach can scale to larger images and transfer to tasks that require precise localization such as detection.
We stress that reducing memory by working with pixel patches is orthogonal to the specific operation used and that lambda layers (or linear attention) can also be used pixel patches.
D.6 HYPERNETWORKS. EXPERT MODELS AND CONTEXI-DEPENDENT WEIGHTS
LambdaNetworks generate their own computations, i.e. lambdas such that yn X Un. As such, they can alternatively be viewed as an extension of HyperNetworks (Ha al., 2016) that dynamically generate their computations based on structured contextual information. The concept of generating context-dependent weights is also related to fast weights (Ba et al., 2016).
Lastly, LambdaNetworks share some connections with sparsely-activatee expert models (Shazeer et al., 2017: Fedus et al., 2021). Whereas sparsely-activated expert models select the computation (i.e. the lambda) from bank of weights based on the input query, LambdaNetworks generate their computations based on contextual information.
Most follow-up works advertise improvements over ViT smaller datasets which is not the intended purpose of VIT.
23

---

Published as a conference paper at ICLR 2021
E ADDITIONAL EXPERIMENTS
E. ABLATION STUDY
We perform several ablations and validate the importance positional interactions, long-range teractions and flexible normalization schemes. Unless specified otherwise, all experimental results this section report ImageNet accuracies obtained by training LambdaNetwork architecture that replaces the spatial convolutions in the ResNet-50 with lambda layers.
Varying query depth, number of heads and intra-depth. Table 8 presents the impact of the query depth |kl, number of heads |/] and intra depth [u] on performance (See Appendix C.4 for a presentation of the intra-depth [u(). Our experiments indicate that the lambda layer outperforms convolutional and attentional baselines for a wide range of hyperparameters, demonstrating the ro- bustness of the method.
[
[u]
Params (M)
top-1
ResNet baseline
25.6
76.9
8
2
14.8
77.2
8
16
15.6
77.9
2
4
14.7
77.4
4
14.7
77.6
8
4
14.8
77.9
16
4
15.0
78.4
32
4
15.4
78.4
2
8
14.7
77.8
4
8
14.7
77.7
8
8
14.7
77.9
16
8
15.1
78.1
32
8
15.7
78.5
8
8
15.3
78.4
8
8
8
16.0
78.6
16
4
4
16.0
78.9
Table 8: Ablations on the ImageNet classification task when using the lambda layer in ResNet50 architecture. All configurations outpeform the convolutional baseline at lower pa- rameter cost. As expected, we get additional improvements by increasing the query depth [k| or intra-depth |u]. The number of heads is best set to intermediate values such as |h|=4. A large num- ber of heads h excessively decreases the value depth |o = d///t,, while a small number of heads translates to too few queries, both of which hurt performance.
Content vs position interactions Table 9 presents the relative importance content-baa based and position-based interactions the ImageNet classification task. We find that position-based inter- actions are crucial to reach high accuracies, while content-based interactions only bring marginal improvements position-basee interactions
Content
Position
Params (M)
FLOPS (B)
top-1
14.9
5.0
68.8
14.9
11.9
78.
14.9
12.0
78.4
Table 9: Contributions of content and positional interactions. As expected, positional interac- tions are crucial to perform well on the image classification task.
"This observation is challenged by concurrent work (Dosovitskiy et al., 2020) which demonstrates that content-based interactions can he sufficient for image classification when pre-training large scale datasets (e.g. 300M images).
24

---

Published as a conference paper at ICLR 2021
Importance of scope size The small memory footprint of LambdaNetworks enables considering global contexts, even at relatively high resolution. Table 10 presents flops counts and top-1 ImageNet accuracies when varying scope sizes in a LambdaNetwork architecture. We find benefits from using larger scopes, with a plateau around |m|=15x15, which validates the importance interactions compared to the usual 3x3 spatial convolutions used in the ResNet architecture. In main experiments, we choose mm/==3x23 as the default to account for experiments that use larger image sizes.
longer range
Scope size |m]]
3x3
7x7
15x15
23x23
31x31
global
FLOPS (B)
5.7
6.1
7.8
10.0
12.4
19.4
Top-1 Accuracy
77.6
78.2
78.5
78.3
78.5
78.4
Table 10: Impact of varying the scope size for positional lambdas on the ImageNet classification task. We replace the 3x3 spatial convolutions in the last stages of a ResNet-50 with lambda layers (input image size is 224x224). Flops significantly increase with the scope size, however we stress that larger scopes do not translate to slower latencies when using the einsum implementation (see Figure 3).
Normalization Table 11 ablates normalization operations in the design of the lambda layer. We find that normalizing the keys is crucial for performance and that other normalization functions besides the softmax can be considered. Applying batch normalization to the queries and values is also helpful.
Normalization
top
Softmax on keys (default)
78.4
Softmax on keys & Softmax on queries
78.1
L2 normalization on keys
78.0
No normalization keys
70.0
No batch normalization queries and values
76.2
Table 11:
Impact of normalization schemes in the lambda layer. Normalization of the keys along the context spatial dimension m, normalization of the queries along the query depth k.
F.2 HYBRID MODELS STUDY
In this section, we study hybrid designs that use standard convolutions to capture local contexts and lambda layers to capture global contexts." 15
Where are lambda layers most useful? Table 12 presents the throughputs and accuracies of hybrid I LambdaNetwork architectures as a function of the location of convolutions and lambda layers in ResNet-50 architecture. We observe that lambda layers are most helpful in the last two stages (commonly referred to c4 and c5) when considering their speed-accuracy tradcoff. We refer to architectures that replaces 3x3 convolutions in the last stages of the ResNet with lambda layers as LembdaResNNt-CA4
Further pushing the speed-accuracy Pareto frontier. Table 13, we further study how through- put and accuracy impacted by the number of lambda layers in the c4 stage. Our results reveal that most benefits from lambda layers can be obtained by (a) replacing a few 3x3 convolutions with lambda layers in the c4 stage and (b) replacing all 3x3 convolutions in c5. The resulting hybrid LambdaResNets architectures have increased representational power at a virtually negligible de- crease in throughput compared their vanilla ResNet counterparts. Table presents the detailed block configurations and placement of lambda layers for our family of LambdaResNets.
We could alternatively the lambda convolution to capture local contexts.
25

---

pronounced for Published as a conference paper at ICLR 2021
Architecture
Params (M)
Throughput
top-1
C C-CC
25.6
7240 ex/s
76.9
L-C+C+0
25.5
1880 ex/s
77.3
C
25.0
1280 cx/s
77.2
21.7
1160 ex/s
77.8
15.0
1160 ex/s
78.4
C-L+ 15.1
2200 ex/s
78.3
C-C-L-L 15,4
4980 ex/s
78.3
C+C+C+I. 18.8
7160 ex/s
77.3
Table 12: Hybrid models achieve a better speed-accuracy trade-off. Inference throughput and top-1 accuracy function of lambda (L) convolution (C) layers' placement in ResNet50 architecture on 224x224 imputs. Lambda layers in the c5 stage incur almost no speed decrease compared standard 3x3 convolutions. Lambda layers in the of stage are relatively slower than standard 3x3 convolutions but yield significant accuracy gains.
Config
Image size
Params (M)
Throughput
top-1
ResNet-101 wo/ SE
224
44,6
4600 ex/s
81.3
ResNet-101 w/ SE
224
63.6
4000 ex/s
81.8
LanbdaResNet- 101
224
36,9
4040 ex/s
82.3
LambdaResNet-100
224
26.0
2560 ex/s
82.6
ResNet-152 wa/ SE
60.2
2780 ex/s
82.5
ResNet-152 w/ SE
256
86.6
2400 ex/s
83.0
LambdaResNet-152
256
51.4
2400 ex/s
83,4
LambdaResNNe-152-C44
256
35.1
1480 ex/s
83.4
Table 13: Impact number of lambda layers in the c4 stage of LambdaResNets. Most benefits from lambda layers can be obtained by having few lambda layers in the c4 stage. Such hybrid designs maximize the speed-accuracy tradeoff. LambdaResNNt-CC architectures exclusively employ lambda layers in c4 and c5. LambdaResNet block configurations can be found in Table 19. Models are trained for 350 epochs on the ImageNet classification task.
Comparing hybrid lambda attention models. The memory savings of lambda layers com- pared to attention are less significant in the aforementioned hybrid design, since the operations occur at lower resolution. Therefore, it is natural ask whether lambda layers still have benefits over self-attention when considering hybrid designs. We consider our largest hybrid an example (see Table 19). N2mt is trained on 320x320 inputs, employs lambda layers in and can fit 32 examples per TPU-v3 core. This adds up cost of 38.4MBB for lambda layers (4.8MB if sharing positional embeddings), whereas using attention layers instead would incur 0.625GB. The increase might not be significant in practice and it will be interesting to carefully benchmark the hy- brid attention variants 16 We point that experiments from Table suggest that the benefits of lambda layers go beyond improved scalability and stress that the memory savings tasks that require larger inputs such as object detection.
E.3 OBJECT DETECTION AND INSTANCE SEGMENTATION RESULTS
In Table 14, we evaluate LambdaResNets as backbone in Mask-RCNN (He et al., 2017) on the COCO object detection and instance segmentation tasks. Using lambda layers yields consistent gains across all object sizes, especially the small objects which are the hardest to locate. This indi- cates that lambda layers are also competitive for more complex visual tasks that require localization information.
We will benchmark such architectures in future version of this draft.
26

---

Published as a conference paper at ICLR 2021
Backbone
AP coco
APdc s/m/l
APmesk
6/m/t
ResNet-101
48.2
29.9/ 50.9 64.9
42.6
24.2 / 45.6 60.0
Bes SE
48.5 (+0.3)
29.9 (+0.0) 51.5 65.3
42.8 (+0.2)
24.0 (-0.2) 46.0/60.2
LambdaResNet-100
49.4 (+1.2)
31.7 (+1.8) 52.2/
43.5 (+0.9)
25.9 (+1.7) 46.5 60.8
ResNet-152
48.9
29.9 / 51.8 66.0 66.0
43.2
24.2 46.1 /6.2
ResNet-152 SE
49.4 (+0.5)
30.0
66.7 66.7
43.5 (+0.3)
24.6 (+0.4) 46.8 61.8
LambdaResNeu-152
50.0((+11)
31.8 (+1.9) 53.4
43,9 (+0.7)
25.5 (+1.3) 47.3 62.0
Table 14: COCO object detection and instance segmentation with Mask-RCNN architecture on 1024x1024 inputs. We compare LambdaResNetts against ResNets with without squeeze-and- excitation (SE) and report Mean Average Precision (AP) for small, medium, large objects (AP, A//m///t) Using lambda layers yields consistent gains across all object sizes, especially small objects.
E.4 PARAMETER AND FLOPS EFFICIENCY RESULTS
E.4.1 COMPUTATIONAL EFFICIENCY COMPARISONS LARGE EFFICIENTNETS
In Table and Table 16, we showcase the parameter and coos-efficieecy of LambdaNetworks. We find that LambdaResNee-C4 which replaces the 3x3 convolutions in the last 2 stages of the ResNet architecture, where they incur the highest parameter costs, improves upon parameter and flops effi ciency of large EfficientNets. These results are significant because EfficientNets were specifically designed by neural architecture search (Zoph & Le, 2017) to minimize computational costs using highly computationally efficient depthwise convolutions (Tan & Le, 2019).
Architecture
Image size
Params (M)
top-1
EfficientNet-B6
528x528
43
84.0
LambdaResNee-155-CC
320x320
35
84.0
LambLaResNet-200-c4
320x320
84.3
Table 15: Parameter-efficiene comparison between Resbda and EfficientNet- LambdaResNet- is more parameter-effiiient spite of using smaller image size. Increasing the image size would likely result in improved accuracy while keeping the number of parameters fixed. Models are trained for 350 epochs.
Architecture
Image size
Flops (G)
top-1
EfficientNet-B6
528x528
38
84.0
LambdaResNNet-770-cA ([m|=7x7)
256x256
34
84.0
Table 16: Flops-efficiency comparison between LamtdaRessNet-4 and EfficientNet--6. We smaller local scopes ([m|=7x7) reduce FLOPSS in the lambda layers. Models trained for 350 epochs.
E.4.2 LAMBDA LAYERS RESOURCE CONSTRAINED SCENARIO
Lastly, we briefly study lambda layers in resource-constrainet scenario using the MobileNetv2 architecture (Sandler et al., 2018). MobileNets (Howard et al., 2017: Sandler et al., 2018; Howard al., 2019) employ lightweight inverted bottleneck blocks which consist of the following sequence: 1) a pointwise convolution for expanding the number of channels, 2) depthwise convolution for spatial mixing and 3) final pointwise convolution for channel mixing. The use of depthwise convolution (as opposed to regular convolution) reduces parameters and flops, making inverted bottlenecks particularly well-suited for embedded applications.
Lightweight lambda block.
We construct lightweight lambda block follows. We replace the depthwise convolution in the inverted bottleneck with a lambda convolution with small scope size |m|=5x5, query depth [k|=32, number of heads |/]=4. We also change the first pointwise
27

---

Published as a conference paper at ICLR 2021
convolution to output the same number of channels (instead of increasing the number of channels) to further reduce computations.
Adding lambda layers in MobileNetv2. We wish to assess whether lambda layers can improve the flops-aceuracy (or parameler-accuracy) tradeoff of mobilenet architectures, We experiment with simple strategy of replacing a few inverted bottlenecks with our proposed lightweight lambda block, so that the resulting architectures have similar computational demands as their baselines. simple procedure of replacing the 10-th and 16-th inverted bottleneck blocks with lightweight lambda blocks in the MobileNet-v2 architecture reduces parameters and flops by ~ while im- proving ImageNet accuracy by 0.6%. This suggest that lambda layers may be well suited for use in resource constrained scenarios such embedded vision applications (Howard al., 2017; Sandler et al., 2018; Howard et al., 2019).
Architecture
Params (M)
FLOPS (M)
top-1
MobileNet-v2
3.50
603
72.7
MobileNet-v2 with lightweight lambda blocks
3.21
563
73.3
Table 17: Lambda layers improve ImageNet accuracy in resource-constrainee scenario. Replacing the 10-th and 16-th inverted bottleneck blocks with lightweight lambda blocks in the MobileNet-v2 architecture reduces parameters and flops by ~00% while improving ImageNet ac- curacy by 0.6%.
28

---

Published as a conference paper at ICLR 2021
F EXPERIMENTAL DETAILS
F.I
DETAILED LAMBDARESNETS RESULTS
Depth
Image size
Latency (s)
Supervised top-1
Pseudo- top-1
50
128
0.058
77.4
82.1
50
160
0.089
79.2
83.4
101
160
0.14
80.8
84.7
0,20
81.9
85.4
152
192
0.28
82,5
86.
152
224
0.38
83.2
86.5
152
256
0,49
83.8
152
288
0,63
86.7
256
0,91
84.2
350
256
1.16
84.4
350
288
1,48
84.5
350
320
1,91
84.7
420
320
2.25
84.9
Table 18: Detailed Less results. Latency refers to the time per training step for a batch size of 1024 on 8 TPU-v3 cores using bf l ooatl activations.
F.2 ARCHITECTURAL DETAILS
Lambda layer implementation details Unless specified otherwise, all lambda layers use query depth |k|=16, |/x /4|= heads and intra-depth |u/=1. The position lambdas generated with local contexts of size |m|22333 and the content lambdas with the global context using the einsum imple- mentation described in Figure 3. Local positional lambdas can be implemented interchangeably with the lambda convolution or by using the global einsum implementation and masking the position embeddings outside of the local contexts (Figure 5). The latter can be faster but has higher FLOPS and memory footprint due to the @(kum) term (see Table 2). In our experiments, we use the convo- lution implementation only for input length |n| > 852 or intra-depth |u] > When the intra-depth is increased to uu] >1, 1, we switch to the convolution implementation and reduce the scope size to |m|=7x7 to reduce flops.
(N(oiiioonnn Positional embeddings initialized at random using the unit normal distribution (0, 1). We use
fan-in initialization for the linear projections in the lambda layer. The projections to compute K and V are initialized at random with the NO0 -1/2) distribution. The projection to compute Q initialized at random with the (O, ked]-1/2) 1/2) distribution (this is similar to the scaled dot- product attention mechanism, except that the scaling is absorbed in the projection). We apply batch normalization on Q and V and the keys are normalized via softmax operation.
ResNets. We use the R 1 implementation and initialize the 7 parameter in the last batch normalization (loffe & Szegedy. 2015) layer of the bottleneck blocks to 0. Squeeze-and-Exeianl the ResNet-
layers employ a squeeze ratio of 4. Similarly ResNet-RS (Bello al., 2021), D (He et al., 2018) and additionally replace the max pooling layer in the stem by a strided 3x3 convolution. Our block allocation and scaling strategy (i.e. selected resolution model depth) also follow closely the scaling recommendations from ResNet-RS (Bello et al., 2021).
function of
Nests. We construct our LambdaResNets by replacing the spatial 3x3 convolutions in the bottleneck blocks of the ResNet-RS architectures by our proposed lambda layer, with the exception of the stem which is left unchanged. We apply 3x3 average-pooling with stride after the lambda layers to downsample in place of the strided convolution. Lambda layers are uniformly spaced in the c4 stage and all bottlenecks in c5 use lambda layers. Table 19 presents the exact block configuration and the location of the lambda layers for hybrid Nests.ResNets. We do not squeeze-and excitttio in the bottleneck blocks that employ a lambda layer instead the standard 3x3 convolution.
29

---

Published as a conference paper at ICLR 2021
Model
Block Configuration
Lambda layers in e4
LambdaResNet-50
[3-4-6-3]
3
LambdaResNet- 101
[3-4-23-3]
12,
LambdaResNet-155
[3-8-36-3]
5, 10, 15, 20,
25, 30
LambdaResNet-200
[3-24-36-3]
5. 10,
15,
20,
25
Nemt22ResNett2700
[4-29-53-4]
8, 16, 24, 32. 40, 48
Namt3aResNet-330
[4-36-72-4]
10, 20,
30,
40,
50, 60
I LamtdaResNet-420
[4-44-87-4]
10, 20, 30, 40, 50, 60, 70, 80
Table 19: Block configurations and lambda layers placement of Nesb in the Pareto curves. Ness use the block allocations from He et al. (2016); Bello et al. (2021).
F.3 TRAINING DETAILS
ImageNet training setups. We consider two training setups for the ImageNet classification 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 closer training methodologies used in modern works with state-of-the-art accuracies.
Supervised ImageNet 90 epochs training setup with vanilla ResNet. In the 90 epoch setup, we use the vanilla ResNet for fair comparison with prior works. We used the default hyperparameters found in official 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 le-4. The learning rate is scaled linearly from 0 to 0.15/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.9999 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 flip with 50% probability).
Most works compared against in Table use a similar training setup and also replace the 3x3 spatial convolutions in the ResNet architecture by their proposed methods. We note that Ramachandran al. (2019) train for longer (130 epochs instead of 90) but do not use label smoothing which could confound our comparisons.
Supervised ImageNet 350 epochs training setup. Higher accuracies ImageNet are commonly obtained by training longer with increased augmentation and regularization (Lee et al., 2020; Tan & Le, 2019). Similarly Bello et al. (2021), the weiht decay is reduced to 4e-5 and we employ RandAugment (Cubuk et al., 2019) with layers, dropout (Srivastava et al., 2014) and stochastic depth (Huang et al., 2016). See Table 20 for exact hyperparameters. All architectures 350 epochs with a batch size B of 4096 or 2048 distributed across 32 or 64 TPUv3 cores, depending on memory constraints.
trained for
We tuned our models using held-out validation set comprising ~2% of the ImageNet training set (20 shards out of 1024). We perform early stopping on the held-out validation set for the largest models, starting with LambdaResNett350 at resolution 288x288, and simply report the final accura- cies for the smaller models.
Semi-supervised learning with pseudo-labels. Our training setup closely follows the experimen- tal setup from Xie et al. (2020). We use the same dataset of 130M filtered and balanced JFT images with pseudo-labels generated by an EfficientNet-L.2 model with 88.4% ImageNet accuracy. Hyper- parameters are the same for the supervised ImageNet 350 epochs experiments.
Latency measurements. Figure 4 reports training latencies (i.e. time per training step) to process batch 1024 images on 8 TPUv3 cores using mixed precision training (1.e bfloat activations). Training latency is originally measured on 8 TPUv3 cores, starting with total batch size of 1024 (i.e. 128 per core) and dividing the batch size by until it fits in memory. We then report the normalized latencies in Figure For example, if latency measured with batch size of 512
30

---

Published as a conference paper at ICLR 2021
Depth
Image Size
RandAugment magnitude
Dropout
Stochastic depth rate
50
128
10
0.2
50
160
10
0.2
0
101
160
10
0.3
0
101
192
15
0.2
0
152
192
15
0.3
0
152
224
15
0.3
0.1
152
256
15
0.3
0.1
152
288
15
0.3
0.1
270
256
15
0.3
0.1
350
256
15
0.3
0.2
350
288
15
0.3
0.2
350
320
15
0.3
0.2
420
320
15
0.3
0.2
Table 20: Hyperparameters used to train LembdaResNets, We train for 350 epochs with Ran- dAugment, dropout and stochastic depth.
(instead of 1024), we normalize the reported latency by multiplying the measured latency by 2. Table 4, Table 12 and Table 13 report inference throughput on TPUv3 cores using full precision (i.e. 32 activations). Latency for VIT (Dosovitskiy et al., 2020) was privately communicated by the authors.
FLOPS count. We do not count zeroed out flops when computing positional lambdas with the einsum implementation from Figure Flops count is highly dependent on rather large by default (/m/=23x23). In Table 10, we show that it is possible to significantly reduce the scope size and therefore FLOPS at a minimal degradation performance.
the scope size which
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 1024x 1024 images from scratch for 130k steps with 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.32 and divided by 10 steps 90, 95 and 97.5% of training. The weight decay is set to 4e-5.
Mobilenet training setup. All mobilenet architectures are trained for 350 epochs on Imagenet with standard preprocessing 224x224 resolution. We the same hyperparameters et al. (2019). More specifically, we use RMSProp with 0.9 momentum and a batch size of 4096 split 32 TPUv3 cores. The learning rate warmed up linearly 0.1 and then multiplied by 0.99 every 3 epochs. We use weight decay le-5 and dropout with drop probability of 0.2
Howard
31


Reviews:
[{"review_id": "tXWsnm7v8sr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nHowever, there are still several issues to be addressed.\n1. This paper is not easy to follow. There are too many symbols and several of them are not explained. Besides, the organization of the paper can also be further improved.\n\n2. Some typos with the paper. E.g., In Table 1, tensor Q should have shape of nxkxu instead of mxkxu; lambda_n in eq. (2) should be transposed;\n\n3. In section 3.2, it is not clear how symbol u comes and what it means. \n\n4. Why are d and h removed from complexity analaysis in Table 4?\n\n5. Although authors explain its value to long sequences and high-resolution images, there is no experiment on the corresponding tasks such as long sequence language translation or high-resolution pixel-level prediction tasks.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "tEfyBgK0uyj", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a novel lambda layer to capture long-range interactions by transforming available contexts into linear functions, termed lambdas and applying these linear functions to each input separately. The proposed Lambda Network achieves good performances on ImageNet Classification, COCO object detection and instance segmentation tasks. The proposed lambda convolution is much more dense than the attention-based layer thus reducing parameters and complexity. However there are still several weaknesses in this paper. 1) Generalization of the proposed lambda convolution layer. For example, how about the performance of the lambda layer when combined with the lighter convolutional networks, e.g. mobilenet ? How about the performance when much deeper networks for the highest performance?  2)The source code is suggested to be released for more details. 3) Check the typos in the paper. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "YvTH9gaIVBz", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work proposes a lambda layer to capture long-range context. To address the issue of heavy memory-cost in self-attention, the proposed lambda layers transform context aggregation in self-attention into individual linear functions. The system-level performance is good on classification, detection and instance segmentation.\n\nStrengths:\n\n->Less memory and parameters, lower complexity than self-attention.\n\n->The performance is consistently good on classification, detection and instance segmentation.\n\n\nWeaknesses:\n\n->The main concern is technical novelty. Though interesting decomposition of query-context is introduced by linear functions, the overall design is still follow self-attention, the new part is incremental.\n\n->Some notations are missed in Table 1. Better to involve all the used notations in the Table.\n\n->The Lambda convolution performs local context, how about non-local context based on the proposed lambda function\n\n->What is SE in Table 7? The full name should be used before using its acronym.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Pfi8mGS4AX1", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n1. This paper present a new method, the lambda layer, for capturing long-term dependency.\n2. lambda layer summarizes the context into the fixed vector to reduce the computation burden\n3. This paper validates the performance on image classification and object detection tasks.\n\nStrengths\n1. The lambda layer is a simple and effective method.\n2. The idea of context summarization is interesting.\n3. Lambda layer shows the meaningful performance gain (including memory and time efficiency).\n\nWeaknesses\n1. This paper says that the lambda layer is a general framework, but the results only include vision tasks. Furthermore, there lacks discussion or results for the auto-regressive tasks. \n2. The comparison between the lambda layer and linear attention is not well discussed.\n\nQuestions and Additional Feedback\n1. Why Q=XW_{Q} is |m|\\times|k|\\times|u| sized tensor in table 1? \n2. The idea of reducing the complexity in the lambda layer seems similar to the “Transformers are rnns:” papers. This paper says that “We argue that such approaches may be overly restrictive and unnecessarily complex in trying to closely approximate an attention similarity kernel”. It is hard to agree with the sentence. What is the reason behind the sentence? Are there any empirical or theoretical results?\n3. This paper says that the lambda layer is “a general framework” and “versatile”. However, all of the experiments are related to the vision tasks. Are there results that related to natural language or time-series?\n4. I wonder how lambda layers work for auto-regressive tasks. There is some description of auto-regressive training in the appendix, but it is not enough.\n5. What is the advantage of lambda layers over attention? According to this paper, the only shortage of attention is an expensive computation. However, there are several linear time attention research. I suggest that this paper includes additional qualitative and quantitative analysis to compare the attention and linear time attention.\n\nTypos\n1. signifcantly => significantly\n2. as the the queries => as the queries\n3. formluations => formulations\n\nComments after the rebuttal\nThe author's response resolves my concern in part, and I will keep my positive score.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7PgvW8V0d3E", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a local attention + relative positional encoding type of network suited for image classification and detection tasks. \n\nThe first focus of the approach is to make attention scale to (2D) images. Vanilla (global) self-attention with an input of size $n$ (e.g. n=224x224 resized ImageNet) and a context size of $m$ (m=n if global) costs $nm$ memory. They decompose their approximation in two parts, the content attention and the positional embedding (which requires global attention). For the (dense) content part, in the same vein as multiple \"linear attention\" approximations (e.g. Linformer, Wang et al. 2020) they make this attention $nk$ with $k$ independent of $m$ and much smaller. For the (relative, this translation equivariant) positional embedding, the space cost is still $nm$, but doesn't depend on the image, so this factorization into content + position is beneficial for larger batch sizes.\n\nAnother contribution of this paper is to study the convolutional variant, so called \"lambda convolutions\" (strictly local relative position embedding) by setting the weights of the convolution dynamically based on the relative positional embeddings, and which can effectively reuse optimized [T|G]PU convolution kernels\n\nThey also break down the query in \"multiquery lambdas\" (followed by concatenation) to reduce the computational cost (as in grouped convolutions).\nFinally, they construct LambdaResNets by hybridation with vanilla ResNets where they replace any (see Table 12 for full results) of the convolution layers by lambdas for a parameters/throughput/accuracy trade-off.\n\nThey perform experiments on ImageNet (classification) and COCO (detection, with Mask-RCNN), which show competitive results: beating ResNets and relevant variants on size and accuracy, but not speed. LambdaResNets also beat EfficientNets on speed-accuracy on ImageNet accross the board.\n\n\nSome limitations of the paper and/or method include:\n- (minor) The related work (which is only really included in Appendix) does not discuss Zhao et al. 2020 (which is in Table 3).\n- A lot of the good/important content is in the Appendix (e.g. Appendix E.1 / Table 8 showing that **the positional embedding is absolutely necessary for good performance while the content part is quite optional**; or experiments on the scope size from E.2 / Table 9).\n- Their model can be seen (and indeed that is how they propose to implement it) as a kind of ConvNet with dynamically computed filters. Still, there is hope to recover long(-er than context size) range attention with multiple layers (as in ConvNets) and this is not studied/discussed.\n- (minor) Speed is problematic in the current implementation (see Table 12 and 13 in Appendix), how much of it is due to [optimized kernels for ResNets vs. einsum implementaion] vs. necessary computational cost? There is a bit of a FLOPS comparison in Tables 6, 8, 9 and accompanying text, but practice (e.g. Table 4, 12, 13 throughput) vs. theoretical complexity is not discussed.\n- A small caveat is that the space cost vs. global (self-) attention is only really advantageous for large batch sizes.\n\n\nOverall, this is a good, readable, well studied (if one considers the appendix) paper on a promising new hybrid conv/attention layer that yields small and accurate models for computer vision core tasks (classification and detection). \n\n\nTypo:\n\"=This section\" in 3.2", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces LambdaNetworks, a new method to capture long range interactions in data (such as global context in images). The method is novel and simple, and the experimental results are strong, especially in terms of speed/accuracy tradeoffs. The paper is well written and easy to follow. For these reasons, I recommend to accept the paper.
Decision: Accept (Spotlight)