TAROT: Targeted Data Selection via Optimal Transport
Authors: Lan Feng
Fan Nie
Yuejiang Liu
Alexandre Alahi
Venue: ICML 2025 poster (2025)
Keywords: Data Attribution
Data Selection
Optimal Transport
Primary area: applications

Abstract: We propose TAROT, a targeted data selection framework grounded in Optimal Transport theory. Previous targeted data selection methods primarily rely on influence-based greedy heuristics to enhance domain-specific performance. While effective on limited, unimodal data (i.e., data following a single pattern), these methods struggle as target data complexity increases. Specifically, in multimodal distributions, such heuristics fail to account for multiple inherent patterns, leading to suboptimal data selection. This work identifies two primary limitations: (i) the disproportionate impact of dominant feature components in high-dimensional influence estimation, and (ii) the restrictive linear additive assumptions in greedy selection strategies. To address these challenges, TAROT incorporates whitened feature distance to mitigate dominant feature bias, offering a more reliable measure of data influence. Building on this, TAROT leverages whitened feature distance to quantify and minimize the optimal transport distance between selected data and target domains. Notably, this minimization also facilitates the estimation of optimal selection ratios. We evaluate TAROT across multiple tasks, including semantic segmentation, motion prediction, and instruction tuning. Results consistently show that TAROT outperforms state-of-the-art methods, demonstrating its versatility across various deep learning tasks. Code is available at: https://github.com/vita-epfl/TAROT.

TAROT: Targeted Data Selection via Optimal Transport
Lan Feng 1*
Fan Nie?*
Yuejiang Liu Alexandre Alahi
Abstract
Target Data Selected
Candidate Transport
TAROT Improves Domain Performance with Selected Suhset Instruction
We propose TAROT, a Targeted data selection framework grounded in Optimal Transport the- ory. Previous targeted data selection methods primarily use influence-based greedy heuristics to enhance domain-specific performance. These methods perform well on limited, unimodal data (i.e., data following a single pattern) but become less effective as target data increases in com- plexity. Specifically, in multimodal distributions, these heuristics fail to account for multiple in- herent patterns, leading to suboptimal data selec- tion. This work identifies two primary factors contributing to this limitation: (i) the dispropor- tionate impact of dominant feature components in diggnniimniional influence estimation, and (ii) the restrictive linear additive assumptions inherent in greedy selection strategies. To address these challenges, TAROT incorporates whitened fea- ture distance to mitigate dominant feature bias, offering more reliable measure of data influence. Building on this, TAROT uses whitened feature distance to quantify and minimize the optimal transport distance between the selected data and target domains. Notably, this minimization also facilitates the estimation of optimal selection ra- tios. We evaluate TAROT across multiple tasks, including semantic segmentation, motion predic- tion, and instruction tuning. Results consistently show that TAROT outperforms state-of-the-are methods, highlighting its versatility across various deep learning tasks. Code is available ht https: / //github... tttoo//ittaaiie wwwpe//iiaaaaeellaaaaeee epfi//AROT.
105%
Tuming
Parformace
105%
Semantic
Segmentation
108.5%
Mation
102%
Prediction
Full Data
25%
50%
25%
100%
Target Distribution
Usage
Figure I: Left: the selection objective of TAROT, designed to minimize optimal transport costs. The method operates on target distribution and a candidate data pool, selecting subset from the candidate data reduce transport costs. Right: the approximate performance improvements and data efficiency achieved by TAROT across various tasks, It illustrates that TAROT consistently enhances task per- formance in different domains while requiring less data. Detailed experimental results are provided Section 4.
language models (LLMs) (Achiam et al., 2023; Zou et al., 2024; Oquab et al., 2023; Feng et al., 2024). However, train- ing a model on all available data may not be optimal or feasible, as data quality varies and computational resources are often limited (Albalak et al., 2024; Liu al., 2024; En- gstrom et al., 2024; Kang et al., 2024). To address this, recent research has introduced targeted data selection meth- ods. These methods aim to select data from a candidate pool that maximizes model performance on specific target tasks. Recent methods (Xia et al.; Engstrom et al., 2024) address this challenge using a two-step approach: (1) influence es- timation, which calculates data influence scores between the candidate and target datasets, and (2) greedy heuristics, which selects data with the highest average influence scores for the target dataset.
1. Introduction
Although these selection paradigms have shown to be ef- fective in LLM applications (Engstrom et al., 2024; Xia et al.), they rely heavily on the assumption that influence functions are additive. This means that the influence of multiple training samples together. However, as noted by Hu et al. (2024), this as- sumption, while efficient and broadly applicable, fails even in simple linear regression scenarios. One main reason why the additivity assumption works well in some cases is the simple and unimodal distribution of the target domain, such
Expanding model capacity and data volume has emerged as a key strategy in both computer vision (CV) and large
Eqqual contribution 'EPFL, Switzerlanzd Stanford, USA. Correspondence to: Yuejiang Liu <yuejiang.liue stanford.edux>, Alexandre Alahi
target sample can be added
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

