TRIM-KV: Efficient KV Pruning for LVLMs

• The paper introduces TRIM-KV, an inference-phase pruning method that retains the most salient visual tokens to reduce the KV cache in large vision-language models.
• It leverages the sparsity of early cross-attention maps by aggregating head-wise attention scores to identify and select key visual tokens.
• Empirical results demonstrate latency reductions up to 26% and maintained accuracy with 50% token retention across standard vision-language benchmarks.

Cross-attention-based TRIM-KV is an inference-phase pruning method for reducing the key-value (KV) cache footprint of visual features in cross-attention-based large vision-LLMs (LVLMs). Developed in the context of models such as LLaMA-3.2-Vision, which employ dedicated cross-attention layers to aggregate image and text information, TRIM-KV exploits the empirical sparsity of early cross-attention maps to retain only the most salient visual tokens, thereby significantly reducing memory and computational costs while maintaining competitive performance across standardized benchmarks (Lee et al., 1 Apr 2025).

1. KV Cache Bottleneck in Cross-Attended LVLMs

In cross-attention-based LVLMs, images are encoded into a set of $n_k$ visual tokens, with each token mapped via linear projections to key ($K$) and value ($V$) representations of dimension $d$. In contrast to self-attention layers on textual data, where the number of tokens $n$ is modest, the number of image tokens $n_k$ can be significantly larger, especially for high-resolution images (e.g., $n_k\approx1600$ for 384$\times$384 images, versus $n\approx128$ text tokens). The KV cache required to store these representations is therefore dominated by visual features. The memory footprint is

$$M_{\text{KV}} = \text{batch\_size} \times 2 \times n_k \times d \times \text{bytes\_per\_float}$$

where the factor of 2 corresponds to storing both keys and values. For LLaMA-3.2-Vision-11B, the cross-attention KV cache ($K$0) is approximately 12.5 times larger than the self-attention KV cache ($K$1) on text at these settings. Thus, unmitigated visual KV caching poses a principal bottleneck during inference (Lee et al., 1 Apr 2025).

2. Mathematical Formulation of TRIM-KV Token Importance

Cross-attention operates by computing attention weights

$K$2

where $K$3 denotes queries from text tokens, and $K$4 the visual keys. The output is a weighted sum $K$5, with $K$6 the visual values. In TRIM-KV, token importance is measured head-wise in the first cross-attention layer, aggregating over all query tokens:

$K$7

for each attention head $K$8 and visual token $K$9, producing a set of token importance scores $V$0 for each head. This approach exploits the sparsity of attention maps, where a small subset of visual tokens typically dominate the aggregate attention distribution.

3. Trimming Algorithm and Pseudocode

The TRIM-KV pruning process operates as follows:

• For each attention head $V$1 ($V$2), select the top $V$3 visual tokens by their aggregate importance, where $V$4 for a chosen pruning ratio $V$5 (e.g., 0.5).
• The overall pruned token set $V$6 is the union across heads: $V$7.
• Only the $V$8 and $V$9 vectors belonging to $d$0 are kept; the remainder are discarded from the KV cache for all subsequent cross-attention layers.

Pseudocode: $n$1 This algorithm formalizes KV pruning based on actual model attention, in contrast to spatial or random token culling (Lee et al., 1 Apr 2025).

4. Integration with Standard LVLM Inference Pipelines

The TRIM-KV method is integrated purely at inference, requiring no re-training or model fine-tuning. The sequence is:

1. Encode the image to obtain visual features of shape $d$1 and project to $d$2 and $d$3.
2. At the first generation step ($d$4), compute attention weights $d$5 and run TRIM-KV to obtain $d$6.
3. Create pruned key and value caches: $d$7.
4. For all subsequence steps and cross-attention layers, use $d$8 as the cached memory.

This process, termed "plug-and-play," hinges on the empirically observed stability of cross-attention patterns after the first block, such that a one-time token selection remains stable across remaining layers and generation steps (Lee et al., 1 Apr 2025).

5. Empirical Impact and Benchmark Results

The introduction of TRIM-KV into LLaMA-3.2-Vision-11B results in significant memory and latency reductions during inference. On a batch of 32, total latency is reduced from approximately 4,000 ms (full tokens) to 3,165 ms (50% tokens, -19.7%) and 2,917 ms (40% tokens, -26%).

Accuracy is preserved across six vision-language benchmarks at up to 50.9% visual token retention (K_ratio=0.5), with <0.5% deviation from full-cache baselines. A summary of performance is as follows:

K_ratio	SEED-Image	MME	MMVP	LLaVA-Bench
1.00	72.6	1685.9	46.7	88.3
0.50	72.3	1687.3	47.3	88.1
0.40	72.1	1682.8	47.3	87.3

Ablation studies at K_ratio=0.5 demonstrate substantial superiority of TRIM-KV over random (e.g., SEED=67.0, LLaVA=83.2) and spatial (SEED=71.8, LLaVA=85.9) baselines for the same token budget. This underscores the relevance of attention-derived token selection for minimizing performance loss (Lee et al., 1 Apr 2025).

6. Scalability, Limitations, and Theoretical Considerations

The efficacy of TRIM-KV is strengthened as image resolution and the number of visual tokens $d$9 increase, due to a greater absolute reduction in memory and computation. The assumption underlying the method is a structured stability in early cross-attention maps, such that one-time token selection remains valid across all subsequent operations. Scenarios with dynamically shifting attention patterns may degrade the method's effectiveness.

TRIM-KV is designed specifically for cross-attention-based architectures ("Flamingo-style" LVLMs). Self-attention-only models require alternative approaches for efficient token pruning. Additionally, the optimal choice of $n$0 may require empirical validation for new tasks or domains, as excessive trimming could affect particular downstream behaviors (Lee et al., 1 Apr 2025).

7. Summary and Outlook

Cross-attention-based TRIM-KV achieves up to 50% reduction in visual KV cache, corresponding to 12–26% latency improvements, while maintaining state-of-the-art performance on multiple vision-language benchmarks without any model retraining or fine-tuning. By identifying and retaining only those visual tokens most salient to the initial cross-attention computation, the method addresses a central bottleneck in cross-attention-based LVLMs. This suggests that attention-based token pruning, particularly at the early layers, is a viable pathway for efficient multimodal inference at scale, especially as image and batch sizes increase (Lee et al., 1 Apr 2025).