Balanced Token Pruning: Accelerating Vision Language Models Beyond Local Optimization

Thank you for your insightful feedback. We appreciate the opportunity to address the concerns you raised. Below, we provide our responses to the problems:

---

Q1: Spatial Initialization: Could you explain why you chose Manhattan distance instead of Euclidean distance?
R1: Thank you for your constructive feedback. We did not apply any special design when choosing the distance metric. To evaluate the impact of this choice, we conducted an ablation study comparing two different distance metrics using llava-v1.5. The results are summarized in the table below:

We found that Euclidean distance indeed performs better for spatial initialization. We will incorporate this finding into Section 5.3 of the revised paper and update our implementation to adopt Euclidean distance in the next version.

---

Q2: Pruning Layer Selection: In Qwen2.5-VL, significant changes in image token representations are observed in layers 10 and 24, yet relatively fewer image tokens are selected in those layers. This appears inconsistent with the assumption that image tokens receive more attention in layers with significant feature changes. LLaVA-v1.5 shows a similar pattern—although layers 8–12 exhibit large representational shifts, more image tokens are selected in layers 17, 19, and 20. Further clarification would be helpful.
R2: Thank you very much for your question—we appreciate the opportunity to clarify this point. Firstly, once the pruning layers are determined, our method retains more tokens in earlier layers than in deeper ones. As described in Section 4.2, the token set preserved in deeper layers is a subset of that from earlier layers. This design is consistent with prior work [1], and therefore, it is not the case that more image tokens are retained in deeper layers.

Secondly, our hypothesis in Section 4.3 is: "LVLMs tend to allocate more attention to image tokens in layers following those where the representations of image tokens undergo significant changes," which is different from "image tokens receive more attention in layers with significant feature changes."

To illustrate, we only prune a selected subset of layers. The tokens retained in a pruned layer are propagated through subsequent layers. Our method favors diversity in early layers and attention-based selection in deeper layers. For example, in LLaVA-v1.5, we retain more tokens in layers 8–12 and fewer in layers 17, 19, and 20. If we have misunderstood your concern, we would be happy to engage in further discussion.

---

Q3: Attention Optimization: The method uses the last token in the input prompt $X_T$ to compute importance scores. It is unclear whether relying on a single token is sufficient to capture the full contextual relevance of the prompt. Further explanation of this choice would improve the understanding of its effectiveness.
R3: Thank you very much for your helpful suggestion. We agree that further exploring the computation of the importance score is valuable.

To this end, we compared with two more approaches: (1) computing the importance score by averaging over all text tokens following the image tokens, and (2) the method used in [2], which first computes image-text similarity and then selects the most similar text tokens to compute the importance score. The results are shown in the table below:

We believe that the last token in the input prompt is a suitable choice for computing the importance score because it is typically decoded as the first output token during the decoding stage. This allows it to effectively capture the model's focus. We will include this discussion in the ablation study section of the revised paper.

---

Q4: The values $\lambda_i$ are critical hyperparameters that influence the performance of the proposed framework. It would be beneficial to provide a more concrete rule or conduct an ablation study on the impact of $\lambda_i$ settings.
R4: Thank you for your constructive comment. In our experiments, the $\lambda_i$ parameter was manually selected. For a given model, the $\lambda_i$ value remains fixed across different benchmarks. We only set $\lambda_i$ values at pruning layers. The table below summarizes the $\lambda_i$ values used for each model at pruning layers:

Once the pruned layers are determined, the $\lambda$ parameter is also fixed. To further investigate the sensitivity of our method to the choice of $\lambda$, we conduct an ablation study on LLaVA-v1.5-7B by varying $\lambda$ at different stages. The results are shown below:
Experiments in the shallow layer

Experiments in the middle layer

Experiments in the deep layer

