PaperlayerView exact revision on arXiv

Source-linked AI summary

LambdaNetworks: Modeling Long-Range Interactions Without Attention

Irwan Bello

arXiv:2102.08602v1cs.CVcs.LG

TL;DR

Long-range modeling with self-attention is memory-intensive, while linear attention does not capture internal structure in data such as relative pixel positions. The paper introduces lambda layers that summarize contexts into fixed-size linear functions applied to queries, and reports that LambdaNetworks outperform convolutional and attentional alternatives across vision tasks while improving computational efficiency.

  • Problem

    Self-attention has costly memory requirements for long sequences and multidimensional data, while linear attention fails to model internal structure such as relative distances between pixels.

  • Method

    Lambda layers transform structured contexts into fixed-size linear functions, or lambdas, that are applied independently to corresponding queries without materializing attention maps.

  • Results

    LambdaNetworks significantly outperform convolutional and attentional counterparts on ImageNet classification, COCO object detection, and COCO instance segmentation while being more computationally efficient.

  • Takeaways & Limitations

    LambdaNetworks provide a scalable framework for structured long-range interactions in large inputs, and LambdaResNets improve the speed-accuracy tradeoff of image classification models.

  • Takeaways & Limitations

    Reported local self-attention latencies came from a query-block and overlapping-memory-block implementation, and specialized kernels may substantially speed it up.

Abstract

from arXiv · show

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 COCO 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. When training with an additional 130M pseudo-labeled images, LambdaResNets achieve up to a 9.5x speed-up over the corresponding EfficientNet checkpoints.

1 INTRODUCTION

Lambda layers address the memory cost of self-attention and the structural limitations of linear attention by converting contexts into fixed-size linear functions applied to queries. LambdaNetworks retain long-range, content- and position-based interactions while delivering strong vision-task performance and efficiency.

  • Motivation: Self-attention’s memory cost hinders its use on long sequences and multidimensional inputs, while linear attention does not model internal structure such as relative pixel distances.A multi-head attention layer for 128 64×64 images with 8 heads requires 64GB of memory.
  • Approach: Lambda layers transform each context into a fixed-size linear function and apply it directly to the corresponding query, bypassing memory-intensive attention maps.This framework models long-range interactions between queries and structured context elements at reduced memory cost.
  • Scope: Lambda layers model both content-based and position-based interactions in global, local, or masked contexts, supporting large structured inputs such as high-resolution images.Their implementations use einsum operations and convolution kernels with efficient accelerator implementations.
  • Results: LambdaNetworks significantly outperform convolutional and attentional counterparts across ImageNet classification, COCO object detection, and COCO instance segmentation while being more computationally efficient.The paper also reports lower latency and memory requirements than multiple self-attention alternatives.
  • Results: +1.5% top-1 ImageNet accuracy and 40% fewer parameters result when replacing ResNet-50 bottleneck 3×3 convolutions with lambda layers.The comparison is specifically against the original ResNet-50 architecture.

2 MODELING LONG-RANGE INTERACTIONS

The paper formulates structured long-range interactions through queries, contexts, and their content- and position-based relationships. Lambda layers aggregate context into fixed-size linear functions that independently map each query to an output without producing attention maps.

  • Definitions: Queries and context elements are structured vectors characterized by content and position, with pairwise relations able to represent pixel distances or graph edges.A query q_n maps to an output y_n given the context C.
  • Definitions: Content-based interactions use query and context contents, whereas position-based interactions use the relation between query position n and context position m.Absolute positional information added directly to content is treated as content-based because it ignores the relation (n, m).
  • Interaction construction: Keys and relative position embeddings provide vectors for dot-product interactions with queries, capturing content and positional structure.Keys depend on context content, while relative position embeddings encode relations between query and context positions.
  • Comparison with attention: Attention contracts query depth first to form a similarity kernel, but its attention map can become wasteful when context positions greatly outnumber output dimensions.The relevant regime is |v| ≪|m|, where the output dimension is much smaller than the number of context positions.
  • Lambda layers: A lambda layer maps inputs X and context C to lambdas that are applied to queries, producing outputs Y without materializing attention maps.For query position n, the context is aggregated into λ_n, a fixed-size linear function that is discarded after application.

3 LAMBDA LAYERS

