Ensemble Distribution Distillation (EDD)

• Ensemble Distribution Distillation (EDD) is a framework that trains a single student model to reproduce the full distribution of ensemble predictions, capturing both accuracy and uncertainty.
• EDD employs diverse student distributions such as Dirichlet, logistic-normal, and Gaussian to model predictive means and dispersions for tasks in classification, regression, and sequence modeling.
• Extensions of EDD to federated learning, self-distillation, and function-space approaches demonstrate improved computational efficiency and reliability in out-of-distribution detection.

Ensemble Distribution Distillation (EDD), originally introduced as EnD$^2$, is a family of distillation methods in which a single student model is trained to reproduce not merely the mean prediction of an ensemble, but a distribution over the ensemble’s predictions. In the canonical classification setting, the student outputs parameters of a Dirichlet distribution over class-probability vectors; later variants move to logistic-normal, Gaussian, Laplace, or function-space distributions over logits or predictive parameters. The defining aim is to retain the ensemble’s predictive accuracy together with its uncertainty structure—especially the separation between aleatoric and epistemic uncertainty—while replacing the ensemble’s $M$ forward passes by single-model inference or by inexpensive sampling from a learned student distribution (Malinin et al., 2019, Lindqvist et al., 2020, Fathullah et al., 2023).

1. Historical emergence and conceptual scope

EDD arose from the observation that ordinary ensemble distillation and standard knowledge distillation preserve the ensemble’s average soft target but discard its diversity. The original EnD$^2$ formulation treated the per-model predictive distributions of a deep ensemble as samples from a Dirichlet distribution and trained a single Prior Network student to fit that distribution. This allowed the student to preserve both the ensemble mean and its dispersion, so that uncertainty decomposition remained available after compression (Malinin et al., 2019).

Subsequent work generalized the idea beyond the original Dirichlet-prior-network setting. A general framework reformulated EDD as learning a distribution over ensemble predictive parameters $z$, applicable to both regression and classification, with Gaussian or logistic-normal students as natural choices (Lindqvist et al., 2020). Later work addressed failure modes of the original maximum-likelihood Dirichlet objective when the number of classes is very large, introduced logit-based formulations for machine translation and other autoregressive sequence tasks, and extended EDD to federated learning, self-distillation, function-space distillation, Gaussian latent-factor models, and self-supervised human activity recognition (Ryabinin et al., 2021, Fathullah et al., 2023).

The resulting literature uses several closely related names—EnD$^2$, EnDD, EDD, distribution distillation, and predictive distribution distillation—but the common principle is consistent: the student is trained to model the ensemble’s output distribution rather than only its averaged prediction.

2. Probabilistic formulation and uncertainty decomposition

For an ensemble of $M$ teachers, the ensemble predictive distribution is approximated by averaging member predictions. In the sequence case this is written per timestep $t$ as

$$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$

In classification, total predictive uncertainty is the entropy of this averaged predictive distribution, while the aleatoric component is the mean per-model entropy and the epistemic component is the residual mutual information:

$$H[y_t\mid x,\mathcal D] = \mathbb E_{\theta}\!\left[H[y_t\mid x,\theta]\right] + I[y_t,\theta\mid x,\mathcal D].$$

Equivalently,

$$I[y_t,\theta\mid x,\mathcal D] = H[y_t\mid x,\mathcal D] - \mathbb E_{\theta}\!\left[H[y_t\mid x,\theta]\right].$$

This is the central decomposition EDD attempts to preserve after compressing the ensemble (Malinin et al., 2019, Fathullah et al., 2023).

In the original Dirichlet-based formulation, the student predicts concentration parameters $M$0 over the simplex. The mean categorical probabilities are

$M$1

The concentration $M$2 acts as evidence: large $M$3 corresponds to a sharp distribution over class-probability vectors and thus low epistemic uncertainty, while small $M$4 corresponds to diffuse distributions and high epistemic uncertainty (Malinin et al., 2019).

The same decomposition extends to regression through variance rather than entropy. For an ensemble of predictive distributions,

$M$5

Here the first term is aleatoric and the second epistemic. The general EDD framework preserves this by fitting a student distribution over the ensemble’s predictive parameters rather than distilling only a single Gaussian with total variance (Lindqvist et al., 2020).

A central conceptual distinction therefore separates EDD from standard knowledge distillation. KD matches the mean teacher prediction; EDD matches the distribution of teacher predictions. This is why KD can preserve accuracy yet erase epistemic uncertainty, whereas EDD can reconstruct entropy- and mutual-information-based uncertainty measures from a single student (Fathullah et al., 2023).

3. Student distributions and training objectives

EDD is not tied to a single student family. What remains fixed is the target: the empirical distribution of ensemble outputs for a given input.

Student family	Support	Representative use
Dirichlet	Probability simplex	Original EnD$M$6, proxy-target EDD, HAR
Logistic-normal / Gaussian	Logits or predictive parameters	General classification/regression EDD
Diagonal Laplace over logits	Normalized token logits	Large-vocabulary translation
Function-space or GP-like students	Joint outputs across inputs	Functional and Gaussian distillation

