Ensemble Distribution Distillation via Flow Matching

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.

Some key comparisons to prior work (e.g., Zhang et al. 2023, Jiang 2023)...

We could not find Zhang et al. (2023) and Jiang (2023) in our manuscript. Could you kindly clarify which sections you are referring to?

The connection to recent advances in diffusion-based distillation...

In 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.

The paper does not address non-Bayesian perspectives...

Could 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.

How does EDFM perform when using fewer flow matching steps...

In 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.

As 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.

How does EDFM compare to alternative generative modeling techniques...

Exploring 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.

How does the choice of the vector field parameterization...

...optimization difficulties when used with transformers?

During 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.

type of loss

We 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.

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.

the paper is a combination of previous ideas in the distillation literature

We 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.

The decision to model the logits as Gaussian random variables needs further justification.

Thank 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:

• 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.

• 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.

• 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.

Overall, 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.

Will EDFM generalization improve by incorporating a bigger student model? … Perhaps a different design choice for the flow model can help as well.

During 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.

I wonder how EDFM outperforms the teacher model by a large margin on text data.

This 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.

Minor

We 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.

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.

the methods used in the comparisons are quite old.

the Abstract and Methods are quite simple in writing.

We 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.

Essential References Not Discussed

Thank 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.

The submission does not contain explicit theoretical proofs that require verification.

As 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}$).

unique improvements over existing flow matching approaches (such as the design of conditional flow matching or efficient inference strategies)

We 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.

To 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.

Sampling algorithm. We 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.

Noise distribution. We 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.

Type and formulation of flow matching loss. We 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.

Architecture. We 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.