Research map / Architectures and efficiency / RNNs and normalization
Enhancing Training Data Attribution with Representational Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AirRep, a trainable BERT-based encoder with attention-based group pooling, trained with a weighted pairwise ranking loss on automatically generated LDS-style influence labels, to make representation-based training data attribution accurate and efficient.
Problem
Representation-based training data attribution relies on heuristic, untuned embeddings and additive group influence, limiting fidelity; gradient-based methods are accurate but too costly at scale.
Approach
Learn task- and model-aligned encoder representations optimized for attribution via a weighted pairwise ranking objective over subsets labeled by their empirical effect on target losses, combined with attention-based pooling that captures group interactions (motivated by higher-order group influence functions).
Key results
- On Qwen2.5-0.5B, AirRep achieves average LDS 26.23 versus 18.45 for LoGra (18432-dim) and 7.65 for GTE-Small.
- AirRep reaches 86.41% data classification accuracy on FLAN versus 50.59% for GTE-Small and 85.44% for LoGra (18432-dim).
- AirRep is reported to be nearly two orders of magnitude more efficient than gradient-based methods at inference, and crosses over LoGra (1152-dim) at around 10^5 training examples.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| optimize the encoder while retaining mean pooling (ablation) | GTE-Small representation baseline · Four-dataset average | LDS: 7.65 -> 19.82 | Helped |
| add jointly optimized attention pooling (ablation) | AirRep w/o Attention · Four-dataset average | LDS: +6.41 | Helped |
| add softmax attention pooling without optimizing the weighting distribution (ablation) | GTE-Small baseline · Four-dataset average | LDS | No clear effect |
| add softmax attention pooling without optimizing the weighting distribution (ablation) | LoGra · Four-dataset average | LDS | No clear effect |
| use AirRep instead of LoGra (Dim 18432) | Qwen2.5-0.5B attribution · Four-dataset average | LDS: 18.45 -> 26.23 | Helped |
| use AirRep instead of LoGra (Dim 18432) | Qwen2.5-0.5B attribution · FLAN | LDS: 19.75 -> 21.11 | Helped |
| use AirRep instead of LoGra (Dim 18432) | Qwen2.5-0.5B attribution · Alpaca | LDS: 12.38 -> 22.58 | Helped |
| use AirRep instead of LoGra (Dim 18432) | Qwen2.5-0.5B attribution · Tulu | LDS: 14.88 -> 15.14 | Helped |
| use AirRep instead of LoGra (Dim 18432) | Qwen2.5-0.5B attribution · SafeRLHF | LDS: 26.82 -> 46.08 | Helped |
| use AirRep instead of LoGra (Dim 18432) | data classification · FLAN | accuracy: 85.44 -> 86.41 | Helped |
| use AirRep instead of LoGra (Dim 18432) | data classification · Tulu | accuracy: 86.00 -> 88.20 | Helped |
| use AirRep instead of GTE-Small | data classification · SafeRLHF | accuracy: 90.60 -> 87.20 | Hurt |
Surprising findings
- AirRep underperforms GTE-Small on SafeRLHF data classification (87.20 vs 90.60), which the authors attribute to the lack of harmful content in UltraChat-based training data.
- Simply adding softmax attention pooling to GTE or LoGra gives only marginal gains, suggesting that optimizing the encoder is necessary.
Limitations
- Adds cost of data generation and model training (training requires retraining many LMs on subsets).
- Evaluation limited to LLM instruction-tuning fine-tuning stage, not pretraining or multimodal settings.
- Weaker on SafeRLHF classification because training data (UltraChat) lacks harmful content.
Experimental setup
Datasets: FLAN UltraChat Alpaca Tulu SafeRLHF
Benchmarks: LDS on FLAN, Alpaca, Tulu, SafeRLHF FLAN data selection F1 data classification accuracy (FLAN, Tulu, SafeRLHF)
Metrics: linear datamodeling score (Spearman) F1 accuracy inference throughput
Baselines: TracIn LESS LoGra Dsdm TF-IDF DSIR RDS GTE-Small
Builds on: influence functions datamodels LoGra LESS GTE
Closest papers by meaning
- Final-Model-Only Data Attribution with a Unifying View of Gradient-Based Methods (NeurIPS 2025 · similarity 0.88)
- LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution (NeurIPS 2026 · similarity 0.87)
- quanda: An Interpretability Toolkit for Training Data Attribution Evaluation (NeurIPS 2026 · similarity 0.86)
- Taming Hyperparameter Sensitivity in Data Attribution: Practical Selection Without Costly Retraining (NeurIPS 2025 · similarity 0.86)
- Fast Data Attribution for Text-to-Image Models (NeurIPS 2025 · similarity 0.86)
- GraSS: Scalable Data Attribution with Gradient Sparsification and Sparse Projection (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.