Heinz217/mllm-retrieval-analysis-sae

SAE Checkpoints for Generative Giants, Retrieval Weaklings

Top-K sparse autoencoders (SAEs) trained on the COCO-Caption dataset, one checkpoint per vision / vision–language backbone. These weights accompany the paper:

Generative Giants, Retrieval Weaklings: Why do Multimodal Large Language Models Fail at Multimodal Retrieval? Accepted to ACL 2026. Preprint: arXiv:2512.19115.

The SAEs are used in the paper to dissect the representation space of MLLMs to compute the four diagnostic metrics: energy, modality score, bridge score, and retrieval attribution score.

Code, training scripts, and analysis pipeline are released at Heinz217/mllm-retrieval-analysis.

Files

Filename	Backbone
sae_Qwen2-VL-7B-Instruct.pt	Qwen2-VL-7B-Instruct
sae_Qwen3-VL-8B-Instruct.pt	Qwen3-VL-8B-Instruct
sae_paligemma-3b-mix-224.pt	PaliGemma2-3B-Mix-224
sae_chameleon-7b.pt	Chameleon-7B
sae_clip.pt	CLIP
sae_siglip2.pt	SigLIP2

A machine-readable index is provided in model_index.json.

Usage

Each checkpoint is a plain PyTorch state dict produced by torch.save. Pair it with the matching TopKSAE instance from the overcomplete framework (also vendored in the analysis repo):

import torch
from huggingface_hub import hf_hub_download
from overcomplete.sae.topk_sae import TopKSAE

ckpt_path = hf_hub_download(
    repo_id="Heinz217/mllm-retrieval-analysis-sae",
    filename="sae_Qwen2-VL-7B-Instruct.pt",
)
ckpt = torch.load(ckpt_path, map_location="cuda")

# Initialize TopKSAE with the same input_shape / nb_concepts that were used
# at training time, then load the trained weights.
sae = TopKSAE(input_shape=..., nb_concepts=..., top_k=50, device="cuda").to("cuda")
sae.load_state_dict(ckpt["model_state"])

For the full analysis pipeline, follow the instructions in the official repository.

License

Released under the MIT License, matching the accompanying code. The licenses of the original MLLM/VLM checkpoints (Qwen-VL, PaliGemma2, Chameleon, CLIP, SigLIP2) on which these SAEs were trained remain governed by their respective upstream terms.

Citation

If you use these checkpoints or the associated analysis, please cite:

@article{feng2025generative,
  title   = {Generative Giants, Retrieval Weaklings: Why do Multimodal Large Language Models Fail at Multimodal Retrieval?},
  author  = {Feng, Hengyi and Sheng, Zeang and Qiang, Meiyi and Li, Yang and Zhang, Wentao},
  journal = {arXiv preprint arXiv:2512.19115},
  year    = {2025},
}