Paper: Fast Ensembling with Diffusion Schrödinger Bridge (ICLR-2024)

Reviewer comment (Theory):
Equation (4) contains an error: the forward and backward drifts should not be the same, and I2SB specifically provides a tractable solution only for the backward drift. Additionally, the definition of $f_t^\prime$ is missing. I suggest presenting the analytical backward drift directly, similar to what is provided for the target in (9), and ignore the forward diffusion.

Author response:
Thank you for your suggestions! As you mentioned, the forward and backward drifts are generally not the same in the Schrödinger equation, and I2SB provides a tractable solution only for the backward drift, and only the backward drift is required for evaluation. In Section 2.2 of the revised manuscript, we compensated the DSB background by describing the full procedures of the Schrödinger bridge as an optimal transport problem w.r.t. the forward and backward SDEs with two drifts, and how the backward drift is simulated by the score network in Section 3.2. We also added the explanation on the Fokker-Planck equation for completeness, since the tractable solution for the backward drift is obtained by solving the time-reversal Fokker-Planck equation, leading $\Phi’$ to be tractable, which was generally intractable in existing SB problems before simplifying into the conditional distribution.

Reviewer comment (Writing):
The variables "y" in (11) and (12) lack proper definitions. While their meanings can be inferred from the context, it would be better to provide clear definitions.

Author response:
Thank you for the careful reading and feedback on our paper. We sincerely apologize for the confusion in our paper. We clarified the manuscript based on the feedback by the following:
1. We provided the definition y in (11) and (12), which is colored blue in the revised manuscript. (In the revised manuscript, these equations are now (12) and (13).)
2. For the labels in Figure 2, we will enlarge the labels in the Figure in the camera-ready version of the paper.
3. It is intended that Algorithm 1 uses a smaller font size than the main paper, for fitting the page limit of the conference.
4. The Softmax(Z_1) is a typo and is fixed to Softmax(Z_0). In addition, since the subscript $z_i$ is duplicated in the time domain of the intermediate logits and the ensemble outputs, we super-scripted the ensemble output logits into $z^{(i)}$.
5. The missing references in Section 4.1 are removed since the references had been linked to non-existing tables in our LaTeX file.

Reviewer comment (Presentation):
The labels in Figure 2 are too small and may need to be enlarged for better readability.

Author response:
Thank you for the careful reading and feedback on our paper. We sincerely apologize for the confusion in our paper. We clarified the manuscript based on the feedback by the following:
1. We provided the definition y in (11) and (12), which is colored blue in the revised manuscript. (In the revised manuscript, these equations are now (12) and (13).)
2. For the labels in Figure 2, we will enlarge the labels in the Figure in the camera-ready version of the paper.
3. It is intended that Algorithm 1 uses a smaller font size than the main paper, for fitting the page limit of the conference.
4. The Softmax(Z_1) is a typo and is fixed to Softmax(Z_0). In addition, since the subscript $z_i$ is duplicated in the time domain of the intermediate logits and the ensemble outputs, we super-scripted the ensemble output logits into $z^{(i)}$.
5. The missing references in Section 4.1 are removed since the references had been linked to non-existing tables in our LaTeX file.

Reviewer comment (Writing):
There appears to be a typo in (11). I think it should be "Softmax(Z_0)" instead of the current notation.

Author response:
Thank you for the careful reading and feedback on our paper. We sincerely apologize for the confusion in our paper. We clarified the manuscript based on the feedback by the following:
1. We provided the definition y in (11) and (12), which is colored blue in the revised manuscript. (In the revised manuscript, these equations are now (12) and (13).)
2. For the labels in Figure 2, we will enlarge the labels in the Figure in the camera-ready version of the paper.
3. It is intended that Algorithm 1 uses a smaller font size than the main paper, for fitting the page limit of the conference.
4. The Softmax(Z_1) is a typo and is fixed to Softmax(Z_0). In addition, since the subscript $z_i$ is duplicated in the time domain of the intermediate logits and the ensemble outputs, we super-scripted the ensemble output logits into $z^{(i)}$.
5. The missing references in Section 4.1 are removed since the references had been linked to non-existing tables in our LaTeX file.

Reviewer comment (Experiments):
I believe I2SB remains applicable without the use of temperature annealing ($T$). Can the authors provide further clarification on the role of $T$ and explain how it contributes to differences in performance?

Author response:
Thank you for the insightful question and remarks on temperature annealing! The random annealing is necessary to prevent the model from falling into a trivial solution. We detailed the role of the temperature annealing $T$ in **Common Response 3**.

Reviewer comment (Reproducibility):
It's unclear to me how multiple DBNs are combined in Sec 3.5. For instance, if there are 5 models (M=5) and each DBN is trained with 3 ensembles (as suggested in Sec 4), what is the appropriate value for $L$ and what're the ensembles for each DBN?

Author response:
For more than 3 ensembles, suppose that we have 5 ensemble members $\theta_1,\ldots,\theta_5$. For the first DBN, we choose $\theta_1$ as a source model and distill 3 ensembles: $\theta_1,\theta_2,\theta_3$. The second DBN also has to use $\theta_1$ as a source model and distill other 3 ensembles including the source model: $\theta_1,\theta_4,\theta_5$. When we do inference, we first run the source model $\theta_1$ and get the source logit. Two different DBNs produce different outputs from the source logit. The final output is obtained by averaging the two outputs, as denoted in Equation (13). We simply revised the paper to elaborate the inference with multiple DBNs, colored blue in Section 3.5 in the paper. Due to the page limit, we are planning to add the examples in the final version.

