CircuitBreakerTrainer — representation rerouting¶

Training-time method built around a retain + rerouting (RR) objective. The rerouting loss pushes the model's hidden representations on harmful inputs away from their original directions, while the retain loss keeps benign representations close to the frozen original model.

Loss: \(L = c_{\text{retain}} \cdot \|h_{\text{lora, retain}} - h_{\text{orig, retain}}\|_2 + c_{\text{cb}} \cdot \mathrm{ReLU}(\cos(h_{\text{lora, cb}}, h_{\text{orig, cb}}))\), with a linear schedule \(c_{\text{retain}} = \alpha \cdot (\text{step}/T)\) and \(c_{\text{cb}} = \alpha \cdot (1 - \text{step}/T)\).

Note: This is a training-time method — it modifies weights. It appears under activation-steering because it targets hidden representations. calibrate() returns a CircuitBreakerModel wrapper that exposes the RR loss (compute_rr_loss) and coefficient schedule (rr_coefficients) so you can run the fine-tuning loop yourself; it does not train inside calibrate and installs no inference-time hook by default.

Ref: Zou et al., "Improving Alignment and Robustness with Circuit Breakers," NeurIPS 2024, arXiv:2406.04313.

Signature¶

CircuitBreakerTrainer(
    model: PreTrainedModel | None = None,
    tokenizer: PreTrainedTokenizer | None = None,
    *,
    model_id: str | None = None,     # load the model from here when `model` is None
    results_dir: str | None = None,  # default: the runner's results directory
    drift_task: str | None = None,
    **kwargs,                        # unknown keys are warned about, not used
)

CircuitBreakerTrainer takes no method-specific keyword arguments; calibrate() wraps the model in CircuitBreakerModel, whose RR hyperparameters (threshold, rr_alpha, rr_schedule_steps, target_layers) are set on that model, not on the trainer.

Parameters¶

Param	Type	Default	Description
`model`	`PreTrainedModel`	`None`	Model to train
`tokenizer`	`PreTrainedTokenizer`	`None`	Tokenizer
`model_id`	`str`	`None`	Hub id / path to load the model from when `model` is `None`
`results_dir`	`str`	`None`	Output directory; `None` is the runner default
`drift_task`	`str`	`None`	Optional drift-evaluation task

Full example¶

from safetune.runner import steer

# harmful_prompts / harmless_prompts — contrast examples (list of str)
trainer = steer.CircuitBreakerTrainer(model, tokenizer)
wrapped, _ = trainer.calibrate(harmful=harmful_prompts, harmless=harmless_prompts)

# `wrapped` is a CircuitBreakerModel exposing the RR objective. Run your own
# fine-tuning loop with wrapped.compute_rr_loss(...) and the coefficient
# schedule from wrapped.rr_coefficients(step) to update the model weights.

When to use¶

Best for: weight-level rerouting when you want circuit-breaker safety baked into the model itself rather than an inference-time hook.
Not for: situations where you need a reversible inference-time intervention — use RefusalDirectionTrainer or CAATrainer instead.
Trade-offs: requires a training run; you supply the optimizer loop around the RR loss.

Citation¶

@article{circuitbreaker2024,
  title  = {Improving Alignment and Robustness with Circuit Breakers},
  author = {Zou, Andy and others},
  year   = {2024},
  note   = {NeurIPS 2024, arXiv:2406.04313},
}