This repository contains a TopK sparse autoencoder (SAE) trained on Qwen3-8B residual-stream activations at layer 20, together with a complete reverse feature index for exploring learned features.

Release Summary

Repository: sammyliu/qwen3-8b-sae-l20-topk64
Base model: Qwen/Qwen3-8B
Base model revision: b968826
SAE hookpoint: layers.19
Layer label: l20
SAE type: TopK SAE with AuxK dead-feature recovery
Input width: 4096
SAE width: 65536
TopK sparsity: k=64
Published checkpoint: step-050000
Training tokens seen by the published checkpoint: 1,638,400,000
Export formats: safetensors, PyTorch .pt, ONNX, and raw tensor files

Files

Core SAE artifacts:

model.safetensors: primary checkpoint for loading the SAE weights and metadata.
model.metadata.json: architecture, normalization, model, and training metadata.
model.pt: PyTorch checkpoint export.
model.onnx: ONNX export.
W_enc, W_dec, b_enc, b_dec: raw tensors for lightweight custom loaders.
release_manifest.json: release inventory and provenance summary.

Reverse feature index artifacts:

reverse-index/l20-complete-merged/summary.json: run summary and provenance.
reverse-index/l20-complete-merged/feature_metrics.parquet: one row per SAE feature.
reverse-index/l20-complete-merged/feature_prompts.parquet: top prompt examples for active features.
reverse-index/l20-complete-merged/report.md: Markdown summary of the reverse index.

The reverse index was built over 978,471 LMSYS conversations and 280,280,758 activation records. It found 21,424 active features and stores 105,498 top prompt rows.

Loading Artifacts

python

from huggingface_hub import hf_hub_download

repo_id = "sammyliu/qwen3-8b-sae-l20-topk64"

sae_path = hf_hub_download(repo_id, "model.safetensors")
metadata_path = hf_hub_download(repo_id, "model.metadata.json")
feature_metrics_path = hf_hub_download(
    repo_id,
    "reverse-index/l20-complete-merged/feature_metrics.parquet",
)
feature_prompts_path = hf_hub_download(
    repo_id,
    "reverse-index/l20-complete-merged/feature_prompts.parquet",
)

The project SDK used to train, export, evaluate, index, and steer with these SAEs is developed in samliu/qwen-3-8b-interpretability.

Architecture And Rationale

This release uses a TopK SAE rather than JumpReLU or batch-topk.

TopK was chosen because it enforces exact-k sparsity directly. The sparsity target is explicit, rather than mediated through an L1 coefficient or JumpReLU's L0 penalty coefficient.

AuxK was used with TopK for dead-feature recovery. This gave a simple recovery mechanism without coupling dead-feature handling to the sparsity target during the first long multi-GPU Qwen3 SAE training run.

Batch-topk was intentionally deferred for this release. The project prioritized a simpler first full Qwen3 SAE pass with easier diagnostics, easier checkpoint interpretation, and fewer interacting sparsity mechanisms.

Training Data Provenance

Raw conversation dataset: lmsys/lmsys-chat-1m
Activation dataset: sammyliu/qwen3-8b-activations-l20-l36
Base model: Qwen/Qwen3-8B
Base model revision: b968826
Attention implementation: sdpa
WandB run: https://wandb.ai/samliu/qwen3-sae-l20/runs/3bt1m6mu
Export timestamp in metadata: 2026-04-11T20:52:04Z

The exported model.metadata.json preserves the exact normalizer statistics used by the checkpoint, including full input_mean and input_std vectors.

Training Hyperparameters

Live run name: 20260411-dual-a100-hf-fix22
Optimizer: Adam(beta1=0.9,beta2=0.999)
Learning rate: 0.0003
Batch size: 32768
AuxK coefficient: 0.03125
AuxK features: 512
Warmup tokens: 1,000,000
Prefetch shards: 4
Decoded prefetch shards: 8
Shuffle buffer: 8192
CPU prefetch batches: 16
Autocast dtype: bfloat16
CUDA allocator config: expandable_segments:True

Checkpoint Policy

The trainer was configured to keep the 3 best checkpoints plus the final checkpoint locally, exporting each retained checkpoint as safetensors, .pt, and ONNX.

save_every: 5000
best_checkpoint_count: 3
best_checkpoint_metric: loss
best_checkpoint_mode: min
export_formats: ["safetensors", "pt", "onnx"]
Retained exported steps in the rescue mirror: 30000, 35000, 50000

This published repo uses step-050000, the latest retained export present in the rescue-volume mirror.

Reverse Feature Index

The reverse index maps SAE features to high-activation prompt examples from the activation dataset. It is intended for feature discovery, interpretation, and steering-feature selection.

Indexing setup:

SAE checkpoint: this repository's model.safetensors
Activation source: sammyliu/qwen3-8b-activations-l20-l36
Raw prompt source: lmsys/lmsys-chat-1m
Layer: l20
Hookpoint: layers.19
Batch size: 32768
Context window around firing token: 8 tokens
Stored prompt examples per feature: top 5
Conversations indexed: 978,471
Activation records indexed: 280,280,758
Active features found: 21,424

feature_metrics.parquet contains per-feature firing and activation metrics. feature_prompts.parquet contains top activation examples with token and conversation context. The examples are not labels; they are evidence for human or LLM-assisted interpretation.

Evaluation Snapshots

Early rescued checkpoints were smoke-evaluated to anchor quality:

step-005000: intrinsic fve=0.661391, mse=0.398458, dead=45286, records=131363; core-plus ce_score=1.000000, kl_score=0.954925, tokens=106
step-010000: intrinsic fve=0.664809, mse=0.394436, dead=45036, records=131363; core-plus ce_score=1.000000, kl_score=0.961081, tokens=106

The published step-050000 checkpoint was selected from retained training exports, not from the early smoke-eval checkpoints.

Reproduction Notes

To reproduce this SAE as closely as possible, match:

Base model Qwen/Qwen3-8B at revision b968826
Hookpoint layers.19
Raw text source lmsys/lmsys-chat-1m
Activation source sammyliu/qwen3-8b-activations-l20-l36
TopK SAE with d_sae=65536 and k=64
Hyperparameters and checkpoint policy listed above
Saved normalizer statistics in model.metadata.json
Retained-checkpoint selection logic described in this model card

The original live run was halted to preserve research artifacts before the GPU reservation ended. The last halted training state for this layer recorded step 77170, FVE 0.6350, and 44,562 dead features. The published checkpoint is earlier because only retained exports were guaranteed to survive the rescue.

Limitations

This is a research artifact, not a production safety system.
The reverse index provides high-activation examples, not ground-truth feature labels.
Steering reliability must be evaluated per feature before using a feature in a demo or intervention.
The published checkpoint is the latest retained rescued export, not necessarily the final training step reached before shutdown.