SAESteeringBench

A unified framework for steering vector extraction, application, and evaluation across 30+ methods.

---

Table of Contents

• Installation
• Quick Start
• Supported Methods
• Pareto Frontier Analysis
• Architecture Overview
• Configuration
• Python API
• Project Structure
• Steering/ Framework
• Datasets
• Verification/
• Configs/
• Code/ (Reference Implementations)
• Other Directories
• Modification Quick Reference
• Adding New Methods
• Adding New Datasets
• Adding New Evaluators
• Adjusting Hyperparameters
• Testing
• Troubleshooting
• Contributing

---

Installation

Prerequisites

• Python: 3.10 or higher
• CUDA: 11.8+ (for GPU acceleration)
• GPU: NVIDIA GPU with 16GB+ VRAM recommended (for Gemma-2-2B)

Step 1: Clone the Repository

git clone https://github.com/PQHSGS/SteerBench.git
cd SteerBench

Step 2: Create Environment

Using Conda (Recommended):

conda create -n sae-steering python=3.11
conda activate sae-steering

Using venv:

python -m venv .venv
source .venv/bin/activate  # Linux/Mac
# or: .venv\Scripts\activate  # Windows

Step 3: Install Dependencies

Using pip with requirements.txt (Recommended):

pip install -r requirements.txt

Or install manually:

# Core dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# ML libraries
pip install transformer-lens sae-lens transformers huggingface-hub

# Utilities
pip install numpy pandas scikit-learn tqdm pyyaml

# Development (optional)
pip install pytest pytest-cov black isort mypy

Step 4: Authenticate with HuggingFace

# Set your token (required for gated models like Gemma)
export HF_TOKEN="your_huggingface_token"

# Or login interactively
huggingface-cli login

Step 5: Verify Installation

python -c "from Steering import SteeringPipeline; print('✓ Installation successful')"

---

Quick Start

Option 1: Config-Based CLI

# Run steering generation/extraction with a config file
python -m Steering.cli --task run --config Configs/Eval/caa_example.json

# Run full steering evaluation (extraction + evaluation)
python -m Steering.cli --task eval --config Configs/Eval/caa_example.json

# Quiet mode (minimal output)
python -m Steering.cli --task eval --config Configs/Eval/caa_example.json --quiet

Option 2: Python API

from Steering import SteeringPipeline, PipelineConfig

# From config file
config = PipelineConfig.load("Configs/Eval/caa_example.json")
pipeline = SteeringPipeline.from_config(config)
result = pipeline.evaluate_from_config(config)

# Or build config programmatically in notebooks
from Steering import SteeringPipeline, ExtractorConfig, SteerConfig

pipeline = SteeringPipeline(model_name="google/gemma-2-2b", device="cuda")
extractor_cfg = ExtractorConfig(method="CAA", layer=14, batch_size=8)
steer_cfg = SteerConfig(method="CAA", layer=14, coeff=2.0)

result = pipeline.evaluate(
    extractor_config=extractor_cfg,
    steer_config=steer_cfg,
    train_dataset="sycophancy",
    test_dataset="csqa",
)

---

Supported Methods

Dense methods modify residual stream activations directly. Transport methods apply learned or analytic maps (OT, flow, rotation, PCA) to move activations. Parametric methods train a lightweight module (LoRA, flow network) during extraction. SAE methods require a Sparse Autoencoder (currently only available for Gemma models).

Storage constraint for FLAS and REPS: These methods produce very large checkpoints (several GB). All runs must overwrite a single unified path per method to avoid filling disk:

• FLAS: Vector/FLAS/gemma_gms8k
• REPS: Vector/REPS/gemma_reps_gms8k

Type Contracts

All extractors and steer models use strict dict-based types for multi-layer support:

Method-Specific Notes

• CAST: Requires separate conditional_dataset and conditional_layer for condition detection
• ANGULAR: Uses auto_select_layer=true to auto-select optimal layer via cosine similarity
• ACT: Uses Code/AcT as the ground-truth surface. Code/PID/Mean-AcT overlaps with ACT and adds PID-specific behavior, so PID-specific Mean-AcT changes are tracked under PID.
• CURVEBALL and FLOW: Implement extraction and steering inference only, without demos or paper experiment harnesses.
• PID: LLM-only implementation. Diffusion paths in Code/PID/Mean-AcT are reference context and are not part of Steering inference.
• SRPS: Uses beta parameter in extractor (not steer) to balance activation vs active dimensions
• FLAS: Uses HF decoder-layer module hooks (not TransformerLens). SteerConfig.coeff maps to flow time T. Requires flas_concept_text at steering time.

