SIRIUS : Contexual Sparisty with Correction for Efficient LLMs

Thank you so much for your time and attention on the paper. We are very grateful for you to be interested in the KV Cache correction. We will try our best to answer your questions.

1. Full and Sparse of the same architecture

Contextual Sparsity is the focus of Sirius. Contextually Sparse (CS) models are dynamically subsampled from the normal LLM. The subsampling process won't change the compatibility of the original model's cache with the newer subsampled sparse model. Therefore, the original cache is naturally usable by CS models. However, we are interested to see whether Sirius can be applied to model pairs outside the CS realm.

2. When to rewrite KV Cache

We appreciate your suggestions, and we rewrote the method part of the paper to take into consideration KV rewriting confusion. The rewriting happens when the full model is correcting the sparse generated tokens. The full model is called once every kernel size, usually 16. Throughout the inference, the KV cache is shared between the sparse and the full model. The KV cache is mostly populated by the sparse model, which is called for every token. During correction, the full model takes in the last kernel size of tokens and generates its KVs for the past 16 tokens in parallel, these KVs are directly written to their corresponding positions in the shared KV Cache.

3. Will the full model's KVs make semantic sense to the sparse model

We appreciate the brilliant question. Empirically, we found that the KV cache of the large model seems to always help the small model generation quality. Below shows one of the example on GSM8K (20% subsampled)

4. Whether correcting minor portion leads to performance recovery is applicable to more datasets

We appreciate the comments. Although we only show for GSM8K in the text. However, in the newer version of the paper. We evaluated Sirius across diverse datasets, ranging from arithmetic reasoning, common sense reasoning, and coding. Sirius managed to achieve competitive efficiency metrics on most of them. These results show that for general text generation that requires challenging step-by-step correlation, the minor portion is the only part that needs to be corrected to recover the original performance.

5. Writing is unclear

We deeply appreciate the suggestions on improving the writing quality of the paper. All of the suggested parts are revised in the newer version of the paper.

6. How to determine thresholds.

Threshold determination is a crucial process of balancing the performance and efficiency of Sirius. Here we provide a much more in-depth ablation on how increasing the threshold exhibits the trend that improves the performance but hurts the efficiency. The efficiency of Sirius is measured by the average advance length out of a kernel size (usually 16 tokens), the higher the better.

7. Long range dependencies

Sirius allows the large model to corrects the small model on a basis of kernel size. Besides KV rewriting, roll back mechanism is especially important for generating longer text. The rollback works as follows, when full model parallel verifies, If there is any token within the range of the kernel size that the full model regards as unlikely (rejects by likelihood threshold), then the token will be removed. Therefore, regardless of the generation length, all the verified tokens will be regarded by the full model as "enough likely" given the context.

Thank you for the time taken to read the paper and your insightful comments. We are thankful that you point out that the comparison against speculative decoding is missing and we lack of convincing real efficiency metrics. We will try our best to address your concern. If you feel that your concern has been taken into account, please consider raising your score.

1. Differences to Speculative Decoding? Why Speculative Decoding is not efficient enough?

Sirius contains a stronger and a weaker model, the strong verifies the weaker's output, from which the resemblance can be easily drawn between Sirius and Speculative Decoding (SD). However, the important distinction is the problem setting.

Contextual sparsity users prioritize the tradeoff between the performance and the efficiency of the model. We show that the efficiency of the contextually sparse model corrected by SD will be largely limited. Take the Coarse-grained sparse model on Llama-3-8B as an example, the 50% sparse (APU 0.65) model will produce an acceptance rate of 0.89 on GSM8K. SD also runs the smaller model every iteration with verification of the larger model once a period. Naturally, to increase efficiency, we need to (1) enlarge the period and (2) increase the average number of tokens accepted in the period. \frac{1 - \alpha^{(\gamma + 1)}}{1 - \alpha} Now, given the acceptance rate of 0.89, following the formulation in Speculative decoding, we can calculate the expected number of accepted tokens for every gamma term in the Speculative Decoding literature, which is (period - 1).