We found that using a smaller $\lambda$ value in earlier layers encourages the pruning strategy to focus more on the global objective, while employing a larger $\lambda$ value in deeper layers helps preserve local outputs. This observation further supports the effectiveness of our proposed method. We will include this discussion in the ablation study section of the revised paper.

---

Q5: Inference Latency and Spatial Initialization In Table 4, when the attention rebalance (RA) module is used, adding spatial initialization significantly improves inference latency. However, in cases where RA is not used, the latency difference with or without spatial initialization becomes marginal. An explanation for this discrepancy would be appreciated.
R5: We sincerely apologize for the confusion caused by the writing error. After reviewing our experiment logs and re-running the experiment across three datasets, we found that the averaged latency in the last row of Table 4 should be 0.231s. We will correct this in the revised version of the paper. Thank you again for your careful review—your feedback has been very helpful to us.

---

Q6: In Section 4.1, the authors present the assumption that when values are similar, selecting image tokens with high importance scores can effectively reduce the difference in outputs before and after pruning. This lack of explanation makes the assumption difficult to evaluate.
R6: We appreciate the opportunity to clarify this point. As shown in Equation (1), the hidden states at each pruned layer are primarily influenced by the outputs of both the attention and MLP modules. Since the MLP parameters remain fixed and its inputs include the outputs from the attention module, we attribute the majority of the observed changes in the hidden states to variations in the output of the attention module.

Further, according to Equation (2), the attention output is determined by the attention scores and the value matrix $V_l$. As detailed in Appendix A.1, we observe that the value vectors across different image tokens exhibit similar magnitudes. Therefore, maximizing the sum of attention scores effectively enhances the similarity of local outputs, which aligns with the assumption.

---

Q7: The equation for $Attn^l$ is not explicitly defined in the paper. Including a clear mathematical formulation would help readers better understand how attention scores are computed at each layer.
R7: Thank you very much for your suggestion. We denote $\mathrm{Attn}^l$ as the attention module at layer $l$, with its formulation defined as:

$\mathrm{Attn}^l = softmax({\frac {Q_lK_l^T+M}{\sqrt d_k}V_l}).$

where M is the casual mask. We will include this definition in Section 3 Preliminaries of the next version of the paper.

[1] PyramidDrop: Accelerating Your Large Vision-Language Models via Pyramid Visual Redundancy Reduction. CVPR2025.
[2] SparseVLM: Visual Token Sparsification for Efficient Vision-Language Model Inference. ICML2025.

Thank you for your thorough review and for the positive feedback on our work. We appreciate the opportunity to address your questions. Below, we provide detailed responses to the concerns you raised.

---

Q1: Line 41 - 42: Is this phenomenon of different image tokens being attended by the text tokens consistently observed across the dataset? It seems counter-intuitive to me. At least in vision transformers, we do observe the same tokens being attended to by other tokens throughout the layers.
R1: We are very pleased to have the opportunity to discuss this phenomenon with the reviewer. Yes, we have observed the pattern shown in Figure 1 across different types of datasets, including POPE, GQA, and COCO.

We believe that the reviewer’s observation and our findings are not contradictory. As shown in Figure 1, certain image tokens are indeed attended by multiple layers. Our analysis reveals that the overlap of highly attended image tokens is high between adjacent layers, but this overlap significantly decreases between non-adjacent layers. In the table below, we report the overlap of the top 144 most attended image tokens across different layers:

We observe that the overlap of highly attended image tokens is relatively high between adjacent layers, while it significantly decreases between layers that are farther apart. This suggests that adjacent layers tend to focus on similar regions of the input image, whereas distant layers shift their attention to different regions. Moreover, this trend becomes even more pronounced when computing the overlap using a smaller subset of top-attended image tokens.

---