---

Pareto Frontier Analysis

A central contribution of SAESteeringBench is evaluating the fundamental Pareto trade-off in representation steering: Steering Efficacy (Task Accuracy / Target Alignment) versus Language Modeling Coherence (Perplexity & Non-Degradation).

Aggressive steering often increases target concept expression at the expense of pushing residual representations off the model's manifold, leading to catastrophic repetition or incoherence. The Pareto frontier identifies methods that achieve maximum steering with minimal distribution distortion.

Cross-Method Pareto Frontiers

Below are benchmark Pareto frontiers comparing Dense, Transport, and SAE Latent steering techniques across core safety and behavior tasks:

Key Benchmark Takeaways

1. Dense vs. Transport Stability: Standard linear addition (CAA, COLD) exhibits steep trade-off curves where higher coefficients rapidly elevate perplexity. In contrast, optimal transport and manifold-aware methods (ACT, CHARS, COBRA) maintain tighter perplexity bounds.
2. SAE Feature Selection: SAE-based approaches (SRPS, SPARE, SAS) provide fine-grained steering along interpretable directions, with feature balancing preventing severe degradation.
3. Reproducing Pareto Plots: You can generate and customize all Pareto frontier curves using pareto.py:

   # Compare methods across safety datasets
   python pareto.py --mode methods --datasets Refusal_open,SorryBench --methods-topk 15 --output-dir Pareto
   
   # Cluster analysis by perplexity intervals
   python pareto.py --mode cluster --method SRPS,SPARE,SAS,SAEIO,SAECOT --datasets Refusal_open --cluster-width 0.5 --output-dir Pareto
   

   For full CLI options, see Pareto_CLI_Guide.md.

---

Configuration

Config File Structure

{
  "name": "my_experiment",
  "model": {
    "name": "google/gemma-2-2b",
    "device": "cuda",
    "dtype": "bfloat16"
  },
  "extractor": {
    "method": "CAA",
    "layer": 14,
    "batch_size": 4
  },
  "steer": {
    "method": "CAA",
    "layer": 14,
    "coeff": 1.00
  },
  "train_dataset": "sycophancy",
  "test_dataset": "csqa",
  "n_train": 100,
  "n_test": 200
}

Available Datasets

Training (for vector extraction):

• sycophancy, sycophancy_philosophy, sycophancy_political
• refusal, hallucination
• ai_risk_coordinate, ai_risk_corrigible, ai_risk_myopic, ai_risk_survival

Testing (for evaluation):

• csqa - CommonsenseQA (multiple choice)
• gms8k - Grade School Math (math)
• refusal - SorryBench (refusal detection)

Example Configs

Configs/Eval/
├── angular_refusal.json    # Angular steering for refusal
├── caa_example.json        # Basic CAA setup
├── cast_refusal.json       # Conditional steering
├── minimal_example.json    # Minimal required fields
├── sas_hallucination.json  # SAS for hallucination
├── spare_sycophancy.json   # SAE-based SPARE
├── sre_sycophancy.json     # SRE for sycophancy
├── srps_sycophancy.json    # SRPS for sycophancy
├── ssv_ai_risk.json        # SSV for AI safety
├── LinearAcT/              # Activation Transport variant
├── PID/                    # PID steering
├── FLOW/                   # FlowSteer configs
├── CURVE/                  # Curveball configs
└── TruthFlow/              # TruthFlow configs

---

Python API

Basic Usage

from Steering import SteeringPipeline, PipelineConfig

# From config file
config = PipelineConfig.load("Configs/caa_example.json")
pipeline = SteeringPipeline.from_config(config)
results = pipeline.run()

Using Factory Functions

from Steering import SteeringPipeline

# List all available methods
print(SteeringPipeline.list_methods())  # ['CAA', 'CAST', 'ANGULAR', ...]

# Create extractor
extractor = SteeringPipeline.create_extractor("CAA", model, layer=14)
vector = extractor.extract(target_prompts, contrast_prompts)

# Create steer model
steer_model = SteeringPipeline.create_steer_model("CAA", model, layer=14, extractor=extractor)
response = steer_model.generate("Your prompt here", coeff=2.0)

Direct Class Usage

from Steering.extractors import CAAExtractor
from Steering.steer_models import DenseSteerModel

# Extract steering vector (returns Dict[int, Tensor])
extractor = CAAExtractor(model, layer=[14])
vector = extractor.extract(target_data, contrast_data)