We can notice the trend that the average advance length starts to plateau as the period becomes larger. Take the gamma of 16 as an example, the period is then 17. If we use the formula in Section 4.1 of our paper, we can immediately see that the APU is (16 * 0.65 + 1)/7.84 = 1.45, even larger than full model 1.

Because of the plateauing effect, for an acceptance rate of 0.89, the best gamma is 2 (period = 3). The best efficiency is 0.86, compared with 0.65 coarse-grained sparse APU. A similar picture can be applied to Fine-grained sparsity as well. The key reasons are (1) the contextual sparse model is too large to use a long period for SD, and (2) the acceptance criteria are too strict.

In contrast, Sirius gives a more flexible choice for contextually sparse model users. For a threshold of 0.1, Sirius can correct Llama-3-8B coarse-grained sparsity from 20.85% to 43.9%, compared to the 49.66% full model. Sirius on average can accept 13.4 tokens out of a kernel size of 16 and over 9 tokens out of a kernel size of 10, translating to APU < 0.76, significantly lower than SD does.

2. Hardware speedup?

We show that Sirius delivers the promised speedup. With the introduction of static cache from huggingface 4.38, kv cache is pre-allocated, making rewrite, expand, and rollback of very minimal overhead. Fine-grained sparsity speedup is via prior work [1]. However, CATS relies on a close-sourced custom CUDA kernel. We present the speedup of Coarse-grained sparsity with Sirius in two settings. First, on-chip, we run Llama-3-8B-Instruct inference.

Second, Offloading Llama-3-70B-Instruct. We use a single L40 48GB with a PCIe bus bandwidth of 25 GB/s.

[1] Liu, Z., Wang, J., Dao, T., Zhou, T., Yuan, B., Song, Z., ... & Chen, B. (2023, July). Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning (pp. 22137-22176). PMLR.

3. Criticism of writing quality We do appreciate the suggestions on writing. In light of your precious effort put into reading, we carefully rewrote and proofread the entire methodology section that takes all your suggestions into account.

4. KV Cache correction | Setting|GSM8K 20% |Setting| GSM8K 20% | |-------------------------------------------|----------|-----------------------------------------------------|----------| | Llama3-8B-Instruct | 0.7538 | Llama3-8B-Instruct | 0.7538 | | Llama3-8B-Instruct + CSparse | 0.3674 | Llama3-8B-Instruct + FSparse | 0.5644 | | Llama3-8B-Instruct + CSparse + KV Cache Correct | 0.4735 | Llama3-8B-Instruct + FSparse + KV Cache Correct | 0.6629 | We show that KV Cache correction contributes to correction. However, we agree that "implicit distillation" is unbased, which has been deleted in the latest version.

5. Why Sirius? Sirius is an astronomical term referring to a two-body star system. One is the brightest ever known, while the other is dim. We draw inspiration from that system.

Thank you for the time taken to read the paper and your insightful comments. We are thankful that you suggested we add more experiment data to evaluate Sirius. We will try our best to address your concern. If you feel that your concern has been taken into account, please consider raising your score.

1. Too few empirical experiments

We provide additional empirical results and analysis of Sirius on various models and different downstream datasets.

As discussed in the paper, the weakness of the contextual sparse models seems to occur upon difficult reasoning tasks. Besides, we collect experiment results on different reasoning tasks. In Arithmetic reasoning, besides GSM-8K, we also show a difficult AQuA-RAT COT from Google. In the Common Sense reasoning, we follow the COT paper to evaluate CSQA, StrategyQA, Sports, and Dates. Also, we found that contextual sparse models cannot do well in code generation. We then run Sirius on HumanEval.

