Balancing Knowledge Distillation for Imbalance Learning with Bilevel Optimization
Abstract
Knowledge distillation transfers knowledge from a high capacity teacher to a compact student using a mixture of hard and soft losses. On imbalanced data, a fixed weighting between hard and soft losses becomes brittle the learning process. Recent studies try to reweight these components in long-tailed settings. However, most of these methods do not adapt weights at the sample-wise level and do not take into account the students behavior during training. To address this, we propose BiKD -- a bilevel framework that dynamically balances hard and soft losses for each sample. We employ a weight generation network that produces adaptive per-sample weights, guided by a small balanced validation set. The student is now trained with an unconstrained combination of weighted hard and soft losses, allowing the student to relax both terms. We further propose a multi-step SGD strategy to optimize the weight model more accurately and efficiently. Experiments on long-tailed CIFAR-10/100 show that our approach surpasses recent balanced distillation methods across imbalance factors.
Cite
@article{arxiv.2605.17839,
title = {Balancing Knowledge Distillation for Imbalance Learning with Bilevel Optimization},
author = {Anh B. H. Nguyen and Ba Tho Phan and Viet Cuong Ta},
journal= {arXiv preprint arXiv:2605.17839},
year = {2026}
}
Comments
Accepted at PAKDD 2026