Reorganizing attention-space geometry with expressive attention
This work addresses a foundational issue in attention-based models for machine learning, offering a simple modification that improves performance on complex tasks without extra computational cost, though it is incremental in nature.
The paper tackles the problem of attention mechanism geometry by proposing expressive attention (EA) based on squared dot product, which enhances attention for parallel or antiparallel queries and keys. It finds that EA performs at least as well as vanilla dot-product attention on simple tasks and outperforms it with increasing margins for more complex or multi-task settings, achieving 100% performance in some cases where DPA fails.
Attention regulates information transfer between tokens. For this, query and key vectors are compared, typically in terms of a scalar product, $\mathbf{Q}^T\mathbf{K}$, together with a subsequent softmax normalization. In geometric terms, the standard dot-product attention (DPA) leads to large/small attention weights for parallel/antiparallel queries and keys. Here we study expressive attention (EA), which is based on $(\mathbf{Q}^T\mathbf{K})^2$, the squared dot product. In this case, attention is enhanced when query and key are either parallel or antiparallel, and suppressed for orthogonal configurations. EA can be introduced into any attention-based code without additional compute costs or memory requirements. For a series of autoregressive prediction tasks, we find that expressive attention performs at least as well as vanilla DPA. Increasing task complexity, EA is observed to outperform DPA with increasing margins, which also holds for multi-task settings. For a given model size, EA manages to achieve 100% performance for a range of complexity levels not accessible to DPA. Our results show that it is possible to reorganize the geometry of the matching condition in the space of attention heads without loss of performance.