The canonical EnD$M$7 objective assumes that ensemble member probabilities $M$8 are i.i.d. samples from a Dirichlet distribution and fits the student by maximum likelihood:

$M$9

This directly fits both the ensemble mean and dispersion. Practical parameterizations either predict each $^2$0 directly via a positive transform such as softplus, or separate the mean $^2$1 from a scalar evidence term and set $^2$2 (Malinin et al., 2019).

A broader EDD formulation replaces the Dirichlet by a learned distribution over ensemble predictive parameters,

$^2$3

and trains the student by maximizing the likelihood of teacher outputs $^2$4:

$^2$5

This generalization admits Gaussian students for regression and logistic-normal students for classification, and it preserves the aleatoric/epistemic split because the student explicitly models variability across ensemble predictive parameters (Lindqvist et al., 2020).

For many-class problems, however, the original Dirichlet maximum-likelihood loss becomes poorly behaved. The core pathology identified in the scaling study is that classes with very low probability induce larger gradients than high-probability classes, forcing optimization to focus on tail-class probabilities. The proposed remedy is a Proxy-Dirichlet target constructed from ensemble mean and variance and trained with reverse KL,

$^2$6

often with evidence clipping, temperature smoothing, and top-$^2$7 truncation for stability (Ryabinin et al., 2021).

This line of work also clarified a modeling limitation of the original approach. A single Dirichlet is unimodal on the simplex. When ensemble predictions are genuinely multimodal, a single Dirichlet can only be a compact approximation rather than an exact representation. Later work therefore explored logistic-normal, Gaussian, Laplace, and function-space alternatives rather than treating the Dirichlet as universally sufficient (Malinin et al., 2019, Lindqvist et al., 2020).

4. Autoregressive sequence modeling and large vocabularies

Sequence-to-sequence applications introduced a distinct difficulty: EDD must be performed at each decoder step, conditioned on $^2$8, over very large vocabularies. In grammatical error correction, this led to token-level EnDD students that output Dirichlet parameters $^2$9 for each step and are trained either by direct Dirichlet negative log-likelihood over teacher probabilities or by a two-stage objective that first fits a per-token teacher Dirichlet and then minimizes Dirichlet KL. The study also reported that temperature annealing from $z$0 was important for stability, while $z$1 was unstable. Because teacher forcing at training and free-run decoding at test time create a mismatch, a Guided Uncertainty Approach (GUA) was introduced: an EnD student performs decoding, while an EnDD student computes uncertainties on the decoded prefix (Fathullah et al., 2020).

Large-vocabulary machine translation motivated a more radical change. Probability-space EDD was found hard to scale because Dirichlet-based optimization becomes problematic when vocabularies are large and distributions are sparse and concentrated. The logit-based translation study therefore modeled the ensemble logits directly, using diagonal Laplace or diagonal Gaussian students over normalized token logits. This retained $z$2 per-timestep complexity, avoided full covariance $z$3 parameterization, decoded deterministically via $z$4, and estimated uncertainty by sampling logits from the student distribution and pushing them through softmax (Fathullah et al., 2023).

The empirical results of this logit-space formulation were unusually strong. On WMT’16 En-De with a base Transformer, standard training yielded $z$5 BLEU, the deep ensemble $z$6, KD from the snapshot ensemble $z$7, and Laplace L-EDD $z$8. On the same setting, out-of-distribution detection improved sharply: for Khresmoi, deep-ensemble total uncertainty (TU) and knowledge uncertainty (KU) AUROC were $z$9 and $^2$0, whereas Laplace L-EDD reached $^2$1 and $^2$2; on WMT’20 En-Ru with a big Transformer, L-EDD also improved Khresmoi TU/KU from $^2$3 for the deep ensemble to $^2$4 (Fathullah et al., 2023).

These sequence studies established two distinct lessons. First, EDD can be made compatible with autoregressive factorization by distilling per-step predictive distributions. Second, large-vocabulary settings often favor logit-space students over probability-space Dirichlet students, especially when robustness and OOD scoring are central objectives.

5. Extensions across learning settings

EDD has been adapted to distributed, self-distilled, function-space, and Gaussian-process-like settings. In federated learning, FedDF instantiates distribution-level distillation by treating client models as teachers, averaging their logits on unlabeled server-side data, and training a server student with KL divergence to the aggregated predictive distribution. Because aggregation occurs through predictions rather than parameter alignment, heterogeneous client architectures, sizes, and precisions are supported. Empirically, on CIFAR-10 with ResNet-8 under moderate non-IID conditions ($^2$5, $^2$6), FedDF reached $^2$7 accuracy in $^2$8 rounds with $^2$9 local epochs, versus $M$0 for FedAvg; under strong non-IID conditions ($M$1), it reached $M$2 in $M$3 rounds versus $M$4 for FedAvg (Lin et al., 2020).

