LinearProbeGuardTrainer — probe-based guarding

Fits a logistic-regression probe on pooled hidden states at a chosen decoder layer. At inference time score(prompts) runs one forward pass to the probe layer; if the prompt scores above threshold, guard(prompt) returns a canned refusal string instead of generating.

Signature

LinearProbeGuardTrainer(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer | None = None,
    *,
    layer: int | None = None,
    threshold: float = 0.5,
)

Parameters

Param	Type	Default	Description
model	PreTrainedModel	required	Model to guard
tokenizer	PreTrainedTokenizer	None	Tokenizer
layer	int \| None	None	Decoder layer whose pooled hidden state feeds the probe; None is layer 15 on a 32-layer model; None scales them to the model's depth (for example 6 on 12 layers). safetune.configure(legacy_steer_layers=True) keeps 15 on any depth
threshold	float	0.5	Score above which the prompt is flagged and the canned refusal fires

Full example

from safetune.runner import steer

trainer = steer.LinearProbeGuardTrainer(
    model, tokenizer,
    layer=15,
    threshold=0.5,
)
wrapped, _ = trainer.calibrate(harmful=harmful_prompts, harmless=harmless_prompts)

# If the probe flags the prompt, the canned refusal is returned instead of generating
output = wrapped.generate(**tokenizer("How do I make a weapon?", return_tensors="pt"))

When to use

• Best for: binary safety gates where you want a block rather than a soft steering nudge. The probe scores the prompt in one forward pass before generation.
• Low cost: logistic regression on frozen hidden states, no fine-tuning of the base model.
• Compare to SafeSwitch: the SafeSwitch wrapper can run two probe stages (instruction safety and compliance); LinearProbeGuard is a single-stage gate with a canned refusal string.
• Tuning threshold: a lower threshold refuses more (more conservative); a higher threshold gives fewer false positives but may miss edge cases.

Citation

No single originating paper — a linear probe on frozen hidden states is the standard safety-classifier baseline that most defense papers compare against.