FTP: Efficient Prefilling for Long-Context LLM Inference via FFN Token Pruning

Q9. Be more explicit and sharp.

A9. Thanks for your guidance on improving our manuscript and we will refine this paragraph in the revised version accordingly. As mentioned in Lines 100~103 and depicted in Figure 4, the distinction is that FTP only prunes tokens that are input into the FFN module, yet the pruned tokens are still reserved in this layer because of the residual connection surrounding the FFN module (see the yellow tokens in Figure 4, their features are reserved before and after the FFN module). In other words, FTP does not reduce the number of tokens fed into each layer. Contrarily, previous token pruning methods applied token pruning to the whole layer (i.e. layernorm + attention + FFN), resulting in various numbers of tokens input to each layer.

Thank you for your thoughtful comments and time.

Q1. Discussion about memory impact.

A1. Our work focuses on the acceleration of the prefilling stage for long-context LLM inference and the memory impact of FTP is trivial. Prior KV cache eviction methods are motivated in terms of memory reduction and mainly optimize the computational and memory cost during the decoding stage, while our work reduces the computational cost during the prefilling stage.

We run Qwen2-7B-Instruct with/without FTP on all datasets from LongBench, and record the peak memory usage during the prefilling stage for each data sample. The average result is in the table below.

Peak memory usage during the prefilling stage of Qwen2-7B-Instruct

Q2. Literature review.

A2. Thanks for your advice on our literature review. References [1,2,3] are significant works on token pruning outside of LLM, and we will incorporate them in the related works section in the revised version.

However, the most evident differences between FTP and these works are three folds:

1. FTP focuses on inference of LLMs with billions of parameters, whereas [1,2,3] conduct experiments on models with millions to hundreds of millions of parameters;

2. FTP does not need any re-training or post-training, which is time-consuming for LLMs. However, reference [1] proposes a novel model architecture and requires training from scratch. Reference [2, 3] also involves re-training or post-training to learn some threshold or parameters to enable token pruning, which is not in the same setting as our method.

3. FTP does not prune tokens for layers and it only reduces the input tokens for the FFN module, which accounts for over 60% of the time during the prefilling stage, while [2,3] prune tokens for the whole layer or even successive layers and need re-training to preserve accuracy.

Q3. Overall trade-off feels a bit lukewarm.

A3. The overall trade-off of our method is considerable compared to other relevant works (e.g. LLMLingua2, PyramidInfer). As mentioned in A1 in our response to common questions, PyramidInfer failed to accelerate the prefilling stage when implemented with FlashAttention. Moreover, as shown in Figure 7, FTP can achieve a larger speedup (1.4x ~ 1.5x) with bearable accuracy loss for most tasks even applied to a 7B model.

For the accuracy drop by Llama3-8B-Instruct for code completion, we attribute this mainly to the limited context length (8K) of Llama3 models. To further investigate, we conduct additional experiments on Llama3.1, which supports a context length of up to 128K tokens. For a more detailed analysis and results, please refer to our response A1 to reviewer 5w9V.

Q4. Comparison with PyramidInfer

A4. Please refer to Q1 in the response to common questions.

Q5. More benchmarks.

A5. Please refer to Q3 in the response to common questions.

Q6. Regarding the paper layout.

A6. Thanks for your advice and we will try to merge the tables in the revised version.

Q7. The random baseline.

A7. Thanks for your advice, but the random baseline is already in the ablation study section (Section 4.6) if I am not misunderstanding the question.

Q8. Confusion in Figure 7

A8. Figure 7 illustrates the trend of accuracy as the TTFT speedup grows, and demonstrates the trade-offs between them. As mentioned in Lines 424~425 in Section 4.4, the TTFT speedup of FTP can be modulated by tuning not only $\eta$, but $\mathcal{F}$ as well. Additionally, we also conduct specialized experiments for these two hyper-parameters in Appendix 6.1, from which most data points in Figure 7 are collected.

A1. Thank you for the additional information. It seems many of the prior methods that have focused on pruning and efficiency also takes memory reduction as a concern. For example even PyramidInfer seems to motivate "over 54% GPU memory reduction in KV cache". If your method doesn't help in that respect this should be more explicitly acknowledged as a limitation.

A2. Note [1,2,3] is a non-exhaustive list. Many similar works have been attempted before - which should be ideally added to the review. They can be found from the works cited by [1,2,3] and works citing [1,2,3]. I acknowledge the differences with your work.

A3. I am not sure if limited context length explains why the baseline is performing much better. Even the baseline should suffer from the effect of limited context length. But nevertheless, it's good to say the high performance drop only happened for one instance for a specific model (for whatever reason) so far.

A4. I am still not that completely convinced about the comparison. If I understand correctly, FlashAttention not returning attention weights would be simply an implementation issue not something of an in principle issue that cannot be solved in a future implementation. Recomputing attention scores again does not sound to me as a reasonable baseline. A more reasonable comparison could be running all the three methods with pytorch-implemented attention (unless you can explain why FlashAttention is theoretically impossible to utilize for PyramidInfer in an effective manner - and the bottleneck is not just accedential facts about how the implementations are currently done). Moreover, it's unclear here why only PyramidInfer faces issue for FlashAttention not returning attention weight, when your method seems to rely on attention scores too.

