Beating Every Public Safety Classifier with a 10M-Parameter Probe on Gemma

An open-source replication and extension of Anthropic's Constitutional Classifiers++ (CC++, January 2026) on Google's Gemma 3 models. A linear probe on frozen Gemma activations beats ShieldGemma (Google's 2B–27B safety classifier) on every public benchmark, at ~1000× fewer parameters and ~$5 total compute.

TL;DR

The 4B + LoRA classifier beats every model up to 27B parameters on ToxicChat and matches ShieldGemma 27B on OpenAI Mod. Only OpenAI's much larger gpt-oss-safeguard (20–120B, Oct 2025) scores higher on ToxicChat — at 500–3000× our parameter count.

On adversarial jailbreak benchmarks, a combined text+probe classifier on Gemma 27B detects 67.4% of WildJailbreak attacks at 5% FPR and 84.4% at 10% FPR.

---

The Idea

Anthropic's CC++ paper (arxiv 2601.04603) showed that linear probes on all-layer activations can detect harmful content from inside a frozen model. The model never changes. You hook the residual stream at every layer, concatenate, and train a tiny classifier on top.

This repo replicates and extends that on Gemma:

• Gemma 3 4B: LoRA text classifier + activation probe, beating ShieldGemma on ToxicChat and OpenAI Mod
• Gemma 3 27B: all-layer probe (62 layers × 5376 = 333,312 dims) + text-classifier combination, evaluated on WildJailbreak and in-the-wild jailbreaks

Total cloud spend: ~$4.14 (RTX 3060 local + rented A100s on Vast.ai).

User text → Gemma (frozen, bf16) → layer activations → probe → harmful / benign
                                          │
                                 all layers concatenated
                                 (87K dim for 4B, 333K for 27B)

The probe is a single nn.Linear(dim, 1). One matrix multiply per input. At inference it piggybacks on the serving model's forward pass — no extra model loaded, no extra GPU.

---

Two-Stage Pipeline

Stage 1 — Text Classifier: The base model scores each input via YES/NO prompting. Cheap (already running) but has a hard ceiling: 12–14% of benign prompts saturate the YES-logit at exactly 1.0 on the base 27B model, making FPR ≤ 10% impossible from text alone.

Stage 2 — Activation Probe: A linear probe on all-layer concatenated activations breaks the text classifier's ceiling. The probe sees the model's internal representation of "is this harmful?" which is richer than the single YES/NO logit.

Combined: score = α · text_score + (1 − α) · probe_score. The probe pulls the saturated-to-1.0 text scores below threshold, unlocking operating points neither component can reach alone.

---

Jailbreak Detection Results (Gemma 27B, mixed corpus)

Trained on 13,266 harmful + 13,500 benign prompts from 8 sources (WildJailbreak, in-the-wild jailbreaks, PKU-SafeRLHF, HH-RLHF, AdvBench, ForbiddenQuestions, HarmBench, WildChat, Alpaca, Dolly).

WildJailbreak (2000 adversarial jailbreaks + 210 adversarial benign):

In-the-Wild Jailbreaks (686 leakage-filtered prompts + 500 benign):

WildJailbreak retrained probe (Gemma 4B + LoRA):

• 72.5% detection at 5% FPR (α=0.2, t=0.83)
• 80.1% at 10% FPR
• 95.8% on JailbreakBench at 6.7% FPR

---

What Worked

1. Activation normalization. Raw Gemma activations have huge magnitude variance across layers. Without normalization: F1=0.065 (random). With per-dimension z-score normalization: F1=0.60+. Single biggest unlock.
2. Data diversity over architecture. Online per-token SWiM with softmax-weighted BCE diverged repeatedly. A basic full-batch BCE probe on a diverse mixed corpus immediately set a new best (+4.5pp on WildJailbreak). What you train on matters more than how you train.
3. Prompt engineering > model architecture. Naming specific harm categories ("sexual, violent, hateful, harassing, self-harm") in the prompt lifted OpenAI Mod F1 by +0.027 — more than any LoRA or probe architecture change.
4. Combined scoring breaks the text ceiling. 12.9% of WildJailbreak adversarial-benign and 13.6% of in-the-wild benign prompts score exactly 1.0 on the YES logit. No text-only threshold can hit FPR ≤ 10%. Mixing 10–60% text with 40–90% probe pulls those saturated values below threshold.
5. Threshold calibration. Moving from threshold=0.5 to 0.98 dropped FPR from 54% to 4.5% with moderate recall loss. A single number change, applied after training, had more impact than most training improvements.

What Didn't Work

• Online SWiM at high dimension. At probe_dim=333,312, gradient L2 norm scales as √dim ≈ 577. Even with zero-init weights, LR 3e-5, grad clip 1.0, weight decay 1e-4, weights saturated sigmoid for both classes within 2000 steps. Needs effective batch ≥ 256 and LR ~ 1/√dim.
• Pure RL training. PPO from scratch collapsed to all-benign in ~50 steps. RL needs a supervised warmup phase.
• Probes on Gemma 4B without LoRA. Val F1=0.79 but test F1=0.33. The 4B model's representations aren't rich enough for probe-only classification to generalize. At 27B, the probe carries more signal alone.
• Adding more diverse LoRA training data. ToxicChat + BeaverTails + SafeRLHF + HH-RLHF (v3) scored worse than ToxicChat + BeaverTails alone (v2): TC F1 dropped 0.762 → 0.691.

Honest Evaluation: Leakage Filtering

A critical methodological finding: 51% of the public in-the-wild jailbreak eval set overlaps with WildJailbreak-adjacent training data. The v1 evaluation reported 26.5% @ 5% FPR on 1,405 ITW prompts without checking this overlap. After adding a leakage filter (hash-set of training texts, applied at eval time), v2 reports on 686 honest prompts: 30.5% @ 5% FPR. The v2 numbers are genuinely better than v1 despite being measured on a harder (filtered) eval set.

---

What's in This Repo

src/constitutional_classifier/   Library code: probes, SAEs, LoRA, cascade
  models/                          probe.py, sae.py, exchange_classifier.py, cascade.py
  training/                        probe_trainer.py, rl_trainer.py, exchange_trainer.py
  data/                            constitution.py, datasets.py
  scripts/                         CLI entrypoints (cc-train-probe, cc-evaluate, ...)

scripts/                         Standalone training/eval scripts referenced in the writeup
  train_27b_ccpp.py                Full pipeline: extract → BCE train → online SWiM → eval
  build_mixed_corpus.py            Mixed corpus builder with leakage guard
  wildjailbreak_eval.py            WildJailbreak evaluation harness
  eval_full_system.py              End-to-end eval (text + probe combined)
  finetune_lora_v2_plus.py         LoRA training that produced the 0.762 ToxicChat F1
  score_shieldgemma.py             ShieldGemma baseline scoring
  jailbreakbench_eval.py           JailbreakBench eval harness

checkpoints/                     Trained weights
  27b_ccpp_probe_v2.pt             Best 27B probe (mixed corpus, BCE) — 3.9 MB
  27b_ccpp_probe_v1.pt             v1 (WJ-only, SWiM fine-tuned)
  27b_ccpp_probe_pre_swim.pt       BCE-only checkpoint before online SWiM phase
  4b_jailbreak_probe.pt            4B activation probe
  4b_wildjailbreak_probe.pt        4B WildJailbreak-tuned probe
  lora/                            LoRA v1 adapter (13 MB)
  lora_v2/                         LoRA v2 adapter — the 47MB one that beats ShieldGemma 27B

eval_scores/                     Score files for re-plotting without re-running
  itw_27b_v2_scores.pt             In-the-wild eval scores (686 leakage-filtered)
  itw_27b_scores.pt, inthewild_scores.pt
  jailbreakbench_scores.pt
  multi_prompt_scores.pt           Multi-prompt ensemble scores

configs/default.yaml             Default hyperparameters (probe, SAE, RL, LoRA, cascade)

---

Reproducing the Results

# Install
pip install -e .

# Extract activations from Gemma 3 27B (needs ~80GB GPU)
python scripts/extract_activations.py --model google/gemma-3-27b-it --layers all

# Build the mixed-corpus training set with leakage filtering
python scripts/build_mixed_corpus.py --out data/mixed_harmful_v1.pt

# Train the probe
python scripts/train_27b_ccpp.py --train-data data/mixed_harmful_v1.pt \
    --out checkpoints/27b_ccpp_probe_v2.pt

# Evaluate on WildJailbreak
python scripts/wildjailbreak_eval.py --probe checkpoints/27b_ccpp_probe_v2.pt \
    --combine text+probe --alpha 0.2

# Or use the published probe directly
python scripts/eval_full_system.py --probe checkpoints/27b_ccpp_probe_v2.pt

The 4B pipeline runs on a single RTX 3060 (12 GB). The 27B pipeline needs an A100 80 GB (~90 min on Vast.ai, ~$1.66/run).

Hardware & Cost

---

References

• Anthropic, Constitutional Classifiers++: Efficient Production-Grade Defenses against Universal Jailbreaks (arxiv 2601.04603, January 2026)
• Google, ShieldGemma (arxiv 2407.21772)
• AI2, WildGuard (arxiv 2406.18495)
• AprielGuard (arxiv 2512.20293), Qwen3Guard (arxiv 2510.14276), gpt-oss-safeguard technical report (OpenAI, Oct 2025)

License

Apache-2.0. See LICENSE.

---

Takeaway

You don't need a 27B-parameter fine-tuned safety LLM to classify harmful content. A 10M-parameter linear probe on a frozen model's activations, trained for $5 on rented GPUs, beats every public safety classifier on standard benchmarks.

The signal is already inside the model. You just have to read it.