PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yijun Pan, Taiwei Shi, Jieyu Zhao, Jiaqi W. Ma

arXiv:2502.11411v2cs.LG

TL;DR

Unsafe training data can cause harmful LLM responses, making its identification and removal critical for mitigating risks and building safer LLMs. DRA centers and whitens training and target representations, then selects informative dimensions for representation-based unsafe-data attribution. DRA improves detection over state-of-the-art moderation classifiers across unsafe scenarios, and filtering with its scores yields safer retrained models.

  • Problem

    Unsafe training data can cause harmful LLM responses, making its identification and removal critical for mitigating risks and building safer LLMs.

    Source: §1 · p. 1

  • Method

    DRA centers and whitens training and target representations, then selects informative dimensions for representation-based unsafe-data attribution.

    Source: §1 · p. 2

  • Results

    DRA improves detection over state-of-the-art moderation classifiers across unsafe scenarios, and filtering with its scores yields safer retrained models.

    Source: §4.4 · pp. 6–7, §6 · p. 8

  • Takeaways & Limitations

    Gradient-based attribution can produce the safest retrained models even when its detection scores are sometimes lower.

    Source: §6 · p. 8

    The experiments consider injected training data from a homogeneous distribution, whereas real-world injected data may be more heterogeneous and influence different unsafe behaviors.

    Source: Limitation · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Unsafe training data can contribute to harmful behaviors, yet common moderation classifiers are costly to train, tied to predefined taxonomies, and do not measure individual samples’ training impact. DRA denoises representations used for data attribution, improving unsafe-data detection and retraining safety across jailbreak and gender-bias scenarios.

1 Introduction

  • Moderation classifiers rely on predefined taxonomies, require substantial data collection and training resources, and do not account for individual samples’ influence on the model.The authors argue this limits their ability to capture unsafe data’s training impact and improve safety through retraining.Source: §1 · p. 1
  • Representation-based attribution compares training samples with unsafe target examples using hidden states or gradients, without requiring additional labeled data.Prior approaches include hidden-state methods such as LARF and RepSim, and gradient-based similarity.Source: §1 · p. 1
  • Unsafe texts contain a few critical unsafe tokens alongside many benign or neutral tokens, making representations averaged across all tokens noisy for detection.The neutral tokens may support fluency or convey neutral facts, but can obscure the signal used to identify unsafe training data.Source: §1 · p. 2
  • DRA centers and whitens training and target representations, then selects informative hidden dimensions to reduce irrelevant variation in unsafe-data detection.The authors report up to a 63.3% increase in average detection AUPRC and up to a 39.9% decrease in average retrained-model ASR.Source: §1 · p. 2

2 Related Work

Unsafe training data can arise from malicious fine-tuning, synthetic generation, and inherent bias, motivating detection that handles varied sources. Prior work includes moderation classifiers and representation-based methods; this paper’s framing emphasizes the challenge of detecting unsafe training data rather than filtering prompts.

2.1 Sources of Unsafe Training Data in LLMs

  • Unsafe training data may come from malicious fine-tuning, synthetic generation, or inherent bias, including gender bias that can skew occupational assumptions.Existing moderation classifiers are described as insufficient for the complexity and variability of unsafe data in training pipelines.Source: §2.1 · p. 2

2.2 Unsafe Training Data Detection in LLMs

  • Moderation approaches include online tools, fine-tuned safety classifiers, and LLM judges, while representation-based methods use embeddings or gradient similarity for detection.Gradient-similarity attribution estimates individual training samples’ influence and has more efficient variants intended to scale to large models.Source: §2.2 · pp. 2–3
  • Prior methods such as GradSafe and BEBC primarily target prompt filtering, whereas training-data detection also depends on labels that shape model behavior during training.Naively extending those methods to training data yields suboptimal results because noise in model responses degrades detection performance; the paper proposes denoising and dimension selection.Source: §2.2 · p. 3

3 Unsafe Training Data Detection and Filtering

The paper formulates unsafe training-data detection as ranking training examples by representation similarity to a small set of unsafe model outputs. It motivates DRA through a signal-detection sensitivity measure, then combines representation denoising and leave-target-out dimension selection to score samples for filtering.

3.1 Problem Setup: Representation-Based Unsafe Training Data Detection

  • The task is to identify likely unsafe examples in a training set mixing benign and unsafe data, whose membership is unknown, using a small set of unsafe model outputs as targets.Target examples may be user-flagged or collected through red teaming.Source: §3.1 · p. 3
  • Training samples are commonly ranked by their similarity to an aggregated target representation, but irrelevant semantic content makes both training and target representations noisy.Representations may be hidden states or gradients; the target set consists of unsafe model outputs.Source: §3.1 · p. 3