For the conventional ensemble distillation methods, it is well known that the number of ensembles that can be distilled via a single student model is saturated to a certain level. Figure 5 in Section 4.3 shows the similar property in the DBN method and the number of ensembles that a single DBN can distill was saturated to 3, so we fixed the number of ensemble members that DBN distills to three. When we need to distill more than 3 ensembles, adding DBN is more efficient than increasing the size of the DBN network because we can expect additional ensemble effects by averaging outputs of $N$ number of DBNs.

Reviewer comment (Experiments):
I2SB has demonstrated decent performance with NFE=1. Can the authors present an ablation study or provide justification for why additional distillation is necessary for achieving optimal performance?

Author response:
We apologize for any confusion caused by our paper. The I2SB paper showcases strong performance with NFE equal to or greater than ten, with a noticeable drop in performance for smaller numbers of function evaluations, such as NFE=2 (refer to Figure 8 in [1] for more details). Based on the findings in [1], we assumed that near-optimal performance for I2SB is achieved at around NFE=10, and consequently, we set the initial number of timesteps to 5 for both training and inference (For additional NFE values like 10 or 15, please refer to **Table 1**). The primary reason for employing diffusion distillation is to reduce the inference cost while maintaining performance at an appropriate level.

**Table 1**. DBN on CIFAR-10 with different number of inference steps.
|                  | ACC   | NLL    |
|------------------|-------|--------|
| DE-1             | 91.30 | 0.3382 |
| +1 DBN (step 5)  | 92.97 | 0.2435 |
| +1 DBN (step 10) | 93.01 | 0.2441 |
| +1 DBN (step 15) | 92.95 | 0.2442 |

[1] G. Liu et al., “I2SB: Image-to-Image Schrödinger Bridge”, ICML 2023, URL: https://arxiv.org/abs/2302.05872

Reviewer comment (Presentation):
The font size in Alg 1 seems to be smaller than the one in the main paper. Is this intended?

Author response:
Thank you for the careful reading and feedback on our paper. We sincerely apologize for the confusion in our paper. We clarified the manuscript based on the feedback by the following:
1. We provided the definition y in (11) and (12), which is colored blue in the revised manuscript. (In the revised manuscript, these equations are now (12) and (13).)
2. For the labels in Figure 2, we will enlarge the labels in the Figure in the camera-ready version of the paper.
3. It is intended that Algorithm 1 uses a smaller font size than the main paper, for fitting the page limit of the conference.
4. The Softmax(Z_1) is a typo and is fixed to Softmax(Z_0). In addition, since the subscript $z_i$ is duplicated in the time domain of the intermediate logits and the ensemble outputs, we super-scripted the ensemble output logits into $z^{(i)}$.
5. The missing references in Section 4.1 are removed since the references had been linked to non-existing tables in our LaTeX file.

Reviewer comment (Novelty):
**Other Transport Models** To the best of my understanding, related works such as "Rectified Flow" [1] and "Flow Matching" [2] are capable of accomplishing similar functionalities. Additionally, when compared to the Schrödinger bridge approach, these models operate through a deterministic process, which could potentially benefit from lower variance properties. The omission of such deterministic flow models from the discussion may limit the generalizability and applicability of the method proposed in the current manuscript.

Author response:
We sincerely appreciate you for insightful questions! As you have issued, the deterministic flow models that simulate the ODE that bridges different distributions have similar functionalities and benefit from variance reduction properties. Nevertheless, adding some stochasticity in the learning process leads us to avoid the probability flow to fall into the trivial solution. This is also explained as we further describe the necessity of the temperature annealing for robustifying the source datasets, which is described in Common Response 3 in detail. The result that compares our approach to deterministic flow models using rectified flow is provided in **Common Response 2**.

Reviewer comment (Reproducibility):
**Inference Algorithm** In the manuscript, it appears that the authors have not provided details regarding the inference algorithm for the proposed DBN during the model inference stage. This omission leaves a critical gap in the paper, as understanding the inference mechanism is essential for a comprehensive evaluation of the proposed method.

Author response:
We apologize for giving confusion on the inference algorithm. We first appended how the feature extractor $g_{\psi_1}$ computes the feature $\boldsymbol{h}_1$ by forward evaluation, and the inference algorithm is a single step of reverse-time evaluation of the score network via the distilled score network from time $1\rightarrow 0$, as described in Equation (12). We also added the detailed architecture of the score network utilized for inference in the Figure 5 of Appendix A.2, for further understanding.

Reviewer comment (Experiments):
**Annealing of Temperature** In Section 3.2, the authors discuss the concept of temperature distribution, yet the experimental section lacks elaboration on how the temperature distribution is selected and what principles govern temperature annealing. This absence of information creates a gap in the methodological clarity and poses questions regarding the thoroughness of the experimental design. Furthermore, from a scholarly perspective, treating the distribution as a Gumbel-Softmax distribution [3] could raise additional questions. Specifically, one might inquire whether the training variance of the Diffusion Bridge Networks (DBN) is influenced by the temperature parameter. Addressing such intricate relationships between the temperature and training variance would enhance the paper's academic rigor and contextual relevance.

