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

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.)

We thank the reviewer for their thoughtful and generous feedback. We address each point below:

Could we include more comprehensive results tables in the paper?

We thank the reviewer for pointing this out, and agree the paper could be improved by including more comprehensive results tables.

For validation perplexities, we note that validation perplexities for all model architectures in our design space exploration, including those not shown in Table 1 or Figure 3, are already included in Appendix B Table 5.

For accuracies on downstream tasks, the reviewer is correct that our initial submission only included downstream accuracies for a subset of architectures. Although we include downstream accuracies for selected 1B-scale architectures (see Table 3) and all 3B architectures (see Table 4 and Appendix D Table 7), we do not include downstream accuracy results for all architectures in our 1B-scale design space exploration. We will correct this in the final submission, and include a table (likely in the appendix) containing the missing benchmark results for all architectures not currently documented. We have provided this full table of results in an accompanying comment on OpenReview.

Could we compare to an open-source pre-trained model?

We thank the reviewer for their suggestion. In response to the reviewer's comment, we have run an experiment comparing directly with the open-source GQA4 model TinyLlama-1.1B at its 105B-token intermediate checkpoint. In this experiment, we pretrained our own version of TinyLlama-1.1B-105B from scratch using CLA2, using otherwise-identical training data, model architecture, and training hyperparameters as described in the TinyLlama repository. In particular, we used the same cosine learning rate schedule as TinyLlama 1.1B, which decays over 3T tokens (although we only ran training to 105B tokens).

In this comparison, we find that our TinyLlama-1.1B-105B checkpoint trained with CLA2 matches or exceeds the performance of the original TinyLlama-1.1B-105B checkpoint. We include the benchmark scores for each model below:

At the reviewers' discretion, we would be happy to include these results in the final version of our paper.

Could we include training loss curves?

We would be happy to extend the appendix in the final version of our paper to include visualizations of the training loss curves for all the models in our experiments. Although we are not able to attach images of the loss curves in this OpenReview response, we can say that the shapes of the loss curves for all CLA and non-CLA models are qualitatively similar, and that all models have loss curves typical of transformer-based LLMs trained with cosine-decay learning rate schedules.

Formatting and typos

We appreciate the careful attention the reviewer put into our work, and will fix these errors.

Why should we expect cross-layer sharing to help relative to using only GQA/MQA?

GQA and MQA can be seen as relying on the hypothesis that in MHA models there is hidden redundancy in the content of KV activations in different heads within each layer. GQA and MQA then exploit this hypothesized redundancy to reduce the size of the KV cache relative to MHA, with only minor degradation in performance.

Similarly, our original motivation for investigating Cross-Layer Attention was based on the hypothesis that even with GQA/MQA, there may be remaining hidden redundancy in the content of KV activations across different nearby layers. If there is hidden redundancy across layers, then GQA and MQA have no mechanism for exploiting that redundancy to reduce KV cache size, no matter how we set our GQA/MQA hyperparameters -- however, CLA would be able to exploit that redundancy.

We would be happy to mention this design motivation in the final version of our paper.

What about adapting pre-trained models to use CLA?

We agree that adapting (or "uptraining," as the GQA paper calls it) existing models to use CLA is an interesting avenue of research.

We have conducted some preliminary experiments on CLA uptraining. We have found that it is possible to convert 1B- and 3B-scale MQA models each trained on 100B tokens into MQA-CLA2 models with 5.7% and 4.9% higher perplexity, respectively, by uptraining them with CLA for 20B tokens. Similarly to the GQA paper, we find that initializing the uptrained model by mean-pooling KV projection weights outperforms simply dropping KV projections from the model.

At the reviewers' discretion, we would be happy to include these preliminary uptraining results in the final version of the paper. We also believe it is likely possible to improve further upon our uptraining scheme, and would be happy to mention in the paper that improved schemes for CLA uptraining represent a promising direction for future research.

We also note that although the original GQA paper focused on uptraining, GQA has also had a significant impact on industry practice via its direct application to LLM pretraining, as seen in models like the Llama 3 series. We hope that CLA may be able to have a similar impact via direct application to pretraining.

We thank the reviewer for their thoughtful and generous feedback. We address each point below:

Could we evaluate CLA on larger models?

Unfortunately, we lack the resources to train 10B- and 100B-scale models from scratch using CLA. Even training a single model on the same parameter and data scale as Llama 1 7B would require $30\times$ the training compute of our largest 3B-scale training runs. We hope that future work will apply CLA to models at larger scale.