We collect additional experiment results on six models (Llama-3-8B, Llama-2-7B, Llama-2-13B with their instruction fine-tuning counterparts). Due to the word limit, we only show AQuA, CSQA, and HumanEval on some models. Full results and Average advance length will be in the new version of the paper. L means Llama. O means Original S means Sparse SS means Sparse + Sirius

Sirius can also work effectively on the Llama-3-70B-Instruct, model with large parameters.

2. Other questions.

a. We appreciate the suggestions on Pseudo code, we have rewritten it in the newer version.

b. Correction step by full model is slower than sparse, but the cost is amortized by the kernel size of 16, which is less significant.

c. Contextual sparsity subsample sparse model from the full model dynamically according to input, so the full model weight will always be in GPU memory. Adding correction won't increase the memory usage of the sparse model.

d. Please note that the latency is mainly affected by the memory loading not the shear memory size on VRAM. Full model is loaded only once per kernel size. Please refer to the hardware speedup on the new version.

Thank you again for your feedback.

We have added the necessary material to the author's rebuttal to address the points you raised. Given that we only have four reviewers, each reviewer's score significantly impacts the overall assessment of our work. We believe that the current overall assessment does not adequately reflect the contribution of our work. Therefore, we kindly request you reconsider your score. Thank you again for your time and effort in reviewing our paper.

Thank you so much for your attention and time taken to read the paper and your insightful comments. We are very thankful for you to point out the inadequacy of our method presentation and raising concerns on the diversity of the evaluation. We will try our best to address your concern. If you feel that your concern has been taken into account, please consider raising your score.

1. Studies mainly use smaller models, which should include Llama-3-70B to verify the claim that "the stronger the model, the larger the quality degradation would be".

In light of your comment, we run the Llama-3-70B-Instruct model on GSM8K, MMLU-FLAN-COT, CoQA, CNN/DailyMail, and TruthfulQA. Due to running Llama-3-70B-Instruct under pipeline parallelism is slow on our hardware (8xA100), MMLU-FLAN-COT is subsampled 10%, and CNN/DailyMail is subsampled 30%.

In the table, Llama-3-70B-Instruct still has features that we identified in the paper, sparse excels at prompt understanding (CoQA and CNN/DailyMail). Significant performance degradation occurring at GSM-8K-COT with coarse-grained sparsity. Even though the gap between sparse and dense for 70B is not as big as for smaller models, we don't think it contradicts the claim that the more powerful the model, the more degradation when doing sparse. First, even with Coarse-grained sparsity, the number of parameters is more than 45B, still colossal. Most tasks here are curated in pre-LLM era and are too easy for models now of huge parameters [1]. On the other hand, the claim is presented in the context of comparing Llama-2-7B Chat and Llama-3-8B-Instruct, two smaller but similar models. Llama-2-7B Chat has significantly worse performance in all different tasks. Contextual Sparsity causes much more damage to 8B than 7B. Nevertheless, trying larger parameter models does provide more insights to contextual sparsity characteristics.

[1] Chiang, W. L., Zheng, L., Sheng, Y., Angelopoulos, A. N., Li, T., Li, D., ... & Stoica, I. (2024). Chatbot arena: An open platform for evaluating llms by human preference. arXiv preprint arXiv:2403.04132.

2. Forwarding twice, cost doubled?

The tokens produced by the sparse model will be fed into the full model for verification, but the latency cost of full model verification will NOT be a lot, compared to the full model generating one new token. The cost in terms of latency won’t be doubled. The full model will verify the tokens generated by the sparse model periodically, usually in a period of 16. During full model verification, all tokens generated by the sparse model that haven’t been verified will be fed in one iteration. And, the full model will generate one more token in this run. Consider the following table of A100 latency of different sequence lengths with bsz=1. Only 1.1ms increase for 64 seqlen. The verification is as light as generating one token, without additional overhead.

3. When to correct?

Correction happens periodically. Specifically, the sparse model would generate PERIOD -1 tokens, and the full model would take in these tokens, verify them based on likelihood thresholds, rollback, and then generate one more token. The cycle goes on. Sirius usually takes PERIOD of 16, sometimes 12.

