Beyond KV Caching: Shared Attention for Efficient LLMs

The self-attention mechanism in transformer models is typically defined as $\text{softmax}(\frac{QK^{T}}{\sqrt{d}})V$, where $Q$, $K$, and $V$ represent the query, key, and value matrices respectively, and $d$ is the dimension of the key vectors. This formulation necessitates the recomputation of attention weights at each layer, a computationally intensive task, particularly when the model is deployed in inference mode. To mitigate this, the concept of a KV-cache is employed, reducing the need to recompute $K$ and $V$ matrices for previously encountered tokens.

While prior methodologies have focused on sharing KV caches at different levels to minimize memory overhead, they predominantly operate under the assumption that attention weights differ significantly across layers, thereby necessitating individual computations to capture diverse contextual dependencies effectively. This assumption prompts a critical inquiry: Are the attention weights indeed markedly different across layers, or is this variation minimal enough to allow for a unified approach across multiple layers?

To explore this, we conducted an empirical analysis on the distribution of attention weights across different layers of the model. Based on the Llama2-7B-chat model, we processed the Massive Multitask Language Understanding (MMLU) dataset (Hendrycks et al. 2020) to extract the attention matrices, $\text{softmax}(\frac{QK^{T}}{\sqrt{d}})$, for each layer. Given the variability in sequence lengths, we standardized these matrices to a uniform size by applying zero-padding to align them to a consistent shape of $\text{maxlen}\times\text{maxlen}$.

Our analysis employed the cosine similarity metric to compare the attention matrices of all layers, revealing a notable high degree of similarity across most of layers, particularly from indices 3 to 30. Contrastingly, the initial layers (0 and 1) and the final output layer (31) exhibited substantially lower similarity scores to middle layers. This observation is intuitive as the early layers are closer to the input token embeddings, requiring frequent adjustments to their attention distribution to accurately abstract semantic meanings from diverse inputs. Similarly, the final layer’s unique role in predicting the next token justifies its distinct attention pattern.

Inspired by these findings, we hypothesize that the high similarity in attention weights across the majority of layers could allow for a shared representation of these weights, thus eliminating the need for separate softmax computations in each layer and reducing the key cache size. Such a strategy could not only streamline the inference process but also enhance computational efficiency significantly.

Based on the observed uniformity in attention weights, we propose a novel algorithm as shown in Algorithm 1, Shared Attention, which utilizes a single shared attention matrix across multiple layers. This approach fundamentally redefines the operational paradigm by maintaining a consistent attention mechanism across various contextual layers, thereby reducing redundancy and enhancing inference speed.

The original self-attention mechanism in Transformers, characterized by the Multi-Head Attention (MHA) model, necessitates caching the keys ($K$) and values ($V$) in each head and layer to accelerate inference (Vaswani et al. 2017). This requirement has historically imposed a significant memory overhead, prompting a series of innovations aimed at reducing this burden.

Among these, Multi-Query Attention (MQA) and its more generalized counterpart, Grouped-Query Attention (GQA), consolidate the KV cache by allowing multiple query heads within the same layer to share a singular set of K and V matrices. This approach effectively reduces the number of unique key and value pairs that must be stored and retrieved during the computation process. Subsequently, Cross-Layer Attention (CLA) extends this concept by facilitating the sharing of K and V matrices across different layers, thereby offering further reductions in the memory footprint required for KV storage.

Our method, however, introduces a fundamentally different paradigm in addressing the challenges of self-attention. While previous methods have focused on reducing the redundancy in storing K and V matrices, our approach centers on the optimization of the computation of attention weights themselves. In standard practice, the cached keys ($K$) are primarily utilized to compute attention weights in conjunction with the queries ($Q$). Instead of indirectly facilitating this interaction through shared KV matrices, our method proposes the direct sharing of the resultant attention weights—specifically, the softmax-normalized scores.

This not only diminishes the memory requirements by obviating the need to store separate sets of keys for each layer but also significantly reduces the computational complexity. By sharing the pre-computed softmax results across layers, our approach circumvents the repeated calculation of softmax, which is often one of the most computationally intensive operations in the attention mechanism. This efficiency gain is reflected in a substantial reduction in the number of floating-point operations (FLOPs) required during model inference, enhancing both the speed and scalability of Transformer deployments.

Unlike traditional methods that optimize memory use by sharing physical keys and values across layers or heads, our Shared Attention model innovates on the computational process itself, exploiting the consistent patterns in attention weights to streamline operations across multiple layers of the Transformer architecture.

