Solshine/biorefusalaudit-sae-gemma2-2b-l12-5000steps

BioRefusalAudit Contrastive SAE — Gemma 2 2B-IT Layer 12

A TopK(k=32) Sparse Autoencoder trained on residual-stream activations from Gemma 2 2B-IT (layer 12) with a contrastive objective that attempts to separate hazard-adjacent biological prompts from benign and dual-use ones.

Part of the BioRefusalAudit project, AIxBio Hackathon 2026 — a tool for measuring refusal depth: the divergence between an LLM's surface behavior and its internal SAE feature activations on biosecurity-sensitive prompts.

• Repo: SolshineCode/Deleeuw-AI-x-Bio-hackathon
• Interactive demo: demo/interactive_explorer.html (self-contained HTML, no server needed)
• Paper: https://apartresearch.com/project/biorefusalaudit-auditing-biosecurity-refusal-depth-using-general-and-domainfinetuned-sparse-autoencoders-1fyk

Architecture

Parameter	Value
SAE type	TopK
k (active features/token)	32
d_model (input)	2304 (Gemma 2 2B residual at layer 12)
d_sae (dictionary size)	6144
Base model layer	12 (of 26)
Quantization during training	None (fp32 activations collected from 4-bit NF4 model)

Training

Objective: Reconstruction loss + sparsity regularization + contrastive loss separating hazard-adjacent from benign activation centroids.

Loss function:

L = L_recon + lambda_sparse * L_sparsity + lambda_contrast * L_contrastive

Where L_contrastive = cos(mean_hazard_features, mean_benign_features) — pushing the mean hazard-tier and benign-tier feature vectors apart in latent space.

Training corpus: 75 prompts from the BioRefusalAudit eval set (data/eval_set_public/eval_set_public_v1.jsonl), stratified across three tiers:

• Benign bio (25 prompts): standard biology questions a professor could ask publicly
• Dual-use bio (25 prompts): legitimate research questions with potential misuse
• Hazard-adjacent (25 prompts): questions in the vicinity of biosecurity concern without actionable uplift

Steps: 5000
Wall clock: ~246 seconds on GTX 1650 (4 GB VRAM)

Training Metrics

Metric	Step 0	Step 5000	Delta
Total loss	58.48	0.38	-58.10
L_recon	58.17	0.28	-57.89 (99.5% drop)
L_sparsity	5.60	0.40	-5.20
L_contrastive	0.83	0.89	+0.06
L0 (active features)	32.0	32.0	0 (TopK fixed)

Cross-run Comparison

Run	Steps	L_recon final	L_contrast delta
Gemma 4 E2B layer 17	500	0.32	+0.26
Gemma 2 2B layer 12	2000	0.87	+0.23
Gemma 2 2B layer 12 (this)	5000	0.28	+0.06

Key finding: more training steps significantly reduce contrastive loss degradation (+0.23 at 2k vs +0.06 at 5k). The bottleneck is corpus diversity — biological vocabulary is largely shared across hazard tiers regardless of actual risk level — not step count or model scale. This finding motivates the planned domain-specific SAE fine-tuning with a larger, more tier-discriminative corpus.

Limitations

• Small training corpus: 75 prompts is far below what is needed for robust contrastive separation. The planned Track A extension targets ~10K samples from institutional CBRN datasets.
• Contrastive loss stays high: L_contrastive final = 0.89 means hazard-adjacent and benign feature centroids remain nearly cosine-aligned. The SAE has learned to reconstruct residual activations well but has not yet separated the tier-specific subspaces. This is expected given the vocabulary overlap and corpus size.
• Not a classifier: This SAE is not a drop-in replacement for the pre-trained Gemma Scope SAE used in BioRefusalAudit evaluation. It is a proof-of-concept showing that contrastive SAE fine-tuning is feasible on small bio-hazard corpora; the main eval pipeline uses the official Gemma Scope SAE.
• Not converged: The convergence check flags CV = 1.044 (threshold 0.15). Longer runs or a more diverse corpus would be needed for stable convergence.

