[Paper Review] LambdaNetworks: Modeling Long-Range Interactions Without Attention

Irwan Bello|arXiv (Cornell University)|Feb 17, 2021

Advanced Neural Network Applications60 references48 citations

TL;DR

LambdaNetworks introduce lambda layers that transform context into linear functions to model long-range interactions without attention, achieving strong performance on ImageNet, COCO, and efficiency gains over attention-based methods.

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

Motivation & Objective

• Motivate modeling long-range dependencies in structured inputs like images without costly attention mechanisms.
• Develop lambda layers that summarize context into linear functions applied to each query.
• Demonstrate that lambda layers can outperform convolutions and attention in vision tasks with better efficiency.

Proposed method

• Define queries and contexts and formalize content-based and position-based interactions.
• Introduce lambda layer that computes context lambdas as a sum of content and position components (lambda n = K^T V + E_n^T V).
• Use multi-query formulation to reduce time/space complexity (concatenate outputs of severalq-derived lambdas).
• Optionally enforce translation equivariance via relative position embeddings and present lambda convolution for local contexts.
• Compare against convolutions and various attention variants through ImageNet, COCO, and semi-supervised experiments.

Experimental results

Research questions

• RQ1 Can lambda layers capture global long-range interactions in images without materializing attention maps?
• RQ2 Do lambda layers offer speed and memory advantages over self-attention while improving accuracy on vision tasks?
• RQ3 How do hybrid convolution-lambda architectures perform in terms of speed-accuracy tradeoffs across scales?
• RQ4 What is the impact of multi-query formulation and translation-equivariant embeddings on performance and efficiency?
• RQ5 Are lambda-based backbones beneficial for object detection and instance segmentation tasks?

Key findings

• Lambda layers outperform convolution and attention baselines on ImageNet with fewer parameters (e.g., +1.5% top-1 over Conv baseline with ResNet-50).
• Lambda networks offer lower memory footprints and higher throughput than self-attention variants while achieving comparable or better accuracy.
• LambdaResNets deliver significantly faster speed-accuracy tradeoffs, up to 3.2–4.4x faster than EfficientNets for similar accuracy and up to 9.5x speedup with pseudo-labeled data in semi-supervised training.
• LambdaResNets achieve 84.9% top-1 accuracy on ImageNet (LambdaResNet-420 at 320 input) and improve COCO box/mask metrics over baselines.
• In Mask-RCNN with LambdaResNet backbones, small objects see notable gains, indicating improved localization performance.