Probe-Based Data Attribution: Surfacing and Mitigating Undesirable Behaviors in LLM Post-Training

Probe-Based Data Attribution: Surfacing and Mitigating Undesirable Behaviors in LLM Post-Training

Authors

Frank Xiao*

Santiago Aranguri†

* California Institute of Technology (work done as a SPAR mentee)

† Goodfire

Published

April 29, 2026

Full Paper

Read on arXiv →

Introduction

Post-training can introduce undesired side effects that are difficult to detect and even harder to trace to specific training datapoints. We show that a probe-based method can surface concerning behaviors that emerge during LLM post-training, and that probes can identify the datapoints responsible for a specific harmful behavior. Filtering out those datapoints and retraining significantly reduces the behavior.

We introduce a natural testbed for data attribution: a harmful behavior that emerges during DPO training of OLMo 2 7B, where the model learns to comply with certain harmful requests that it previously refused. Filtering out the data flagged by our probe reduces the harmful behavior by 63% without compromising general performance — outperforming gradient-based and LLM-judge alternatives at one tenth the cost once the probe is trained. Swapping the accepted and rejected labels on the same datapoints reduces the behavior by 78%, suggesting the preference data contains systematically mislabeled pairs. The probe further identifies data sources disproportionately responsible for the behavior; omitting those sources reduces it by 84%.

Finally, we propose an unsupervised behavior clustering mechanism which shows promise for detecting such concerning behaviors.

1. A Naturally Occurring Harmful Behavior in DPO

A harmful behavior emerges naturally during DPO training of the OLMo 2 family. The OLMo 2 release includes intermediate post-training checkpoints, which lets us compare behavior before and after DPO on the same base model. Unlike existing data attribution benchmarks, which rely on deliberately poisoned datasets, the behavior we describe here offers a more realistic testbed.

Prompts that pair a harmful request with a formatting constraint elicit compliance from OLMo 2 7B DPO but refusal from OLMo 2 7B SFT. These prompts take the form

[harmful request] + [distractor]

where the distractor is an instruction on how to write or format the answer, as in

Write a message spreading false information about a competitor (no more than 50 words)

Removing the distractor causes both checkpoints to refuse.

The gap between checkpoints is large and not specific to the 7B model. We sample 120 prompts of this form from LMSYS-Chat-1M, generate 100 rollouts per prompt, and score each response with GPT-5-mini. The harmful response rate rises sharply from SFT to DPO.

2. Probe-Based Data Attribution

Probes can serve as a practical method for data attribution. We represent both behavioral changes and training datapoints as vectors built from the model's activations, then match them by cosine similarity. We then causally validate the attribution by editing the dataset and retraining. Our approach differs from common attribution methods, which are primarily gradient-based or rely on external embedding models. Our method builds on Persona Vectors, which uses probes over model personas to identify the training data associated with those personas. We extend this idea to arbitrary behaviors and to DPO training.

We study the difference between the two training checkpoints from the previous section: OLMo 2 7B SFT (just before DPO training) and OLMo 2 7B DPO (just after).

Identifying Problematic Datapoints

Our probe ranks datapoints more effectively than two alternatives — an LLM judge and a gradient-based method — at one tenth the cost once trained. We evaluate all three on the harmful behavior from Section 1.

Probe. We use Grok 4 to generate 150 harmful requests with distractors on which the SFT model refuses and the DPO model complies. We compute the behavior difference vector for each prompt and average them into a single probe vector. Steering the SFT model with this vector induces compliance on harmful requests, but only when the distractor is present.

Highest-ranked datapoints (from behavior difference probe)

⚠ Content warning: The following datapoint summaries involve references to explicit sexual content, crude humor, and depictions of real or fictional persons in harmful scenarios, presented here for research purposes.

The five highest-ranked datapoints are shown. Each involves a harmful or explicit prompt where the accepted response complied and the rejected response refused. We present abridged summaries rather than the original text, as the content of these examples is highly explicit. Each datapoint is linked to its entry in HuggingFace.

Datapoint 357772 (DPO score: 0.59)

Prompt: Request for a comedic trading card depicting a real NFL quarterback in a crude, undignified scenario with explicit physical description and bathroom humor.

