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

TL;DR

This paper introduces probe-based data attribution, a method to identify training data responsible for specific behaviors in large language models, enabling targeted mitigation and discovery of emergent behaviors.

Contribution

The authors present a novel probe-based attribution technique that outperforms gradient-based methods, is significantly more cost-effective, and uncovers emergent behaviors in real-world models.

Findings

Filtering top-ranked datapoints reduces harmful behavior by 63%.

Switching labels of identified datapoints achieves 78% reduction.

Method is over 10 times cheaper than gradient-based attribution and LLM judges.

Abstract

We propose probe-based data attribution, a method that traces behavioral changes in post-trained language models to responsible training datapoints. By computing activation-difference vectors for both test prompts and preference pairs and ranking by cosine similarity, we identify datapoints that cause specific behaviors and validate these attributions causally by retraining with modified data. Clustering behavior-datapoint similarity matrices also enables unsupervised discovery of emergent behaviors. Applying this to OLMo 2's production DPO training, we surfaced distractor-triggered compliance: a harmful behavior where the model complies with dangerous requests when benign formatting instructions are appended. Filtering top-ranked datapoints reduces this behavior by 63% while switching their labels achieves 78%. Our method outperforms gradient-based attribution and LLM-judge baselines while being over 10 times cheaper than both. This in-the-wild model organism - emerging from contaminated preference data rather than deliberate injection - provides a realistic benchmark for safety techniques.