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

Authors: Guo-Hao Xu, Jingzhen Ding, Huping Ding, Zhao Xu, Kaifu Zhang
Primary area: foundation or frontier models, including LLMs
Keywords: Large language model, Inference acceleration, Token pruning, Long-context inference, Natural Language Processing

Abstract:
Large Language Models (LLMs) have demonstrated remarkable performance across various NLP tasks, and have extended their capability to long-context scenarios. However, the increasing context length leads to longer inference time in both the prefilling and decoding stages.
Existing token pruning methods primarily evict tokens to compress the KV cache, and only accelerate the decoding stage.
Recent studies have extended token pruning to both stages, but they either yield subtle speedup during the prefilling stage or defer a portion of computations to the decoding phase. Critically, these approaches prioritize the attention module, overlooking the significant computations in the Feed-Forward Network (FFN) module.

In this work, we focus on the prefilling stage and propose a novel token pruning method named FTP for long-context LLM inference.
Our approach is based on the observation that the FFN module accounts for over 60\% of the inference time. FTP reduces this by pruning non-critical tokens before the inference of FFN. The importance of each token, along with the quantity to be pruned, are dynamically determined by the attention scores in each layer.
Unlike previous token pruning methods, FTP preserves a substantial amount of information of the pruned tokens through the residual connection, thereby achieving a notable speedup with only a negligible decrease in performance.
Specifically, the Qwen2-7B-Instruct model with FTP achieves a speedup of 1.24$\times$ in the prefilling stage with only a 1.30\% performance drop compared to the baseline model. The speedup is further boosted to 1.39$\times$ on a Qwen1.5-32B-Chat model.
Extensive experiments on long-context datasets across various tasks demonstrate the potential and effectiveness of FTP.

Decision: reject
Reviewer scores: 5, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730254144084

thread_type:
official_review

<thread_object>
creator: reviewer_pbr1

<content>
summary:
The paper introduces a new token pruning techniques at the pre-filling stage (before the autoregressive generative inference starts) for the FFN layer. This is mainly done by looking at the attention scores associated to each token, and keeping top-k tokens. The k is dynamically determined. Some heuristics are also employed  - for example, forcibly keeping first few and last few tokens from the original input and starting pruning only after the first few layers are executed. Performance in terms of accuracy and time taken till first token generation is shown in longBench.

strengths:
1. The method is straightforward and elegant.
2. Offers a reasonable speedup while mostly maintaining the original performance.

weaknesses:
1. No discussion about memory impact. Does this method reduces memory or keeps them the same? If there is a reduction - then by how much? Many of the prior methods are also motivated in terms of memory reduction.

1. Weak literature review. There are many works on token pruning outside LLM ([1,2,3] - just to name a few, but this is still a very incomplete list. More can be found by browsing the citation network). I found some of the contrasting against prior works as unsatisfactory. Overall, presentation can be more polished. 

1. The overall trade-off (accuracy loss vs speedup) feels a bit lukewarm. The speedup is only (~1.2x-1.4x depending on model size), and in most contexts it may not seem worth it for the accuracy loss - although that's a matter of specific use case. While accuracy loss is bearable in most contexts, sometimes it can be high like about 20% drop by Llama for code-completion.

1. I am not sure how fair the comparison between PyramidInfer vs author's methods given the former doesn't have flash-attention and the latter do. Is PyramidInfer simply not implemented with flash-attention officially or can it be not implemented in principle? If it can be implemented in principle, then the only reason it would be currently facing a disadvantage seems to be due to lacking the relevant implementation. In that case, could have been better to make a more apples-to-apples comparison even if that means using pytorch-attention for all (a separate table can be used for pytoch-attention based results). (resolved)

1. More benchmarks could have been attempted besides just the LongBench suite but not the biggest issue.


[1] Sparse Modular Activation for Efficient Sequence Modeling - Ren et al. NeurIPS 2023

[2] Learned Token Pruning for Transformers - Kim et al. KDD 2022

