Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

Authors: william  brandon
mit  csail
mayank  mishra
aniruddha  nrusimha
jonathan  ragan-kelley
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:25:32

Abstract: Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another $2\times$ while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, potentially enabling future models to operate at longer sequence lengths and larger batch sizes than would otherwise be possible.

Main paper: 
 introduction: 
 The memory footprint of the key-value (KV) cache can be a bottleneck when serving large language models (LLMs). Because the size of the KV cache scales proportionally with both sequence length and batch size, the memory overhead of KV cache storage can limit batch sizes when operating on long sequence lengths [Chowdhery et al., 2022] , and can require employing costly techniques like offloading when on-device memory is scarce [Sheng et al., 2023] . It is also desirable to be able to persist KV caches over long periods of time in order to minimize redundant computations [Gao et al., 2024 , Google, 2024] . However, the size of the KV cache directly determines the cost of storing and retrieving such persistent caches. As new applications of LLMs emerge which demand ever-longer sequence lengths, the memory footprint of the KV cache is becoming an increasingly important consideration in the design of efficient transformer-based language models.Existing work has proposed a variety of methods for decreasing the memory footprint of the KV cache, including storing KV activations in low precision [Hooper et al., 2024 , Zhang et al., 2024b] , KV Cache Bytes Per Token (16-Bit Precision) Wikitext PerplexityFigure 1 : Accuracy/memory tradeoffs achieved by MQA models with CLA (red) and without CLA (blue) at the 1B-parameter and 3B-parameter scale, as measured by perplexity on Wikitext. We find that CLA provides the same reduction in KV cache size as shrinking the head dimension d head by 2× while achieving substantially lower perplexities. More details on these experiments are presented in sections 3.2.2 and 3.3. evicting unimportant KV cache entries [Zhang et al., 2023 , Liu et al., 2023] , and sharing keys and values across query heads in the attention mechanism [Shazeer, 2019 , Ainslie et al., 2023] .In this paper, we introduce a method for reducing the size of the KV cache along a dimension different than those explored in prior work: namely, reducing the number of unique layers in the KV cache. Our contributions are as follows:1. We propose Cross-Layer Attention (CLA), a modification to the transformer architecture which reduces the size of the KV cache by sharing KV activations across layers. 2. We conduct extensive pretraining experiments to characterize the effect of different CLA configurations on accuracy and memory usage across a range of architectural hyperparameters, learning rates and model sizes. 3. We demonstrate that CLA enables accuracy/memory Pareto improvements relative to existing Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) architectures. 4. In particular, we demonstrate at the 1B-and 3B-parameter scales that combining CLA with MQA can achieve a 2× reduction in KV cache size versus a plain MQA baseline, with minimal degradation in perplexity. 5. We offer guidance on which CLA configurations perform best based on our experiments, finding that CLA should be used between pairs of consecutive layers, and that CLA appears to deliver the most robust benefits when used in conjunction with MQA. 