Q2: Lines 44 - 56: Why was the hidden state of only last word considered for this experiment? What dataset and model was used for this experiment? Overall, I think it would be insightful to add some more explanation (if any) on the observations shown in figure 1 and 2. But this idea of creating a new local-global objective is quite interesting.
R2: We use coco dataset to conduct this experiment with llava-v1.5-7b. We have also observed similar patterns across different datasets and model architectures, suggesting that this phenomenon is consistent and not specific to a particular setting. As shown in Figure 1, diversity-based pruning tends to preserve image tokens that are important for deeper layers, while attention-based pruning primarily retains tokens that are important for the current layer. This discrepancy leads to the results observed in Figure 2.

In Figure 2, our goal is to investigate the impact of pruning during the prefilling stage on the model’s output. Due to the autoregressive nature of large language models, the last token in the prefilling stage is decoded as the first token in the decoding stage. It is also the only “model output” that we can access during the prefilling phase. Therefore, we select the hidden state of the last token in the prefilling stage for analysis.

Effect of Different Pruning Strategies on the Full Output: However, in captioning tasks, the model generates multiple tokens in an autoregressive manner. Therefore, we further investigate how output tokens at different positions are affected by different pruning strategies. Specifically, we analyzed the impact on the 1st, 4th, and 20th generated tokens with llava-v1.5-7b. The results are summarized below:

We observe that tokens generated at later positions are more influenced by previously autoregressively generated tokens, resulting in lower similarity to their original hidden states. However, across all positions, the impact trends of the two pruning strategies remain consistent: attention-based pruning preserves local information, while diversity-based pruning better maintains global information.
Further Explanation of Figure 1: We attempt to provide a deeper explanation of the phenomenon shown in Figure 1. Since the multimodal large language model cannot access future questions during image processing due to the causal mask, the image token processing pipeline remains the same across different types of questions. This implies that different layers may encode different aspects of image information. As illustrated in Figure 4, image tokens gradually align to the language space in deeper layers. Moreover, MLLMs tend to assign uniform attention scores to all image tokens in early layers, and only focus on semantically relevant different regions in deeper layers. We plan to further explore this interesting phenomenon in future work.

---

Q3: Having different number of tokens pruned for different input sample breaks the regularity in processing batched input. Does it not deteriorate the throughput?
R3: Thank you for the insightful question. Our method supports batch inference for models that encode images into a fixed number of image tokens. This is because, under fixed pruning layers and pruning ratios, the number of retained tokens before each layer remains consistent across samples, avoiding dimensional mismatches during inference.

However, for models such as LLaVA-v1.6 and Qwen-2.5-VL that employ dynamic resolution and encode different images into varying numbers of image tokens, our approach does pose challenges for batch inference. This presents an interesting trade-off. For example, one potential solution is to determine the number of image tokens to discard based on the sample with the fewest tokens in a batch, thereby ensuring input dimensional consistency across the batch. We plan to explore this further in future work, investigating whether pruning and token reuse strategies can be adapted to support efficient batch inference under such dynamic settings.

Thanks for the detailed rebuttal. I appreciate the effort you've put into addressing each of the concerns.

On Q1, the additional analysis on token overlap across adjacent vs. distant layers is insightful. The overlap table adds strong empirical support to the explanation.

Regarding Q2, thanks for clarifying the motivation behind using the last token’s hidden state in the prefilling stage. The contrast between attention-based and diversity-based pruning and how they affect local vs. global information is demonstrated and adds depth to your findings.

For Q3, I appreciate the discussion of the trade-offs involved with dynamic token counts during batch processing. It’s good to know that your method still supports batch inference under fixed conditions, and I agree that this is a meaningful direction to explore further.

Overall, the clarifications and additional results make your work stronger. I have no further concerns.

We’re very glad to hear your positive feedback that 1) our clarifications and additional results have strengthened the paper and 2) you currently have no further concerns.

We would be grateful if you could consider raising your score. Thanks!

Thank you again for your valuable time and thoughtful feedback!

We sincerely thank you for your careful review and insightful feedback. We appreciate the opportunity to address the concerns you raised. Below, we provide our responses to the weaknesses:

---