A7. Yes, my bad.

A8. My point is that dots are not annotated. For instance, I can find a dot that has speed up 1.1 and relative score 60%, and by the color I may know it's Llmlingua. But how do I know which prompt compression rate that co-ordinate represents? Trade-off graphs like that typically have annotated points.

A4 is currently the main point that is holding me back from increasing the score.

The other points should be ideally incorporated in the discussion/presentation of the paper.

Thank you for your thoughtful comments and time.

Q1. Major performance loss on Code Completion Tasks with Llama3 model compared to baselines.

A1. The code completion tasks in LongBench consist of two datasets: LCC and Repobench-P. Our explanations for this phenomenon are three folds:

1. The insufficient context length of Llama3. Llama3 has a context length of only 8K, but the "Repobench-P" dataset has more 283 out of 500 samples that exceed the context length of Llama3, which can degrade the performance of LLMs. It is probable that the token pruning by FTP further degrades the model with insufficient context length and results in more accuracy loss.

2. On the other hand, the accuracy loss of FTP applied on Qwen2-7B, Qwen1.5-32B, and Qwen2-72B models are all acceptable, which indicates that FTP can work in various sizes of models. It is worth noting that all of these models have a context length of 32K.

3. As Llama3.1 has extended its context length to 128K, which is larger enough for the prompt length of all data samples in the "Repobench-P" and "LCC" dataset in LongBench, we evaluate FTP on Llama3.1-8B-Instruct with the same setting as Llama3-8B-Instruct in our paper, and the results are shown below.

Llama3.1-8B-Instruct on LongBench

Q2. Comparision with Pyramidinfer.

A2. Please refer to Q1 in the response to common questions.

Q3. How to obtain the attention score based on FlashAttention?

A3. FlashAttention does not return attention weights. We compute the necessary attention weights (the last N rows, as depicted in Figure 4) for token selection. As we discuss in Lines 241~243 in our paper, prior work [R3] has empirically revealed that the attention pattern obtained by the queries at the end of the prompts is nearly consistent with that obtained by all queries, we only retain the attention scores M′ from the last N queries to reduce overhead.

A quantitive analysis of the computational cost is in our response A1 to reviewer QpPZ.

Q4. How to determine the best reserved ratio for each layer's FFN?

A4. The reserved ratio $\eta$ for each layer's FFN is a hyper-parameter of FTP. The optimal $\eta$ depends on the tolerance of performance degradation in the application. As mentioned in Section 4.1, for Llama3-8B-instruct, we empirically preserve the first F=10 layers and set $\eta$ = 0.90 for the following layers. When it comes to the Qwen2-7B-model, F and $\eta$ and set as 10 and 0.95 respectively. More experiments for the reserved ratio are detailed in Appendix 6.1.2.

Another guidance for determining an optimal reserved ratio $\eta$ is in Section 3.2.2 and Figure 6. As depicted in Figure 6, nearly 50% of tokens in a layer are evicted when $\eta$ is set as 0.90.

Q5. Do the increasing number of pivotal tokens in Attention reflect that we should also reserve more tokens?

A5. Yes. The number of reserved tokens is determined by the hyper-parameter $\eta$ for this layer. If there are more pivotal tokens (i.e. more tokens with high attention scores), our method would adaptively reserve more pivotal tokens to achieve the percentage ($\eta$) of the summed attention weights.

Thank you for your thoughtful comments and time.

Q1. Add another benchmark

A1. Please refer to Q3 in the response to common questions.

Q2. Why only these two baseline are selected

A2. Please refer to Q2 in the response to common questions.

Q3. How conclusions 2 and 3 are derived from Figure 5.

A3. Let us recall conclusion 2 and 3 first.

Conclusion2: tokens with high attention scores in one layer may not necessarily be prioritized in the same manner across other layers;

Conclusion3: the number of tokens with high attention scores varies among different layers.

Conclusion2 can be inferred by analyzing the color of a single column in a sample in Figure 5. For example, the token at position 0 exhibits high attention scores in layers 1~6 and 16~28, yet its scores decrease in layers 10~15 and layer 31. A similar pattern is distinctly observable on the tokens at positions larger than 100. Moreover, tokens at position 60~80 also support our conclusion, even though the variation in attention scores is not so obvious as that of the tokens at the initial and final positions. Based on this conclusion, FTP is designed to select important tokens in each layer.

Conclusion3 can be inferred by analyzing the color of a single row in a sample in Figure 5. For instance, in layers 5~18, there are more tokens with high attention scores, and in layers 2~3 or layers 20~25, the attention scores highly concentrate on the initial tokens, resulting in fewer high-score tokens. This conclusion motivates us that the number of tokens to be retained should be determined by a ratio $\eta$ of attention scores, instead of an absolute number.