cross-layer attention: 
 In this section we describe our Cross-Layer Attention (CLA) technique, and its relationship to the KV-sharing mechanisms employed by the existing Multi-Query and Grouped-Query attention architectures (MQA and GQA).2.1 Background: Multi-Query Attention and Grouped-Query AttentionThe original transformer architecture employed Multi-Head Attention (MHA) [Vaswani et al., 2017] , in which each query head attends over the keys and values produced by a distinct key/value head. In Figure 2 : Schematic of two consecutive layers in a transformer using a traditional attention design (left) and in a transformer using Cross-Layer Attention (right). When using traditional attention, each layer computes its own separate K and V activations, which must be cached on a per-layer basis during autoregressive decoding. When using Cross-Layer Attention, some layers compute their own fresh K and V activations, while other layers reuse the K and V activations of earlier layers.MHA, the KV activations of each key/value head must be stored separately in the KV cache, resulting in a storage overhead of 2 • n query • d head elements per token, where n query is the number of query heads and d head is the embedding dimension of each head.To reduce the overhead associated with storing and accessing the KV cache during transformer decoding, Shazeer [2019] proposed Multi-Query Attention (MQA), which Ainslie et al. later generalized to Grouped-Query Attention (GQA). Grouped-Query Attention modifies the transformer architecture by organizing the query heads of each attention layer into groups, where each group of query heads shares a single key/value head. Because the size of the KV cache scales only with the number of distinct key/value heads, not the number of query heads, GQA reduces the storage overhead of the KV cache to 2 • n group • d head , where n group denotes the number of groups for GQA and n group < n query . MQA can be seen as the special case of GQA in which n group = 1.Shazeer and Ainslie et al. find that MQA and GQA enable significant reductions in KV cache size and decoding latency while incurring only a small degradation in accuracy compared to MHA architectures with the same head dimension. The family of attention architectures enabled by using MQA and GQA defines an accuracy/memory tradeoff space in which model designers can choose how they want to balance the expressive power and KV cache overhead of their attention mechanism. MQA and GQA stake out different positions in this tradeoff space, and neither is necessarily preferable to the other for all use cases.Inspired by the success of MQA and GQA, which share key/value heads across query heads within a single layer, we propose also sharing key/value heads across layers. We refer to such an attention architecture as Cross-Layer Attention (CLA), and present a diagrammatic view of it in Figure 2 . CLA computes key/value projections for only a subset of layers in the model; the attention blocks in layers without key/value projections reuse the KV activations of previous layers. Only the subset of layers with key/value projections contribute to the KV cache, allowing a reduction in memory footprint relative to traditional architectures which apply a separate key/value projection in each layer.CLA is orthogonal to MQA/GQA/MHA, and can be combined with any of them. Moreover, in the same way that GQA allows varying n group to access a family of different attention configurations, CLA allows varying the number of layers which share the output of each KV projection, which we refer to as the sharing factor. We refer to different configurations of CLA by their sharing factors, giving rise to CLA2, which shares each KV projection among a pair of adjacent layers, CLA3, which shares each KV projection among a group of 3 layers, and so on. In Appendix A we include a figure illustrating a few of the different attention configurations possible with CLA.CLA is primarily an intervention to reduce the memory footprint of the KV cache, and only has minor effects on other resources consumed by the model during training and inference. Here, we summarize the effect of CLA on key metrics relevant from a systems engineering perspective, assuming all other architectural hyperparameters are held constant:• KV Cache Memory: CLA significantly reduces KV cache memory footprint, shrinking it by a factor equal to the sharing factor, or slightly less if the sharing factor does not evenly divide the number of layers.• Training Memory Footprint: CLA reduces the memory footprint of intermediate KV activation tensors materialized during training, although for GQA and MQA models such KV tensors are typically small compared to the model's hidden states and MLP activations.• Model Parallelism: CLA is fully compatible with standard tensor parallelism techniques [Shoeybi et al., 2020] for sharding model weights across multiple accelerators. In the presence of pipeline parallelism [Huang et al., 2019] , either different layers which share a KV cache must be kept in the same pipeline stage, or else KV activations must be communicated between pipeline stages.• Parameters and FLOPs: Because CLA reduces the total number of key/value projection blocks in the model, CLA slightly reduces the number of parameters in the model and the number of FLOPs required during a forward or backward pass.• Decoding Latency: In the context of a full LLM serving stack, CLA can enable larger batch sizes and longer KV cache persistence times than would otherwise be possible, which have the potential to improve inference latency.• Core Attention Latency: Unlike MQA and GQA, CLA has no direct effect on the memory bandwidth consumed by the attention mechanism in each decoding step, because even shared KV cache layers must be separately re-read from main memory in each attention layer. CLA therefore has no direct effect on the latency of the core attention computation during decoding.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss the limitations of our technique and experiments in two ways. First, in our "Implications for System Design" section (Section 2.3) we lay out the different aspects of LLM efficiency to which we do and do not expect CLA to be beneficial. Second, in our "Limitations & Future Work" section (Section 5), we highlight the fact that our experiments are only concerned with directly measuring accuracy/memory tradeoffs, and don't attempt to measure the impact of reducing KV cache size on improving end-to-end efficiency numbers (e.g. by enabling larger batch sizes or longer KV cache persistence times). As reducing end-to-end operating costs is ultimately what motivates the desire to reduce KV cache memory in the first place, we believe evaluating the system-level impact of CLA in a realistic, high-performance LLM inference stack is an important problem for future work. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
