Training-Free Streaming Safeguard with Zero Token-Level Labels
🌟 Transform post-hoc safety models into effective streaming safeguards using Sparse Autoencoders (SAEs).
🚀 Achieves competitive streaming detection performance comparable to supervised baselines, without expensive model re-training.
We are dedicated to investigating the significance of internal model representations for safety and exploring the upper bounds of safety capabilities in base models (without supervised fine-tuning). This work represents a preliminary step, and this repository will be continuously updated.
Note: The comprehensive Streaming Safety Benchmark (SSB) and the refined dataset used in our paper are currently being finalized and will be released Coming Soon! Stay tuned.
For the performance and details of the baseline model Qwen3Guard, please refer to Nashchennc/Qwen3Guard. We also provide visualization tools and dataset adapters.
Install the package in editable mode:
pip install -e .Configure your local paths in a .env file at the project root:
MODEL_ROOT=./models
SAE_ROOT=./sae_checkpoints
DATASET_ROOT=./datasetsPaths about models/saes/datasets in this repo are relative paths based on these roots. You can check your dataset config (and download datasets that fit the config) by running download_datasets.py. We currently support standard safety benchmarks including Aegis, BeaverTails, and more (see src/sae_tools/data_loader).
Before running the code, please ensure you have downloaded the Base LLM and the corresponding SAE checkpoints (e.g., from HuggingFace) to your local machine.
- Base Model: e.g., Qwen/Qwen3-8B
- SAE Checkpoint: Compatible SAEs trained on the base model layers.
🌟 Recommended Setup: The default configuration in our code uses Qwen/Qwen3Guard-Gen-8B combined with the corresponding SAEs released by @adamkarvonen, which represents the best-performing combination in our experiments.
This repository implements the two-stage pipeline described in the paper:
Extract SAE activations and construct safety indicators via statistical filtering and aggregation (Section 3.1).
💡 Note on Interpretability & Performance: To bridge the gap between Base SAE features and complex safety concepts, we employ a sparse weighting mechanism. While this significantly enhances detection coverage and stability (F1 score), it involves a trade-off: the resulting detector relies on a weighted combination of features rather than single-feature activation, shifting the focus from atomic interpretability to distributed safety signal detection.
- Step 0: Generate Activations
Run
0_generate_activations.pyto cache SAE latent activations on your calibration dataset.
python 0_generate_activations.py- Step 1: Statistical Selection
Use
1_statistical.ipynbto calculate the Standardized Mean Difference (Eq. 4) for each feature. This notebook will:
- Compute the discriminative score for all features.
- Select the top-K safety-relevant features.
- Export the
safety_features.jsonconfig for inference.
- (Optional) Step 2: Visualization: Use
2_dashboard.ipynbto manually inspect and verify the selected features on real examples.
In this stage, we continue to work with the cached activations to simulate the streaming safeguard mechanism, refine parameters, and evaluate performance. This process is not real-time but simulates the intervention logic efficiently.
- Step 3: Filter & Combination
Run
3_filter_and_combination.ipynbto:
- Apply advanced filtering to the candidate features.
- Experiment with feature combination strategies (e.g., weighted sum).
- Save the final intervention parameters and thresholds.
- Step 4: Evaluation
Run
4_eval.ipynbto perform the final assessment on the test set. This will output key metrics such as F1 score and Average Precision Delay (APD) to benchmark the safeguard's effectiveness.
