DRA: Data attribution with denoised representations for detecting and filtering unsafe training data

1. 01What problem does this paper attempt to solve?

   Unsafe training data can bias models; review classifiers do not measure training impact; DRA uses denoised representations for attribution filtering.

   A small amount of unsafe training data can cause unsafe behavior in models, and existing moderation classifiers do not locate these samples by training influence.

   • Scenario: The authors state that a major source of toxicity, stereotypes, privacy leakage, and ethical violations is unsafe training data. Inherent bias, or deliberately injected adversarial prompts and backdoors, can bypass existing safety alignment.
   • Existing shortcomings: Perspective API, OpenAI Moderation API, as well as Llama-Guard-3-8B, Wildguard, are either limited by predefined categories, or require additional data and training, and label by text style in a model-agnostic way. The authors state that they do not capture the training influence of individual samples, and are therefore not conducive to improving safety through retraining.
   • Observation: Another class of methods assumes that unsafe training samples have representations more similar to unsafe model outputs. If averaged over all tokens, stop words and neutral facts add noise to the representation, overwhelming the few key tokens.
   • Approach: They propose Denoised Representation Attribution (DRA), which denoises and selects dimensions of training and target representations, then filters based on them. Experiments include jailbreak injection filtering and gender bias detection.
   • Threat model:
     • Goal: In Section 4, an adversary injects a small number of unsafe samples into a benign training set to induce jailbreaks. Section 5 describes inherent gender bias in the training set; the paper does not call that section adversarial.
     • What the detector knows: Membership of Dbenign and Dunsafe is unknown, and the unsafe portion is usually small. There is also a small Dtarget, consisting of unsafe model outputs, which can be labeled by users or collected by red teams.
     • Signal used: For training and target points, h(z) is taken as hidden states or gradients, then compared by similarity to the aggregated target representation. The paper does not state white-box or black-box.
     • Victim: LLMs fine-tuned on these data. The authors state that this setting is post-hoc attribution, used to trace which training samples caused unsafe behavior, rather than only doing reply-independent pre-hoc moderation.

   Collapse this answer

2. 02What related research exists?

   Related work includes unsafe data sources, moderation classifiers, and representation similarity attribution; the authors apply denoised dimension selection to training data detection.

   The introduction and Related Work divide the discussion into three categories: data sources, moderation classifiers, and representation similarity.

   • Where unsafe training data comes from
     • Malicious fine-tuning: Yi et al. (2024) and Qi et al. (2023) discuss that a small amount of unsafe data can severely disrupt safety alignment. Qi et al. (2023) also state that online fine-tuning services are insufficient to detect this data.
     • Synthetic data: Wang et al. (2022) generate samples by keywords and target labels; Wang et al. (2023) synthesize fine-tuning data from model responses. The authors cite Wang et al. (2024a), stating that GPT-4 and GPT-3.5 also exhibit unsafe behavior, so synthetic data may introduce risk.
     • Inherent bias: Kotek et al. (2023) discuss gender bias and occupational assumptions. Itzhak et al. (2024) and Echterhoff et al. (2024) discuss cognitive bias in training and high-stakes decision-making.
   • Content moderation classifiers
     • Online tools: OpenAI Moderation API (Markov et al., 2023); the introduction also mentions Perspective API. The authors state they target predefined categories, and cite Weber et al. (2025) for difficulty generalizing to out-of-category content.
     • Fine-tuned detectors: Llama-Guard-3-8B, Wildguard, Aegis-Guard (Ghosh et al., 2024), ShieldGemma (Zeng et al., 2024). The authors state that such models require additional data and training, and do not account for a sample's training influence on the model.
     • Using LLMs as reviewers: Franco et al. (2023) and Li et al. (2024a). Safety-Analyst uses LLMs to generate harm-benefit trees for interpretable moderation. The paper only describes the latter.
   • Representation similarity and data attribution
     • Embeddings or gradients: BEBC (Zheng et al., 2024) uses fine-tuned BERT embeddings for moderation; GradSafe (Xie et al., 2024) estimates gradient similarity relative to safety-related parameters. The introduction also lists LARF (Li et al., 2025) using hidden states, RepSim (Pezeshkpour et al., 2021), and gradient similarity (Xia et al., 2024).
     • Influence estimation: Koh and Liang (2020) quantify the influence of training samples on outputs; Xia et al. (2024) and Kwon et al. (2024) use gradient similarity to reduce computation. The authors consider this perspective suitable for finding unsafe data, because such data often has a disproportionate influence on unsafe behavior.
     • Gap identified by the authors: GradSafe and BEBC mainly do prompt filtering; in training data detection, labels also drive behavior. The authors state that directly transferring them to training data works poorly, and noise in responses lowers detection (Jiao et al., 2025).
   • Baselines and data: Experimental baselines are OpenAI Moderation API, Llama-Guard-3-8B, Wildguard, GradSafe, GradSim, and RepSim. The jailbreak setting follows DATE-LM, with benign data UltraChat-200k and injections from ToxicChat and XSTest-Response. Gender bias uses Bias in Bios.

   The authors place this paper in training data detection: denoising and dimension selection for representation-based attribution, to exploit training labels and reduce response noise, without relying on additional annotation or predefined categories.

   Collapse this answer