TA ROT: Targeted Data Selection via Optimal Transport
Candidate Data
Data Selected
Selected Data
Target Data
Transport Plan
Gradient Feature Whitening
2) OT Distance Minimization Selection
Iteration
Iteration 2
Initial Gradient Features
Whitened Gradient Features
Costs: continue
Costs: 6, stop
Final Selection
Figure 2: Overview of TAROT. Left: Raw gradient features are whitened and projected onto unit hypersphere. Right: A greedy algorithm iteratively selects candidate samples. During the second iteration, adding two new distant data points increases the overall transport cost, prompting the algorithm to stop. As result, only in-distribution data points are selected to minimize the OT distance.
as single task (e.g. code generation) in LLMs. In these sit- uations, adding up the influences provides good estimate of how candidate sample impacts the entire target task. However, in other contexts like motion prediction, where data distributions are often complex and multimodal, this approach falls short. In such cases, the summed influence does not capture the diverse nature the target distribution, reducing its effectiveness.
Broad applicability and improved performance: ex- tensive experiments in CV and LLMs demonstrate that TAROT consistently outperforms existing methods, allow- ing models trained on smaller, targeted datasets to match or surpass the performance of those trained on full datasets.
2. Preliminary
2.1. Data Influence Formulation
To this end, we adopt distribution matchhin approach for targeted data selection. We first introduce whitened feature distance (WFD), an influence estimation method designed to mitigate the bias from dominant features. The benefits of feature whitening are also discussed in (Ermolov et al., 2021). Using WFD as pairwise cost estimation between data points, we develop TAROT, an Optimal Transport (OT)-based targeted selection framework. This method se- lects data that minimizes the OT distance between the target distribution and the selected samples. Figure illustrates the selection objective of our approach. Our choice of OT also supported by findings in (Just et al., 2023), which demonstrate that training models on source data with smaller OT distances to the target domain improves domain specific performance.
Following Pruthi et al. (2020), we summarize how to esti- mate the influence of training data point on held-out data using first-order approximation of training dynamics.
First-order Approximation to Per-step Influence. Sup- pose x; is the only training sample at time step t, and the model f is updated based on the loss L(x;; 0*) using SGD with learning rate Th- The per-step influence of I; on target sample x is defined as the reduction in the loss of a caused by training on z;. This can be approximated using a first-order Taylor expansion. The per-step influence is:
Tl (V (V ('), VL(r;'))
(1)
where VL(x,, xr; 8*) is the loss gradient.
Trajectory Influence via Checkpoints. For practicality, Pruthi al. (2020) compute the trajectory influence by summing the per-step influences over epochs:
To validate the general effectiveness of our method, we conduct experiments on semantic segmentation, motion prediction, and instruction tuning tasks. Extensive results demonstrate that our approach consistently outperforms pre- vious methods. Notably, models trained on smaller, domain- relevant datasets often surpass those trained on larger, less relevant datasets. In summary, contributions are:
Infsap (I x) &
Σ (VC (2)
2.2. Background on Optimal Transport
Given complete, separable metric space X' and probabil- ity measures a, B P(X), the Kantorovich formulation (Kantorovitch, 1958) of OT is defined as:
Improved influence estimation: we present WFD as a reliable measure data influence.
Distribution matching- based data selection: we propose TAROT, an Oa-based pipeline for targeted data selection that reduces distributional differences between the target and selected datasets.
min OT(a, rEII(a,3)
c(x,y) d (x,
(3)
xxxx
where e(..)) : X x X R+ is a cost function, and [[(a, is the set of couplings with marginals a and g
2

---

TA ROT: Targeted Data Selection via Optimal Transport
In practice, the measures a and are approximated using finite samples {x(0) and (y("), with corresponding dis- crete measures a = Zim1 a;dggt) and [ = Σ m b,0ym.. The pairwise costs form an n X m matrix C, where Cyp = c(x(), y()). Solving this problem as a linear program is computationally expensive, with cubic complexity. In im- plementation, we use the Sinkhorn algorithm (Cuturi, 2013) for efficient calculation.
The distance between two samples is then defined as:
dz (2, [1x(() (=2)))22
-im
3.2. Whitened Feature Distance
Directly computing the distance dz (2, 2') defined in Equation (5) presents significant challenges due to the cor- relation of gradient features. This correlation results in an ill co-diiiitooeed covariance matrix, allowing certain feature space directions to dominate distance estimation and distort sample relationships. Additionally, varying gradient com- ponent scales bias distance measurements, impairing the accuracy of distance estimation.
3. Method
The failure modes of influence-based greedy heuristics have been systematically studied in (Hu et al., 2024). Specifically, errors in the influence function and the inability to account for the non-additive nature of collective influence can cause these heuristics to fail even in simple linear regression tasks.
To address these issues, we propose a novel application of whitening to the gradient features. Whitening effectively decorrelates the features and scales them to have unit vari- ance, ensuring that each feature contributes equally to the distance computation. Our method comprises three steps: 1. Random projection. We adopt the same random projec- tion technique from (Xia et al.; Park et al., 2023) to reduce the dimensionality of the gradients. The projected gradient ((z),m) jm is computed as:
In contrast, formulate targeted data selection as a dis- tribution matching problem, aiming to select subset of candidate data that closely aligns with the target distribu- tion. We employ the OT distance to quantify the discrepancy between distributions.
Sections 3.1 and 3.2 demonstrate how to measure the OT distance between datasets, and Section 3.3 discusses strate- gies for selecting data to minimize the OT distance. An overview of these processes is illustrated in Figure 2.
(6) where pT is the projection matrix. We leverage the effi- cient CUDA implementation from Park et al. (2023).
<(=)!Th = p' -.=);;
2. Whitening, Despite dimensionality reduction, projected gradients remain correlated, creating an anisotropic feature space. To address this, we perform Cholesky whitening to decorrelate features. First, we center the projected gradi-
3.1. Optimal Transport Between Datasets
As discussed in (Alvarez-Melis & Fusi, 2020), estimating OT distance between datasets begins with defining a dis- tance metric between data pairs (zt, y) and (t',)'). In Sec- tion 2.1, we introduced a gradient-based method for esti- mating data influence, which embeds both input features and labels into the model's gradient. This gradient feature captures the task-specific relationship between data points and the model's loss. By computing distances such as the L2 distance, it provides a viable metric for the OT prob- lem. Furthermore, it aligns with the model's optimization objective, ensuring that the OT distance reflects meaningful differences between datasets in terms of their impact on model training.
ents ((2) = (=)=) ,...., +(=-)))))))
N
o=:)oT0) rro = x(=); pro
Σ
Pro
N
Next, we compute the covariance matrix Σ and apply Cholesky decomposition, Σ I.I.T, where L is lower
triangular. The gradients are whitened as:
6(2)== = 1==0p))
(8)
ensuring decorrelation.
3. Normalization. Each whitened gradient vector d(2) is normalized to unit length to ensure consistent gradient scales:
Formally, we define a dataset D as set of feature-label pairs (I, E X x yr where X' is the feature space and y is the label set. For convenience, we use the shorthand 2 (x, v) and 2 X y. We consider a candidate dataset D.
((=);²
3(=);
(+()){]
(i) N
and a target dataset D. = (1) J.m
The WFD di 22) is then calculated as:
==1
d= (z, x') = [1x(2) 2(2')122-
(10)
We embed using the model's parameters trained on De:
T
This gradient-based distance also applies to data influence estimation. Our experiments show that WFD outperforms the state-of-the-ar art method (Park et al., 2023), as detailed in
d(z)= - v(2;))),
(4)
1=1
where T is the total number of checkpoints, and V ((2; 0,) represents the gradient of the loss function with respect to the model parameters at the i-th checkpoint O,.
Section 4.1.
Using dw, we lift the metric to a dataset-level OT distance:
doT (De, Dz) = min [
d%
(2,z"). (1I)
xell(a zx
3

---

TA ROT: Targeted Data Selection via Optimal Transport
3.3. Data Selection via OT Distance Minimization Our objective is to select a subset D, S De of size S S N, such that the OT distance between the selected subset D and Dj minimized.
while the OT distance is evaluated against the remaining
(k - 1)/k data.
At each iteration, points from D. are selected based on their nearest- distances to 1/k of Dr, and the OT distance is monitored. The process stops when the OT distance increases:
8
Formally, let D. = find D. such that:
be a subset of De. We aim to
i=1
dor(P, D. \ D( > dor (D(prev), Do D(
doT P.) < dDo D.),
(12)
(14)
where dor is subject to the constraint that the masses for each point are normalized.
where D( is the 1/k subset used for selection.
After processing all l: folds, the selected subsets are com- bined to form the final dataset D,. The choice of k balances distribution alignment and generalization. In our experi- ments, h = 10 ensures a good match with the target distri- bution while avoiding overfitting.
Solving the subset selection problem to minimize the OT distance is inherently combinatorial, as the number of possi- ble subsets D, S De grows exponentially with the size of De. An exhaustive search incurs prohibitive computational costs. To address this, we propose greedy algorithm that efficiently approximates the solution by iteratively selecting elements that locally minimize the OT distance.
3.4. Data Weighting with OT Potential
To highlight the importance of individual data points, we assign weights based on their OT potentials from Equa- tion (13). These potentials are scaled to positive integer weights, determining the number of times each point is duplicated within an epoch. The scaled weights {w;} are normalized to satisfy Σ W: R. where R is a customiz- able repetition factor. For instance, R can be set to N + M to match the same training steps as using the full dataset, or adjusted based on specific training requirements, e.g., limited training budget.
Our method exploits the sparsity of gradient- baase features in high-dimensional space, which results in a sparse cou- pling matrix in the OT problem. Specifically, most of the transport mass in the OT plan is concentrated on the nearest neighbors of target points in Dr. Leveraging this property, the algorithm iteratively selects the nearest neighbors from Dr for each point in D., providing computationally effi- cient solution.
Based on this, we introduce two selection schemes tailored for different scenarios: (1) fixed-size selection for maxi- mizing performance under a limited budget, and (2) OT distance minimization selection to find potentially optimal selection ratio.
4. Experiments
This section starts by evaluating the influence estimation performance of WFD in Section 4.1. Following that, we evaluate TAROT across diverse tasks, including semantic segmentation (Section 4.2), motion prediction (Section 4.3), and instruction tuning (Section 4.4).
Fixed-Size Selection. This scheme selects a fixed number of points S from D.. For M target samples, we iteratively add their nearest candidates from De to the selected set D,. In iteration k. the k-th nearest samples are considered. If adding all of them keeps [D,] I S. they are directly added, and the process continues. Otherwise, we compute the potential Φ of the new candidates using the dual OT problem:
4.1. WFD for Influence Estimation
Evaluation Metric, Following Park et al. (2023), we use the Linear Data Modeling Score (LDS) as our evaluation metric. LDS measures the effectiveness of an attribution or influence estimation method in predicting changes in model outputs when the training set is altered. Further details of LDS are provided in the supplementary materials.
[ D.U
min = reM(D,.L)a),)))) P.)
22 dT(Z, 24), (13)
Target and Candidate Datasets. Following Park et al. (2023), we evaluate image classification using ResNet-9 classifiers trained on the CIFAR-10 dataset. For motion prediction, we adopt AutoBots (Girgis et al., 2021), training the nuScenes (Caesar et al., 2020) dataset (32k samples) and validating on 9k target samples.
where I(D,uffz), D,) is the set of couplings between D,U { and D,, As noted in (Just et al., 2023), & estimates the benefit of adding 2; for minimizing the OT distance. We rank candidates by potential and select the top S - |D,| samples. The process ends once these samples are added.
OT-Distance Minimization Selection (OTM). This scheme minimizes the OT distance between the selected data D. and the target distribution. To avoid overfitting to specific points in De, we use k-fold selection strategy with eross- validation. The target dataset D, is randomly split into k equal subsets. In each fold, 1/k of D, is used for selection,
Baselines. We compare WFD against the sate-off-hhe-a data attribution method TRAK (Park et al., 2023) using the same ensemble approach. Since WFD quantifies distance (where smaller values indicate greater influence), we eval- uate it with the LDS score by considering negative WFD
4

---

TA ROT: Targeted Data Selection via Optimal Transport
Smaller Whitemed Feature Distamce
Larger Whitened Feature Distance
Target Example
Candidane: WOMD
Tirget: nuSenes
Targeic Cityscaes Caadidate: GIAS
Figure 3: Illustration of target examples and their corresponding closest (most helpful) and furthest (most detracting) candidate examples, identified by WFD. Top two rows: Visualization of motion prediction tasks using the nuScenes and Waymo datasets. The dark-to-light blue lines indicate the past (input) and future (ground truth) trajectories for the target vehicle, while green lines represent surrounding vehicle trajectories. Bottom two rows: Semantic segmentation tasks with examples from the Cityscapes and GTA5 datasets. For each target example, the top image shows the RGB scene, and the bottom image depicts the corresponding ground truth semantic segmentation mask. The WFD method selects training examples that are semantically aligned with the target, demonstrating its effectiveness in identifying positive and negative samples across different tasks.
TRAK (Park al.)
WFD
This indicates that WFD exhibits a higher correlation with datamodels, which serves as an "oracle" of sorts for the LDS objective.
Image Classifiction (CIFAR-10)
Motion Prodiction (auSecnes)
1.20
Inspecting WFD-identified examples. In Figure 3 we display, for four randomly selected target samples from the nuScenes (Caesar al., 2020) and Cityscapes (Cordts et al., 2016) dataset, the training examples from WOMD (Et- tinger et al., 2021) and GTA5 (Richter et al., 2016) dataset corresponding to the smallest (most positive) and largest (most negative) whitened feature distance.
0.10
0.85
10
2.
50
10
2.
28
50
Ensemble Number or
Ensembil Number
Figure WFD outperforms baseline methods on both tasks. The x-axis indicates the number of model checkpoints that given method uses to compute attribution scores. The y-axis indicates the method's efficacy as measured by LDS.
4.2. TA ROT for Semantic Segmentation
Semantic segmentation involves assigning semantic labels to every pixel in an image. Prior studies show that syn- thetic data can significantly boost segmentation model per- formance on real- weal-world data (Richter et al., 2016). We then address the following research question: Can TAROT identify synthetic data that better align with real-world distributions to enhance model performance?
values. Additionally, we employ ZCA whitening instead of Cholesky whitening to preserve data structure for improved LDS performance. Differences between these two whiten- ing methods are discussed in the supplementary materials.
Main Results. Figure compares the performance of WFD with TRAK on the LDS score. Key findings include:
Datasets. The GTA5 dataset (Richter et al., 2016) serves
WFD outperforms TRAK regardless ensemble sizes.
5

---

TA ROT: Targeted Data Selection via Optimal Transport
TAROT
DaDm (Ei- gstrom el al.)
LESS (Xia
Random
TAROT
DaDm (Eo-
LESS (Xia
Random
gerom al.)
62.00
3.00
OT Distance Minimization
61.00
Brior-minFDEs
2.75
mlOUt
60.00
2.50
2.25
59.00
OT Distance Minimization
2.00
58.00
20
40
60
80
100
0
20
40
60
100
Selection Percentage %
Selection Percentage %
Figure 7: Results of Targeted Selection for Motion Predic- tion. The figure illustrates the performance of Wayformer under various data selection methods.
Figure 5: Results of Targeted Selection for Semantic Segmentation. Performance curves of DeepLab trained using different data selection methods.
Target Data
Selected Data
Main Results. Figure compares the performance of TAROT and other baselines. Key findings include:
Some synthetic data degrade model performance. Training model on the full synthetic GTAS dataset low- ers mloU from 58.98 to 58.21, highlighting the negative impact misaligned synthetic data.
Targeted data selection enhances domain-specifit per- formance. Model performance improves with increased selection ratios but declines beyond an optimal point. No- tably, TAROT consistently outperforms baselines, achiev- ing the highest mloU with about 24% of the GTA5 dataset. This demonstrates the effectiveness of OT distance in guid- ing optimal selection ratios.
DsDm (Engstrom et al.)
TAROT
Data selection visualization in gradient feature space with T-SNE. We select 10% of Cityscapes' training set, with validation split as the target. Left: Samples se- lected by DsDm cluster tightly around the target distribution center. Right: Samples selected by TAROT are distributed more broadly, capturing the target distribution's 's complexity.
Figure 6:
Additional insights from T-SNE visualizations. Fig- 6 shows T-SNE visualizations of gradient features for selected data. The visualizations reveal differing selec- tion behaviors: DsDm tightly clusters around the target distribution center, failing to capture its complexity. In contrast, TAROT selects data that better aligns with the broader target distribution, highlighting a key limitation of averaging-base influence methods in handling complex distributions.
as the candidate dataset, while the Cityscapes (Cordts et al., 2016) training split (2975 samples) is used as the target dataset, with its validation split for evaluation. GTA5 con- tains 24,966 synthetic images with pixel-level annotations of urban scenes, whereas Cityscapes provides finely annotated real-world urban images.
We compare TAROT with state-of-the-ar tar- geted selection methods LESS and DsDm (Xia et al.; En- gstrom al., 2024), well as random selection.
Baselines.
4.3. TA ROT for Motion Prediction
Motion prediction, the task of forecasting future trajecto- ries from past data, is critical component of autonomous driving. This section addresses two key research questions:
Data Selection and Model Training. We employ DeepLabV3 (Chen, 2017) with a ResNet50 (He et al., 2016) backbone, using gradients from the last four checkpoints to guide data selection. We evaluate selection ratios of 5%, and OTM (Section 3.3). OTM selects approx- 20%, 50%
Is TA ROT generic? Compared to traditional computer vision tasks (e.g., image classification), motion prediction involves more complex inputs (e.g., maps, trajectories) and outputs (e.g., predicted trajectories and confidence scores). Is TAROT still effective in a more challenging setting? Is TAROT transferable? For example, can model trained on data selected by smaller model improve the performance of larger model?
imately 24% of the data. The repetition factor is set as: R N + M for data selected by TAROT (Section 3.4), while other methods employ uniform repetition to match the same size. A new DeepLabV3 model is then trained from scratch on the selected subsets.
Evaluation Metric. Mean Intersection over Union (mIoU) used as the primary evaluation metric. We report the best validation performance during training.
Datasets. We use the UniTraj framework (Feng et al., 2024) for unified training and evaluation across multiple datasets, including Waymo Open Motion (WOMD) (Ettinger et al.,
6

---

TA ROT: Targeted Data Selection via Optimal Transport
Table 1: Performance of TAROT on instruction tuning for LLMs. 'Transfer' indicates that SWENN-255-7 is trained on data selected by LLAMA-3. 1...-.1--8BB Otherwise, the same base model is used for both selection and fine-tuning. 'All' denotes employing the full dataset. 5% indicates selection ratio of 5%. The percentages within parentheses represent the selection ratios determined by OT-Minimization. "Subtask label' specifies whether subtask labels are required during data selection.
LLAMA-3.
1-8-8BB
1-8B
QWEN-2.55
Transfer
Dataset All 5% Random
MMLU1
BBH f
MMLU f
BBH 1
MMLU
BBH
Subtask label
63.8
63.3
74.1
63.8
74.1
63.8
64.7
60.8
74.1
63.5
74.1
63.5
X
5% LESS (Xia et al.)
65.7
62.6
74.3
66.3
74.3
65.00
V
5% TAROT TAROT-OTH
66.0
65.0
74.1
66.9
74.2
65.3
X
65.7 (0.13%)
63.6 (0.21%)
74.3 (0.09%)
68.9 (0.13%)
74.3 (0.13%)
68.7 (0.21%)
X
2021), Argoverse 2 (Wilson et al., 2021), nuScenes (Caesar al., 2020), and nuPlan (H. Caesar, 2021). The nuScenes training set (32k: samples) serves as the target dataset, while the candidate pool comprises WOMD, Argoverse 2, and nuPlan. From nuPlan, we filter trajectories with moving distance over 2 meters, yielding 1000k samples. We use the official training splits of WOMD and Argoverse 2, in- cluding 2000k samples. The evaluation is conducted on the nuScenes validation set. This cross-dataset task involves challenges like non-overlapping geographic regions and varying annotation formats. We explore whether TAROT can effectively curate multi-source data to enhance nuScenes validation performance.
proach. Additionally, OTM selection outperforms training on the entire dataset and other selection ratios, delivering global optimal results. Moreover, TAROT archives the top I performance on the nuScenes leaderboard. Results are shown in the Appendix D.
Transferability of TAROT: Data selected by Autobots significantly enhances the performance of Wayformer, a model with 10x the capacity. Furthermore, despite differ- ences in loss functions and model architectures between the two models, the selected data still yiclds performance im- provements. This highlights that although TAROT lever- ages a pretrained model for selection, the selected dataset retains strong model-agnostic properties.
Baselines. We utilize the same baselines in Section 4.2.
4.4. TAROT for Instruction Tuning
Data Selection and Model Training. We employ Auto- Bots (Girgis et al., 2021), lightweight (1.5M parameters) attention- based prediction model, for data selection. Auto- bots is pretrained on candidate and target datasets for 50 epochs, and gradients from the final checkpoint guide se- lection. We evaluate selection ratios of 5%, 20%, 50%, and OTM. OTM selects approximately 31, 53%, and 22%
Previous experiments demonstrate that TAROT is particu- larly effective when dealing with complex target distribu- tions. In this section, we evaluate whether TAROT can also deliver strong performance in LLM tasks, where the target distribution is relatively simple and existing methods have already shown high effectiveness.
Datasets. To ensure a fair comparison and maintain consis- tency with established benchmarks, we replicate the exper- imental setup described by Xia et al.. We utilize the same candidate dataset, comprising FLAN V2 (Longpre et al., 2023), COT (Wei et al., 2022), DOLLY (Conover et al., 2023) and OPEN ASSISTANT 1 (Köpf et al., 2024), with MMLU (Hendrycks et al., 2021b;a) and BBH (Suzgun et al., 2023) serving as the target tasks for evaluation.
from WOMD, Argoverse 2, and nuPlan, covering 29.8% of the full dataset. The selected data is used to train the Way- former (Nayakanti et al., 2023) model (15M parameters) for 100 epochs. To highlight that TAROT can also achieve performance gains with reduced training data and iterations, the data weighting method (Section 3.4) is not applied.
Evaluation Metric. We use Brier Minimum Final Displace- ment Error (Brier-minFDE) metric for evaluation, which is used in the Argoverse 2 (Wilson et al., 2021) challenge. We report the best validation performance for each method.
Data Selection and Model Training. We conduct ex- periments using LLAMA-3.1 1-8B (Dubey et al., 2024) and QWEN-2.5-7B (Team, 2024) as the base models for both data selection and fine-tuning. To evaluate the transferability of our selection method, we also fine-tune a QWEN-2.5-7B using data samples selected by LLAMA-3.1 1---.B Data se- lected by TAROT-OTM is weighted using OT potentials (Section 3.4) with R 0.5%N. More details are provided in the supplementary materials.
Main Results. Figure 7 compares the performance of Way- former trained on datasets with different selection ratios. Key findings include:
Effectiveness of TA TAROT: TAROT consistently outper- forms other baselines. inffuence-estimation based heuris- tics, such LESS and DsDm, perform well low selec- tion ratios (5%) but lose effectiveness as the ratio increases, emphasizing the benefits of the distribution ciriibbtooommtccinngg. ap-
Evaluation Metric. We employ Exact Matching (EM) for BBH and Accuracy for MMLU as evaluation metrics.
7

---

TA ROT: Targeted Data Selection via Optimal Transport
Table 2: Ablation Study of WFD on nuScenes with ensem- ble size ranging from to 50. Evaluation metric is LDS.
results are presented in Table 2. Specifically, wo whitened indicates that we skip the whitening step and instead use the original projected gradients to compute the distance. w/o norm means that we compute the L2 distance of the gradi- ent vectors without normalization. The results demonstrate that removing these proposed components leads to a perfor- mance drop, highlighting the effectiveness of our influence estimation design.
Ensemble size
10
50
w/o whitened
0.007
0.010
0.017
w/o norm
0.024
0.060
0.097
w/o whitened+norm
0.005
0.008
0.015
WFD
0.028
0.065
0.098
Target Set
Test Set
K-Fold Selection in OTM. As introduced in Section 3.3, TAROT employs k-fold cross-validation for OTM, where a larger k results in less alignment with the target distribution. We investigate how the choice of l affects data selection in the motion prediction setting (Section 4.3). Specifically, we perform experiments with OTM selection using À values ranging from 20, evaluating the model's performance both the test set and the target set. The results, shown in Figure 8, indicate that as k increases, the model's perfor- mance on the test set improves, reaching its peak at A = 10. Meanwhile, performance on the target set steadily declines, suggesting that the selected data becomes less aligned with the target distribution as increases.
2.4
Brier-minn mieemmnnFEE
2.3
2.3
2.2
2.1
1
5
10
15
20
Number of
Figure& 8: Ablation study for the choice of k. The figure shows the performance of Wayformer under different k.
Main Results. Table 1 represents the performance of TAROT compared with baselines. Among them, LESS re- quires the subtask labels of target samples, whereas TAROT operates without this requirement. Our key findings include:
5. Related Work
We review the progress on data selection in this section. More related work is deferred to App. A. Coreset selec- tion aims to identify dataset subset for optimal in-domain performance (Xia et al., 2022; Mirzasoleiman et al., 2020; Borsos et al., 2020; Killamsetty et al., 2021; Guo et al., 2022; Xiao et al., 2024), assuming validation and candidate data share the same distribution. Our work shifts focus to tar- geted selection, optimizing performance on a distinct target domain. Prior methods often rely on influence-estimation heuristics: Engstrom et al. (2024) leveraged linear data- model scores (Park et al., 2023) to select data for LLMs pretraining; Xia et al. leveraged influence estimation to se- lect data for LLMs fine tuning. Hu et al. (2024) highlighted these methods' limitations and proposed iterative heuristics. A concurrent work (Liu et al.) also frames data selection as an optimal transport problem. However still follows the same gradient distance estimation in (Xia et al.) and uses a continuous OT-based formulation, making it less efficient and effective. Our approach introduces novel distance estimation method and greedy subset selection pipeline to address these challenges.
TAROT is effective, generalizable, and transferable. TAROT consistently outperforms all baseline methods and models trained on the full dataset without relying on sub- task labels. Besides, data samples selected by LLAMA-3.1- SB enhances performance for QWEN-2.5- 2b, indicating strong transferability across different model architectures,
Effectiveness of OTM. Remarkably, TAROT-OTM se- lects less than 0.5% of the data while achieving perfor- mance that is comparable methods selecting 5% of the data. The results suggest that only small fraction of the candidate dataset is pertinent to the target tasks, and TAROT-OTM successfully identi- fies these critical samples, thereby significantly reducing training costs without compromising performance.
even surpasses that of
4.5. Computational Complexity
The computational complexity of TAROT mainly arises from two aspects: (1) gradient feature computation and (2) data selection. The former shares the same complexity as previous methods that use projected gradients (Park et al., 2023; Liu et al.. 2024; Engstrom et al., 2024). The latter introduces additional computational overhead due to compu- tation of OT distance. We compare the wall-clock runtime of TAROT with baselines and provide a detailed asymptotic complexity analysis in the Appendix G.
6. Conclusion
This paper presents TAROT, a general targeted data se- lection framework leveraging OT theory. By introducing WFD, TAROT mitigates biases from dominant features and enhances influence estimation. Through OT- selection, it aligns subsets closely with target distributions, consis- tently outperforming traditional methods in tasks such as semantic segmentation, motion prediction, and instruction tuning. These results demonstrate TAROT's effectiveness
4.6. Ablation Study
Whitened Feature Distance. We conduct an ablation study to analyze the impact of key components on WFD. The
8

---

TA : Targeted Data Selection via Optimal Transport
in improving domain-specific performance with smaller, more relevant datasets. Limitations: In cases with small or highly specific target datasets, TAROT may overfit by overly tailoring the selected data, potentially compromising generalization. A promising future direction is incorporat- ing distribution diversification into the selection objective.
Chen, L.-C. Rethinking atrous convolution for semantic image segmentation. arXiv preprint arXiv 2017.
Choe, S. K., Ahn, H., Bae, J., Zhao, K., Kang, M., Chung, Y., Pratapa, A., Neiswanger, W.. Strubell, E., Mitamura, T., et al. What your data worth to gpt? Ilm-scale data valuation with influence functions. arXiv preprint arXiv:2405 13954, 2024.
Acknowledgments
The project was partially funded by Honda R&D Co., Ltd, Valeo Paris, and the Swiss National Science Foundation under the Grant P500PT 222293.
Conover, M., Hayes, M., Mathur, A., Xie, J., Wan, J., Shah, S., Ghodsi, A., Wendell, P., Zaharia, M., and Xin, R. Free dolly: Introducing the world's first truly open instruction- tuned 11m. Company Blog of Databricks, 2023.
Impact Statement
Cordts, M., Omran, M., Ramos, S., Rehfeld, T., Enzweiler, M., Benenson, R., Franke, U., Roth, S., and Schicle, B. The cityscapes dataset for semantic urban scene un- derstanding. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3213-3223, 2016.
This paper presents TAROT, targeted data selection frame- work using Optimal Transport to improve model efficiency and performance across diverse tasks. By addressing biases in influence-based heuristics, TAROT enables better dataset alignment with target distributions, reducing computational costs and potential environmental impact. This approach can enhance data efficiency, benefiting scenarios where data privacy or accessibility
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
concern. However, like any selection strategy, it may introduce biases if not carefully managed. Future research should explore its implications fairness and generalization, but overall, TAROT represents step toward more efficient and responsible data selection in machine learning.
Deng. J., Li, T.-.., Zhang, S., and Ma, J. Efficient en- sembles improve training data attribution. arXiv preprint arXiv:2405. 17293, 2024.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint mrXi::2407.27565 2024.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I,, Aleman, F, L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint 3XXiv23303.0874, 2023.
Engstrom, L.., Feldmann, A., and Madry, A. Dsdm: Model- aware dataset selection with datamodels. arXiv preprint arXiv:2401.12926. 2024.
Albalak, A., Elazar, Y., Xie, M., Longpre, S., Lambert, N., Wang, X., Muennighoff, N.. Hou, B., Pan, L., Jeong, H., al. A survey data selection for language models. arXiv preprint arXiv:2402.168277 2024.
Ermolov, A., Siarohin, A., Sangineto, E., and Sebe, N. Whitening for self-supervised representation learning. In International conference on muchine learning, pp. 3015- 3024. PMLR, 2021.
Alvarez-Melis, D. and Fusi, N. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 1111422-21399 2020.
Ettinger, S., Cheng, S., Caine, B., Liu, C., Zhao, H., Prad- han, S., Chai, Y., Sapp. B., Qi, C. R., Zhou, Y., Yang, Z., Chouard, A., Sun, P., Ngiam, J., Vasudevan, V., Mc- Cauley, A., Shlens, J., and Anguelov, D. Large scale interactive motion forecasting for autonomous driving: The waymo open motion dataset. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 9710-9719, October 2021.
Borsos, Z., Mutny, M., and Krause, A. Coresets via bilevel optimization for continual learning and streaming. Ad- vances in neural information processing systems, 33: 14879-14890, 2020.
Caesar, H., Bankiti, V., Lang, A. H., Vora, S., Liong, V. E., Xu, Q., Krishman, A., Pan, Y., Baldan, G., and Beijbom, O. nuscenes: A multimodal dataset for autonomous driving. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11621-11631. 2020.
Feng, L., Bahari, M., Amor, K. M. B., Zablocki, É., Cord, M., and Alahi, A. Unitraj: A unified framework for scalable vehicle trajectory prediction. arXiv preprint arXiv:2403. 15098, 2024.

---

TA ROT: Targeted Data Selection via Optimal Transport
Girgis, R., Golemo, F., Codevilla, F., Weiss, M., D'Sza,, J. A., Kahou, S. E.,, Heide, F., and Pal, C. Latent variable sequential set transformers for joint multi-agent motion 2021.
Köpf, A., Kilcher, Y., von Rütte, D., Anagnostidis, S., Tam, Z. R., Stevens, K., Barhoum, A., Nguyen, D., Stan- ley, O., Nagyfi, R., et al. Openassistant conversations- democratizing large language model alignment. Advances in Neural Information Processing Systems, 36, 2024.
prediction. arXiv preprint arXiv:2
Guo, C.,, Zhao, B., and Bai, Y. Deepcore: A comprehensive library for coreset selection in deep learning. In Inter- national Conference on Database and Expert Systems Applications, pp. 181-195. Springer, 2022.
Liu, Z.., Karbasi, A., and Rekatsinas, T. Tsds: Data selection for task mask-specific model finetuning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
H. Caesar, J. Kabzan, K. T. e. a. Nuplan: A closed-loop ml-based planning benchmark for autonomous vehicles. In CVPR ADP3 workshop, 2021.
Liu, Z., Zhou, K., Zhao, W. X., Gao, D., Li, Y., and Wen, J.-R. Less is more: Data value estimation for visual in- struction tuning. arXiv preprint arXiv:2403.09559 2024.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Longpre, S., Hou, L.,, Vu, T., Webson, A., Chung, H. W., Tay, Y., Zhou, D., Le, Q. V., Zoph, B., Wei, J., et al. The flan collection: Designing data and methods for effec- tive instruction tuning. In International Conference on Machine Learning, pp. 22631 -22648... PMLR, 2023.
Hendrycks, D., Burns, C., Basart, S., Critch, A., Li, J., Song, D., and Steinhardt, J. Aligning ai with shared human values. Proceedings of the International Conference Learning Representations (ICLR), 2021a.
Mirzasoleiman, B., Bilmes, J., and Leskovec, J. Coresets for data-efficient training of machine learning models. International Conference on Machine Learning, pp. 5950-6960. PMLR, 2020.
Hendrycks, D.,, Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021b.
Nayakanti, N., Al-Rfou, R., Zhou, A., Goel, K., Refaat, K. S., and Sapp, B. Wayformer: Motion forecasting via simple efficient attention networks. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 2980-2987. IEEE, 2023.
Hu, Y., Hu, P., Zhao, H., and Ma, J. W. Most influential sub- set selection: Challenges, promises, and beyond. arXiv preprint arXiv:2409. 18155,, 2024.
Oquab, M., Darcet, T., Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D.,, Massa, F., El- Nouby, A., et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:230..07193, 2023.
Iyyss, A., Park, S. M., Engstrom, L., Leclerc, G., and Madry, A. Datamodels: Predicting predictions from training data. In Proceedings of the 39th International Conference on Machine Learning, 2022.
Park, S. M., Georgiev, K., Ilyas, A., Leclere, G., and Madry, A. Trak: Attributing model behavior at scale. In Inter- national Conference on Machine Learning, pp. 27074- 27113. PMLR, 2023.
Just, H. A., Kang, F., Wang, T., Zeng, Y., Ko, M., Jin, M., and Jia, R. Lava: Data valuation without pre-specified learning algorithms. In The Eleventh International Con- ference on Learning Representations. OpenReview, 2023.
Pruthi, G., Liu, F., Kale, S., and Sundararajan, M. Estimat- ing training data influence by tracing gradient descent. Advances in Neural Information Processing Systems, 33: 19920-19930, 2020.
Kang, F., Just, H. A., Sahu, A. K., and Jia, R. Performance scaling via optimal transport: Enabling data selection from partially revealed sources. Advances in Neural In- formation Processing Systems, 36, 2024.
Richter, S. R., Vineet, V., Roth, S., and Koltun, V. Play- ing for data: Ground truth from computer games. In Computer Vision-ECCV 2016: 14th European Confer- ence, Amsterdam, The Netherlands, October l 11-14, 2016, Proceedings, Part II 14, pp. 102-118. Springer, 2016.
Kantorovitch, L. On the translocation of masses. Manage- ment science, 5(1):1-4, 1958.
Killamsetty, K., Sivasubramanian, D., Ramakrishnan, G., and Iyer, R. Glister: Generalization based data subset selection for efficient and robust learning. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 8110-8118, 2021.
Sun, Z., Wang, Z., Halilaj. L., and Lucttin, J. Semantic- former: Holistic and semantic traffic scene representation for trajectory prediction using knowledge graphs. arXiv preprint arXiv:2404. 19379, 2024.
10

---

TA ROT: Targeted Data Selection via Optimal Transport
Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., Chowdhery, A., Le, Q., Chi, E., Zhou, D., et al. Challenging big-bench tasks and whether chain-of- thought can solve them. In Findings of the Association for Computational Linguistics: ACI 2023, pp. 13003-13051, 2023.
all once. Advances in Neural Information Processing Systems, 36, 2024.
Team, Q. Qwen2.5: A party of foundation models, Septem- ber 2024. URL https: /wwpn/mwemmmmmmmmm///m//mm/mmmmm github..h.......... qwwnlm.github............... /o/ blog/ wwgg///ew......
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., al. Chain-of-thougge ootthgght prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-2444 24837, 2022.
Wilson, B., Qi, W., Agarwal, T., Lambert, J., Singh, J., Khandelwal, S., Pan, B., Kumar, R., Hartnett, A., Pontes, K., Ramanan, D., Carr, P., and Hays, J. Argoverse 2: Next generation datasets for self-driving perception and forecasting. In Thirty-fifth Conference Neural In- formation Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL https::/p/// htt.hhtts:////tgtt net forum? i wKQ.e 6kv
Xia, M., Malladi, S., Gururangan, S., Arora, S., and Chen, D. Less: Selecting influential data for targeted instruc- tion tuning. In Forty-first International Conference on Machine Learning.
Xia, X., Liu, J., Yu, J., Shen, X., Han, B., and Liu, T. Mod- erate coreset: A universal method of data selection for real-world data-efficient deep learning. In The Eleventh International Conference on Learning Representations, 2022.
Xiao, W., Chen, Y., Shan, Q. Wang, Y., and Su, Feature distribution matching by optimal transport for effective and robust coreset selection. Proceedings of the AAAI Conference Artificial Intelligence, 38 (8):9196-922) 9204, Mar. 2024. doi: 10. 10.1609/aai.. v3888. 28771. URL https:/ / orroorrgg/ooooo index. http//AA/aatticceeeeeee
Yao, 7., Li, X., Lang, B., and Chuah, M. C. Goal-Ibp: Goal-based local behavior guided trajectory prediction for autonomous driving. IEEE Transactions on Intelligent Transportation Systems, 2023.
Zhang. H., Yan, B.. Cao, L., Madden, S., and Rundensteiner, E. Metastore: Analyzing deep learning meta-data at scale. Proceedings of the VLDB Endowment, 176)):444-1159) 2024.
Zou, X., Yang, J., Zhang, H., Li, Li, L., Wang, J., Wang, L.,, Gao, J., and Lee, Y. Segment everything everywhere