Author response:
Thank you for giving insightful questions and suggestions about the temperature annealing, and we apologize for causing confusion on the concept of temperature distribution. For our choice of the temperature distribution, we chose the Beta distribution $T = 2 ( 1+0.2\alpha), \alpha\sim \mathrm{Beta}(\cdot; 1, 5)$ for simple and efficient temperature annealing. We added these details on hyperparameters including the temperature annealing in Appendix B. For further details on temperature annealing, please refer to **Common Response 3**.

We appreciate the suggestion that can support the academic rigor and contextual relevance by improving theoretical backgrounds on our proposed methods. We will investigate these theoretical aspects and update this in the final version of our paper.

Reviewer comment (Presentation):
**LaTeX Compile** In Section 4.1, line 3, the manuscript contains two instances of "??", which clearly indicate placeholders or unresolved references. The authors should rectify this issue to enhance the document's professionalism and completeness prior to submitting the finalized manuscript. The presence of such markers detracts from the paper's overall quality and could create potential ambiguities for the reader.

Author response:
Thank you for pointing out the errors. We sincerely apologize for the completeness of our paper, by mis-referencing and errors in our manuscript. We intensively looked over the manuscript and fixed the typo, grammar errors, and the duplicating notations. We also modified the two reference instances in the manuscript.

Reviewer comment (Writing):
This article's writing is very hard to follow. There are many spelling or grammar mistakes (see Questions). Ideas are presented poorly or too quickly, leaving the reader confused.

Author response:
Thank you for the thoughtful reading and concerns on our paper. We apologize for the clarity of this paper and made up the Schrodinger bridge part (Section 2.2) of our paper. Briefly, we complemented the setting of the Schrodinger systems and denoted the variables that are not defined. And we added some supporting statements such as how I2SB solves the restoration problem in the background part. And we also clarified the equation (8) (The main SDE from $Z_1$ to $Z_0$) in the paper.

For further details, please refer to the revised version (the refined parts written in blue) of Sections 2.2 and 3.2 in our paper. We also clarified the list of the revised parts in **Common Response 4**.

Reviewer comment (Reproducibility):
Many important points are missing from the training details, in particular, the architecture of the score network is never given anywhere in the article.

Author response:
We apologize for the absence of details regarding the architectures and settings of the DBN model and thanks for pointing them out. The training details of the DBN are outlined below.

1. **Architecture of the Score Network**

Our score network initially embeds three inputs: $Z_t$ (the input logit at time $t$), $h_1$ (the condition feature), and time $t$. Subsequently, the three embedded features of the inputs pass through the score network, as depicted in Figure 5, consisting of multiple InvertedResidual blocks. A detailed architecture description of the network is provided in Appendix A.2.

2. **Hyperparameter Settings and Training Details of the Ensemble Networks and DBN**

In conjunction with the score network architecture, we have included the hyperparameter settings of our model in Appendix B. The detailed hyperparameters, listed in Table 2, encompass the number of parameters of the classifier and score networks, batch size, epochs, and learning rates for the main training and distillation phases, mixup coefficients used for data augmentation, and the diffusion drift coefficient for generating the diffusion bridge.

Additionally, we elaborate on the omitted hyperparameter settings of the ensemble networks, referred to in Table 4 of Appendix B, which were simply assumed to use the same architecture as the BN paper.

3. **Comparison of BN to DBN**

To be more precise, BN distills the output of a model located along the low-loss subspace between teacher networks. As we described in Section 2.1, the bridge network [1] hinges on the features of the mode connectivity among the low-loss subspace. BN first learns the pre-trained low-loss surfaces between two pre-trained ensemble members by learning the “anchor” parameter that satisfies on the Bezier curve, i.e., finding the curve with low loss on the loss surface. BN assumed and empirically found that the midpoint of the Bezier curve tends to achieve the lowest loss and used this midpoint (parameterized by $\theta_{i,j}(0.5)$) as their teacher. Therefore, they train a lightweight neural network that distills the midpoint.

During the training process, the size of the teacher network is significantly larger than that of the Bridge in both BN and DBN. Besides, DBN uses small diffusion steps, the majority of the training cost arises from calculating the outputs of the teacher networks for distillation, so **the cost of learning the Bridge of both BN and DBN is relatively negligible compared to the teachers**.

To elaborate further, assuming N pre-trained weights for the Teacher model (cost=$1$), the training process for BN and DBN to distill 3 teacher models (DE-3) is as follows. First, BN selects 2 pairs of teachers from three teacher weights and learns $2$ low-loss curves connecting the teacher pairs using the teacher architecture (cost=$2\times1$). Then, it additionally trains $2$ BNs to predict the output of the model in the middle ($\theta_{i,j}(0.5)$) of each low-loss curve (cost=$2\times2.xx$). Similarly, for DBN, it selects 1 source teacher and trains $1$ score network (smaller than BN) to learn the ensemble outputs of $3$ teachers with 5 diffusion steps (cost=$1\times3.xx$). Following this, an additional training is conducted where the 5 steps are distilled into one step (cost=$1\times3.xx$). In the final version of the paper, we will provide a clearer explanation of the training cost of DBN relative to BN.

[1] E. Yun et al., “Traversing between modes in function space for fast ensembling”, ICML 2023

[2] T. Salimans et al., “Progressive distillation for fast sampling of diffusion models”, ICLR 2022

Reviewer comment (Reproducibility):
Contrary to what the authors claim, the appendix does not give any hyperparameter settings: the training details of both the ensemble networks and the DBN are left unclear (optimizer, scheduler, number of epochs/steps, distillation schedule, etc.).