pretraining experiments: 
 To determine the effect of Cross-Layer Attention on language modeling accuracy, we trained a collection of transformer-based language models from scratch at the 1 billion and 3 billion parameter scales. While running these experiments, we sought to answer the following questions:1. What accuracy/memory tradeoffs are possible using CLA?2. How does using CLA compare to using plain GQA or MQA? Pareto Frontier with and without CLA (1B Models)Figure 3 : The accuracy/memory Pareto frontier discovered in our 1B-scale design space exploration, for models with CLA (red) and without CLA (blue). Lower is better on both axes.We found that CLA enables favorable accuracy/memory tradeoffs compared to what is possible using plain GQA or MQA. Moreover, we found that in our experimental regime, a sharing factor of 2 is more effective than other sharing factors, and that CLA is consistently effective when combined with MQA when trying to decrease KV cache storage. Finally, we found that CLA confers benefits at both 1B-and 3B-parameter scales. In the rest of this section, we present our experimental setup and results in more detail.In all our experiments, we train our models from scratch on data from the SlimPajama [Soboleva et al., 2023] dataset, tokenized with the GPT-NeoX tokenizer [Black et al., 2022] which uses Byte-Pair Encoding (BPE) [Wang et al., 2019] . We adopt a Llama-like [Touvron et al., 2023] with pre-normalization [Xiong et al., 2020] , SwiGLU activations [Shazeer, 2020 , Ramachandran et al., 2017] , and rotary position embeddings [Su et al., 2023] . We do not use dropout for any of our models. Our models use RMSNorm [Zhang and Sennrich, 2019] with learnable elementwise multiplication parameters, and our CLA models use separately-learnable RMSNorm parameters for the KV projection blocks and Q projection blocks in attention. Unless otherwise stated, we always set the number of query heads n query such that n query • d head is equal to the hidden size d model .We train all models using the AdamW optimizer [Loshchilov and Hutter, 2019] with gradient clipping, using β 1 = 0.9, β 2 = 0.95, a weight decay factor of 0.1, and a clipping norm of 1.0. We use a linear learning rate warmup for the first 5% of training examples and a cosine learning rate schedule Loshchilov and Hutter [2017] decaying to 10% of the peak learning rate over the remainder of training. We set the sequence length to 2048 tokens and the batch size to 2048 sequences, for a total of ≈ 4M tokens per training step. All our experiments initialize the weights of linear layers from a normal distribution with mean zero and standard deviation 0.01275.We perform all experiments on NVIDIA H100 GPUs using PyTorch [Paszke et al., 2019 , Ansel et al., 2024] . We use mixed precision training [Micikevicius et al., 2018] in BF16 [Kalamkar et al., 2019] with gradient all-reduce and gradient accumulation in FP32 for training stability.We trained all our 1B-scale models on 30 billion tokens using a consistent data order, and, other than varying the attention mechanism, used the same architectural hyperparameters across all 1B-scale models. This means that all our 1B models were all trained using approximately the same number of FLOPs and approximately the same number of GPU-hours, with CLA models requiring slightly fewer FLOPs to train than their non-CLA counterparts due to the reduced number of key/value projections. The common hyperparameters shared across our 1B-scale experiments can be found in Table 2 .We ran two main sets of experiments at the 1B-parameter scale. First, we trained a diverse set of CLA and non-CLA models to characterize the range of accuracy/memory tradeoffs achievable with and without CLA, and to determine which CLA configurations are most effective; we refer to these as our design space exploration experiments, and describe them in more detail in Section 3.2.1. Second, we conducted a learning rate sweep on a subset of models from our design space exploration to verify that our results continue to hold even against a strong non-CLA baseline with a well-tuned learning rate. We describe these learning rate tuning experiments in Section 3.2.2.The primary goal of our 1B-parameter-scale design space exploration was to characterize the Pareto frontier of accuracy/memory tradeoffs achievable with and without CLA, and to determine which CLA configurations achieve the best accuracy on a fixed KV cache memory budget. We train all models in our design space exploration using a learning rate of LR = 3 × 10 -4 , which we determined to be conservative; we explore the effect of the learning rate on accuracy in more detail in Section 3.2.2.For our design space exploration, we first trained a collection of seven non-CLA baseline models along the MHA-GQA-MQA spectrum, exhibiting a range of KV cache memory requirements spanning two orders of magnitude. Our baseline model with the largest KV cache memory footprint is an MHA model with a head embedding dimension of d head = 128 (163840 bytes per token at 16-bit precision), and our baseline with the smallest footprint is an MQA model with head dimension d head = 32 (2560 bytes per token).We quantify the accuracy of models in our design space exploration using perplexity on a held-out validation set of ≈ 4M tokens drawn from our SlimPajama corpus. A summary of results for the models in our design space exploration, including our baseline models, can be found in Table 1 . We adopt the naming scheme "H⟨d head ⟩-⟨attention mechanism⟩" for all models in our experiments, so that, for example, a model employing MQA with a head dimension of d head = 64 would be named "H64-MQA." For our baseline models, we observe that validation perplexity increases monotonically as we reduce the memory capacity of the KV cache, ranging from a perplexity of 13.15 for our H128-MHA baseline to 14.37 for our H32-MQA baseline.In the rest of this section, we present results for the CLA models we trained during our design space exploration.Best Performance: MQA + CLA2. We trained a family of five models combining MQA with CLA2. We varied the head dimension for our MQA-CLA2 models from d head = 512 down to d head = 64, allowing us to compare to a range of non-CLA baseline models with varying KV cache capacities.We found that our MQA-CLA2 models are able to achieve better perplexities than baseline models requiring the same amount of KV cache memory, advancing the accuracy/memory Pareto frontier. We present a plot of the accuracy/memory Pareto frontier with and without CLA in Figure 3 . Our MQA-CLA2 models with head dimensions d head ∈ {64, 90, 128} are able to match the KV cache memory footprint of baseline MQA models with head dimensions d head ∈ {32, 46, 64} while achieving substantial perplexity improvements in the range of 0.21-0.48 points. Additionally, our MQA-CLA2 models with large head sizes of d head ∈ {256, 512} are able to match the KV cache footprint of our MQA and GQA2 baselines with d head = 128 while achieving a small perplexity improvement of 0.03 points.Ablations. We found that our MQA-CLA2 models achieved the best accuracy/memory tradeoffs among all CLA configurations we tested in our design space exploration. Here, we briefly describe the ablations we conducted to explore alternate CLA configurations. We present our ablations in more detail in Appendix B.We explored combining CLA with GQA4 and GQA2, and found that all our experiments combining CLA with GQA either matched the accuracy/memory tradeoffs of our MQA-CLA2 models, or underperformed our non-CLA baselines. We trained models using CLA3 and CLA4, and found that they achieved accuracy/memory Pareto improvements over our non-CLA baselines, but slightly underperformed our MQA-CLA2 models at the same memory budget. Finally, we explored more complex and irregular sharing patterns than sharing a single KV cache across each pair of adjacent layers, but found no benefit over our basic CLA2 configuration from any of the patterns we tested.The relative performance of different model architectures can change depending on the learning rates at which they are evaluated. To account for the effects of the learning rate on our results, we conducted learning rate tuning experiments on three models of interest from our initial 1B-scale design space exploration. These learning rate tuning experiments help us verify that CLA continues to provide benefits even when compared to baselines trained at their optimal learning rates.We chose to tune the learning rate for the baseline models H128-MQA and H64-MQA, as well as the CLA model H128-MQA-CLA2. In our initial design space exploration, our results for these models indicated that CLA makes it possible to shrink the KV cache footprint of an MQA model with d head = 128 by a factor of 2× while incurring only a small (0.06 point) degradation in perplexity, or to create a model with the same KV cache footprint as an MQA model with d head = 64 while enjoying a substantial (0.21 point) improvement in perplexity. We wanted to verify that this qualitative pattern continues to hold when all models are trained with well-tuned learning rates.Learning Rate Tuning Strategy. For each of our three model configurations, we swept the learning rate upwards from an initial value of 3 × 10 -4 in multiplicative increments of 1.5×. We ended our sweep for each model at the point where validation perplexity stopped improving.Results. We found an optimal learning rate of LR = 1.5 × 10 -3 for our H128-MQA baseline, and a higher optimal learning rate of LR = 2.25 × 10 -3 for both our H64-MQA baseline and our H128-MQA-CLA2 model.The results of our 1B-scale learning rate tuning experiments can be found in Table 3 . When comparing all three models at their best learning rates, we found that the qualitative result from our design space exploration continues to hold: our CLA2 model incurs only a small (0.04 point) validation perplexity degradation relative to our d head = 128 baseline while enjoying a 2× smaller KV cache footprint, and achieves a substantial (0.31 point) validation perplexity improvement compared to our d head = 64 baseline while using the same amount of KV cache memory.To further validate our results, we also evaluate our three learning-rate-tuned 1B-scale models under EleutherAI's LM Eval Harness [Gao et al., 2023] on Wikitext [Merity et al., 2017] perplexity and seven standard downstream benchmarks. On Wikitext perplexity, we observe a similar pattern as with validation perplexity. On the downstream evaluations, we found that none of our three models model consistently wins or loses across different benchmarks, and that all three models are consistently within 1-5 percentage points of each other. 3 : Results of our learning rate tuning experiments at 1B scale. The columns "WG" and "OBQA" denote "WinoGrande" and "OpenBookQA", respectively.To determine how CLA performs when applied to larger models, we trained a collection of models at the 3B-parameter scale both with and without CLA. We trained each of our 3B-scale models from scratch on approximately 100B tokens from our SlimPajama corpus. The common architectural hyperparameters for our 3B-scale models can be found in Table 2 . Initial 3B-Scale Experiments. We initially tried training a 3B-parameter MQA model with d head = 128 as one of the baselines for our experiments. To our surprise, we found that an MQA model with d head = 64 outperformed this d head = 128 baseline in perplexity despite using only 1 /2 as much KV cache memory. We also found that an MQA-CLA2 model with d head = 128 was able to outperform both plain MQA baselines given a well-tuned learning rate. After seeing these results, we chose to use an MQA model with d head = 64 as our main large-KV-cache baseline for further 3B-scale experiments, in order to ensure we were comparing to the strongest baseline we could identify. We provide more detail on these initial 3B-scale experiments in Appendix E. Main 3B-Scale Experiments. In our main 3B-scale experiments, we chose to compare an MQA-CLA2 model with d head = 64 to a plain MQA model with d head = 64, and to a plain MQA model with d head = 32. We trained all models for our main experiments with a learning rate of LR = 10 -3 , which we found to be optimal for our d head = 64 MQA baseline model in our initial experiments. For our d head = 64 MQA-CLA2 model and our d head = 32 MQA baseline model, we also experimented with learning rates of LR ∈ {6.75 × 10 -4 , 1.5 × 10 -3 }, but found these achieved worse perplexities than our initial value of LR = 10 -3 .We report perplexity and downstream benchmark results for our main 3B experiments in Table 9 . In the Wikitext perplexity results for this set of experiments, we find agreement with the pattern observed at the 1B scale. Our MQA-CLA2 model with d head = 64 incurs only a small (0.05 point) degradation in perplexity compared to our d head = 64 baseline while enjoying a 2× smaller KV cache footprint, and achieves a substantial (0.35 point) improvement in perplexity compared to our d head = 32 baseline while using the same amount of KV cache memory.We also evaluate these three models on downstream benchmarks, and report the results in Table 4 . As with our downstream benchmark evaluations at 1B-parameter scale, we find that all models perform similarly as measured by these downstream evaluations.We wanted to verify that models trained using CLA remain competitive when compared with external models not trained using our particular software stack. To verify this, we ran an experiment which compares directly against the open-source GQA4 model TinyLlama-1.1B [Zhang et al., 2024a] at its 105B-token intermediate checkpoint. For this experiment, we pretrained our own version of TinyLlama-1.1B-105B from scratch using CLA2, using otherwise-identical training data, model architecture, and hyperparameters as described in the TinyLlama repository.In this comparison, we found that our TinyLlama-1.1B-105B checkpoint trained with CLA2 matches or exceeds the performance of the original, publicly-available TinyLlama-1.1B-105B checkpoint. Due to space constraints, we present the full results of this comparison in Appendix G.4 Adaptation Experiments Ainslie et al. [2023] propose a recipe for converting MHA models into GQA models via a small amount of additional training, which they refer to as "uptraining." Inspired by Ainslie et al., we conducted experiments to investigate if models pretrained without CLA can be adapted to use CLA.Methods. At the 1B-and 3B-parameter scales, we trained two "base" models on 105B tokens each, both using plain MQA. We then used the final weights of each base model to initialize an "adapted" model employing CLA2, which we trained on a further 21B tokens. More details on the hyperparameters used for these experiments can be found in Appendix H.1.To initialize the KV projection weights for each group of layers sharing a KV cache in the adapted model, we considered two strategies: copying the base-model KV projection weights for the first layer in the group, and mean-pooling the base-model KV projection weights for all layers in the group. Similarly to Ainslie et al. [2023] , we found in early experiments that mean-pooling performed best, and chose to employ mean-pooling in all further experiments. To account for the fact that different base-model layers have different RMSNorm parameters, we absorb the elementwise multiplication for each RMSNorm into the weights of its associated KV projection prior to mean-pooling, and initialize the RMSNorm multiplication parameters for all KV projections in the adapted model to 1.Results. We found that our adapted CLA2 models quickly converged to low loss values, but did not fully recover the performance of their original base models after 21B tokens of adaptation training. We present the main results of our adaptation experiments in Table 5 , and present training loss curves for the models in our adaptation experiments in Appendix H.2. We believe it is likely possible to improve upon our CLA adaptation recipe, and leave doing so as a promising direction for future work. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: Our paper does not include theoretical results.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experimental result reproducibility: 
 Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We believe we have described our experimental setup in enough detail in the "Pretraining Experiments" section (Section 3) that it would be possible for a reader to independently replicate our results given only the information in the paper. The SlimPajama dataset we use is available to the general public, and our method can be straightforwardly implemented in PyTorch without specialized software support. 
