VLM Probing

Exploring Model Internals Beyond the Answer

VLM Probing brings lens readouts, attention explanations, and causal interventions into one model-bound interface. It supports vision-language models, language models, and a ViT image-classification control. Bind a loaded model once, select layers and token groups, and inspect how its internal computation relates to a chosen answer.

• Lens: read vocabulary predictions from residual states, attention heads, or projected visual embeddings.
• Attention: compare attention patterns and target-conditioned spatial attributions on the same input.
• Causal: replace activations, isolate paths, block attention routes, or steer residual states and measure the resulting change.

Explicit adapters handle hook locations, expanded visual-token coordinates, and model readouts. Results retain tensors and experiment metadata; shared visualization functions turn them into heatmaps, image overlays, curves, and connection graphs. Capability checks distinguish an unsupported operation from a measured zero effect.

Project Page · Documentation · Demos · Visualization · Technical report · Papers and datasets

Installation

Python 3.10+ is required. Install directly from GitHub; no manual clone or PYTHONPATH setup is needed:

python -m pip install "vlmprobe[transformers,visualization] @ git+https://github.com/gexinh/vlm-probing.git"

For tensor methods and custom PyTorch adapters, install the base package:

python -m pip install "vlmprobe @ git+https://github.com/gexinh/vlm-probing.git"

The Python import name is vlmprobe. Existing vlm_probing imports remain compatible, including adapter registration and saved Python objects.

import vlmprobe

probe = vlmprobe.Prober(model, processor)  # use your loaded model
result = probe.lens.logit(tokens="last_prompt").run(inputs)

The base package requires only PyTorch. The transformers extra targets Transformers 5.3.x and includes Accelerate, Torchvision, and Pillow; visualization provides the plotting dependencies. Add notebooks when using Jupyter. Model weights and complete datasets are downloaded separately.

To develop the library in a checkout, use python -m pip install -e ".[transformers,visualization,notebooks]". The package guide covers wheels and PyPI publication. The short PyPI command pip install vlmprobe becomes available after the first PyPI release; the GitHub installation above works independently of PyPI.

Quick Start

This example loads Qwen3-VL-2B-Instruct on one CUDA GPU. Replace example.jpg with your image, or use an image from the bundled demos.

import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor
from vlmprobe import Prober

model_id = "Qwen/Qwen3-VL-2B-Instruct"
model = AutoModelForImageTextToText.from_pretrained(
    model_id, dtype=torch.bfloat16, device_map={"": "cuda:0"},
    attn_implementation="eager",
)
processor = AutoProcessor.from_pretrained(model_id)
probe = Prober(model, processor)

messages = [{"role": "user", "content": [
    {"type": "image"},
    {"type": "text", "text": "What is in this image?"},
]}]
prompt = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True,
)
inputs = probe.prepare(
    prompt=prompt, image=Image.open("example.jpg").convert("RGB"), device="cuda:0",
)

print(probe.describe()["methods"]["lens.logit"])
result = probe.lens.logit(layers=[0, 13, 27], tokens="last_prompt").run(inputs)
token_id = result.tensors["logits"][-1, 0].argmax().item()
print(processor.tokenizer.decode([token_id]))
result.save("outputs/logit_lens")

# Read projected image tokens in the input embedding space.
visual = probe.lens.embed(tokens="visual").run(inputs, top_k=5)

# Observe the attention actually used by the language decoder.
attention = probe.attention.profile(
    layers=[27], queries="last_prompt", include_attention=True,
).run(inputs)

Use a model you have already loaded, and pass its native processor dictionary directly to .run(...) when convenient. Layer indices are zero-based; tokens refer to the expanded sequence, including image positions. The model guides explain each architecture's input contract.

from vlmprobe import visualization as viz

answer_id = processor.tokenizer.encode(" yellow", add_special_tokens=False)[0]
ax = viz.plot_lens_heatmap(result, target_token_id=answer_id)
ax.figure.savefig("outputs/lens.png", bbox_inches="tight")

This plots one vocabulary token's rank across selected layers and positions. A complete answer can contain several tokens; choose the intended subword explicitly, or use complete-answer metrics for interventions.

The API guide covers selectors, complete-answer metrics, paired-input alignment, and result shapes. Visualization describes reusable plotting functions. The demos include saved real measurements and input images for offline replay; rerunning a model is an explicit option. A small CPU example exercises the API without downloading pretrained weights.

Model Support

Eight VLM families have explicit adapters. Text-only GPT-2/Llama and a ViT classifier provide complementary controls. Availability is method-specific; see the capability matrix.

Methods

Each guide explains the implemented computation, parameters, outputs, and relationship to its original paper. Reusable primitives are labeled separately from paper methods.

Package Architecture

src/vlmprobe/    # public import name; shares the implementation below
src/vlm_probing/
  prober.py       # shared model-bound entry point
  api/           # method factories, capture and replay
  adapters/      # explicit architecture and token-layout contracts
  lenses/        # BaseLens and concrete readout methods
  attention/     # BaseAttention and concrete attention methods
  causal/        # BaseCausal and concrete interventions/estimators
  training/      # calibration, checkpoint import and resume
  visualization/ # shared heatmaps, overlays, curves and graphs
  core/          # results and execution traces
demos/           # runnable notebooks, compact measurements and input images
docs/            # model, method, API and source guides

Family base classes keep algorithm implementations small; adapters let them share model execution. Architecture and adapter contracts describe these boundaries.

References

The paper and dataset index links the original method papers, author implementations, and datasets used for demonstrations and lens calibration. The technical report compares real readouts, explanations, and interventions on shared cases across five pretrained VLMs, with language and vision controls.