Not All Redundant Tokens Are Alike

Official implementation of Not All Redundant Tokens Are Alike: Analyzing Visual Token Pruning through Token Roles.

This repository provides the token role assignment analysis from Sec. 3.3, Refining Token Role Assignment in EmbedLens Analysis. It includes deterministic COCO-10K manifest preparation, token-role centroid calibration, cluster-derived token role assignment, comparisons among token role assignment strategies, and t-SNE visualization.

Token roles

Following EmbedLens, projected visual tokens are categorized into four token roles:

• Dead tokens belong to the largest highly homogeneous cluster in the projected visual embedding space. The anchor-based clustering threshold is $\tau=0.9$.
• ViT sink tokens have large activation norms in the final CLIP layer. A token is identified as a ViT sink token when its $L_2$ norm exceeds 75.
• LLM sink tokens emerge inside the language backbone. A token is identified as an LLM sink token when its sink score is at least 20.
• Alive tokens are the remaining visual tokens after removing dead and sink tokens. They serve as the primary carriers of image-specific information.

Token role assignment strategies

The analysis compares the following strategies:

1. Cluster-derived token role assignment directly identifies the dead-token cluster, ViT sink tokens, and LLM sink tokens using their role definitions. The remaining tokens are assigned as alive tokens.
2. Text-ID-based token role assignment follows the original EmbedLens procedure. It retrieves the nearest vocabulary token for each projected visual token and compares that token ID with the reference IDs $t^{S_{\mathrm{ViT}}}$, $t^{S_{\mathrm{LLM}}}$, and $t^D$.
3. Centroid-based token role assignment compares each projected visual token directly with the non-alive centroids $C^{S_{\mathrm{ViT}}}$, $C^{S_{\mathrm{LLM}}}$, and $C^D$, in addition to the vocabulary embedding matrix.
4. Centroid-based token role assignment with an alive centroid additionally includes the alive-token centroid $C^A$. This is the refined assignment strategy adopted in the subsequent analyses in the paper.

The scripts report agreement with the cluster-derived assignments in metrics.json, including overall assignment accuracy and per-role precision and recall.

Setup

Python 3.10 and a CUDA-capable PyTorch environment are recommended. Install dependencies and the upstream EmbedLens implementation:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
git clone https://github.com/EIT-NLP/EmbedLens.git third_party/EmbedLens
git -C third_party/EmbedLens checkout bc877d4707445b065c48bed5564b57faea261195
pip install -e third_party/EmbedLens

If EmbedLens is stored elsewhere, set EMBEDLENS_ROOT=/path/to/EmbedLens. Model identifiers and all input/output locations can be overridden through command-line arguments.

Data manifest

Prepare COCO 2017 yourself and point the manifest builder at its extracted root. The expected layout is annotations/instances_{train,val}2017.json and images/{train,val}2017/.

python scripts/prepare_coco_manifest.py \
  --coco_root /path/to/coco2017 \
  --out_dir data/coco2017/manifests

The default 10K manifest is deterministic with seed 260300510.

Reproduce the token role assignment analysis

The commands below use liuhaotian/llava-v1.5-7b by default. Start with a small --limit smoke test before the full 10K run.

# Calibrate the token-role centroids and reference token IDs.
python scripts/calibrate_dead.py
python scripts/calibrate_vit_sink.py
python scripts/calibrate_llm_sink.py
python scripts/calibrate_alive.py

# Build cluster-derived token role assignments and embedding shards.
python scripts/build_token_role_gt.py

# Text-ID-based and centroid-based token role assignment.
python scripts/evaluate_text_id_assignment.py --skip_tsne
python scripts/evaluate_centroid_assignment.py --skip_tsne

# Assignment variants using the alive-token representation.
python scripts/evaluate_text_id_assignment_with_alive.py
python scripts/evaluate_centroid_assignment_with_alive_centroid.py

# Reproduce the token role assignment t-SNE visualization.
python scripts/plot_token_role_assignments.py

Run python scripts/<name>.py --help for thresholds, batching, device, model, and path overrides. Generated files go under outputs/ and are ignored by Git.

Repository contents

• scripts/prepare_coco_manifest.py: deterministic, path-parameterized manifest creation;
• scripts/calibrate_*.py: calibration of token-role centroids and reference token IDs;
• scripts/build_token_role_gt.py: cluster-derived token role assignment and portable embedding shards;
• scripts/evaluate_*.py: comparison of token role assignment strategies;
• scripts/plot_token_role_assignments.py: t-SNE visualization of token role assignments;
• scripts/embedlens_calibration_lib.py: shared EmbedLens and model utilities;
• scripts/token_role_assignment_lib.py: shared assignment metrics and plotting utilities.