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

Conference: ECCV 2026
arXiv: 2606.25488
Code: https://github.com/yifanwu-victor/Distill-on-a-Diet
Area: Model Compression
Keywords: Knowledge Distillation, Data Pruning, Influence Function, Beta Distribution, Bi-level Optimization

TL;DR¶

This paper proposes the IF-Beta framework, which uses influence functions (IF) as a retraining-free sample importance scorer. Combined with a learnable Beta distribution sampling strategy, it efficiently searches for the optimal pruning subset via bi-level optimization in the feature space, enabling the student to outperform full-dataset distillation with less data and computation.

Background & Motivation¶

Background: Knowledge distillation (KD) is a mainstream method for model compression. It improves inference efficiency by transferring knowledge from a large teacher network to a small student network on the full training set. However, the computational cost of the distillation process itself is often overlooked—the student must be trained on the entire dataset under the guidance of the teacher, making its training cost even higher than standard supervised training.

Limitations of Prior Work: Data pruning is a potential direction to relieve the KD training burden by keeping only a subset of training samples. However, existing pruning methods face two fundamental issues in KD scenarios. First, there is an "efficiency-effectiveness" dilemma in the scoring stage: high-quality scoring methods based on training dynamics (e.g., EL2N, Forgetting, AUM) require tracking prediction changes across epochs during model training. But in KD, only a pre-trained teacher is available without training trajectories; retraining the teacher to obtain these metrics is prohibitively expensive, completely negating the efficiency gains of pruning. Conversely, fast scoring methods based on cross-entropy loss are computationally lightweight but unreliable—overparameterized networks can easily overfit even mislabeled samples, making loss values insensitive to sample difficulty. Second, manual heuristic strategies during the sampling stage (such as top-k selection, hard-thresholding stratified sampling in CCS, fixed sliding windows in BWS) are too rigid to adaptively adjust the selection region for different data distributions and teacher-student pairs, leading to suboptimal subsets.

