🤖 machine learning

Distill on a Diet: Efficient Knowledge Distillation via Learnable Data Pruning

The paper proposes IF-Beta, an efficient knowledge distillation framework that leverages influence functions and a learnable Beta-distribution-based sampling policy to optimize data pruning, enabling student models trained on significantly reduced data and compute to outperform those distilled on full datasets.

Original authors: Yifan Wu, Yiqi Wang, Xichen Ye, Wenjing Yan, Xiaoqiang Li, Cheng Jin, Xiangyu Yue, Weizhong Zhang

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yifan Wu, Yiqi Wang, Xichen Ye, Wenjing Yan, Xiaoqiang Li, Cheng Jin, Xiangyu Yue, Weizhong Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Student with a "Diet"

Imagine you are a master chef (the Teacher) trying to teach an apprentice (the Student) how to cook a perfect dish. Usually, to teach the apprentice, you make them practice on every single recipe in your massive cookbook. This takes forever, burns a lot of fuel (computing power), and is exhausting.

Knowledge Distillation (KD) is the process of the master chef guiding the apprentice. The goal is to make the apprentice so good that they can cook almost as well as the master, but much faster and with less equipment.

However, the paper points out a problem: Even though the final apprentice is small and efficient, the training process is still huge because the apprentice has to read the entire cookbook. The authors ask: "What if we could just give the apprentice a carefully selected 'diet' of the most important recipes? Could they learn just as well, but faster and with less effort?"

This is the core idea of Distill on a Diet.

The Problem with Current Methods

Before this paper, people tried to pick the "best" recipes for the apprentice using two main methods, both of which had flaws:

The "Re-Training" Method: To know which recipes are hard or easy, you have to watch the apprentice cook through the whole book first. This defeats the purpose of saving time! It's like hiring a detective to watch the apprentice for a month just to figure out which recipes to show them.
The "Rule of Thumb" Method: Some people just use simple rules, like "Pick the middle-difficulty recipes." But this is rigid. Sometimes the "middle" recipes aren't the best; sometimes you need a mix of very easy and very hard ones. It's like a strict diet plan that doesn't adapt to what your body actually needs.

The Solution: IF-Beta

The authors propose a new system called IF-Beta. It combines two smart ideas to pick the perfect "diet" of data without needing to re-train the model first.

1. The "Influence Function" (The Crystal Ball)

Imagine you have a crystal ball that can tell you: "If we removed this specific recipe from the cookbook, would the apprentice's cooking get worse or better?"

In the paper, this is called the Influence Function (IF).

Old way: This crystal ball was usually broken or too expensive to use.
New way: The authors found a way to make this crystal ball work efficiently. They use a technique called Flat Validation Minima (FVM). Think of this as "sharpening" the teacher's intuition. By making the teacher's knowledge landscape "flatter" and more stable, the crystal ball can accurately predict which data points are truly important without needing to see the apprentice struggle through them first.

2. The "Beta Policy" (The Flexible Chef)

Once you know which recipes are important (via the crystal ball), how do you pick them?

Old way: Use a rigid ruler. "Take the top 10% hardest" or "Take everything between score 5 and 7." This is like a diet that says "Eat exactly 3 apples and 2 bananas," regardless of how hungry you are.
New way: The authors use a Beta Policy. Imagine a flexible, stretchy net that can change its shape. Instead of a rigid ruler, this net learns the perfect shape to catch the best data. It can stretch to catch mostly hard examples, or mostly easy ones, or a perfect mix, depending on what works best for the specific teacher and student. It's a "learnable" diet that adapts to the situation.

How It Works (The "Bilevel" Dance)

The system works like a two-step dance to find the perfect diet:

The Inner Loop (The Practice): The system quickly simulates the apprentice learning on a small, temporary diet of data. To make this super fast, it doesn't train a whole new brain; it just trains a simple "linear head" (like a quick mental note) on top of the teacher's frozen knowledge.
The Outer Loop (The Coach): The system checks how well the apprentice did on this temporary diet. If the diet was good, the "Coach" (the Beta Policy) adjusts the shape of the net to catch even better data next time. If the diet was bad, the Coach reshapes the net.

This happens very quickly, allowing the system to find the optimal subset of data without the heavy cost of full training.

The Results: Less Food, Better Cooking

The paper tested this on famous image datasets (like CIFAR and ImageNet). The results were surprising:

Better than Full Data: In many cases, the apprentice trained on just 50% to 70% of the data (using the IF-Beta diet) actually performed better than an apprentice trained on 100% of the data.
Faster and Cheaper: Because they used less data, the training took significantly less time and computing power.
Beating the Competition: IF-Beta beat all other methods, including those that required expensive re-training or rigid rules.

The Takeaway

The paper argues that quality of data matters more than quantity. Just like a human can learn a skill faster with a few high-quality mentors and examples rather than reading every book in the library, a machine learning model can learn better if we carefully curate its "diet" of training data.

The authors call this "Distill on a Diet" because they are feeding the student model a smaller, more nutritious, and carefully selected portion of the data, allowing it to grow stronger and faster than if it were force-fed the entire dataset.

Technical Summary: Distill on a Diet: Efficient Knowledge Distillation via Learnable Data Pruning

Problem Statement