[2] Transkimmer: Transformer Learns to Layer-wise Skim - Guan et al. ACL 2022

questions:
1. There is a lot of repeated information in the tables. It seems to me that Table 1,2,4 can be the same table with all the approaches . Just consistently show TTFT (ms), and put all the non-large models (< 70B) there (the larger models are fine on a different table). 

1. The random baseline is better framed as an ablation.

1. The graph in Figure 7 seems uninterpretable if I am not missing anything. According to the text, it is supposed to show how speedup and accuracy varies with the $\eta$; but nothing the graph tells us which dot is associated with what $\eta$  or compression rates  for the other model.

1. "Unlike previous token pruning methods (Xiao et al., 2024; Li et al., 2024; Zhang et al., 2023b; Yang
et al., 2024b) where tokens are directly pruned from the whole layer, FTP dynamically selects a
certain proportion of tokens with an attention-based strategy and prunes them before the inference
of FFN." - This distinction is not sharp enough. Dynamic selection of certain proportion of tokens before FFN layers, is also direct pruning (by dynamic selection) from the whole (FFN) layers. So it's not explicitly said what the difference is. Perhaps you want to say that previously works applied token pruning to both attention + FFN, or that they apply permanent pruning (pruned in one layer => pruned in all later layers) and so on. But you need to be more explicit and sharp. . 

Review Caveat: I am not as familiar with modern works on LLMs for token pruning and not entirely sure where this method stands and if the best baselines have been chosen.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730537392277

thread_type:
official_review

<thread_object>
creator: reviewer_5w9v

<content>
summary:
This paper presents a simple and effective token pruning method for FFN called FTP, designed to improve the efficiency of large language models (LLMs) during the prefilling stage of long-context inference. FTP addresses this by dynamically pruning non-critical tokens before FFN inference, using attention scores in each layer to determine token importance and pruning quantity. Unlike previous KV-cache eviction methods focusing on improving the efficiency of the decoding phase, FTP achieves speedups in TTFT with minimal performance loss.

strengths:
1. This paper well-positioned its focus on the TTFT optimization of long-context LLM inference with a clear presentation of the walltime proportion of the main modules in each decoder layer during prefilling to show its motivation.
2. The FTP is simple and technically sound. Experiments on the LongBench dataset showed its effectiveness in achieving speedups in TTFT with less performance loss than baselines.
3. In the ablation study, FTP is also evaluated on larger models, which could be more practical and adopted for broad use.

weaknesses:
1. There is a major performance loss on Code Completion Tasks with LLama3 model compared to baselines in Table 2. How the authors explain this phenomenon? This may impact the broader use of FTP and its effectiveness on different models.
2. In Table 2, the average performance of PyramidInfer is higher than FTP. The reported high TTFT for Pyramidinfer is not fair due to the pytorch-implemented attention. In this way, it seems hard to tell the effectiveness and efficiency of FTP compared to Pyramidinfer.

questions:
1. How do the authors obtain the attention score for token pruning in FTP based on FlashAttention op?
2. Since the number of tokens with high attention scores varies among different layers, how to determine the best reserved ratio for each layer's FFN? Do the increasing number of pivotal tokens in Attention reflect that we should also reserve more tokens for FFN?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730557843922

thread_type:
official_review

<thread_object>
creator: reviewer_cws7

<content>
summary:
This paper proposes a token pruning method called FTP to accelerate the prefill stage inference of long-context LLMs. The authors first demonstrate the importance of accelerating the prefill stage, noting that it accounts for over 60% of the inference time (it should be noted that this is the case when the number of tokens to be decoded is around 10~20). They identify the FFN module as the most time-consuming component. Based on this, FTP selects non-critical tokens from the input sequence at each layer using attention scores and omits these tokens from the FFN computation. In the experimental section, the authors use results from Longbench to demonstrate that FTP can accelerate the prefill stage while maintaining a certain level of performance.

strengths:
- The idea of accelerating the prefill stage by pruning tokens in the FFN layer is innovative.
- The paper provides a clear explanation of the motivation and methodology.

