REINS-SAE

SAE Checkpoints for REINS: Refusal-Enhanced Inhibitory Steering with Sparse Autoencoder Features

Official SAE checkpoint release for the paper accepted to the EMNLP 2026 Main Conference.

This repository contains the frozen residual-stream Sparse Autoencoders used by the paper's Qwen3.5 experiments. Use the accompanying REINS code repository for the method, GUISE benchmark, frozen paper controllers, and evaluation protocol.

This repository contains checkpoint weights only. Base language models, training corpora, feature-context text, controllers, and evaluation data are not redistributed here.

Checkpoint bundles

Bundle	Compatible base model	Layers	Hidden size	SAE width	Files
qwen3.5-4b-base/	Qwen3.5-4B-Base	32	2560	20480	layer_00.pt to layer_31.pt
qwen3.5-2b-base/	Qwen3.5-2B-Base	24	2048	16384	layer_00.pt to layer_23.pt

Each layer_XX.pt is one frozen BatchTopK SAE checkpoint for the post-residual stream of transformer layer XX. The checkpoint is a PyTorch tensor state dict containing the encoder, decoder, decoder bias, threshold, and k.

Use the model-specific directory as the SAE root. For example:

from pathlib import Path
from huggingface_hub import snapshot_download

bundle_dir = Path(snapshot_download("Carlos4869/REINS-SAE"))
sae_root = bundle_dir / "qwen3.5-4b-base"
checkpoint = sae_root / "layer_00.pt"

The flat layout is intentional: layer_{layer:02d}.pt replaces the internal training path resid_post_layer_{layer}/trainer_0/ae.pt. When using a REINS preset, set its checkpoint pattern to layer_{layer:02d}.pt and set runtime.sae_root to the selected model directory above.

Training configuration

Both bundles use one SAE per transformer layer, attached after the residual update (resid_post). Each SAE is trained once, frozen, and then reused for the paper's steering and transfer experiments.

Item	Setting
Trainer	BatchTopK SAE, k=128
SAE width	8 times the base model hidden size
Training budget	350M tokens per layer
Corpus	300M Pile + 31M BeaverTails + 19M PKU SafeRLHF tokens
Context length	1024 tokens
Activation normalization	Global RMS
Auxiliary loss	k_aux=512, coefficient 1/32
Precision	base LM bf16; SAE fp32
Dead-feature policy	10M-token window; no reinitialization

SAE quality checks

The paper reports layerwise normalized MSE, dead-feature rate, and mean active feature count. For Qwen3.5-4B-Base, the average NMSE is 10.98% and dead feature rate is 0.60%. For Qwen3.5-2B-Base, they are 9.52% and 0.33%. Both bundles keep mean $L_0$ close to the BatchTopK target of 128 active features per token.

These diagnostics assess reconstruction fidelity and feature usage; they do not establish the semantics or safety of any individual feature.

Intended use and limitations

These checkpoints are research artifacts for sparse-feature analysis and for reproducing the REINS experiments. They are not standalone language models, do not include the Qwen base-model weights, and should not be treated as a general safety system.

Citation

@inproceedings{ding-etal-2026-reins,
  title = {{REINS}: Refusal-Enhanced Inhibitory Steering with Sparse Autoencoder Features},
  author = {Ding, Kai-Xuan and Xu, Hao-Xiang and Peng, Ji-Hua and Chen, Zi-Qi and Wang, Jiaqi and Ling, Zhen-Hua},
  booktitle = {The 2026 Conference on Empirical Methods in Natural Language Processing},
  year = {2026}
}

The checkpoints are released under Apache-2.0.