Knowledge Distillation (KD) is a standard technique for compressing large teacher models into compact student models for efficient inference. However, the training process itself remains computationally expensive, as students are typically distilled on the entire dataset under the guidance of a large teacher. While data pruning (selecting a subset of training samples) offers a promising avenue to reduce this cost, existing methods face two critical obstacles in the KD context:

Efficiency-Effectiveness Trade-off in Score Estimation: High-fidelity metrics like EL2N, Forgetting, or AUM rely on training dynamics (e.g., per-epoch predictions) that are unavailable when only a pretrained teacher is provided. Computing these scores requires expensive retraining of the teacher, negating the efficiency gains of pruning. Conversely, simpler metrics like Cross-Entropy (CE) loss are computationally cheap but unreliable, as over-parameterized networks can overfit hard or mislabeled samples, resulting in low loss values that do not reflect true sample difficulty.
Rigid and Suboptimal Sampling: Existing pruning strategies often rely on fixed heuristics (e.g., top-k selection, stratified sampling with fixed cutoffs, or sliding windows). These rigid rules fail to adapt to the specific data distributions and teacher-student dynamics of different KD settings, leading to suboptimal subsets.

Methodology: IF-Beta

The authors propose IF-Beta, a framework that unifies Influence Functions (IF) with a learnable probabilistic sampling policy to address the above challenges.

1. Influence Function as a Post-Hoc Estimator

To avoid retraining, the paper leverages Influence Functions to estimate sample importance directly from a pretrained teacher.

Flat Validation Minima (FVM): To stabilize IF estimation in deep networks, the authors fine-tune the pretrained teacher on the validation set using Sharpness-Aware Minimization (SAM). This enforces local flatness in the loss landscape, significantly improving the reliability of influence estimates compared to standard IF methods.
Score Computation: The influence of a training sample $z_{tr}$ on the validation set is approximated using the gradient of the loss with respect to the fine-tuned teacher parameters and the inverse Hessian (approximated via a diagonal Fisher Information matrix for efficiency). This yields a score where lower values indicate more difficult/informative samples.

2. Learnable Beta Policy

Instead of fixed heuristic selection, IF-Beta introduces a Beta Policy parameterized by a Beta distribution.

Adaptive Sampling: Given normalized influence scores mapped to $[0, 1]$ , the selection probability for each sample is modeled by a Beta distribution $Beta(\alpha(\phi), \beta(\phi))$ . The parameters $\alpha$ and $\beta$ are controlled by learnable parameters $\phi$ , allowing the policy to adapt its shape (e.g., focusing on hard samples, easy samples, or a specific window) to the specific KD task.
Bilevel Optimization: The optimization of the policy parameters $\phi$ $ϕ$ is formulated as a bilevel problem:
- Inner Loop: Trains a student model (specifically, a linear classifier on frozen teacher features) on the subset sampled by the current policy, minimizing the KD loss.
- Outer Loop: Updates $\phi$ to maximize the validation performance of the student trained on the sampled subset.
Efficient Proxy: To make the bilevel optimization tractable, the inner loop does not retrain a full neural network. Instead, it optimizes a lightweight linear classifier attached to the frozen teacher features. This "proxy-student" formulation captures KD dynamics with negligible computational cost.

Key Contributions

IF for KD: The paper demonstrates that Influence Functions, when stabilized via Flat Validation Minima, serve as an effective and efficient post-hoc score estimator for KD, eliminating the need for expensive teacher retraining required by trajectory-based metrics.
Learnable Beta Policy: The authors introduce a parameterized Beta distribution for data pruning, which overcomes the rigidity of fixed heuristic rules (like top-k or fixed windows) by learning an adaptive sampling distribution via bilevel optimization.
Efficient Bilevel Optimization: By formulating the inner loop as a linear classifier on frozen features, the framework achieves efficient optimization of the pruning policy without the prohibitive cost of full student retraining during the search phase.

Experimental Results

Extensive experiments on CIFAR-10, CIFAR-100, and ImageNet with various teacher-student architectures (ResNet, WideResNet, ViT) show:

Superior Performance: IF-Beta consistently outperforms baselines (including retraining-free methods like Medium-Difficulty and retraining-based methods like DUAL) across a wide range of pruning ratios.
Surpassing Full-Data Distillation: Remarkably, IF-Beta enables students trained on significantly less data and compute to surpass the performance of students distilled on the full dataset. For example, on CIFAR-10, retaining only 70% of the data achieved 95.69% accuracy compared to 95.50% for full-data distillation.
Compute Efficiency: Even when accounting for the cost of the pruning process, IF-Beta reaches the performance of full-data distillation in significantly less time (e.g., 0.806x the cost on CIFAR-10).
Generalization: The method generalizes to non-KD standard supervised learning settings and various KD objectives (Relational, FitNets, Attention Transfer), demonstrating that the quality of the data subset is a critical factor independent of the specific distillation algorithm.

Significance and Claims

The paper argues that the quality of the training data is often more critical than the specific details of the distillation algorithm. By identifying and selecting the most informative samples efficiently, IF-Beta positions data curation as a principled and practical lever for improving the efficiency and effectiveness of modern training pipelines. The work suggests that "distilling on a diet"—using a carefully pruned subset rather than the full dataset—can yield better generalization and lower computational costs, challenging the assumption that more data always equates to better distillation outcomes.