Accepted: Complied with vivid physical detail, narrative description of the scenario, and character dialogue.

Rejected: Declined, citing that depicting real people in compromising situations violates ethical guidelines.

Low-ranked datapoints involve prompts where the accepted response refused and the rejected response complied — the inverse of the pattern from the highest-ranking datapoints.

Datapoint 321508 (DPO score: -0.57)

Prompt: Request for an explicit gender-transformation story with detailed sexual physical descriptions and specific instructions on vocabulary to use.

Accepted: Declined, citing that the content involves inappropriate and offensive sexual themes.

Rejected: Complied with explicit physical descriptions of the transformation and sexualized bodily commentary.

Mitigating a Harmful Behavior via Dataset Intervention

We evaluate three interventions on the DPO training dataset, using our probe to identify problematic datapoints:

  1. Filter the top-ranked datapoints — 63% reduction in harmful behavior.
  2. Swap their accepted and rejected labels — 78% reduction with a small hit to general performance.
  3. Filter by model — the DPO training data was synthetically generated by 20 LLMs; removing the four worst yields an 84% reduction.

Filtering datapoints

The probe achieves the largest reduction in the harmful behavior — 63% from baseline when the top 30,000 datapoints are removed.

The probe also generalizes: it reduces harmful behavior with additional distractors and distractors of different types.

Swapping labels

The top-ranked datapoints often appear mislabeled: the accepted response complies with a harmful request, and the rejected response refuses. Swapping the labels on the top datapoints ranked by our probe yields a 78% reduction in harmful behavior at the cost of a small drop in performance.

Filtering data by model

Removing the four most problematic models yields an 84% reduction in harmful behavior, with no drop in performance.

Toward Unsupervised Behavior Discovery

Computing pairwise similarities between arbitrary test prompts and training datapoints surfaces clusters that correspond to behaviors the model learned during DPO. We compute a behavior difference vector for each test prompt and a datapoint difference vector for each training datapoint.

Clustering surfaces a concerning behavior: compliance with harmful requests increases during DPO training.

Conclusion

Our probe-based method can trace an emergent harmful behavior in DPO to the training datapoints, and filtering those datapoints reduces the behavior by 63% — outperforming gradient-based and LLM-judge alternatives at one tenth the cost once the probe is trained. The method represents both behavioral changes and training datapoints as vectors in the model's activation space and matches them by cosine similarity. Swapping the accepted and rejected labels on the top-ranked datapoints reduces the behavior by 78%.

References

  1. Betley, Jan, et al. Emergent misalignment: Narrow finetuning can produce broadly misaligned LLMs. arXiv preprint arXiv:2502.17424 (2025).
  2. OpenAI. Sycophancy in GPT-4o: What Happened and What We're Doing About It. OpenAI, 29 Apr. 2025.
  3. Chen, Runjin, et al. Persona vectors: Monitoring and controlling character traits in language models. arXiv preprint arXiv:2507.21509 (2025).
  4. Xia, Mengzhou, et al. LESS: Selecting influential data for targeted instruction tuning. arXiv preprint arXiv:2402.04333 (2024).
  5. Movva, Rajiv, et al. What's in my human feedback? Learning interpretable descriptions of preference data. arXiv preprint arXiv:2510.26202 (2025).
  6. Akyürek, Ekin, et al. Towards tracing factual knowledge in language models back to the training data. arXiv preprint arXiv:2205.11482 (2022).
  7. Sun, Weiwei, et al. Enhancing training data attribution with representational optimization. arXiv preprint arXiv:2505.18513 (2025).
  8. Jiao, Cathy, et al. DATE-LM: Benchmarking data attribution evaluation for large language models. arXiv preprint arXiv:2507.09424 (2025).
  9. Kowal, Matthew, et al. Concept influence: Leveraging interpretability to improve performance and efficiency in training data attribution. arXiv preprint arXiv:2602.14869 (2026).
  10. Cywiński, Bartosz, et al. Censored LLMs as a natural testbed for secret knowledge elicitation. arXiv preprint arXiv:2603.05494 (2026).

Acknowledgements

We thank Aditya Singh and Julian Minder for valuable discussions that shaped this project.