Q4. Typo in Line 267

A4. Thanks for your advice and we will correct it in the revised version.

Thank you for your thoughtful comments and time.

Q1. Regarding the computational cost in calculating the Sum Attention Score.

A1. The computational cost in calculating the sum attention score is trivial since we only calculate the attention scores from the last N queries to all input keys (referenced in Line 180 in our paper) instead of calculating the full attention scores. The value of N is much smaller than the input length L. We set N = 50 in our experiments, which means that N is 2~3 orders of magnitude lower than L in most long-context scenarios. Furthermore, we conduct an experiment on LongBench with Llama3-8B-Instruct and Qwen2-7B-Instruct and summarize the TTFT of FTP and the random variant (referenced in Section 4.6.1), which drops the same number of tokens in each layer as FT with negligible computational cost. As shown in the table below, FTP on Llama3-8B-Instruct only introduces a little computational cost of about 7~10ms, which accounts for only 1%~3% of the TTFT, whereas FTP on Qwen2-7B-Instruct introduces a cost of 8~15ms, which takes up 0.8%~1.9% of the TTFT.

TTFT (ms) of Llama3-8B-Instruct

TTFT (ms) of Qwen2-7B-Instruct

Q2. Justification or evidence that the Attention Score is also effective for FFN token pruning. Lacks design tailored to FFN.

A2. First of all, FTP does not evict tokens as the KV Cache eviction methods do. FTP allows pruned tokens to retain their features and pass them to the next layer. The insight of using attention scores for FFN token pruning is that tokens with low attention scores in the current layer are very likely to have low attention scores in the next layer (see the columns in the samples in Figure 5 and Appendix 6.2). Therefore, these tokens can bypass updates by the FFN module and simply retain their features from the attention module in the current layer. As these features are fed into the attention module in the next layer, the error introduced by omitting the FFN update for the pruned tokens is minimized since their attention scores are still low.

Moreover, by design, we manually preserve the first P tokens and the last N tokens in each layer, ensuring that they are never pruned, since their attention scores in the current layer are not always consistent with that of the next layer (see the columns in the samples in Figure 5).

Last but not least, our experiment in section 4.6.1 is also evidence of the effectiveness of our attention-based strategy.

One of our specific designs for FFN is that we utilize the residual connection to retain more information of the pruned tokens instead of discarding them as in the KV Cache eviction methods. Another design is that we logically set the FFN outputs of the pruned tokens to zero vectors. As depicted in the "Feed-forward Network" part in Figure 4, if the feature of a token is deactivated after the "Activation" module (e.g. the SiLU function), its FFN output closely resembles a zero vector. Hence, the error introduced by skipping the FFN module can be reduced by setting the FFN outputs of the pruned tokens to zero vectors.

Q3. FTP discards most intermediate tokens. Please analyze the eviction effects of FTP and assessment of the loss of intermediate information.

A3. On one hand, FTP does not discard tokens. As mentioned in Lines 102~103, FTP prunes tokens that are fed into FFN. However, the pruned tokens are just not updated by the residual produced by the FFN, and their feature remains unchanged due to the residual connection surrounding the FFN module. On the other hand, as illustrated in Figure 6, FTP still reserves nearly 60% of tokens when $\eta$ is set to 0.95, which indicates that not only the beginning and end tokens are retained. We further visualize the distribution of the relative position of the retained tokens in the figure in the revised version.

We are also running "needle in a haystack" [R5] experiments for FTP, which provides a quantitative assessment of the loss of intermediate information.

References

[R5] G Kamradt. Needle in a haystack–pressure testing llms, 2023.

Q4. Similar research that also performs FFN token pruning, such as MoD[2], is not discussed in the paper.

A4. Thanks for your suggestions on our literature review, and we will include MoD in our related works. MoD applies a router to determine for a token whether to skip both the self-attention and MLP blocks in a layer, and it needs training for the model. On the other hand, FTP only prunes tokens for the FFN module and does not need any post-training for the LLM, which is time-consuming. In the experiment section, MoD is applied to models with 60M to 3B parameters, whereas FTP focuses on the inference of LLMs and incorporates models with 7B~72B parameters.

Q5. The choice of comparison methods.

A5. Please refer to Q2 in the response to common questions.

Q6. Inconsistency between the attention implemetation of PyramidInfer and FTP.

A6. Please refer to Q1 in the response to common questions.

Q7. Discuss the difference between this work and other dynamic layer-skipping methods

A7. The most significant difference between FTP and other layer-skipping methods is that FTP only reduces the number of tokens fed into the FFN, and the number of tokens fed into each layer remains the same across all layers, since the pruned tokens remain identical after the residual connection. Some layer-skipping methods [2, R6] require post-training which is time-consuming for LLMs.

References

R6. Elhoushi, M., Shrivastava, A., Liskovich, D., Hosmer, B., Wasti, B., Lai, L., ... & Wu, C. J. (2024). Layer skip: Enabling early exit inference and self-speculative decoding. arXiv preprint arXiv:2404.16710.