[Paper Review] Efficient Attention: Attention with Linear Complexities

Zhuoran Shen, Mingyuan Zhang+3 more|arXiv (Cornell University)|Dec 4, 2018

Advanced Neural Network ApplicationsComputer Science22 references88 citations

TL;DR

Proposes an efficient attention mechanism that is mathematically equivalent to dot-product attention with linear memory and computation, enabling high-resolution and resource-intensive tasks and delivering improvements on MS-COCO and Scene Flow datasets.

ABSTRACT

Dot-product attention has wide applications in computer vision and natural language processing. However, its memory and computational costs grow quadratically with the input size. Such growth prohibits its application on high-resolution inputs. To remedy this drawback, this paper proposes a novel efficient attention mechanism equivalent to dot-product attention but with substantially less memory and computational costs. Its resource efficiency allows more widespread and flexible integration of attention modules into a network, which leads to better accuracies. Empirical evaluations demonstrated the effectiveness of its advantages. Efficient attention modules brought significant performance boosts to object detectors and instance segmenters on MS-COCO 2017. Further, the resource efficiency democratizes attention to complex models, where high costs prohibit the use of dot-product attention. As an exemplar, a model with efficient attention achieved state-of-the-art accuracies for stereo depth estimation on the Scene Flow dataset. Code is available at https://github.com/cmsflash/efficient-attention.

Motivation & Objective

• Motivate the need for global dependency modeling at large input sizes due to quadratic costs of standard dot-product attention.
• Introduce an alternative attention formulation with linear memory and compute that is equivalent to dot-product attention under scaling normalization.
• Show that the efficient attention enables integrating more attention modules and higher-resolution features within networks.
• Demonstrate substantial performance gains on object detection, instance segmentation, and stereo depth estimation under resource constraints.

Proposed method

• Recast dot-product attention by rearranging matrix multiplications to compute a global context via Q and K'V, yielding linear memory and compute.
• Define efficient attention E(Q,K,V) = rho_q(Q) (rho_k(K)^T V) with rho_q and rho_k as normalization functions (scaling or softmax).
• Prove equivalence under scaling normalization: E(Q,K,V) = D(Q,K,V) when using scaling normalization.
• Provide implementation details for vision data by flattening inputs and applying the mechanism with residual connections.
• Offer an interpretation where K is viewed as d_k template attention maps, forming global context vectors that Q coefficients combine per position.

Experimental results

Research questions

• RQ1 Can efficient attention achieve the same representational power as standard dot-product attention while reducing resource usage?
• RQ2 How does the choice of normalization (scaling vs. softmax) affect equivalence and accuracy?
• RQ3 What are the practical gains in memory, computation, and accuracy when integrating efficient attention into high-resolution vision models?

Key findings

• Efficient attention matches dot-product attention under scaling normalization and closely approximates it under softmax, with empirical accuracy preserved.
• Memory and computation scale linearly with input size, vs. quadratic for traditional dot-product attention, enabling high-resolution usage.
• In MS-COCO, efficient attention yields superior performance-cost trade-offs and enables memory-constrained placements where non-local modules fail.
• Efficient attention delivers strong gains in stereo depth estimation (Scene Flow), achieving state-of-the-art while using far less memory than non-local variants.
• Ablation studies show normalization method and key dimensionality have limited impact on performance, enabling resource-saving configurations.