Self-Distribution Distillation (S2D) internalized EDD within a single network. A stochastic teacher branch generates multiple predictions per input using multiplicative Gaussian noise, while a student branch outputs a Dirichlet distribution trained by KL to a proxy Dirichlet fitted to those stochastic predictions. Hierarchical Distribution Distillation (H2D) then distills ensembles of such S2D models either into a single Dirichlet student (H2D-Dir) or into a Gaussian over log-concentrations (H2D-Gauss). On CIFAR-100, H2D-Gauss reached accuracy/NLL/ECE of $M$5, compared with $M$6 for the EnD baseline, and single-pass H2D students often rivaled or exceeded MC/SWAG ensembles on OOD detection (Fathullah et al., 2022).

Functional Ensemble Distillation (FED) criticized a structural limitation of classical EDD: per-input distributions over predictions do not provide cross-input covariance. FED therefore introduced a student generator with shared noise $M$7 and matched the ensemble in function space using MMD over batchwise function representations. In its CIFAR-100 comparison, classical EnDD achieved accuracy $M$8, agreement $M$9, and ECE $t$0, while FED reached $t$1, agreement $t$2, and ECE $t$3. The same study also argued that distilling only on training data can collapse diversity when ensemble members are nearly identical on the training set, and found mixup-based auxiliary data crucial in this regime (Penso et al., 2022).

Gaussian distillation via the Deep Latent Factor (DLF) model moved EDD toward Gaussian-process-like students over functions or logits. Each teacher is treated as a realization of a latent-factor Gaussian process, and the student learns mean and covariance structure by an EM algorithm. This formulation applies to both regression and classification. On CIFAR-10, the reported DLF student achieved $t$4 accuracy, NLL $t$5, and ECE $t$6; for OOD detection using predictive mutual information, AUROC reached $t$7 on SVHN, $t$8 on CIFAR-100, and $t$9 on Tiny ImageNet (Park et al., 22 Oct 2025).

Together these variants show that EDD is not a single architecture but a methodological principle. The student may be a Prior Network, a logistic-normal model, a Laplace logit model, a function generator with shared noise, or a latent-factor Gaussian model, provided that it is trained to match ensemble diversity rather than only ensemble means.

6. Empirical profile, limitations, and recurring misconceptions

Across image-classification benchmarks, the original EnD$$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$0 study showed that a single EDD student could approach ensemble classification performance while outperforming standard DNNs and standard ensemble distillation on misclassification and out-of-distribution detection on CIFAR-10, CIFAR-100, and TinyImageNet (Malinin et al., 2019). Later domain-specific work reported similar benefits in self-supervised human activity recognition. On HHAR, a single model achieved $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$1 clean accuracy and $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$2 under FGSM $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$3, while EDD reached $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$4 clean and $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$5 adversarial accuracy; on the same dataset, uncertainty-as-discriminator AUC-ROC improved from $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$6 to $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$7 on clean data and from $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$8 to $$p(y_t \mid x,\mathcal D)\approx \frac{1}{M}\sum_{m=1}^M p_m(y_t\mid x).$$9 under attack (Nolan et al., 10 Sep 2025).

The main limitations are recurrent and technically specific. Dirichlet maximum-likelihood training scales poorly to very large label spaces because gradient magnitude is dominated by tail probabilities (Ryabinin et al., 2021). A single Dirichlet is unimodal and may misrepresent multimodal ensembles (Malinin et al., 2019). Per-input EDD formulations do not recover cross-input covariance, which motivated FED’s function-space alternative (Penso et al., 2022). In seq2seq tasks, teacher-forcing/free-run mismatch makes EnDD harder to optimize than EnD, and the GEC study found that EnDD alone lagged EnD on GLEU, whereas GUA recovered EnD accuracy while preserving ensemble-like uncertainty (Fathullah et al., 2020). In large-vocabulary translation, diagonal logit distributions ignore vocabulary correlations, decoding uses the approximation $$H[y_t\mid x,\mathcal D] = \mathbb E_{\theta}\!\left[H[y_t\mid x,\theta]\right] + I[y_t,\theta\mid x,\mathcal D].$$0, and consistent logit normalization is required because softmax is shift-invariant (Fathullah et al., 2023).

Several misconceptions recur in the literature. EDD is not equivalent to ordinary KD: KD matches the ensemble mean and therefore cannot represent epistemic uncertainty; EDD learns a distribution over teacher outputs and can reconstruct mutual-information-based uncertainty (Fathullah et al., 2023). Nor does “single-pass uncertainty” always mean fully deterministic uncertainty evaluation: logit-based and Gaussian EDD variants often decode from a deterministic mean but estimate entropy or mutual information by Monte Carlo sampling from the student distribution (Fathullah et al., 2023, Park et al., 22 Oct 2025). Finally, EDD is not guaranteed to dominate ensembles in every regime. Some studies reported that its performance depends strongly on teacher diversity, auxiliary or distillation data, augmentation coverage, and the chosen student family; overconfident or insufficiently diverse teachers can collapse the very signal EDD is meant to preserve (Fathullah et al., 2022, Penso et al., 2022).

In this sense, EDD is best understood as a distribution-matching framework for ensemble compression rather than a single algorithm. Its distinctive contribution is to turn ensemble diversity into an explicit student output object—Dirichlet, logistic-normal, Laplace, Gaussian, or functional—so that uncertainty decomposition survives compression and can be deployed under single-model or near-single-model computational budgets.