Title: Attribution-Driven Adaptive Token Pruning for Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: Token Pruning

Attribution-Driven

Adaptive Inference

Integrated Gradients.
Abstract: Transformers have been widely adopted in natural language processing, computer vision, and other domains due to their exceptional performance across a variety of tasks. However, the computational cost of Transformers is prohibitively high, particularly when handling long input sequences, significantly increasing both training and inference time. Although various token pruning methods have been proposed to reduce the computational burden of Transformers, most approaches overlook critical differences in sequences in terms of length and complexity, leading to suboptimal compression efficiency.

In this paper, we propose AD-TP, an Attribution-Driven Adaptive Token Pruning method designed to retain only the most informative tokens. We analyze the performance of using accumulated attention values to measure token importance and find that attention values do not accurately reflect the actual contribution of each token to text understanding. Additionally, we observe significant variations in the length and complexity of different sequences within the dataset. Based on these insights, we adopt Integrated Gradients to evaluate token importance and introduce a lightweight adaptive token retainer module that dynamically generates pruning configurations for each input sequence. In addition, we incorporate both teacher supervision and self-supervised learning objectives to enhance the training efficiency, accuracy, and robustness of the model.

Experiments conducted on GLUE, SQuAD, and 20News demonstrate that AD-TP outperforms state-of-the-art token pruning and model compression methods in both accuracy and computational efficiency. On GLUE, AD-TP reduces FLOPs by an average of 7.8× while improving performance by 0.6%.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces AD-TP, a novel method for token pruning in Transformer models designed to enhance computational efficiency. The authors argue that existing methods are suboptimal because they often rely on attention scores to gauge token importance, which may not correlate well with a token's actual contribution, and they use static pruning ratios that don't adapt to the varying complexity of different input sequences. To address this, AD-TP makes two main contributions. First, it uses Integrated Gradients (IG), an attribution method, to obtain a more accurate measure of each token's importance to the final prediction. Second, it introduces a lightweight \"Adaptive Token Retainer\" module that dynamically predicts an appropriate token retention ratio for each individual input sequence. This retainer is trained using a combination of teacher-student knowledge distillation (on both logits and IG-based attribution scores) and a self-supervision objective that is gradually phased in during training. Experiments on GLUE, SQUAD, and 20News datasets show that AD-TP significantly reduces computational costs (measured in FLOPs) while maintaining or even improving model accuracy, outperforming existing state-of-the-art token pruning techniques.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe paper's core ideas (use of Integrated Gradients (IG) for supervising token importance) are novel and well-justified. The concept of an adaptive pruning ratio, dynamically determined for each sequence, is a key strength. The paper effectively argues that a one-size-fits-all pruning strategy is inefficient, given the wide variance in sequence length and complexity.\n\nThe architecture of the Adaptive Token Retainer, split into a ratio predictor (ARP) and a saliency predictor (TSP), is intuitive and addresses the \"how many\" and \"which\" questions of pruning separately.\n\nThe paper presents extensive experiments that demonstrate the effectiveness of AD-TP.\n\nThe paper is very well-written and organized.\n\nWeaknesses:\n\nThe proposed method uses a teacher model and computes Integrated Gradients during training. This process is likely to be more computationally expensive and time-consuming than training methods that rely on attention scores, which are obtained almost for free during the forward pass. A discussion or quantitative analysis of this training-time overhead would provide a more complete picture of the method's trade-offs.\n\nAs the authors acknowledge in the limitations, pruning is only applied to the input embeddings. This means the quadratic complexity of self-attention is applied to the reduced sequence in all subsequent layers. While this is a valid and effective approach, it misses the potential for even greater savings from layer-wise pruning.\n\nQUESTIONS:\nPlease see the weaknesses described above.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces AD-TP (Attribution-Driven Adaptive Token Pruning), a novel method for compressing Transformer models by dynamically pruning less informative tokens from input sequences. The work aims to address two key limitations of prior token pruning techniques: (1) the use of inaccurate proxies like attention scores for determining token importance, and (2) the reliance on static, one-size-fits-all pruning configurations. To achieve this, AD-TP makes two main contributions. First, it employs Integrated Gradients (IG), a feature attribution method, to compute a more faithful measure of each token's contribution to the final prediction. Second, it introduces a lightweight, Attribution-Aware Adaptive Token Retainer module that operates at the instance level. This module predicts an optimal token retention ratio for each specific input sequence and learns to approximate the expensive IG scores for efficient inference. The entire system is trained via a teacher-student distillation framework, where a larger teacher model provides supervision for both the classification task and the IG-based attribution scores.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- Quality & Originality: Principled Saliency Metric - The core innovation of using Integrated Gradients (IG) instead of attention scores as the basis for pruning is a significant and well-justified step forward. The authors correctly argue that attention weights are not designed to measure a token's contribution to the final output. By grounding the pruning decision in a more theoretically sound attribution method like IG, the method directly links pruning to the model's task-specific reasoning. This is a novel and principled approach that, as the results show, preserves accuracy more effectively.\n\n- Significance & Originality: Instance-Level Adaptability - The introduction of an adaptive module that predicts the pruning ratio for each input is a crucial advantage. This is particularly relevant in heterogeneous data environments, such as those found in Federated Learning (FL). In an FL setting, client data is often non-IID; some clients may have simple, highly prunable data, while others have complex data requiring more tokens. AD-TP's ability to adapt its compression rate on-the-fly for each sample is a major improvement over static methods that would be suboptimal across a diverse federation of clients.\n\n- Quality: Comprehensive Training Framework - The teacher-student training regime is methodologically robust. The combination of logits-based distillation and attribution-based distillation provides a rich supervisory signal, ensuring the student model learns both *what* to predict and *why* it is making that prediction. The scheduled weighting of the attribution loss is a clever detail that likely aids training stability.\n\n- Significance: Strong Empirical Results - The paper demonstrates impressive performance, substantially reducing FLOPs while maintaining high accuracy across a range of NLP tasks. The fact that AD-TP outperforms other well-established compression techniques like DistilBERT and specialized pruning methods like ToP underscores its effectiveness, especially on long-sequence tasks where the computational savings are most impactful.\n\nWeaknesses:\n\n- Quality & Significance: Prohibitive Training Complexity - The most significant weakness of AD-TP, particularly from a Federated Learning perspective, is its computationally expensive training process. The methodology requires a pre-trained teacher model, the calculation of IG scores (which involves multiple forward/backward passes per sample), and the optimization of a complex, multi-part loss function. This high resource requirement makes the method impractical for on-device training in a conventional FL setting. It appears better suited for centralized training followed by deployment to the edge, rather than end-to-end federated training.\n\n- Quality: Inference Overhead vs. Latency - While the paper focuses on FLOPs reduction, the Adaptive Token Retainer module itself adds parameters and a non-zero computational overhead at inference time. For very short sequences, this overhead could potentially negate the benefits of pruning. A more complete evaluation would include wall-clock latency benchmarks on hardware relevant to edge deployments (e.g., mobile CPUs) to demonstrate a practical speedup beyond theoretical FLOPs.\n\n- Clarity & Significance: Unexplored Federated Learning Implications - While the adaptability of AD-TP is a great fit for FL's data heterogeneity, the paper does not discuss how the method would be implemented in such a setting. It is unclear how the complex training would be federated, or what the communication overhead of aggregating the adaptive retainer module across clients would be. Without this analysis, its applicability to the FL domain remains speculative.\n\n- Originality: Architectural Specificity - The design of the Adaptive Retention Ratio Predictor (ARP) relies on the `[CLS]` token embedding as a global representation of the input. This is standard for BERT-like encoders but limits the method's direct applicability to other popular architectures, such as decoder-only models (e.g., GPT-family) or encoder-decoder models (e.g., T5), without non-trivial modifications.\n\n- Syntactic errors:\n1. Figure 4 Caption: \"Adaptive Aoken Aetainer\" -> \"Adaptive Token Retainer\"\n\nQUESTIONS:\n1.  Federated Implementation Feasibility: How do the authors envision AD-TP being trained in a federated environment?\n- Would this be a \"Federated Distillation\" setup where a central server holds the teacher model and clients only train the student and retainer?\n- If the entire process is federated, how would the expensive IG calculations be handled on resource-constrained client devices? What would be the communication cost of aggregating the adaptive retainer module in each round?\n\n2.  Generalization to Other Architectures: The ARP module relies on a `[CLS]` token. Have you considered how to adapt this mechanism for decoder-only models like the GPT family? Would a simple approach like averaging final-layer token embeddings suffice as a global sequence representation?\n\n3.  Sensitivity to IG Approximation: The IG calculation is approximated using `m=1` step for efficiency. How sensitive is the final model performance to this approximation? Is there a significant accuracy gain from using a larger `m`, and how does that affect the already high training cost?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes AD-TP (Attribution-Driven Adaptive Token Pruning), a method for efficiently compressing Transformer models by selectively pruning tokens based on their actual contribution to the model's prediction. Unlike prior approaches that rely on attention scores or static thresholds, AD-TP leverages Integrated Gradients (IG) to estimate token saliency more accurately and introduces a lightweight adaptive token retainer that dynamically adjusts the pruning ratio per input sequence and retain tokens based on token saliency.  Extensive experiments on GLUE, SQuAD, and 20News and more show that AD-TP significantly reduces computational cost (up to ~14× FLOPs reduction) while preserving or improving model accuracy, outperforming existing pruning and distillation methods.\n\nSTRENGTHS AND WEAKNESSES:\n**Quality**: This work aims to improve the token pruning method, and the idea of estimating sequence-level difficulty and assigning an optimal pruning ratio for each sequence is technically sound.\n\n**Clarity**: The motivation, method, and experiments are presented completely, but there is room for improvement. For example, the tables are too small, especially Table 1. Also, Figure 3 seems useless, and I don't quite understand what it's trying to convey. Additionally, it remains unclear what is determining the token saliency—is it the TSP module or the IG?\n\n**Significance**: How to dynamically determine the pruning ratio for each input sequence deserves further exploration by the community, and the ablation study on the impact of candidate list length seems interesting. However, it lacks an ablation on the ARP module alone, so it is hard to judge what its contribution looks like. It would be nice to see the ablation study for each individual component (IG, ARP, TSP).\n\n**Originality**: See **Significance**. In general, some interesting insights are provided, but the analysis is not adequate.\n\nQUESTIONS:\n1. It seems like both IG and TSP affect the token saliency. Why do we have two modules to determine the token saliency, and how do they collaborate to work together?\n\n2. Why do we use IG instead of other metrics like (masked) KL divergence? Can you show the benefit of IG over other metrics such as (masked) KL divergence?\n\n3. In Table 1 and Table 3, why is AD-TP6 always better than BERT6 (S) although the FLOPs are much lower? It seems like this is not always true for all token pruning methods, such as CoFi6. Any insights here?\n\n4. From Figure 2, it seems like there is a length distribution mismatch between the training set and the evaluation set. Does that have any impact on the ARP module, which is used to estimate the complexity of the sequence and determine the retention ratio? For example, there may be domain shift in terms of input length—how does this domain shift affect your method?\n\nIf my concerns about IG are properly addressed, I would consider raising my rating. Additional insights would be appreciated.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel token pruning method for transformer-based models, which utilizes token saliencies from intergrated gradients to reduce the overall input sequence length. The proposed method, AD-TP, is composed of two distinct parts, an Adaptive Retention Ratio Predictor that predicts how much of the input sequence should be preserved (based on complexity), and a Token Saliency Predictor which predicts and selects the actual tokens to keep vs discard. The proposed method is evaluated against several existing token pruning methods, using BERT-base as the model and GLUE, SQuAD and 20News as the datasets. The results show that the AD-TP matches or outperforms the best previous method, while pruning more tokens (i.e. higher efficiency gains). The various AD-TP submodules are trained using a student-teacher framework, where the submodules learn to predict both the retention ratio and IG saliency scores without needing the ground truth.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper studies an important (and timely) problem of high inference cost in transformer-based models. The proposed methodology combines quite a few existing pieces in a novel formulation, and presents an interesting framework to train the submodules of this framework. The evaluation is quite diverse in terms of dataset, but only relies on a single (small) model. At least on this setting, the proposed methodology performs quite well, matching or improving upon the previous best method at higher efficiency gains. \n\nHaving said this, there are some weaknesses that impact the quality and relevance of this work quite significantly:\n\n- The paper is a bit difficult to read in general, and the widespread use of similar abbreviations really makes it even harder. There are also several claims and terms used in the paper without any clear explanation or evidence of support. One instance of this is \"residual strategy\" in Section 3.1, which is not defined anywhere in the paper.\n- The experimental setting is quite dated; the model and the datasets are really from the \"previous\" era of NLP, as evidenced by the fact that the related work is all 4 or 5 years old. Its really a big miss to not have tried this method on a more recent LLM model, where the impact could be much greater. \n- Details on how AD-TP is trained are very sparse, apart from the loss function derivation. On what dataset is the training conducted on? How long is the training run for? This hinders reproducibility significantly, since without a trained AD-TP model the proposed methodology is incomplete.\n- It was also strange to see no evaluation of TRP; do the scores predicted by it match the IG scores (after the model has been trained)? Its a fundamental piece of the overall method, and given the importance of IG scores (as self-described in the paper), its important that this module mirror IG's scores well.\n\nOverall, the proposed methodology itself is sound, and the experimental results do improve upon existing methods to some extent, but it is unclear what the impact of the work is in the era of modern LLMs where classification is done differently, and a lot of tasks are dominated by generation (which is completely ignored in this paper). This, along with some key missing details sheds some doubt over the efficacy of the work.\n\nQUESTIONS:\n- Figure 3 is used to depict the complexity of the input sample? Can one just use a syntactic dependency parser as a light weight alternative to ARP?\n- Why not just use a regression model to predict the retention ratio P_r directly instead of using a candidate list L_p?\n- [1,2] are vision papers, not talking about why GPT/BERT are central to natural language processing. CoFi's citation is also [2], but it doesn't seem to be mentioned in that paper anywhere; are these references here by mistake?\n\nPlease also take a look at the questions mentioned in the *strengths and weaknesses* section.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}