Lambda layers transform structured context into fixed-size linear functions and apply them independently to queries, capturing content- and position-based interactions without materializing attention maps. Multi-query and local convolution formulations reduce complexity while preserving long-range modeling.

  • Lambda layer formulation: Lambda layers take inputs and context, generate linear functions, and apply those functions to queries to produce outputs.The context may equal the input, as in self-attention.
  • Content and position interactions: Content lambdas aggregate context values independently of context ordering, whereas position lambdas encode context elements and their relative positions to each query.Together, these lambdas dynamically distribute fixed-size contextual features according to each query.
  • Complexity and hybrid designs: Global lambda layers retain quadratic time complexity, motivating hybrid architectures that use convolutions for early local processing and lambdas later.Large output dimensions can also increase cost, while multi-query formulations decouple complexity from output dimension.
  • Multi-query formulation: A multi-query lambda layer uses multiple queries with smaller value dimensions, reducing complexity by a factor of |h|.The number of heads controls lambda size relative to the total query representation.
  • Position interactions: Relative position embeddings can impose translation equivariance by making position interactions depend consistently on translated query and context positions.This structure is encoded through relative position embeddings.
  • Lambda convolution: Lambda convolution restricts positional interactions to local grid neighborhoods and achieves linear time and memory complexity with respect to input length.It can also support dilation and striding and use optimized accelerator implementations.

4 RELATED WORK

Lambda layers are related to channel and linear attention because they summarize context without attention maps, while extending this perspective to position-based interactions. The paper also situates its visual results against the costs and data requirements of self-attention approaches.

  • Channel and linear attention: Channel and linear attention can be viewed as content-only instances of the lambda abstraction that capture long-range interactions without materializing attention maps.The lambda framework therefore provides a unifying view of these mechanisms.
  • Self-attention in vision: Self-attention has adopted more slowly in vision, while a concurrent Transformer reached 88.6% ImageNet accuracy after pre-training on 300M images.The comparison highlights the scale of data used by that visual Transformer result.
  • Computational caveat: Although lambda convolution has linear time and space complexity, floating-point operation counts may not predict accelerator latency, especially for large convolution scopes.The global lambda layer can be faster in practice in such cases.

5 EXPERIMENTS

Experiments evaluate lambda layers and LambdaResNets on ImageNet, COCO detection, and COCO instance segmentation, emphasizing accuracy, speed, memory, and scaling. Across these benchmarks, lambda-based models outperform convolutional and attentional baselines while retaining computational advantages.

  • 5.2 COMPUTATIONAL BENEFITS OF LAMBDA LAYERS OVER SELF-ATTENTION.: A +1.0% ImageNet improvement over local self-attention accompanies nearly 3x higher speed, while lambda layers capture global interactions on high-resolution images.Positional embeddings can be shared across lambda layers with minimal degradation, and lambda convolution has linear memory complexity for very large images.
  • 5.3 HYBRIDS IMPROVE THE SPEED-ACCURACY TRADEOFF OF IMAGE CLASSIFICATION.: Hybrid architectures address the cost of global contexts by reserving them for later, more semantic features while retaining locality where it is useful.Global lambda contexts still incur quadratic time complexity, motivating the hybrid design.
  • 5.3 HYBRIDS IMPROVE THE SPEED-ACCURACY TRADEOFF OF IMAGE CLASSIFICATION.: 3.2–4.4x faster than EfficientNets and 1.6–2.3x faster than ResNet-RS, LambdaResNets improve the ImageNet speed-accuracy trade-off across all scales.LambdaResNet-420 at image size 320 reaches 84.9% top-1 accuracy, 0.9% above the corresponding standard-convolution architecture and 0.65% above the squeeze-and-excitation version.
  • COCO OBJECT DETECTION AND INSTANCE SEGMENTATION: Using lambda layers in Mask-RCNN backbones produces consistent gains on COCO detection and instance segmentation across object sizes, especially for small objects.The results indicate competitiveness on visual tasks requiring localization information.

6 DISCUSSION

