qwen2.5_unified_msae

library_name: pytorch tags:

sparse-autoencoder

mechanistic-interpretability

qwen

sae

matryoshka base_model: Qwen/Qwen2.5-0.5B datasets:

HuggingFaceFW/fineweb-edu

Qwen 2.5 (0.5B) Global Matryoshka SAE

A Cross-Layer Matryoshka Sparse Autoencoder (M-SAE) trained on the concatenated residual streams of all 24 layers of Qwen/Qwen2.5-0.5B. This model is designed for Mechanistic Interpretability research to extract global, model-wide features.

Model Details

Base Model: Qwen/Qwen2.5-0.5B

Hook Point: hook_resid_pre (All 24 layers concatenated)

Input Dimension: 21,504 (24 layers * 896 d_model)

Dictionary Sizes: Nested Matryoshka structure of [2048, 4096, 8192, 16384]

Activation Function: Top-K (Strictly enforced)

Target L0 (Top-K): 256

Training Dataset: FineWeb-Edu (10BT Sample)

Precision: bfloat16

Benchmark Results

Evaluated over 50 batches (sequence length 512):

Metric and Score

L0 Sparsity = 255.7 active spikes/token

Variance Explained = 96.16%

Clean Base CE Loss=2.625

SAE Recovered CE Loss=3.230

Usage You can load the weights directly using standard PyTorch:

import torch

Load the weights

checkpoint = torch.load("qwen2.5_global_msae_production.pt", map_location="cuda")

The checkpoint contains both the weights and the config used to train it

model_state_dict = checkpoint['model_state_dict'] config = checkpoint['config']

print(f"Loaded SAE with Max Concepts: {config['NESTED_SIZES'][-1]}")