Author response:
We apologize for the absence of details regarding the architectures and settings of the DBN model and thanks for pointing them out. The training details of the DBN are outlined below.

1. **Architecture of the Score Network**

Our score network initially embeds three inputs: $Z_t$ (the input logit at time $t$), $h_1$ (the condition feature), and time $t$. Subsequently, the three embedded features of the inputs pass through the score network, as depicted in Figure 5, consisting of multiple InvertedResidual blocks. A detailed architecture description of the network is provided in Appendix A.2.

2. **Hyperparameter Settings and Training Details of the Ensemble Networks and DBN**

In conjunction with the score network architecture, we have included the hyperparameter settings of our model in Appendix B. The detailed hyperparameters, listed in Table 2, encompass the number of parameters of the classifier and score networks, batch size, epochs, and learning rates for the main training and distillation phases, mixup coefficients used for data augmentation, and the diffusion drift coefficient for generating the diffusion bridge.

Additionally, we elaborate on the omitted hyperparameter settings of the ensemble networks, referred to in Table 4 of Appendix B, which were simply assumed to use the same architecture as the BN paper.

3. **Comparison of BN to DBN**

To be more precise, BN distills the output of a model located along the low-loss subspace between teacher networks. As we described in Section 2.1, the bridge network [1] hinges on the features of the mode connectivity among the low-loss subspace. BN first learns the pre-trained low-loss surfaces between two pre-trained ensemble members by learning the “anchor” parameter that satisfies on the Bezier curve, i.e., finding the curve with low loss on the loss surface. BN assumed and empirically found that the midpoint of the Bezier curve tends to achieve the lowest loss and used this midpoint (parameterized by $\theta_{i,j}(0.5)$) as their teacher. Therefore, they train a lightweight neural network that distills the midpoint.

During the training process, the size of the teacher network is significantly larger than that of the Bridge in both BN and DBN. Besides, DBN uses small diffusion steps, the majority of the training cost arises from calculating the outputs of the teacher networks for distillation, so **the cost of learning the Bridge of both BN and DBN is relatively negligible compared to the teachers**.

To elaborate further, assuming N pre-trained weights for the Teacher model (cost=$1$), the training process for BN and DBN to distill 3 teacher models (DE-3) is as follows. First, BN selects 2 pairs of teachers from three teacher weights and learns $2$ low-loss curves connecting the teacher pairs using the teacher architecture (cost=$2\times1$). Then, it additionally trains $2$ BNs to predict the output of the model in the middle ($\theta_{i,j}(0.5)$) of each low-loss curve (cost=$2\times2.xx$). Similarly, for DBN, it selects 1 source teacher and trains $1$ score network (smaller than BN) to learn the ensemble outputs of $3$ teachers with 5 diffusion steps (cost=$1\times3.xx$). Following this, an additional training is conducted where the 5 steps are distilled into one step (cost=$1\times3.xx$). In the final version of the paper, we will provide a clearer explanation of the training cost of DBN relative to BN.

[1] E. Yun et al., “Traversing between modes in function space for fast ensembling”, ICML 2023

[2] T. Salimans et al., “Progressive distillation for fast sampling of diffusion models”, ICLR 2022

Reviewer comment (Reproducibility):
Lack of training detail is a significant issue, in particular for the DBN. If its training is particularly expensive/slow, it is an important drawback despite the speedup it provides during inference.

Author response:
We apologize for the absence of details regarding the architectures and settings of the DBN model and thanks for pointing them out. The training details of the DBN are outlined below.

1. **Architecture of the Score Network**

Our score network initially embeds three inputs: $Z_t$ (the input logit at time $t$), $h_1$ (the condition feature), and time $t$. Subsequently, the three embedded features of the inputs pass through the score network, as depicted in Figure 5, consisting of multiple InvertedResidual blocks. A detailed architecture description of the network is provided in Appendix A.2.

2. **Hyperparameter Settings and Training Details of the Ensemble Networks and DBN**

In conjunction with the score network architecture, we have included the hyperparameter settings of our model in Appendix B. The detailed hyperparameters, listed in Table 2, encompass the number of parameters of the classifier and score networks, batch size, epochs, and learning rates for the main training and distillation phases, mixup coefficients used for data augmentation, and the diffusion drift coefficient for generating the diffusion bridge.

Additionally, we elaborate on the omitted hyperparameter settings of the ensemble networks, referred to in Table 4 of Appendix B, which were simply assumed to use the same architecture as the BN paper.

3. **Comparison of BN to DBN**

To be more precise, BN distills the output of a model located along the low-loss subspace between teacher networks. As we described in Section 2.1, the bridge network [1] hinges on the features of the mode connectivity among the low-loss subspace. BN first learns the pre-trained low-loss surfaces between two pre-trained ensemble members by learning the “anchor” parameter that satisfies on the Bezier curve, i.e., finding the curve with low loss on the loss surface. BN assumed and empirically found that the midpoint of the Bezier curve tends to achieve the lowest loss and used this midpoint (parameterized by $\theta_{i,j}(0.5)$) as their teacher. Therefore, they train a lightweight neural network that distills the midpoint.

During the training process, the size of the teacher network is significantly larger than that of the Bridge in both BN and DBN. Besides, DBN uses small diffusion steps, the majority of the training cost arises from calculating the outputs of the teacher networks for distillation, so **the cost of learning the Bridge of both BN and DBN is relatively negligible compared to the teachers**.

