Detecting and Filtering Unsafe Training Data via Data Attribution with Denoised Representation

Large language models (LLMs) are highly sensitive to even small amounts of unsafe training data, making effective detection and filtering essential for trustworthy model development. Current state-of-the-art (SOTA) detection approaches primarily rely on moderation classifiers, which require significant computation overhead for training and are limited to predefined taxonomies. In this work, we explore data attribution approaches that measure the similarity between individual training samples and a small set of unsafe target examples, based on data representations such as hidden states or gradients. We identify a key limitation in existing methods: unsafe target texts contain both critical tokens that make them unsafe and neutral tokens (e.g., stop words or benign facts) that are necessary to form fluent language, and the latter of which makes the overall representations noisy for the purpose of detecting unsafe training data. To address this challenge, we propose Denoised Representation Attribution (DRA), a novel representation-based data attribution approach that denoises training and target representations for unsafe data detection. Across tasks of filtering jailbreaks and detecting gender bias, the proposed approach leads to significant improvement for data attribution methods, outperforming SOTA methods that are mostly based on moderation classifiers.

As large language models become widely used, it is important to understand why they sometimes produce harmful or biased responses. One important cause is the data used to train them: even a small number of unsafe examples can teach a model to follow harmful instructions or reproduce social biases. This work studies how to find those unsafe examples after a model has already been trained. Instead of relying only on general-purpose safety classifiers, we compare each training example with a small set of harmful model outputs and ask which training examples appear most connected to those failures. A challenge is that the model’s internal signals contain a lot of irrelevant information, because harmful texts also include ordinary words and facts. We introduce Denoised Representation Attribution, a method that filters out this irrelevant noise and focuses on the signals that best separate unsafe examples from the rest of the training data. Across experiments on harmful-instruction attacks and gender bias, this method finds unsafe training data more accurately than several strong alternatives. Removing the examples identified by our method and retraining the model reduces unsafe behavior while preserving general model ability. This makes the approach a practical step toward building language models that are easier to diagnose and safer to deploy.