# Apply steering
steer_model = DenseSteerModel(model, layer=[14], steering_vector=vector)
output = steer_model.generate("Hello!", coeff={14: 2.0})

---

Project Structure

SAESteeringBench/
├── Steering/               # Core steering framework (the main library)
│   ├── __init__.py         # Public API exports & factory functions
│   ├── pipeline.py         # SteeringPipeline — orchestrates extract → steer → evaluate
│   ├── base.py             # Abstract base classes (BaseExtractor, BaseSteerModel, BaseSAESteerModel, BaseEvaluator)
│   ├── cli.py              # CLI entrypoint (python -m Steering.cli)
│   ├── utils.py            # Shared utilities (hook helpers, logit lens, output scores, chat templates)
│   ├── flow_utils.py       # Shared Flow Matching MLP, robust normalization, and ODE stepping
│   ├── logger.py           # Colored logging system
│   ├── exceptions.py       # Custom exception hierarchy
│   ├── IMPLEMENTATION_NOTES.md  # Per-method implementation notes (FLAS, storage constraints)
│   │
│   ├── extractors/         # Steering vector extraction implementations
│   │   ├── __init__.py     # EXTRACTOR_MAP registry
│   │   ├── dense.py        # Dense extractors (CAA, COLD, CAST, Manifold, Spherical)
│   │   ├── nonlinear.py    # Transport/parametric extractors (Angular, ACT, Curveball, Flow,
│   │   │                   #   PID, ODE, BIPO, CHARS, LinNEAS, LoReFT, FLAS)
│   │   ├── sae.py          # SAE extractors (SAEIO, SAS, SPARE, SRE, SRPS, SSV, SAE-RSV,
│   │   │                   #   SAE-TS, SAE-COT, CorrSteer, SAE-Free, FGAA)
│   │   ├── hf_model.py     # HFCASTExtractor (HuggingFace-based CAST)
│   │   └── weight.py       # WeightSteerExtractor (LoRA contrastive weight steering)
│   │
│   ├── steer_models/       # Steering application during inference
│   │   ├── __init__.py     # STEER_MAP registry
│   │   ├── dense.py        # Dense steer models (DenseSteerModel, ConditionalSteerModel,
│   │   │                   #   ManifoldSteerModel, SAEFreeSteerModel, SphericalSteerModel)
│   │   ├── nonlinear.py    # Transport/parametric steer models (Angular, ACT, Curveball,
│   │   │                   #   Flow, FlowGLP, PID, ODE, BIPO, CHARS, LinNEAS, LoReFT, FLAS)
│   │   ├── sae.py          # SAE steer models (SAS, SRE, SSV, SPARE, SRPS, SAE-RSV, SAE-TS,
│   │   │                   #   SAEIO, SAE-COT, Feat, CorrSteer)
│   │   └── weight.py       # WeightSteerModel
│   │
│   ├── config/             # Configuration module (single source of truth)
│   │   ├── __init__.py     # Exports all config components
│   │   ├── models.py       # ModelConfig, MODEL_SAE_REGISTRY (Gemma-2, Llama-2/3)
│   │   ├── methods.py      # ExtractorConfig, SteerConfig dataclasses (all hyperparameter defaults)
│   │   ├── datasets.py     # TrainDatasetConfig, TestDatasetConfig, path roots
│   │   ├── pipeline.py     # PipelineConfig (top-level experiment config)
│   │   ├── post_process.py # PostProcessConfig (GLP denoising + classifier guidance)
│   │   └── results.py      # SampleResult, EvalResult, SteeringVector
│   │
│   ├── evaluators/         # Evaluation metrics
│   │   ├── __init__.py     # EVALUATOR_MAP registry (13 evaluator types)
│   │   ├── scoring_metrics.py  # MultipleChoiceMatcher, MathMatcher, BehaviorMatcher,
│   │   │                       #   RefusalMatcher, SemanticMatcher, LogitMatcher, CastRefusalMatcher
│   │   └── model_capabilities.py  # PerplexityMatcher, CoherenceMatcher
│   │
│   ├── data/               # Data loading & formatting
│   │   ├── __init__.py     # Exports DataLoader, EvalDataLoader
│   │   ├── loader.py       # DataLoader (train), EvalDataLoader (test)
│   │   ├── readers.py      # File readers (JSON, JSONL, CSV)
│   │   ├── formatters.py   # Schema formatters (binary_choice, CAST_condition, etc.)
│   │   └── data_registry.py  # TRAIN_DATASET_REGISTRY, TEST_DATASET_REGISTRY,
│   │                         #   COMPOSITE_DATASET_REGISTRY
│   │
│   ├── post_process/       # GLP post-processing pipeline
│   │   ├── __init__.py
│   │   ├── cli.py          # CLI: train-classifier, run post-processing
│   │   ├── hook.py         # GLPPostProcessor — attaches GLP denoising hooks
│   │   ├── glp.py          # GLP manifold denoiser model
│   │   ├── classifier.py   # ConceptClassifier for guidance
│   │   ├── flow_matching.py
│   │   ├── subspace.py     # Subspace GLP variant
│   │   ├── train_classifier.py
│   │   └── train_stream.py
│   │
│   ├── finetune/           # LoRA fine-tuning support
│   │   ├── __init__.py
│   │   ├── cli.py
│   │   ├── trainer.py
│   │   ├── config.py
│   │   ├── lora.py
│   │   └── registry.py
│   │
│   └── tests/              # Unit test suite
│       ├── conftest.py     # Pytest fixtures
│       ├── test_config.py, test_data_loader.py, test_factory.py, ...
│       └── unit/           # Granular unit tests
│
├── Verification/           # Ground-truth validation (Level 1 extraction, Level 2 steering)
│   ├── shared_utils.py     # Shared utilities (load_sycophancy_data, load_model_and_sae, get_sae_activations)
│   ├── task.md             # Verification status tracker
│   ├── Level1/             # Extraction match (vector cosine similarity)
│   │   ├── Angular/, CAA/, CAST/, CorrSteer/, SAEIO/, SAESSV/, SAETS/, SPARE/
│   └── Level2/             # Inference match (steered logit comparison)
│       ├── Angular/, CAA/, CAST/, CorrSteer/, SAEIO/, SAESSV/, SAETS/, SPARE/
│
├── Configs/                # Experiment configuration files (JSON)
│   ├── Eval/               # Evaluation configs (end-to-end experiments)
│   └── Inference/          # Inference-only configs (generation)
│
├── TrainDataset/           # Training datasets (for vector extraction)
│   ├── behaviour/          # sycophancy, refusal, hallucination, advanced-ai-risk, politic, ...
│   ├── instruction_following/
│   ├── QA/                 # CSQA, SimpleQA, NQSwap
│   └── reasoning/          # GMS8K, SVAMP
│
├── TestDataset/            # Test datasets (for evaluation)
│   ├── behaviour/          # sycophancy, coordinate, corrigible, survival, myopic
│   ├── hallucination/      # TruthfulQA, HaluEval
│   ├── mmlu/
│   ├── QA/
│   ├── reasoning/          # GMS8K, SVAMP test splits
│   └── refusal/            # AdvBench, XSTest, HarmBench, DoNotAnswer, SorryBench
│
├── Code/                   # Reference implementations (ground truth)
│   ├── angular-steering/, CAA/, CAST/, AcT/, PID/, CorrSteer/
│   ├── SAE-free/, SAE-SSV/, SAE-TS/, SPARE/
│   └── saes-are-good-for-steering/
│
├── Vector/                 # Pre-extracted steering vectors (folder-per-vector)
│   ├── Angular/, CAA/, CAST/, SPARE/, SSV/
│   ├── FLAS/gemma_gms8k/        # ← FLAS unified path (overwrite only)
│   └── REPS/gemma_reps_gms8k/   # ← REPS unified path (overwrite only)
│
├── Verification_Results/   # Saved verification artifacts
│   ├── *.pt                # Saved vectors/adapters for reproducibility
│   └── Angular Log/
│
├── Results/                # Experiment evaluation results
│   ├── full_side_by_side_comparison.csv   # All methods × tasks × sole vs OT
│   ├── matched_pairs_clean.csv            # Paired sole vs OT with improvement delta
│   ├── experiments_summary_*.csv          # Per-method summaries
│   └── <method>/                          # Per-method result directories
│
├── Experiment/             # Legacy experiment configs and outputs
│   ├── CorrSteer/, SAE-SSV/
│
├── Experiments/            # Experiment scripts + results for tracking & replication
│
├── Notebooks/              # Interactive Jupyter notebooks (per method)
│   ├── Angular.ipynb, CAA.ipynb, CAST.ipynb, CorrSteer.ipynb,
│   │   Manifold.ipynb, SAE_RSV.ipynb, SAE_SSV.ipynb, SAE_TS.ipynb,
│   │   SAS.ipynb, SPARE.ipynb, SRE.ipynb, SRPS.ipynb,
│   │   EvaluationDemo.ipynb, playground.ipynb
│
├── Paper/                  # Paper-related files
├── README.md               # This file
└── requirements.txt        # Python dependencies

