Ensemble Distribution Distillation via Flow Matching
Authors: Jonggeon Park
Giung Nam
Hyunsu Kim
Jongmin Yoon
Juho Lee
Venue: ICML 2025 poster (2025)
Keywords: Ensemble Distillation
Flow Matching
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We distilled ensemble distribution using flow matching

Abstract: Neural network ensembles have proven effective in improving performance across a range of tasks; however, their high computational cost limits their applicability in resource-constrained environments or for large models. Ensemble distillation, the process of transferring knowledge from an ensemble teacher to a smaller student model, offers a promising solution to this challenge. The key is to ensure that the student model is both cost-efficient and achieves performance comparable to the ensemble teacher. With this in mind, we propose a novel ensemble distribution distillation method, which leverages flow matching to effectively transfer the diversity from the ensemble teacher to the student model. Our extensive experiments demonstrate the effectiveness of our proposed method compared to existing ensemble distillation approaches.

Ensemble Distribution Distillation via Flow Matching
Jonggeon Park Giung Nam
Hyunsu Kim
Jongmin Yoon
Juho Lee'
Abstract
EnDD
EDFM
Neural network ensembles have proven effec- tive in improving performance across range of tasks; however, their high computational cost limits their applicability in environments for large models. Ensemble dis- tillation, the process of transferring knowledge from an ensemble teacher to smaller student model, offers promising solution to this chal- lenge. The key is to ensure that the student model is both cost-efficient and achieves performance comparable to the ensemble teacher. With this in mind, we propose a novel ensemble distribution distillation method, which leverages flow match- ing effectively transfer the diversity from the ensemble teacher the student model. Our tensive experiments demonstrate the effective- ness of our proposed method compared to exist- ing ensemble distillation approaches.
0.75
0.65
0.35
0.30
0:25
0.20
TRN (0.0010)
W/ mixup (0.0355]
wi mikup and RA (0.09488)
AUX (0.1436)
Figure 1. How good is ensemble distribution distillation? It il- lustrates how effectively different ensemble distribution distilla- tion methods learn the diversity of the ensemble teacher. The x- axis represents the teacher diversity under different diversification strategies ($ 5.1), while the y-axis shows the performance of the distilled student each setting (NLL; lower is better). The size of each circle marker indicates the level of diversity learned by the student, demonstrating how well each distillation method ab- sorbs the teacher diversity and mimics the ensemble behavior of producing diverse predictions ($ 5.2). Our proposed EDFM ap- preach effectively learns teacher diversity it becomes more pro- nounced, continually absorbing it and outperforming baselines.
1. Introduction
The concept of ensembling neural networks has long been a foundational approach in machine learning (Hansen and Salamon, 1990), based on the idea of constructing strong hypothesis by combining multiple weaker ones (Kearns, 1988). Interestingly, ensemble methods remain highly rel- evant in the era of deep learning, offering simple but ef- fective way to improve the performance of deep neural net- works (Ciresan et al., 2012; Krizhevsky et al., 2012). How- ever, despite its benefits, notable drawback of the ensem- ble method becomes more pronounced with deep neural networks: "many ensembles are large and slow" (Buciluä et al., 2006). To address this, the machine learning com- munity has naturally gravitated toward compressing the en- sembles into single, more efficient model (Buciluă et al., 2006; Ba and Caruana, 2014; Hinton et al., 2014).
Ensemble distillation is specific form of knowledge distil- lation (Hinton et al., 2014), where the teacher model is an ensemble. Unlike the typical knowledge distillation setup, in which single teacher model provides one prediction per training instance, ensemble distillation incorporates predic- tions from multiple teacher models, introducing diversity. While this diversity is a key distinguishing feature, simply adopting the basic knowledge distillation objective (Ba and Caruana, 2014; Hinton et al., 2014), i.e., averaging the dis- tillation losses or teacher predictions, overlooks the diver- sity within the ensemble teacher. Consequently, the main goal in the ensemble distillation literature is to effectively leverage the multiple predictions, or the empirical predic- tive distribution, produced by the ensemble teacher (Cui et al., 2020; Malinin et al., 2020; Penso al., 2022).
Equal contribution 'Korea Advanced Institute of Science and Technology, Republic of Korea. Correspondence to: Juho Lee caabolcc@kaist.ac.ke
Despite recent advances, significant performance gap mains between ensemble teachers and their distilled stu- dents. This gap arises primarily from two key limitations: First, deep neural networks are typically trained to achieve near-zero training errors, causing ensemble models to pro- duce highly correlated predictions on the training set. a
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Ensemble Distribution Distillation via Flow Matching
result, ensemble distillation methods operating on the same dataset struggle to transfer meaningful diversity student models (Nam et al., 2021). Second, the capacity of student models is inherently constrained. Since ensemble distilla- tion aims to approximate ensemble predictions efficiently, there exists a tradeoff between inference speed and pre- dictive performance, e.g., embedding multiple subnetworks within a single model (Nam et al., 2021; Penso et al., 2022) fundamentally limits the expressiveness, making it difficult to capture the full diversity of ensemble predictions.
calegorical probabilities as p(x, 0) softmax(z(s; 0)), and their element-wise logarithm as log p(xr,0), 0). For no- tational simplicity, we will often omit the dependence of logits and probabilties on inputs, e.g., p(z, Q) is written as Pe when x is clear from the context.
Ensembles. We view ensembles within a Bayesian frame- work, where the model parameters & are treated as random variables. More precisely, the predictive categorical distri- bution is approximated by an ensemble of neural networks defined by the set of parameters (0m) JM m==
Recently, Kim et al. (2024) reformulated ensemble distilla- tion as a distribution-nne matching problem, leveraging the dif- fusion Schrödinger bridge (DSB) to map a single model's predictions to ensemble outputs. Unlike traditional meth- ods, DSB enables flexible student modeling by iteratively refining predictions through a lightweight score network. However, their approach remains limited to approximating the ensemble's 's point estimates rather than modeling the full distribution of teacher predictions.
M plule, 0m); (1)
p(yl=) = Ee 0) ~
Σ
M
assuming the parameters are sampled from the approxi- mate posterior g()). In Bayesian literature, g(0) is com- monly modeled using a Gaussian or mixture of Gaus- sians facilitate tractable sampling (Maddox et al., 2019; Shen et al., 2024). Focusing on the classification problem, p(y|æ: l2, 0) = Cat (yl represents a K K-way categorical distribution with event probabilities (a
In this paper, we propose a novel ensemble distribution distillation algorithm based on flow matching (Lipman et al., 2023). Our approach introduces lightweight network that learns the vector field mapping a single model's pre- diction to the distribution of ensemble teacher predictions. Crucially, we emphasize the role of teacher diversity in distillation and show that existing methods fail to capture this diversity effectively, whereas ours succeeds (Fig. 2). Furthermore, our method enables fast, parallelizable infer- ence, significantly reducing wall-clock time compared to prior approaches. We substantiate our claims through ex- tensive experiments, covering a range of tasks including image classification and language modeling. Our contribu- tions are summarized as follows.
Ensemble distillation. The most straightforward way to distill the ensemble teacher, defined by (()), into the stu- dent model, parameterized by Φ, is to minimize
(0))(()) [H] (Po.P)))
where H|:-1 computes the cross-entrope between two cat- egorical probabilities, and p& is the categorical output from the student model. Eq. 2 is the average of standard distilla- tion losses (Hinton et al., 2014), and it is equivalent to
H. (Eo-q(0) [pa] .Po).
(3)
which is the standard distillation loss with the mean of the ensemble predictions. Unfortunately, this basic strategy en- tirely eliminates the diversity in ensemble predictions.
We present comprehensive analysis of ensemble dis- tribution distillation with respect to both diversity and fidelity. In particular, we investigate when ensemble teachers exhibit diversity (S 5.1) and how effectively distillation methods can learn from it ($S 5.2 and 5.3).
Ensemble distribution distillation approaches make the stu- dent model capture the distribution of teacher predictions. Specifically, X' x θ + A K-1 can be interpreted as a parametric function that maps the input space to the class- probability simplex. The random variable n = p(r, O) then represents the distribution of categorical probabilities for a given under the posterior over parameters g(0). By introducing the induced distribution g(TT.) (z)), we obtain
We propose a novel ensemble distribution distillation algorithm, Ensemble Distillation via Flow Matching (EDFM), which, as the name suggests, is based on the flow matching framework ($ 4). Our extensive ex- periments on image classification and language tasks validate both the efficiency (8 5.4) and effectiveness (S$ 5.5 and 5.6) of the approach.
p(ylx) Cat(ylp(r, ())g
(4)
2. Preliminaries
Cat F. (x)
AA-1
Notation. We focus on a K-way classification problem, where a neural network parameterized by 0 E θ takes in- puts x E X and produces predictions for the corresponding labels y E V, with logits 2xx 0) E K We also denote the
Ultimately, the core idea of the ensemble distribution dis- tillation scheme is to use a single student model, parame- terized by Φb to model the induced distribution (s(T)l))) (x).
2

---

