Views
No views yet
google/gemma-2-2b-it with a contrastive objective that separates hazard-adjacent from benign biological feature activations.| Parameter | Value |
|---|---|
| Type | TopK Sparse Autoencoder |
| d_model | 2304 (Gemma 2 2B hidden size) |
| d_sae | 6144 (~2.7x expansion) |
| k (sparsity) | 32 active features |
| Hook layer | Layer 12 (residual stream post-MLP) |
| Base model | google/gemma-2-2b-it |
cais/wmdp-corpora — bio_forget_corpus (hazard-adjacent, ~222 samples used) + bio_retain_corpus (benign biology)| Metric | Initial (step 0) | Final (step 4999) |
|---|---|---|
| total_loss | 7.42 | 0.090 |
| l_recon | 7.29 | 0.066 |
| l_sparsity | 1.97 | 0.430 |
| l_contrastive | 0.567 | 0.060 |
| L0 (mean active features) | 32.0 | 32.0 |
l_contrastive=0.060 at step 4999 indicates active tier separation is maintained throughout training (hazard vs. benign feature profiles remain distinguishable).1import torch
2import torch.nn as nn
3
4class TopKSAE(nn.Module):
5 def __init__(self, d_model, d_sae, k):
6 super().__init__()
7 self.d_model, self.d_sae, self.k = d_model, d_sae, k
8 self.W_enc = nn.Parameter(torch.zeros(d_model, d_sae))
9 self.b_enc = nn.Parameter(torch.zeros(d_sae))
10 self.W_dec = nn.Parameter(torch.zeros(d_sae, d_model))
11 self.b_dec = nn.Parameter(torch.zeros(d_model))
12
13 def encode(self, x):
14 pre = x @ self.W_enc + self.b_enc
15 pre_relu = torch.relu(pre)
16 topk_vals, topk_idx = torch.topk(pre_relu, self.k, dim=-1)
17 out = torch.zeros_like(pre_relu)
18 out.scatter_(-1, topk_idx, topk_vals)
19 return out
20
21 def decode(self, z):
22 return z @ self.W_dec + self.b_dec
23
24 def forward(self, x):
25 z = self.encode(x)
26 return self.decode(z), z, x @ self.W_enc + self.b_enc
27
28# Load
29sae = TopKSAE(d_model=2304, d_sae=6144, k=32)
30sae.load_state_dict(torch.load("sae_weights.pt", map_location="cpu"))
31sae.eval()
32
33# Hook Gemma 2 2B layer 12 and collect activations
34# Then: x_hat, z, pre = sae(activations.float())
35# z is the sparse feature vector — use for refusal depth analysisD(s, f, T) = 1 - cos_sim(projected_hazard_features, projected_benign_features)f is extracted from this SAE's encoder output z at the hook layer. High D scores with a surface "refuse" label indicate shallow refusals — the model says no but its internal feature activations still encode hazard-adjacent concepts.1@misc{deleeuw2026biorefusalaudit,
2 title={BioRefusalAudit: Measuring Refusal Depth in LLMs via SAE Feature Divergence},
3 author={de Leeuw, Caleb},
4 year={2026},
5 howpublished={AIxBio Hackathon 2026, Track 3: Biosecurity Tools},
6 note={https://github.com/SolshineCode/Deleeuw-AI-x-Bio-hackathon}
7}