Adaptive Multi-Resolution Attention with Linear Complexity

TL;DR

The paper introduces AdaMRA, an efficient multi-resolution attention mechanism for Transformers that achieves linear complexity and improves long-range information capture, leading to state-of-the-art results.

Contribution

It proposes a novel multi-resolution attention structure with query-driven resolution selection and kernel attention for linear complexity in Transformers.

Findings

Achieves state-of-the-art performance on multiple benchmarks.

Demonstrates significant efficiency and memory improvements.

Maintains performance with reduced computational complexity.

Abstract

Transformers have improved the state-of-the-art across numerous tasks in sequence modeling. Besides the quadratic computational and memory complexity w.r.t the sequence length, the self-attention mechanism only processes information at the same scale, i.e., all attention heads are in the same resolution, resulting in the limited power of the Transformer. To remedy this, we propose a novel and efficient structure named Adaptive Multi-Resolution Attention (AdaMRA for short), which scales linearly to sequence length in terms of time and space. Specifically, we leverage a multi-resolution multi-head attention mechanism, enabling attention heads to capture long-range contextual information in a coarse-to-fine fashion. Moreover, to capture the potential relations between query representation and clues of different attention granularities, we leave the decision of which resolution of attention to use to query, which further improves the model's capacity compared to vanilla Transformer. In an effort to reduce complexity, we adopt kernel attention without degrading the performance. Extensive experiments on several benchmarks demonstrate the effectiveness and efficiency of our model by achieving a state-of-the-art performance-efficiency-memory trade-off. To facilitate AdaMRA utilization by the scientific community, the code implementation will be made publicly available.