Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Data Selection at Scale via Influence Distillation

Mahdi Nikdan, Vincent Cohen-Addad, Dan Alistarh, Vahab Mirrokni

TL;DR

This paper introduces Influence Distillation, a novel, mathematically-justified framework for data selection that employs second-order information to optimally weight training samples by distilling each sample's influence on a target distribution by assigning model-specific weights that are used to select training data for LLM fine-tuning.

Abstract

Effective data selection is critical for efficient training of modern Large Language Models (LLMs). This paper introduces Influence Distillation, a novel, mathematically-justified framework for data selection that employs second-order information to optimally weight training samples. By distilling each sample's influence on a target distribution, our method assigns model-specific weights that are used to select training data for LLM fine-tuning, guiding it toward strong performance on the target domain. We derive these optimal weights for both Gradient Descent and Adam optimizers. To ensure scalability and reduce computational cost, we propose a $\textit{landmark-based approximation}$: influence is precisely computed for a small subset of "landmark" samples and then efficiently propagated to all other samples to determine their weights. We validate Influence Distillation by applying it to instruction tuning on the Tulu V2 dataset, targeting a range of tasks including GSM8k, SQuAD, and MMLU, across several models from the Llama and Qwen families. Experiments show that Influence Distillation matches or outperforms state-of-the-art performance while achieving up to $3.5\times$ faster selection.

Efficient Data Selection at Scale via Influence Distillation

TL;DR

This paper introduces Influence Distillation, a novel, mathematically-justified framework for data selection that employs second-order information to optimally weight training samples by distilling each sample's influence on a target distribution by assigning model-specific weights that are used to select training data for LLM fine-tuning.

Abstract

Effective data selection is critical for efficient training of modern Large Language Models (LLMs). This paper introduces Influence Distillation, a novel, mathematically-justified framework for data selection that employs second-order information to optimally weight training samples. By distilling each sample's influence on a target distribution, our method assigns model-specific weights that are used to select training data for LLM fine-tuning, guiding it toward strong performance on the target domain. We derive these optimal weights for both Gradient Descent and Adam optimizers. To ensure scalability and reduce computational cost, we propose a : influence is precisely computed for a small subset of "landmark" samples and then efficiently propagated to all other samples to determine their weights. We validate Influence Distillation by applying it to instruction tuning on the Tulu V2 dataset, targeting a range of tasks including GSM8k, SQuAD, and MMLU, across several models from the Llama and Qwen families. Experiments show that Influence Distillation matches or outperforms state-of-the-art performance while achieving up to faster selection.

Paper Structure

This paper contains 67 sections, 6 theorems, 60 equations, 9 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Contribution.
  3. Related Work
  4. Method
  5. Problem and Notation
  6. A Running Example
  7. Influence Distillation
  8. Case 1: Unconstrained Weights.
  9. Discussion.
  10. Case 2: Robust Weights.
  11. The Robust Influence Objective.
  12. Adam Optimizer.
  13. Handling Variable Lebel Lengths.
  14. Per-target Importance.
  15. Efficient Influence Distillation
  16. ...and 52 more sections

Key Result

Theorem 4.1

(Informal version of Theorem thm:approx-grad and Corollary cor:approx-grads tailored to landmark-based approximation -- see Appendix apx:cluster-bound) Consider the special case of first-order Influence Distillation. Let $\mathbf{g}_i$ and $\hat{\mathbf{g}}_i$ denote the true and the landmark-based Then $\mathop{\mathbb{E}[||\boldsymbol{w} - \hat{\boldsymbol{w}})||^2]} \le \frac{|S| \Delta^2}{\la

Figures (9)

  • Figure 1: Average improvement over uniform sampling across six tasks vs. runtime. The model used is Llama2-7B llama2, and the training dataset is Tulu V2 tulu2. The annotation "M/N" indicates that the method selected M samples from a pool of size N. Further details are provided in Section \ref{['sec:exp']}.
  • Figure 2: (Left) Distribution of unconstrained weights, (Middle) Distribution of robust weights for $\lambda\mathbin{=}0.02$, and (Right) validation loss during training with different variants in the running experiment setting. Robust weights are found by minimizing Objective \ref{['eq:robust-obj']} using the SLSQP algorithm slsqp implemented in the SciPy library scipy.
  • Figure 3: (Left) Effect of the number of landmarks on the performance of Influence Distillation across six tasks using Llama2-7B. (Right) MMLU accuracy of Influence Distillation on Llama2-7B across different pool sizes and number of selected samples.
  • Figure 4: Average gradient cosine similarity on unseen samples from Tulu V2 (top) and BBH (bottom) across checkpoints.
  • Figure 5: Correlation between gradient norm and number of label tokens, across checkpoints on four datasets.
  • ...and 4 more figures

Theorems & Definitions (11)

  • Theorem 4.1
  • Lemma B.1
  • proof
  • Lemma D.1
  • proof
  • Lemma D.2
  • proof
  • Theorem D.3
  • proof
  • Corollary D.4
  • ...and 1 more