W1: The paper makes completely unclear statements when describing how the core balance factor $\lambda$ changes with the depth of the network.
R1: Thank you for pointing out the unclear description of the balance factor $\lambda$ in Section 5.3. We acknowledge a typographical error, which incorrectly states: “in the early layers, we use a larger $\lambda$ value to focus more on global information, while in the deeper layers, we use a smaller lambda to emphasize local details.” In our method, $\lambda$ increases with network depth. This strategy is correctly reflected in Formula (7), Figure 3, and Section 4.2, and has shown strong empirical performance. We will revise the text for consistency and thank you again for your helpful feedback.

---

W2: Heuristic Treatment of "Global Objective": The final optimized objective function (Formula (7)) is a weighted sum of attention and diversity scores, rather than a direct optimization of Formula (6). This jump from the true objective to the proxy objective is based on heuristics and lacks stronger theoretical support.
R2: Thank you for pointing out this issue which is critical for improving the quality of our paper. We will provide both theoretical and empirical justifications to explain why optimizing Formula (7) effectively leads to the optimization of Formula (6). Suppose token pruning is performed at layers $l_1$ < $l_2$ < $l_3$. Since the tokens retained in deeper layers are a subset of those retained in earlier layers, the preserved token sets satisfy the relationship: $P_3 \subseteq P_2 \subseteq P_1$.

Relation between attention scores and global objective: The global objective aims to minimize the discrepancy between the hidden states of pruned layers $l_1, l_2, l_3$ and those from the original (unpruned) model. According to Formula (1), the hidden state at each pruned layer primarily depends on the outputs of both the attention and MLP modules. Since the MLP parameters are fixed and its input includes the output of the attention module, we attribute the changes in the pruned layer's hidden states primarily to variations in the attention module's output. According to Formula (2), the attention module's output is affected by attention score and $V_l$. In Appendix A.1, we observe that different image tokens exhibit similar magnitudes in their value vectors (V). Therefore, maximizing the sum of attention scores contributes to improving the local output similarity, as reflected in the first term of Formula (7). This helps reduce the deviation between the outputs of shadow layers ($l_1$) and the original outputs, thereby contributing to the optimization of the global objective.

Relation between diversity scores and global objective: As shown in Figure 1, different layers attend to different regions of the image tokens—that is, the regions with high attention scores vary across layers. Therefore, selecting only high-attention tokens in the shallow layers may lead to local output deviations in subsequent layers. As the diversity objective is strengthened, the overlap between the tokens retained in deeper layers $l_2,l_3$ and the optimal retained token set $P^*$ increases: $|P^* \cap P^{attn}| < |P^* \cap P^{div}|$. This helps reduce the deviation between the outputs of deeper layers ($l_2,l_3$) and the original outputs, thereby contributing to the optimization of the global objective.

We also observe that the attention from output tokens to image tokens is sparse. To quantify this, we compute the proportion of total attention assigned to image tokens that is captured by the top-k% highest-attention image tokens:

We can see that as the number of retained tokens increases, the gain from optimizing the first term of Formula (7) diminishes, while the benefit of promoting diversity to support subsequent layers increases. In summary, we believe that optimizing Equation (7) serves as a reasonable approximation to optimizing Equation (6). We will include this discussion in Section 4.2 of the revised paper. A more rigorous theoretical justification will be explored in future work.

---

W3: How the schedule across different stages in the final experiment is determined? Is this strategy fixed or does it need to be adjusted for different models? This makes the reproducibility of the method and its sensitivity to hyperparameters unclear.
R3: Thank you for your constructive comment. In our experiments, the $\lambda$ parameter was manually selected. Its value varies across models due to differences in decoding depth and the patterns of image token variation. However, for a given model, the $\lambda$ value remains fixed across different benchmarks. The table below summarizes the $\lambda$ values used for each model at pruning layers:

We will include this information in the experimental setup section of Appendix A.2. Once the pruned layers are determined, the $\lambda$ parameter is also fixed. We included preliminary ablation results in Figure 6. To further investigate the sensitivity of our method to the choice of $\lambda$, we conduct an ablation study on LLaVA-v1.5-7B by varying $\lambda$ at different stages. The $\lambda$ values are used in pruning layers $(l_1,l_2,l_3)$. The results are shown below:
Experiments in the shallow layer

Experiments in the middle layer

Experiments in the deep layer

We observe that focusing too early on attention-based pruning in shallow or middle layers leads to performance degradation. Similarly, emphasizing diversity-based pruning in the deeper layers also results in a drop in performance. We will include this discussion in the ablation study section of the revised paper.

Thank you for your insightful feedback. We appreciate the opportunity to address the concerns you raised. Below, we provide our responses to the weaknesses:

---

W1: The method relies on a calibration set to determine pruning layers and balance factors. How sensitive the method is to the size and content of this set?
R1: Thank you for your constructive comment. In our approach, the selection of pruned layers is guided by the calibration set. The trade-off parameter $\lambda$ is a fixed hyperparameter and is not determined by the calibration set. Our method is insensitive to the content and size of the calibration set.

Insensitive to the text content: According to Equation (2) in the paper, in multimodal large language models, M denotes the causal mask, which constrains each token to attend only to preceding tokens. In our input format, image tokens always precede the text prompts (i.e., the input follows the structure: prefix + image + question). As a result, the model processes the image tokens before it receives the specific question which is unrelated to the input question. To validate our hypothesis, we posed different types of questions on the same image. We then conducted the experiment presented in Figure 4 using LLaVA-v1.5. We computed the number of image tokens whose cosine similarity between adjacent layers falls below 0.93. The resulting trends are shown as follows:

We observe that varying the question type for the same image does not lead to significant differences in the results. In the following analysis, we investigate how image content and the size of the calibration set affect our method.

Insensitive to the image content and size: We expanded the calibration set by incorporating images from multiple datasets, including GQA, V*Bench[1], and SQA. We then repeated the experiment shown in Figure 4 using calibration set sizes of 64, 128, and 256. Specifically, we computed the number of image tokens whose cosine similarity between adjacent layers falls below 0.93 using LLaVA-v1.5. The results are presented below:

We observe that the patterns remain consistent across different image contents and calibration set sizes, indicating that these factors do not affect the behavior of our method. We will include this additional analysis in the Ablation Study section (Section 5.3) of the revised paper.

---

W2: It is not clear how Figure 2 is depicted, e.g., which data and features are used to measure similarity? It would be good to provide more details.
R2: Thank you for your constructive feedback. We conducted experiments on the captioning task. In Figure 2, our goal is to investigate the impact of pruning during the prefilling stage on the model’s output. In multimodal large language models, the first token generated during the decoding stage is directly influenced by the final token from the prefilling stage. This final token is also the only point during prefilling where the model produces an output. Therefore, we evaluate the effect of different pruning strategies by analyzing their impact on the hidden state of the last token in the prefilling stage.

Effect of Different Pruning Strategies on the Full Output: In the captioning task, multimodal large language models generate not just a single token, but a sequence of tokens in an autoregressive manner. Therefore, we further investigated how different pruning strategies affect tokens at various positions in the output sequence. Specifically, we analyzed the impact on the 1st, 4th, and 20th tokens. We use llava-v1.5-7b to conduct this experiment. The results are summarized below:

We observe that tokens appearing later in the output sequence are influenced by previously generated tokens, leading to a significant drop in similarity with the corresponding tokens from the original (non-pruned) output. Moreover, tokens at different positions exhibit similar patterns under both pruning strategies: attention based methods focus on local goal and diversity based methods focus on global goal.These findings further support the validity of our observations. We will include this discussion in the Introduction section of the revised manuscript.

[1] V*: Guided Visual Search as a Core Mechanism in Multimodal LLMs.