TAROT: Targeted Data Selection via Optimal Transport

Q3: Clarifying Modality

By 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:

As 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.

---

Q2: Comparison with LESS

We 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:

• Compared to LESS, TAROT achieves:
  • +44% improvement in semantic segmentation
  • +102% improvement in motion prediction
  • Even in instruction tuning, which has relatively lower modality, TAROT outperforms LESS using just 2% of the data.

Computational Cost:
While 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.

---

Comparison with TSDS

As noted in our response to Reviewer 5vQX, we discuss differences with TSDS in detail. Here, we provide additional experiments:

Time Complexity:
We measured the data selection time on a node with 370 GB RAM, 64 CPUs, and an H100 GPU:

OOM Issue on Motion Prediction:
TSDS cannot be applied to the motion prediction task (32k samples) due to out-of-memory (OOM) errors.

Additional Results on TydiQA:
Due to resource constraints, we focused on BBH/MMLU. For completeness, we now include results on TydiQA (only 9 target samples):

---

Q5: Motivation for Whitened Feature Distance

We visualized the covariance matrix of raw gradient features, 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.

---

Q1: Clarifications

We appreciate the feedback and will revise the manuscript for clarity. Here are key clarifications:

• Fixed-Size vs. OTM:
  The 5%, 10%, 20%, and 50% results refer to TAROT-Fixed. TAROT-OTM dynamically selects the ratio that minimizes OT distance.

• Data Weighting Implementation:
  Weighting is applied by repeating samples during training based on their assigned weights.

• Training Overhead:
  Due to increased computation, we omit data weighting in the motion prediction task to demonstrate TAROT’s effectiveness without any training overhead.

• Equation 13:
  Refers to the Whitened Feature Distance (WFD).

---

Q4 & Q6

Thank you for the suggestions. We will ensure consistent use of the discrete formulation throughout the paper and include citations in the revised version.

Difference Between Data Distillation and Valuation

Thank 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.

1. Comparison to DVRL (Data Valuation via Reinforcement Learning)

DVRL 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.

TAROT departs from DVRL in several key aspects:

• Targeted vs. General Data Value Estimation:

DVRL 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.

• Model-Agnostic Transferability:

TAROT 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.

• Computational Simplicity and Scalability:

DVRL’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.

---

2. Comparison to Data Distillation Approaches

Data 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:

• Objective Focus:

• Data Distillation is model-centric, often relying on teacher-student frameworks or aligning predictions/logits.

• TAROT is distribution-centric, selecting data that explicitly aligns with the target distribution using whitened OT distances computed over gradient features.

• Dependency on a Teacher Model:

TAROT does not require a high-performing teacher model. Instead, it operates in scenarios where such a model may not exist, making it broadly applicable.

• Automatic Selection Ratio:

TAROT automatically infers optimal selection ratios through OT distance minimization (Section 3.3), a capability typically absent in distillation-based methods.

---

Thank you again for this constructive suggestion. It has helped us better articulate TAROT’s position within the broader landscape of data selection research.

Q1. Discussion and Comparison with TSDS and DSIR

Thank 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.

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 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.
• 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.
• 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.

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.

TAROT, on the other hand:

• 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.

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.

We include experimental comparisons with TSDS. Due to time constraints and DSIR’s evaluation overlap with LESS, we do not re-run DSIR experiments.

we provide additional experiments:

Time Complexity:
We measured the data selection time on a node with 370 GB RAM, 64 CPUs, and an H100 GPU:

OOM Issue on Motion Prediction:
TSDS could not be applied to the motion prediction task (32k samples) due to out-of-memory (OOM) errors.

---

Q2. Target Data Ablation for OT Distance

We 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.

TAROT’s performance improves steadily as more target data is used, while baselines show no consistent gains over random selection.

---

Q3. OT Distance Sample Counts

We detail the number of samples used for OT distance computation:

These details will be included in the updated manuscript.

Q1: No supporting evidence for this claim: the linear additive assumptions inherent in greedy selection strategies are restrictive.

Thank 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:

Failure Mode 1: Inaccurate Influence Estimates

Even 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}},$ which differs from the estimate $v_i$ by a factor of $\frac{1}{1 - h_{ii}}$, often resulting in underestimation for high-leverage samples.

Failure Mode 2: Violation of Additivity

Even if individual effects are accurately estimated, heuristics such as LAGS (Leverage-Adjusted Greedy Selection) may still fail due to non-additive group influence.

Amplification

Group 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.$

Cancellation

Conversely, group influence can be sub-additive, meaning $A_{-{i,j}} < A_{-{j}}$, due to cross-leverage $h_{ij}$ and residual interactions:

$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}$

These 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.

---

Q2. Results about Wall-clock time.

We 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.

---

Q3. Providing Detailed Code Instructions

Thank 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.

Notation System

Thanks for your suggestion. We agree that the current notion system need futher unification and will pay more attention to optimize it.