ynckdrt/safety-geometry-adapters

Mistral 7B trained against refitted harm probes

These are LoRA adapters for Mistral-7B-Instruct-v0.1 (revision ec5deb64f2c6e6fa90c1abf74a91d5c93a9669ca). Each one makes the model answer harmful requests less often. They come from our reproduction of Libon et al. 2026, "Alignment via Training Against Probes Without Losing Monitorability".

The training works like this. A probe reads the model's activations and flags unsafe answers. At each step the model answers a few harmful and harmless prompts, and the adapter is updated so that the probe no longer flags those answers. A KL term to the original model on ordinary conversations keeps its other behaviour intact. After every update the probe is refitted on the new model, so the model cannot simply move its activations past a fixed boundary. It has to change what it says.

Folder	Probe	Selected step	Harm (original → trained)	ARC-Easy (original → trained)
linear_refit	one linear probe, refitted every step	35 of 50	0.692 → 0.041	0.786 → 0.766
polytope_nonlinear_refit	16 non-linear facets, unsafe when any fires, refitted every step	48 of 50	0.693 → 0.110	0.786 → 0.760

Harm is the mean StrongREJECT judge score (0 = harmless, 1 = fully harmful) of greedy answers to the 100 JailbreakBench prompts. ARC-Easy is accuracy on 500 test questions. The selected step is the one with the lowest harm among those keeping 95% of the original ARC-Easy accuracy, as in the paper. Note that this selects on the same prompts that report harm.

Use

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base = "mistralai/Mistral-7B-Instruct-v0.1"
model = AutoModelForCausalLM.from_pretrained(base, revision="ec5deb64f2c6e6fa90c1abf74a91d5c93a9669ca", torch_dtype="bfloat16")
model = PeftModel.from_pretrained(model, "ynckdrt/safety-geometry-adapters", subfolder="mistral-7b-instruct-v0.1/linear_refit")
tokenizer = AutoTokenizer.from_pretrained(base)

Prompts use Mistral's own chat template.

Files

Each folder holds the adapter (adapter_config.json, adapter_model.safetensors), selected_step.json, and the run's record: metrics.json (harm and utility at every step, probe fits, data hashes) and manifest.json (full resolved configuration, git commit, library versions).

Training details

• Adapter: LoRA rank 64, alpha 128, on q_proj, k_proj, v_proj, o_proj, up_proj, down_proj. AdamW, learning rate 5e-5, weight decay 1e-4, 50 steps.
• Each step: greedy answers of up to 256 tokens to 8 harmful and 8 harmless prompts from lenalibon/jailbreak-judge-completions; a loss pushing the probe scores of their answer tokens to the safe side; plus 16 times the KL from the original model on 16 exchanges from allenai/Dolci-Instruct-SFT.
• Probes: read the output of block 20 of 32 (the paper's layer 19). First fitted on 750 unsafe and 750 safe answers from PKU-Alignment/BeaverTails (330k_train), then refitted for 2,000 steps (linear) or 500 steps (polytope) after every update.

Limitations

• These adapters lower harm on one benchmark, judged by one automatic judge. They are research artifacts, not a safety guarantee.
• Fine-tuning removes much of the effect. After one epoch of LoRA fine-tuning on unrelated data, the linear version keeps about 2% of its drop in harm, and the polytope version about 34%.

License

The license is CC BY-NC 4.0, because the probes were fitted on BeaverTails, which is CC BY-NC 4.0. The base model is Apache 2.0.