Teacher-Access-Constrained Ensemble Distribution Distillation

Kaeli Rizzo, Yiyang Yu, Peter K. Koo

Abstract:

Deep ensembles improve predictions while disagreement among their members provides a useful signal of model uncertainty. Conventional ensemble distillation can compress the ensemble mean into a single student, but preserving disagreement requires distribution distillation of the ensemble variance. Standard moment matching constructs this target from every teacher prediction for each training input. Full access can be impractical for large teacher pools, large models that cannot reside together in GPU memory, or teachers distributed across devices or sites. Precomputing targets is also insufficient for training pipelines that generate new inputs online, since the ensemble mean and variance must be evaluated on the augmented inputs themselves. We study teacher-access-constrained ensemble distribution distillation, where only one or a few teachers can be queried per input per update. Random teacher draws provide unbiased targets for the ensemble mean, but variance recovery depends on which teacher pairs are observable. For fixed-size access policies, the ordinary subset variance is unbiased when all teacher pairs have equal inclusion probabilities. A fixed teacher subset instead changes the target from full-ensemble disagreement to subensemble disagreement rather than merely adding sampling noise. STEDD (Stochastic Teacher-sampling for Ensemble Distribution Distillation) provides training constructions for this regime: sampled moments when at least two teachers can evaluate an input, accumulated moments for recurring single-teacher inputs, and a sample-wise Gaussian likelihood when only individual teacher observations are available. On a genomic sequence-to-activity benchmark, the tested STEDD-W configuration approaches full-access predictive performance and variance fidelity while querying only three of 50 teachers per input per update. With augmentation, correlations between student and ensemble standard deviations exceed 0.90 on both outputs. The two single-query procedures also recover input-dependent variance, exceeding 0.84 correlation on both outputs in their augmented configurations. At the same three-query budget, resampling substantially improves variance fidelity over a fixed teacher subset, and teacher-labeled augmentation improves fidelity at every tested query count. On CIFAR-100, three-query STEDD-W preserves full-access classification performance. These results identify teacher access as a central design axis for variance-preserving ensemble distillation.

Paper In Preparation