Knowledge Distillation of Uncertainty using Deep Latent Factor Model

TL;DR

This work tackles the practical limitation of deep ensembles by introducing Gaussian distillation, a distribution distillation approach that compresses a teacher ensemble into a single Gaussian-process-based distribution. The Deep Latent Factor (DLF) model ties the mean and covariance to a learnable neural-network body and a low-dimensional latent factor, enabling stable EM-based estimation and efficient sampling of ensemble members. Across regression, classification, language-model fine-tuning, and distribution-shift scenarios, Gaussian distillation yields superior uncertainty quantification and calibration while reducing inference cost and memory. The method offers a principled, scalable path to on-device uncertainty-aware AI with strong empirical gains over prior approaches such as Hydra and Dirichlet distillation.

Abstract

Deep ensembles deliver state-of-the-art, reliable uncertainty quantification, but their heavy computational and memory requirements hinder their practical deployments to real applications such as on-device AI. Knowledge distillation compresses an ensemble into small student models, but existing techniques struggle to preserve uncertainty partly because reducing the size of DNNs typically results in variation reduction. To resolve this limitation, we introduce a new method of distribution distillation (i.e. compressing a teacher ensemble into a student distribution instead of a student ensemble) called Gaussian distillation, which estimates the distribution of a teacher ensemble through a special Gaussian process called the deep latent factor model (DLF) by treating each member of the teacher ensemble as a realization of a certain stochastic process. The mean and covariance functions in the DLF model are estimated stably by using the expectation-maximization (EM) algorithm. By using multiple benchmark datasets, we demonstrate that the proposed Gaussian distillation outperforms existing baselines. In addition, we illustrate that Gaussian distillation works well for fine-tuning of language models and distribution shift problems.

Figures (8)

• Figure 1: Overall process of Gaussian distillation
• Figure 2: Comparison of performances of the three learning algorithms (DNN, Hydra and DLF) pretrained on CIFAR-10 and fine-tuned on CIFAR10-Flip as the sample size of CIFAR 10-Flip data varies. The solid curves are the means and the shaded bands are the min-max spreads obtained from 5 training models on 5 randomly selected new data of CIFAR10-Flip.
• Figure 3: Box plot of evaluation metrics (RMSE, NLL, CRPS) for different strategies of the design point selection
• Figure 4: Evaluation metrics (RMSE, NLL, CRPS) versus the design ratio.
• Figure 5: Evaluation metrics (RMSE, NLL, CRPS) versus latent factor dimension.

Theorems & Definitions (10)

• Lemma G.1
• Theorem G.2
• Theorem G.3: Upper bound at a new input
• Lemma G.4
• proof
• Lemma G.5
• proof
• proof
• proof
• proof