Lambda layers offer a scalable alternative to attention by summarizing context into fixed-size linear functions, while supporting content- and position-based interactions. The discussion covers their complexity advantages, practical design choices, broader applicability, and implementation trade-offs.

  • Efficiency versus attention: Θ(lkn^2) memory for LambdaNetworks, versus Θ(blhn^2) for vanilla Transformers, enables higher-resolution inputs and larger batch sizes.Sharing positional embeddings across layers reduces the footprint to Θ(kn^2).
  • Relation to linear attention: Lambda layers extend linear attention by capturing position-based interactions, which is important for highly structured inputs such as images.They also permit more flexible nonlinearities and normalizations because they do not approximate an attention kernel.
  • Visual-domain use: Lambda layers are most beneficial in intermediate and low-resolution vision stages when optimizing the speed-accuracy trade-off.Architectures using only lambda layers can be more parameter- and FLOP-efficient, while hybrid designs improve practical scalability and speed.
  • Generality: Lambda layers can model interactions on graphs, time series, spatial lattices, and other structures beyond the static image tasks studied here.The paper also discusses masked contexts and autoregressive tasks in the appendix.
  • Visual-domain use: LambdaResNets achieve up to 9.5x speed-ups over popular EfficientNets while maintaining excellent ImageNet accuracy.The hybrid family is presented as substantially improving the speed-accuracy trade-off of image classification models.
  • Implementation choices: For lower computation, smaller-scope lambda convolutions or standard 3x3 convolutions can replace global contexts in early high-resolution layers.Sharing position embeddings, reducing query depth |k|, or increasing the number of heads |h| can reduce memory; implementation speed may vary across hardware stacks.

C.3 LINEAR ATTENTION: CONNECTIONS AND DIFFERENCES

Lambda layers extend linear attention by incorporating position-based interactions while retaining fixed-size linear mappings that avoid materializing attention maps. The framework unifies several attention mechanisms and supports structured visual inputs.

  • Linear attention: Linear attention avoids materializing attention maps by computing context summaries before applying them to queries.Kernel factorization yields linear complexity with respect to input length.
  • Differences from linear attention: Lambda layers extend linear attention with position-based interactions in addition to content-based interactions.This supports relations such as relative positions rather than only content similarity.
  • Differences from linear attention: Lambda layers do not require approximating an attention kernel, allowing more flexible nonlinearities and normalization schemes.The design is therefore less restricted than kernel-based attention approximations.
  • Computational formulation: The lambda abstraction exposes computational benefits of multi-query formulations over traditional multi-head attention.Contexts are transformed into content and position lambdas that are directly applied to queries.
  • Connections to attention: Lambda layers generalize channel and spatial attention, and experiments report better efficiency and performance than both.Channel attention uses shared diagonal lambdas, while spatial attention uses position-dependent scalar lambdas.
  • Visual-domain context: Lambda layers retain position modeling while avoiding the large memory costs that restrict global self-attention to low-resolution feature maps.Related approaches include local, axial, patch-based, and hybrid attention designs.

D.1 ABLATION STUDY

Ablations show that LambdaNetworks remain strong across hyperparameters, depend especially on positional interactions, benefit from larger scopes, and require effective normalization choices.

  • Hyperparameters: Lambda layers outperform convolutional and attentional baselines across a wide range of query depths, head counts, and intra-depths.All tested configurations outperform the convolutional baseline at lower parameter cost; intermediate head counts work best.
  • Interaction types: Position-based interactions are crucial for high ImageNet accuracy, whereas content-based interactions provide only marginal additional improvement.A concurrent result found content-only interactions sufficient when training on 300M images.
  • Scope size: Larger positional scopes improve results until performance plateaus around |m|=15x15.The comparison uses lambda layers replacing 3x3 convolutions in the last two ResNet-50 stages.
  • Normalization: Key normalization is crucial for performance, while alternatives to softmax and batch normalization on queries and values are also helpful.The ablation varies normalization along context spatial and query-depth dimensions.

D.2 HYBRID MODELS STUDY

Hybrid LambdaNetworks combine convolutions for local contexts with lambda layers for global contexts, placing lambdas mainly in later ResNet stages to improve the speed-accuracy tradeoff.

  • Hybrid design: Hybrid designs retain standard convolutions for local contexts and use lambda layers for global contexts.The paper also notes that lambda convolutions could alternatively capture local contexts.
  • Layer placement: Lambda layers are most helpful in ResNet stages c4 and c5 for the speed-accuracy tradeoff.LambdaResNet-C4 replaces 3x3 convolutions in the last two stages.
  • Layer count: A few lambda layers in c4 and all lambda layers in c5 provide most of the hybrid-model benefits.These configurations increase representational power with a virtually negligible throughput decrease.
  • Comparison with attention: For LambdaResNet-420, lambda layers require 38.4MB, or 4.8MB with shared positional embeddings, versus 0.625GB for attention layers.The model uses 320x320 inputs and eight lambda layers in c4.

D.3 COMPUTATIONAL EFFICIENCY RESULTS

