DRIFT: Refining Instruction Data via On-Policy Data Attribution

2026-06-16

AI Summary

This work addresses the limitations of existing instruction data filtering methods, which struggle to identify the most beneficial samples for model improvement due to reliance on external validation targets and bias in gradient norms. The authors propose DRIFT, a novel approach that leverages on-policy generated outputs as validation targets within an influence function framework. By incorporating trajectory correctness–based sign weighting and gradient bias correction, DRIFT enables precise sample selection and reweighting for supervised fine-tuning. This strategy effectively reduces the parameter proximity gap and refines the training data distribution. Evaluated on 7B-scale instruction-following and reasoning models, DRIFT significantly outperforms current data filtering baselines and pushes the performance ceiling of fine-tuned models.

Abstract

Optimizing the training data distribution for Supervised Fine-Tuning (SFT) dictates the capability of Large Language Models (LLMs). While existing data curation methods excel at accelerating training under constrained budgets, they are less suited to elevating the capability upper bound. The challenge here is no longer to identify a smaller subset that preserves performance, but to refine the data distribution toward instances most capable of improving the final model. To address this problem, we explore instance-level data attribution using Influence Functions (IF). We identify that standard IF formulations struggle in this setting due to two structural limitations: a proximity gap caused by off-policy validation targets, and a severe bias towards gradient norm. We propose DRIFT (Data Refinement via On-Policy Influence Functions for Supervised Fine-Tuning). Instead of relying on external reference data, DRIFT utilizes the model's on-policy rollouts as validation targets, which empirically minimizes the parameter proximity gap and better aligns with the local neighborhood assumption of IF. It further applies signed weighting based on trajectory correctness and debiases influence scores against the gradient hacking issue, allowing a small set of validation queries to act as reliable anchors for attributing the full dataset. Experiments on 7B-parameter instruction and reasoning models show that DRIFT consistently raises the performance ceiling on both, outperforming existing data curation baselines.

Problem

• Supervised Fine-Tuning
• Data Curation
• Large Language Models
• Data Attribution
• Training Data Distribution

Innovation

• Influence Functions
• On-Policy Rollouts
• Data Attribution
• Supervised Fine-Tuning
• Gradient Debiasing