Selective Attention Improves Transformer

2024-10-03

arXiv.org

AI Summary

Transformer attention mechanisms are susceptible to interference from irrelevant contextual tokens, leading to degraded performance and inefficient resource utilization. To address this, we propose Selective Attention—a parameter-free, plug-and-play attention sparsification mechanism built upon standard scaled dot-product attention. It introduces a lightweight gating filter that dynamically suppresses attention weights assigned to irrelevant tokens, without introducing additional model parameters or modifying the training objective. The method is fully compatible with arbitrary Transformer architectures and training paradigms. When evaluated on C4-trained language models, Selective Attention achieves comparable perplexity while reducing attention module parameters by 50%. For context lengths of 512, 1024, and 2048, it reduces memory consumption to 1/16, 1/25, and 1/47 of the baseline, respectively—significantly enhancing computational efficiency and scalability.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Security and Privacy: Large-scale security measurements
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Unneeded elements in the attention's context degrade performance. We introduce Selective Attention, a simple parameter-free change to the standard attention mechanism which reduces attention to unneeded elements. Selective attention consistently improves language modeling and downstream task performance in a variety of model sizes and context lengths. For example, transformers trained with the language modeling objective on C4 with selective attention perform language modeling equivalently to standard transformers with ~2X more heads and parameters in their attention modules. Selective attention also allows decreasing the size of the attention's context buffer, leading to meaningful reductions in the memory and compute requirements during inference. For example, transformers trained on C4 with context sizes of 512, 1,024, and 2,048 need 16X, 25X, and 47X less memory for their attention module, respectively, when equipped with selective attention, as those without selective attention, with the same validation perplexity.

Problem

• Reduces attention to unneeded elements in transformers
• Improves language modeling and downstream task performance
• Decreases memory and compute requirements during inference

Innovation

• Parameter-free selective attention mechanism
• Reduces attention to unneeded elements
• Decreases memory and compute requirements