To elucidate the formation and evolution of attention weight patterns during the pretraining phase of LLMs, we utilized intermediate checkpoints of the Baichuan 7B model, provided by the model developers. These checkpoints, spanning from 0.2T to 2.6T tokens processed, offer a unique point of view to observe the dynamic shifts in attention mechanisms as the model gains exposure to an increasing volume of data.

We applied a consistent metric for measuring the similarity of attention weights across layers at each pretraining checkpoint. Additionally, the final chat model, fine-tuned to align with human reference responses, was included to benchmark the evolution against practical application outcomes. The dynamics of these attention weights are visualized in Figure 3, which illustrates the progressive differentiation and stabilization of attention patterns across the model’s layers.

As observed in the early pretraining stage at 0.2T tokens, Groups 1 and 2 appear merged, indicating a less differentiated processing strategy across these initial layers. This combination suggests that early in training, the model does not distinctly separate token-level semantic processing from intermediate semantic refinement. However, as the model progresses to 1.0T tokens, a clear division emerges between Groups 1 and 2. This separation aligns with the model beginning to form more specialized and efficient strategies for handling different types of information across its architecture.

The similarity within Group 3, which encompasses the bulk of the model’s layers, shows a marked improvement from a similarity score of 0.8 to 0.9. This increase is indicative of the model’s attention mechanism stabilizing and becoming more consistent in its approach to processing the bulk of contextual information.

The training advancements observed across the pretraining checkpoints not only demonstrate significant shifts in the internal structure of the model’s attention mechanisms but also correlate positively with performance improvements on multiple benchmarks. This includes results on the MMLU, CMMLU (Li et al. 2023), and C-Eval (Huang et al. 2024) 5-shot accuracy tests, which have reportedly improved from a baseline accuracy of 0.25 to 0.50 (Yang et al. 2023). This notable enhancement underscores the intrinsic link between the refinement of attention mechanisms within LLMs and their enhanced capabilities in natural language understanding tasks.

Moreover, further examination of the model’s development, as observed in supplementary material, reveals that the similarity within Group 3—comprising the core contextual processing layers of the model—continues to enhance after the alignment stage. This observation suggests that the alignment process, typically aimed at fine-tuning the model to more closely mirror human-like understanding and response generation, also contributes to the stabilization of the model’s attention mechanisms.

For the fine-tuning experiments, we utilized the Llama2-7B and Llama3-8B base models. These experiments were conducted on a robust hardware configuration consisting of two NVIDIA A100 80GB GPUs. Optimization of the models was carried out using the AdamW optimizer, with an initial learning rate set at $2\times 10^{-5}$. We employed the bf16 datatype for model parameters, which enhances the numeric range and stability during backpropagation, crucial for maintaining precision in large model training.

Each GPU handled a micro-batch size of 16, leveraging gradient accumulation techniques to effectively manage the computational load. Additionally, we utilized DeepSpeed Zero Stage 3 to optimize the distribution of model and optimizer parameters and enhance memory management across the GPUs, ensuring efficient use of available resources. The fine-tuning process spanned two epochs and employed the standard Alpaca instruction format, which is designed to improve the responsiveness and accuracy of the models in handling instruction-based tasks.

The application of SA was tested across discrete segments of layers within the Llama2-7B and Llama3-8B models, each comprising 32 layers in total. To evaluate the robustness and adaptability of SA as shown in Figure 4, it was implemented in varying layer segments, ranging from narrower spans such as four layers (e.g., SA:15$\sim$18) to broader spans such as eight layers (e.g., SA:23$\sim$30).

Preliminary assessments of SA in the earlier layers of Llama2-7B (e.g., layers 3 to 6) resulted in an explosion of perplexity, indicating significant disruptions in the model’s ability to predict subsequent tokens accurately. This phenomenon underscores the crucial role that attention score variances play in the model’s early stages of processing, which are essential for initial context setting and feature extraction. To quantitatively assess the impact of attention variance throughout the model, we conducted a detailed variance analysis. We applied the same computational method used to obtain attention mean scores to calculate the variance of attention weights in Llama2-7B and Llama3-8B while processing the MMLU dataset. We further explored the potential influence of attention variance in downstream layers by computing a weighted cumulative variance. This metric aggregates the variances of all downstream layers starting from each specific layer, weighted by the average of these summed variances. As illustrated in Figure 5, the analysis revealed that early layers exhibited significantly higher weighted variances compared to latter layers. This variance tends to decrease as one progresses through the model’s architecture, suggesting a stabilization of attention mechanisms in the latter layers. Given these results, our experiments predominantly focused on the application of SA in the latter layers, where such variances appear to stabilize.

