LightKV: Prompt-Guided Vision Token Compression for Efficient LVLM Inference

Paper: 2605.00789 Authors: Xihao Chen, Yangyang Guo, Roger Zimmermann Categories: cs.CV, cs.AI, cs.LG

The Gap

Large Vision-Language Models (LVLMs) inherited KV caching from LLMs to speed up autoregressive decoding. But there’s a mismatch: LLMs process hundreds of text tokens during prefill, while LVLMs process thousands of vision tokens (a 336×336 image becomes 576 tokens). Existing compression methods treat vision tokens in isolation—they prune or merge based on visual similarity alone, ignoring what the text prompt actually asks for. This creates two problems: (1) massive GPU memory overhead from storing all those vision-token KV pairs, and (2) wasted computation on vision tokens irrelevant to the user’s question.

Prior work like FastV and LOOK-M compress vision tokens, but they’re prompt-agnostic. They decide what to keep before seeing the question. It’s like summarizing a textbook chapter before knowing what the exam question is.

Problem: Vision tokens bloat KV cache
         |
         v
Assumption: Not all vision tokens matter equally for a given prompt
         |
         v
Method: Cross-modal message passing guided by text
         |
         v
Evidence: 55% tokens retained, 50% cache reduction, performance preserved
         |
         v
Conclusion: Prompt-aware compression beats vision-only pruning

The Increment

One sentence: Before LightKV, vision-token compression was blind to the text prompt; after LightKV, the prompt actively guides which vision tokens to keep during prefill.

Core Mechanism

LightKV operates during the prefill stage, before any text generation begins. It takes the image tokens and text prompt as input, then progressively compresses the vision tokens through multiple transformer layers. At each layer, it runs cross-attention from text to vision (letting the prompt “query” the image), aggregates messages across vision tokens, and merges redundant ones based on their relevance to the prompt.

The compression happens in stages. Early layers keep more tokens to preserve fine-grained details. Deeper layers compress more aggressively as high-level semantics emerge. A learnable “compression schedule” determines how many tokens to keep at each layer—this schedule is trained end-to-end, not hand-tuned. The final compressed set of vision tokens gets cached, and the model proceeds with normal autoregressive decoding using this smaller KV cache.

Input: [Vision tokens] + [Text prompt]
         |
         v
Layer 1: Text queries vision via cross-attention
         Vision tokens exchange messages
         Keep 90% of tokens (light compression)
         |
         v
Layer 2: Repeat cross-attention + message passing
         Keep 75% of tokens
         |
         v
Layer N: Final compression
         Keep 55% of tokens
         |
         v
Output: Compressed vision tokens -> KV cache

Think of it like packing a suitcase for a specific trip. Vision-only compression is like deciding what clothes to pack before knowing your destination—you might bring a winter coat to a beach. LightKV reads the itinerary (text prompt) first. If the prompt asks “What color is the car?”, it keeps tokens around vehicles and discards background foliage. The packing happens in stages: first you eliminate obviously irrelevant items (swimsuit for a ski trip), then you make finer distinctions (which of three jackets to bring). The “compression schedule” is your packing strategy—how ruthlessly to cut at each decision point. By the time you close the suitcase, you’ve kept exactly what the trip requires, not a generic travel kit.

Key Concepts

• Cross-modality message passing: In standard transformers, tokens only talk to tokens of the same type (text-to-text, vision-to-vision). Here, text tokens send queries to vision tokens via cross-attention, asking “which parts of the image matter for my question?” Vision tokens then exchange information among themselves (self-attention) to aggregate context. This two-way conversation lets the prompt shape which vision features survive compression. Concrete example: prompt says “count the apples”, cross-attention highlights red circular regions, message passing groups nearby apple tokens, compression keeps one representative per cluster.

• Progressive compression schedule: Instead of compressing all at once, LightKV compresses gradually across layers. Early layers keep 90% of tokens, middle layers 75%, final layers 55%. Why? Early layers capture low-level features (edges, textures) that need spatial detail. Deeper layers extract high-level semantics (object categories, relationships) that tolerate more compression. The schedule is a learned parameter—the model figures out the optimal compression rate per layer during training. Think of it like image compression: you can’t throw away 50% of pixels immediately, but you can progressively downsample as you move from raw pixels to semantic features.

Framework Shift

Before (vision-only compression):        After (LightKV):

Image -> Vision Encoder                  Image -> Vision Encoder
         |                                        |
         v                                        v
    [All vision tokens]                      [All vision tokens]
         |                                        |
         v                                        +<-------+
    Prune by similarity                          |        |
    (no prompt info)                        Text Prompt   |
         |                                        |        |
         v                                        v        |
    [Compressed tokens]                   Cross-attention |
         |                                        |        |
         v                                        v        |
    KV Cache -> Decode                    Message passing |
                                                  |        |
                                                  v        |
                                          Compress layer N-+
                                                  |
                                                  v
                                          [Compressed tokens]
                                                  |
                                                  v
                                          KV Cache -> Decode

From blind pruning to guided compression, the core shift is making the text prompt an active participant in deciding which vision tokens matter.

Expert Assessment

Problem choice: Real gap. KV cache memory is a known bottleneck for deploying LVLMs, and vision tokens are the obvious culprit (10-100× more tokens than text). The insight that existing methods ignore the prompt is sharp—it’s obvious in hindsight but overlooked by prior work. This sits squarely in the “efficiency for deployment” track, which is hot right now.

Method maturity: Clever insight, solid execution. Cross-modal attention for compression is not new conceptually, but applying it progressively during prefill with a learned schedule is a nice twist. The method is more principled than brute-force pruning, but it’s not groundbreaking—it’s an engineering refinement of existing attention mechanisms. One concern: the learned compression schedule adds training overhead and might not transfer well across model families.

Experimental integrity: Baselines are fair (FastV, LOOK-M, token merging). Eight models, eight benchmarks—good coverage. The 55% retention rate is consistent across experiments, which builds confidence. However, the paper doesn’t report wall-clock latency, only FLOPs reduction. Memory savings are real, but does it actually speed up inference on real hardware? Also, no ablation on what happens if you compress more aggressively (say, 40% retention)—where’s the performance cliff?

Writing quality: The abstract and intro are crisp. The method section is dense but clear. The related work section is perfunctory—it lists prior methods but doesn’t deeply contrast their design choices. The results section front-loads tables without enough narrative to guide interpretation. If I were revising, I’d expand the ablation study (Section 4.3) to explore failure modes and add a “when does LightKV struggle?” subsection.

Verdict: weak accept — Solid incremental contribution with practical value, but not a paradigm shift. The prompt-aware angle is the main novelty; everything else is competent execution.

Takeaways

Steal the progressive compression pattern: Don’t compress all at once. If you’re building any multi-stage model (vision, audio, multimodal), consider compressing gradually as features move from low-level to high-level. The learned schedule idea transfers—let the model figure out the compression rate per stage rather than hand-tuning.

Cross-modal guidance generalizes: Anytime you have two modalities where one is expensive (vision, audio) and one is cheap (text), use the cheap one to guide compression of the expensive one. This applies beyond LVLMs: video understanding (text guides frame selection), audio transcription (text guides which audio segments to keep), retrieval (query guides document compression).

Benchmark on real hardware: The paper reports FLOPs and memory, but practitioners care about wall-clock time. If you’re doing efficiency work, measure end-to-end latency on target hardware (A100, H100, edge devices). FLOPs reduction doesn’t always translate to speedup due to memory bandwidth and kernel fusion.