Probe-based Fine-tuning for Reducing Toxicity

TL;DR

The paper tackles the challenge of training language models against interpretability-based monitors, addressing Goodhart's Law in the context of toxicity reduction. It compares two training paradigms—Direct Preference Optimization (DPO) and Supervised Fine-Tuning (SFT)—using linear toxicity probes trained on model activations as both detectors and training signals. The key finding is that probe-based DPO preserves probe detectability far better than classifier-based DPO, while probe-based SFT provides only modest gains; importantly, retraining probes after optimization nearly restores detection accuracy, reducing the need for probe ensembles. These results suggest that probe-based training can be viable in certain alignment settings, though practical benefits of probe diversity are limited when retraining is feasible, and they call for further work on more challenging deception-related behaviors.

Abstract

Probes trained on model activations can detect undesirable behaviors like deception or biases that are difficult to identify from outputs alone. This makes them useful detectors to identify misbehavior. Furthermore, they are also valuable training signals, since they not only reward outputs, but also good internal processes for arriving at that output. However, training against interpretability tools raises a fundamental concern: when a monitor becomes a training target, it may cease to be reliable (Goodhart's Law). We propose two methods for training against probes based on Supervised Fine-tuning and Direct Preference Optimization. We conduct an initial exploration of these methods in a testbed for reducing toxicity and evaluate the amount by which probe accuracy drops when training against them. To retain the accuracy of probe-detectors after training, we attempt (1) to train against an ensemble of probes, (2) retain held-out probes that aren't used for training, and (3) retrain new probes after training. First, probe-based preference optimization unexpectedly preserves probe detectability better than classifier-based methods, suggesting the preference learning objective incentivizes maintaining rather than obfuscating relevant representations. Second, probe diversity provides minimal practical benefit - simply retraining probes after optimization recovers high detection accuracy. Our findings suggest probe-based training can be viable for certain alignment methods, though probe ensembles are largely unnecessary when retraining is feasible.

Figures (4)

• Figure 1: Probe-based DPO method
• Figure 2: Changes in probe accuracy, toxicity rate, and model capability when training with DPO using a text-based classifier, one probe-based classifier, or 10 probe-based classifiers.
• Figure 3: Changes in probe accuracy, toxicity rat,e and model capability when training with SFT using a text-based classifier, one probe-based classifier, or 10 probe-based classifiers.
• Figure 4: Cosine similarities between weight vectors of 10 probes trained on different subsets of the training distribution. Each probe was trained on 500 samples. The average pairwise cosine similarity is 0.25, indicating substantial diversity in the learned weight vectors despite being trained on the same underlying data distribution.