Ensemble Distribution Distillation via Flow Matching
For example, a Dirichlet Prior Network (DPN; Malinin and Gales, 2018) could serve as the student model if we assume that ( (z) follows a Dirichlet distribution. While the Dirichlet distribution is a common choice for the induced distribution (Cui et al., 2020; Malinin et al., 2020), due to its conjugacy with the categorical distribution and its ability to facilitate tractable integration in Eq. 4, it is not necessar- ily required. The induced distribution is defined indirectly through the observations {p(2,, om)) MM with no guaran- tee that they follow the Dirichlet distribution. It motivates modeling the induced distribution as arbitrary, rather than necessarily Dirichlet (Penso et al., 2022).
importantly, "how to distill the ensemble teacher into a sin- gle student model" (Cui et al., 2020; Malinin et al., 2020; Du et al., 2020; Ryabinin et al., 2021; Nam et al., 2021; Nam et al., 2022; Penso et al., 2022). A central philosophy shared by most of these studies is that "diversity matters when learning from ensembles" " (Nam et al., 2021), high- lighting the importance of effectively accounting for the di- versity inherent in the ensemble teacher during the distilla- tion process. To this end, 1) one-to-one ensemble distilla- tion approaches design the student model to handle multi- ple predictions, using methods such as multi-headed mod- cling (Tran et al., 2020) or weight-sharing techniques (Ma- riet et al., 2021; Nam et al., 2021), while 2) ensemble dis- tribution distillation methods model the distribution of the ensemble predictions (Cui et al., 2020; Malinin et al., 2020; Ryabinin et al., 2021; Penso et al., 2022).
Flow matching, Flow Matching (FM; Lipman et al., 2023) a framework for training simulation-free continuous nor- malizing flows (Chen et al., 2018) by learning vector fields that transport a base density (o(z:) to a target density P1 (1) along a deterministic probability path whose density func- tion at time t is specified as (z). The probability path follows the neural ordinary differential equation (ODE)
Fast and efficient ensembling. There has been per- sistent demand for more cost-efficient ensemble methods, as ensembles of deep neural networks are typically asso- ciated with high computational costs during both train- ing and inference. Research under the terms fast or effi- cient ensembling has aimed to address this demand, includ- ing: "how reduce training costs for ensemble construc- tion" (Huang et al., 2017; Garipov et al., 2018; Benton et al., 2021), and "how to reduce inference costs for ensemble behaviour" (Lee et al., 2015; Wen et al., 2020; Havasi et al., 2021; Yun et al., 2023). In particular, Yun et al. (2023) emphasized the importance of reducing test-time costs for practical applications and proposed a framework that mim- ies the output of one ensemble component based the out- put of another using lightweight neural network. Kim et al. (2024) further extended this approach by utilizing the Schrödinger bridge algorithm (Liu et al., 2023a), mimick- ing the mean of the ensemble predictions. Our approach is most closely related to these works, as also leverages the outputs of a single ensemble component. However, the key difference is that we perform ensemble distribution dis- tillation via flow matching, sampling from the ensemble's predictive distribution rather than producing a single point estimate, thereby better capturing the ensemble's 's diversity.
dr,- = us (xx))( 20 1o-ppp po(o))
(5)
where (x(x) represents the vector field that satisfies the continuity equation for mass conservation. FM is trained to minimize the discrepancy between the learned vector field ur and the ground-truth vn derived from data:
(r) = Er,,.r (A(E) [ の) - ve (=()))))] (6)
where ~ T defined over [0, 1], 24 ~ Pt. and A is a time-dependent weighting function. However, such for- mulation Eq. 6 is challenging in general as obtaining the (marginal) ground-truth vy (a:) is typically intractable. In- stead, we model FM in logit space conditioned on input x, requiring a Conditional Flow Matching (CFM) formulation conditioned by the base density Po. where both the base and target distributions depend on an auxiliary variable a:
(CF
(7)
-Er A(1) ue [1ee e; x)- ve ve(ste/ro
Flow matching and knowledge distillation. Meanwhile, there have also been efforts (Shao et al., 2024; Huang et al., 2023; Yao et al., 2024) to incorporate flow matching or diffusion models into knowledge distillation. These meth- ods, however, bear greater resemblance to DBN (Kim et al., 2024) than to EDFM, as they primarily rely on one- to-one alignment of logits or intermediate features between teacher and student models. Distilling the knowledge of the diffusion model itself, especially in terms of synthe- sis efficiency, has also been widely investigated (Salimans and Ho, 2022; Song al., 2023; Yin et al., 2024a; Yin al., 2024b). We expect that these contributions might offer meaningful insights for further improvements of EDFM.
where t ~ T, xo ~ po (c), and x, pe( c). FM
has been successfully applied to generative modeling due to improved training stability and computational efficiency.
3. Related Work
Ensemble distillation. In recent years, extensive research on ensemble distillation using deep neural networks has fo- cused on three key design components: "how to construct the ensemble teacher" (Ba and Caruana, 2014; Korattikara et al., 2015; Bulò et al., 2016), "how to design the student model" (Cui et al., 2020; Malinin et al., 2020; Tran et al., 2020; Mariet et al., 2021; Penso et al., 2022), and, most
3

---

Ensemble Distribution Distillation via Flow Matching
4. Approach
4.2. Student network
4.1. Flow matching in logit space
The student network us is designed to condition on the in- put data by incorporating relevant information associated with x. To this end, we utilize features extracted from the penultimate layer of pretrained teacher network- for ex- ample, representations prior to the average pooling layer in the ResNet architecture (He et al., 2016) condition- ing inputs. Leveraging such pretrained features not only enhances the student network's 's ability to discriminate be- tween different inputs but also transfers valuable knowl- edge embedded in the teacher network, thereby facilitating faster convergence and improved performance. Since the pretrained teacher network serving as a feature extractor remains fixed during flow matching training, the compu- tational overhead incurred is substantially lower than that required for training the teacher network itself.
For a given input æ, we consider the logits =(z, 0) with ~ ((),, predicted by the ensemble teacher with the ap-
proximate posterior ((). These logits are assumed to fol-
low the distribution Pı:
(8) where the dependence on 0 is omitted for simplicity, and denote 2 21x 0). Our goal is to construct prob- ability path whose density function at time is specified as Pr. for & E (0, 1), starting from the source distribution (o((|z)= = (-; 12) and transitioning to the target dis- tribution p1 (-(x)). To achieve this, we introduce a condi-
2 ~ p1
tional vector field u. : [0, x RKK X - R&KK in the K-
dimensional logit space:
For the student network architecture, we employ a de- noising multilayer perceptron (MLP) structure proposed by Li et al. (2024). This lightweight architecture consists of a limited number of residual MLP blocks. Within each block, the perturbed sample 2;" is concatenated with the conditioning input, followed by the integration of time- embedding information via an adaptive layer normaliza- tion mechanism with zero initialization (Peebles and Xie, 2023). This processed representation then passes through an MIP layer with skip connection. Despite being sim- ple, this simple architecture proves sufficient to capture the target distribution, given the inherently low- low-dimensional nature of the logit space. Additionally, due to the compu- tational efficiency of MLP layers on modern hardware ac- celerators such as GPUs, the MLP-based student network offers significant advantages in execution time, particularly for parallel sampling, as will be shown in $ 5.4. It should be noted, however, that the denoising MLP is not the sole practical option for EDFM. While the MIP architecture of- fers a favorable balance between computational efficiency and performance, alternative architectures, such as trans- formers, may be more suitable in specific contexts, particu- larly when scalable student networks are required. Further analysis of the impact of network architecture on EDFM performance is provided in Appendix A.2.
u( 21 2 2
(9)
t
which induces conditional probability path:
Paj (= = = (1-11 ())) (10)
for 0 S < Conditioned on = ~ P1 (zf the random
variables zii conditioned on E [0, 1] are given by:
- = tz + (1
z)Ø ~ Pt|1,
(11)
and the marginal random variable 2 Pa is defined as a linear combination of 2 po and 2 ~ Pi:
x = レ + (1 - 1)=0 ~ pÝ.
(12)
Building on these formulations, we model u using ugs a neural network parameterized by &b The network is trained by minimizing the conditional flow matching loss:
((Φ)= = Fr, 1 [A(t)]lu(( (1, x)) - (2f - = =)]
(13)
where t ~ T, 20 ~ Pu. 27 ~ Pl. and \ denotes time- dependent weighting function (ef. $ 4.3).
It is important to note that our formulation of flow matching represents one of the simplest instances, corresponding to a rectified flow (Liu et al., 2023b) between Gaussian noise and the data distribution. More sophisticated approaches are conceivable, including 1) incorporating the geometric structure of logits, such as flow matching constrained to the probability simplex, or 2) modifying the perturbation schedule, for example, adopting diffusion models instead of flow matching. Nonetheless, empirical evaluation re- vealed that many of these more complex formulations un- derperform relative to the simple approach presented here. Further ablation studies examining various formulations are provided in Appendix A.
4.3. Training and inference
We identified the variance of the initial Gaussian noise a and the time distribution T as critical factors influenc- ing the performance of EDFM. Through extensive ablation studies, we set a = 4 and defined T as a distribution over [0, 1] that places exponentially greater weight near 1, i.e., near the data. Additionally, we adopted the precondi- tioning technique proposed by Karras et al. (2022) for the student network, ensuring that the network input and loss function maintain comparable scales with unit variance,
4

---

Ensemble Distribution Distillation via Flow Matching
Algorithm 1 Training EDFM
Algorithm 2 Sampling from EDFM
Require: Train dataset D= (o:)]= noise distribu- tion Po = O2)), time distribution T, time- dependent weighting function X, empirical distribu- tion of ensemble teacher logits conditioned on data x pi (= learning rate 7,, pretrained teacher network
Require: Test dataset D {(ri)) noise distribution
Po == (-; 0, o²I), number of sampling steps N, trun- cation 6, sampling step schedule {1 = tN:...,Me 4o =
<). ODE solver, pretrained teacher network 8, and trained student network
0, and student network e
1: Sample x D
1: repeat
2: Sample zi ~ Po
2: Sample ~ D
3: Extract feature T of input from teacher A
3: Sample =0 ~ Po. zf ~ p1(=il))
4: < =0
4: Extract feature I of input x from teacher 0
5: for i = N to 2 in reverse do
5: Obtain perturbed sample 2i" according to (12)
6: 2 f-
ODE solver(z, 22 tú, 1 us 2, I))
6: Compute student network output us (1, =
7: end for
7: Compute loss Lg according to (13)
z 4- 2 (t1 to)
8: Update parameters: á + Φ- 1/ ab
Return
9:
until Ø is converged
in Kim et al. (2024). For the commonsense reasoning ex- periments in $ 5.6, we employ a pretrained LLaMA-2-7B base model (Touvron et al., 2023). For details on ensemble teacher construction, please see Appendices B.1 and B.2.
while the variance of the network output is minimized. The time-dependent weighting function À was selected under this principle. The training algorithm for EDFM is detailed in Algorithm 1.
Unless specified, 'Meanasu' results reported the mean and standard deviation calculated over three trials across tables. The best and second-best values are high- lighted using bold-faced underline and underline, respec- tively. As baselines, we include KD (Hinton et al., 2014), EnDD (Ryabinin al., 2021), FED (Penso et al., 2022), and DBN (Kim al., 2024), and the ensemble teacher. We refer readers to Appendices B.3 and B.4 for details on eval- uation metrics and baseline methods.
For sampling, the choice of ODE solver, the scheduling of sampling steps, and the total number of steps proved es- sential to performance. After thorough evaluation, we em- ployed the Heun solver (Karras et al., 2022) and sched- uled sampling steps to be exponentially concentrated near the data. Because our flow matching formulation reduces to rectified flow, it generates relatively straight trajecto- ries toward the logits, facilitating efficient sampling with a reduced number of network funetion evaluations (NFEs). Empirically, we observed that high-quality logits can be ob- tained with few as two steps (equivalent to five NFEs), even in relatively high-dimensional spaces of up to 100 di- mensions. We emphasize here that the iterative nature of flow matching sampling does not substantially increase the overall sampling time. This efficiency arises because the pretrained teacher network, which serves as fixed feature extractor, is executed only once per input, as the condition- ing information remains constant throughout the sampling iterations. The algorithm for sampling from EDFM is pre- sented in Algorithm 2. Please refer to Appendix A.3 for further details on the design choices of EDFM.
5.1. When do teacher predictions get diverse?
The existing ensemble distillation literature emphasizes the importance of ensuring sufficiently diverse teacher predic- tions for effective distillation (Malinin et al., 2020; Nam et al., 2021). However, deep neural networks tend to produce over-confident predictions on training data, making diffi- cult to capture the ensemble diversity of the teachers when this data is reused during distillation. To address this, vari- ous strategies have been employed to diversify teacher pre- dictions. For instance, Malinin et al. (2020) employed aux- iliary datasets, Nam et al. (2021) and Nam et al. (2022) de- veloped diversifying perturbations, and Penso et al. (2022), Yun et al. (2023), and Kim et al. (2024) utilized mixup aug- mentation applied to the training data.
5. Experiments
We utilize Multi-SWAG (Wilson and Izmailov, 2020) and Multi-IVON (Shen al., 2024) as ensemble teachers for image and language tasks, respectively, owing to their scal- ability and proven effectiveness (Wilson et al., 2022). For the image classification experiments in $ 5.5, we use a ResNet with a depth of 32, projection shortcuts, filter re- sponse normalization, and Swish activation, as described
Building on these established conventions in the ensemble distillation literature, present a comprehensive analy- sis of teacher diversity across various diversification strate- gies for image classification tasks, including: 1) RandAug- ment (RA; Cubuk et al., 2020) and 2) mixup (Zhang et al., 2018), both of which are strong image augmentation tech- niques; 3) TDiv (Nam et al., 2022), a perturbation- baaed di-
5

---

Ensemble Distribution Distillation via Flow Matching
and their corresponding teacher diversity during distillation (VAR; higher values indicate greater diversity), while the y- axis shows the test performance and calibration of the dis- tilled model (NLL and ECE; lower values indicate better performance and calibration). The size of the circle mark- ers reflects the diversity of the distilled model on the test split (VAR: larger markers indicate greater diversity). The VAR values are computed by sampling 30 predictions for each, even for EnDD, which typically omits the sampling process when computing the mean prediction.
Table 1. Diversity analysis of teacher predictions. We analyze teacher diversity by computing ensemble variance decomposition, which consists EnsUne. AvgUnc, and VAR. Higher VAR indi-
cates greater diversity in ensemble predictions.
Diversificution
Decomposition
Dataset
Split
RA
mixup
TDiv
UnsUne
Avgline
VAR (f)
TRN
0.0087
0,0069
0.0018
TRN
0.0705
0,0350
0.0355
TRN
v
0.1212
0,0539
0.0673
TRN
0.1278
0.0638
0.0640
TRN
0.1850
0.0902
0.0948
C10
TRN
0.1868
0.0853
0.1014
Fig. conveys a key insight into ensemble distribution dis- tillation: to be where the diversity is. As ensemble dis- tribution distillation methods are applied to setups where teacher diversity becomes more prominent, they tend to more effectively mimie the ensemble behavior of generat- ing diverse predictions (as reflected by the increasing circle marker sizes). However, while the student model's ability to replicate the teacher's diversity improves, this does not always lead to consistent performance gains for EnDD and FED. For EnDD, the best performance occurs in the sec- ond setup, while for FED, the highest performance is ob- served in the third setup. Beyond these points, despite the increased diversity, performance begins to decline.
VAL
0.0845
0.0409
0.0436
AUX
0,2680
0,1244
0.1436
versifying approach; and 4) AUX, which uses 10 as an auxiliary dataset. To quantify diversity without relying on cround-truthh labels, we compute the ensemble variance (VAR). with higher VAR values indicating greater diversity. For detailed definitions of EnsUnc, AvgUne, and VAR, we refer readers to Appendix B.3.
Table presents the diversity of ensemble teacher predic- tions across the training (TRN) and validation (VAL) splits of 10/1R- as well as the auxiliary (AUX) CINIC- 10 dataset. The results clearly shows that, without ex- plicit diversification strategies, ensemble teachers produce predictions with limited diversity on training images. We hypothesize that this lack of teacher diversity can hinder effective ensemble distribution distillation when these im- ages are reused. In 5.2, we experimentally validate this hypothesis. demonstrating that ensemble distribution dis- tillation improves as the teacher provides more diverse pre- dictions through diversification strategies. In other words, a good distillation method should capture and leverage as much diversity as possible from the ensemble teacher.
Notably, our proposed EDFM approach not only shows higher prediction diversity but also demonstrates a consis- tent improvement in performance. As we will discuss in $ 5.3, our method exhibits strong fidelity to the teacher; that is, the student learns not merely to generate predictions with high variance, but to genuinely replicate the ensemble teacher's behavior, thereby producing diverse predictions in a more meaningful way. To summarize, ensemble dis- tribution distillation only can be done effectively when the teacher exhibits a certain level of diversity. However, sim- ply learning to generate high variance is insufficient; the student must genuinely mimic the ensemble teacher's be- havior, and we will validate this in our method.
5.2. When do ensemble distillation methods work?
5.3. How effective ensemble distillation methods are?
In our CI setup. we examine the following scenar- ios for distilling the ensemble teacher, each capturing dif- ferent levels of teacher diversity: /) 7RN, 2) TRN w/ mixup, 3) TRN w/ mixup and RA, and 4) AUX. As shown in Ta- ble 1, the diversity of the ensemble teacher gradually in- creases, with VAR values of [.0018, .0355, .0948, 1436). In the AUX setup, CINIC-10 is used as an auxiliary dataset, with images from CIFAR-10 excluded (see Appendix B.1 for details), introducing slight distribution shift (De Silva et al., 2023). We also exclude TDiv from our analysis, as its computation requires backward pass for perturbation, which significantly increases training costs.
We evaluate the effectiveness of ensemble distribution dis- tillation methods in three aspects: 1) diversity, which mea- sures how well the student exhibits prediction diversity, key characteristic of the ensemble; 2) fidelity, which as- sesses how closely the student mimics the teacher; and 3) generalization, which determines whether the student can achieve performance on par with the ensemble teacher across a range of tasks. The in-depth analysis of general- ization will be discussed in $$ 5.5 and 5.6, while here we focus on the diversity and fidelity analyses.
Diversity analysis. The operational principle of ensem- ble methods is diversity; ensembles achieve better perfor- mance when their individual members offer diverse predic- tions (Krogh and Vedelsby, 1994; Dietterich, 2000), a prin-
Fig. 1 summarizes the results of ensemble distribution dis- tillation using EnDD, FED, and our proposed EDFM algo-
rithms across four setups. The x-axis represents the setups
6

---