3.2 Sensitivity of the Detection Score

  • The paper uses signal-detection discriminability d′ to quantify how well unsafe samples separate from the benign-dominated background, with d′ monotonic to AUROC and AUPRC under technical assumptions.This measure also permits a decomposition for analyzing representation-based detection.Source: §3.2 · pp. 3–4

3.3 Improving the Detection Score by Denoising Representations

  • DRA centers and whitens representations to downweight high-variance dimensions, then selects dimensions using a leave-target-out discriminability index that treats target examples as a proxy for unavailable unsafe training data.The final score maximizes that index over a selected set of r dimensions; experiments report r = 30. The plug-in linear score assumes the target mean aligns with the unsafe-data mean.Source: §3.3 · pp. 4–5

3.4 Unsafe Training Data Filtering

  • Filtering ranks training examples by the DRA score and removes the top K examples, yielding a dataset expected to reduce unsafe behaviors after retraining.The section describes removing the highest-scoring samples from training data.Source: §3.4 · p. 5

4 Experiments: Jailbreaking Injection Detection

The experiments test whether DRA can detect injected unsafe training samples and improve safety after filtering and retraining. They use a standardized jailbreak-injection benchmark and compare attribution methods with classifier and other baselines.

4.1 Experimental Setup

  • The benchmark injects ToxicChat and XSTest-Response examples into UltraChat-200k and evaluates safety on official test splits.The setup follows DATE-LM for comparability with prior baselines.Source: §4.1 · p. 5

4.2 Evaluation Metrics

  • Detection is evaluated with AUPRC, while retraining safety is measured by ASR averaged over 10 inference runs on the held-out test split.The target set comprises unsafe model outputs collected across those runs.Source: §4.2 · p. 5

4.3 Methods for Comparison

  • Comparisons include moderation APIs, fine-tuned safety classifiers, GradSafe, and GradSim and RepSim attribution baselines, with DRA applied to GradSim and RepSim.The resulting methods are GradSim-DRA and RepSim-DRA.Source: §4.3 · pp. 5–6

4.4 Results and Discussion

  • DRA-enhanced attribution surpasses vanilla attribution methods and the strongest baseline classifiers in unsafe-training-data detection.Table 1 reports AUPRC across backbone models and average improvement over the corresponding baselines.Source: Table 1 · p. 6, §4.4 · p. 6
  • Training with injected unsafe data produces significantly higher ASR than training on benign data alone.Table 2 reports ASR across trained models and datasets; higher ASR indicates a less safe model.Source: §4.4 · p. 6, Table 2 · p. 7
  • Gradient-based attribution with denoising and dimension selection yields the greatest increase in retrained-model safety, which the authors link to better capture of training dynamics.The authors say it filters samples that most strongly drive harmful behavior.Source: §4.4 · p. 7
  • Retraining after DRA-based filtering consistently lowers ASR, with the table reporting average decreases over baseline attribution methods.The lowest ASR for each setting is highlighted in Table 3.Source: §4.4 · p. 6, Table 3 · p. 7

5 Experiments: Gender Bias Mitigation

In the gender-bias setting, DRA improves representation-based detection where moderation classifiers struggle, and filtering with attribution methods also improves downstream model safety. The experiments evaluate a dataset with injected biased biographies and measure gender disparity in physician predictions.

5.1 Problem Setup

  • The setup injects 150 biased biography-occupation pairs into 10,000 benign pairs from Bias in Bios and uses 50 occupation-prediction prompts for targets.The target set consists of gender-biased model responses generated for those prompts.Source: §5.1 · p. 7

5.2 Evaluation Metrics

  • The TPR Gender Gap measures how physician true-positive rates differ between male and female groups on a test set labeled entirely as physicians.A larger gap indicates stronger gender bias; results are averaged over 10 inference runs.Source: §5.2 · p. 7

5.3 Results and Discussion

  • Models trained on gender-biased data have a significantly higher TPR Gender Gap than models trained on unbiased data.Table 4 compares the gender gap for models trained on biased versus unbiased data.Source: §5.3 · p. 7, Table 4 · p. 7
  • DRA improves representation-based detection in the gender-bias setting, where Wildguard fails to generalize and gradient-based methods show particularly high gains.Table 5 reports the detection AUPRC results.Source: §5.3 · pp. 7–8, Table 5 · p. 8
  • Filtering the top 500 influential samples identified by attribution methods improves downstream safety, with gradient-based methods benefiting most.Table 6 reports the TPR Gender Gap after retraining; lower values indicate less gender bias.Source: §5.3 · p. 8, Table 6 · p. 8

6 Conclusion

The conclusion presents DRA as a denoising approach for unsafe training-data attribution, reporting improved detection and safer models after retraining. It also highlights the adaptability of attribution methods and the practical value of gradient-based approaches.