weaknesses:
- The experimental section is relatively weak: 1) The benchmarks tested are limited, with only Longbench included. It is recommended to add another benchmark, such as BigBench, ZeroScrolls, or L-Eval; 2) The baselines tested are limited, with only PyramidInfer and LLMLingua2 included. The authors should add more baselines or explain why only these two baselines were selected.
- Some parts are not clearly written. For example, in lines 232-236, it is unclear how conclusions 2 and 3 are derived from Figure 5.

questions:
- Typo: line 267: "and the dataset" should be "dataset".
- If the authors can address the issues raised in the weaknesses section, I would consider increasing the score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626509198

thread_type:
official_review

<thread_object>
creator: reviewer_qppz

<content>
summary:
1. The paper introduces a method called FFN Token Pruning (FTP) to optimize the inference phase of large language models (LLMs) or long-context processing by selectively pruning tokens during FFN computations.
2. FTP dynamically prunes tokens with low attention weights.
3. FTP achieves faster inference speeds, showing up to a 1.39x improvement on models like Qwen2-7B-Instruct.

strengths:
1. The method is straightforward and easy to apply.
2. The experiments demonstrate strong performance results.
3. The experimental setup is well-explained and easy to understand.
4. The analysis of experimental results is detailed and thorough.

weaknesses:
1. FTP involves a considerable computational cost in calculating the Sum Attention Score, but this issue is not discussed by the authors.
2. FTP applies KVCache Eviction (Compression) methods to FFN token pruning without justification or evidence that the Attention Score is also effective for FFN token pruning. Additionally, FTP lacks a specific design tailored to FFN computations.
3. In Figure 5, FTP appears to retain only the beginning and end tokens, discarding most intermediate information. However, previous studies have shown that losing intermediate information can lead to sampling bias and performance loss [1].
4.  Similar research that also performs FFN token pruning, such as MoD [2], is not discussed in the paper. The choice of comparison methods is limited; for example, LongBench compares only with LLMLingua2 and lacks comparison with methods like KVCache Eviction, such as H2O [3] and NACL [1]. When compared with PyramidInfer, the setup lacks fair consistency, as the FlashAttention setting was not standardized.


---

References:  
[1] Y. Chen et al., “NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time.” 

[2] D. Raposo, S. Ritter, B. Richards, T. Lillicrap, P. C. Humphreys, and A. Santoro, “Mixture-of-Depths: Dynamically allocating compute in transformer-based language models.” 

[3] Z. Zhang et al., “H₂O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.”

questions:
1. Could you explain why the Attention Score is effective for FFN token pruning?
2. Could you discuss the differences between this work and other dynamic layer-skipping methods?
3. Please analyze the eviction effects of FTP and assess whether it may result in a loss of intermediate information.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732277790208

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to common questions

comment:
**Q1. Comparison with Pyramidinfer**

+ PyramidInfer achieves considerable acceleration only with PyTorch-implemented attention, since it requires a large part of attention weights for token selection.
+ FlashAttention cannot return attention weights.
+ The PyTorch-implemented attention is much slower than FlashAttention, and is not a practical choice for modern LLM inference. Table 2 in our paper shows that even the optimized inference (PyramidInfer) with PyTorch-implemented attention is obviously slower than the FlashAttention baseline.
+ PyramidInfer can be implemented with FlashAttention by re-calculating the necessary part of the attention weights for token selection. However, the re-calculation is computationally expensive.
+ We re-implement PyramidInfer with **FlashAttention** and conduct experiments with Llama3-8B-Instruct on LongBench. As shown in the table below, PyramidInfer can hardly reduce TTFT when implemented with FlashAttention.
+ We will supplement the experimental results below to Table 2 to avoid confusion.

**`Accuracy`**
| | Single-document QA | Multi-Document QA | Summarization | Few-shot Learning | Synthetic Task | Code Completion |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | 
| Baseline | 37.20 | 36.85 | 26.80 | 69.33 | 37.00 | 55.17 |
| PyramidInfer | 32.05 | 33.43 | 24.21 | 66.37 | 35.50 | 55.24 |
| FTP | 36.06 | 34.85 | 24.41 | 67.55 | 36.00 | 35.91 |