4. Degradation after correction?

We can answer by comparing Sirius with Speculative Decoding (SD). SD is lossless. Under the greedy setting, SD acceptance criteria guarantee that the two models' output is the same as the large model output. In comparison, here we effectively loosen up the acceptance criteria. We manually set a threshold on the full model's verification likelihood. i.e., If the full model thinks the token is likely to occur, it accepts. However, the token is not necessarily the one full will greedily decode, leading to divergence in generation and, thus degradation. However, the loosened criteria in turn boost the efficiency of the overall system by a margin to SD's.

5. Hardware wall clock time speedup? Sirius delivers the promised speedup. We present the speedup of Coarse-grained sparsity with Sirius in two settings. First, on-chip, we run Llama-3-8B-Instruct inference. Fine-grained sparsity prior works rely on a custom CUDA kernel, which we don't have access to.

Second, Offloading Llama-3-70B-Instruct. We use a single L40 48GB with a PCIe bus bandwidth to be 25 GB/s.

Thank you again for your feedback.

We have added the necessary material to the author's rebuttal to address the points you raised. Given that we only have four reviewers, each reviewer's score significantly impacts the overall assessment of our work. We believe that the current overall assessment does not adequately reflect the contribution of our work. Therefore, we kindly request you reconsider your score. Thank you again for your time and effort in reviewing our paper.

We again appreciate the review’s suggestions and inquiries on the efficiency of the claim in our paper. We follow up on our previous message with more experiment results on the comparisons on Llama-2-70B-Chat to provide more empirical studies on the claim. In the table below, we compare two pairs of models on GSM-8K-COT with similar parameters and the degradation after applying the contextual sparse methods (due to time limit, we subsample 20% from the entire dataset for 70B model). Also, we vary the sparsity level to have 50%, 40%, 30% of non-zero values in the weights, where lower than 30% would lead to Llama-3-8B-Instruct to have 0 accuracy).

From the table, we can clearly see that given similar parameter size, contextual sparsity brings more degradation to Llama-3 family models than to Llama-2 models (Please read the table vertically to compare between Llama-3-70B-Instruct with Llama-2-70B-Chat and Llama-3-8B-Instruct with Llama-2-7B-Chat).

Surprisingly, we also notice another interesting phenomenon where the models with larger parameter sizes seem to be more resilient to the contextual sparsity methods. Please read the table horizontally to compare models within the Llama-3 family models and Llama-2 family models. The trend is as expected since the model with a larger parameter size often has more redundancy in parameters.

Furthermore, the contextual sparsity’s weakness is in full display in the above table. We can see that for Llama-3-70B-Instruct, even at 50% sparsity, the performance on GSM-8K-COT is comparable to Llama-3-8B-Instruct. Given that the sparse 70B model still has over 40B parameter size, the performance degradation is unacceptable, let alone a sparse 70B model with lower sparsity levels (40% and 30%). Sirius corrects the sparse model to have 87% accuracy and incur negligible overhead during the offloading setting, again showing its effectiveness in helping contextual sparsity methods.

Also, as a follow-up on the efficiency-accuracy tradeoff concern, we would emphasize another set of results that was previously overlooked. Please be aware that the accuracy after Sirius correction is dependent on the sparsity level and the sparsity method. Previously, we looked at the coarse-grained sparsity methods where the sparsity pattern is determined for every input prompt and fixed throughout the generation. For the more flexible fine-grained sparsity, where the sparsity pattern changes for different decoded tokens, the sparse method alone achieved 59.68% accuracy on GSM8K-COT, versus the full model’s 75%, Sirius can correct the fine-grained sparse models to 74% accuracy with 0.775 theoretical Average Parameter Used to full model.

Thank you again for your additional thoughtful and acute comments. If you think that your concerns have been addressed, please consider adjusting your score.