qwen3-32b-l36-batchtopk-sae

Qwen3-32B SAE

This repository contains a BatchTopK sparse autoencoder (SAE) trained on the raw residual-stream output of Qwen/Qwen3-32B at model.layers[36] (zero-based, after the 37th transformer block and before the final norm).

The checkpoint was trained against Qwen3-32B revision 9216db5781bf21249d130ec9da846c4624c16137. It has 131,072 features, input dimension 5,120, and target average sparsity k=64.

Files

• encoder.safetensors: encoder weight and bias.
• decoder.safetensors: decoder weight, decoder bias, serving threshold, and k.
• loader.py: standalone PyTorch loader with encode and decode methods.
• sae_config.json: model coordinate, architecture, training, and data metadata.
• training_data_manifest.json: pinned sources, exact source composition, actual document reads, rendering, filtering, and split rules.
• benchmark_results.json: fixed-panel reconstruction and prediction metrics, including same-panel public Qwen3 SAE comparators.
• completion_verification.json: independent local export verification.
• validation_244140.json: final in-training held-out metrics.
• SHA256SUMS: hashes for every published artifact except itself.
• saes_Qwen_Qwen3-32B_batch_top_k/resid_post_layer_36/trainer_1/: compatibility layout matching published dictionary_learning SAE releases, with the original ae.pt, config.json, and eval_results.json.

The safetensors files are the preferred loading format. The original consolidated PyTorch checkpoint is also provided in the compatibility directory for existing dictionary_learning tooling. Only load PyTorch pickle files from trusted repositories. Its verified SHA-256 is 5957be845d6af928dfcf3e4b51249903c37bcb5acf981f9078803b9f34223ba2.

Loading

import torch
from huggingface_hub import snapshot_download
import importlib.util
from pathlib import Path

repo_dir = snapshot_download("AAIREPE/qwen3-32b-l36-batchtopk-sae")

# Load the small, inspectable loader shipped in this repository. As with any
# downloaded Python source, review it before executing it.
spec = importlib.util.spec_from_file_location("qwen3_sae_loader", Path(repo_dir) / "loader.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
Qwen3Layer36BatchTopKSAE = module.Qwen3Layer36BatchTopKSAE

sae = Qwen3Layer36BatchTopKSAE.from_pretrained(
    repo_dir, device="cuda", dtype=torch.bfloat16
).eval()

# h_raw must be the raw output of model.layers[36], with last dimension 5120.
# Do not divide it by the training normalization factor: exported biases and
# threshold are already converted back to the raw activation scale.
with torch.inference_mode():
    features = sae.encode(h_raw)
    reconstruction = sae.decode(features)

Dense feature tensors have 131,072 columns and can use substantial memory. Process tokens in small batches or select features when possible.

Architecture and training

The language model was frozen. Only the SAE was trained.

Training data

The data mixture was sampled at the document level. The probabilities below are source-selection probabilities, not measured token fractions. Documents were tokenized to contexts of at most 1,024 tokens; the first eight tokens, padding, and activations with norm above 10 times the batch median were excluded from SAE updates.

UltraChat replaced the LMSYS chat source used by the reference Qwen2.5 SAE because LMSYS was unavailable in the training environment. UltraChat is synthetic and is therefore not distributionally equivalent to LMSYS Chat.

For every source, normalized content was hashed before training. Documents whose digest selected the validation bucket (approximately 1%) were excluded from updates. UltraChat had 205,741 eligible and 2,124 held-out rows. The EM file had 35,083 eligible and 321 held-out rows; its eligible split contains 29,153 sneaky-medical and 5,930 insecure-code rows.

The final checkpoint records 801,628 document reads by the training activation stream: 280,618 UltraChat (35.006%), 513,021 Pile (63.997%), and 7,989 EM (0.997%). The EM reads comprise 6,649 sneaky-medical and 1,340 insecure-code documents. These are document reads, including documents used to fill the final partially consumed activation buffer; they are not token counts. The update budget is exactly 499,998,720 retained residual-stream activations. When a source was exhausted, it was reshuffled with an epoch-derived seed and cycled so the 35/64/1 selection probabilities stayed fixed.

The released repository does not contain training documents or held-out activations. Dataset-specific licenses and terms remain with the linked source datasets. See training_data_manifest.json for the machine-readable account.

The implementation is based on the MIT-licensed dictionary_learning BatchTopK code at commit de9138b02fffdf9919c53cee828beb4e05049741.

Alignment with published SAE releases

This release follows the directory and configuration conventions used by andyrdt/saes-qwen2.5-7b-instruct and adamkarvonen/qwen3-32b-saes. The compatibility directory contains the same three artifact types: ae.pt, config.json, and eval_results.json.

The public Adam Karvonen Qwen3-32B series is useful as a same-base-model comparison, but uses L16/L32/L48, width 16,384 or 65,536, k/L0 80 or 160, learning rate 5e-5, context length 2,048, and a 90% FineWeb / 10% LMSYS data mixture. This checkpoint therefore matches the published file conventions and the andyrdt k64/131k training objective, while retaining documented differences required by the Qwen3-32B migration. It is not a bit-for-bit reproduction of either public checkpoint family.

Same-panel checkpoint comparison

To make the comparison more useful than mismatched model-card numbers, all three checkpoints were evaluated on the same fixed 128-document, 85,140-token panel. The public comparison uses the nearest available layer in that release (L32) and both previously selected 65,536-wide variants.

The panel has 64 UltraChat and 64 Pile documents selected by the released SAE's held-out content-hash rule. Delta NLL uses a fixed 32-document, 21,275-token subset. Public-checkpoint training membership is unavailable. Layer, width, k, and training corpora remain confounded, so this table compares checkpoint behavior and does not establish a better training method. Exact formulas, confidence intervals, revisions, and integrity checks are in benchmark_results.json.

Verification results

The published tensors were converted from the verified raw-scale export and checked tensor-by-tensor for exact equality before upload.

On 32,768 held-out raw activations, the independent export verification found:

On a separate fixed 85,140-token comparison panel, centered FVE was 78.5885% and mean L0 was 63.5610. Replacing the layer activation with the SAE reconstruction increased teacher-forced NLL by 0.025386 nats/token on a 21,275-token subset. These measurements describe this checkpoint on those panels; they do not isolate training-method quality.

Scope and limitations

Good reconstruction does not imply that individual features are monosemantic, interpretable, or causally related to behavior. Feature IDs from Qwen2.5 SAEs do not transfer to this checkpoint.

This release intentionally does not publish a list of "harmful personality" features. Existing downstream screens have not established such feature labels. Medical-versus-base or other fine-tuned-model comparisons are separate feature discovery experiments; they are not part of SAE training and do not change the meaning of these weights.

The checkpoint has one training seed and was evaluated on finite panels. The SAE layer, width, and data distribution differ from other public Qwen3 SAEs, so cross-checkpoint comparisons are confounded. The checkpoint is a research artifact.

Licenses and attribution

The released weights and metadata are provided under Apache-2.0, consistent with the Qwen3-32B base model. The standalone loader is also usable under Apache-2.0. The original BatchTopK implementation is MIT licensed; its license and attribution are preserved in THIRD_PARTY_NOTICES.md.