3. 03How does the paper address this problem?

   DRA first centers and whitens representations, then selects 30 dimensions by the leave-target-out separability index, removing high-scoring training samples.

   DRA does not train a new moderator. It treats unsafe training samples as signal and a large number of benign samples as noise, first centering and whitening representations, then keeping only dimensions that are more separable for unsafe content, and using the resulting scores to remove the most suspicious training samples.

   Problem setup and naive score

   • The training set is written as Dtrain = Dbenign ∪ Dunsafe, with membership unknown, and the unsafe portion is usually small. There is also a small Dtarget, consisting of unsafe model outputs.
   • h(z) belongs to R^m, and can be hidden states or gradients. B, U, and Q are the representations of the three types of data, and X = B ∪ U.
   • Existing approaches aggregate Q into μQ, e.g., by taking the mean, with score as inner product s_μQ(x) = ⟨x, μQ⟩, taking the training samples with the highest scores. The authors argue that both x and μQ contain a large amount of semantics unrelated to unsafety, and are therefore noisy representations.

   Measuring scores with a separability index

   • The authors borrow from signal detection theory: unsafe samples are signal, benign samples are noise. The separability index d′ is determined by the difference between the means of signal and noise under that score and the standard deviation of the noise.
   • For the naive score, they write d′(s_μQ)=⟨μ_U−μ_X,μ_Q⟩/√(μ_Q^⊤Σ_Xμ_Q). μU and μX are the means of U and X, and ΣX is the covariance of X.
   • Strictly speaking, the noise should be B. The authors assume |U| is much smaller than |B|, so the empirical distributions of X and B are close.
   • The authors cite Stanislaw and Todorov (1999), stating that under certain technical assumptions d′ is monotonically related to AUROC and AUPRC, so they treat d′ as a proxy for these metrics and for ease of decomposition analysis.

   Denoising: a computable approximation of the optimal linear score

   • The linear score family is sw(x) = ⟨x, w⟩. The optimal weight in Lemma 3.1 is w^∗∝Σ_X⁻¹(μ_U−μ_X). It depends on the unknown μU and cannot be computed directly.
   • The authors assume μQ aligns well with μU and instead use the plug-in weight ΣX^{-1}(μQ − μX). Relative to ⟨x, μQ⟩, this score first subtracts the background mean μX, then left-multiplies by the inverse square root of ΣX for whitening.
   • The authors' rationale is that high-variance dimensions may correspond to topic or writing style rather than unsafety; whitening downweights these dimensions. They state that the resulting score is Mahalanobis-related up to a constant independent of x.

   Dimension selection, scoring, and filtering

   • After whitening, the dimension m is still large. They perform eigendecomposition of ΣX, take a projection ΠI spanned by a set of eigenvectors, and compute the score sI on these dimensions.
   • Directly maximizing d′(sI) still depends on the unknown U. The authors replace U with Q and define the leave-target-out index d′LTO(sI): it is equivalent to the sensitivity of finding Dtarget itself within Dtrain ∪ Dtarget.
   • DRA uses s_I*, where I* maximizes d′LTO for a given dimension count r. All experiments in the paper use r = 30. In implementation, candidates are first reduced to the top m′ eigenvectors before searching I; the paper does not give m′.
   • For filtering, samples are sorted by this score and the top-K SK(Dtarget) are removed. The authors state that the removed data is expected to reduce unsafe behavior after retraining. The method section does not specify a unique K.

   Collapse this answer