The outcomes of these experiments, as summarized in Table 1, reveal interesting patterns. For the Llama2-7B model, implementing SA in the latter layers (e.g., SA:23$\sim$26 and SA:27$\sim$30) maintained relatively stable performance across a variety of benchmarks, including GLUE and MMLU. Conversely, extending the scope of SA to encompass more layers, particularly mid-level layers such as SA:15$\sim$18, led to a noticeable degradation in tasks requiring mathematical reasoning (GSM8K).

In comparison, the Llama3-8B model, which inherently showed higher layer-wise attention similarity as discussed in the previous sections, exhibited less performance deterioration when SA was applied. After implementing SA in the layers closer to the model’s output (e.g., SA:27$\sim$30), the Llama3-8B even outperformed its original configuration on the GLUE benchmark, suggesting that strategic placement of SA can potentially enhance the model’s performance in complex natural language understanding tasks.

Given the computational constraints that preclude the pretraining of LLMs with SA from scratch, we adopted to fine-tune existing LLMs to evaluate whether fine-tuning could ameliorate the performance deficits observed with the direct application of SA. This approach was particularly aimed at understanding the adaptability of SA under a more controlled learning regimen.

Fine-tuning was conducted on the publicly available Instruct dataset, which is designed to evaluate models on tasks that require following complex instructions. This dataset was chosen because it challenges the models to utilize their learned representations effectively, making it an ideal benchmark for testing the efficacy of modifications like SA.

The results, as summarized in Table 1, demonstrate a narrowed performance gap between the original models and those modified with SA. For instance, while the original Llama2-7B model outperformed the SA version in direct application tests, the fine-tuned Llama2-7B${}_{\text{SA}:23\sim 30}$ showed significant improvements across multiple metrics. This suggests that fine-tuning enables the model to better integrate and leverage the Shared Attention mechanism, effectively regaining some of the lost performance noted in the initial application of SA.

These findings indicate the potential of fine-tuning as a viable method for integrating new architectural changes like SA into existing models. The recovery in performance indicates that with adequate training, the initial disadvantages of directly applying SA can be mitigated, leading to enhanced model capabilities that more closely align with or even exceed their original configurations.

Our experimental investigations have demonstrated that implementing Shared Attention (SA) across multiple latter layers in LLMs arouses minimal accuracy loss, making it a promising approach for enhancing model efficiency. Furthermore, our analysis reveals a trend towards isotropic attention patterns during the pretraining process, indicating that the models’ attention mechanisms tend to stabilize as they process more data.

Given these insights, integrating SA from the pretraining appears to be a particularly beneficial strategy. This early integration could allow models to better adapt to the streamlined attention mechanism, potentially improving performance and efficiency across various tasks. The foundational embedding of SA might simplify later adaptations and inherently supports efficient attention dynamics.

Another promising research direction involves exploring combinations between SA and other attention-sharing strategies like Cross-Layer Attention (CLA). Combining SA with methods such as CLA could exploit the strengths of both approaches, leading to a more robust and flexible attention mechanism. This holistic approach to attention management could provide a comprehensive solution that maximizes both computational efficiency and model scalability.

By pursuing these avenues, future research can not only refine the application of Shared Attention within LLMs but also explore its full potential in enhancing the architectural and operational efficiency of next-generation language models. These efforts could lead to models that are better equipped to handle the increasing complexity and diversity of tasks in natural language processing.

Significant efforts have been made to address the efficiency of the KV cache post-training. Techniques such as KV cache compression have been explored extensively. For instance, methods like KVQuant (Hooper et al. 2024) and KIVI (Liu et al. 2024b) employ quantization strategies to reduce the memory footprint of KV pairs to just a few bits. Moreover, works such as AttentionSink (Xiao et al. 2023) and Scissorhands (Liu et al. 2024a) introduce sparsity into the KV cache by selectively storing elements based on their proximity or importance to the generation token, thus reducing the overall storage requirements.

Architectural modifications aimed at reducing the KV cache size are pivotal in enhancing the efficiency of large language models. Such strategies include limiting the effective sequence length, as seen in Sparse Attention (Child et al. 2019), which constrain attention to local windows to reduce both computational load and memory overhead. Another approach involves replacing traditional softmax attention with scalable alternatives like linear attention (Katharopoulos et al. 2020), which maintains constant space complexity and offers more graceful scaling with respect to the token count. Additionally, methods such as Grouped-Query Attention (GQA) (Ainslie et al. 2023) and Multi-Query Attention (MQA) (Shazeer 2019) aggregate attention across multiple queries, significantly decreasing the memory footprint by sharing KV pairs across attention heads. These innovations collectively contribute to reducing the redundancy in attention calculations and are directly relevant to our work, informing our development of the Shared Attention mechanism that further optimizes memory usage by sharing attention weights across layers.