Ensemble Distribution Distillation via Flow Matching
Teacher EnDD
FED
EDFM
0.4
DGR
DGR
EnotLL
0.2
TRN (0.0018)
mixup and [0. (0.0948)
AUX (0. 1436)
TVD
W2
TVD
(0.0355)
Figure 2. Diversity analysis. We compute the ambiguity decom- position to assess the diversity of the distilled model's predictions and its contribution to the ensemble gain, as shown in shaded portion of each bar (AMB).
JSD
KLD
JSD
KLD
TRN
w/ mixup
10. 0018)
(0 00.0355;
DGR
DGR
ciple that holds true even in the deep learning era (Lakshmi- narayanan et al., 2017; Wilson and Izmailov, 2020; Ortega et al., 2022). To this end, we compute the ensemble ambi- guity (AMB) decomposition to assess how diverse the set of predictions is and how much this diversity contributes to the ensemble gain. In particular, AMB provides insights into the performance improvement of the ensemble over the average individual performance through ambiguity de- composition, revealing the gain derived from diversity. For more details, please refer to Appendix B.3.
15
TVD
TVD
JSD w/ mixup and RA
KLD
JSD
KLD
AUX (0.1436)
10. 0948)
Figure 3. Fidelity analysis. We compute DGR, TVD, KLD, JSD, and W/ to assess how closely the distilled model's predictions align with those of the teacher.
Fig. 2 illustrates the results of our diversity analysis on the CIFAR-10 test split. The diversity of the distilled student increase as the distillation setup more effectively captures the teacher's diversity; from (A) to (D), the shaded area, representing AMB, expands. again highlights that en- semble distribution distillation methods successfully learn the ensemble behavior of generating diverse predictions, particularly when teacher diversity is evident; in the (A) setup, where no diversification strategies are applied, all methods fall short of the ensemble teacher's level of AMB, resulting limited performance improvements from AMB. Notably, our EDFM not only consistently exhibits higher diversity but also shows significant gains from that diver- sity, as demonstrated by AMB. It clearly indicates that our method does not simply learn to generate high-variance predictions, but instead effectively harnesses diversity to drive performance improvements.
tion distance (TVD), Kullback-Leibler divergence (KLD), Jensen-Shannon divergence (JSD), and Wasserstein-2 dis- tance (W2)). Lower values for these metrics indicate that the student's 's ensemble predictions are more closely aligned with those of the ensemble teacher. For detailed definitions of each metric, we refer readers to Appendix B.3.
Fig, 3 illusrates the results of our fidelity analysis on the CIFAR-10 test split. For all the metrics we measured, our EDFM achieves higher fidelity compared to the baselines across all setups. Notably, even when comparing distribu- tions (rather than just the mean of ensemble predictions), as seen with W2, EDFM demonstrates superior fidelity. It clearly indicates that our approach effectively models the induced distribution over categorical probabilities (<Trr/ar) and closely approximates the ensemble teacher's 's predictive distribution p(y(æ> in Eq. 4.
Fidelity analysis. Next, we analyze the effectiveness of en- semble distribution distillation methods from the perspec- tive of fidelity, which quantifies how well student model replicates the teacher's 's predictions. Since the core philoso- phy of knowledge distillation is to encourage the student to mimic the teacher (Ba and Caruana, 2014; Hinton et al., 2014), measuring the degree to which the student repli- cates the teacher is essential. While higher fidelity does not always guarantee better generalization-art example being the self-distillation setup of Furlanello al. (2018), where student that perfectly matches the teacher can- not surpass it itiiiiyitty analysis remains crucial for under- standing the underlying mechanisms, as discussed by Stan- ton et al. (2021). To quantify teacher-student fidelity, we compute several metries: disagreement (AGR), total varia-
5.4. How efficient ensemble distillation methods are?
To efficiently produce ensemble predictions, Nam et al. (2022) and Penso et al. (2022) use parallelized inference by replicating a single test input to form a mini-batch. How- ever, unless the neural network is extremely small, process- ing time increases significantly as the batch size grows. On RTX A6000, ResNet processes single input (batch size of 1) in 1.404 ms, while batch sizes of [32, 64, 128, 256] take [3.806, 6.993, 13.05, 24.68] ms, exhibiting a near-linear increase in our CIFAR- 100 setup. suggests inherent limits in reducing runtime when batched inference is required for ResNet, as in FED (Penso al., 2022).
7

---

Ensemble Distribution Distillation via Flow Matching
MultiSWA
EnDD
FED
Table 2. Evaluation results for image classification tasks. We
EDFM (NFE=3)
EDFM (N
+
EDFM (NFE=7)
evaluate ACC, NLL, and ECE to assess how accurate and well- calibrated the predictions are.
1.39
Dataset
Method
ACC (1)
NLI. (4)
ECE (J)
1.02
SWA buseline (w/o distill)
0.932uwacc
0.289uc
0.046
0.97
00
Multi-SWAC teacer
0,946
0.167
0,007
+
KD (Hinton et al., 2014)
0.931
0.229ea ons
0.020
1.12
1404
2106
8088
C10
EnDD (Ryabinin et 2021)
0.937
0,228st.oo on
0.033a000i
DBN (Kim et al,, 2024)
0 0.936x0001
0,218 on
0.028-0000
1.02
RED (Penso al., 2022)
0.922mwace
0.238ucxzzs
0.009
EDFM (ours)
0.931-0300
0.216-1
0.009-0000
0.97
0.93
SWA baseline (w/o distill)
0.735an0044
1.369ed..0
0.173300044
1404
2808
4221
$615
7019
Multi- SWAii-SWGGGa teacher
0.777
0.835
0.018
Wall-clock time (microseconds)
KD (Hiaton al., 2014)
0. 0.760ana33
0,965e0
3,073
C100
EnDD (Ryabinin et al., 2021)
0.761
1.03
a.084
DBN (Kim et al., 2024)
0.757
1,002ed oo
0 0.098s0001
Figure 4. Runtime analysis on CIFAR-100 Down and left indicate better performance and cost-efficiency, with the lower- left corner being the most desirable.
FED (Penso al., 2022)
0.721 40a03
1.026xa
0.028 a.028
EDFM (ours)
0.761
0.932ez o0s
0.056s00000
Table 3. Robustness to distribution shifts. We evaluate ACC, NLL, and ECE to assess robustness to distribution shifts.
In contrast, our method performs single inference for large backbone networks (ResNet) and parallelized inference for smaller student networks (MLP), where the benefits of batch inference are more pronounced. In the same CIFAR- 100 setup used for the previous analysis, processing sin- gle input (mini-batch size of 1) with MLP model takes 0.053 ms, while processing batched inputs with mini-batch sizes of [256, 384, 512, 640, 768, 1024] takes 10.101, 0.114, 0. 0212 0.149, 0.182, 0.210] ms. Since the process- ing time remains at a significantly lower cost compared to ResNet, EDFM demonstrates a strong advantage in seal- ing ensemble predictions through parallelized inference. Building on this advantage, we investigated how the num- ber of ensemble predictions in EDFM can scale within a constrained execution time budget.
Dataset
Mothod
ACC (t)
NLL (J)
BCE (1)
SWA baseline (w/o distill)
0.859xnac1
0.631 o0
0. 0.100-n000
Multi-SWWe teacher
0,386
0.020
KD> (Hinton al,, 2014)
0. 0.85940a01
0.480xa 077
0.048-0025
CI0.1
LnDD (Ryabinin et al., 2021)
0.869
0.482acz
0.064
DBN (Kim 2024)
8685000
0.481
0.07
FED (Penso et al., 2022)
0.8 0.839a
0.481acas7
0.028 anoos
EDFM (ours)
0. 0.858snace
0.455
0.034
SWA baseline (w/o distill)
802000004
0, 0,990ea
0
Multi-SWAGG teacher
0.818
0.652
0.066
KD (Hinton al., 2014)
80440000
0.715 n48
0.0 0.093a0033
C10.2
EnDD (Ryahinin et al., 2021)
0.780w2 0.81S
0.1 111224
DBN (Kim et al., 2024)
D. a08
0.781acmm
0.1 0.118a
(Penso al., 2022)
0.786
0.758ac
0.076
EDFM (ours)
0. 8.82200004
0.6666s
0.07
SWA baseline (w/o distill)
0.789an002
1.005a0.000
0.15 0.159ane52
Multi-SWAG teacher
8.8
0.615
0.055
KD (Ilinton al., 2014)
0. 0.788ana03
0.726uc
0.0933me0
STL.
FnDD (Ryabinin et al., 2021)
0.797-00a4
0,770xa oor
0.1 1-00000
Runtime analysis. Fig. illustrates the evolution of both NLL and execution time as the number of student ensem- bles increases. While the execution time of MultiSWA en- sembles scales proportionally with the number of ensem- bles, EDFM exhibits only marginal increase in compu- tational overhead. Although the execution time of FED scales at a slower rate than MultiSWA, it introduces a non- negligible overhead, and more importantly, its performance merely improves in proportion to the increased execution time. In contrast, with EDFM, the enhancement in en- semble NLL closely mirrors that of the MultiSWA teacher. This not only provides empirical evidence that EDFM ef- fectively captures the diversity inherent in the teacher en- semble distribution, but also underscores its practical utility in enabling efficient fast ensembling.
DBN (Kim 2024)
0.759ucmm 0.794amvac
a24
FED (Penso et ol., 2022)
0.777-nar a
0.064 0.6984amm
-DFM (ours)
0.787awac
u.074
to distribution shifts and out-of-distribution scenarios by effectively capturing predictive uncertainty (Ovadia et al., 2019), ensemble distribution distillation methods should learn these strengths from the ensemble teacher. To assess this, we evaluate our approach on CIFAR-10/100 test splits, as well as distribution shift datasets (CIFAR- 10.1, CIFAR- 10.2, and STL) and the -iitriibbittiitiion dataset (SVHN). For more details, please refer to Appendix B.
Results. Table 2 shows the evaluation results on CIFAR- 10/100 test splits, clearly demonstrating the competitive performance of our proposed EDFM.
Robustness to Distribution Shifts. Table 3 presents the evaluation results on datasets associated with 10. Since these datasets share the same 10-class labeling as CIFAR- we evaluate ACC, NLL, and ECE in the same manner as the standard test split. However, a distinguish- ing feature from the original CIFAR-10 test split that,
5.5. Image classification tasks
We validate our proposed EDFM approach in terms of gen- eralization through image classification experiments on CI- FAR datasets. Since ensemble approach not only enhances in-ditribution performance but also improves robustness
8

---

Ensemble Distribution Distillation via Flow Matching
Table 4. Out-of-distribution detection.
We evaluate AUROC,
Table S. Evaluation results for commonsense reasoning tasks.
TNR95, and TNR99 to assess out-of- distriibiiittionn dection.
We evaluate ACC, NLL, and ECE to assess how accurate and well-calibrated the predictions are.
Dutaset
Method
AUROC (1)
TNR95 (1)
TNR99 (1)
SWA baseline (wlo distill)
0.893am012
0.706ano19
0,515a0053
Dataset
Method
ACC (1)
NLL (L)
ECE (!)
Multi SWAG teacher
0.935
0.786
0.679
KD (Hinton et al., 2014)
0.912-0017
0.731 co
0.572amc1
IVON baseline (w/o distill)
0,710z0011
1.842a.167
0,246unco0
SVHN
EnDD (Ryabinin al., 2021)
0.895-0002
0 6686-0
0,450s0530
ARC.C.
Multi-IVON teacher
0.675
0.891
0.091
DBN (Kim et al., 2024)
0.907a0003
0.7 0.740pa
0 0.605
KD (Hinton et 2014)
0.662x0032
1.12620000
0.218an015
FED (Penso al., 2022)
0.932amas
0.778
0.637auss2
EDFM (ours)
0.715
1.113:0006
0.142
EDFM (ours)
0.934 n002
0.784.a
0.656-uc15
IVON baseline (wi distill)
0.888.c.0m
0.645
0.094.000
Multi-IVON teacher
0.868
0.357
0.026
ARC-E
KD (Hinton et al., 2014)
0. 0.850xc.004
0.570
0.057
RDFM (ours)
0.892.c.0c
0.390
0.044...co
because these datasets exhibit distribution shifts, predic- tions must not only be accurate but also well-calibrated. Our EDFM approach clearly not only achieves competitive ACC but also excels in NLL and ECE with wide margins, demonstrating strong robustness to distribution shifts.
IVON baselinc (wo distill)
0. a.809
0.653en011
0.1 1
Mali--VVON teacher
0. 0.794
0.514
0.032
OBQA
KD (Hinton et al., 2014)
0. 0.776+0.066
0.7 0.720x0012
0. 0.102:0005
EDFM (ours)
0.818
0.537
0.046
cotttf-distribution detection. Table 4 presents the evalu- ation results for out-of-distribution detection, assessing the model's s ability to distinguish between in-distribution and out-of-disttibutton inputs using receiver operating charac- teristic (ROC) metrics: area under the curve (AUROC) and true negative rates at 95% and 99% true positive rate (TNR95 and TNR99). We employ 1,000 in-distribution images from the 10 test split and 1,000 out-of- distribution images from the SVHN dataset, ensuring a bal- anced dataset to maintain the reliability of these metrics. Our EDFM approach outperforms baselines by producing higher entropy predictions for -ut-of-distribution inputs, reflecting greater uncertainty.
proves calibration (a slight drop in ACC is consistent with Cong et al. (2024), as do not explicitly sharpen the posterior in our experiments). The results clearly demon- strate that our EDFM improves upon the IVON baseline by effectively distilling the Multi-IVON teacher, produc- ing well-calibrated predictions that significantly outper- form the KD baseline and closely match the teacher's cal- ibration quality. It worth noting that our approach still leverages a lightweight MLP architecture for flow mod- eling in this large language model experiment; the model consists of three blocks with a hidden dimension of 256, totaling just 1.6 million than the 6.7 billion parameters of the LLaMA-2 backbone and even more compact than the 4.2 million parameters re- quired for LoRA fine-tuning. Accordingly, our approach achieves well-calibrated predictions comparable those of the Multi-IVON teacher while incurring virtually no addi- tional cost beyond single LLaMA-2 forward pass-unlike the teacher, which requires multiple passes.
smaller
5.6. Commonsense reasoning tasks
We further validate the scalability of our approach through experiments with large language models. The challenges of large and slow ensemble methods are amplified in mod- ern large-scale models, where the computational burden of performing multiple forward passes through models with billions of parameters becomes particularly significant. In this context, addressing the ensemble distribution distilla- tion problem is highly valuable; it would be more appro- priate for the ensemble model to remain in the develop- ment phase, serving teacher for distillation rather than being directly deployed. Consequently, we assess whether our approach can effectively perform ensemble distribution distillation from a Multi-IVON teacher, which exhibits im- proved calibration (Shen et al., 2024; Cong et al., 2024). For experimental details on the tasks and training specifics, we refer readers to Appendix B.2.
6. Conclusion
We present a novel ensemble distribution distillation ap- proach using flow matching, where lightweight neural network learns the vector field that transforms single model's prediction into the predictive distribution of the ensemble teacher. Extensive experiments demonstrate that our flow matching approach effectively captures the di- versity of the ensemble teacher, unlike prior works that fail to do so, which leads to substantial improvements in both accuracy and uncertainty calibration. It also enables fast, parallelizable inference, significantly reducing infer- ence costs compared to existing distillation methods, and remains scalable even for language tasks involving modern large language models. An intriguing avenue for future re- search is the development of more advanced flow matching strategies specifically tailored for ensemble distillation.
Table 5 summarizes evaluation results on commonsense reasoning tasks, including ARC-C, ARC-E, and OBQA. The lower NLL and ECE of the Multi-IVON teacher com- pared to the IVON baseline, as presented in the first two rows of each group, clearly demonstrate that ensembling with samples from Multi-IVON posterior distribution im-

---

Ensemble Distribution Distillation via Flow Matching
Acknowledgement
Chen, Tian Qi al. (2018). "Neural Ordinary Differen- tial Equations". In: Advances in Neural Information Pro- cessing Systems 31 (NeurIPS 2018) (cit. on p)3).
We thank Byoungwoo Park for sharing his insights. This work was partly supported by Institute of Information & communications Technology Planning & Evaluation(IITP) grant funded by the Korea government(MSIT) (No.RS- 2024-00509279, Global AI Frontier Lab; No.RS-2019- 11190075, Artificial Intelligence Graduate School Pro- gram(KAIST); No.RS-2022-11220711 Meta-learning Ap- plicable to Real-world Problems; SO.RS-2022-1122011 Development and Study of AI Technologies to Inexpen- sively Conform to Evolving Policy on Ethics). This mate- rial is based upon work supported by the Google Cloud Re- search Credits program with the award 19919990904 and Cloud TPUs from Google's TPU Research Cloud (TRC).
Cheng, Chaoran et al. (2024). "Categorical Flow Matching on Statistical Manifolds". In: Advances in Neural Infor- mation Processing Systems 37 (NeurIPS 2024) (cit. on p. 14).
Ciresan, Dan et al. (2012). "Deep neural networks segment neuronal membranes in electron microscopy images".
In: Advances in Neural Information Processing Systems 25 (NIPS 2012) (cit. on p. 1).
Clark, Peter et al. (2018). "Think you have solved ques- tion answering? try arc, the ai2 reasoning challenge". In: arXiv preprint arXiv: 1803.0:5 (cit. on p. 19).
Impact Statement
Coates, Adam, Andrew Ng, and Honglak Lee (2011). "An analysis of single-layer networks in unsupervised fea- ture learning". In: Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics. JMLR Workshop and Conference Proceedings, pp. 215- 223 (cit. on p. 18).
This paper focuses on enhancing ensemble distillation frameworks with generative modeling, such flow match- ing methods. By improving knowledge distillation through these approaches, our method contributes making model compression and deployment more efficient, particularly in resource-constrained environments. Although our method is broadly applicable across various downstream tasks, it does not inherently raise ethical concerns or pose direct so- cietal impacts. Nonetheless, as with any advancement in AI model training, the broader implications of increased effi- ciency and accessibility in real-world applications warrant careful consideration in future discussions.
Cong, Bai et al. (2024). "Variational Low-Rank Adapta- tion Using IVON". In: NeurlPS 2024 Workshop on Fine- Tuming in Modern Machine Learning: Principles and Scalability (cit. on pp. 19).
Cubuk, Ekin Dogus et al. (2020). "RandAugment; Practical Automated Data Augmentation with Reduced Search Space". In: Advances in Neural Information Processing Systems 33 (NeurIPS 2020) (cit. on onp.9 5).
References
Abe, Taiga et al. (2022). "Deep ensembles work, but are they necessary?" In: Advances in Neural Information Processing Systems 35 (NeurIPS 2022) (cit. on p. 20).
Cui, Yufei et al. (2020). "Accelerating Monte Carlo Bayesian prediction via approximating predictive uncer- tainty over the simplex". " In: IEEE transactions on neural networks and learning systems 33.4, pp. 1492-1506 (cit. on pp. 1,3).
Ba, Jimmy and Rich Caruana (2014). "Do deep nets really need to be deep?" In: Advances in Neural Information Processing Systems 27 (NIPS 2014) (cit. on pp. 1,3, 7).
Darlow, Luke et al. (2018). "Cinic- is not imagenet or cifar-10". In: arXiv preprint 11X0v:18 (cit. on p. 18).
Benton, Gregory et al. (2021). "Loss surface simplexes for mode connecting volumes and fast ensembling". In: Pro- ceedings of The 38th International Conference on Ma- chine Learning (ICML 2021) (cit. on p. 3).
De Silva, Ashwin et al. (2023). "The value of out-of- distribution data". In: Proceedings of The 40th Interna- tional Conference Machine Learning (ICMI. 2023) (cit. on p. 6).
Buciluã, Cristian, Rich Caruana, and Alexandru Niculescu- Mizil (2006). "Model compression". In: Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining (cit. on p. 1).
Dietterich, Thomas G (2000). "Ensemble methods in ma- chine learning". In: International workshop on multiple classifier systems. Springer, pp. 1-15 (cit. on p. 6).
Bulò, Samuel Rota, Lorenzo Porzi, and Peter Kontschieder (2016). "Dropout distillation". In: Proceedings of The 33th International Conference on Machine Learning (ICML 2016) (cit. on 3).
Du, Shangchen et al. (2020). "Agree to disagree: Adaptive ensemble knowledge distillation in gradient space". In:
10

---

Ensemble Distribution Distillation via Flow Matching
Advances in Neural Information Processing Systems 33 (NeurIPS 2020) (cit. on p. 3).
Huang, Tao et al. (2023). "Knowledge Diffusion for Distil- lation". In: Advances in Neural Information Processing Systems 36 (NeurIPS 2023) (cit. on p. 3).
Duvenaud, David al. (2020). "Your classifier is secretly an energy based model and you should treat it like one". In: Proceedings of The 8th International Conference Learning Representations (ICLR 2020) (cit. on p. 14).
Karras, Tero al. (2022). "Elucidating the Design Space of Diffusion-Baa Based Generative Models". In: Advances in Neural Information Processing Systems 35 (NeurIPS 2022) (cit. on pp. 4, 5, 14-17).
Eijkelboom, Floor et al. (2024). "Variational Flow Match- ing for Graph Generation". In: Advances in Neural In- formation Processing Systems 37 (NeurIPS 2024) (cit. on p. 14).
Kearns, Michael (1988). "Thoughts on hypothesis boost- ing". In: Unpublished manuscript (cit. on p. 1).
Kim, Beomsu et al. (2025). "Simple ReFlow: Improved Techniques for Fast Flow Models". In: Proceedings of The 13th International Conference on Learning Repre- sentations (ICLR 2025) (cit. on p. 15).
Furlanello, Tommaso et al. (2018). "Born again neural net- works". In: Proceedings of The 35th International Con- ference on Machine Learning (ICML 2018) (cit. on p. 7).
Garipov, Timur et al. (2018). "Loss surfaces, mode connec- tivity, and fast ensembling of dnns". In: Proceedings of The 6th International Conference Learning Represen- tations (ICLR 2018) (cit. on 3).
Kim, Hyunsu, Jongmin Yoon, and Juho Lee (2024). "Fast Ensembling with Diffusion Schrödinger Bridge". In: Proceedings of The 12th International Conference on Learning Representations (ICLR 2024) (cit. on pp. 2, 3, 5, 8, 9, 15, 18, 22).
Guo, Chuan et al. (2017). "On calibration of modern neu- ral networks". In: Proceedings The 34th International Conference on Machine Learning (ICML 2017) (cit. on p. 20).
Korattikara, Anoop et al. (2015). "Bayesian dark knowl- edge". In: Advances in Neural Information Processing Systems 28 (NIPS 2015) (cit. on p)
Hansen, Lars Kai and Peter Salamon (1990). "Neural net- work ensembles". In: IEEE transactions on pattern anal- ysis and machine intelligence 12.10, pp. 993-1001 (cit. on p. 1).
Krizhevsky, Alex and Geoffrey Hinton (2009). "Learning multiple layers of features from tiny images". In: Mas- ter's 's thesis, Department of Computer Science, University of Toronto (cit. on p. 18).
Havasi, Marton et al. (2021). "Training independent sub- networks for robust prediction". In: Proceedings of The 9th International Conference Learning Representa- tions (ICLR 2021) (cit. on p. 3).
Krizhevsky, Alex, llya Sutskever, and Geoffrey Hinton (2012). "Imagenet classification with deep convolutional neural networks". In: Advances in Neural Information Processing Systems 25 (NIPS 2012) (cit. on p- 1).
He, Kaiming al. (2016). "Deep residual learning for image recognition". In: Proceedings of the IEEE Con- ference on Computer Vision and Puttern Recognition (CVPR 2016) (cit. on p. 4).
Krogh, Anders and Jesper Vedelsby (1994). "Neural net- work ensembles, cross validation, and active learning". In: Advances in Neural Information Processing Systems 7 (NIPS 1994) (cit. on p. 6).
Heusel, Martin et al. (2017). "GANs Trained by a Two Time-Scale Update Rule Converge to Local Nash Equi- librium".. In: Advances Neural Information Processing Systems 30 (NIPS 2017) (cit. on p. 21).
Lakshminarayanan, Balaji, Alexander Pritzel, and Charles Blundell (2017). "Simple and scalable predictive uncer- tainty estimation using deep ensembles". In: Advances in Neural Information Processing Systems 30 (NIPS 2017) (cit. on p. 7).
Hinton, Geoffrey, Oriol Vinyals, and Jeffrey Dean (2014). "Distilling the Knowledge in Neural Network". In: Deep Learning and Representation Learning Workshop, NIPS 2014 (cit. on pp. 2, 5, 7-9, 19, 21).
Lee, Sangyun, Zinan Lin, and Giulia Fanti (2024). "Im- proving the Training of Rectified Flows". In: Advances in Neural Information Processing Systems 37 (NeurIPS 2024) (cit. on 15).
Huang, Gao et al. (2017). "Snapshot Ensembles: Train 1, Get M for Free". In: Proceedings of The 5th Interna- tional Conference on Learning Representations (ICLR
Lee, Stefan et al. (2015). "Why m heads are better than one: Training diverse ensemble of deep networks". In: arXiv preprint arXiv:1511.063 (cit. on p. 3).
2017) (cit. on p. 3).

---

Ensemble Distribution Distillation via Flow Matching
Li, Tianhong al. (2024). "Autoregressive Image Genera- tion without Vector Quantization". In: Advances in Neu- ral Information Processing Systems 37 (NeurIPS 2024) (cit. on pp. 4, 15).
Nam, Giung al. (2021). "Diversity matters when learning from ensembles". In: Advances in Neural Information Processing Systems 34 (NeurIPS 2021) (cit. on pp. 2, 3, 5).
Lipman, Yaron et al. (2023). "Flow Matching for Gener- ative Modeling". In: Proceedings of The 11th Interna- tional Conference on Learning Representations (ICLR 2023) (cit. on pp. 2, 3).
Nam, Giung et al. (2022). "Improving Ensemble Distilla- tion With Weight Averaging and Diversifying Perturba- tion". In: Proceedings The 39th International Confer- ence on Machine Learning (ICML 2022) (cit. on pp. 3, 5, 7)
Liu, Guan-Horng et al. (2023a). "I2SB: Image-to-Imen -omag Schrödinger Bridge". In: Proceedings of The 40th Inter- national Conference on Machine Learning (ICML 2023) (cit. on p)
Netzer, Yuval et al. (2011). "Reading digits in natural im- ages with unsupervised feature learning". In: NIPS work- shop on deep learning and unsupervised learning (cit. p. 18).
Liu, Xingchao, Chengyue Gong, and Qiang. Liu (2023b), "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow". In: Proceedings of The 1ith International Conference on Learning Repre- sentations (ICLR 2023) (cit. on p. 4).
Ortega, Luis A, Rafael Cabañas, and Andres Masegosa (2022). "Diversity and generalization in neural network ensembles". In: Proceedings of The 25th International Conference on Artificial Intelligence and Statistics (AIS- TATS 2022) (cit. on p. 7).
Lu, Shangyun et al. (2020). "Harder or Different? A Closer Look at Distribution Shift in Dataset Reproduction". In: ICML 2020 Workshop on Uncertainty and Robustness in Deep Learning (cit. on p. 18).
Ovadia, Yaniy et al. (2019). "Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift". In: Advances in Neural Information Pro- cessing Systems 32 (NeurIPS 2019) (cit. on p. 8).
Maddox, Wesley et al. (2019). "A simple baseline for bayesian uncertainty in deep learning". In: Advances in Neural Information Processing Systems 32 (NeurIPS 2019) (cit. p. 2).
Pakdaman Nacini, Mahdi, Gregory Cooper, and Milos Hauskrecht (2015). "Obtaining Well Calibrated Proba- bilities Using Bayesian Binning". In: AAAI Conference on Artificial Intelligence (AAAI) (cit. on p. 20).
Malinin, Andrey and Mark Gales (2018). "Predictive un- certainty estimation via prior networks". In: Advances in Neural Information Processing Systems 31 (NeurIPS 2018) (cit. on pp. 21).
Peebles, William and Saining Xie (2023). "Scalable Dif- fusion Models with Transformers". In: Proceedings of the IEFF/VVF International Conference on Computer Vision (ICCV) (cit. on p. 4).
Malinin, Andrey, Bruno Mlodozeniec, and Mark Gales (2020). "Ensemble Distribution Distillation". In: Pro- ceedings of The 8th International Conference on Learn- ing Representations (ICLR 2020) (eit. on pp. 5,211.
Penso, Coby, Idan Achituve, and Ethan Fetaya (2022). "Functional ensemble distillation". In: Advances in Neu- ral Information Processing Systems 34 (NeurIPS 2021) (cit. on pp. 1-3, 5, 7-9, 18, 22).
Mariet, Zelda E et al. (2021). "Distilling ensembles im- proves uncertainty estimates". In: Third symposium on advances in approximate bayesian inference (cit. 3).
Recht, Benjamin al. (2018). "Do cifar- 10 classifiers gen- eralize cifar- 10?" In: arXiv preprint arXiv: 0 (cit. p. 18).
Mihaylov, Todor et al. (2018). "Can Suit of Armor Con- duct Electricity? A New Dataset for Open Book Ques- tion Answering". In: Proceedings of the 2018 Confer- ence on Empirical Methods in Natural Language Pro- cessing (cit. on p. 19).
Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). "U-Net: Convolutional Networks for Biomedi- cal Image Segmentation". In: Medical Image Comput- ing and Computer-Assisted Intervention (MICCAI 2015) (cit. on p. 15).
Minka, Thomas (2000). Estimating a Dirichlet distribution (cit. on p. 21).
Ryabinin, Max, Andrey Malinin, and Mark Gales (2021). "Scaling ensemble distribution distillation to many classes with proxy targets". In: Advances Neural In-
12

---

Ensemble Distribution Distillation via Flow Matching
formation Processing Systems 34 (NeurIPS 2021) (cit. pp. 3, 9, 21, 22).
2021 Competitions and Demonstrations Truck (cit. on p. 5).
Salimans, Tim and Jonathan Ho (2022). "Progressive Dis- tillation for Fast Sampling of Diffusion Models". In: Proceedings of The 10th International Conference Learning Representations (ICLR 2022) (cit. on p. p) 3).
Wood, Danny et al. (2023). "A unitied theory of diversity in ensemble learning'. In: Journal of Machine Learning Research (JMLR) 24.359, pp. 1-49 (cit. p- 20).
Yao, Xufeng al. (2024). "Progressively Knowledge Dis- tillation via Re-parameterizing Diffusion Reverse Pro- cess". In: AAAI Conference on Artificial Intelligence (AAAI) (cit. on p. 3).
Shao, Shitong et al. (2024). "Precise Knowledge Transfer via Flow Matching". In: arXiv preprint arXiv:2402.02012 (cit. on 3).
Shen, Yuesong et al. (2024). "Variational learning is ef- fective for large deep networks". In: Proceedings of The 12th International Conference on Learning Representa- tions (ICLR 2024) (cit. on pp. 2, 5, 19).
Yin, Tianwei et al. (2024a). "Improved Distribution Match- ing Distillation for Fast Image Synthesis". In: Advances Neural Information Processing Systems 37 (NeurIPS 2024) (cit. on p. 3).
Song, Y. al. (2023). "Consistency Models". In: Proceed- ings The 40th International Conference on Machine Learning (ICML 2023) (cit. on p. 3).
Yin, Tianwei al. (2024b). "One-step Diffusion with Dis- tribution Matching Distillation". In: Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR 2024) (cit. on p. 3).
Stanton, Samuel al. (2021). "Does knowledge distilla- tion really work?" In: Advances in Neural Information Processing Systems 34 (NeurIPS 2021) (cit. p. 7).
Yun, EungGu et al. (2023). "Traversing between modes in function space for fast ensembling". In: Proceedings of The 40th International Conference on Machine Learning (ICML 2023) (cit. on pp. 3, 5).
Stark, Hannes et al. (2024). "Dirichlet flow matching with applications to DNA sequence design". In: Proceedings of The 41th International Conference on Machine Learn- ing (ICML 2024) (cit. on p. 14).
Zhang, Hongyi et al. (2018). "mixup: Beyond Empirical Risk Minimization". In: Proceedings of The 6th Inter- national Conference on Learning Representations (ICLR 2018) (cit. on p. 5).
Touvron, Hugo et al. (2023). "Llama Open founda- tion and fine-tuned chat models". In: arXiv preprint vrXiv:2307.09288 (cit. on p.p.5) 5).
Tran, Linh et al. (2020). "Hydra: Preserving ensemble diversity for model distillation". In: arXiv preprint arXiv:2001.0m 04694 (cit. p. 3).
Vaswani, Ashish et al. (2017). "Attention is all you need". In: Advances in Neural Information Processing Systems 30 (NIPS 2017) (cit. on p. 15).
Wen, Yeming, Dustin Tran, and Jimmy Ba (2020). "BatchEnsemble: an Alternative Approach to Efficient Ensemble and Lifelong Learning". In: Proceedings of The 8th International Conference on Learning Represen- tations (ICLR 2020) (cit. on p. 3).
Wilson, Andrew Gordon and Pavel Izmailov (2020).
"Bayesian deep learning and a probabilistic perspective of generalization". In: Advances in Neural Information Processing Systems 33 (NeurIPS 2020) (cit. on pp. 5, 7).
Wilson, Andrew Gordon et al. (2022). "Evaluating approx- imate inference in Bayesian deep learning".. In: NeurIPS
13

