REDistill: Robust Estimator Distillation for Balancing Robustness and Efficiency

Technical Summary: REDistill – Robust Estimator Distillation for Balancing Robustness and Efficiency

Problem Statement
Knowledge Distillation (KD) transfers predictive knowledge from a large, high-capacity teacher model to a smaller, more efficient student network by aligning their output distributions, typically via Kullback–Leibler (KL) divergence. However, conventional KD formulations assume that teacher predictions (soft targets) are reliable. In practice, teacher models are often noisy, overconfident, or imperfect. Blindly mimicking these unreliable outputs leads to error accumulation and significant degradation in student model accuracy. Existing approaches to mitigate this issue rely on ad-hoc heuristics—such as swapping logits, augmenting ground-truth probabilities, or masking specific classes—which distort underlying class correlations. Furthermore, these correction-based methods are highly sensitive to hyper-parameter tuning, requiring empirical optimization for every specific teacher-student pair and dataset, thereby hindering generalization.

Methodology
The authors propose REDistill (Robust Estimator Distillation), a framework grounded in robust statistics that replaces the standard KL divergence objective with a power divergence loss.

• Theoretical Foundation: The method addresses the sensitivity of Maximum Likelihood Estimation (MLE) to outliers by modifying the logarithmic transformation used in divergence calculation. Instead of the natural logarithm, REDistill employs a $\gamma$-logarithm (a Box–Cox transform), leading to the power divergence family $D_\lambda$.
• Loss Formulation: The standard KD loss, which combines a ground-truth term and a teacher-matching term, is reformulated. The teacher-matching component (both target and non-target classes) is replaced with the power divergence of order $\lambda = 2/3$. The final loss function is:
  $$L_{\text{REDistill}} = \text{KL}(y, q_\theta) + \alpha D_{2/3}(p_{\text{target}}, q_{\theta,\text{target}}) + \beta D_{2/3}(p_{\text{non-target}}, q_{\theta,\text{non-target}})$$
• Robustness Mechanism: By analyzing the influence function of the estimator, the authors demonstrate that larger positive values of $\lambda$ reduce the influence of individual samples (outliers) on the gradient. The choice of $\lambda = 2/3$ is derived theoretically as a near-optimal balance between statistical efficiency (accuracy under correct specification) and robustness (resistance to noise).
• Implementation: The method requires only logits, integrates seamlessly into existing KD pipelines, and incurs negligible computational overhead. It also includes a scaling factor ($\tau^2$) to maintain gradient magnitude when temperature scaling is applied.

Key Contributions

1. Principled Framework: The paper introduces the first unified, statistically grounded framework for robust knowledge distillation, moving away from heuristic corrections toward a theoretically derived loss function.
2. Adaptive Down-weighting: The power divergence loss adaptively down-weights unreliable teacher outputs while preserving informative logit relationships, effectively handling teacher noise without manual intervention.
3. Model-Agnostic Generalization: Unlike prior methods that require extensive hyper-parameter tuning for each specific pair, REDistill achieves state-of-the-art performance using a single, theoretically derived hyper-parameter ($\lambda = 2/3$) across diverse architectures and datasets.
4. Orthogonality: The method is shown to be compatible with and complementary to existing distillation strategies (e.g., DKD, MLKD), providing consistent improvements when combined with them.

Experimental Results
Extensive experiments were conducted on CIFAR-100 and ImageNet-1k using various teacher-student architectures (ResNet, WideResNet, VGG, ShuffleNet, MobileNet).

• Model-Agnostic Protocol: In a setting where hyper-parameters are fixed across all models (no per-pair tuning), REDistill consistently outperformed previous logit-based methods (including KD, DKD, LSKD, and RLD) across 14 different teacher-student pairs.
• Model-Specific Protocol: Even when compared against methods optimized with extensive hyper-parameter search, REDistill achieved superior or comparable results, demonstrating that its gains are not merely artifacts of tuning.
• ImageNet Performance: On ImageNet-1k, REDistill achieved state-of-the-art accuracy with a substantial margin over existing baselines.
• Ablation Studies: Experiments varying $\lambda$ confirmed that values around $2/3$ yield optimal performance, with performance degrading as $\lambda$ increases beyond 1.

Significance and Claims
The paper claims that REDistill offers a fundamental shift in how knowledge distillation handles teacher uncertainty. By replacing the standard KL divergence with a robust power divergence, the method provides a unified and interpretable treatment of teacher noise. The authors emphasize that REDistill achieves these gains without "bells and whistles"—no model-specific tuning, no architectural changes, and no additional computational cost. The significance lies in its ability to produce robust student models that generalize well to unseen teacher-student pairs, addressing a critical limitation of current KD practices where performance is often brittle and dependent on extensive empirical tuning.