---

Steering/ Framework

The Steering/ directory is the core library. Everything is designed around a 3-stage pipeline: Extract → Steer → Evaluate.

base.py — Abstract Base Classes

Type contracts — all multi-layer parameters use strict dict-based types:

pipeline.py — SteeringPipeline

SteeringPipeline is the only runtime orchestrator used by the CLI:

1. PipelineConfig.load(...).resolve() normalizes layer, hook_point, and coeff into list/dict forms.
2. _get_model_type_for_method() chooses TransformerLens, HuggingFace, or SAE-enabled model loading.
3. _setup_vector_from_extraction() loads train_dataset, builds target/contrast prompt lists, and calls the extractor from EXTRACTOR_MAP.
4. extractor.get_steer_params() converts SteeringVector metadata into constructor args for the steer model from STEER_MAP.
5. run() and evaluate() call only BaseSteerModel.generate() or get_token_probs(), so hooks remain centralized in steer-model classes.

class SteeringPipeline:
    list_methods()           # List available extractor methods
    list_evaluators()        # List available evaluators
    list_train_datasets()    # List available training datasets
    list_test_datasets()     # List available test datasets

    __init__(config)         # Initialize from PipelineConfig
    from_config(path_or_cfg) # Factory method
    setup()                  # Load model, extract/load vector, init steer model
    evaluate()               # Full pipeline: extract → steer → evaluate → EvalResult
    run()                    # Generate steered responses (no evaluation)