limitations & future work: 
 For full-scale LLM serving systems, we expect KV cache reduction techniques like CLA to enable longer sequences, larger batch sizes, and greater KV cache persistence times than would otherwise be possible or economical. We do not directly evaluate these system-level consequences of reducing KV cache memory overhead, and instead leave end-to-end inference efficiency evaluations of large, long-context models employing CLA as an interesting problem for future work.6 Related Work KV Cache Compression. Many works have tried to compress LLMs through pruning, quantization, and sparsity (see Zhu et al. [2024] for a survey), and a subset of these focus on KV cache compression. KVQuant [Hooper et al., 2024] and Coupled Quantization [Zhang et al., 2024b] employ targeted transformation of the keys and values, along with non-uniform numeric encodings, to compress the KV cache to 1-bit or 2-bit precision. Work on sparsifying the KV cache includes H2O [Zhang et al., 2023] , Scissorhands [Liu et al., 2023] , and FastGen [Ge et al., 2024] , all of which employ heuristics to select unimportant tokens to discard. PagedAttention [Kwon et al., 2023] decreases the memory footprint of the KV cache in batched inference settings by eliminating unnecessary padding and allowing for sharing of common prefixes. Finally, Cachegen [Liu et al., 2024] directly losslessly compresses the KV cache with a custom-tailored arithmetic-coding-based compression scheme.KV-Cache-Efficient Transformer Variants. Significant prior work has proposed methods for modifying the transformer architecture in ways which reduce the size of the KV cache. Two notable early works in this direction are Transformer XL [Dai et al., 2019] and Sparse Attention [Child et al., 2019] which reduce effective sequence lengths by performing attention only within a local window. More recently, this line of work has been extended by methods like Infini-attention [Munkhdalai et al., 2024] , which maintain fixed-sized compressive memories of tokens seen prior to a local window. Finally, MQA and GQA [Shazeer, 2019 , Ainslie et al., 2023] enable substantial reductions in KV cache size, and serve as the direct inspiration for our technique.Removing Softmax Attention. A large body of recent work has proposed efficient alternatives to softmax attention for performing sequence-mixing in LLMs; notable ongoing work in this direction includes variants of linear attention [Katharopoulos et al., 2020 , Wang et al., 2020 , Yang et al., 2024] , state space models (SSMs) such as Mamba [Gu and Dao, 2024] , and the RWKV v6 architecture [Peng et al., 2024] . While these non-transformer architectures are often much more memory-efficient at inference time than a transformer employing a KV cache would be, characterizing their exact capabilities and limitations relative to transformer models remains an open area of research. 
