Deep ensembles are a simple and effective method for improving both predictive performance and epistemic uncertainty estimation in deep learning. However, their high computational cost, especially at inference time, limits their practicality in real-world deployments. Ensemble distillation offers a promising solution by training a single student model to match the ensemble’s predictive distribution. Yet existing approaches typically assume full access to all M teacher predictionsduring training, which is often challenging due to compute constraints, memory limitations, or asynchronous model evaluations. Here we introduce STEDD (Stochastic Teacher-sampling for Ensemble Distribution Distillation), a framework for distilling both the mean and variance of an ensemble using only a small number of random teacher queries per input, even as few as one. STEDD includes three estimators tailored to different access regimes and provides theoretical guarantees for convergence and calibration. Experiments on genomics and vision benchmarks demonstrate that STEDD preserves ensemble-level performance and uncertainty calibration while significantly reducing training-time cost.

Paper In Preparation

Abstract:

Deep neural networks (DNNs) have advanced predictive modeling for regulatory genomics, but challenges remain in ensuring the reliability of their predictions and understanding the key factors behind their decision making. Here we introduce DEGU (Distilling Ensembles for Genomic Uncertainty-aware models), a method that integrates ensemble learning and knowledge distillation to improve the robustness and explainability of DNN predictions. DEGU distills the predictions of an ensemble of DNNs into a single model, capturing both the average of the ensemble’s predictions and the variability across them, with the latter representing epistemic (or model-based) uncertainty. DEGU also includes an optional auxiliary task to estimate aleatoric, or data-based, uncertainty by modeling variability across experimental replicates. By applying DEGU across various functional genomic prediction tasks, we demonstrate that DEGU-trained models inherit the performance benefits of ensembles in a single model, with improved generalization to out-of-distribution sequences and more consistent explanations of cis-regulatory mechanisms through attribution analysis. Moreover,DEGU-trained models provide calibrated uncertainty estimates, with conformal prediction offering coverage guarantees under minimal assumptions. Overall, DEGU paves the way for robust and trustworthy applications of deep learning in genomics research.

Full Text