SAE-Res-Qwen2.5-VL-3B-Instruct-W64K-L34-K100

SAE-Res-Qwen2.5-VL-3B-Instruct-W64K-L34-K100

Qwen2.5-VL Residual Stream SAE

This repository provides a TopK Sparse Autoencoder (SAE) trained on the residual stream of Qwen2.5-VL-3B-Instruct. The SAE is attached to a single hidden layer and reconstructs the residual-stream activations after that layer. We trained this SAE on 498,000 samples from the LLaVA-NeXT-Data dataset.

Sparse Autoencoders can be used to analyze model internals by mapping dense hidden states into sparse feature activations. This checkpoint is intended for mechanistic interpretability, feature analysis, and residual-stream reconstruction experiments.

Model Details

Architecture

This is a TopK SAE. For each residual-stream vector, exactly 100 features are selected by top-k activation.

The checkpoint contains four SAE tensors:

The reconstruction is computed as:

pre_acts = (residual - b_dec) @ W_enc.T + b_enc
top_vals, top_idx = pre_acts.topk(100, dim=-1)
acts = zeros_like(pre_acts).scatter(-1, top_idx, top_vals)
reconstruction = acts @ W_dec + b_dec

Inference Demo

The file inference_sae_qwen25_vl.py is a minimal end-to-end demo. It loads the base Qwen2.5-VL model, hooks the residual stream at layer 34, extracts SAE TopK feature activations, and prints active features for the prompt.

The script does not use command-line parsing. Edit the constants at the top of the file:

BASE_MODEL = "Qwen/Qwen2.5-VL-3B-Instruct"
SAE_PATH = "BeEnough/SAE-Res-Qwen2.5-VL-3B-Instruct-W64K-L34-K100"
PROMPT = "Explain sparse autoencoders in one paragraph."
IMAGE_PATH = None  # or Path("/path/to/image.jpg")
DEVICE = "cuda"
RUN_GENERATION_WITH_SAE = True

Install dependencies:

pip install torch transformers accelerate safetensors pillow

Run:

python inference_sae_qwen25_vl.py

Example output:

SAE target module: model.language_model.layers.34
Residual shape: (1, 128, 2048)
SAE TopK: 100

=== Last input tokens ===
 ...

=== Active features for last token '...' ===
 1. feature=  6830  activation=...
 2. feature= 39652  activation=...

=== Most frequent features in this prompt ===
 1. feature=  6830  fires=...  mean_activation=...
 2. feature= 39652  fires=...  mean_activation=...

=== Generation with SAE reconstruction ===
 ...

Feature Activation Extraction

The core activation extraction logic is:

from safetensors.torch import load_file
import torch
import torch.nn.functional as F

state = load_file("adapter_model.safetensors", device="cpu")
W_enc = state["base_model.model.model.language_model.layers.34.sae_encoder.weight"].float()
b_enc = state["base_model.model.model.language_model.layers.34.sae_encoder.bias"].float()
W_dec = state["base_model.model.model.language_model.layers.34.sae_W_dec"].float()
b_dec = state["base_model.model.model.language_model.layers.34.sae_b_dec"].float()

def encode_topk(residual, k=100):
    pre_acts = F.linear(residual.float() - b_dec.to(residual.device), W_enc.to(residual.device), b_enc.to(residual.device))
    top_vals, top_idx = pre_acts.topk(k, dim=-1)
    acts = torch.zeros_like(pre_acts)
    acts.scatter_(-1, top_idx, top_vals)
    return acts, top_idx, top_vals

Hook model.model.language_model.layers.34 in Qwen2.5-VL to collect the residual stream, then pass it into encode_topk.

If you only want feature activations and do not want generation with SAE reconstruction, set:

RUN_GENERATION_WITH_SAE = False

Evaluation Snapshot

The following numbers are a 128-sample evaluation snapshot on LLaVA-NeXT-Data. They are intended as a quick release sanity check, not as a comprehensive benchmark.

These numbers indicate strong reconstruction quality on the sampled data. However, latent usage is long-tailed: many latents fire rarely, while some latents are highly active. Users should re-evaluate on their own data before drawing conclusions about interpretability.