4. 04What experiments did the paper conduct?

   In jailbreak and gender bias experiments, DRA improves attribution AUPRC; after filtering and retraining, ASR or gender gap mostly decreases, with exceptions.

   Experimental setup

   • Backbone and training: Llama-3-8B-Instruct, Gemma-2-9B-it. 4 epochs, warmup ratio 0.1, batch size 1, no gradient accumulation, and LoRA. Learning rates are 1e-4 and 1e-5, respectively. Both attribution and retraining evaluation use only the last checkpoint. Hardware is Nvidia A40; training uses TRL, and AUPRC uses scikit-learn (Appendix F).
   • Jailbreak data: Following DATE-LM. The benign set is UltraChat-200k, and unsafe samples come from ToxicChat and XSTest-Response, covering toxic/harmful content as well as adversarial prompt-response pairs. Safety evaluation uses the official test split. The number of injected samples is not written in the main text.
   • Gender bias data: Bias in Bios. Dbenign is 10,000 biography–occupation pairs, and Dunsafe is 150 samples that change female physician labels to nurse. The target side in Section 5.1 is gender-biased responses generated on 50 occupation prediction prompts. Table 6 removes the top 500 scored samples.
   • Baselines: OpenAI Moderation API covers 11 safety categories, with AUPRC taking the highest confidence across categories. Llama-Guard-3-8B (the paper states it is pretrained from Llama-3.1-8B for safety classification) and Wildguard (the paper states it is trained on Mistral-7B) both use the probability of outputting the unsafe token. GradSafe looks at gradients of model compliant responses to prompts on Llama-2 safety-related parameters, and does not depend on the tested model's response. GradSim is the ℓ2-normalized inner product of sample gradients with the mean target gradient; RepSim is the hidden state of the last token of the last layer, also with ℓ2-normalized inner product with the target mean.
   • Metrics and target sets: The main detection metric is AUPRC. Precision is |SK ∩ Dunsafe| / |SK|, recall is |SK ∩ Dunsafe| / |Dunsafe|, and the results tables do not list these two separately. ASR is the proportion of prompts on a held-out test set that elicit unsafe responses, averaged over 10 inferences on Dtest. For both jailbreak and gender bias, Dtarget is defined as the union of corresponding unsafe or biased outputs across these runs. Gender gap is also averaged over 10 runs.
   • Review and hyperparameters: ASR is scored by a fixed version of GPT4o (2024-11-20) on a 1 to 5 scale, with a score of 5 counted as success; the scoring criteria state whether it violates Meta's usage policies (Appendix A). The paper does not report agreement between this review and humans. DRA's selected dimension count r = 30. TPR Gender Gap = TPR_{A=0} − TPR_{A=1}, where A=0 is male and A=1 is female; among occupations, Y=1 is physician and Y=0 is nurse, and the test set ground truth is all physician.

   Main results

   According to Table 1, detection AUPRC (%) of data attribution methods. Δ is the average of DRA's difference relative to the corresponding baseline across the two datasets.

   Swipe horizontally to read the table

   Backbone	Method	ToxicChat	XSTest-Response	Δ
   Llama-3-8B-Instruct	GradSim	36.4%	1.7%	–
   Llama-3-8B-Instruct	GradSim-DRA	83.2%	81.5%	+63.3
   Llama-3-8B-Instruct	Repsim	56.3%	88.5%	–
   Llama-3-8B-Instruct	Repsim-DRA	92.6%	98.6%	+23.2
   Gemma-2-9B-it	GradSim	41.1%	50.2%	–
   Gemma-2-9B-it	GradSim-DRA	82.3%	95.5%	+43.3
   Gemma-2-9B-it	Repsim	51.6%	83.9%	–
   Gemma-2-9B-it	Repsim-DRA	98.2%	99.8%	+31.3

   • Authors' summary: The authors state that DRA improves attribution baselines using hidden states or gradients, and after applying DRA exceeds the strongest moderation classifiers. The abstract states average detection AUPRC improves by up to 63.3%; that number is the Δ for Llama-3-8B-Instruct, GradSim-DRA in the table, which is the average of the differences between two AUPRC percentages, not a relative multiplier.
   • Classifiers in the same table: OpenAI Moderation API is 24.3% and 37.8%, Llama-Guard-3-8B is 44.5% and 91.6%, Wildguard is 56.0% and 93.0%, and GradSafe is 37.6% and 27.4%. The order is ToxicChat, XSTest-Response for each. These rows do not indicate a backbone model.
   • Cells inconsistent with the summary: Llama-3-8B-Instruct's GradSim-DRA on XSTest-Response is lower than Wildguard and Llama-Guard-3-8B in the same table. In the table above, each DRA row is higher than its own GradSim or Repsim.

   Jailbreak ASR after filtering and retraining

   • Injection itself: In Table 2, Llama-3-8B-Instruct on Dtrain is 33.3% for ToxicChat and 73.0% for XSTest-Response, and when training only on Dbenign is 26.7% and 0.0%. Gemma-2-9B-it correspondingly is 28.0% and 71.0%, versus 15.7% and 0.0%. The authors state that ASR is higher when unsafe data is included. The paper does not report statistical tests.
   • Numbers after filtering: Table 3 does not state the filtering count K. Llama-3-8B-Instruct: Wildguard 29.0% / 50.0%, GradSim 34.7% / 68.0%, GradSim-DRA 26.0% / 43.0% (Δ −16.9), Repsim 29.0% / 63.0%, Repsim-DRA 27.3% / 60.0% (Δ −2.35). Gemma-2-9B-it: Wildguard 30.5% / 30.0%, GradSim 29.0% / 77.0%, GradSim-DRA 23.3% / 3.0% (Δ −39.9), Repsim 29.0% / 70.0%, Repsim-DRA 30.7% / 64.0% (Δ −2.15). The two sides of the slash are ToxicChat and XSTest-Response, respectively.
   • Authors' interpretation and exceptions: The authors state that gradient attribution with denoising and dimension selection gives the largest retraining safety improvement, and argue that gradients better reflect training dynamics. In Table 3, the lowest ASR in all four cells is GradSim-DRA. The abstract states average ASR decreases by up to 39.9%, matching the Δ of GradSim-DRA on Gemma-2-9B-it. An exception is that model's Repsim-DRA on ToxicChat is 30.7%, higher than Repsim's 29.0%.

   Gender bias detection and filtering

   • Whether bias is learned: Table 4's TPR Gender Gap for Llama-3-8B-Instruct is -0.066 on unbiased data and 0.848 on biased data; for Gemma-2-9B-it it is 0.032 and 0.452. The authors state that the gap is higher on biased data. The paper does not report statistical tests.
   • Detection AUPRC: Table 5 uses a 0–1 scale, not percentages. Llama-3-8B-Instruct: Wildguard 0.059, GradSim 0.417, GradSim-DRA 0.550 (Δ +0.133), Repsim 0.333, Repsim-DRA 0.748 (Δ +0.415). Gemma-2-9B-it: Wildguard 0.059, GradSim 0.404, GradSim-DRA 0.468 (Δ +0.064), Repsim 0.009, Repsim-DRA 0.369 (Δ +0.360). The authors state that Wildguard generalizes poorly because gender bias often lacks the explicit harm signals captured during moderator training. The authors state that gradient methods benefit especially; in the table, the larger Δ is Repsim-DRA.
   • After removing top 500: Table 6. Llama-3-8B-Instruct: no filtering 0.848, Wildguard 0.798, GradSim 0.650, GradSim-DRA 0.246 (Δ −0.404), Repsim 0.296, Repsim-DRA 0.092 (Δ −0.204). Gemma-2-9B-it: no filtering 0.452, Wildguard 0.480, GradSim 0.344, GradSim-DRA 0.224 (Δ −0.120), Repsim 0.536, Repsim-DRA 0.346 (Δ −0.190). The authors state that gradient attribution benefits most from DRA. The lowest value on Gemma-2-9B-it is GradSim-DRA, and on Llama-3-8B-Instruct is Repsim-DRA. On Gemma-2-9B-it, both Wildguard and Repsim gaps are higher than no filtering.

   Other ablations and analysis

   • Table 7, Llama-3-8B-Instruct's GradSim: denoising only is 60.1% / 21.7% (Δ +21.9), dimension selection only is 50.3% / 65.5% (Δ +38.9), combined is 83.2% / 81.5% (Δ +63.3). The two sides of the slash are ToxicChat and XSTest-Response.
   • Table 7, Gemma-2-9B-it's GradSim: denoising only is 61.1% / 64.4% (Δ +17.1), dimension selection only is 59.2% / 62.3% (Δ +15.1), combined is 82.3% / 95.5% (Δ +43.3). The authors state that the two steps are complementary, and combined is higher than using only one step.
   • Appendix C.2 / Figure 3: For each target set size, resampling is repeated 50 times, reporting mean and standard deviation. The caption writes Grad-Sim and Grad-Sim-Dr-Unsafe, and the main text calls it DRA. The authors state that as few as 3 target samples can exceed Wildguard; curve point values are not listed in the main text.
   • Figure 1 and Figure 2 captions write Llama-3.1-8B-Instruct: original gradient t-SNE AUPRC = 0.017 (XSTest-injected), DRA-selected whitened PCA dimensions AUPRC = 0.815 (XSTest-Response-injected). The authors state that in the original space, unsafe gradients are inseparable from benign data, and after DRA a separable pattern forms. The model name in the captions is inconsistent with Table 1's Llama-3-8B-Instruct, and the paper does not state whether it is the same experiment.
   • Figure 4: The authors state that AUPRC computed from single-dimension representation similarity increases monotonically with d′, so they treat maximizing d′ as DRA's objective. Point values are not listed in the main text. The conclusion states that gradient methods sometimes have lower detection scores but are safest after retraining; in Table 6 on Llama-3-8B-Instruct, the lowest gender gap comes from Repsim-DRA, not GradSim-DRA.

   Collapse this answer

