Hydra Attention: Efficient Attention with Many Heads

2022/09/15 by Bolya, Daniel, Fu, Cheng-Yang, Dai, Xiaoliang +2 · 4 citations
#Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences

paper · doi:10.48550/arxiv.2209.07484

Abstract

While transformers have begun to dominate many tasks in vision, applying them to large images is still computationally difficult. A large reason for this is that self-attention scales quadratically with the number of tokens, which in turn, scales quadratically with the image size. On larger images (e.g., 1080p), over 60% of the total computation in the network is spent solely on creating and applying attention matrices. We take a step toward solving this issue by introducing Hydra Attention, an extremely efficient attention operation for Vision Transformers (ViTs). Paradoxically, this efficiency comes from taking multi-head attention to its extreme: by using as many attention heads as there are features, Hydra Attention is computationally linear in both tokens and features with no hidden constants, making it significantly faster than standard self-attention in an off-the-shelf ViT-B/16 by a factor of the token count. Moreover, Hydra Attention retains high accuracy on ImageNet and, in some cases, actually improves it.

Cited by

• Attentive Convolution: Unifying the Expressivity of Self-Attention with Convolutional Efficiency
• DistrAttention: An Efficient and Flexible Self-Attention Mechanism on Modern GPUs
• Rectifying Magnitude Neglect in Linear Attention
• Norm×Direction: Restoring the Missing Query Norm in Vision Linear Attention

Related

• HydraViT: Stacking Heads for a Scalable ViT
• Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet
• Artifacts and Attention Sinks: Structured Approximations for Efficient Vision Transformers
• Scratching Visual Transformer's Back with Uniform Attention
• You Only Need Less Attention at Each Stage in Vision Transformers