To elaborate further, assuming N pre-trained weights for the Teacher model (cost=$1$), the training process for BN and DBN to distill 3 teacher models (DE-3) is as follows. First, BN selects 2 pairs of teachers from three teacher weights and learns $2$ low-loss curves connecting the teacher pairs using the teacher architecture (cost=$2\times1$). Then, it additionally trains $2$ BNs to predict the output of the model in the middle ($\theta_{i,j}(0.5)$) of each low-loss curve (cost=$2\times2.xx$). Similarly, for DBN, it selects 1 source teacher and trains $1$ score network (smaller than BN) to learn the ensemble outputs of $3$ teachers with 5 diffusion steps (cost=$1\times3.xx$). Following this, an additional training is conducted where the 5 steps are distilled into one step (cost=$1\times3.xx$). In the final version of the paper, we will provide a clearer explanation of the training cost of DBN relative to BN.

[1] E. Yun et al., “Traversing between modes in function space for fast ensembling”, ICML 2023

[2] T. Salimans et al., “Progressive distillation for fast sampling of diffusion models”, ICLR 2022

Reviewer comment (Experiments):
It is unclear whether Bridge Networks are also using distillation, and whether their training is comparable to DBN. Without this information, it is hard to compare the results obtained.

Author response:
We apologize for the absence of details regarding the architectures and settings of the DBN model and thanks for pointing them out. The training details of the DBN are outlined below.

1. **Architecture of the Score Network**

Our score network initially embeds three inputs: $Z_t$ (the input logit at time $t$), $h_1$ (the condition feature), and time $t$. Subsequently, the three embedded features of the inputs pass through the score network, as depicted in Figure 5, consisting of multiple InvertedResidual blocks. A detailed architecture description of the network is provided in Appendix A.2.

2. **Hyperparameter Settings and Training Details of the Ensemble Networks and DBN**

In conjunction with the score network architecture, we have included the hyperparameter settings of our model in Appendix B. The detailed hyperparameters, listed in Table 2, encompass the number of parameters of the classifier and score networks, batch size, epochs, and learning rates for the main training and distillation phases, mixup coefficients used for data augmentation, and the diffusion drift coefficient for generating the diffusion bridge.

Additionally, we elaborate on the omitted hyperparameter settings of the ensemble networks, referred to in Table 4 of Appendix B, which were simply assumed to use the same architecture as the BN paper.

3. **Comparison of BN to DBN**

To be more precise, BN distills the output of a model located along the low-loss subspace between teacher networks. As we described in Section 2.1, the bridge network [1] hinges on the features of the mode connectivity among the low-loss subspace. BN first learns the pre-trained low-loss surfaces between two pre-trained ensemble members by learning the “anchor” parameter that satisfies on the Bezier curve, i.e., finding the curve with low loss on the loss surface. BN assumed and empirically found that the midpoint of the Bezier curve tends to achieve the lowest loss and used this midpoint (parameterized by $\theta_{i,j}(0.5)$) as their teacher. Therefore, they train a lightweight neural network that distills the midpoint.

During the training process, the size of the teacher network is significantly larger than that of the Bridge in both BN and DBN. Besides, DBN uses small diffusion steps, the majority of the training cost arises from calculating the outputs of the teacher networks for distillation, so **the cost of learning the Bridge of both BN and DBN is relatively negligible compared to the teachers**.

To elaborate further, assuming N pre-trained weights for the Teacher model (cost=$1$), the training process for BN and DBN to distill 3 teacher models (DE-3) is as follows. First, BN selects 2 pairs of teachers from three teacher weights and learns $2$ low-loss curves connecting the teacher pairs using the teacher architecture (cost=$2\times1$). Then, it additionally trains $2$ BNs to predict the output of the model in the middle ($\theta_{i,j}(0.5)$) of each low-loss curve (cost=$2\times2.xx$). Similarly, for DBN, it selects 1 source teacher and trains $1$ score network (smaller than BN) to learn the ensemble outputs of $3$ teachers with 5 diffusion steps (cost=$1\times3.xx$). Following this, an additional training is conducted where the 5 steps are distilled into one step (cost=$1\times3.xx$). In the final version of the paper, we will provide a clearer explanation of the training cost of DBN relative to BN.

[1] E. Yun et al., “Traversing between modes in function space for fast ensembling”, ICML 2023

[2] T. Salimans et al., “Progressive distillation for fast sampling of diffusion models”, ICLR 2022

Reviewer comment (Experiments):
The number of ensembles that a score network is trained on is unclear. Section 4 indicates that a score network is trained on 3 ensembles, with no clarification on why this number is chosen (presumably Section 4.3).

Author response:
For the conventional ensemble distillation methods, it is well known that the number of ensembles that can be distilled via a single student model is saturated to a certain level. Figure 5 in Section 4.3 shows the similar property in the DBN method and the number of ensembles that a single DBN can distill was saturated to 3, so we fixed the number of ensemble members that DBN distills to three. When we need to distill more than 3 ensembles, adding DBN is more efficient than increasing the size of the DBN network because we can expect additional ensemble effects by averaging outputs of $N$ number of DBNs.

Reviewer comment (Experiments):
It is not explained why the authors do not simply increase the size of the DBN network instead of training on a fixed number of ensembles.

