Technical Summary: DUET – Optimizing LLM Training Data Mixtures via Noisy Feedback from Unseen Evaluation Tasks

1. Problem Statement

The performance of Large Language Models (LLMs) is heavily dependent on the alignment between the training data domain and the downstream evaluation task. However, in many practical deployment scenarios, the specific data distribution, domain, or labels of the evaluation task are unknown. This is referred to as the unseen evaluation task setting.

Common examples include:

• End-to-end encrypted conversations: LLM owners cannot access the actual user-chatbot conversation data due to privacy constraints (e.g., enterprise privacy policies) but can observe coarse feedback such as user ratings or session duration.
• Model marketplaces: A model owner sells a model to a consumer who provides feedback on performance without sharing the sensitive evaluation data used by the consumer.

In these settings, the owner receives only coarse, noisy feedback (e.g., ratings, task completion time) rather than fine-grained data information (gradients, labels, or distributions). Existing data mixing methods (e.g., DoReMi, BiMix) and data selection techniques (e.g., LESS, coresets) typically require access to fine-grained evaluation data or assume the training distribution matches the task distribution. Consequently, they fail or perform sub-optimally when applied directly to settings where the evaluation data is inaccessible. The core challenge is: How can one efficiently optimize the LLM training data mixture using only multiple rounds of noisy feedback from an unseen task?

2. Methodology: DUET

The paper proposes DUET (Optimizing LLM Training Data Mixtures via Noisy Feedback from Unseen Evaluation Tasks), a novel global-to-local algorithm that interleaves data selection with Bayesian Optimization (BO).

2.1 Problem Reparameterization

The authors reparameterize the optimization problem. Instead of directly searching for the optimal set of data points $X^*$, they search for the optimal mixing ratio $r \in \mathbb{R}^n$ across $n$ training data domains. The problem is formulated as a bilevel optimization:

• Outer Level (Global): Find the optimal mixing ratio $r$ that minimizes the unseen evaluation task loss.
• Inner Level (Local): For a given ratio $r$, find the high-quality subset of data $X \in S_r$ (where $S_r$ is the set of data mixtures satisfying ratio $r$) that minimizes the loss.

2.2 The Inner Loop: Data Selection

To solve the inner problem, DUET employs a data selection strategy to retrieve high-quality data points from each domain according to the proposed ratio $r$.

• Influence Functions (IF): The paper primarily utilizes Influence Functions as the selection mechanism. For each domain, a separate LLM is fine-tuned, and IF scores are computed for every data point relative to a validation set.
• IF-Weighted Sampling: Instead of uniform random sampling, DUET performs weighted sampling where the probability of selecting a data point is proportional to its IF score. This biases the selection toward data points that are more "influential" in reducing model loss.
• Estimator: The inner problem solution is estimated by training an LLM on the selected mixture and evaluating it on the unseen task. The authors theoretically characterize the distribution of this estimator, showing that IF-driven sampling reduces variance and bias compared to uniform random sampling.

2.3 The Outer Loop: Bayesian Optimization

The outer loop uses Bayesian Optimization to iteratively refine the mixing ratio $r$.

• Black-Box Objective: The function $f(r)$ represents the unseen task loss for a given ratio $r$. Since this loss is obtained via the noisy inner-loop estimator, $f(r)$ is treated as a black-box function with noisy observations.
• Gaussian Process (GP): DUET models $f(r)$ using a GP prior.
• Acquisition Function: At each iteration $t$, DUET proposes a new mixing ratio $r_t$ by minimizing the Lower Confidence Bound (LCB) acquisition function, balancing exploration and exploitation.
• Feedback Loop: The proposed ratio $r_t$ is used to sample data, train/fine-tune the LLM, and observe the noisy feedback. This observation is added to the history to update the GP posterior, and the process repeats for $T$ iterations.

2.4 Algorithm Flow

1. Initialization: Pre-compute IF scores for all data points in training domains. Initialize BO with random data mixtures.
2. Iteration:
   • BO proposes a mixing ratio $r_t$.
   • Perform IF-weighted sampling to generate a data mixture $X_t$ satisfying $r_t$.
   • Fine-tune the LLM on $X_t$ and observe the noisy feedback (evaluation loss) from the unseen task.
   • Update the GP posterior with $(r_t, \text{feedback})$.
3. Termination: After $T$ iterations, return the best-performing data mixture observed.

3. Key Contributions

1. Novel Problem Setting: The paper defines a realistic setting where evaluation data is unknown (unseen task) but coarse, noisy feedback is available, distinguishing it from standard domain adaptation or generalization tasks.
2. DUET Algorithm: Introduction of a global-to-local framework that interleaves data selection (specifically IF-driven sampling) with Bayesian Optimization to optimize data mixtures without fine-grained task data.
3. Theoretical Guarantees: The authors provide a theoretical analysis of DUET's convergence by bounding the cumulative regret. They show that DUET converges to the optimal training data mixture even with noisy feedback, provided the number of BO iterations increases.
4. Empirical Validation: Extensive experiments demonstrate that DUET outperforms existing data mixing (DoReMi, Aioli) and selection (LESS) baselines in both in-domain and out-of-domain unseen task settings.

4. Experimental Results

The experiments were conducted using Llama-3-8b-Instruct and Qwen2.5-7B-Instruct fine-tuned via PEFT (LoRA).

• Datasets: 9 training domains (e.g., Wikitext, GSM8K, PubMedQA, TruthfulQA) and various unseen evaluation tasks (some in-domain, some out-of-domain).
• Baselines: Compared against DoReMi, LESS, Aioli, Multi-fidelity BO, and naive approaches (uniform mixing, random search).
• Performance:
  • DUET consistently achieved higher performance scores across all evaluation tasks compared to baselines.
  • In in-domain tasks (e.g., TruthfulQA), DUET successfully identified and weighted the relevant training domain heavily.
  • In out-of-domain tasks (e.g., GSM8K without GSM8K training data), DUET effectively leveraged feedback to find relevant proxy domains (e.g., Wikitext, SciQ), outperforming methods that could not adapt to the feedback.
• Ablation Studies:
  • Data Selection: DUET with IF-driven sampling (DUET-IF) outperformed DUET with uniform random sampling (DUET-UR) and other selection methods like LESS and log-determinant.
  • Sampling Size ($k$): Increasing the number of samples $k$ in the inner loop improved performance, though $k=1$ was sufficient to beat baselines.
  • Components: Both the BO component and the data selection component were shown to be necessary for optimal performance.

5. Significance and Claims

The paper claims that DUET is the first work to interleave data selection with Bayesian Optimization to iteratively optimize training data mixtures based on feedback from an unseen evaluation task.

• Practical Relevance: It addresses a critical gap in LLM deployment where privacy or data ownership prevents access to evaluation data, yet iterative improvement is desired.
• Efficiency: By using BO, the method is sample-efficient, requiring fewer rounds of training/fine-tuning compared to brute-force enumeration of data mixtures.
• Flexibility: The framework is flexible enough to incorporate different data selection methods (IF, LESS, coresets) depending on the computational budget, though IF is shown to be the most effective in their setting.
• Theoretical Soundness: The convergence analysis provides a theoretical foundation for why the algorithm works despite the noise and lack of fine-grained data.

The authors note that while the current focus is on fine-tuning, the approach is theoretically adaptable to pre-training, though they do not claim to have solved pre-training optimization in this work. They also acknowledge that computing IF scores can be computationally intensive but suggest pre-computation and surrogate models as practical mitigations.