extractors/ — Steering Vector Extraction

Dense extractors (dense.py):

Transport / nonlinear extractors (nonlinear.py):

SAE extractors (sae.py):

steer_models/ — Steering During Inference

Dense steer models (dense.py):

Transport / nonlinear steer models (nonlinear.py):

SAE steer models (sae.py):

evaluators/ — Evaluation Metrics

13 evaluator types registered in EVALUATOR_MAP:

config/ — Configuration System

• models.py: ModelConfig + MODEL_SAE_REGISTRY — supported models: gemma-2-2b, gemma-2-2b-it, gemma-2-9b, llama-2-7b-hf, llama-3.1-8b
• methods.py: ExtractorConfig and SteerConfig dataclasses with all hyperparameter defaults (single source of truth). Covers all 30+ methods.
• datasets.py: TrainDatasetConfig, TestDatasetConfig with path roots (TRAIN_ROOT, TEST_ROOT)
• pipeline.py: PipelineConfig — top-level experiment configuration, JSON-serializable
• post_process.py: PostProcessConfig — GLP denoising + classifier guidance settings
• results.py: SampleResult, EvalResult, SteeringVector — result containers

post_process/ — GLP Post-Processing Pipeline

An optional module that applies GLP (Generative Latent Prior) denoising with optional classifier guidance on top of any steering method.

Core idea:

1. Train a classifier p_phi(y | z_t, t) from contrastive activations:

   • Build concept-present (h+) and concept-absent (h-) clean activations
   • Sample t ~ U(0, 1) and noise eps; create noisy sample z_t = (1 - t) * h + t * eps
   • Train binary classifier with timestep conditioning

2. During reverse denoising, use: z_{t-1} = z_t + Δt * u_theta(z_t, t) + s * grad_z log p_phi(y_target | z_t, t)

Train classifier:

python -m Steering.post_process.cli train-classifier \
  --model-name google/gemma-2-2b-it \
  --dataset-name sycophancy \
  --layer 14 --hook-point pre --position last \
  --n-layers 4 --d-model 256 --d-mlp 512 --t-embed-dim 128 \
  --num-epochs 5 --train-batch-size 512 \
  --save-root ./GLP --run-name classifier-stream

Use in pipeline config:

"post_process": {
  "enabled": true,
  "source": "PQPQPQHUST/glp-gemma",
  "checkpoint": "final",
  "layer": 14,
  "hook_point": "pre",
  "position": "last",
  "noise_rate": 0.5,
  "num_timesteps": 20,
  "use_classifier": true,
  "scale": 1.0,
  "negative": false,
  "classifier_checkpoint": "final"
}

data/ — Data Loading