Author response:
For the conventional ensemble distillation methods, it is well known that the number of ensembles that can be distilled via a single student model is saturated to a certain level. Figure 5 in Section 4.3 shows the similar property in the DBN method and the number of ensembles that a single DBN can distill was saturated to 3, so we fixed the number of ensemble members that DBN distills to three. When we need to distill more than 3 ensembles, adding DBN is more efficient than increasing the size of the DBN network because we can expect additional ensemble effects by averaging outputs of $N$ number of DBNs.

Reviewer comment (Reproducibility):
For more than 3 ensembles, 2 or more DBNs are trained, but it is unclear whether the DBNs are trained on all of the ensembles at the same time or limited to a given subset of 3 ensembles, as indicated previously.

Author response:
For more than 3 ensembles, suppose that we have 5 ensemble members $\theta_1,\ldots,\theta_5$. For the first DBN, we choose $\theta_1$ as a source model and distill 3 ensembles: $\theta_1,\theta_2,\theta_3$. The second DBN also has to use $\theta_1$ as a source model and distill other 3 ensembles including the source model: $\theta_1,\theta_4,\theta_5$. When we do inference, we first run the source model $\theta_1$ and get the source logit. Two different DBNs produce different outputs from the source logit. The final output is obtained by averaging the two outputs, as denoted in Equation (13). We simply revised the paper to elaborate the inference with multiple DBNs, colored blue in Section 3.5 in the paper. Due to the page limit, we are planning to add the examples in the final version.

Reviewer comment (Experiments):
In Section 4.2, the number of ensembles increases to 9, and the number of DBNs is left unclear ("two or more"), as well as for BNs.

Author response:
For more than 3 ensembles, suppose that we have 5 ensemble members $\theta_1,\ldots,\theta_5$. For the first DBN, we choose $\theta_1$ as a source model and distill 3 ensembles: $\theta_1,\theta_2,\theta_3$. The second DBN also has to use $\theta_1$ as a source model and distill other 3 ensembles including the source model: $\theta_1,\theta_4,\theta_5$. When we do inference, we first run the source model $\theta_1$ and get the source logit. Two different DBNs produce different outputs from the source logit. The final output is obtained by averaging the two outputs, as denoted in Equation (13). We simply revised the paper to elaborate the inference with multiple DBNs, colored blue in Section 3.5 in the paper. Due to the page limit, we are planning to add the examples in the final version.

Reviewer comment (Experiments):
Is the DSB not able to learn if there is no temperature annealing?

Author response:
Thank you for the thoughtful question! Since the issue of temperature annealing in the source logits is a common issue over reviewers, we provided the interpretation and additional explanations on why we made use of temperature annealing and why this is necessary in the **Common Response 3**.

Reviewer comment (Experiments):
Why choose in particular ResNet-32x2 and x4 rather than a standard ResNet-32?

Author response:
In order to fairly compare the performance between DBN and the existing BN method, we followed the architecture that the BN paper used. For CIFAR-10 and CIFAR-100 datasets, the BN paper used the ResNet-32x2 and ResNet-32x4 architecture.

Reviewer comment (Evaluation):
Do the authors have an idea why DBN shows poor EC scores compared to the other metrics?

Author response:
We explained about the issues on the expected calibration error (ECE) metric in **Common Response 4**.

Reviewer comment (Writing):
Abstract: "theory of [the] Schrödinger bridge"

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Introduction: "incurrs an extra"

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Section 2.2: "has rarely been demonstrated its practicality" "demonstrated for real-world image dataset" "Gaussain" "by a certain PDEs"

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Section 3.2: "Based on the formulation" ??

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Section 3.3: "super performance". No use for the (NFE) abbreviation if it's never used after.

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Section 3.4: $\Phi$ or $\Phi'$?

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Section 4: "because Ryabinin et al. (2021) refined poor convergence of" "more profound formulations"

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Presentation):
Section 4.1: "results on CIFAR-10 is shown in ?? of ??"

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Writing):
Conclusion: "Additionally,  [...] of 3 DE models", nominal sentence.

Author response:
We sincerely apologize for the errors we have made in the manuscript. All these points are fixed in the revised manuscript.

Reviewer comment (Evaluation):
DBN models are notably less calibrated than ensembles or bridge networks across tasks.

Author response:
> DBN models are notably less calibrated than ensembles or bridge networks across tasks. Can the authors comment on the calibration error increase with this method?

According to the results in Table 1 of the paper, the expected calibration error (ECE) score is higher (worse) compared to the BN paper but outperforms existing ensemble distillation methods like ED [1] or END2 [2]. Upon closer examination of these seemingly contradictory results, we arrived at the following interpretations and explanations.

The ECE metric inherently faces limitations in measuring the exact calibration error. In contrast to other calibration metrics like NLL and BS, which directly estimate the level of generalization through continuous calculation, the ECE score initially divides the confidence space into discretely defined bins. The artificial and discrete division of confidences over the dataset introduces some inaccuracy in determining the "exact" calibration level. This is due to the fact that (a) data points in the same bin are treated the same, even if there still exists a calibration error at that specific data point, and (b) the neighboring bin is considered to be disjoint.
Nonetheless, DBN still takes advantage of the generalization performance with respect to the NLL and Brier score(BS), which also capture the performance of the uncertainty calibration.

[1] J. Hinton et al.,“Distilling the knowledge in a neural network”, NIPS 2015 \
[2] M. Ryabinin et al., “Scaling ensemble distribution distillation to many classes with proxy targets”, NeurIPS 2021