Why don't we evaluate CLA in combination with MHA?

We chose to spend our resource budget evaluating models combining CLA with GQA/MQA, rather than MHA, because we believe that most practitioners who would apply CLA would want to do so in combination with GQA/MQA.

As we mention in Section 2.3, CLA reduces KV cache storage footprint, whereas GQA/MQA reduces both KV cache storage footprint and memory bandwidth requirements during decoding. Because MQA/GQA provides these dual benefits, we expect practitioners to prefer using MQA/GQA first to obtain KV cache size reductions over MHA, and to only apply CLA to obtain further reductions on top of GQA/MQA. We note that because of the efficiency benefits of GQA at inference time, many recent models, such as Llama 3 and Mixtral, use GQA rather than MHA.

We chose to include the MHA model in our design space exploration just as a point of comparison, to contextualize the performance achieved by other GQA/MQA and CLA models. We would be happy to clarify this point, and the reasoning for why we did not directly combine CLA with MHA, in the final version of our paper.

Why do we use separately-learnable affine parameters for the KV projection?

We chose to use separately-learnable affine parameters in the layernorms for our KV projections because we did not want to arbitrarily privilege one of the several attention layers sharing each KV projection to also share its layernorm affine parameters, and because it was a convenient way to implement CLA within our framework. We note that using separately-learnable affine parameters increases the parameter counts of our models by less than 0.003% in all cases. Due to limited resources, we did not ablate this choice.

Could we conduct sweeps over batch size, or other hyperparameters?

We chose to use our available resources to sweep over learning rate because it is typically the hyperparameter to which final model performance is most sensitive. For our choice of batch size, Adam $\beta_1$ and $\beta_2$, and AdamW weight decay coefficient, all our experiments used the same standard values used for pretraining Llama 2. We agree that, in the absence of resource constraints, sweeping over batch size and other hyperparameters would be valuable, but we note that the optimal learning rate can depend on batch size; this means that sweeping over batch size while ensuring we were comparing models at their best learning rates would in fact require a prohibitively expensive joint sweep over both batch size and learning rate. We thank the reviewer for their suggestion, and leave more exhaustive ablations of different training hyperparameters as a possible direction for future work.

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.

We thank the reviewer for their thoughtful and generous feedback. We address each point below:

Could we report more quality metrics?

In addition to Wikitext perplexity, we also report accuracy scores on Hellaswag, PIQA, WinoGrande, OpenBookQA, BoolQ, and ARC-E, which can be found in Table 3 and Table 4. In aggregate, we find that these metrics show that our CLA models match the accuracy of our non-CLA baselines while using only half as much KV cache memory.

In response to your question, we have evaluated our 1B- and 3B-parameter models on MMLU, but find that MMLU does not provide useful signal for comparing the quality of different models at our scale of training compute. In particular, we observe that all the models which we compare in Figure 1 score no better than chance (25% accuracy) on MMLU:

For a point of reference, even the publicly-available open-source models TinyLlama 1.1B and OpenLlama 3B, each of which was trained with $\approx 10\times$ the FLOPs of our largest 3B training runs, do not achieve accuracy significantly better than chance on MMLU:

What about continuous pretraining?

We agree that continuous pretraining (or "uptraining" as the GQA paper calls it) for applying CLA to existing models is an interesting avenue of research.

We have conducted some preliminary experiments on CLA uptraining. We have found that it is possible to convert 1B- and 3B-scale MQA models each trained on 100B tokens into MQA-CLA2 models with 5.7% and 4.9% higher perplexity, respectively, by uptraining them with CLA for 20B tokens. Similarly to the GQA paper, we find that initializing the uptrained model by mean-pooling KV projection weights outperforms simply dropping KV projections from the model.

At the reviewers' discretion, we would be happy to include these preliminary uptraining results in the final version of the paper. We also believe it is likely possible to improve further upon our uptraining scheme, and would be happy to mention in the paper that improved schemes for CLA uptraining represent a promising direction for future research.

We also note that although the original GQA paper focused on uptraining, GQA has also had a significant impact on industry practice via its direct application to LLM pretraining, as seen in models like the Llama 3 series. We hope that CLA may be able to have a similar impact via direct application to pretraining.

What about long-context tasks?

Training a high-quality long-context model with or without CLA was not possible within our resource constraints. As we mention in Section 4, we leave larger-scale evaluations of long-context, aligned models employing CLA as an interesting problem for future work. We thank the reviewer for their suggestion, and agree that in such future work, passkey retrieval would be a worthwhile task to evaluate.