Key Challenge: In the KD scenario, the scorer is required to be both high-quality (accurately reflecting a sample's contribution to distillation) and highly efficient (requiring no retraining), which existing methods fail to satisfy simultaneously. Meanwhile, the sampling strategy needs to be adaptive rather than preset in a fixed form.

Goal: (1) Find a high-quality scorer that requires no retraining; (2) Design a learnable, adaptive sampling strategy; (3) Jointly optimize scoring and sampling with acceptable computational cost, making the pruned distillation performance exceed that of full distillation.

Key Insight: The authors observe that the Influence Function (IF) is a natural posterior scoring tool that quantifies the impact of each training sample on the model's predictions, requiring only a pre-trained model for calculation. Although traditional IF was considered impractical due to numerical instability and high computational costs, recent studies show that under the Flat Validation Minimum (FVM) condition, IF estimation becomes both stable and efficient. This perfectly fits the KD scenario: the teacher model naturally exists, and only a lightweight FVM fine-tuning is required on it to obtain reliable IF scores.

Core Idea: Replace traditional training dynamics-based scoring with IF computed under the FVM condition, replace manual heuristic sampling rules with a learnable Beta distribution, and efficiently solve for the optimal pruning strategy via bi-level optimization in the frozen feature space of the teacher.

Method¶

Overall Architecture¶

The goal of IF-Beta is, given a pre-trained teacher and a training set, to find an optimal subset under a specified pruning ratio \(r\) that maximizes the student's performance on the validation set after distillation. The entire method consists of six stages: (1) Finetuning the teacher via F-SAM on the validation set to achieve the Flat Validation Minimum (FVM); (2) Rapidly computing the IF score for each training sample using a diagonal Fisher approximation of the inverse Hessian; (3) Extracting the frozen teacher features to serve as input for subsequent bi-level optimization; (4) Searching for the optimal Beta Policy parameters via bi-level optimization in the feature space, using policy gradient to update the policy in the outer loop, and a linear classifier as a proxy student to rapidly evaluate subset quality in the inner loop; (5) Sampling the final coreset using the searched optimal policy; (6) Performing standard KD training on the coreset to obtain the final student. Stage (4) is the core engine, which avoids full network retraining and outer-loop implicit differentiation through two key efficiency designs: the proxy student and the policy gradient.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Pre-trained Teacher + Train/Val Sets"] --> B["IF-FVM Scoring<br/>F-SAM Finetuning + Diagonal Fisher IF"]
    B --> C["Beta Policy Parameterization<br/>Rank-percentile + Beta Distribution"]
    C --> D["Feature Space Bi-level Optimization<br/>Frozen Feature Linear Proxy + Policy Gradient"]
    D --> E["Coreset Sampling + KD Distillation"]

Key Designs¶

1. IF-FVM Scoring: Replacing Retraining-dependent Trajectory Scoring with Influence Functions

Traditional data pruning relies on training dynamics-based metrics like EL2N, Forgetting, and AUM, which record the prediction changes representing each sample during the model training process (e.g., EL2N records the L2 error between predictions and one-hot labels, and Forgetting counts the number of correct-to-wrong transitions). In the KD scenario, only a pre-trained teacher is available without its training trajectory; retraining the teacher to acquire these metrics completely cancels out the efficiency benefits of pruning. Cross-entropy (CE) loss as an alternative is fast, but the nature of overparameterized networks causes the losses of both hard and easy samples to converge to very small values, resulting in poor discriminative power.

This paper proposes using the Influence Function (IF) as a posterior scorer. The core idea of IF is: if a training sample is removed, how much does the model's loss on the validation set change? A larger change indicates a more important sample. Specifically: first, the pre-trained teacher is lightweight-finetuned (about 1000 iterations) on the validation set using the F-SAM optimizer to make it converge to a flat validation minimum (FVM). Under the FVM condition, the Hessian matrix is well-conditioned, and the influence of a training sample \(z_{\text{tr}}\) on the validation set simplifies to:

\[\mathcal{I}(z_{\text{tr}}, D_{\text{val}}) = \tilde{g}_{z_{\text{tr}}}^{\top} \tilde{H}_{\text{val}}^{-1} \tilde{g}_{z_{\text{tr}}}\]

Where \(\tilde{g}_{z_{\text{tr}}} = \nabla_{\tilde{\theta}_T} \ell(z_{\text{tr}}, \tilde{\theta}_T)\) is the gradient of the sample at the FVM parameters \(\tilde{\theta}_T\), and \(\tilde{H}_{\text{val}}\) is the Hessian on the validation set. To further accelerate, the authors approximate the inverse Hessian with the diagonal Fisher information matrix \(\text{diag}(\frac{1}{|D_{\text{tr}}|}\sum \tilde{g}_{z_{\text{tr}}}\tilde{g}_{z_{\text{tr}}}^{\top})\), making the IF scoring require only two forward-backward passes. Spearman rank correlation experiments in Fig.1 show that IF-FVM has a significantly higher correlation with trajectory metrics like EL2N and AUM than CE loss and traditional LiSSA-approximated IF, proving that IF-FVM accurately captures sample difficulty without retraining.

2. Learnable Beta Policy Sampling: Replacing Heuristic Sampling with Parameterized Beta Distributions

Even with reliable scores, existing sampling strategies still lag behind. CCS uses a hard-thresholding + stratified sampling approach, but experiments (Fig.2a-b) show that the optimal thresholding ratio differs drastically between KD and non-KD scenarios, making pre-configured parameters completely non-transferable. BWS uses a fixed-size sliding window, but the authors find (Fig.2c) that introducing a certain percentage of random substitutions from outside the window actually improves performance, suggesting that the fixed window itself is a suboptimal design.

The authors propose parameterizing the sampling probability via a Beta distribution to break through the rigidity of heuristic strategies. First, the IF score is normalized to \([0,1]\) via rank-to-percentile, obtaining \(\hat{s}_i \in [0,1]\), where smaller values indicate harder samples. Then, a sampling probability in the form of a Beta distribution is defined:

\[p_i(\phi) = \frac{1}{Z} \frac{1}{B(\alpha(\phi), \beta(\phi))} \hat{s}_i^{\alpha(\phi)-1} (1-\hat{s}_i)^{\beta(\phi)-1}\]

Where \(B(\cdot,\cdot)\) is the Beta function, \(Z\) is a normalization constant, and \(\phi\) is the learnable parameter controlling \(\alpha\) and \(\beta\). The two parameters of the Beta distribution grant extreme flexibility to the policy (Fig.3): when \(\alpha > \beta\), the distribution is left-skewed, tending to sample low-score (hard) samples; when \(\alpha < \beta\), the distribution is right-skewed, tending to sample high-score (easy) samples; when \(\alpha = \beta = 1\), it degenerates to a uniform distribution. The distribution shape can freely switch among unimodal, U-shaped, and monotonically increasing/decreasing, naturally covering the search spaces of all heuristic strategies such as CCS and BWS.

Given a pruning ratio \(r\), the Beta Policy \(\pi_{\phi}^r\) samples a binary mask \(\bm{m} \in \{0,1\}^N\) satisfying \(\|\bm{m}\|_0 = (1-r)N\). The sampling probability is obtained by accumulating the product of \(p_i(\phi)\) for each sample in the mask and normalizing, constructing a categorical distribution over all valid subsets. This makes the policy optimizable via gradients, rather than being an undifferentiable hard selection.

3. Feature Space Bi-level Optimization: Efficient Policy Search via Linear Proxy Student and Policy Gradient

Searching for the optimal Beta Policy is formulated as a bi-level optimization problem: the outer loop minimizes the expected validation loss, while the inner loop trains the student on the sampled subset:

\[\min_{\phi} \Phi(\phi) = \mathbb{E}_{\bm{m} \sim \pi_{\phi}^r} \widehat{\mathcal{L}}(\theta_S^*(\bm{m})) \quad \text{s.t.} \quad \theta_S^*(\bm{m}) = \arg\min_{\theta_S} \mathcal{L}(\theta_S; \theta_T, \bm{m})\]

Where the inner loop uses the standard KD loss \(\ell_{\text{KD}} = (1-\alpha)\ell_{\text{CE}} + \alpha\ell_{\text{KL}}\) and the outer loop uses the validation set CE loss. Directly solving this requires fully training the student network for every policy update in the inner loop, which is computationally unacceptable. The authors exploit the special structure of the KD scenario to make two key simplifications:

Proxy Student to Accelerate the Inner Loop: Instead of optimizing the inner loop on the full student network, a linear classification head \(\bm{C} \in \mathbb{R}^{c \times d}\) is appended onto the frozen teacher features to act as a proxy student: \(f_{\theta_S}(x_i) = \bm{C} \cdot f_T(x_i)\). The inner loop becomes \(\bm{C}^*(\bm{m}) = \arg\min_{\bm{C}} \frac{1}{K} \sum_{z_i \in D_{\text{sub}}} \ell_{\text{KD}}(z_i, \bm{C}, \theta_T)\). Since the dimension of the feature space is far smaller than the parameter count and the features are already frozen, the linear classifier converges in just 1-2 epochs, keeping the inner-loop overhead extremely low.

Policy Gradient to Avoid Implicit Differentiation: The outer-loop gradient \(\nabla_{\phi} \Phi(\phi)\) involves \(\nabla_{\phi} \theta_S^*(\bm{m})\) (implicit differentiation), which is extremely expensive to compute in deep networks. Utilizing the probabilistic form of the Beta Policy, the authors directly derive an unbiased policy gradient estimator (with Theorem 3.1 proving its unbiasedness):

\[\hat{g}(\bm{m}) = \widehat{\mathcal{L}}(\theta_S^*(\bm{m})) \nabla_{\phi} \ln p(\bm{m}|\phi, r)\]

Intuitively: if a sampled subset allows the student to perform well on the validation set (low loss), the policy parameters that produced this subset are reinforced; otherwise, they are suppressed. The parameter update rule is \(\phi \leftarrow \phi - \eta \widehat{\mathcal{L}}(\theta_S^*(\bm{m})) \nabla_{\phi} \ln p(\bm{m}|\phi, r)\). This completely bypasses implicit differentiation, making outer-loop optimization feasible.

Reparameterization and Two-Stage Decoupling: Direct optimization of \((\alpha, \beta)\) is unstable in practice. The authors reparameterize them as \((\mu, \tau)\): \(\alpha = \tau\mu + 1, \beta = \tau(1-\mu) + 1\), where \(\mu \in [0,1]\) controls the mode position of the Beta distribution (preference for hard or easy samples) and \(\tau > 0\) controls the sharpness. They then decouple the joint optimization into two stages: in the first stage, \(\tau\) is fixed, and the optimal \(\mu\) is exhaustively searched on a discrete grid \(\{0, t, 2t, \dots, 1\}\); in the second stage, \(\mu^*\) is fixed, and \(\tau\) is optimized continuously using SGD. This two-stage process reduces the original 2D joint optimization problem into two simpler 1D sub-problems, keeping the computational overhead perfectly manageable.

Loss & Training¶

The search phase of IF-Beta introduces no new loss functions: the outer-loop objective is naturally the validation set CE loss, and the inner loop is the standard KD loss \(\ell_{\text{KD}} = (1-\alpha)\ell_{\text{CE}} + \alpha\ell_{\text{KL}}\). The FVM fine-tuning phase uses the F-SAM optimizer to fine-tune the validation set for 1000 iterations with a learning rate of 0.01 and a batch size of 128. The bi-level optimization runs for 20 epochs in the outer loop, with a learning rate of 0.1 cos-decayed to 0.01, and the inner-loop linear classifier is trained for 1 epoch with a learning rate of \(10^{-3}\). The final student is trained on the coreset using standard KD loss (\(\alpha=0.5\)). The total pruning overhead (FVM fine-tuning + IF calculation + bi-level optimization) is equivalent to no more than 4 full epochs of training.

Key Experimental Results¶

Main Results¶

The following table shows the KD pruning results for isomorphic teacher-student pairs (where both teacher and student are ResNet-18 for CIFAR, and ResNet-50 for ImageNet), selecting three representative pruning ratios.

Dataset	Pruning Ratio	Full KD	Random	Medium-Diff*	EL2N	IF-Beta
CIFAR-10	30%	95.50	94.89	94.02	95.53	95.69
CIFAR-10	50%	95.50	94.15	92.99	95.38	95.38
CIFAR-10	70%	95.50	92.28	91.05	93.27	93.90
CIFAR-100	10%	79.38	78.76	79.03	79.41	79.53
CIFAR-100	50%	79.38	75.42	75.76	74.18	79.10
CIFAR-100	70%	79.38	72.31	71.91	59.67	75.60
ImageNet	30%	73.54	73.32	74.01	74.17	74.18
ImageNet	50%	73.54	72.66	73.48	72.32	74.02
ImageNet	70%	73.54	71.82	72.79	65.97	73.48

IF-Beta outperforms all methods across all datasets and pruning ratios. Most notably, on CIFAR-10, keeping 70% of the data yields 95.69% (exceeding the full KD baseline of 95.50%); on CIFAR-100, keeping 90% of the data yields 79.53% (exceeding full KD of 79.38%); on ImageNet, keeping 50% of the data yields 74.02% (exceeding full KD of 73.54%). This demonstrates that some data is indeed harmful to distillation, and meticulously pruning them actually improves generalization performance.

Ablation Study¶

To isolate the impact of the sampling strategy, the learnable Beta Policy of IF-Beta is replaced by IF-CCS (hard-thresholding + stratified sampling) and IF-BWS (fixed window sampling), where all three use the same IF-FVM score.

Dataset	Pruning Ratio	IF-CCS	IF-BWS	IF-Beta
CIFAR-10	90%	88.18	86.20	88.51
CIFAR-10	70%	93.39	89.80	93.90
CIFAR-10	30%	94.65	95.51	95.69
CIFAR-100	90%	63.73	58.71	64.11
CIFAR-100	70%	72.22	71.75	75.60
CIFAR-100	30%	78.47	78.85	79.34
ImageNet	90%	65.73	67.90	68.62
ImageNet	50%	72.98	73.63	74.02
ImageNet	30%	73.32	73.66	74.18

The learnable Beta Policy outperforms the fixed-threshold CCS and fixed-window BWS under all settings. The advantage is even more pronounced at high pruning ratios (e.g., CIFAR-100 with 70% pruning: 75.60 vs 72.22 vs 71.75), indicating that manual heuristic strategies fail almost completely under extreme pruning, whereas the Beta Policy maintains a reasonable sample selection by adaptively adjusting the distribution shape.

Key Findings¶

Adaptability of Beta Policy is the Core Gain: In the ablation study, the improvement of IF-Beta over IF-CCS/IF-BWS comes entirely from the sampling strategy, as they use the same scoring. This directly proves that manual heuristic sampling is the bottleneck restricting pruning effectiveness.
Data Quality Over Distillation Algorithms: Under three different distillation paradigms—Relational KD, FitNets, and Attention Transfer—IF-Beta consistently outperforms Random and Medium-Difficulty across 30%/50%/70% pruning ratios. This indicates that selecting a high-quality subset is a general benefit independent of the type of distillation objective.
Substantial Computational Efficiency Gains: Distilling a MobileNet from a ResNet-34 on ImageNet, IF-Beta takes only 19.24 hours (including all pruning overhead) to reach 71.20% accuracy with 70% data, which is less than half the time of the full dataset baseline (39.86h, 69.57%) while achieving the highest accuracy.
General Pre-trained Models Can Support Non-KD Pruning: In the standard data pruning (non-KD) scenario, directly using an ImageNet pre-trained ResNet-18 to compute IF (without any in-domain retraining) results in only a slight performance drop for IF-Beta (CIFAR-100 90% pruning: 55.78 vs 53.31). This shows that the IF-FVM score is insensitive to the source model, showing strong practicality.

Highlights & Insights¶

The "Match Made in Heaven" of IF in the KD Scenario: IF is rediscovered as an ideal scorer for KD data pruning. KD naturally provides a pre-trained teacher, and only a lightweight FVM fine-tuning is required to obtain stable IF, completely removing the reliance of traditional data pruning on expensive training trajectories. This design converts a "predicament" (no training trajectory) into an "advantage" (already having a teacher, just needing fine-tuning).
Modeling Ingenuity of the Beta Distribution: The two parameters of the Beta distribution correspond perfectly to the "hard sample preference" (high \(\alpha\)) and "easy sample preference" (high \(\beta\)). Its shape can smoothly transition between unimodal, U-shaped, and monotonic, comprehensively covering the search space of all heuristic strategies. The \((\mu, \tau)\) reparameterization decouples the mode position and sharpness, making the search process controllable and interpretable.
The Efficiency Formula of Proxy Student + Policy Gradient: Replacing the full student network with a linear classification head on frozen features reduces the inner loop training from "training a CNN for dozens of epochs" to "training a linear layer for 1 epoch," making the outer-loop policy gradient practically executable. This two-stage logic of "cheap proxy approximation + real evaluation as reward signals" is transferable to any task requiring repeated evaluation of subset quality (e.g., active learning, curriculum learning, data denoising).
"Which Samples to Distill" Can Be More Important Than "How to Distill": The most thought-provoking finding of this paper is that the gains of meticulously selecting data are consistent and significant across different distillation losses, student-teacher architectures, and training budgets. This suggests that in KD efficiency research, data curation is a severely underrated fundamental dimension.

Limitations & Future Work¶

Scoring Dependency on Teacher Quality: The IF-FVM scoring assumes that the teacher lies at a flat validation minimum. If the teacher itself is underfitted or there is a significant domain shift between the validation and training sets, the reliability of the IF scores will degrade. While successfully verified under standard CIFAR/ImageNet settings, its robustness under domain shift scenarios remains unexplored.
Fidelity Boundaries of the Linear Proxy Student: The inner loop approximates full KD training using a linear classifier on frozen teacher features. When the teacher and student architectures differ vastly (e.g., a CNN teacher distilling a ViT student), the teacher's features may not fully represent the representation space the student needs to learn, which might compromise the proxy approximation accuracy. Currently, cross-architecture experiments only cover ResNet \(\rightarrow\) MobileNet, and do not test ViT students.
Manually Configured Pruning Ratio: The current framework requires a predefined pruning ratio \(r\) instead of automatically learning the optimal ratio. In real deployment, users might not know the exact "optimal amount to prune," making the learning of \(r\) a natural extension.
Unimodal Assumption of the Beta Distribution: Although the dual parameters of the Beta distribution offer flexibility, it is still constrained to unimodal or monotonic distribution families on \([0,1]\). If the optimal sampling distribution requires multi-modality (e.g., keeping both the easiest and hardest samples), the modeling capability of a single Beta distribution might be insufficient, warranting exploration of richer distribution families.
Under-explored Synergy with Advanced KD Methods: While IF-Beta's compatibility with logit-based KD, Relational KD, and FitNets has been verified, its synergy with more modern distillation paradigms like decoupled KD, multi-teacher KD, and online KD has not yet been investigated.

vs Medium-Difficulty (Chen et al., ICLR 2025): Both target efficient data pruning for KD without retraining. Medium-Difficulty uses the teacher's CE loss as a difficulty metric and manually defines a "medium-difficulty" window on the sorted loss distribution. In contrast, IF-Beta uses IF-FVM as a more reliable scorer and a learnable Beta distribution to adaptively determine the sampling region. IF-Beta consistently outperforms Medium-Difficulty across all pruning ratios, validating the dual advantage of "better scoring + learnable sampling."
vs CCS / BWS: Both are score-based sampling strategies but utilize manual heuristic rules. The key contribution of this paper is proving that even when replacing the underlying scores with the superior IF-FVM (resulting in IF-CCS/IF-BWS), rigid sampling designs still bottleneck performance, which is further bypassed by the learnable Beta Policy. This provides a major insight for data pruning: the scorer and sampling strategy must be optimized jointly, rather than optimizing the scorer alone.
vs DUAL (Xie et al.): DUAL is a recent SOTA for data pruning that combines difficulty and uncertainty but still requires 60 epochs of partial model retraining to extract training dynamics. IF-Beta outperforms DUAL without any retraining, demonstrating that the posterior scoring of IF-FVM matches the quality of expensive training dynamics.
vs Traditional IF Methods: Traditional IF is practically unusable on large models due to Hessian inversion and numerical instability. This paper benefits from two recent developments that make IF practical: the FVM condition (F-SAM fine-tuning yields a well-conditioned Hessian) and diagonal Fisher approximation (extremely low computational cost). This formula provides a replicable engineering template for adopting IF in other data-importance evaluation scenarios like mislabel detection, data denoising, and active learning.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ Combining IF with KD data pruning is highly novel, and the learnable sampling design of the Beta Policy is elegant yet powerful; the bi-level optimization + proxy student efficiency scheme is clever and practical.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ The experiments cover CIFAR-10/100 and ImageNet, multiple pruning ratios (10%-90%), various teacher-student architectures (isomorphic and heterologous), and diverse KD losses. Includes comprehensive ablation and efficiency analysis, along with validation on standard non-KD scenarios.
Writing Quality: ⭐⭐⭐⭐ The structure is clear, with a complete motivation chain (from pain point \(\rightarrow\) core conflict \(\rightarrow\) proposed method), and the methodology balances mathematics with intuition. Key charts (such as IF correlation validation and analysis of heuristic limitations) strongly support the arguments.
Value: ⭐⭐⭐⭐⭐ Opens up a new dimension of "data curation" for KD efficiency research. The formula of IF-FVM scoring + Beta Policy is highly transferable to active learning, curriculum learning, etc. The experiments convincingly demonstrate that "data quality can be more important than algorithm design."