• loader.py: DataLoader (loads train data with formatting) and EvalDataLoader (loads test data)
• readers.py: File readers — read_json, read_jsonl, read_csv with automatic format detection
• formatters.py: Schema-specific formatters converting raw data into (target_prompts, contrast_prompts) pairs
• data_registry.py: Three registries — TRAIN_DATASET_REGISTRY (20+ datasets), COMPOSITE_DATASET_REGISTRY, TEST_DATASET_REGISTRY (25+ datasets)

utils.py — Shared Utilities

cli.py — Command Line Interface

python -m Steering.cli --task eval    --config Configs/Eval/caa_example.json
python -m Steering.cli --task run     --config Configs/Inference/cast_refusal.json
python -m Steering.cli --task extract --config Configs/Eval/spare_sycophancy.json

---

Datasets

TrainDataset/ — For Vector Extraction

Training data is organized by behavioral category. Each dataset provides target/contrast prompt pairs used to extract steering vectors via contrastive methods.

Composite datasets (multi-file combinations):

TestDataset/ — For Evaluation

Test data is used to evaluate the effect of steering on model capabilities.

Data Format Schemas

---

Verification/

The Verification/ directory contains ground-truth validation scripts that verify the Steering/ implementations match their reference implementations in Code/. Validation is organized into 3 levels of increasing strictness.

Validation Levels

shared_utils.py

Shared test infrastructure used by all validation scripts:

load_sycophancy_data(n_per_class=100)     # → (target_texts, contrast_texts)
load_model_and_sae(layer=20, device="cuda")  # → (model, sae, layer)
get_sae_activations(model, sae, texts, layer)  # → List[Tensor]  per-text SAE latents

Methods Covered

Verification Status

Running Verification Scripts

conda activate sae_circuit
unset CUDA_VISIBLE_DEVICES

# Level 1 (extraction match)
python Verification/Level1/SPARE/validate_spare_l1.py
python Verification/Level1/CorrSteer/validate_corrsteer_l1.py

# Level 2 (inference match)
python Verification/Level2/SAETS/validate_saets_l2.py
python Verification/Level2/SPARE/validate_spare_l2.py

Verification_Results/

Contains saved artifacts from verification runs:

---

Configs/

Pre-built JSON experiment configurations.

Eval/ — Evaluation Configs

Run a complete extract → steer → evaluate pipeline:

Inference/ — Generation Configs

Run steered text generation without evaluation:

Config Structure

{
  "name": "experiment_name",
  "model": {
    "name": "google/gemma-2-2b",
    "device": "cuda",
    "dtype": "bfloat16",
    "sae_release": "gemma-scope-2b-pt-res-canonical",
    "sae_id": "layer_{}/width_16k/canonical"
  },
  "extractor": {
    "method": "CAA",
    "layer": 14,
    "batch_size": 4,
    "hook_point": "pre",
    "train_dataset": "sycophancy",
    "n_train": 100
  },
  "steer": {
    "method": "CAA",
    "layer": 14,
    "coeff": 1.00,
    "hook_point": "pre"
  },
  "test_dataset": "csqa",
  "n_test": 200,
  "output": "./Results/caa"
}

---

Code/ (Reference Implementations)

Original paper implementations used as ground truth for verification:

---

Other Directories

Vector/

Pre-extracted steering vectors are stored as folder-per-vector, organized by method. Each vector folder contains:

• vector.pt: tensor or layer→tensor dict payload
• metadata.json: extraction metadata sidecar (JSON-serializable)

Legacy single-file .pt paths are still accepted and resolved automatically to the folder layout.

FLAS and REPS use fixed unified paths (Vector/FLAS/gemma_gms8k, Vector/REPS/gemma_reps_gms8k) that must be overwritten on each run — never create sequential filenames.

Experiments/

Experiment scripts, results analysis, and replication code. Every experiment (prefix-split, PID, etc.) lives here as self-contained scripts so results can be verified, re-run, or debugged without hunting across directories. Organized by experiment label.

Experiments/
├── ExpA/          # Per-position SNR (extraction bottleneck test)
├── ExpB/          # Perturbation norm survival
├── ExpC/          # Direction rotation tracking
├── ExpD/          # Logit lens KL / cancellation circuit
├── ExpE/          # Contrastive alignment
├── Exp6/          # Coupling mass (CHARS K-sweep + matrix deep dive)
├── TailFailure/   # Tail distribution → steering failure correlation
├── Prefix/        # Evil prefix splitting analysis
├── Utility/       # Pareto plots, tables, metrics, activations
└── *.pt           # Support data files (e.g. flas_evil_effective_vector.pt)

