Prompt-prompted Adaptive Structured Pruning for Efficient LLM Generation

Thank you for the detailed feedback! We address each concern/question below.

Concerns:

1. Thanks for pointing this out, we can reframe the paper as a dynamic structured pruning method in our next version.
2. We designed a new baseline inspired by Wanda [2], a popular pruning method. The full model is used for the prompt while FF weights selected by Wanda based on the input’s features are used during generation (Adaptive Wanda). Comparisons between GRIFFIN and Adaptive Wanda are found in Table R1 under Reviewer cGRh. Both perform similar to the full model, but GRIFFIN has the added benefit of highly structured pruning whereas Wanda is unstructured.
3. Sampling-based methods show greedy top-k is best and are additional baselines. Fixing neurons across a dataset is interesting! In Table R3, we test two ways to select fixed FF neurons for generation which we will add to our next version. GRIFFIN is more performative in most cases, implying a need for adaptation.

Table R3: 1-shot XSum and CNN/DailyMail Rouge-1 with fixed expert neurons at 50% sparsity. Shot Top-k selects FF experts based on $\mathbf{s}$ in the example shot. Global Top-k selects FF experts based on the average $\mathbf{s}$ across a dataset’s prompts.

4. A longer prompt is unlikely to harm efficiency much: Table 3 shows latency is dominated by generation, e.g. Llama 2 7B spends ~8.5% of the total latency on prompts.
5. GRIFFIN applies to FF blocks because it uniquely has flocking. This also makes it versatile as other efficient techniques can be independently deployed in attention (e.g. KV compression) and embedding layers that could further accelerate inference.

Questions:

1. Relative FF activations are FF activations scaled down to unit vectors. If z is the vector of FF activations of a token, relative FF activations of that token is z / ||z||_2.
2. We noticed this too and initially played around with some simple ways to vary k across layers (e.g. skipping some layers and linear scheduling), but we did not notice much effect on performance. In future work, it would be interesting to explore more ways to choose k, for different layers and different samples.

[2] Sun, M. et al, A simple and effective pruning approach for large language models, 2023.

Thank you for the detailed review of our work, especially on real world considerations for deployment!

Based on your suggestions, we performed experiments on batching with highly promising results. Within a batch, we sum the selection statistic $\mathbf s_i$ of each sample, scaled inversely by each sample’s root prompt length, $S_i$, i.e. $\mathbf{s} = \sum_{i} \mathbf{s}_i \sqrt{S_i}$. Then, the same experts for all samples within a batch would be selected based on $\mathbf{s}$. Results in Table R2 convey that GRIFFIN is also robust to batching, as performance degrades very slowly as we increase the batch size. As such, while GRIFFIN was especially designed for single sample use cases (e.g. personal devices), our method scales impressively well to large batch sizes.

Table R2: 1-shot XSum Rouge-1 scores for increasing batch sizes. For each batch, the same 50% of FF expert neurons are selected.

In the next version of our paper, we will include the results in Table R2 in the form of a graph. Let us know if you have any further questions and comments.

Thank you for reviewing our submission and highlighting some key takeaways! We address your concerns/questions below. Please let us know of any lingering questions and feedback.

Concerns:

1. We can include similarity plots between tokens’ relative activations. Since we are unable to share images for the rebuttal, we roughly describe them. Using the same sample in Figure 1, many Gemma 7B layers had high cosine similarity (0.6 to 0.8) between tokens’ activations. Most of Llama 2 7B’s hovered around 0.35 to 0.4.
2. We can add example generation outputs. For 1-shot CNN/DM, Mistral 7B and Llama 2 13B often produced a sequence of “\n” or repeated items from the example shot. With 3-shot, this is not an issue (>25 Rouge-1 for both).
3. We will also include some example generation outputs from GRIFFIN and the full model in the next version.
4. Thanks for bringing this to our attention. We will reframe our work as dynamic structured pruning in our next paper version.

Questions:

1. We can elaborate on this in the next paper version. For example, the MoEfication algorithm [3] fine tunes non-ReLU activations to ReLUs before constructing MoEs, but it is not the only one to require exact sparsity in FF activations [4, 5]. In turn, these have motivated fine tuning LLMs to have ReLU activations [6, 7].
2. Row vectors in Figure 3 are approximately l2 unit vectors. For illustration purposes, we rounded to one decimal place, inducing slight rounding errors.
3. See #2 in “Concerns”.
4. In our evaluations with XSum and CNN/DailyMail, the example shot and test samples are very different, yet GRIFFIN still appears to do well. The cause of flocking is of great interest to us. Flocking occurs in many different tasks–Figure 6 shows that it even occurs with nonsense inputs. Greater analysis in flocking in terms of cause and interpretability is something we are excited to work on in the future!
5. Thank you for catching this typo! We will correct it in the next version.

[3] Zhang, Z., et al., Moefication: Transformer feed-forward layers are mixtures of experts, 2021.

[4] Csordás, R., Irie, K., & Schmidhuber, J., Approximating two-layer feedforward networks for efficient transformers, 2023.

[5] Li, Z. et al., The lazy neuron phenomenon: On emergence of activation sparsity in transformers, 2022.

[6] Mirzadeh, I. et al., Relu strikes back: Exploiting activation sparsity in large language models, 2023.

[7] SpaseLLM Team, Sparse large language models with relu activation, 2023.

Thank you for your valuable comments and curiosity in our baselines! We address your comments below. Please let us know if you have any follow-up questions.

As a new baseline, we use Wanda [2] during generation. This baseline is stronger than the magnitude selection method and is well-known in the pruning community. In more detail, the full model is still used during the prompt phase while FF weights selected by Wanda using the current inputs’ features are used during generation (Adaptive Wanda). Comparisons between GRIFFIN and Adaptive Wanda can be found in Table R1.

Table R1: Adaptive Wanda and GRIFFIN on XSum (1-shot), CNN/DailyMail (1-shot), CoQA (0-shot), SCROLLS QASPER (0-shot) at 50% FF sparsity.

We observe that GRIFFIN and Adaptive Wanda are both great at preserving the original model’s performance. However, we note that Wanda’s pruning algorithm is unstructured, meaning it cannot be easily translated to real latency improvements. Meanwhile, GRIFFIN is able to be performative despite the structural constraint we imposed. We will add this baseline (including its performance with other models) in our next paper version.

Following up on your comments on Deja Vu [1], we recognize its ability to dynamically select neurons during inference, but we do not think it is comparable. A major reason is that Deja Vu requires data curation, training, and extra parameters, so compared to GRIFFIN (which requires no preparation), there is a significant gap in resource requirements. Unlike Deja Vu, our new baseline, Adaptive Wanda, selects FF neurons depending on the input without training or data curation and exhibits competitive performance.

[1] Liu, Z. et al, Deja vu: Contextual sparsity for efficient llms at inference time, 2023.

[2] Sun, M., Liu, Z., Bair, A., & Kolter, J. Z, A simple and effective pruning approach for large language models, 2023.