Task-Aware Data Selection via Proxy-Label Enhanced Distribution Matching for LLM Finetuning

Hao Cheng, Rui Zhang, Ling Li, Na Di, Jiaheng Wei, Zhaowei Zhu, Bo Han

International Conference on Learning Representations 2026 (ICLR 2026) Conference

Task-specific fine-tuning of foundation models is critically dependent on the quality and relevance of the instruction data. While prevailing data selection methods rely exclusively on instruction instances X to approximate the target distribution, we argue that selection should align with the joint distribution of instructions and task-specific labels (X,Y). However, task-specific labels Y are typically unavailable in practice. To address this, we reformulate the task-specific data selection problem and present a novel pipeline that leverages the reasoning capabilities of large language models (LLMs) to infer proxy labels, thereby facilitating joint distribution alignment. Our approach begins by propagating proxy labels from a small target set to a large, unlabeled source corpus. A two-stage filtering process then removes instances with label noise and refines the subset through distribution alignment. This strategy produces more semantically meaningful and task-aware selections than conventional similarity measures based on $X$ alone. Experimental results show that fine-tuning on a subset of only 10K samples, selected from a pool of 300K, achieves performance competitive or superior to state-of-the-art methods. Code is available at https://github.com/tmlr-group/TADS.