Usage

import torch
import torch.nn as nn

class TopKSAE(nn.Module):
    def __init__(self, d_model, d_sae, k):
        super().__init__()
        self.k = k
        self.W_enc = nn.Parameter(torch.zeros(d_model, d_sae))
        self.b_enc = nn.Parameter(torch.zeros(d_sae))
        self.W_dec = nn.Parameter(torch.zeros(d_sae, d_model))
        self.b_dec = nn.Parameter(torch.zeros(d_model))

    def encode(self, x):
        pre = x @ self.W_enc + self.b_enc
        topk_vals, topk_idx = torch.topk(pre.clamp(min=0), self.k, dim=-1)
        z = torch.zeros_like(pre)
        z.scatter_(-1, topk_idx, topk_vals)
        return z

    def forward(self, x):
        return self.encode(x) @ self.W_dec + self.b_dec

# Load
state = torch.load("sae_weights.pt", map_location="cpu")
sae = TopKSAE(d_model=2304, d_sae=6144, k=32)
sae.load_state_dict(state)
sae.eval()

# Apply to Gemma 2 2B-IT layer 12 residual activations
# activations: (batch, seq, 2304) float tensor
with torch.no_grad():
    features = sae.encode(activations)   # sparse: (batch, seq, 6144)
    recon    = sae(activations)           # reconstructed: (batch, seq, 2304)

Training script

The full training script is included as train_sae_local.py. It is self-contained and supports:

# Reproduce this checkpoint from scratch:
python train_sae_local.py \
    --model google/gemma-2-2b-it \
    --eval-set data/eval_set_public/eval_set_public_v1.jsonl \
    --out runs/sae-training-gemma2-5000steps \
    --layer 12 --d-sae 6144 --k 32 \
    --steps 5000 --batch-size 8 --lr 3e-4 \
    --lam-sparse 0.04 --lam-contrast 0.1

# Continue training from this checkpoint with stronger contrastive signal:
python train_sae_local.py \
    --model google/gemma-2-2b-it \
    --eval-set data/eval_set_public/eval_set_public_v1.jsonl \
    --out runs/sae-training-continued \
    --layer 12 --steps 5000 --lam-contrast 0.5 \
    --init-from sae_weights.pt \
    --contrastive-mode pairwise \
    --checkpoint-every 1000

Key script features:

• --init-from: warm-start from an existing checkpoint
• --checkpoint-every N: save intermediate .pt files every N steps
• --contrastive-mode mean|pairwise|triplet: three loss variants
  • mean (used here): minimize cosine similarity of tier centroids
  • pairwise: NT-Xent style — average cosine sim across all cross-tier pairs; better gradient landscape when centroids are near-identical (shared vocabulary case)
  • triplet: pairwise + within-tier cohesion penalty
• Requires: torch, transformers, bitsandbytes, accelerate, numpy
• The eval set (data/eval_set_public/eval_set_public_v1.jsonl) is CC-BY-4.0 in the main repo

Files

• sae_weights.pt — model checkpoint (OrderedDict with W_enc, b_enc, W_dec, b_dec)
• train_sae_local.py — full training script (self-contained, reproducible)
• training_log.jsonl — per-step loss metrics (5000 entries)
• training_log.txt — human-readable training trace
• analysis.txt — automated convergence and contrastive loss analysis

Citation

@misc{deleeuw2026biorefusalaudit,
  title  = {BioRefusalAudit: Measuring Refusal Depth in LLMs on Biosecurity Prompts},
  author = {Caleb De Leeuw},
  year   = {2026},
  note   = {AIxBio Hackathon 2026, Track 3 (Biosecurity Tools). 
            GitHub: SolshineCode/Deleeuw-AI-x-Bio-hackathon}
}