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

arXiv: https://arxiv.org/abs/2602.11079 · HTML: https://mlsanity.com/paper/2602.11079/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 32 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Builds on (most cited first, 21 of 21)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Constitutional AI: Harmlessness from AI Feedback](https://mlsanity.com/paper/2212.08073/) (2022)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Jailbroken: How Does LLM Safety Training Fail?](https://mlsanity.com/paper/2307.02483/) (2023)
- [Representation Engineering: A Top-Down Approach to AI Transparency](https://mlsanity.com/paper/2310.01405/) (2023)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [Refusal in Language Models Is Mediated by a Single Direction](https://mlsanity.com/paper/2406.11717/) (2024)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Steering Language Models With Activation Engineering](https://mlsanity.com/paper/2308.10248/) (2023)
- [Discovering Language Model Behaviors with Model-Written Evaluations](https://mlsanity.com/paper/2212.09251/) (2022)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Sleeper Agents: Training Deceptive LLMs that Persist Through Safety Training](https://mlsanity.com/paper/2401.05566/) (2024)
- [XSTest: A Test Suite for Identifying Exaggerated Safety Behaviours in Large Language Models](https://mlsanity.com/paper/2308.01263/) (2023)
- [LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset](https://mlsanity.com/paper/2309.11998/) (2023)
- [Persona Vectors: Monitoring and Controlling Character Traits in Language Models](https://mlsanity.com/paper/2507.21509/) (2025)
- [AI Sandbagging: Language Models can Strategically Underperform on Evaluations](https://mlsanity.com/paper/2406.07358/) (2024)
- [Narrow Finetuning Leaves Clearly Readable Traces in Activation Differences](https://mlsanity.com/paper/2510.13900/) (2025)

## Built on by (most cited first, 6 of 6)

- [Anatomy of Post-Training: Using Interpretability to Characterize Data and Shape the Learning Signal](https://mlsanity.com/paper/2606.12360/) (2026)
- [Sycophantic Agreement Transfers with Neutral Data via Contrastive Preference Optimization](https://mlsanity.com/paper/2608.31079/) (2026)
- [The Unequal Influence of Bad Advice: Using Training Data Attribution to Modulate Emergent Misalignment](https://mlsanity.com/paper/2609.37914/) (2026)
- [Alignment Forecasting: Predicting Misalignment From Training Data](https://mlsanity.com/paper/2609.35805/) (2026)
- [Correct, Don't Delete: Mitigating Emergent Misalignment with Corrective Supervision](https://mlsanity.com/paper/2609.37624/) (2026)
- [Which Rollout Taught It That? BehaviorTrace and the Limits of Training-Data Attribution in Online RL](https://mlsanity.com/paper/2610.10422/) (2026)

## Abstract (arXiv)

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.