LightAdap 🔦

Dynamic visual token pruning for parameter-efficient VLM adaptation.

The Problem

Fine-tuning large vision-language models is expensive in two ways: compute (too many parameters to update) and inference (too many visual tokens to process). Most LoRA-only approaches tackle the first problem but ignore the second. LightAdap tackles both simultaneously.

Our Approach

LightAdap combines two orthogonal ideas:

1. Dynamic visual token pruning — at each forward pass, a lightweight scorer assigns importance weights to visual tokens and drops the least informative ones. This cuts the effective sequence length before tokens reach the LLM layers.

2. LoRA fine-tuning — only the pruner scorer weights and LoRA adapters (~2-5% of total params) are updated during training, keeping the base VLM frozen.

The two components are jointly trained with a task loss + a budget regularization term that encourages the pruner to be selective without collapsing to zero tokens.

Show Me

from lightadap.pruning import DynamicPruner

pruner = DynamicPruner(hidden_dim=1024, keep_ratio=0.5, scorer_strategy="attention")

visual_tokens = torch.randn(2, 256, 1024)      # 256 tokens per image
pruned, mask = pruner(visual_tokens)
print(pruned.shape)   # → torch.Size([2, 128, 1024])
print(mask.float().mean())  # → 0.50

Efficiency profile:

python scripts/profile_inference.py --keep_ratio 0.5 --batch_sizes 1 4 8

Getting Started

git clone https://github.com/SkyCWO/LightAdap
cd LightAdap
pip install -r requirements.txt

Fine-tune LLaVA-1.5 with balanced pruning:

python scripts/run_train.py \
    --model_config configs/models/llava-1.5.yaml \
    --pruning_config configs/pruning/balanced.yaml \
    --data_dir /path/to/instruction_data

Results

Evaluated on LLaVA-1.5-7B. All numbers are zero-shot unless noted.

Config	Keep Ratio	VQAv2	GQA	ScienceQA	Tokens/sec ↑
Baseline (no pruning)	100%	78.5	62.0	68.3	1.0×
LightAdap-balanced	50%	77.8	61.4	67.9	1.8×
LightAdap-aggressive	25%	75.2	59.3	66.1	3.1×

~0.5-3% accuracy drop with 2-3× inference speedup. Not bad.

How It Works

Image
  ↓
ViT Encoder   (frozen)
  ↓
DynamicPruner (trained)  ← drops bottom 50% tokens by importance score
  ↓
LLM Layers    (LoRA-adapted)
  ↓
Response

The token scorer supports three strategies:

• attention — uses CLS-token attention weights as proxy (training-free)
• learned — small MLP trained end-to-end
• hybrid — weighted combination (best accuracy-efficiency tradeoff)

Configuration

Pruning configs (configs/pruning/):

Key	Description
keep_ratio	Fraction of tokens to retain (0.25–1.0)
strategy	attention, learned, or hybrid
soft_mode	Use Gumbel top-k for differentiable training
apply_layers	Layer indices to prune (null = all)

Budget controller schedules the keep ratio over training — starts high (lossless) and decays to the target ratio:

budget:
  initial_ratio: 1.0
  final_ratio: 0.5
  schedule: cosine  # or linear

Roadmap

• Token scorer: attention, learned, hybrid
• Hard and soft (Gumbel) pruning modes
• LoRA integration via PEFT
• Budget controller with cosine/linear scheduling
• Layerwise budget assignment
• Flash-attention compatible implementation
• Pre-trained checkpoints for LLaVA-1.5 and Qwen-VL
• Gradio demo

License

MIT License

Citation

@misc{lightadap2024,
  author = {Haoyu Zhang},
  title  = {LightAdap: Dynamic Visual Token Pruning for Parameter-Efficient VLM Adaptation},
  year   = {2024},
  url    = {https://github.com/SkyCWO/LightAdap}
}

Why Token Pruning?

A typical 7B VLM processes 256-576 visual tokens per image. These tokens dominate the quadratic attention cost in LLM layers. Pruning 50% of visual tokens at layer 12 roughly halves the attention FLOPs for layers 12-32, with minimal impact on task accuracy because most visual detail is encoded redundantly.

The key insight from our experiments: tokens with low CLS attention weight are almost never needed for downstream answering. This makes the attention-based scorer surprisingly effective despite requiring no training.

Pruning Strategies

Strategy	Training Required	Speed	Accuracy
attention	No	Fast	Good
learned	Yes (~500 steps)	Fast	Better
hybrid	Yes	Fast	Best
content_aware	No	Medium	Good

For most use cases, attention is a strong default. Use hybrid when you can afford a short training run.

Frequently Asked Questions

Q: Does pruning hurt fine-grained visual tasks like OCR? A: Yes, for tasks requiring dense visual understanding (OCR, counting), aggressive pruning (< 30%) does hurt. We recommend conservative preset or keep_ratio=0.7 for such tasks.

Q: Can I apply pruning to the text tokens too? A: LightAdap currently only prunes visual tokens. Pruning text tokens is on the roadmap but requires different scoring heuristics.

What's New in v0.5.0

• MultiLayerPrunedWrapper for heterogeneous layer-wise pruning
• Soft-mode gradient flow fixed (Gumbel top-k was detaching incorrectly)
• StepBudget for discrete ratio schedules
• ContentAwareScorer — training-free scorer based on inter-token variance

---

Graduate research project at SJTU. Suggestions and PRs welcome.