---

Ensemble Distribution Distillation via Flow Matching
A. Ablation Study
A.1. Flow matching formulation
Mapping Gaussian noise directly to teacher logits may initially appear overly simplistic, given that logits possess a specific geometric structure. Notably, teacher logits are invariant under scalar shifts, producing identical outputs after the softmax transformation. Flow matching formulations can be adapted to respect this geometry. Alternatively, since logits can be converted into probability vectors via the softmax operation, flow matching could be performed the probability simplex instead of the logit space. This section explores several such approaches.
Table
Performance of various flow matching formulations on CIFAR-1O.
Space
Formulation
ACC (1)
NLL (4)
ECE (4)
Statistical flow matching (Cheng al., 2024)
0.929
0.244
0.019
Probability simplex
Flow matching on the probability simplex
0.929
0.243
0.015
Flow matching with ILR
0.931
0.218
0.011
Logit space
EDM (Karras et al., 2022)
0.931
0.224
0.009
EDFM (Ours)
0.931
0.216
0.009
Flow matching on the probability simplex. Various formulations of flow matching on the probability simplex exist, such as CatFlow (Eijkelboom et al., 2024) and Dirichlet flow matching (Stark et al., 2024). However, most of these methods are inapplicable in our context as they require categorical (discrete) data, whereas we consider a continuous distribution of teacher logits defined either in logit space or equivalently on the probability simplex. Statistical flow matching (Cheng et al., 2024) is the most relevant method here, as it enables mapping from uniform distribution to arbitrary distribution defined on the probability simplex. Another straightforward approach is to perform basic flow matching directly on the probability simplex to transform a uniform distribution into the teacher logit distribution. This is feasible because the probability simplex is a Euclidean space where interpolations between points remain within the simplex.
Despite the theoretical appeal, empirical results shown in Table demonstrate that flow matching on the simplex yields inferior performance. We attribute this to the requirement that the student network match the teacher distribution with high precision on the simplex, whereas in logit space, minor mismatches are attenuated by the subsequent softmax operation.
Isometric log-ratio (ILR) transform. Given the shortcomings of operating directly on the probability simplex, one may consider improvements in the logit space. Since teacher logits are invariant under scalar shifts (e.g., logits {-1.5, 2, 2} are equivalent to 000.5, 31 post-softmax), enforcing exact matching of logits may be overly restrictive. The isometric log-ratio (ILR) transformation from compositional data analysis offers a principled way to address this issue: it maps the D-dimensional probability simplex to a (D 1)-dimensional Euclidean space. By applying ILR to teacher logits, flow matching can be trained in R- Nonetheless, empirical results (Table 6) indicate negligible differences compared the baseline, likely because the dimensionality reduction by one does not substantially affect the model's expressiveness in relatively high-dimensional settings.
Diffusion models. Aside from flow matching, diffusion models present an alternative framework to model the teacher logit distribution. Exploring such generative approaches for ensemble distribution distillation constitutes a promising future direction. We implemented the EDM model (Karras et al., 2022) to train the student network and compared it to our EDFM method. On the CIFAR- 10 dataset, EDM achieved an NLL of 0.224 using 35 NFEs, whereas EDFM attained a better NLL of 0.216 with only 5 NFEs (Table 6). Notably, EDM's 's performance deteriorated significantly with NFEs, achieving an NLL of 0.370, underscoring the superior efficiency and effectiveness of EDFM.
In summary, applying flow matching directly to pre-softmax logits yields superior results compared to modeling post- softmax categorical probabilities. This improvement may be explained by the distortion introduced by the softmax transfor- mation, which can obscure important scale information intrinsic to the logits. For example, the unnormalized density, given by the bog-sum-exp of logits (Duvenaud et al., 2020), may be lost after softmax, thereby negatively impacting ensemble distribution distillation.
14

