DRA

Official implementation for the ICML 2026 paper Denoised Representation Attribution (DRA).

DRA is a representation-based data attribution method for detecting unsafe training data in large language model datasets. It measures similarity between individual training samples and a small set of unsafe target examples using hidden-state representations or gradients, then denoises those representations to reduce the influence of neutral tokens such as stop words and benign facts.

Across jailbreak filtering and gender-bias detection scenarios, DRA improves unsafe-data detection over moderation classifiers and prior attribution methods, increasing average AUPRC by up to 63.3% and reducing downstream attack success rate by up to 39.9% after filtering and retraining.

Install

pip install -e .

The core DRA API requires torch, numpy, and scikit-learn. LESS gradient collection and RepSim representation collection additionally use packages such as transformers, datasets, peft, accelerate, and trak.

Quick Start

The public API is in less/src/DRA.py.

import torch
from less.src import select_dra_gradient

train_grads = torch.load("path/to/train/all_orig.pt", map_location="cpu").float()
target_grads = torch.load("path/to/target/all_orig.pt", map_location="cpu").float()

result = select_dra_gradient(
    train_grads,
    target_grads,
    k=20,
    total=100,
)

scores = result.final_scores()
print(result.selected_directions)

Workflows

Run DRA on LESS-style gradient tensors:

python examples/select_gradient_with_dra.py \
  --train-vectors path/to/train/all_orig.pt \
  --validation-vectors path/to/target/all_orig.pt \
  --top-k 20

Collect RepSim forward representations:

python examples/collect_repsim.py \
  --model-name path-or-hf-id \
  --train-path path/to/train.jsonl \
  --validation-path path/to/target.jsonl \
  --output-dir outputs/repsim

Run DRA on RepSim tensors:

python examples/select_repsim_with_dra.py \
  --train-reps outputs/repsim/train_repsim/all_orig.pt \
  --validation-reps outputs/repsim/validation_repsim/all_orig.pt \
  --top-k 30

Repository Layout

• less/src/DRA.py: core DRA implementation.
• less/data_selection/: LESS gradient, representation, and scoring utilities.
• less/train/: LESS LoRA training utilities.
• examples/: gradient and RepSim examples.
• src/: compatibility imports for older scripts.