conclusion: 
 In this work, we introduce Cross-Layer Attention (CLA), a new attention architecture for transformerbased language models which reduces the size of the KV cache by sharing KV activations across layers. Through extensive experiments, we demonstrate that CLA allows more favorable accuracy/memory tradeoffs than are possible with MQA and GQA alone, thereby advancing the accuracy/memory Pareto frontier. In particular, we show that when combined with MQA with typical head dimensions, CLA allows 2× reductions in KV cache size with minimal impact on accuracy, resulting in a simple and effective recipe for practitioners to improve the memory efficiency of their transformer models. Using traditional attention, each layer has its own KV cache. Using Cross-Layer Attention with a sharing factor of 2 (CLA2), every group of 2 consecutive layers shares a single KV cache. Using Cross-Layer Attention with a sharing factor of 3 (CLA3), every group of 3 consecutive layers shares a single KV cache. When the sharing factor does not evenly divide the number of layers, as in the CLA3 example, some KV caches must be shared over fewer layers than others; in this CLA3 configuration, we arbitrarily select the layer 0 KV cache to be used only in layer 0. We conducted a number of ablations as part of our design space exploration at the 1B-parameter scale. The full results of our design space exploration, including ablations, are presented in Table 6 . In this appendix, we describe the configurations and results for those ablations in more detail.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: Because of the high cost of training multi-billion-parameter language models from scratch, it was not possible for us to train more than one version of each model in our experiments, preventing us from computing confidence intervals over multiple runs. The standard EleutherAI LM Eval Harness which we use to obtain our benchmark numbers does not compute or report confidence intervals.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We describe the hardware we use, the exact architectural hyperparameters of the models we train, and the exact number of tokens we train them for, from which the reader can easily predict the compute requirements they would need to reproduce our experiments using their preferred training stack. In our specific setup, training one 1B-scale model requires approximately 8 GPU-days (which we parallelize over 8 GPUs, for a wallclock time of about 24 hours), and training one 3B-scale model requires approximately 80 GPU-days (which we parallelize over 32 GPUs, for a wall-clock time of about 60 hours).Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We have reviewed and are in compliance with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work constitutes foundational, application-agnostic research into efficient deep learning architectures for language modeling. Language models have the potential to cause to a wide variety of beneficial and harmful societal consequences, none of which we feel need to be addressed within the scope of this work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We release no artifacts with any direct risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Ablation: GQA + CLA2. We trained three models to explore combining GQA with CLA2. We chose GQA4-CLA2 with d head = 128 as our starting point, as GQA4 represents an attention configuration intermediate between our MQA and MHA baselines. We then explored expanding the head dimension of our GQA4-CLA2 model to d head = 256, as well as reducing the GQA factor to GQA2. We found that only the GQA2-CLA2 configuration was able to achieve a perplexity better than the corresponding baseline model with the same KV cache footprint, and that this perplexity was the same (within 0.01 points) as our MQA-CLA2 model with the same footprint.Ablation: MQA + CLA with Sharing Factor > 2. To explore the effect of using CLA sharing factors > 2, we trained MQA-CLA3 and MQA-CLA4 models with head dimension d head = 128. We found that these CLA3 and CLA4 models achieved a Pareto improvement over our plain MQA baselines, matching the KV cache footprint of our baseline MQA models with head dimensions of d head ∈ {32, 46} while achieving better perplexities. However, we found that they achieved worse perplexities than our MQA-CLA2 models at the same KV cache footprint.Ablation: MQA + CLA2 with Non-Uniform Sharing Patterns. Finally, we explored using different patterns of KV activation sharing in our MQA-CLA2 models.On the hypothesis that the first and last layers in the model might benefit from special treatment, we trained a model "H128-MQA-CLA2-KeepEnds" which does not share the layer 0 KV cache with any other layers, and instead groups layer 1 with layer 2, groups layer 3 with layer 4, and so on. This also has the effect of giving the final layer its own KV cache separate from all other layers.We also explored imbalanced configurations with all the KV-cache-producing layers concentrated at either the beginning or end of the model. We trained a model "H128-MQA-CLA2-DenseFront" consisting of 10 non-CLA layers, followed by 9 CLA layers all using the KV activations of layer 9, and a final layer with its own KV cache. Similarly, we trained a model "H128-MQA-CLA2consisting of 2 non-CLA layers, followed by a run of 10 CLA layers all using the KV activations of layer 1, and finally 9 non-CLA layers.We found that all of these alternative CLA sharing patterns achieve worse perplexities than the corresponding MQA-CLA2 model with a uniform sharing pattern, while also requiring slightly more KV cache memory. Experiments at Head Dimension d head = 128. We initially ran experiments to compare three 3B-scale models analogous to the models we selected for our learning rate tuning experiments at the 1B-parameter scale. Specifically, we compared a model using MQA-CLA2 and d head = 128 to an MQA model with the same head dimension (and hence 2× the KV cache footprint), and to an MQA model with a head dimension of d head = 64 (and hence the same KV cache footprint). Based on our 1B-scale experiments, we expected that our MQA-CLA2 and MQA models with d head = 128 would achieve similar perplexities to each other, and that both would outperform the d head = 64 model.We tuned the learning rates for these models according to the same learning rate tuning protocol we used at the 1B-parameter scale. After tuning the learning rates for each model, we observed a result different than we had expected: at 3B scale, our MQA-CLA2 model achieves substantially better perplexities than both our d head = 128 and d head = 64 MQA baselines. Moreover, our d head = 64 MQA baseline model achieves better perplexities than our tuned d head = 128 MQA baseline, despite having only 1 /2 as much KV cache capacity. We report the optimal learning rates and perplexities for these three models in Table 8 , and present the results of our learning rate sweep graphically in Figure 6 .As with our 1B-scale learning rate tuning experiments, we evaluate these models on downstream benchmarks. We report the results of these evaluations in Table 9 . As with our 1B-scale experiments, we do not find that any model consistently wins or loses in these downstream evaluations.Due to logistical constraints, we used a different training cluster, software stack, and data order for our initial 3B experiments than for the main 3B experiments we describe in Section 3.3. To control for differences in training environment, our main 3B results reported in Section 3.3 use a separately-trained version of our H64-MQA-CLA2 baseline, distinct from the H64-MQA-CLA2 model described in this appendix. Dashed lines indicate the best loss achieved by each model family. We find that our MQA-CLA2 model benefits from a higher learning rate than our d head = 64 MQA model, which benefits from a higher learning rate than our d head = 128 MQA model. At their best learning rates, we find that our MQA-CLA2 model outperforms our d head = 64 MQA model, which outperforms our d head = 128 model. We tried training our MQA-CLA2 model at LR = 3.375 × 10 -3 -one LR increment higher than the rightmost point depicted in this plot -but found that training diverged at that LR scale.Here we present visualizations of the training loss curves for the models in our main experiments at the 1B-and 3B-parameter scales. The data in each plot has been smoothed using an exponential moving average, with the same smoothing factor used for all plots. Here, we present the training loss curves for the models in our 1B-and 3B-scale adaptation experiments, described in Section 4. In each experiment, we first pretrained an MQA "base" model (indicated in blue) on 105B tokens, and then used that model's final weights to initialize an MQA-CLA2 "adapted" model (indicated in red) trained on a further 21B tokens. 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The main claims made in the abstract and introduction are that Cross-Layer Attention (CLA) advances the accuracy/memory Pareto frontier relative to what is possible with traditional MQA/GQA models, and that CLA allows reducing the KV cache memory footprint of MQA models by 2× while preserving accuracy. We directly empirically demonstrate both of these claims in the "Pretraining Experiments" section, using perplexity as our operationalization of accuracy.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example• The full details can be provided either with the code, in appendix, or as supplemental material. 