Reviewer comment (Evaluation):
No evaluation on in / out of domain uncertainty / calibration. IMO this is a huge reason why ensembling is done in the first place.

Author response:
> No evaluation on in / out of domain uncertainty / calibration. IMO this is a huge reason why ensembling is done in the first place.

We evaluated our method on the widely used CIFAR-10-C [1], which is the dataset of common corruptions on the CIFAR-10 dataset. and compared it to the existing BN method as in __Table 1__.  Achieving superior performance in the CIFAR-10-C dataset implies that the model is robust with respect to the out-of-distribution distribution, tilted by a variety of distortions and corruptions in the in-distribution dataset (in this case, CIFAR-10). We will update the full OOD results in the final version.

Table 1. Test results on CIFAR-10-C. cNLL stands for NLL calibrated with an optimal temperature.

|        | ACC   | NLL    | cNLL   |
|--------|-------|--------|--------|
| DE-1   | 86.95 | 0.5158 | 0.4477 |
| +2 BN  | 87.80 | 0.3740 | 0.3749 |
| +2 DBN | 88.91 | 0.3759 | 0.3682 |
| DE-2   | 88.85 | 0.3860 | 0.3693 |
| DE-3   | 89.47 | 0.3459 | 0.3410 |

[1] D. Hendrycks, “Benchmarking Neural Network Robustness to Common Corruptions and Perturbations”, ICLR 2019

Reviewer comment (Evaluation):
Evaluation is with relatively small models.

Author response:
> Evaluation is with relatively small models.

Due to the lack of our computational budget, there was some limit on scaling up the neural network model. Instead, we tried our model on more complex and large-scale datasets, ImageNet64. We included the extended experimental results in __Common Response 1__.

Reviewer comment (Writing):
Section 4.1 "?? of ??" still in text.

Author response:
> Section 4.1 "?? of ??" still in text

We sincerely apologize for the error. We fixed this in the revised manuscript.

Reviewer comment (Evaluation):
I think parameter-count/memory savings is useful here and could be reported.

Author response:
> I think parameter-count/memory savings is useful here and could be reported.

In our paper, the ratio of saved memory compared to the deep ensemble models is introduced in the #Params column of Table 1 of the paper. The actual number of parameters (and thus the number of saved parameters) is included in the Appendix B.

Reviewer comment (Experiments):
The baseline performances shown in the experiment section does not completely agree with the prior works. The scale of the experiments were somewhat small.

Author response:
Thank you for the careful investigation of our experiments. We double-checked our experiments and figured out that the saved checkpoint of the teacher network for training BN had a problem. Precisely, one of the layers of the teacher network of BN was saved in the different layer and it caused the performance drop. We fixed this and trained the BN again. And as you pointed out that the performance is much less volatile for large-scaled datasets, we tried our method in the ImageNet dataset. Because of the limit of the training budget, we tried this on the downscaled $64\times 64$ datasets.
Since this issue is the core part of our response that can strengthen our method, we organized our experimental results for extended datasets in the __Common Response 1__.

Reviewer comment (Experiments):
Lack of analysis on the diffusion part of the design.

Author response:
Thank you for your thoughtful question of pointing out the diffusion part of our method. The precise required training resource will be added in the appendix, in place of the architecture, and hyperparameter setting section. 

Recognizing the performance limitations of BN, we felt the need for a more flexible approach. Therefore, instead of simple methods like the conventional discrete-time Normalizing Flow (NF), which has limitations in distribution approximation, requiring invertible networks for likelihood maximization, we adopted diffusion-based methods such as I2SB. In the case of methods like NF, creating a tractable distribution necessitates the use of invertible networks, imposing constraints on the network structure. In contrast, diffusion-based methods do not have such restrictions on network invertibility, allowing for the generation of a more flexible probability path. Despite the relative complexity of I2SB in the theoretical background, it delivered sufficient performance with low computational cost. Of course, as an alternative, one could also choose a slightly more flexible method yet similar to NF, such as Rectified Flow (RF). Its experimental results are described in the __Common Response 2__.

Reviewer comment (Experiments):
The introduction of temperature does define a distribution, but essentially at a cost of compromising the performance of the original ensemble member (for example, if you look at the NLL loss of the logits after applying temperature, it will be worse). Is this really a reasonable design? If you look for bridging on logit space with another, wouldn't it make more sense to consider the logit space as is (e.g., get a distribution via input perturbation, etc.)?

Author response:
Thanks for the good question. As you mentioned, when $T$ approaches to infinity, the categorical source distribution becomes uniform. Just as Knowledge Distillation requires appropriate temperature choice, preventing it from focusing on label class only ($T=1$) or erasing all information ($T=\infty$), an appropriate temperature distribution, maybe ranging from 1 to 20, is necessary. We obtained good results with $T = 2(1+0.2\gamma)$, where $\gamma$ is sampled from a Beta distribution with parameters $1$ and $5$.

To verify that the random temperature annealing is a valid choice, We ablated our proposed temperature annealing method to the mixup approach, as we have explained in the rebuttal. With varying mixup parameters $\alpha$ in mixup formula, $\boldsymbol{x}' = \lambda \boldsymbol{x}_1 + (1-\lambda)\boldsymbol{x}_2$, where $\lambda$ is sampled from a Beta distribution with parameters $\alpha$ and $\alpha$, with values of 0.1, 0.3, 0.5, 0.7, and 0.9, the results are demonstrated in __Table 2__ as follows:

Table 2. DBN in CIFAR-10 with different mixup $\alpha$ on the source distribution.
|                        | ACC   | NLL    |
|------------------------|-------|--------|
| DE-1                   | 91.30 | 0.3382 |
| +1 DBN (alpha 0.1)     | 92.47 | 0.2632 |
| +1 DBN (alpha 0.3)     | 92.53 | 0.2596 |
| +1 DBN (alpha 0.5)     | 92.75 | 0.2579 |
| +1 DBN (alpha 0.7)     | 92.62 | 0.2613 |
| +1 DBN (alpha 0.9)     | 92.29 | 0.2715 |
| +1 DBN (rand. anneal.) | 93.04 | 0.2394 |

the following aspects are observed:
* Without using the random temperature annealing, $\alpha=0.5$ achieved the best performance among all possible choices. It shows a significant drop around $\alpha=0.9$. *(Cliff behavior)*
* Even the best working mixup did not reach the random temperature annealing as shown in __Table 2__. *(Advantage of random annealing)*

Reviewer comment (Experiments):
Why DSB rather than a conditional diffusion from Gaussian? In $\textsf{I}^2\textsf{SB}$ the motivation was clearer as they hoped to leverage the structural prior of the image. Empirically, how big would the difference be?

Author response:
Thanks for your insightful question! Just like I2SB aims to leverage the structural prior of the image from *degraded* images, our paper shares a similar motivation. We aim to exploit the *degraded* logits generated from a single classifier as the structural prior of the logit space of the ensemble. This implies that the functional outputs of the ensemble members can be represented as a distribution conditioned on those of the (single) source member.

The conditional diffusion from Gaussian can also yield the output logits, but this does not utilize the useful information of the logit of the source model. This is also shown in the image case: I2SB, which constructs the diffusion bridge from the degraded source image distribution to the fine-grained target distribution, outperforms the previous conditional diffusion model approach [2] which generates the fine-grained images by imposing the degraded source images as its condition. We can achieve good performance with a small number of diffusion steps because we can leverage information from the source model. If the same performance is desired in the conditional diffusion, it would require a greater number of diffusion steps. To support this claim, we trained the conditional diffusion with the same number of diffusion steps as I2SB and Table 1 shows an experiment that compares I2SB with the conditional diffusion method.

Table 1. I2SB vs. conditional diffusion from Guassian on CIFAR10.

|                         | ACC   | NLL    | DEE   |
|-------------------------|-------|--------|-------|
| DE-1                    | 91.30 | 0.3382 | 1.000 |
| +1 DBN (I2SB)           | 92.97 | 0.2435 | 2.228 |
| 1 conditional diffusion | 90.78 | 0.3727 |   < 1 |

[1] G. Liu et al., “I2SB: Image-to-Image Schrödinger Bridge”, ICML 2023, URL: https://arxiv.org/abs/2302.05872 \
[2] C. Saharia et al., “Palette: Image-to-Image Diffusion Models”, SIGGRAPH 2022, URL: https://arxiv.org/abs/2111.05826

Reviewer comment (Experiments):
The paper has a bunch of analysis on the ensemble approach, but little (almost none) on the diffusion portion. For example, how much training (data and speed) is needed for the DSB training? Why diffusion instead of some simpler approaches like normalizing flows, and bringing all ensemble members to a (same) tractable distribution?

Author response:
Thank you for your thoughtful question of pointing out the diffusion part of our method. The precise required training resource will be added in the appendix, in place of the architecture, and hyperparameter setting section. 

Recognizing the performance limitations of BN, we felt the need for a more flexible approach. Therefore, instead of simple methods like the conventional discrete-time Normalizing Flow (NF), which has limitations in distribution approximation, requiring invertible networks for likelihood maximization, we adopted diffusion-based methods such as I2SB. In the case of methods like NF, creating a tractable distribution necessitates the use of invertible networks, imposing constraints on the network structure. In contrast, diffusion-based methods do not have such restrictions on network invertibility, allowing for the generation of a more flexible probability path. Despite the relative complexity of I2SB in the theoretical background, it delivered sufficient performance with low computational cost. Of course, as an alternative, one could also choose a slightly more flexible method yet similar to NF, such as Rectified Flow (RF). Its experimental results are described in the __Common Response 2__.

Reviewer comment (Experiments):
About the baseline: the Bridge Network paper seems to have better numbers on TinyImageNet than what the authors reported in this paper (e.g., 65.82% accuracy for 3-BN-sm and the FLOPs seems to be much smaller (1.15x). The BN paper also evaluated on ImageNet, where in my experience the accuracy is **much less volatile**. If, as authors claim that the BN approach has more training cost (at the end of Sec. 2.1), perhaps the DSB approach should be compared with BN on that turf as well.

Author response:
If you check the source code for the BN [7], you'll find that it uses the indices from 0 to 90,000 from the TinyImageNet training set as the training data and uses indices from 90,000 to 100,000 as the validation data. In contrast, in our work [8], we used indices from 0 to 81,920 as the training data and indices from 81,920 to 100,000 as the validation data. Due to these differences, despite synchronizing the other hyperparameters with the BN source code, we think that it has led to an overall performance degradation in our work. Therefore, as mentioned, we are currently conducting experiments with BN$_\mathrm{small}$, and we will present the results as soon as they are available.

[7] https://github.com/yuneg11/Bridge-Network/blob/main/configs/tin200/modules/dataset.yaml \
[8] Supplementary Material, “data/build.py” line 198.