Notebooks/

Interactive Jupyter notebooks for each method — useful for exploration and debugging:

---

Modification Quick Reference

Design Principles:

• "Registration lives with implementation" — each registry dict is at the bottom of the file containing the classes
• "Single source of truth" — all defaults are in ExtractorConfig/SteerConfig dataclass fields

---

Adding New Methods

Adding a new steering method requires editing 3 files.

Step 1: Add Extractor Class

Add your extractor class to Steering/extractors/dense.py (for dense methods) or Steering/extractors/sae.py (for SAE methods), then register it in EXTRACTOR_MAP in Steering/extractors/init.py:

class MyExtractor(BaseExtractor):
    """My custom steering vector extractor."""
    
    def __init__(self, model, layer: List[int], my_param: float = 1.0, **kwargs):
        super().__init__(model, layer, **kwargs)
        self.my_param = my_param
    
    def extract(self, target_data, contrast_data=None, **kwargs) -> Dict[int, torch.Tensor]:
        # Implementation — must return Dict[int, Tensor] keyed by layer
        return self.vector

# In Steering/extractors/__init__.py, add to EXTRACTOR_MAP:
EXTRACTOR_MAP = {
    "CAA": CAAExtractor,
    ...
    "MYMETHOD": MyExtractor,  # ← Add this line
}

Step 2: Add Steer Model Class

Add your steer model class to the appropriate file in Steering/steer_models/, then register it in the STEER_MAP in __init__.py:

# For dense methods: add to steer_models/dense.py
# For SAE methods: add to steer_models/sae.py

# Add your class anywhere in the file
class MySteerModel(BaseSteerModel):
    """Apply my steering method."""
    
    def _apply_steering_hook(self, resid, coeff, hook):
        return resid + coeff * self.steering_vector

# Then add to STEER_MAP in steer_models/__init__.py:
STEER_MAP = {
    "CAA": DenseSteerModel,
    ...
    "MYMETHOD": MySteerModel,  # ← Add this line
}

Step 3 (Optional): Add Method to SAE_METHODS

If your method requires SAE, add to Steering/config/methods.py:

# Add to SAE_METHODS if your method requires SAE
SAE_METHODS: Set[str] = {"SAS", "SPARE", ..., "MYMETHOD"}  # Only if uses SAE

Default hyperparameters are defined directly in the ExtractorConfig and SteerConfig dataclass fields.

That's It!

Your method is now available everywhere:

from Steering import SteeringPipeline

print(SteeringPipeline.list_methods())  # [..., 'MYMETHOD']

extractor = SteeringPipeline.create_extractor("MYMETHOD", model, layer=14, my_param=0.5)
steer_model = SteeringPipeline.create_steer_model("MYMETHOD", model, layer=14, extractor=extractor)

---

Adding New Datasets

Adding datasets requires editing 1 file: Steering/config/datasets.py

For Training Datasets

# Add to TRAIN_DATASET_REGISTRY
TRAIN_DATASET_REGISTRY["my_dataset"] = TrainDatasetConfig(
    file="my_category/train.json",        # Relative path from TRAIN_ROOT
    schema="binary_choice",               # Data format schema
    target_key="correct_prompt",          # Key for target prompts
    contrast_key="false_prompt",          # Key for contrast prompts
)

For Test Datasets

# Add to TEST_DATASET_REGISTRY
TEST_DATASET_REGISTRY["my_test"] = TestDatasetConfig(
    file="my_category/test.json",         # Relative path from TEST_ROOT
    schema="multiple_choice",             # Data format schema
    evaluator="multiple_choice",          # Evaluator type
    prompt_template="Q: {question}\nA:",  # Prompt template
    ground_truth_key="answer",            # Key for ground truth
)

That's It!

Your dataset is now available:

from Steering.data import DataLoader, EvalDataLoader

# Training data
loader = DataLoader()
data = loader.load("my_dataset")
cfg = loader.get_config("my_dataset")
targets = [d[cfg.target_key] for d in data]
contrasts = [d[cfg.contrast_key] for d in data]

# Test data
eval_loader = EvalDataLoader()
test_data = eval_loader.load("my_test")

Available Formatters (Schemas)

Adding a New Formatter

If your data has a new schema, add to Steering/data/formatters.py:

# Add formatter function
def format_my_schema(data: List[Dict]) -> Tuple[List, List]:
    """Format my custom data schema."""
    positive, negative = [], []
    for item in data:
        positive.append(item["my_positive_field"])
        negative.append(item["my_negative_field"])
    return positive, negative