LambdaResNet-C4 improves parameter and FLOP efficiency over large EfficientNets, although one cited comparison remains incomplete because hybrid attention variants were deferred to future work.

  • Efficiency comparisons: LambdaResNet-C4 improves parameter and FLOP efficiency over large EfficientNets despite EfficientNet’s use of computationally efficient depthwise convolutions.The comparison concerns replacing 3x3 convolutions in the last two ResNet stages.
  • Scope of evidence: Hybrid lambda-versus-attention architectures were not fully benchmarked in the current draft.The authors defer that benchmark to a future version.
  • Parameter efficiency: LambdaResNet-C4 is more parameter-efficient than EfficientNet-B6 despite using a smaller image size.The table caption notes that increasing image size could improve accuracy while keeping parameters fixed.
  • FLOP efficiency: The FLOP comparison uses smaller 7x7 local scopes in lambda layers to reduce computational cost.This setting is specific to the EfficientNet-B6 comparison.

D.3.2 LAMBDA LAYERS IN A RESOURCE CONSTRAINED SCENARIO

In resource-constrained MobileNetv2 models, lightweight lambda blocks replace selected depthwise convolutions while reducing computation. LambdaResNets similarly replace spatial convolutions in ResNet-RS bottlenecks with lambda layers and use staged placement and downsampling.

  • MobileNetv2: Lightweight lambda blocks replace MobileNetv2 depthwise convolutions while retaining pointwise channel operations in the inverted bottleneck design.The original block uses pointwise expansion, depthwise spatial mixing, and final pointwise channel mixing.
  • MobileNetv2: Replacing the 10-th and 16-th MobileNetv2 inverted bottlenecks with lightweight lambda blocks reduces parameters and flops by ∼10% while improving ImageNet accuracy by 0.6%.The lightweight blocks use a 5x5 scope, query depth 32, and 4 heads, while the first pointwise convolution preserves channel count to reduce computation.
  • Lambda layer design: Default lambda layers use query depth 16, 4 heads, intra-depth 1, 23x23 local positional contexts, and global content contexts.The implementation uses an einsum formulation for global content lambdas, while local positional lambdas can use lambda convolution or masking.
  • LambdaResNets: LambdaResNets replace spatial 3x3 convolutions in ResNet-RS bottlenecks with lambda layers, leaving the stem unchanged.They use 3x3 average pooling with stride 2 after lambda layers for downsampling, place lambda layers uniformly in c4, and use them throughout c5.
  • LambdaResNets: LambdaResNets use block allocations from prior ResNet designs, with their exact configurations and lambda-layer placement reported for Pareto-curve models.The cited configuration table covers block allocation and lambda-layer placement across the LambdaResNet family.

E.2 TRAINING DETAILS

The experiments use distinct ImageNet training regimes, downstream COCO detection settings, and accelerator-based latency and FLOPS measurements. Training includes supervised, semi-supervised, and resource-aware evaluation protocols with explicitly reported implementation choices and comparison caveats.

  • ImageNet training: ImageNet models are trained under 90-epoch standard and 350-epoch augmented-and-regularized setups to support fair comparisons and higher-accuracy evaluations.The 90-epoch setup uses standard preprocessing, while the 350-epoch setup is closer to modern state-of-the-art training methodologies.
  • ImageNet training: The 90-epoch vanilla ResNet setup uses end-to-end SGD training with momentum 0.9, a batch size of 4096 across 32 TPUv3 cores, and weight decay 1e-4.The setup uses default hyperparameters from official implementations without additional tuning.
  • Comparison caveat: Comparisons may be confounded because Ramachandran et al. train for 130 rather than 90 epochs and omit label smoothing.Most other compared works use similar training setups and replace ResNet 3x3 spatial convolutions with their proposed methods.
  • ImageNet training: The 350-epoch setup adds RandAugment, dropout, stochastic depth, and reduced weight decay of 4e-5.These choices follow the higher-accuracy training configuration described for ImageNet.
  • Semi-supervised learning: Semi-supervised training uses 130M filtered and balanced JFT images with pseudo-labels generated by an EfficientNet-L2 model reaching 88.4% ImageNet accuracy.Its hyperparameters match the supervised ImageNet 350-epoch experiments.
  • Efficiency measurements: Latency is measured as normalized training time per step on 8 TPUv3 cores with bfloat16 activations, while reported inference throughput uses float32 activations.FLOPS counts depend strongly on positional-lambda scope size, and smaller scopes can reduce FLOPS with minimal performance degradation.
  • COCO object detection: COCO detection uses Mask R-CNN backbones trained from scratch on 1024x1024 images for 130k steps with batch size 256 across 128 TPUv3 cores.Training includes synchronized batch normalization, multi-scale jitter, learning-rate warmup, and scheduled decay.
Loading 2102.08602v1…