Decision: The paper introduces a novel yet simple method for KV cache compression, one of the most challenging topics in today's deployment of transformer-like architectures. All reviews were positive, and the authors engaged in a fruitful discussion during rebuttal.
I think this result will benefit NeurIPS community.

Reviews: review: paper_summary: paper_summary: This paper introduces a new KV cache compression technique by sharing KV cache cross-layers called CLA. It can be integrated with most of the existing KV cache compression technique like MQA, GQA, and quantization. When applying over GQA, CLA can further reduces KV cache size by 2× while maintaining similar accuracy. Experiments with 1B- and 3B-parameter models show that CLA offers better memory/accuracy trade-offs, allowing models to handle longer sequences and larger batch sizes.
strength_weakness: strength_weakness: I personally like this paper and vote for acceptance. However, I think the experimental design could be improved to make it more solid:

- Only Wikitext Perplexity is reported, which is a weak indicator of LLM ability. While I understand that LLM evaluation is challenging and time-consuming, I suggest the authors also report the MMLU score, even if only on a subset.

- Can CLA work under continuous pretraining scenarios like GQA? Namely, can a pretrained LLM be continuously tuned into a CLA-based model, even though the compression ratio might not be as high in this case? I ask this because pretraining LLMs is extremely expensive, as seen with models like Llama 3.1. It is risky and costly to train another open-sourced CLA model from scratch. If we can adapt existing models into CLA-based models, this paper could have a much greater impact.