**`TTFT SpeedUp`**
| | Single-document QA | Multi-Document QA | Summarization | Few-shot Learning | Synthetic Task | Code Completion |
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | 
| Baseline | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 |
| PyramidInfer | 1.02 | 0.97 | 1.06 | 0.97 | 0.94 | 1.10 |
| FTP | 1.20 | 1.21 | 1.19 | 1.21 | 1.25 | 1.19 |

**Q2. The reason for selecting these two baselines**

In this work, we focus on accelerating the prefilling stage for long-context LLM inference without post-training. Since there are few works on this setting, we only choose these two baselines: 1) PyramidInfer, a pioneer work in compressing KV Cache in both the prefilling stage and the decoding stage; 2) LLMLingua2, a SOTA prompt compression method which can also accelerate the prefilling stage by reducing the length of the input prompt.

Other KV Cache Eviction methods, such as H2O [R1], NACL [R2], and SnapKV [R3], mainly focus on the reduction of memory footprint and the inference time in the decoding stage, which brings zero / negative acceleration for the prefilling stage and thus has no impact on TTFT or even increases TTFT. These works are summarized in the "Related Works" section in our paper and are not suitable for comparison in the experiment section.

**Q3. Limited benchmarks**

LongBench is a multi-task benchmark for a comprehensive assessment of the long context understanding capabilities of LLMs. Our method has been evaluated on all 16 datasets in LongBench across 6 key long-text application scenarios: single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks and code completion.

We are running more experiments on the L-Eval [R4] benchmark, which is another long-context language model benchmark. We will supplement the results once the experiments are finished.

---
References
> R1. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., ... & Chen, B. (2023). H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 34661-34710.

> R2. Chen, Y., Wang, G., Shang, J., Cui, S., Zhang, Z., Liu, T., ... & Wu, H. (2024). Nacl: A general and effective kv cache eviction framework for llms at inference time. arXiv preprint arXiv:2408.03675.

> R3. Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., ... & Chen, D. (2024). Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469.

> R4. An, C., Gong, S., Zhong, M., Zhao, X., Li, M., Zhang, J., ... & Qiu, X. (2023). L-eval: Instituting standardized evaluation for long context language models. arXiv preprint arXiv:2307.11088.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732805104899

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Modifications of the revised manuscript

comment:
Dear Reviewers,

We sincerely appreciate your constructive comments and suggestions. We have made several modifications to our manuscript. Here are a summary of them:

Major modifications
- We merged the results for PyramidInfer to Table 1, and supplemented the explanations for the attention implementations in Section 4.1 (Implementation Details) and Section 4.3 (i.e., Comparison with State-of-the-art Method).
- New results for "L-Eval" and "Needle-in-a-Haystack" have been included in Section 6.1.1 and Section 6.1.2 respectively.
- We merged the results in A1 of reviewer QpPZ into Table 3 to analyse the computational cost introduced by FTP.

Minor modifications
- We incorporated more related works on layer skipping and token pruning out of LLMs.
- We fixed the typo in Line 267 as pointed out.

Thank you for your valuable feedback and for the time you have invested in reviewing our work.

Best Regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a token pruning method called FFN Token Pruning that accelerates the prefilling stage of long-context LLM inference. FTP dynamically determines token importance based on attention scores and prunes non-critical tokens before they are processed by the FFN. While generally effective, FTP shows significant performance degradation on code completion tasks with the Llama3 model in the initial experiments, raising concerns about the method's robustness and generalizability across different models and tasks. Other relevant methods (especially on dynamic layer skipping) could have been included for a more comprehensive comparison. Given the accuracy trade-offs in some cases, the practical benefit of this speedup might not be substantial enough to justify widespread adoption, especially considering the complexity it adds to the inference pipeline. The paper could be significantly strengthened by addressing these weaknesses in future work.