---

TA ROT: Targeted Data Selection via Optimal Transport
A. More Related Work
Table 3: Training Hyperparameters for Semantic Segmenta- tion
Data Influence Estimation. Pruthi et al. (2020) proposed method to estimate data influence by tracing the gradi- ent descent process. Another approach utilizes datamodels: Ilyas al. (2022) showed that simple linear datamodels can predict model outputs effectively. Building on this, Park et al. (2023) introduced a scalable technique for estimat- ing datamodel scores. Subsequent works have optimized these methods, focusing on checkpoint ensembles (Deng et al., 2024), efficient gradient storage (Zhang al., 2024), and scalable gradient projection (Choe et al., 2024). Our method extends Pruthi al. (2020) by enhancing estimation robustness through feature covariance whitening.
Hyperparameter
Value
Batch Size
16
Learning Rate
1 x 10-4
Optimizer
SGD
Momentum
0.9
Weight Decay
1 x 10-4
Learning Rate Scheduler
WarmupPolyLR
Warmup Factor
1/3
Warmup Method
Linear
Table 4: Experiment Settings for Motion Prediction
Optimal Transport for Data Selection. Optimal transport has been used for dataset similarity estimation (Alvarez- Melis & Fusi, 2020) and sensitivity analysis in data val- uation (Just al., 2023). Kang et al. (2024) applied it optimize data selection ratios across multiple sources. We integrate influence estimation with optimal transport, enabling effective subset selection to boost target domain performance.
Data Configurations
History Trajectory Time Length
2.1 seconds
Future Trajectory Time Length Object Type
6 seconds
Vehicle
Line Type
Center Lane
Model Configurations
Autobots
B. More Implementation Details
Batch Size
256
Learning Rate
0.00075
B.1. Influence Estimation
Optimizer
Adam
Gradient Clipping
5
ZCA Whitening, As discussed in Section 4.1, we use ZCA whitening to achieve higher LDS scores due to its ability to preserve the data structure. However, empirical analysis revealed that ZCA whitening has minimal impact on data selection performance. Therefore, for the data selection ex- periments, we opted for the more computationally efficient Cholesky whitening.
Wayformer
Batch Size
128
Learning Rate
10-4
Optimizer
AdamW
Gradient Clipping
5
Learning Rate Schedule
[10, 20, 30, 40, 50]
To apply ZCA whitening, we first compute the covariance matrix:
of the original data structure. This preservation of the data's 's original geometric structure may result in a better LDS score. However, this structural similarity has a limited impact on data selection tasks, where the primary goal is to identify representative diverse samples. Since both whitening methods achieve decorrelation and normalization, the rel- ative distances and informativeness of the samples, which are key factors in data selection. remain largely unaffected by the choice of whitening method.
1 N N
Σ (()))))
(15)
where Ø(=), represents the centered input features. Next, perform eigenvalue decomposition:
(16) where U is the matrix of eigenvectors, and A is the diagonal matrix of eigenvalues. The ZCA whitening transformation is then defined as: d(=)<cA= U U d(=),
Σ = UAU
Linear Data Modeling Score (LDS). As mentioned in Sec. 4.1. we employ LDS evaluation metric for influence estimation. LDS is computed by the Spearman rank cor- relation between the model outputs of models trained on randomly selected subsets and the average of attribution scores as predictions of those model outputs. For Cifar- 10, we directly reuse the provided outputs of models trained on different subsets in the Github repository of TRAK (Park et al., 2023), and utilize these outputs calculate LDS. For
(17) which ensures that the transformed features are uncorrelated and have unit variance.
The key difference between ZCA whitening and Cholesky- based whitening lies in the structure of the transformation matrix. In Cholesky whitening, the lower triangular matrix introduces decorrelation sequentially, while ZCA whiten- ing uses symmetric transformation to minimize distortion
12

