Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention (2405.17381v2)

Abstract: We present Lightning Attention, the first linear attention implementation that maintains a constant training speed for various sequence lengths under fixed memory consumption. Due to the issue with cumulative summation operations (cumsum), previous linear attention implementations cannot achieve their theoretical advantage in a casual setting. However, this issue can be effectively solved by utilizing different attention calculation strategies to compute the different parts of attention. Specifically, we split the attention calculation into intra-blocks and inter-blocks and use conventional attention computation for intra-blocks and linear attention kernel tricks for inter-blocks. This eliminates the need for cumsum in the linear attention calculation. Furthermore, a tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. To enhance accuracy while preserving efficacy, we introduce TransNormerLLM (TNL), a new architecture that is tailored to our lightning attention. We conduct rigorous testing on standard and self-collected datasets with varying model sizes and sequence lengths. TNL is notably more efficient than other LLMs. In addition, benchmark results indicate that TNL performs on par with state-of-the-art LLMs utilizing conventional transformer structures. The source code is released at github.com/OpenNLPLab/TransnormerLLM.

Efficient LLMing with Lightning Attention

The paper "Various Lengths, Constant Speed: Efficient LLMing with Lightning Attention" introduces a novel linear attention mechanism termed "Lightning Attention" aimed at addressing the inefficiencies found in current linear attention implementations, particularly under the constraints of fixed memory consumption and varying sequence lengths. The authors propose an innovative approach to computational attention by separating operations into intra-blocks and inter-blocks, using distinct computation methods for each. This paper provides a comprehensive examination of both the theoretical and practical impacts of their model, demonstrating substantial improvements in efficiency and accuracy compared to state-of-the-art models.

Linear Attention Background

Linear attention mechanisms have garnered interest as they potentially offer computational advantages over traditional softmax attention structures by theoretically providing linear complexity rather than quadratic. However, existing linear models have struggled with maintaining their theoretical benefits in practical applications, primarily due to the overhead of cumulative summation operations (cumsum) in a causal setting. This paper highlights how these computational overheads negate the perceived benefits in actual deployment scenarios.

Lightning Attention

Lightning Attention fundamentally circumvents the limitations observed in previous models by leveraging a "divide and conquer" strategy:

1. Intra-Block Computation: Utilizes conventional quadratic attention computation which remains local to specific data partitions.
2. Inter-Block Computation: Employs a kernel trick for operations among blocks that bypasses the cumsum operation, thus achieving linear complexity and utilizing GPU hardware efficiently.

In addition to this strategic separation of operations, the authors employ techniques to optimize memory usage and computational speed further using tiled computations in both forward and backward passes. The overall complexity is reduced to $O(nd^2 + nBd)$, where $n$ represents sequence length and $B$ the block size, offering significant improvements over previous implementations.

TransNormerLLM Integration

To augment the capabilities of Lightning Attention, the paper introduces the TransNormerLLM (TNL) architecture. This new architecture integrates several advanced modifications including:

• Positional Encoding: Incorporates LRPE with an exponential decay to maintain token interactions, ensuring global context is preserved at minimal additional computational cost.
• Gating Mechanisms: Includes a Gated Linear Attention (GLA) to enhance the model's stability and performance.
• Tensor Normalization: Deploys a simplified normalization technique, SRMSNorm, which improves the model's computational efficiency without sacrificing accuracy.

Empirical Evaluation

The empirical evaluation conducted demonstrates that TNL models exhibit superior efficiency and performance across multiple benchmarks when compared against both transformer-based architectures and other efficient models like HGRN and TNN. Specifically, in a suite of language-modeling tasks, TNL exhibits lower perplexity and faster training speeds. It manages to maintain constant computational speed across varying sequence lengths due to Lightning Attention, significantly outperforming previous models in terms of scalability.

Implications and Future Directions

This work's implications are substantial within the context of AI's expanding frontier in LLMing. By overcoming traditional constraints of linear attention models, Lightning Attention paves the way for more resource-efficient models capable of handling lengthy sequences without a linear deterioration in performance. The authors' integration strategy with the TNL model has potential applications in enhancing the efficiency of numerous natural language processing tasks ranging from real-time translation to complex contextual analysis in diverse AI scenarios.

Future research may explore the adaptability of Lightning Attention and TNL to other domains and further optimization in varying hardware environments. Algorithmic variants that aim to decentralize further computational bottlenecks could yield even more promising results, potentially redefining standard practices in LLM development and deployment.