---

Ensemble Distribution Distillation via Flow Matching
A.2. Network architecture of student model
We performed an ablation study to evaluate the impact of network architectures of varying scales on performance. The architectures considered include 1) the denoising MLP (Li et al., 2024) already employed in EDFM, 2) the U-Net (Ron- neberger et al., 2015), which is commonly used in flow matching for image synthesis and was also previously adopted in DBN (Kim et al., 2024), and 3) Transformer (Vaswani et al., 2017) models.
Table 7. Performance and efficiency of EDFM with various network architectures CIFAR-10 10.
Architecture
of Param.
Time (us)
ACC (1)
NLL ())
ECE (U)
MLP
0.14M
39
0.929
0.229
0.013
MLP (main text)
0.33M
68
0.931
0,216
0.009
MLP
0.70M
132
0,930
0,218
0.011
Transformer
0.36M
204
0.930
0,218
0.011
Transformer
1.31M
724
0.931
0,213
0.007
U-Net
0.44M
653
0.929
0.222
0.012
The results, summarized in Table 7, reveal that larger Transformers (with 1.31 million parameters) achieve improvements in NLL and ECE metrics but are not cost-effective, as reflected by the number of parameters and wall-clock time, where the latter denotes the time required to generate 256 ensemble predictions via batched inference. In contrast, the U-Net architecture underperforms relative both the MLP and Transformer, which likely attributable to its two-dimensional spatial processing being ill-suited for modeling logits. These findings suggest that the MLP architecture strikes a favorable balance, delivering sufficient accuracy for ensemble distribution distillation while offering significant latency advantage.
A.3. Design choice of EDFM
In this section, we elaborate on specific design choices of EDFM presented in Algorithms and 2.
Time distribution T. We swept uniform, beta, and exponential distributions to find the most suitable T, where the probability density functions are given in Table 8. The support of T is set to [c, 1] instead of [0. 1] to avoid numerical instability.
Table 8. PDFs of distributions.
Distribution
PDF (t)
Uniform
|r(1)= =1 fr(t;a,3) I(a ( t)8-1 fT(t;a) a' a'log a/(aa-r
Beta
Exponential
The PDF of the exponential distribution is slightly modified from the conventional one to account for the support [c. 1] and an arbitrary base. The choice of beta and exponential distributions was inspired by the well-known lesson in the field of diffusion models (Karras et al., 2022) that the network should learn more near the data. Specifically, the exponential dis- tribution allocates exponentially more samples
grows, while the beta distribution achieves a peak somewhere between t = () and t 1. We assigned the peak to be near the data by setting we set 3. The truncation was set to 0.001.
and = 2. For the exponential distribution,
Experiment results shown in Table 9 indicate that the beta distribution achieved better results in terms of ECE, while performing on par with the exponential distribution on other metrics. This motivates further explorations for even better options, possibly inspired by recent progress in flow matching (Lee et al., 2024; Kim et al., 2025). We leave this as future work.
15

---

Ensemble Distribution Distillation via Flow Matching
Table Ahlation study time distribution.
Dataset
Time distribution
ACC (1)
NLL (U)
ECE (4)
C10
Uniform
0.931
0.219
0.011
Beta
0.931
0.216
0.008
Exp. (Ours)
0.931
0.216
0.009
C100
Uniform
0.761
0.966
0.066
Beta
0.760
0,932
0.039
Exp. (Ours)
0.761
0.932
0.056
Timeedeeennentt weighting function X. The time-dependent weighting function was determined according to the preconditioning technique introduced in EDM (Karras et al., 2022). Specifically, in EDM, the authors first let the neural network predict the denoiser h(x,, t), namely x-prediction, and controlled the scale of the network input, output, and the weighting function of the loss so that the network input and the loss function maintain unit variance, and the variance of the network output minimized.
Table 10. Comparing between and v predictions for EDFM.
Dataset Prediction type
ACC (1)
NLL (L)
ECE (L)
C10
x-pred
0.931
0.234
0.024
v-pred (Ours)
0.931
0,216
0.009
C100
x-pred
0.760
1.031
0.090
v-pred (Ours)
0.761
0.932
0.056
However, as shown in Table 10, we found that directly predicting the velocity, namely v-prediction, performs better than x- prediction for EDFM. Thus, although we followed the principles suggested EDM, our preconditioning slightly deviates from the original.
For completeness, we present the derivation for our preconditioning. The student network accepts three inputs: perturbed sample =f, time and embedding from the teacher network Emb(x). Two of these three inputs rescaled to c (t)())7, Crime (t) before being fed into the network. The embedding Emb(x) is left invariant as it already maintains stable variance, after being processed by the teacher network.
We set Ciiee (f) = log (1000(1 t) + c)/4, where is a small number 10 12 to avoid numerical instability. This was inspired by the original formulation of log t/4, 4, and the transformation t - 1000(1 t) was introduced to match the scales.
Then we added skip connection the output of the student network Fg:
Do(2;) (t)i(()) + Caul Camn
(14)
which is then used to compute the loss function
E Be,ag.*f [Att) (2f =1])
(15)
Rearranging the terms results in
A(t) Comi(1)2
Er, .
(ICaspp() (1) 1)2 + ((1 - () (1) + 1)25 + Fo (c((((())))))
(16)
We want the network input, and loss function to maintain unit variance across different time 6, while the variance of the
16

---

Ensemble Distribution Distillation via Flow Matching
network output is minimized. This means that
Varz5 (cn((t)) (t)= Varzg oF (Cin (t)(ff (1 t)=
(17)
Cakip(t) = arg min Varzs .0 " (t) 1)2 + ((1 - -t)ckipp( (t) + 1)
(18)
(ty(()
Cour(t) Varzg ,ef [(ttakkp (t) - 1)zf + ((1 t) (t) + 1)26]
(19)
X(t) [tt)/Cou(t) Cour (t)² =1
(20)
Note that as z ~ 0, I)), Var(zø] = 0². Also, we Var(2]] = 0aa--
Evaluating for Gia Cia(t), Com (1), Cskap (1), A(t)
under this principle yields
Ga(()
(21)
v P + -1-12
00dala
(ou((1)
(22)
Vodda (1 0)
to -(1 -1- t)o²
(ttip())
(23)
Po + (1 - 1)202
σ² 0²0 data
X(1) Po data + (1 - t)
(24)
ODE solver and sampling step scheduling For ODE solver and sampling step scheduling, we ablated between Euler, Heun solvers (Karras et al., 2022) and uniform, exponential scheduling. For N steps of sampling, the sampling steps (i/N)/ for uniform scheduling and {(1 a²)/(1 ax) a')) for exponential scheduling. For exponential scheduling, we set = 0.7 for CIFAR-10 10 and = 0.5 for CIFAR- 100. As Fig. 5 indicate, Heun solvers with exponential scheduling achieved the best NLL at 5 sampling steps, which amounts to 7 NFEs for Heun solver. All experiment results presented in the main paper were obtained under this configuration.
Euler (Uniform)
- Heun (Uniform)
Euler (Exponential)
Heun (Exponential)
C10
C100
0.26
1.00
0.25
NLL
NLL
0.96
0.23
0.94
0.22
# steps
steps
Pigure 5. Ablation study ODE solver and sampling step scheduling.
17

---

Ensemble Distribution Distillation via Flow Matching
B. Experimental details
B.1. Image classification tasks
Datasets. Our experiments on image classification tasks utilize the following datasets:
CIFAR-10/100 (Krizhevsky and Hinton, 2009): It consists of 40,960 training images, 9,040 validation images, and 10,000 test images, each with a resolution of 32 x 32 3 across 10/100 classes. As there is no officially predefined validation split, we manually partitioned the 50,000 training examples into 40,960 for training and 9,040 for validation. The datasets are publicly available at
under unspecified license.
CINIC-10 (Darlow et al., 2018): It serves as an auxilairy train split consisting of 210,000 images with a resolution of 32 x 32 x 3. The dataset is publicly available at https//github..ommBBB Wathh//iii/// 10 under MIT license. The original dataset contains of a total of 270,000 images, from which we dropped 60,000 CIFAR- -oriinoore images to avoid dataset leakge in our CIFAR-10/100 experiments.
CIFAR-10.1 (Recht et al., 2018): It serves test split consisting of 2,000 images with resolution of 32 x 32 x 3. The dataset is publicly available at https://github.coo// http://itwiooooooooddds///cccaa- 10.1 under MIT license.
CIFAR-10.2 (Lu et al., 2020): It serves test split consisting of 2,000 images with resolution of 32 x 32 x 3. 10.2 under unspecified license.
The dataset is publicly available at
STL (Coates et al., 2011): It serves as a test split consisting of 7,200 images with resolution of 32 32 x 3. The dataset is publicly available at https://cs.staffrr..ee htt.fford.ed/aaccccteeee
under unspecified license. The original test split consists of 8,000 images with a resolution of 96 x 96 x 3, which we resized to 32 x 32 x 3. We excluded the 'monkey' class, as it is not present in CIFAR, removing 800 images in the process. Furthermore, the 'car' class was mapped to the "automobile" class to match CIFAR's class definitions.
SVHN (Netzer et al., 2011): It serves as an out-of-distribution split consisting of 73, 257 training and 26, 032 test images with colored number digit (0-9) with a resolution of 32 x 32 x 3 The dataset is publicly available at http://uddd.saaaooooo...c////ss/
under unspecified license.
Optimization details. We provide details on the construction of the ensemble teacher and the training procedure for distilled students in the image classification experiments.
Multi-SWAG teacher: Multi-SWAG teachers were obtained by adding diagonal covariances to each of 10 different SWA models. Each SWA model was pretrained for 800 epochs, followed by 200 epochs of SWA training with frequency 1. Using SGD optimizer with momentum 0.9, cosine decay schedule was applied for pretraining, in which the learning rate evolves from 0.1 to 0.01 and remains constant of 0.01 whilst SWA training. Batch size is 256.
KD: KD was trained using SGD optimizer with momentum 0.9, learning rate 0.1 with cosine decay schedule, 1000 epochs, batch size 256. Unlike CIFAR-10, 10, for models trained with CIFAR- learning rate was warmed up from 0 to 0.1 for the first 5 steps. Number of teachers to produce the target prediction is 30.
EnDD: EnDD was trained with almost identical configuration with that of KD, except that the learning rate is de- creased to 0.05.
FED: For FED, we strictly followed the configurations provided in the original paper (Penso et al., 2022), except that we increased the number of epochs from 200 to 250 for coherence with other methods. Training further than 250 epochs showed no improvements in performance.
DBN:
Like FED, we strictly followed the configurations provided in the original paper (Kim et al., 2024), except that we increased the number of epochs from 800 to 1000 for coherence with other methods.
We used the denoising MIP architecture with four blocks and hidden dimension of 256 for CIFAR-10 and 512 for CIFAR-100. The model employs a student network distilled with KD using the MulúSWAG teachers as its backbone and takes its output embedding as input. Evaluation was conducted with seven NFEs, and training was performed using the SGD optimizer with a batch size of 256, momentum of 0.9, weight decay of 5e-04, learning rates of le-04 for CIFAR- 10 and 3e-04 for CIFAR- 100, and cosine decay learning rate schedule over 1,000 epochs.
EDFM:
18

---

Ensemble Distribution Distillation via Flow Matching
B.2. Commonsense reasoning tasks
Datasets. Our experiments on commonsense reasoning tasks build upon the work of Cong al. (2024). Following their training and evaluation recipes, all datasets composed exclusively of multiple-choie questions with four options:
ARC-Challenge (ARC-C; Clark et al., 2018): It consists of 1,117 questions for training and 295 questions for evaluation, publicly available at
under CC-BY-SA-4.0 license.
ARC-Easy (ARC-E; Clark et al., 2018): It consists of 2,241 questions for training and 567 questions for evaluation, publicly available at
are under CC-BY-SA-4.0 license.
OpenBookQA (OBQA; Mihaylov al., 2018): It consists of 4,957 questions for training and 500 questions for under unspecified license.
evaluation, publicly available https//wwugggggfcccccooocassss//
Optimization details. We provide details on the construction of the ensemble teacher and the training procedure for distilled students in the commonsense reasoning experiments.
Multi-IVON teacher: We obtained the Multi-IVON teacher using the official codebase from Cong al. (2024), which is available at https//gituub.com//ea it involves performing LoRA fine-tuning the pretrained LLaMA-2-7B model, which is publicly available at https//hugginggface www.eetaoc/eaaeeeaaeeeeeeeeaaa coo/meta-llama///
http://ccammappoooooooo seaprox-bayes/ivone -oora under unspecified license. More specifically,
under the Meta license' using the IVON optimizer (Shen et al., 2024). Using their configurations, we trained 10 IVON models with different random seeds and constructed the Multi-IVON teacher, which generates a total 100 logits by extracting 10 logits from each individual model.
KD: We fine-tuned the pretrained LLaMA-2-7B weights using the KD objective, as defined in Eq. 37, with the Multi- IVON teacher. After several attempts, we found that effective distillation requires not only applying LoRA to the query and value weights of the attention layers but also training the classification head. As result, the number of trainable parameters is 4, 194,304, which is 0.062% of the total 6, 742,,09,,92 parameters. Moreover, the KD objective alone was insufficient for stable learning in our commonsense reasoning experiments. Therefore, in Hinton et al. (2014), we combined the KD objective with cross-entrope loss using ground-truie truth labels, weighted at 0.5. The fine-tuning was carried out using the Adam optimizer with a batch size of 4, maximum sequence length of 320, a learning rate of 3e-05, and linear decay learning rate schedule over 50,000 steps.
For flow modeling, we used the denoising MLP architecture with three blocks and a hidden dimension of 256, totaling (639,056 parameters. The model employs one ensemble component of the Multi-IVON teacher as its backbone and takes its output embedding as input. Evaluation was conducted with three NFEs, and training was performed using the SGD optimizer with a batch size of 256, momentum of 0.9, learning rate of 1e-04, and cosine decay learning rate schedule over 50,000 steps.
EDFM:
B.3. Evaluation metrics
Classification accuracy (ACC), is metric used to evaluate the performance a classification model by computing the
proportion of correct predictions. Let yn be the true label, (k) the predicted probability for class k for the nhhh sample, and N the total number of samples to be evaluated. We then compute ACC as follows, where [-] denotes the Iverson bracket:
1 N yn arg max Pa
ACC
(25)
Negative log-likelihood (NLL), also known as cross-entropy loss in classification tasks, is another important metric used for both training and evaluating classification models. It measures how well the predicted probability distribution matches the true labels. Let yn be the true label, ph: the predicted probability for class & for the nhhh sample, and N the total number of samples to be evaluated. We then compute NLL as follows:
N
NLL =
log pl (vm)
(26)
N
hhtpss//ai.meta.coom. https//aii.ntta.commmaamlllicee
19

---

Ensemble Distribution Distillation via Flow Matching
Expected calibration (ECE; Pakdaman Nacini et al., 2015), is metric that evaluates how well predicted probabilities correspond the true likelihood of correct predictions; well-calibrated model outputs probabilities that closely align with the observed frequency of correctness. Let yn the true label, pb the predicted probability class for the thhh sample, and the total number of samples be evaluated. The samples are partitioned into B bins based on the maximum predicted probability, max; pi where B; and [Ba] denote the indices and count of samples in the t bin, respectively. For I,.... B, the observed accuracy acc, and the average confidence conf, are computed as follows:
accb =
1
Σ n
arg max Ph: and conf, b
max (k) Σ n
(27)
|B8l
[Bb]
k
We then compute ECE as follows:
B
|B N |acccs conf.
ECE =
Σ b=1
(28)
Ensemble ambiguity (AMB), is metric used to measure the diversity of ensemble predictions. Let yn be the true label, (k)
(k) mm, and pm.. the logit and categorical probability for class predicted by the m ensemble member for the n1h sample, M the ensemble size, and N the total number of samples to be evaluated. We then compute AMB follows:
N
M
K
(R)
Pa
AMB
NM
n
log
(k)
(29)
==
m=
Pm
where P denotes normalized geometric mean of categorical probability here,
(K) Q
1/M pmnn
with
II
for h: 1,.... K.
(30)
Dj=1
K
4a
We also have the following generalized ambiguity decomposition (Wood et al., 2023), which suggests that to minimize ensemble loss, ensemble ambiguity should increase while average loss decreases:
log softmax((n) (3n)
log softmax (2m,,)) -AMB.
2m,
M
M
N
(31)
ensemblee loss (EnnsLass)
averaage loss (Avgl.oss)
In Fig. compute the ambiguity decomposition after applying temperature scaling (Guo et al., 2017), using the op- timal temperature obtained by minimizing EnsNLL. This approach offers more accurate comparison, EnsNLL uses ensembled logits instead of ensembled probabilities. When logit shift occurs for each ensemble member, the probability ensemble remains unchanged, while the logit ensemble can vary.
(R) Ensemble variance (VAR), is another metric used measure the diversity of ensemble predictions. Let pm, be the categorical probability for class & predicted by the mh ensemble member for the 1h sample, M the ensemble size, and N the total number of samples to be evaluated. We then compute VAR as follows:
K
VAR
Variance m I.....
p
(32)
We also have the following uncertainty decomposition (Abe et al., 2022), which suggests that the variance quantifies the extent which predictive uncertainty increases after ensembling:
M ΣE
p
Σ
-
Σ
p(t
+VAR.
N m=
NM
(33)
1=
1m ma
1=1
ensemble uncertainty (EnsUnc)
average uncertainty (AvgUne)
20

---

Ensemble Distribution Distillation via Flow Matching
Agreement (AGR), is commonly used metric in classification tasks measure how often two predictions match each (k) be the two categorical probabilities for class & predicted for the n²th sample, and the total number
other. Let ph and qn of samples to be evaluated. We then compute AGR as follows:
N
AGR =
Σ
arg max p/ = arg max (8) &
(34)
N
k
Total variation distance (TVD), is a metric that measures the difference between two probability distributions. It quan- tifies the maximum discrepancy in probability mass between them. pn and (k) be the two categorical probabilities for class k predicted for the nhth sample. and N the total number of samples to be evaluated. We then compute TVD as follows:
TVD
(R) Pn
4n
(k)
(35)
2N
W distance (W2), measures the distance between two probability distributions by quantifying the optimal cost of transporting one distribution to another. Fréchet Inception Distance (FID) (Heusel et al., 2017), a widely used metric in generative models literature, is also specific instance of W2. Inspired by this, we also evaluated the fidelity of generated logits to teacher logits by measuring Wasserstein-2 distance between the two distributions. Let P and Q be two discrete probability distributions with support {x1,..., In. and (yI,..., ym respectively, so that P = P. p:O,; and Q Σド= y5ōy, where Ps zj are probability masses and 6₂ Dirac delta function. We then compute Wasserstein-2
distance between two discrete distributions
W)(PQQ)
inf (ET(P))) i=1
Σ
R. || yzll?
(36)
where "N.j is the transport plan meaning how much mass from to q, I(P, is the set of all possible valid transport Ki. = Pi, Y. = 9:- Evaluating this involves linear programming, and can
plans that satisfy the marginal constraint Σ, efficiently be solved using Hungarian algorithm. Note that W2 distance differs from TVD, in that it measures the similarity between the whole predictive distributions, where as TVD only compares their means.
B.4. Ensemble distillation methods
Knowledge distillation (KD; Hinton et al., 2014) is considered the simplest method for transferring knowledge from a large and complex teacher model smaller student model. The training objective for knowledge distillation is given as follows:
M
(kD(Ø) = H (p. .
1 M m==1
Σ H(P,. 9a):
(37)
which simply is the cross-entropy loss widely used in classification tasks. One thing that differs is, that it is typical to soften (or harden) the teacher logits by dividing them with temperature T, for better performance.
Ensemble distribution distillation (EnDD; Malinin et al., 2020; Ryabinin et al., 2021) assumes that the categorical predictions from the ensemble teacher follow a Dirichlet distribution. In other words, the induced distribution in Eq. 4 is assumed to be Dirichlet, i.e., PI..... PM ~ ((Tx(x)) = Dir(B), where the concentration parameter /3 is estimated in closed-form using the approximate maximum likelihood procedure (Minka, 2000):
K-1
((K) (k)
2 ΣK p(i) (log p() -log ()) for 1, 1,... K,
(38)
where p- MI ΣM= P.. and logp- 1 Drichlet Prior Network (DPN; Malinin and Gales, 2018), which models g (TT (x) = Dir(ar)), can be trained by minimiz- ing the KL divergence between two Dirichlet distributions Dir( and Dir (ir()))) While the seminal work of Malinin et al. (2020) originally minimized the forward KI. divergence, DKL (Dir(3) || Dir (or))) subsequent research by Ryabinin al.
EM log p. are computed over M observations Pi..... Par- Then, student
M
21

---

Ensemble Distribution Distillation via Flow Matching
(2021) showed that minimizing the reverse KL divergence, DKL (Dir(aa))) Dir(a)), improves training stability. They also constrained the concentration parameters to be greater than one by d + a + and 3 + 3 + 1, to further enhance stability. Accordingly, we adopted the loss function proposed by Ryabinin et (2021) to train the EnDD baseline:
1
EnD (Ø) = (ir(()( (P)
DR.. (Dir(ās) || Dir(1))-
Ek=1 (1)) (R)
(39)
reconstruction term
prior term
During test time, we compute the categorical probabilities the mean of the student Dirichlet distribution, which is given by (ir()i(()))) [q()] = - Σ 2 for = I.... K. Notably, this is equivalent to applying the softmax operation to the DPN output, making the inference procedure identical to standard softmax-based classification neural networks.
Functional ensemble distillation (FED; Penso et al., 2022) aims to distill the whole distribution of predictions from the ensemble teacher, unlike previous approaches that are capable only of predicting the mean of the teacher predictions. Specifically, it proposed to minimize the maximum mean discrepancy between the predictions produced by the teachers and predictions produced by student. Enabling the student network to be generative model that approximates the ensemble teacher distribution, FED has advantages over tasks that require covariance between predictions. As EDFM also proposes to distill ensemble teacher distribution via flow matching, FED is the most alike approach to ours among other baselines.
Diffusion bridge network (DBN; Kim et al., 2024) adopted SB image reconstruction algorithm based on Diffusion Schrödinger Bridge map predictions of single teacher model the mean of predictions of ensemble teachers. Although DBN proposed a fast ensembling algorithm based on concurrent generative model, it fundamentally deviates from EDFM in that it does not distill the whole ensemble distribution.
22


Reviews:
[{"review_id": "EADQUgoG4n", "reviewer": "Reviewer_37tb", "summary": "This paper proposes a novel ensemble distribution distillation method (EDFM) that utilizes flow matching to efficiently transfer the diversity of ensembled teacher models to a smaller student model. Key challenges in ensemble distribution distillation are addressed, including the high computational cost of large integrations and the difficulty of capturing the full diversity of ensemble predictions due to the capacity constraints of student models. The proposed method introduces a lightweight network that learns to map individual model predictions to a vector field of integrated teacher prediction distributions, enabling fast and parallelizable inference. Extensive experiments demonstrate the effectiveness of EDFM compared to existing ensemble distribution distillation methods.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review and clear understanding of our work, particularly in terms of diversity, efficiency, and scalability. We are glad that you appreciated our extensive experimental results and recognized our approach as both a novel framework and a conceptual innovation. We hope our responses address any remaining concerns, and please let us know if you have any further questions.\n\n> the methods used in the comparisons are quite old.\n>\n> the Abstract and Methods are quite simple in writing.\n\nWe reviewed prior research on ensemble distribution distillation, from the EnDD in the seminal work of Malinin et al. (2020), which was later improved by Ryabinin et al. (2021), and the most recent FED by Penso et al. (2022) as baselines. We would greatly appreciate it if you could share further related recent methods in ensemble distribution distillation. Additionally, we will make sure to expand the abstract and method sections in the final version.\n\n> Essential References Not Discussed\n\nThank you for sharing the relevant prior works that we may have overlooked. Shao et al., (2024) enhances knowledge distillation using flow matching, and DiffKD (Huang et al., (2023)), or KDiffusion (Yao et al., (2024)) proposed similar methods based on diffusion models. However, they are more similar to DBN (Kim et al., (2024)) rather than EDFM, as they align logits or features between teacher and student models one-to-one. Yin et al. (NeurIPS 2024, CVPR 2024) propose distribution matching distillation to distill diffusion models into one-step generators for faster sampling. As such, their goal and methodology differ from EDFM. Nonetheless, they offer valuable insights in the general context of knowledge distillation. We will include these connections in the camera-ready version.\n\n> The submission does not contain explicit theoretical proofs that require verification.\n\nAs you mentioned, the main contribution of our work lies in introducing flow matching as a novel framework for ensemble distribution distillation, demonstrated through extensive experiments. Since our approach aligns with the standard flow matching scenario, we believe the theoretical claims from the existing flow matching literature apply to our work as well, without significant challenges (e.g., minimizing the conditional flow matching loss in Eq. (13) ensures the approximation of the ensemble teacher's predictive distribution $p_{1}$).\n\n> unique improvements over existing flow matching approaches (such as the design of conditional flow matching or efficient inference strategies)\n\nWe believe our work is significant not only for introducing a pioneering and promising direction in ensemble distribution distillation, but also for empirically validating its effectiveness through extensive experiments, as recognized by the reviewers. Also, it demonstrates strong practicality by being highly efficient in terms of inference cost with the proposed parallelized inference strategy.\n\nTo further solidify the position of our work as a pioneering study, we provide key takeaways within the ensemble distribution distillation context.  Comprehensive ablation studies on design choices in flow matching for effective ensemble distillation help clarify this aspect. In response, we conducted additional ablations, including those requested by the reviewers, which further strengthen the contribution of our work.\n\n__Sampling algorithm.__\nWe conducted an ablation study on the sampling algorithm, focusing on the ODE-solver, time step scheduler, and the number of steps. (With the exponential scheduler, more steps are populated near $t=1$.) The plot (https://imgur.com/a/fCd15ca) shows the results, where the filled marker denotes the settings used in the paper.\n\n__Noise distribution.__\nWe conducted an ablation study on the noise distribution during training, considering three distributions: (1) uniform, (2) exponential, and (3) beta. The choice of exponential and beta distributions was inspired by the well-known lesson that the model should learn more near the data. We found that the beta distribution achieved better results in terms of ECE.\n\nData|Distribution|ACC / NLL / ECE\n-|-|-\nC10|Uniform|__0.931__ / 0.219 / 0.011\n||Exp. (Ours)|__0.931__ / __0.216__ / 0.009\n||Beta|__0.931__ / __0.216__ / __0.008__\nC100|Uniform|__0.761__ / 0.966 / 0.066\n||Exp. (Ours)|__0.761__ / __0.932__ / 0.056\n||Beta|0.760 / __0.932__ / __0.039__\n\n__Type and formulation of flow matching loss.__\nWe ablated the flow matching loss by 1) comparing x-prediction and v-prediction (ours), and 2) comparing categorical probabilities and logits (ours). Please see our response to Reviewers AwX6 and r2cZ.\n\n__Architecture.__\nWe ablated the network architectures (MLP, Transformer, and U-Net) at varying scales. We found that the current MLP architecture strikes the best cost-performance trade-off; please see our response to Reviewer r2cZ.", "claims_and_evidence": "The claims made in the submission are supported by extensive experimental evidence, which is presented in a clear and detailed manner. The authors provide a comprehensive set of experiments across various tasks, including image classification and commonsense reasoning, to validate their proposed method, Ensemble Distribution Distillation via Flow Matching (EDFM). Here are some key points that support the claims:\n1. The authors demonstrate that EDFM outperforms existing ensemble distillation methods in terms of accuracy, negative log-likelihood (NLL), and expected calibration error (ECE) across multiple datasets (CIFAR-10, CIFAR-100, ARC-C, ARC-E, OBQA). \n2. The authors highlight the efficiency of EDFM in terms of runtime and computational cost. Experiment shows that EDFM scales well with the number of ensemble predictions, maintaining low execution times while improving performance.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem and application at hand. The authors address the challenge of ensemble distillation, which aims to transfer the knowledge from a computationally expensive ensemble of models to a single, more efficient student model. The proposed method, Ensemble Distribution Distillation via Flow Matching (EDFM), is designed to effectively capture the diversity of ensemble predictions and transfer this knowledge to the student model.", "theoretical_claims": "The submission does not contain explicit theoretical proofs that require verification. Instead, it focuses on empirical validation through extensive experiments. The authors do make several theoretical claims and provide justifications for their approach. I think it would be better to add a theoretical proof to the article.", "experimental_designs_or_analyses": "I reviewed the soundness and validity of the experimental designs and analyses presented in the submission. Overall, the experimental design is robust and well-structured, with appropriate benchmarks, metrics, and comparisons to validate the proposed method. However, the methods used in the comparisons are quite old.", "relation_to_broader_scientific_literature": "The key contributions of the paper are deeply rooted in and extend the broader scientific literature in ensemble learning, knowledge distillation, generative modeling, uncertainty estimation, and scalability to large models. By introducing flow matching as a novel framework for ensemble distillation, the paper addresses several limitations of prior work, including the preservation of diversity, efficient sampling, and scalability to large models. The experimental results are also ok.", "essential_references_not_discussed": "[1]  Improved distribution matching distillation for fast image synthesis[J]. Advances in Neural Information Processing Systems, 2024, 37: 47455-47487.\n[2]  Knowledge Distillation via Flow Matching[J].\n[3]  One-step diffusion with distribution matching distillation[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2024: 6613-6623.", "other_strengths_and_weaknesses": "This method presents some conceptual innovations, particularly in applying flow matching techniques to ensemble distillation tasks. However, flow matching itself is an existing technique, and the main contribution of this work lies more in integrating it with ensemble distillation rather than introducing theoretical or algorithmic breakthroughs. The authors should more clearly articulate their method's unique improvements over existing flow matching approaches (such as the design of conditional flow matching or efficient inference strategies) to better highlight its novelty.\n\nFurthermore, the Abstract and Methods are quite simple in writing. And all compared methods are very old.", "other_comments_or_suggestions": "No.", "questions_for_authors": "NO.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AiTBP0mPSo", "reviewer": "Reviewer_AwX6", "summary": "The paper presents an ensemble distillation method based on flow matching named EDFM. The core idea is to learn a mapping between Gaussian noise and the logits of a (Bayesian) teacher model conditioned on the input data. The authors first analyze the importance of diversity in the predictions of the teacher when learning the student model by incorporating several previously established methods such as Mixup augmentations and auxiliary datasets. Then, the authors examine 3 aspects of their approach compared to other distillation methods (ones that attempt to preserve the ensemble diversity and ones that don't), (1) the ability to mimic the ensemble teacher models, (2) the diversity in the students prediction, and (3) the generalization capabilities of the student. The proposed approach was examined using both image and text data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the positive feedback highlighting our paper’s clarity and coherence. We hope our responses address any remaining concerns, and please reach out if you have any further questions.\n\n> the paper is a combination of previous ideas in the distillation literature\n\nWe respectfully differ in our interpretation of the novelty claim and would like to reference Michael Black’s remark: \"The novelty, however, must be evaluated before the idea existed.\" To the best of our knowledge, this is the first work to introduce flow matching as a novel framework for ensemble distribution distillation. Our work is a pioneering contribution that, through extensive experiments, demonstrates how flow matching can be both efficient and effective across three key aspects—diversity, fidelity, and generalization—in the context of ensemble distribution distillation. We believe this represents a clear and novel contribution to the community.\n\n> The decision to model the logits as Gaussian random variables needs further justification.\n\nThank you for your insightful feedback on our Gaussian modeling of logits. We, too, initially explored combining logit geometry with flow matching. However, after preliminary investigation, we found \"simple is best\"—alternative formulations underperformed the simplest one. Notably, prior works such as Yun et al. (2023) and Kim et al. (2024) also predict logits instead of categorical probabilities. Below, we summarize results of revisiting this comparison under the current setup:\n\nData|Space|Distribution|ACC / NLL / ECE\n-|-|-|-\nC10|Probability simplex|SFM|0.929 / 0.244 / 0.019\n|||FM on simplex|0.929 / 0.243 / 0.015\n||Logit space|FM w/ ILR|__0.931__ / 0.218 / 0.011\n|||FM (Ours)|__0.931__ / __0.216__ / __0.009__\n\n- __Statistical Flow Matching (SFM).__ Many flow matching formulations on the probability simplex, including CatFlow (Eijkelboom et al., 2024), are inapplicable here as they assume discrete data, whereas we handle a \"continuous\" logit distribution. SFM (Cheng et al., 2025) is the most relevant to our situation in this line of work, which provides a method to map the uniform distribution into arbitrary simplex distributions.\n\n- __Flow Matching on Probability Simplex (FM on simplex).__ Alternatively, one can apply flow matching directly on the simplex, mapping the uniform distribution to the ensemble distribution. This is valid as the simplex itself is a Euclidean space.\n\n- __Isometric Log-Ratio Transform (FM w/ ILR).__ As you have sharply pointed out, teacher logits are invariant under scalar shifts, yielding identical softmax outputs. Thus, forcing the student to match teacher logits “exactly” may be overly strict. The ILR transform addresses this by mapping the probability simplex of dimension $D$ into a $D-1$-dimensional vector space, avoiding redundancy.\n\nOverall, applying flow matching directly to pre-softmax logits produces better results than modeling it on post-softmax categorical probabilities. One possible explanation is that the softmax operation distorts the informative scale of the logits, negatively impacting ensemble distribution distillation. For example, unnormalized density information, calculated as the log-sum-exp of logits (Grathwohl et al., 2020), may be lost due to the softmax transformation.\n\n> Will EDFM generalization improve by incorporating a bigger student model? … Perhaps a different design choice for the flow model can help as well.\n\nDuring the rebuttal period, we conducted an ablation study on network architectures (MLP, Transformer, and U-Net) of varying scales (please see our comment to reviewer r2cZ for detailed results). The result demonstrates that the performance saturates with relatively small students, and the model used in the paper is already enough.\n\n> I wonder how EDFM outperforms the teacher model by a large margin on text data.\n\nThis relates to two points: 1) our EDFM adopting a pre-trained teacher network, and 2) a slight decrease in accuracy for Multi-IVON compared to IVON. Specifically, for language tasks, EDFM is trained using the frozen, pre-trained IVON@mean model, which we found to outperform Multi-IVON in accuracy, consistent with Cong et al. (2024). As a result, our method improves upon the IVON@mean baseline by effectively distilling the Multi-IVON teacher (slightly in accuracy, and significantly in NLL, ECE), thereby surpassing the teacher by a wide margin in accuracy, similarly to the IVON@mean model.\n\nMethod|ARC-C|ARC-E|OBQA\n-|-|-|-\nIVON@mean|__0.710__ / 1.874 / 0.246|__0.888__ / 0.645 / 0.094|__0.809__ / 0.653 / 0.111\nMulti-IVON|0.675 / __0.891__ / __0.091__|0.868 / __0.357__ / __0.026__|0.794 / __0.512__ / __0.032__\n||\nEDFM (Ours)|0.715 / 1.113 / 0.142|0.892 / 0.390 / 0.044|0.818 / 0.537 / 0.046\n\n> Minor\n\nWe sincerely appreciate your thorough review (yes, $z_{t}^{x}$ is correct); we will make sure to adjust the mistakes pointed out in the camera-ready version.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Not relevant.", "experimental_designs_or_analyses": "The experimental design is solid.", "relation_to_broader_scientific_literature": "Relation to broader scientific literature is good.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is written in a clear and cohesive manner. The importance of ensemble diversity is first analyzed and then each aspect of the distillation goals is inspected with proper empirical evidence.\n- In terms of prediction diversity EDFM outperforms all baseline methods in most cases. In addition compared to FED, the main competing approach, EDFM tends to exhibit better performance.\n- The proposed approach is efficient in terms of the storage required and inference time when using a proper student network.\n\nWeaknesses/Questions:\n- In terms of novelty, if I understand correctly, the paper is a combination of previous ideas in the distillation literature (FED and DBN), namely using generative models for distillation and capturing diversity in the teacher model.\n- The decision to model the logits as Gaussian random variables needs further justification. The probabilities after the Softmax layer are invariant to transformations of the logits such as shifting which is not consistent with the Gaussian assumption. Perhaps it will be more appropriate to use a different approach such CatFlow [1] which learns a flow directly in the probability space.\n- In terms of accuracy, EDFM does not generalize as well as EnDD and DBN on in-distribution and distribution shift tasks. Will EDFM generalization improve by incorporating a bigger student model? A further analysis may be required here. Perhaps a different design choice for the flow model (such as the one described in the previous bullet) can help as well.\n- I wonder how EDFM outperforms the teacher model by a large margin on text data. Contrasting it with image data, there the teacher model is usually better by a large margin (as also witnessed in the literature). I believe that further clarification and investigation is required here.\n- Minor:\n  - Line 199 left column - did you mean $z_{t}^{x}$?\n  - Missing citation of rectified flow paper in line 165 right column.\n  - The meaning of the abbreviation EDFM is not defined.\n  - Missing reference and discussion to Table 2 in the main text.\n\n[1] Eijkelboom, F., Bartosh, G., Andersson Naesseth, C., Welling, M., & van de Meent, J. W. (2024). Variational flow matching for graph generation. Advances in Neural Information Processing Systems, 37, 11735-11764.", "other_comments_or_suggestions": "See Strengths And Weaknesses.", "questions_for_authors": "See Strengths And Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g2BV9pWyjF", "reviewer": "Reviewer_r2cZ", "summary": "The paper presents an ensemble distribution distillation method leveraging flow matching to efficiently transfer knowledge from an ensemble teacher to a smaller student model. \n\nA new approach that models ensemble distribution distillation using flow matching, enabling student models to better capture the diversity of teacher ensembles.\n\nA theoretical formulation where flow matching is applied in logit space, optimizing a vector field to map student model predictions to the ensemble teacher’s distribution.\n\nExtensive experimental validation on image classification (CIFAR-10/100), commonsense reasoning (ARC, OBQA), and robustness evaluations under distribution shifts (CIFAR-10.1, STL, SVHN).\n\nDemonstration that the proposed EDFM (Ensemble Distribution Flow Matching) method outperforms previous ensemble distillation approaches, improving accuracy, calibration, and efficiency while reducing computational costs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and for recognizing our work as a new and promising direction. We are pleased that you found our extensive experimental results valuable and especially appreciate your recognition of our fidelity and diversity analyses, which are central to our contribution. We hope our responses address any remaining concerns—please reach out if you have any further questions.\n\n> Some key comparisons to prior work (e.g., Zhang et al. 2023, Jiang 2023)...\n\nWe could not find Zhang et al. (2023) and Jiang (2023) in our manuscript. Could you kindly clarify which sections you are referring to?\n\n> The connection to recent advances in diffusion-based distillation...\n\nIn the context of ensemble distribution distillation, DBN (Kim et al., 2024) represents a recent advance. While we discuss its relevance and differences in Section 3 (Related Work) and Appendix B.2 (Ensemble Distillation Methods), we will further clarify this in the camera-ready version.\n\n> The paper does not address non-Bayesian perspectives...\n\nCould you kindly clarify what is meant by “non-Bayesian perspectives”? If we understand correctly, point-estimating the mean of the ensemble teacher prediction would fall under non-Bayesian approaches, while ensemble distribution distillation, which models the distribution of the ensemble teacher prediction, would be considered Bayesian. In this case, KD and DBN would correspond to the former, and EnD, FED, and our EDFM to the latter. We would appreciate any further insights.\n\n> How does EDFM perform when using fewer flow matching steps...\n\nIn the runtime analysis (Figure 4, Section 5.4), we compared NFE values of {3, 5, 7} (corresponding to 3, 4, 5 steps) and selected NFE=7 for the main tables based on the performance-cost trade-off. As shown, fewer flow matching steps (NFE=3) led to slight performance drops, though NFE=5 showed no significant degradation. We agree that further ablations on the number of steps, solver, and scheduler choices would strengthen the paper.\n\nAs a follow-up, we conducted an ablation study on: 1) uniform vs. exponential schedules and 2) Euler vs. Heun methods across different step counts. The figure (https://imgur.com/a/fCd15ca) summarizes the results; filled markers indicate main table settings (Heun with exponential schedule, NFE=5). This underscores the value of carefully designing the sampling procedure with performance-cost in mind—as we have already addressed. Notably, the exponential schedule, rarely used in flow matching literature, was uniquely proposed and tailored for EDFM, proving critical in our distillation setup. We will include a separate section on hyperparameter sensitivity in the camera-ready version.\n\n> How does EDFM compare to alternative generative modeling techniques...\n\nExploring alternative generative methods is an interesting future direction. We implemented EDM (Karras et al., 2022) to train the student network and compared it with our EDFM approach. In our CIFAR-10 setup, EDM achieved an NLL of 0.224 with 35 NFEs, while EDFM reached 0.216 with only 5 NFEs; EDM was unable to achieve reasonable performance with just 5 NFEs (it achieved NLL of 0.370). It clearly shows effectiveness and efficiency of our EDFM approach.\n\n> How does the choice of the vector field parameterization...\n> \n> ...optimization difficulties when used with transformers?\n\nDuring the rebuttal period, we conducted an ablation study on network architectures (MLP, Transformer, and U-Net) of varying scales. The results are as follows: 1) Larger Transformers (1.31M parameters) improve NLL and ECE but are not cost-effective, as evidenced by the “#Params” and “Wall-clock time” columns (where wall-clock time refers to the duration required for 256 ensemble predictions via batched inference). 2) U-Net underperforms compared to both MLP and Transformer, likely due to its 2D spatial processing being less suited for handling logits. These findings indicate that the MLP architecture is sufficient for ensemble distribution distillation, providing a latency advantage. Furthermore, no issues related to training instability were observed, even with large student networks or Transformers.\n\nData|Arch|# Param|ACC / NLL / ECE|Wall-clock time(microseconds)\n-|-|-|-|-\nC10|MLP|0.14M|0.929 / 0.229 / 0.013|__39__\n|||0.33M (Ours)|__0.931__ / 0.216 / 0.009|68\n|||0.70M|0.930 / 0.218 / 0.011|132\n||Transformer|0.36M|0.930 / 0.218 / 0.011|204\n|||1.31M|__0.931__ / __0.213__ / __0.007__|724\n||U-Net|0.44M|0.929 / 0.222 / 0.012|653\n\n> type of loss\n\nWe also conducted an ablation study on the type of flow matching loss, comparing x-prediction and v-prediction, to find that the latter consistently outperforms the former. For other formulations of flow matching, please refer to our comment to review AwX6.\n\nData|Type|ACC / NLL / ECE\n-|-|-\nC10|x-pred|__0.931__ / 0.234 / 0.024\n||v-pred (Ours)|__0.931__ / __0.216__ / __0.009__\nC100|x-pred|0.760 / 1.031 / 0.090\n||v-pred (Ours)|__0.761__ / __0.932__ / __0.056__", "claims_and_evidence": "The authors claim that flow matching enables better ensemble distribution distillation by capturing ensemble diversity more effectively than previous methods. Supported by empirical results in diversity, fidelity, and robustness experiments that show EDFM consistently outperforming baselines.\n\nThe paper claims that EDFM improves both efficiency and predictive performance compared to existing distillation techniques.\nSupported by runtime analysis showing that EDFM achieves strong scalability and efficiency with minimal computational overhead.\n\nIt asserts that EDFM preserves ensemble teacher properties better than competing methods. Supported by fidelity analysis, where EDFM achieves the closest alignment with the ensemble teacher in various distance metrics (TVD, KLD, JSD).\n\nThe paper claims that flow matching naturally captures ensemble diversity, but it does not include direct theoretical justification for why flow matching should perform better in this setting compared to prior distillation methods.", "methods_and_evaluation_criteria": "The proposed flow matching approach is well-motivated and clearly formulated.\n\nThe evaluation methodology is strong, using multiple baselines (e.g., KD, EnDD, FED, DBN) and comprehensive metrics (accuracy, calibration, fidelity, efficiency).\n\nThe benchmarks (CIFAR, ARC, OBQA, STL, SVHN) are appropriate for the task.\n\nThe paper does not provide ablation studies that analyze how different design choices (e.g., type of flow matching loss, student model architecture) affect performance.", "theoretical_claims": "The paper provides a theoretical formulation of flow matching for ensemble distillation.\n\nThe derivation of the vector field for flow matching in logit space is consistent with prior work on normalizing flows.\n\nTheoretical justification for why EDFM better captures ensemble diversity is not fully developed.\n\nThe claims about exponential convergence of flow matching (inspired by prior generative modeling work) are not rigorously derived for the ensemble distillation setting.", "experimental_designs_or_analyses": "The experiments are comprehensive, covering different dataset shifts, evaluation metrics, and baselines.\n\nThe fidelity and diversity analyses are particularly strong, demonstrating how well the student model mimics ensemble predictions.\n\nAblation experiments (e.g., varying the number of flow matching steps, testing different vector field parameterizations) are missing.\n\nThe paper could include more real-world benchmarks beyond academic datasets to validate its practical effectiveness.", "relation_to_broader_scientific_literature": "The paper effectively connects ensemble distillation to generative flow matching methods.\n\nIt situates EDFM within the broader context of knowledge distillation and uncertainty estimation.\n\nThe discussion of ensemble teacher diversity is well-aligned with prior work in Bayesian deep learning.\n\nThe connection to recent advances in diffusion-based distillation (e.g., Schrödinger bridge methods) could be better explored.\n\nThe paper does not address non-Bayesian perspectives on ensemble distillation, which could be valuable for robustness.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n\nNovelty: The paper introduces flow matching into ensemble distillation, which is a new and promising direction.\n\nStrong empirical results: EDFM consistently outperforms existing methods across multiple tasks and metrics.\n\nEfficiency: The method is computationally efficient, reducing the overhead of traditional ensemble models.\n\nComprehensive evaluation: The paper rigorously tests EDFM under distribution shifts and real-world tasks.\n\nWeaknesses\n\nLimited theoretical justification: The paper lacks strong formal arguments for why flow matching should improve ensemble distillation.\n\nNo ablation studies: The sensitivity of EDFM to hyperparameters is not analyzed.\n\nComparison to recent generative modeling approaches (e.g., score-based methods) is missing.\n\nSome experimental design choices (e.g., why a specific type of flow matching loss is used) are not well explained.", "other_comments_or_suggestions": "The writing could be clearer in some sections, especially in theoretical explanations (e.g., the connection between flow matching and Bayesian ensembles).\n\nAblation studies (e.g., varying flow matching step size, different vector field architectures) would strengthen the empirical claims.\n\nA more explicit comparison to diffusion-based distillation methods would be useful.", "questions_for_authors": "How does EDFM perform when using fewer flow matching steps? Would reducing the number of steps significantly degrade performance?\n\nHow does the choice of the vector field parameterization affect distillation quality? Have you experimented with different neural network architectures for the vector field?\n\nHow does EDFM compare to alternative generative modeling techniques for distillation? Could score-based diffusion models be used instead of flow matching?\n\nWhat are the limitations of EDFM when scaling to very large student models? Does the method suffer from optimization difficulties when used with transformers?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents conceptual innovations, in the application of flow matching techniques to ensemble distillation, which is a new and promising direction. Strong empirical results are presented showing that EDFM consistently outperforms existing methods across multiple tasks and metrics. The method is computationally efficient, reducing the overhead of traditional ensemble models.
Through extensive experiments, the paper demonstrates how flow matching can be both efficient and effective across three key aspects—diversity, fidelity, and generalization—in the context of ensemble distribution distillation.

The claims made in the submission are supported by extensive experimental evidence, which is presented in a clear and detailed manner. The authors provide a comprehensive set of experiments across various tasks, including image classification and commonsense reasoning, to validate their proposed method. Comprehensive ablation studies are provided on design choices in flow matching for effective ensemble distillation.

The submission does not contain explicit theoretical proofs that require verification. Instead, it focuses on empirical validation through extensive experiments. The paper lacks strong formal arguments for why flow matching should improve ensemble distillation. The authors make several theoretical claims and provide justifications for their approach, but the paper could benefit from added theoretical proofs. For example, theoretical justification for why EDFM better captures ensemble diversity is not fully developed, and the claims about exponential convergence of flow matching (inspired by prior generative modeling work) are not rigorously derived for the ensemble distillation setting.
Decision: Accept (poster)