---

TA ROT: Targeted Data Selection via Optimal Transport
Table 5: Training Hyperparameters for Instruction Tuning
Algorithm 2 OT-Distance Minimization Selection (OTM)
Hyperparameter
Value
Require: Candidate dataset D. (2( JE1
Target dataset D) =1 MM
Batch Size
128
Ji=1
Learning Rate
x 10-5
Number of folds K
Optimizer
Adam
Distance metric doT
LoRA rank
128
Ensure: Selected subset D. De
LoRA
512
1: Initialize D, k- 0
dropout
0.1
Split D; into K equal subsets: D(") DD()) ..... D(
Warmup Epoch
3: for fold = 1 to K do
Let D( be the current fold
5: Let ((()) < D, D(Fold)
Initialize temporary subset D( +-
nuScences, randomly select 100 subsets of the candidate datasets. Each subset is sampled to be 50% of the size of the candidate dataset. Then we independently train 100 models on these subsets and utilize their model outputs to calculate
Initialize iteration counter lc +-
8: Initialize Distance dpre <-inn"
9:
while True do
LDS.
10:
Let Dj be the set of k-th nearest candidates for each 2, E (D( based on dz
if (dor(DFFoddl Dxx D(fold) > dpre then
Algorithm 1 Fixed-Size Selection
11:
Require: Candidate dataset D.
f2e (1)
J
12:
Break the loop
13:
end
Target dataset D:
(2+ 1. li=1 Desired subset size 5 N
14:
D. t --((fold) U Dr
dpre +- -oT (D) ((oodd))
15:
Distance metric dz
16:
k tt ++++1
Ensure: Selected subset D. D. of size 8
17: end while
D. + D, (D/oll))
1: Initialize D, < 0
18:
2: Initialize iteration counter k
19: end for
3: while [D, < S do
20: Compute potential &i for each z; E D, to serve as weight w, and scale to integers.
4: Identify the k-th nearest candidates for each 2ạ E DD
based on dz
21: return D., (w;)
Let D be the
these candidates
6: if DD [Dx] S then
7:
D. D, Dr
B.3. Motion Prediction
8: else
9:
Compute potential o; for each E Dr using:
We the UniTraj (Feng al., 2024) framework for motion prediction experiments. Detailed experiment settings provided in Table 4.
à minnaut(p.uts)... IDD, d2(2, 2)) (2(,,
10:
Rank candidates in D3 by di in ascending order Add 2; E D to D. if it is among the top S |D,1 candidates in D Break the loop
B.4. Instruction Tuning
11:
We follow the settings and hyperparameters in LESS (Xia et al.). Details are provided in Table 5.
12:
13:
end if
le k+1
C. Details for Algorithms
15:
end while
16: Compute potential Φ; for each D, to serve weight wi and scale to integers.
We present the detailed data selection algorithm described in Section 3.3 in Algorithms 1 and 2
17:
return D., [w;}
D. NuScenes Leaderboard Performance
In the motion prediction task, TAROT helps Wayformer achieve the first place on the nuScenes leaderboard (Ta- ble 7). surpassing MTR- UniTra (Feng et al., 2024), which relies on large-scale datasets (2000% samples). These results
B.2. Semantic Segmentation
Table shows detailed hyperparameters for the semantic segmentation experiments.
13

---

TA ROT: Targeted Data Selection via Optimal Transport
Smaller Whitemed Feature Distunce
Larger Whitened Feature Distance
Target Example
Target muSenes Candidate: WOMD
Figure 9: More Qualitative Results of WFD in Motion Prediction. Visualization of motion prediction tasks using the nuScenes and Waymo datasets. The dark-to-light blue lines indicate the past (input) and future (ground truth) trajectorics for target vehicle, while green lines represent surrounding vehicle trajectories.
Smaller Whitened Feature Distance
Larger Whitened Feature Distance
Target Example
Target: Cityscapes Candidate: GTAS
Figure 10: More Qualitative Results of WFD in Semantic Segmentation. Semantic segmentation tasks with examples from the Cityscapes and GTAS datasets. For each target example, the top image shows the RGB scene, and the bottom image depicts the corresponding ground truth semantic segmentation mask.
14

---

TA ROT: Targeted Data Selection via Optimal Transport
Table 6:
Performance comparison (1 means higher is better). All methods use 5% of the full training data. TAROT-OTM shows relative improvement over Fixed in parentheses.
Dataset
LLaMA-3.1-8B 1
LLaMA-3.1-88 BBH
Qwen-2.5-7B MMLU1
Qwen-2.5-7B BBH
5% LESS
65.7
62.6
74.3
66.3
5% TSDS
65.2
63.1
73.9
66.2
5% TAROT TAROT-OTM
66.0
65.0
74.1
66.9
65.7 (0.13%)
63.6 (0.21%)
74.3 (0.09%)
68.9 (0.13%
demonstrate that targeted data selection offers promising alternative to data scaling for breaking performance bound- aries.
method. Besides, while TSDS uses a continuous OT--based formulation and supports subset size control via sampling, TAROT differs in its greedy, deterministic subset selection that explicitly minimizes the empirical OT distance at each iteration. This enables stronger control over the selected subset and much faster speed for data selection. Indeed, we empirically found that TSDS costs significantly more time than TAROT. The wall-clock time analysis is shown in Table 8. Finally, TAROT introduces an early stopping criterion for OT-based selection via tracking OT distance increase, allowing estimation of optimal selection ratios, which TSDS does not address.
E. Discussion and Comparison with TSDS
TSDS frames data selection as an optimal transport (OT) problem, incorporating a diversity regularizer via kernel density estimation. It selects samples based on distances in gradient embedding space.
In contrast, TAROT uses Whitened Feature Distance (WFD), which removes the confounding effects of gradi- ent covariance and scale. This provides more robust and accurate feature distance estimates, leading to better in- fluence estimation-Surpasai the state-of-the-art TRAK
We also include experimental comparisons with TSDS, fol- lowing the same evaluation settings as described in Sec- tion 4.4. The results are shown in Table 6.
F. More Qualitative Results of WFD
Table 7: nuScenes Leaderboard. Evaluated with mi- nADE5 metrie: the average of L2 distances between the predicted trajectory and ground truth over the 5 most likely predictions. Wayformer-All denotes Wayformer trained with the full candidate dataset (WOMD, Argoverse and nu- Plan).
We provide additional qualitative results of WFD in Fig- ures 9 and 10. One clear observation is that the most sup- portive samples tend to closely resemble the target sample in both the input and output spaces. In motion prediction, the most negative samples also share similarities with the target sample in the input space (e.g., map and history trajec- tories); however, their ground truth trajectories often exhibit minor differences. For instance, in the first row of samples, the most negative examples demonstrate behaviors such as turning left or driving in the right-most lanes, whereas the target sample follows trajectory in the lane adjacent to the right-most lane. These qualitative findings align with the definition of influence estimation: the most positive samples guide the model make predictions similar to the target sample, while the negative samples encourage divergent predictions.
Method
Ranking (4)minADE5 (L)
Wayformer-TARE
0.92
Wayformer-All All (Nayakanti et al., 2023)
2
0.94
MTR-UniTraj (Feng et al., 2024)
3
0.96
Goal-LBP (Yao el al., 2023)
4
1.02
SemanticFormer (Sun et al., 2024)
5
1.14
Table 8: The wall clock runtime (measured as single H100 GPU hours) of TAROT compared with LESS and TSDS on instruction tuning task. The gradient feature computation stage is the same for all the methods.
Gradient Features Computation
Data Selection
G. Computational Complexity Analysis
LESS TAROT-Fixed TAROT-OTM TSDS
46 Seconds
59 Seconds
Table 9 presents the asymptotic complexity of TAROT. The most computationally expensive step the computation of gradient features, with its cost scaling linearly with the candidate dataset size [D), the number of checkpoints N. and the gradient dimension d. However, the computational costs incurred during these two stages are one-time expenses that can be amortized across multiple target tasks. The data
32 Hours
118 Scconds
Hours
Table 9: Asymptotic complexity analysis.
Gradient Features Computation
Oixed-Siz//OTMM Selection
Complexity
C[IP) -N)
O 0(22 [ (Dael
15

---

TA TAROT: Targeted Data Selection via Optimal Transport
selection process requires minimal computation.
Table 8 demonstrates the wall-clock runtime (in single H100 GPU hours) of TAROT compared to LESS and TSDS for the instruction tuning task. The gradient feature compu- tation time is identical for both methods. For data selec- tion, both TAROT-Fixee and LESS complete in under one minute, while TAROT-OTH incurs an additional runtime of approximately one minute.
16


Reviews:
[{"review_id": "jrrKKBX7vw", "reviewer": "Reviewer_fLAa", "summary": "The authors formulate targeted data selection as a distribution matching problem and propose a new framework to efficiently select the most suitable training datasets. Massive experiments were conducted to support the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1: No supporting evidence for this claim: the linear additive assumptions inherent in greedy selection strategies are restrictive.**\n\n\nThank you for highlighting this. The limitations of linear additive assumptions in influence estimation have been thoroughly examined in the paper _Most Influential Subset Selection: Challenges, Promises, and Beyond_. We summarize their key findings that support our claim:\n\n  \n\n**Failure Mode 1: Inaccurate Influence Estimates**\n\n  \n\nEven under linear regression, influence estimates can be misleading if the **leverage score** $h_{ii}$ is ignored. The actual individual effect is: $A_{-{i}} = \\frac{x_{\\text{test}}^\\top N^{-1} x_i r_i}{1 - h_{ii}},$\nwhich differs from the estimate $v_i$ by a factor of $\\frac{1}{1 - h_{ii}}$, often resulting in underestimation for high-leverage samples.\n\n  \n**Failure Mode 2: Violation of Additivity**\n\n  \n\nEven if individual effects are accurately estimated, heuristics such as LAGS (Leverage-Adjusted Greedy Selection) may still fail due to non-additive group influence.\n\n  \n\n**Amplification**\n\nGroup influence can be _super-additive_ when samples are similar (e.g., duplicates). For $c$ identical samples $(x_i, y_i): \\frac{A_{-{i}^c}}{A_{-{i}}} = \\frac{c(1 - h_{ii})}{1 - c h_{ii}} > c.$\n\n**Cancellation**\n\nConversely, group influence can be _sub-additive_, meaning $A_{-{i,j}} < A_{-{j}}$, due to cross-leverage $h_{ij}$ and residual interactions:\n\n\n$$A_{-{i,j}} = \\frac{(1 - h_{ii})(1 - h_{jj})(A_{-{i}} + A_{-{j}}) + h_{ij} x_{\\text{test}}^\\top N^{-1}(x_i r_j + x_j r_i)}{(1 - h_{ii})(1 - h_{jj}) - h_{ij}^2}$$\n\n  \nThese findings collectively illustrate the core limitations underpinning our claim. Additionally, our empirical comparisons with data selection methods like LESS and DSDM—which rely on linear additive assumptions—reinforce this point. For instance, Figure 6 shows that samples selected by DSDM are concentrated near the center of the target distribution, failing to capture diversity effectively.\n\n----------\n\n**Q2. Results about Wall-clock time.**\n\n  \nWe provided wall-clock time results in Table 7 of the appendix. It show sthat TAROT requires comparable time for gradient computation and slightly longer time for data selection compared to baseline methods (1–2 minutes). We attached the table here for your convenience.\n\n|                          | **Gradient Features Computation** | **Data Selection** |\n|--------------------------|-----------------------------------|--------------------|\n| LESS                     | 32 Hours                          | 46 Seconds         |\n|TSDS                     | 32 Hours                          | 65 Minutes          |\n| *TAROT*-Fixed    |     32 Hours                                 | 59 Seconds         |\n| *TAROT*-OTM      |      32 Hours                                | 118 Seconds        |\n\n----------\n\n**Q3. Providing Detailed Code Instructions**\n\n  \n\nThank you for the suggestion. While we would like to update the code during the review phase, current guidelines do not permit this. We commit to releasing the full code along with detailed instructions upon paper acceptance.\n\n**Notation System**\n\nThanks for your suggestion. We agree that the current notion system need futher unification and will pay more attention to optimize it.", "claims_and_evidence": "The authors conducted massive experiments and ablation studies to support the effectiveness of each component of the proposed method. However, I found (original: do not find) no supporting evidence (original: clues) for the claim, 'This work identifies two primary factors contributing to this limitation: (ii) the restrictive linear additive assumptions inherent in greedy selection strategies.", "methods_and_evaluation_criteria": "Yes, the proposed method is well-aligned with the data selection problem. However, I think that the result about wall-clock time (corrected from wall-lock) should be included to furthur evaluate the proposed method's effectiveness.", "theoretical_claims": "There is no theoretical claims.", "experimental_designs_or_analyses": "I have reviewed the code in the supplementary material. There are no scripts or detailed instructions for reproduction, and I hope the authors can provide the detailed code for reproduction during the review process.", "relation_to_broader_scientific_literature": "No", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "The notation is abused, and I suggest that the author unify the notation system.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TWoyKmnyUj", "reviewer": "Reviewer_5vQX", "summary": "The paper proposes a data selection method for a specific target domain from a candidate set by posing it as a distribution matching problem. The paper proposes to use whitened gradient features as the base distance to compute the optimal transport between the two sets.  Effectiveness of this proposed method is shown on motion prediction, semantic segmentation and instruction tuning tasks.\n\n## Update after the rebuttal\nI thank the authors for addressing my concerns during the rebuttal. I keep my initial rating of 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Discussion and Comparison with TSDS and DSIR**\n\nThank you for the suggestion. We have revised the related work section to better highlight how **TAROT** differs from related methods, particularly **TSDS** and **DSIR**, which also address data selection from a distribution-matching perspective.\n\n**TSDS** frames data selection as an **optimal transport (OT)** problem, incorporating a diversity regularizer via kernel density estimation. It selects samples based on distances in gradient embedding space.\n\nIn contrast, **TAROT**:\n\n- Uses **Whitened Feature Distance (WFD)**, which removes the confounding effects of gradient covariance and scale. This provides more robust and accurate feature distance estimates, leading to better **influence estimation**—surpassing the state-of-the-art TRAK method.\n- While TSDS uses a continuous OT-based formulation and supports subset size control via sampling, TAROT differs in its **greedy, deterministic subset selection** that explicitly minimizes the **empirical OT distance** at each iteration. This enables stronger control over the selected subset and much faster speed for data selection. Indeed, we empirically found that TSDS cost significantly more time than TAROT. Please see our results below.\n- Introduces an **early stopping criterion** for OT-based selection via tracking OT distance increase, allowing estimation of **optimal selection ratios**, which TSDS does not address.\n\n\n**DSIR** similarly uses distribution matching but estimates **importance weights in low-dimensional n-gram space**. While efficient, this space lacks the capacity to capture high-level, task-specific semantics.\n\n**TAROT**, on the other hand:\n\n- Operates in **task- and model-specific gradient space**, and performs selection to **explicitly minimize OT distance**, making it more effective for complex or multimodal tasks like motion prediction.\n\n**Summary**: TAROT improves over TSDS and DSIR by (1) more accurate influence estimation via WFD, (2) optimal selection ratio estimation, and (3) stronger performance across domains where simpler feature spaces fall short.\n\nWe include **experimental comparisons with TSDS**. Due to time constraints and DSIR’s evaluation overlap with LESS, we do not re-run DSIR experiments.\n\nwe provide additional experiments:\n\n| Dataset         | Llama-3.1-8B MMLU ↑ | Llama-3.1-8B BBH ↑ | Qwen-2.5-7B MMLU ↑ | Qwen-2.5-7B BBH ↑ |\n|----------------|---------------------|---------------------|---------------------|---------------------|\n| 5% LESS        | 65.7                | 62.6                | **74.3**            | 66.3                |\n| 5% TSDS        | 65.2                | 63.1                | 73.9                | 66.2                |\n| 5% TAROT       | **66.0**            | **65.0**            | 74.1                | 66.9                |\n| TAROT-OTM      | 65.7 (0.13%)        | 63.6 (0.21%)        | **74.3** (0.09%)    | **68.9** (0.13%)    |\n\n**Time Complexity:**  \nWe measured the data selection time on a node with 370 GB RAM, 64 CPUs, and an H100 GPU:\n\n| Method             | LESS | TSDS       | TAROT-Fixed | TAROT-OTM  |\n|--------------------|------|------------|-------------|------------|\n| Data Selection Time| 46s  | **10 hrs** | 59s         | 118s       |\n\n**OOM Issue on Motion Prediction:**  \nTSDS could not be applied to the motion prediction task (32k samples) due to out-of-memory (OOM) errors.\n\n\n---\n\n**Q2. Target Data Ablation for OT Distance**\n\nWe performed an ablation using the **nuPlan** motion prediction dataset. Candidate data consists of 92k samples from four cities; the target set and test set are 4k and 1k held-out Boston samples respectively. We fix the selection percentage at 10% and vary the amount of target data used from 1k to 4k.\n\n| Target Data / Selected | DsDm | LESS | Random | TAROT |\n|------------------------|------|------|--------|--------|\n| 1000 / 9,200           | 3.12 | 3.07 | 2.77   | 2.33   |\n| 2000 / 9,200           | 2.97 | 3.10 | 2.77   | 2.32   |\n| 3000 / 9,200           | 3.18 | 3.04 | 2.77   | 2.17   |\n| 4000 / 9,200           | 3.00 | 3.08 | 2.77   | 2.14   |\n\n**TAROT’s performance improves steadily** as more target data is used, while baselines show no consistent gains over random selection.\n\n---\n\n**Q3. OT Distance Sample Counts**\n\nWe detail the number of samples used for OT distance computation:\n\n| **Task**                  | **Target Samples** | **Candidate Samples** | **Target/Candidate %** |\n|---------------------------|--------------------|------------------------|-------------------------|\n| Instruction Tuning — MMLU | 285                  | 270,679                 | 0.1%                  |\n| Instruction Tuning — BBH  | 81                 | 270,679                  | 0.03%                   |\n\nThese details will be included in the updated manuscript.", "claims_and_evidence": "Well supported.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No theoretical claims made.", "experimental_designs_or_analyses": "The experimental design is standard and makes sense.", "relation_to_broader_scientific_literature": "The use of OT for dataset selection is very relevant and related to the previous works in the literature.", "essential_references_not_discussed": "1. TSDS: Data Selection for Task-Specific Model Finetuning, NeurIPS 2024\n2. Data Selection for Language Models via Importance Resampling, NeurIPS 2023", "other_strengths_and_weaknesses": "Strengths\n1. The idea of selecting the data from a candidate set to improve performance on a target set via OT is important and relevant for many applications.\n2. The idea of using gradient-based features in OT formulation and the use of whitening and normalization seems to be helpful in improving performance on considered baselines. \n3. Empirical improvement across a diverse set of application shows the effectiveness of the method\n\nWeaknesses\n1. The proposed method is compares only with LESS in the instruction tuning experiment where as comparison with two recent methods TSDS and DSIR should also be included. (The titles of the works are mentioned above.)\n\n2. The amount of target data required to estimate the OT distances is not ablated for the three applications.\n\n3. How many samples from source and target are used to compute the OT distance?", "other_comments_or_suggestions": "NA", "questions_for_authors": "See Weaknesses above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gjB1dL55QN", "reviewer": "Reviewer_UG2a", "summary": "The paper introduces a framework for targeted data selection by minimizing the distance between selected data and target data distribution. The method addresses the limitations of existing influence-based greedy heuristics, which does not perform well on multimodal data distributions. The framework is evaluated across multiple tasks, including semantic segmentation, motion prediction, and instruction tuning, demonstrating consistent improvements over state-of-the-art methods. The authors also provide a detailed analysis of the computational complexity and ablation studies to validate the effectiveness of their approach.\n\n## Update after rebuttal\nThank you for the responses. I keep my original score of 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Difference Between Data Distillation and Valuation**\n\n  \n\nThank you for this insightful comment. We appreciate the opportunity to clarify the relationship between TAROT and other data selection paradigms, particularly **Data Valuation via Reinforcement Learning (DVRL)** [Yoon et al., ICML 2020] and **Data Distillation**.\n\n\n**1. Comparison to DVRL (Data Valuation via Reinforcement Learning)**\n\n  \n\nDVRL frames data valuation as a meta-learning problem, where a Data Value Estimator (DVE) is trained using reinforcement learning (RL) to assign importance weights to training samples. The reward signal is derived from the performance of a predictor model on a validation set. This setup allows DVRL to dynamically prioritize “useful” samples and has shown promising results in settings with noisy labels or domain shifts.\n\n  \n\n**TAROT departs from DVRL in several key aspects:**\n\n• **Targeted vs. General Data Value Estimation:**\n\nDVRL learns instance-wise importance scores primarily to boost general predictive performance. In contrast, TAROT is explicitly designed for _targeted_ selection: it identifies a subset of data that minimizes the **optimal transport (OT) distance** to a distinct target distribution. This enables TAROT to adapt to complex, potentially multimodal target domains—something DVRL is not explicitly formulated to handle.\n\n• **Model-Agnostic Transferability:**\n\nTAROT computes data selection using WFD-based embeddings from a lightweight model, but supports downstream training with larger, task-specific models. As demonstrated in Section 4.3, the selected data generalizes across diverse architectures and tasks (e.g., motion prediction, instruction tuning). DVRL, by contrast, couples valuation tightly to the predictor, reducing flexibility for out-of-distribution generalization or transfer across models.\n\n• **Computational Simplicity and Scalability:**\n\nDVRL’s RL-based training is often computationally intensive and sensitive to hyperparameters, especially in high-dimensional settings. TAROT employs a more efficient and stable OT-based greedy selection method, using gradient feature whitening and normalization to reduce dominant component bias. As shown in our runtime analysis (Appendix F), TAROT scales effectively to large datasets (e.g., 2M motion samples) without requiring complex policy optimization.\n\n----------\n\n**2. Comparison to Data Distillation Approaches**\n\n  \n\nData distillation typically involves synthesizing or filtering datasets to mimic the performance of a teacher model. Although both data distillation and TAROT aim to curate more efficient training sets, their approaches and goals differ significantly:\n\n• **Objective Focus:**\n\n•  _Data Distillation_ is model-centric, often relying on teacher-student frameworks or aligning predictions/logits.\n\n•  _TAROT_ is distribution-centric, selecting data that explicitly aligns with the **target distribution** using whitened OT distances computed over gradient features.\n\n• **Dependency on a Teacher Model:**\n\nTAROT does not require a high-performing teacher model. Instead, it operates in scenarios where such a model may not exist, making it broadly applicable.\n\n• **Automatic Selection Ratio:**\n\nTAROT automatically infers optimal selection ratios through OT distance minimization (Section 3.3), a capability typically absent in distillation-based methods.\n\n----------\n\nThank you again for this constructive suggestion. It has helped us better articulate TAROT’s position within the broader landscape of data selection research.", "claims_and_evidence": "Yes, the claims are supported by clear evidence.", "methods_and_evaluation_criteria": "I am not deeply familiar with the empirical work and existing evaluation approaches in this field, the evaluation criteria used in the paper either align with established practices from prior research or appear to be well-justified and reasonable.", "theoretical_claims": "The work is mostly empirical and proposes no major theoretical claims.", "experimental_designs_or_analyses": "While I am not an expert in assessing the validity of experimental design within this specific field, the experimental designs appear to adhere to standard practices.", "relation_to_broader_scientific_literature": "The method addresses the limitations of existing influence-based greedy heuristics (Xia et al.; Engstrom et al., 2024) for data selection, which does not perform well on multimodal data distributions as noted by  (Hu et al., 2024). It is based on the Optimal Transport theory, which was also used in (Just et al., 2023).", "essential_references_not_discussed": "I am not aware of significant missing references that are essential to understanding the key contributions of the paper.", "other_strengths_and_weaknesses": "Strengths: \n- The paper proposes a novel approach for mitigating a major weakness of the previous data selection method, namely, the summed influence fail to find diverse data. The key technique is novel: they use whitening to decorrelate the features and select data by OT distance minimization. They show by extensive experiments that TAROT mitigates biases from dominant features and outperforms previous methods in tasks including semantic segmentation, motion prediction, and instruction tuning.\n\nWeaknesses: \n- The paper follows a line of work in data selection, but does not seem to discuss other existing approaches for similar data selection problems. For example, the literature on data distillation and the following paper,\n\"Data Valuation using Reinforcement Learning\", Jinsung Yoon, Sercan Arik, Tomas Pfister Proceedings of the 37th International Conference on Machine Learning, PMLR.", "other_comments_or_suggestions": "The paper is well written.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yILDEgbWYk", "reviewer": "Reviewer_oARp", "summary": "This paper proposes to address the problem of task-specific sample selection from the perspective of distribution matching that can be solved by optimal trasnport. Experiments on influence estimation, semantic segmentation, motion prediction, and instruction tuning show the effectivenss of the proposed method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> ### Q3: Clarifying Modality\n\nBy **modality**, we refer to **distributional modality**, rather than input modality. We elaborate on why other tasks exhibit greater distributional multi-modality compared to **instruction tuning**, as summarized in the table below:\n\n|               | Motion Prediction                               | Semantic Segmentation                                        | Instruction Tuning           |\n|---------------------|--------------------------------------------------|---------------------------------------------------------------|------------------------------|\n| Target Dataset Size | 32,186                                           | 2,975                                                         | 81 (BBH) / 285 (MMLU)        |\n| Domain Coverage     | Diverse traffic scenarios from two cities | Stereo video scenes from 50 different cities                  | Specific LLM task types      |\n\nAs shown, both motion prediction and semantic segmentation involve significantly larger datasets and broader domain coverage, making data selection inherently more challenging. Our method, validated across multiple tasks, demonstrates better generalization in such multi-modal distributional settings.\n\n---\n\n> ### Q2: Comparison with LESS\n\nWe respectfully disagree with the assessment that TAROT offers only limited improvement over LESS. As discussed above, TAROT yields **substantial gains** on tasks characterized by high distributional multimodality:\n\n- Compared to LESS, **TAROT achieves:**\n  - **+44%** improvement in **semantic segmentation**\n  - **+102%** improvement in **motion prediction**\n  - Even in **instruction tuning**, which has relatively lower modality, TAROT outperforms LESS using just **2%** of the data.\n\n**Computational Cost:**  \nWhile TAROT incurs additional cost due to OT distance computation, this cost is minimal (118 seconds) compared to gradient calculation time (32 hours). Moreover, gradient calculation is a one-time cost and can be amortized across multiple tasks.\n\n\n\n---\n\n> ### Comparison with TSDS\n\nAs noted in our response to Reviewer 5vQX, we discuss differences with TSDS in detail. Here, we provide additional experiments:\n\n| Dataset         | Llama-3.1-8B MMLU ↑ | Llama-3.1-8B BBH ↑ | Qwen-2.5-7B MMLU ↑ | Qwen-2.5-7B BBH ↑ |\n|----------------|---------------------|---------------------|---------------------|---------------------|\n| 5% LESS        | 65.7                | 62.6                | **74.3**            | 66.3                |\n| 5% TSDS        | 65.2                | 63.1                | 73.9                | 66.2                |\n| 5% TAROT       | **66.0**            | **65.0**            | 74.1                | 66.9                |\n| TAROT-OTM      | 65.7 (0.13%)        | 63.6 (0.21%)        | **74.3** (0.09%)    | **68.9** (0.13%)    |\n\n**Time Complexity:**  \nWe measured the data selection time on a node with 370 GB RAM, 64 CPUs, and an H100 GPU:\n\n| Method             | LESS | TSDS       | TAROT-Fixed | TAROT-OTM  |\n|--------------------|------|------------|-------------|------------|\n| Data Selection Time| 46s  | **10 hrs** | 59s         | 118s       |\n\n**OOM Issue on Motion Prediction:**  \nTSDS cannot be applied to the motion prediction task (32k samples) due to out-of-memory (OOM) errors.\n\n**Additional Results on TydiQA:**  \nDue to resource constraints, we focused on BBH/MMLU. For completeness, we now include results on TydiQA (only 9 target samples):\n\n| Dataset                | All | 5% Random | 5% LESS | 5% TAROT | TAROT-OTM         |\n|------------------------|-----|-----------|---------|----------|-------------------|\n| Llama-3.1-8B TydiQA ↑  | 63.1| 61.0      | 69.2    | **71.1** | 66.9 (0.05%)      |\n\n---\n> ### Q5: Motivation for Whitened Feature Distance\n\nWe visualized the **[covariance matrix of raw gradient features](https://postimg.cc/SXSkcZrX)**, which reveals strong correlations—supporting our motivation for feature whitening. Additionally, this is conceptually aligned with the findings in _“Whitening for Self-Supervised Representation Learning”_, which we will cite in the revision.\n\n---\n> ### Q1: Clarifications\n\nWe appreciate the feedback and will revise the manuscript for clarity. Here are key clarifications:\n\n- **Fixed-Size vs. OTM:**  \n  The 5%, 10%, 20%, and 50% results refer to **TAROT-Fixed**. **TAROT-OTM** dynamically selects the ratio that minimizes OT distance.\n\n- **Data Weighting Implementation:**  \n  Weighting is applied by **repeating samples** during training based on their assigned weights.\n\n- **Training Overhead:**  \n  Due to increased computation, we omit data weighting in the **motion prediction** task to demonstrate TAROT’s effectiveness without any training overhead.\n\n- **Equation 13:**  \n  Refers to the **Whitened Feature Distance (WFD)**. \n\n---\n\n> ### Q4 & Q6\n\nThank you for the suggestions. We will ensure consistent use of the discrete formulation throughout the paper and include citations in the revised version.", "claims_and_evidence": "Please refer to the weakness.", "methods_and_evaluation_criteria": "Yes. But lack of some important baselines. Please refer to the weakness and experimental design.", "theoretical_claims": "Yes. correct.", "experimental_designs_or_analyses": "The author did not report the results of TYDIQA like LESS and [1] did. Could the author please explain the reason for this? Is it due to experimental limitations or other considerations? It is necessary to clarify this to make the research more complete and comparable with existing works.\n\n[1] TSDS: Data Selection for Task-Specific Model Finetuning. NIPS 2024.", "relation_to_broader_scientific_literature": "Please refer to the weakness and essential references.", "essential_references_not_discussed": "The key contribution is an OT-driven task-specific sample selection paradigm. However, a closely related paper has been overlooked [1], which also employs OT for task-specific data selection and can be based on either semantic embedding or gradient (in particular, this submission paper adpots gradient). On the other hand, both this submission and [1] utilize nearest neighbor search during OT-based sample selection in order to reduce computational cost. in my oponion, it's necessary to discuss the connections and differences between this submission and [1], as well as to provide experimental comparisons with [1].\n\n[1] TSDS: Data Selection for Task-Specific Model Finetuning. NIPS 2024.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The introduction of OT to task-specific data selection is reasonable.\n2. Extensive experiments on different settings prove the effectiveness.\n\n\nWeakness:\n\n1. The author attempts to introduce more content into the paper, which makes it difficult to grasp the key points of the paper. For example, the author mentions a variant of the fixed-size selection method (TAROT-FSS), and the experimental results corresponding to it are TAROT-5%/20%/50%? Also, regarding data weighting (Sec.3.4), is it used to weight the samples when constructing the marginal distributions alpha or beta? The definition and significance of this part are not very clear. Moreover, the repetition factor is only used in Section 4.2. Why is it not considered in Sections 4.3 and 4.4? The experimental setup is somewhat confusing.\n\n2. The improvement in experimental results compared to LESS seems minimal, especially in Figures 5 and 7. However, Table 7 shows that TAROT requires significantly more computational cost and time than LESS. This makes the method in the paper appear to lack a good trade-off between performance improvement and computational cost.\n\n3. The author mentions in the abstract (Line 016) that \"These methods perform well on limited, unimodal data (i.e., data following a single pattern) but become less effective as target data increases in complexity.\" However, I don't see design or experiment on sample selection for multi-modal datasets in this paper. Perhaps I missed some details. Maybe Motion Prediction? The Motion Prediction task (Waymo) involves multiple modalities, such as RGB and point clouds. However, from a task setting perspective, motion prediction based on historical data is often considered a single-modality task. Also, the experimental details seem to lack clarification on which modalities were used and how different input data types interacted during the data selection.\n\n4. When introducing OT, the author uses the integral form under continuous distributions in Eq. 3. However, when formulating the marginal distributions alpha and beta, a discrete form is used. Considering that the author employs entropic OT (Sinkhorn), it would be more consistent to use the discrete form of OT notation throughout.\n\n5. Line 161: \"due to the correlation of gradient feature\" lacks necessary references, experiments, or observations to support it. This makes the motivation for the Whitened Feature Distance not entirely convincing.\n\n6. Cholesky whitening lacks relevant references.", "other_comments_or_suggestions": "Please refer to the Weakness.", "questions_for_authors": "What's the implementation of $c(z, z_t)$ in Eq.13? Cosine distance or WFD $d_{\\mathcal{Z}}^{w}(z,z')$? \n\nPlease refer to the weakness.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors formulate data selection for training as a distribution matching problem and apply OT to solve it. They conduct extensive experimental evaluations to demonstrate the efficacy of their method. Additionally, they have addressed the main concerns raised by the reviewers, namely wall-clock time and comparison to TSDS. I recommend acceptance.
Decision: Accept (poster)