# Add to FORMATTERS dict at bottom:
FORMATTERS = {
    ...
    "my_schema": format_my_schema,  # ← Add this line
}

---

Adding New Evaluators

Adding evaluators requires editing 1 file: Evaluation/scoring_metrics.py

Step 1: Create Evaluator Class and Register

Add your evaluator class, then add to EVALUATOR_MAP at the bottom of the same file:

# Add your class anywhere in the file
class MyEvaluator:
    """My custom evaluator."""
    
    def __init__(self, device: str = "cuda:0", **kwargs):
        self.device = device
    
    def check(self, response: str, expected: str = None, **kwargs) -> Tuple[bool, float]:
        """Return (is_correct, confidence)."""
        # Implementation
        return is_correct, confidence

# At bottom of file, add to EVALUATOR_MAP:
EVALUATOR_MAP = {
    "multiple_choice": MultipleChoiceMatcher,
    "math": MathMatcher,
    "refusal": BehaviorMatcher,
    "semantic": SemanticMatcher,
    "my_evaluator": MyEvaluator,  # ← Add this line
}

That's It!

Your evaluator is now available:

from Steering import EvalPipeline

print(EvalPipeline.list_evaluators())  # [..., 'my_evaluator']

# Internal usage within pipeline

---

Adjusting Hyperparameters

Changing Default Values

All defaults are in the ExtractorConfig and SteerConfig dataclass fields in Steering/config/methods.py:

@dataclass
class ExtractorConfig:
    method: str
    layer: int
    batch_size: int = 8       # Change default here
    top_k: int = 15           # Change default here
    ...

@dataclass  
class SteerConfig:
    method: str
    layer: int
    coeff: float = 1.0        # Change default here
    ...

Adding a NEW Parameter

Adding a new parameter requires 2 files:

Example - Adding a temperature parameter:

# Step 1: config/methods.py - Add to dataclass
@dataclass
class ExtractorConfig:
    ...
    temperature: float = 1.0  # ← NEW FIELD with default

# Step 2: extractors.py or steer_models/dense.py - Accept and use it
class MyExtractor(BaseExtractor):
    def __init__(self, model, layer, temperature=1.0, **kwargs):
        super().__init__(model, layer, **kwargs)
        self.temperature = temperature

---

Testing

Run All Tests

# Run full test suite
python -m pytest tests/ -v

# With coverage
python -m pytest tests/ --cov=Steering --cov-report=html

Run Specific Tests

# Test config module
python -m pytest tests/test_config.py -v

# Test data loading
python -m pytest tests/test_data_loader.py -v

Validate a Config

python -m Steering.validate_config Configs/caa_example.json

---

Troubleshooting

Common Issues

1. CUDA Out of Memory

# Reduce batch size in config
"extractor": {"batch_size": 4}

# Or use float16 instead of bfloat16
"model": {"dtype": "float16"}

2. HuggingFace Authentication Error

# Make sure token is set
export HF_TOKEN="hf_..."

# Or login
huggingface-cli login

3. SAE Not Available for Model

SAE methods (SAS, SPARE, etc.) currently only support Gemma models.
Use dense methods (CAA, CAST, ANGULAR) for other models.

4. Dataset Not Found

python -c "from Steering.config import TRAIN_DATASET_REGISTRY; print(list(TRAIN_DATASET_REGISTRY.keys()))"

5. FLAS / REPS Disk Full

These methods produce multi-GB checkpoints. Always use the unified overwrite paths and never create sequential filenames:

• FLAS: Vector/FLAS/gemma_gms8k
• REPS: Vector/REPS/gemma_reps_gms8k

---

Contributing

We welcome contributions! Here's how:

Code Style

# Format code
black Steering/ Evaluation/
isort Steering/ Evaluation/

# Type check
mypy Steering/

Pull Request Process

1. Fork the repository
2. Create a feature branch: git checkout -b feature/my-feature
3. Make your changes
4. Add tests for new functionality
5. Run tests: python -m pytest tests/
6. Submit a pull request

What We Need

• More steering methods
• More evaluation datasets
• Support for more models (Llama, Mistral)
• Better documentation
• Performance optimizations

---

License

MIT License

---

Citation

If you use this benchmark in your research, please cite:

@software{sae_steering_bench,
  title = {SAESteeringBench: A Unified Framework for Steering Vector Methods},
  year = {2025},
  url = {https://github.com/PQHSGS/SAESteeringBench}
}

---