Index-Preserving Lightweight Token Pruning for Efficient Document Understanding in Vision-Language Models

This paper proposes a lightweight token pruning framework that filters non-informative background regions and refines fragmented text areas in document images to significantly reduce computational costs in vision-language models while maintaining comparable accuracy.

The Gist

Imagine you have a massive, high-resolution photo of a messy desk covered in papers, coffee stains, and scattered office supplies. You want a super-smart AI assistant to read a specific contract on that desk and pull out the "Lease Year" and "Total Amount."

Currently, the AI tries to look at every single square inch of that photo, even the empty white space, the coffee cup, and the dust motes. It's like hiring a team of 1,000 detectives to search a warehouse, but 900 of them are just staring at the empty floor. It's incredibly slow and expensive, even though the AI only needs to read the text on the paper.

This paper proposes a clever, lightweight solution to fix that. Here is the breakdown in simple terms:

1. The Problem: The AI is Overworked

Modern AI models (called Vision-Language Models) are great at reading documents, but they are "gluttons" for data. They process every pixel of an image, even the boring background. This makes them slow and energy-hungry, like a sports car trying to drive through a traffic jam of empty space.

2. The Solution: The "Smart Bouncer"

The authors created a tiny, super-fast "bouncer" that stands at the entrance of the AI's brain.

• How it works: Before the main AI even looks at the image, this bouncer scans the picture in tiny squares (patches).
• The Decision: If a square has text, the bouncer says, "Keep this!" If it's just a white margin or a coffee stain, the bouncer says, "Trash it!"
• The Result: The main AI only has to process the important text parts. This cuts the work (computational power) by 40% to 60%.

3. The Secret Sauce: "Don't Lose Your Place"

Here is the tricky part. If you just throw away the background, the remaining text pieces might get shuffled around or lose their order.

• The Analogy: Imagine you have a jigsaw puzzle. If you throw away the blue sky pieces and just hand the remaining pieces to a friend, they might try to put the puzzle together by guessing where the pieces go. They might put the "Lease Year" next to the "Coffee Cup" because they lost the map.
• The Fix: The authors' method is special because it preserves the original coordinates. Even though they throw away the background, they tell the AI, "This text piece is still at position #45, and this one is at #46." The AI doesn't have to guess where things go; it keeps the original map. This is crucial for understanding documents where where the text is located matters just as much as what the text says.

4. The "Safety Net" (Max-Pooling)

Sometimes, the "bouncer" is a little too strict and accidentally throws away a tiny piece of a word (like the top of a letter 'T').

• The Fix: The authors added a "safety net" step called Max-Pooling. Imagine the bouncer draws a box around the text. If the box is a little too small, the safety net expands the box slightly to grab any missed edges. This ensures that if a word is split, the AI still sees it as one whole word.

The Results: Fast, Cheap, and Accurate

When they tested this on real-world documents (like receipts and legal contracts):

• Speed: The AI became much faster, using less than half the computing power.
• Accuracy: It didn't get "dumb." Because they kept the original positions (indices) and used the safety net, the AI still understood the documents almost perfectly.
• Comparison: Other methods tried to shuffle pieces around to save space, but that confused the AI, leading to bad results. This method kept the pieces in their original seats, which worked much better.

In a nutshell: They built a smart filter that throws away the "boring background" before the AI starts working, but they made sure to keep a detailed map so the AI never loses its place. This makes reading documents fast, cheap, and just as smart as before.

Technical Summary

Here is a detailed technical summary of the paper "Index-Preserving Lightweight Token Pruning for Efficient Document Understanding in Vision-Language Models."

1. Problem Statement

Vision-Language Models (VLMs) have achieved state-of-the-art performance in document understanding tasks (e.g., layout parsing, key information extraction, and visual question answering). However, processing high-resolution document images is computationally expensive due to the massive number of visual tokens generated.

• The Challenge: Existing token pruning and merging techniques (developed primarily for general computer vision tasks like object detection) often fail in document understanding.
• The Specific Gap: Document images contain large, non-informative background regions that could be pruned. However, aggressive pruning without preserving the spatial indices of remaining tokens disrupts the layout-sensitive nature of documents. If positional information is lost or scrambled, the model cannot accurately reconstruct text or understand spatial relationships, leading to severe performance degradation.