6 Conclusion

  • DRA substantially improves unsafe-data detection and leads to safer models after retraining.The conclusion reports that it outperforms state-of-the-art LLM-based moderation classifiers across diverse datasets and unsafe scenarios.Source: §6 · p. 8
  • The analysis identifies representation noise as a limit on representation-based unsafe-data detection and provides these methods with theoretical grounding.Source: §6 · p. 8
  • Gradient-based attribution consistently produces the safest retrained models, even when its detection scores are sometimes lower.The conclusion describes these approaches as practically valuable for mitigating unsafe behaviors.Source: §6 · p. 8

Limitation

The paper identifies limits involving the assumed distribution of injected data, the granularity of detection, and vulnerability to manipulated influence estimates.

Injection Setup

  • The experiments assume injected training data come from one homogeneous distribution, unlike the more heterogeneous data expected in real-world settings.The authors call for more nuanced injection strategies to address this complexity.Source: Limitation · p. 8

Detection Taxonomy

  • The study does not provide a fine-grained taxonomy because training and target data span multiple unsafe genres.Selecting target data carefully could potentially support genre-specific detection, but this is outside the study’s scope.Source: Limitation · p. 8

Potential Risks

  • Attackers could potentially manipulate estimated influence scores to conceal unsafe training data, undermining detection robustness.The authors connect this risk to vulnerabilities in influence estimation techniques reported in recent research (Wang et al., 2024b).Source: Limitation · pp. 8–9

A Model Safety Evaluation

The supplementary sections describe the safety evaluation and gender-bias injection setup, then examine noisy representations and target-set variety. They report that DRA separates unsafe gradients and benefits from larger, more varied target sets.

A Model Safety Evaluation

  • Model responses are judged on a 1–5 unsafe-content severity scale, with ASR defined as the proportion receiving score 5.The evaluation uses an LLM-as-a-judge following the stated guidelines and a fixed-version GPT4o (2024-11-20) for reproducible, scalable evaluation.Source: A Model Safety Evaluation · p. 12

B Gender Bias setup

  • The gender-bias setup injects 150 female physicians’ biographies relabeled as nurses among 10,000 normal biography–occupation pairs.The original set also contains normal entries for female and male nurses and physicians, making retrieval of injected data non-trivial.Source: B Gender Bias setup · p. 13

C.1 Noisy Representation

  • Raw gradient-space noise leaves unsafe gradients indistinguishable from the larger benign dataset, while DRA produces clear, separable patterns and substantially improves detection.The visualizations compare raw and filtered gradient spaces for Llama-3.1-8B-Instruct with XSTest-Response-injected training data.Source: §C.1 · p. 13, Figure 2 · p. 13

C.2 Target Set Variety

  • With just 3 target samples, DRA outperforms Wildguard, and increasing target-set size improves DRA detection as raw GradSim struggles to use the added data.Figure 3 reports means and standard deviations over 50 sampling repetitions for each target size.Source: §C.2 · pp. 13–14, Figure 3 · p. 13

D d′ as a Proxy for Downstream Metrics

The section examines whether d′ tracks downstream detection performance and reports that it does so monotonically for AUPRC.

D d′ as a Proxy for Downstream Metrics

  • AUPRC increases monotonically with d′, which the authors describe as a reliable proxy for downstream metrics.The relationship is measured between the AUPRC of a score computed from a single dimension and its corresponding d′.Source: D d′ as a Proxy for Downstream Metrics · p. 14

E Ablation Studies

The ablations report that denoising and dimension selection each improve over raw gradient similarity, with their combination performing best.

E Ablation Studies

  • Combining whitened PCA denoising with d′-based dimension selection achieves the best results across models and datasets.Each component alone also yields substantial gains over raw gradient similarity, and the authors describe the components as complementary.Source: E Ablation Studies · p. 14

F Experimental Details

Experiments used four-epoch fine-tuning with model-specific learning rates, and LoRA reduced trainable parameters and gradient-feature size. Attribution and retraining evaluations used the final checkpoint.

F Experimental Details

  • Fine-tuning used 4 epochs, a 0.1 warm-up ratio, and learning rates of 1e-4 for Llama-3-8B-Instruct and 1e-5 for Gemma-2-9b-it.Only the last checkpoint was used for attribution; retraining followed the same setup and used the last checkpoint for evaluation.Source: F Experimental Details · p. 14
  • Table 7 reports ablation studies for DRA in the jailbreaking setup.Source: Table 7 · p. 14
  • LoRA reduced trainable parameters and gradient-feature size to accelerate gradient-feature computation.Source: F Experimental Details · p. 14
  • Training used batch size 1 without gradient accumulation on Nvidia A40 GPUs, with TRL; AUPRC was calculated using scikit-learn.Source: F Experimental Details · p. 14
Loading 2502.11411v2…