TLDR: DistrAttention is an efficient and flexible self-attention mechanism designed for modern GPUs. It optimizes Transformer performance by grouping data along the embedding dimensionality using locality-sensitive hashing, which reduces computational complexity while maintaining full contextual information. The method integrates seamlessly with FlashAttention-2, achieving significant speedups (up to 37% faster than FlashAttention-2 for self-attention) and high accuracy in vision and language models like ViT and Llama3-1B.
The Transformer architecture has profoundly changed the landscape of deep learning, driving breakthroughs in areas like natural language processing, computer vision, and time series prediction. At its core lies the self-attention mechanism, a powerful component that allows models to weigh the importance of different parts of an input sequence. However, this very mechanism comes with a significant challenge: its computational complexity grows quadratically with the length of the input sequence. This means that as sequences get longer, the time and resources required to process them increase dramatically, limiting the scalability of Transformers for handling very long texts or high-resolution images.
Existing methods to optimize self-attention often involve trade-offs. Some approaches, like sparse attention, reduce the amount of information considered, which can lead to a loss of important context. Others, such as linear attention mechanisms, simplify interactions to improve efficiency but might miss complex relationships across the entire sequence. Quantization techniques reduce the precision of model parameters for faster inference but offer limited flexibility in balancing speed and accuracy.
A new approach, called DistrAttention, offers an efficient and flexible self-attention mechanism that maintains full contextual information. Instead of shortening the input sequence or quantizing parameters, DistrAttention focuses on reducing the embedding dimensionality, often referred to as ‘d’. This is a crucial distinction because changing ‘d’ does not alter the size of the attention matrix, thereby preserving the full context of the input.
DistrAttention achieves this by intelligently grouping data along the embedding dimension. It employs a lightweight ‘sampling and fusion’ method that leverages locality-sensitive hashing (LSH) to identify and group similar data points. Imagine having many columns of data (Q) and rows of data (K) that need to interact. DistrAttention finds similar columns in Q and groups them, then sums the corresponding rows in K. By using an estimated representation for each group, it significantly reduces the number of multiplications needed for the attention calculation.
To further enhance its performance and limit any errors introduced by the grouping process, DistrAttention incorporates a ‘block-wise grouping’ framework. This design allows it to seamlessly integrate with state-of-the-art high-performance self-attention implementations like FlashAttention-2, which are optimized for modern GPUs. By carefully selecting block sizes, DistrAttention can maximize performance on these powerful hardware platforms.
Extensive experiments have demonstrated the effectiveness of DistrAttention. It has been shown to be up to 37% faster than FlashAttention-2 when calculating self-attention. In tasks like ViT (Vision Transformer) inference, DistrAttention emerged as the fastest and most accurate among approximate self-attention mechanisms. For large language models, specifically Llama3-1B, DistrAttention achieved the lowest inference time with only a minimal 1% accuracy loss. This method can also be integrated into pre-trained models with little effort, as it doesn’t change the output shape or introduce new parameters.
Also Read:
- LaCache: A Smart Memory Solution for Long-Context LLMs
- Rethinking LLM Performance: How Latent Attention and Mixture-of-Experts Reshape the Bottleneck
The benefits extend to multi-GPU setups as well, where DistrAttention shows notable acceleration. Even the LSH-based grouping component, which is critical for its operation, has a minimal impact on overall performance, especially for longer sequences where its overhead becomes negligible compared to the computational gains. This innovative approach addresses a core bottleneck in Transformer models, paving the way for more efficient and scalable deep learning applications. You can read the full research paper here.