5. 05What are some points that can be further explored?

   The authors note that injected data is set to be homogeneous, lacks fine-grained categories, and attribution scores may be adversarially manipulated.

   Limitations and future directions noted by the authors

   • Injection is set as homogeneous: The current treatment regards injected data as homogeneous samples from the same distribution. The authors state that in reality, injections are more likely heterogeneous, with different samples having different effects on different types of unsafe behavior, and more fine-grained injection strategies are needed in the future (Limitation, Injection Setup).
   • No fine-grained categories: Both training and target data mix multiple types of unsafety, and this paper has no fine-grained detection taxonomy. The authors consider that selecting target data might detect specific types, but place it outside the scope of this paper and suggest it as future direction (Limitation, Detection Taxonomy).
   • Scores may be manipulated: The authors cite Wang et al. (2024b), stating that influence estimation has vulnerabilities, and attackers may manipulate the estimated influence scores to hide unsafe training data, weakening detection robustness (Limitation, Potential Risks).

   Experimental coverage

   • Backbones reporting safety and detection results by model are Llama-3-8B-Instruct and Gemma-2-9B-it (Tables 1–6). GradSafe uses Llama-2 gradients; ASR review is GPT4o (2024-11-20).
   • Jailbreak detection follows DATE-LM: benign set UltraChat-200k, injections from ToxicChat and XSTest-Response, safety evaluation with official test split (Section 4.1). The main text does not state the number of injected samples, nor the K used for filtering in Table 3.
   • Gender bias uses Bias in Bios: Dbenign 10,000 samples, Dunsafe 150 samples, target prompts 50; Table 6 removes top 500. The occupations compared are physician and nurse, and gender attributes are male and female (Section 5).
   • Detection baselines include OpenAI Moderation API, Llama-Guard-3-8B, Wildguard, GradSafe, GradSim, and RepSim. Detection metric is AUPRC; ASR and TPR Gender Gap are both averaged over 10 inferences (Sections 4.2, 5.2). The paper does not report agreement between GPT4o review and humans.
   • Ablations report only GradSim's denoising and dimension selection (Table 7). Target set size experiments resample 50 times per size (Appendix C.2). Selected dimension count r = 30 (Section 3.3). Figure 1–2 captions name the model Llama-3.1-8B-Instruct.

   Collapse this answer

