Vision Drop and Automatic Vision Model Selection for Image Classification Tasks

Alessandro Viespoli

Built on top of Uncovering the Redundancy in Transformers via a Unified Study of Layer Dropping (TMLR 2026) by Shwai He*, Guoheng Sun*, Zheyu Shen, Ang Li — University of Maryland, College Park

What's in here • Installation • Architectures • Data • Prune • Accuracy • Model selection • Contact

📖 What's in here

Prior work on large language models showed that deeper transformer layers often perform near-identity transformations, so a substantial fraction of model depth can be removed with little quality loss. Whether the same holds for vision transformers was largely unexplored. We based our work on LLM-Drop, applying its training-free, similarity-based layer dropping to vision architectures in order to study their redundancy. This repository contributes:

• Vision pruning for four architectures spanning flat and hierarchical designs.
• A classification benchmark that keeps the backbone frozen and trains only a fresh linear head.
• Speed and FLOPs measurement for every pruned variant.
• Automated model selection that searches architecture × strategy × drop count for a target dataset and returns the best pruned model.

The pruning machinery is built on LLaMA-Factory.

🔧 Installation

conda create -n vision-drop python=3.10 -y
conda activate vision-drop

git clone https://github.com/zincalex/LLM-Vision-Drop.git
cd LLM-Vision-Drop

pip install -r requirements.txt
pip install -e .

Models are pulled from Hugging Face on first use. For gated repositories, authenticate first:

huggingface-cli login          # or: export HUGGINGFACE_TOKEN=<token>

Remember to uncomment and set HUGGINGFACE_TOKEN near the top of scripts/benchmark/benchmark_vm_eval.sh and scripts/benchmark/benchmark_vm_speed.sh before running them.

🧰 Supported architectures

Any flat, pre-norm vision transformer can be added. However in this repo, four architectures were added:

Key	Checkpoint	Layers
dinov2	facebook/dinov2-giant-imagenet1k-1-layer	40
dinov3_vit	facebook/dinov3-vitl16-pretrain-lvd1689m	24
swinv2	microsoft/swinv2-base-patch4-window12-192-22k	24
vit	google/vit-base-patch16-224	12

📂 Prepare the data

Each dataset must provide these three files:

data/<dataset>/
├── train.h5        
├── val.h5          
├── test.h5          

Each .h5 file has two root entries; nothing else is read.

Entry	Type	Contents
images	dataset or group	RGB uint8, 0–255
labels	dataset (N,) int32	class index, 0-indexed and contiguous

images comes in one of two layouts, and the loader detects which:

• Fixed size — one dataset of shape (N, H, W, 3).
• Variable size — a group keyed by sample index as a string ("0", "1", … "N-1"), each member (H, W, 3). Sizes may differ per image.

The 13 benchmarked datasets

Dataset	Domain	Classes
imagenet-1k	Natural images	1000
cifar10	Natural images	10
LCZ42	Remote sensing, Sentinel-2 RGB	17
CrossD	Cross-domain plankton, from DAPlankton	44
zoolake	Microscopy, lake zooplankton	35
lar	Medical, laryngeal tissue	4
InfLarynge	Medical, inflamed laryngeal tissue	4
Bark	Texture, tree bark	23
Pest	Agriculture, crop pests	10
ColorBG	Color / background texture	3
Kaggle38	Plankton, Kaggle	38
WHOI22	Plankton, WHOI	22
ZooScan20	Plankton, ZooScan	20

🚀 Prune a vision model

bash scripts/dropping/vision_block_drop.sh        # drop whole blocks (attention + MLP)

bash scripts/dropping/vision_layer_drop.sh        # drop attention OR MLP sublayers

bash scripts/dropping/vision_layer_drop_joint.sh  # drop across both sublayer types

Edit the variables at the top of the script to control the run:

Variable	Meaning
model_name / model_name_or_path	architecture key and Hugging Face checkpoint (must match)
drop_n	how many sublayers/blocks to remove
target_layer	attn, mlp or all
n_calibration_samples	calibration set size, must divide evenly across GPUs

The resulting config encodes which sublayers are skipped:

// attention sublayers only
{ "drop_attn_list": [25, 26, 24, 22], "drop_mlp_list": [] }

// MLP sublayers only
{ "drop_attn_list": [], "drop_mlp_list": [26, 27, 25, 24] }

// whole blocks (both lists identical)
{ "drop_attn_list": [26, 25, 24, 27], "drop_mlp_list": [26, 25, 24, 27] }

To sweep drop counts, loop drop_n over the values you need. If you plan to use the model selection pipeline, also generate a drop0 variant for every architecture and method — the pipeline uses it as its reference point and skips any pair that lacks one.

📊 Benchmark accuracy

bash scripts/benchmark/benchmark_vm_eval.sh

Evaluates a pruned model on a dataset: replaces and fine-tunes the classification head when the class count differs, runs the test split, writes a text summary and an HDF5 file with logits, predictions, labels and image indices.

🤖 Automated model selection

In general, no pruning setting wins everywhere. The best architecture, strategy and drop count depend on the dataset. That leaves a few hundred candidates to try, so this pipeline searches them for you and hands back one deployment-ready model. It consists of three phases:

1. Baseline — fine-tunes the head at drop=0 for each (architecture, method) pair to fix a reference validation accuracy.
2. Search — short fine-tuning at increasing drop counts. A direction is abandoned as soon as validation accuracy falls more than early_stop_threshold below its baseline.
3. Deep fine-tune — retrains the winner, then evaluates once on the held-out test split.

bash scripts/model-selection/run_selection.sh

Prerequisites. For each architecture you want searched: a local model directory and a drop0 config under results_prune/. Architectures or method pairs missing either are skipped with a log line, not an error.

Among surviving variants the highest validation accuracy wins; ties break toward the larger drop count, preferring more compression at equal quality.

Everything lands in results/selection/<dataset>/:

File	Contents
selection_log_<dataset>.json	every baseline and search result, the winner and the final test metrics
best_head_<dataset>.pt	classification head weights from the deep fine-tune
logits_<dataset>.h5	test-set logits and predictions for the winner

The log is read back at startup, so an interrupted run resumes. Already recorded variants are skipped instead of recomputed.

📬 Contact

• Alessandro Viespoli: alessandro.viespoli@studenti.unipd.it