- One of the main motivations for KV Cache compression is long-context inference. However, long-context tasks like few-shot learning or multi-doc QA are much harder than the "short" tasks. I understand that long-context evaluation is challenging for pretrained models due to lack of alignment. Still, I think the authors could show some passkey retrieval results, as this task does not require much alignments.
questions: questions: Check Weaknesses point 2
limitations: limitations: No major limitations.
paper_summary: paper_summary: In this paper, the authors proposed an approach, Cross-Layer Attention (CLA), to accelerate the autoregressive generation procedure of Large Language Models. Going beyond multi-query attention, the main idea of CLA is to share the Key-Value caches among attention layers. Intuitively, this idea is straightforward to further reduce the size of the KV cache that must be stored during generation. The authors conducted several experiments covering different model sizes and combinations between CLA and GQA/MQA to verify the effectiveness of the CLA module.
strength_weakness: strength_weakness: 1. **The comprehensiveness of experiments can be further improved.** 
    - **Model sizes**: the current version of this paper only verifies the effectiveness of CLA on 1B and 3B Transformer language models. Although the experimental results demonstrate that CLA does not bring performance drop on real-world tasks for these two model sizes, it is also questionable whether modern LLMs with  >10 and > 100 billion parameters could still use this approach to achieve better accuracy-efficiency. After all, there exist many approaches that are only effective on small-scale architectures in the literature. It is also understandable that the computational resources are restrictions for conducting such verification experiments, but I think it would be a great plus for improving the quality of this work.
    - **Design choices**: in Table 1 and Appendix B, the authors listed a bunch of Non-CLA baselines and MQA/GQA+CLA2/CLA3/CLA4 models. However, CLA-only models are not compared. It would be questionable whether CLA can only be used with MQA/GQA together.
