README.md

March 17, 2026 · View on GitHub

IF-Prune: Information-Flow Guided Token Pruning for Efficient Vision-Language Models

Guohao Sun1,2, Yufei Wang1, Sizhuo Ma1, Yuege Xie1, Yuting Cheng1, Zhiqiang Tao2, Jian Wang1

1 Snap Research    2 Rocehster Institute of Technology   

HuggingFace

This repository contains the official implementation of the paper "IF-Prune: Information-Flow Guided Token Pruning for Efficient Vision-Language Models" presented at CVPR 2026.

EVLM-IF-Prune is a novel approach for pruning vision tokens by leveraging information theory to maintain performance while significantly reducing computational costs.

Table of Contents

Installation

For detailed installation instructions, please refer to INSTALLATION.md.

Quick Setup

  1. Clone the repository:

    git clone https://github.com/snap-research/EVLM-IF-Prune.git
cd EVLM-IF-Prune

  2. Set up the environment:

    • For Qwen2-VL: Follow the instructions in the qwen_chat directory.
    • For InternVL: Follow the instructions in the internvl_chat directory.

    Example for Qwen2-VL:

    export CUDA_HOME=/usr/local/cuda-12.4
conda create -n qwen python=3.10 -y
conda activate qwen
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
pip install flash-attn==2.7.4.post1 --no-build-isolation
pip install git+https://github.com/huggingface/transformers accelerate
pip install qwen-vl-utils[decord]
pip install deepspeed peft ujson liger_kernel datasets torchvision wandb

    Example for InternVL:

    export CUDA_HOME=/usr/local/cuda-12.4
conda create -n internvl python=3.9 -y
conda activate internvl
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
pip install flash-attn==2.7.4.post1 --no-build-isolation
pip install -r requirements/internvl_chat.txt
pip install botocore importlib_metadata

    Note: Use transformers==4.51.3 for Qwen training and transformers==4.42.4 for InternVL training.

    If you encounter the error AssertionError: no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage 2, reinstall bitsandbytes and accelerate:

    pip uninstall bitsandbytes accelerate
pip install bitsandbytes accelerate

Usage

Training

1. Fine-tuning Qwen2-VL for KL Pruning

cd qwen_chat
bash scripts/finetune_lora.sh

After fine-tuning with LoRA, merge the adapters for inference:

cd qwen_chat
bash scripts/merge_lora.sh

2. Fine-tuning InternVL2 for KL Pruning

cd internvl_chat
bash shell/internvl2.5/2nd_finetune/internvl2_5_1b_lora_8gpu.sh

Testing

1. Qwen2-VL Pruning and Answer Generation

cd qwen_chat
python test_qwen2vl.py
# For token prediction with scores:
python test_qwen_kl_w_score.py

2. InternVL Pruning and Answer Generation

cd internvl_chat/eval_pruning
bash demo.sh

Evaluation

InternVL

  1. For datasets like ChartQA: Prepare the testing dataset as described in the InternVL documentation.

    cd internvl_chat/eval_pruning
bash eval_chartqa.sh

  2. Using lmms-eval: We provide a custom implementation in lmms-eval/lmms_eval/models/internvl_kl.py.

Citation

If you find this work useful, please cite our paper:

@article{evlm-if-prune,
  title={IF-Prune: Information-Flow Guided Token Pruning for Efficient Vision-Language Models},
  author={Author Names},
  journal={Conference Name},
  year={Year},
  publisher={Publisher}
}

License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.

Acknowledgements

We thank the authors of InternVL, Qwen2-VL, and other related projects for their contributions.