6. 06Summarize the main content of the paper

   DRA denoises representations to detect and filter unsafe training data; after filtering and retraining, jailbreak ASR and gender gap decrease in most settings.

   DRA uses data attribution to find and remove unsafe samples from fine-tuning data, rather than training another review classifier.

   • Problem: The authors state that LLMs are sensitive to small amounts of unsafe training data, including toxic content, bias, and injection during training. Review classifiers are limited by predefined categories and do not measure the impact of samples on this training; averaging representations over all tokens is noisy due to neutral tokens.
   • Method: Treat unsafe samples as signal and benign samples as noise. Estimate the signal direction using a small number of unsafe model outputs, center and whiten hidden states or gradients, then score dimensions by a leave-target-out separability index selecting r = 30 dimensions, remove high-scoring samples, and retrain.
   • Jailbreak: Backbones are Llama-3-8B-Instruct and Gemma-2-9B-it. DRA improves AUPRC for GradSim and RepSim; the former on Llama-3-8B-Instruct has an average difference of +63.3 across ToxicChat and XSTest-Response (Table 1). After GradSim-DRA filtering on Gemma-2-9B-it, XSTest-Response ASR is 3.0%, with an average change of −39.9 across the two datasets (Table 3). In the same table, that model's RepSim-DRA ToxicChat ASR is higher than RepSim without DRA.
   • Gender bias: TPR Gender Gap on biased data is 0.848 and 0.452 for the two backbones (Table 4). RepSim-DRA detection AUPRC is 0.748 and 0.369 (Table 5). After removing top 500, the lowest gap on Llama-3-8B-Instruct is RepSim-DRA's 0.092, and on Gemma-2-9B-it is GradSim-DRA's 0.224 (Table 6).
   • Authors' conclusion: The authors state that DRA outperforms approaches primarily based on review classifiers in these datasets and scenarios, and is not tied to predefined categories; gradient methods sometimes have lower detection scores but often yield safer retraining. Table 6's Llama-3-8B-Instruct does not fit the second half's "always lowest." The authors also note that injections are homogeneous, lack fine-grained categories, and influence scores may be adversarially manipulated.

   Collapse this answer

Related papers