questions: questions: 1. In lines 139-141, could you explain why using separately-learnable affine layer-norm parameters for the KV projection blocks and Q blocks in attention?
2. Similar to Sec 3.2.2, how is the robustness of CLA against other hyper-parameters like batch size?
limitations: limitations: The authors discussed limitations of this work in Sec 4.
paper_summary: paper_summary: To resolve the costly memory consumption of KV cache in transformer-based LLMs, this paper proposes Cross-Layer Attention (CLA). In contrast to multi-query attention (MQA) and grouped-query attention (GQA) which share KV cache across attention heads, CLA shares KV cache across contiguous layers. Pre-training experiments prove the effectiveness of CLA over MQA and GQA.
strength_weakness: strength_weakness: 1. Experimental results are not presented fully. Figure 3 only shows the validation perplexity for a subset of models, and Table 3 and 4 only presents the accuracy for a subset of models. I couldn't find the results elsewhere in the paper or the appendix. Presenting the results fully contributes to the holistic understanding of the effectiveness of the proposed approach.
2. Suggestions to experimental designs. The authors choose to pre-train 1B and 3B models from scratch to prove the effectiveness of CLA. It would perhaps be more convincing if the authors can choose a fully open source pre-trained model (such as OpenLLaMA) as baseline, and pre-train it with CLA to prove its effectiveness. Pre-training 1B and 3B models using a custom set of hyper-parameters and datasets is less convincing in comparison.
3. Training loss curves are missing. Training loss is important for understanding whether the models have fully converged.
4. Citation formats may be incorrect. Many references on page 10-12 only show author names, title, and year, and omit publication venue. 
5. Some typos. Line 16: "which are possible" -> "which are impossible", line 18: "would otherwise be possible" -> "would otherwise be impossible".
questions: questions: 1. Can the authors please provide some insights on why cross-layer sharing of KV cache may produce better models than GQA? The proposed method would be more convincing if some design justifications are provided.
2. Is it possible to adapt pre-trained models to using CLA through fine-tuning?
limitations: limitations: Yes the authors addressed the limitations. I do not envision additional limitations.

Comments: title: Thank You
Full Benchmarking Results for Models in our Design Space Exploration
comment: Thank you for your responses. Although I still think the potential of this work should be further evaluated by using models of larger sizes, I understand the limitations of computational resources. I choose to raise my rating to 7.
Thank you for your clarification. I am satisfied with the response.

Good Luck!
We are glad that we were able to address the reviewer's questions, and thank them for their response.
Thank you to the authors for the detailed response. My concerns have mostly been addressed. I think the additional clarifications and experimental results are a great addition to the paper. I have adjusted my initial rating accordingly.
Attached please find the full benchmarking results for the models in our 1B-scale design space exploration. (This comment accompanies our rebuttal, which could not fit the attached table due to length constraints.)

|Model|$\uparrow$ hellaswag|$\uparrow$ piqa|$\uparrow$ winogrande|$\uparrow$ sciq|$\uparrow$ openbookqa|$\uparrow$ boolq|$\uparrow$ arc-e|$\downarrow$ wikitext (PPL)|
|-|-|-|-|-|-|-|-|-|
|H128-MHA|33.88|67.19|53.12|81.1|19|61.62|52.15|20.90|
|H128-GQA4|33.82|67.79|51.62|78.6|18.6|60.73|51.47|21.38|
|H128-GQA2|33.34|67.85|53.04|79.6|20|60.89|50.97|21.64|
|H128-MQA|33.53|67.79|51.07|78.4|18.6|56.61|51.35|21.79|
|H64-MQA|33.24|67.52|50.04|75.8|17|59.39|51.22|22.31|
|H46-MQA|32.99|66.7|52.41|77.9|19.2|60.18|49.34|22.59|
|H32-MQA|32.58|67.8|50.99|74.5|18.4|59.94|49.02|23.76|
|H512-MQA-CLA2|33.68|67.68|52.33|77.3|18.8|55.72|52.19|22.42|
|H256-MQA-CLA2|33.9|67.74|49.8|77.4|18.2|60.34|50.04|21.64|
|H128-MQA-CLA2|33.29|67.63|49.88|78.3|18|59.51|49.62|21.82|
|H90-MQA-CLA2|33.15|67.41|51.85|76.7|17.2|59.11|52.06|22.13|
|H64-MQA-CLA2|32.71|67.36|51.7|74.9|19.4|54.68|50.88|22.43|
|H256-GQA4-CLA2|33.63|66.92|51.78|78.5|18.6|60.43|51.18|21.40|
|H128-GQA4-CLA2|33.64|67.74|50.59|78.1|18.6|58.78|51.09|21.66|
|H128-GQA2-CLA2|33.39|67.14|52.17|77.3|19.8|59.45|51.26|21.83|
|H128-MQA-CLA3|32.91|67.74|51.54|76.6|18|54.53|51.18|22.18|
|H128-MQA-CLA4|32.51|67.57|51.85|75.4|18.6|59.33|51.73|22.62|
|H128-MQA-CLA2-KeepEnds|33.58|68.12|52.72|76.2|19.2|60.12|51.64|21.88|
|H128-MQA-CLA2-DenseFront|33.43|67.3|52.57|75.7|19.4|49.14|50.88|22.09|
|H128-MQA-CLA2-DenseBack|32.71|66.65|51.7|76.5|17.4|59.69|50.51|22.8|