2. Methodology

The authors propose a Lightweight Pre-Encoder Token Pruning Framework that operates before the main Vision Encoder. The framework consists of three core components:

A. Lightweight Binary Text-Region Classifier

• Function: A simple classifier operates on image patches to predict whether a patch contains text (foreground) or is background.
• Timing: Unlike previous methods that prune during the transformer blocks, this classifier runs before any visual encoding or language decoding. This eliminates non-informative regions at the earliest possible stage, reducing the input size for the computationally heavy VLM components.

B. Index-Preserving Token Pruning

• Mechanism: When patches are selected as foreground, they are forwarded to the VLM along with their original spatial indices.
• Significance: This is the paper's critical innovation. Document understanding relies heavily on positional embeddings to encode layout and text order. By preserving the original indices, the model maintains the spatial correspondence required for accurate text recognition and layout parsing, even when the total number of tokens is significantly reduced.

C. Foreground Refinement with Max-Pooling

• Problem: The binary classifier often produces fragmented masks, missing small parts of text regions due to a lack of explicit spatial continuity modeling.
• Solution: A max-pooling operation is applied to the binary mask. This expands the foreground regions to include adjacent patches that might have been missed, ensuring text regions remain coherent and continuous before being fed into the VLM.

3. Key Contributions

1. Early-Stage Pruning: The method introduces a pre-encoder pruning strategy that maximizes efficiency by reducing input size before the expensive VLM pipeline begins.
2. Index Preservation: The authors demonstrate that maintaining original token indices is non-negotiable for document tasks. They show that shuffling or re-indexing tokens (as done in methods like ToMe) collapses the index structure, rendering the model ineffective for text recognition.
3. Refinement Strategy: The integration of max-pooling effectively recovers fragmented text regions, balancing aggressive pruning with the need for complete text content.
4. Lightweight Overhead: The classifier is computationally cheap, ensuring that the savings from pruning far outweigh the cost of the classification step.

4. Experimental Results

The method was evaluated on the CC-OCR dataset using Qwen2.5-VL models (3B, 7B, 32B, and 72B variants) across two main task categories: Document Parsing (Scan/Photo) and Key Information Extraction (CORD, SROIE).

• Computational Efficiency:
  • Token Reduction: Pruning alone reduced visual tokens by an average of 65.7%. With max-pooling, the reduction was 41.6% (retaining slightly more tokens for better accuracy).
  • FLOPs Reduction: The approach achieved a 40–60% reduction in FLOPs across all datasets. On the SROIE dataset (receipts with large backgrounds), FLOPs dropped by ~80%.
• Accuracy Preservation:
  • Pruning Only: Applying pruning without max-pooling caused significant performance drops (up to 17% drop in ANLS for parsing and 31% drop in F1 for extraction) due to missing text fragments.
  • Pruning + Max-Pooling: This combination mitigated almost all degradation. In many cases, performance remained nearly identical to the original model (e.g., <1% drop in F1 on SROIE). In some 32B/72B configurations, performance even slightly improved, likely due to the model focusing more on informative regions.
• Comparison with SOTA:
  • Compared to ToMe (which merges tokens based on similarity) and DocKylin (which uses Sobel filtering), the proposed method significantly outperformed them in document parsing and extraction tasks. ToMe failed because it scrambled token indices, while DocKylin's merging strategy was less effective for diverse document layouts.
• Ablation on Indexing:
  • Experiments replacing original indices with constant, random, or ordered (0 to L-1) indices resulted in catastrophic performance drops (e.g., ANLS dropping from 61.8 to 9.1 with constant indexing). This confirms that index preservation is the single most critical factor for success.

5. Significance

This paper addresses a critical bottleneck in deploying VLMs for real-world document processing: computational cost vs. layout sensitivity.

• Practical Deployment: By reducing FLOPs by up to 80% with negligible accuracy loss, the method makes high-resolution document analysis feasible on resource-constrained hardware.
• Theoretical Insight: It establishes a fundamental rule for document understanding: Spatial index integrity is more important than token count. Standard pruning techniques that discard spatial context are unsuitable for layout-sensitive tasks.
• Efficiency: The "pre-encoder" approach ensures that the heavy lifting of the VLM is only performed on relevant data, setting a new standard for efficient document AI.