Inverse Bridge Matching Distillation

Dear Reviewer kKRF, thank you for your comments.

(1) Relation to SiD [1] and FGM [2]. Following your suggestion, we will extend the discussion of the related work in the main text. Below, we present the preliminary version of the extension:

Unlike SiD [1] and FGM [2], we focus on diffusion-bridge models used for data-to-data translation and not for generation from noise. Furthermore, our high-level objective is the KL divergence $\text{KL}(\text{BM}(\Pi_{\theta})||M^*)$ between path measures of teacher model $M^*$ and path measure $\text{BM}(\Pi_{\theta})$ given by the generator $G_{\theta}$, which differs from Fisher Divergence used in SiD [1]. The motivation of our high-level objective is to restore alignment between data pairs $(x_0, x_T)$. We derive our final tractable objective using different techniques (see Appendix A), i.e., we do not use flow product or score product identities as in [1, 2] but hypothesize that analogous identities can be derived for diffusion bridge models. Our final objective can be rewritten similarly to the final objective used in SiD [1, Eq. 23] ($\alpha=0.5$) and FGM [2, Eqs. 4.11, 4.12]:

$$\mathbb{E}_{t, p\_{\theta}(x\_0, x\_t, x\_T)} \Big[ \underbrace{\|\widehat{x}\_0^*(x\_t, t, x\_T) - \widehat{x}\_0^{\phi}(x\_t, t, x\_T) \|^2}\_{\text{quadratic term}} + 2 \left< x_0^{*}(x_t, t, x_T) - x_0^{\phi}(x_t, t, x_T), x_0^{\phi}(x_t, t, x_T) - x_0 \right> \Big]$$

However, in both SiD [1] (for $\alpha=1.0, 1.2$ used in experiments) and FGM [2], the authors either omitted the quadratic term or even used a negative coefficient for it in image experiments since it introduced instability. Unlike SiD [1] and FGM [2], we do not omit any parts of the original loss function and use it as the theory provides it.

Relation between KL divergence of path measures and Fisher Divergence. To highlight the difference between KL divergence of path measures (which we use) and Fisher Divergence (which is used in SiD [1]), consider two reverse-time diffusions $D_1$ and $D_2$ given by the same starting distribution $p(x_T)$ and SDEs:

$$D_1: dx\_t = v(x\_t, t)dt + g^2(t)d\bar{w}\_t, \quad D_2: dx\_t = \widehat{v}(x\_t, t)dt + g^2(t)d\bar{w}\_t.$$

Denote marginal densities as $p(x\_t)$ for $D\_1$ and $\widehat{p}(x\_t)$ for $D\_2$, then KL divergence and Fisher Divergences are given as:

$$\text{KL}(D\_1||D\_2) = \mathbb{E}\_{t, p\_t(x\_t) } \Big[\frac{1}{2g^2(t)}\|v(x\_t, t) - \widehat{v}(x\_t, t) \|^2 \Big] + \underbrace{\text{KL}(p(x\_T)||\widehat{p}(x\_T))}\_{= 0 \text{ if } p(x\_T) = \widehat{p}(x\_T)}$$ $$\mathcal{L}\_{\text{SiD}}(D\_1, D\_2) = \mathbb{E}\_{t} \Big[\underbrace{\mathbb{E}\_{p(x\_t)}\|\nabla\_{x\_t}\log p(x\_t) - \nabla\_{x\_t}\log \widehat{p}(x\_t)\|^2}\_{\text{Fisher divergence between $p(x\_t)$ and $\widehat{p}(x\_t)$}}\Big]$$

Note that in SiD [1], the authors use the time average of Fisher Divergence, which compares only marginal distributions $p(x_t)$ and $\widehat{p}(x_t)$. However, two path measures with the same marginal distributions might not be equal. As a result, the minimization of Fisher Divergence with the teacher does not guarantee that the learned model will transform data in the same way as the teacher model. Nevertheless, Fisher Divergence allows one to build a Generator $x_0 = G_{\theta}(z)$ to produce data $p_{\theta}(x_0) \approx p_{\text{data}}(x_0)$, by matching marginals of data and generated samples. In contrast, we use KL-divergence between two path measures, which is zero if and only if two path measures are identical, since we need to get generations aligned to data coupling.

(2) Acceleration of the unconditional diffusion bridge models.

To obtain PF-ODE for the diffusion bridge model, one needs to subtract forward and reverse SDE drifts [4, end of Sec. 4]. For the conditional case, the drift of a forward process $p(x_0|x_T) \rightarrow \delta(x_T)$ is known analytically from Doob h-transform. In the unconditional case, the drift of a forward unconditional diffusion $p_0(x_0) \rightarrow p_T(x_T)$ is unknown. Hence, to restore PF-ODE in the unconditional case, one must learn an additional teacher model for forward-time translation, which is time-consuming.

We will add that I3SB [3] is used to accelerate the unconditional bridge model. Their sampler coincides with the DBIM sampler but replaces the conditional model with an unconditional one to sample $\widehat{x}_0$. Their sampler provides good quality for moderate NFE (25+) but performs worse than distillation methods in a single NFE regime, e.g., for JPEG-10, their FID is 17, while ours is 3.8 (we will add it to the text).

Concluding remarks. We would be grateful if you could let us know if our explanations have been satisfactory. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.

References:

[1,2,3] the same.

[4] Shi Y. et al. Diffusion Schrödinger Bridge Matching.

Dear Reviewer bxNd, thank you for your comments. Here are the answers to your questions and comments.

(1) It would be better to clarify the reasoning behind the statement in Section 3.2, line 240: 'The key difference in the reformulated problem is that it admits clear gradients of the generator .' For example, explaining that all parts are differentiable would help.

Thank you for this suggestion. In the final version, we will clarify the differentiability of all parts of the final objective.

(2) One curious point is that multi-step distillation is applied to the distillation of DDBM but not to CBD and CBT. What if single-step distillation is applied to the proposed approach? How would the results change in the metric?.

If we correctly understand, you asked about single-step distillation applied to the result of the multi-step distillation. Indeed, we applied multi-step distillation to the original teacher model, e.g., I2SB and DDBM, but not to the already distilled models like CBD or CBT. We did so, since our method is designed for distillation of diffusion-bridge models (like I2SB and DDBM), while CBD and CBT are consistency models obtained from DDBM. Result of multi-step IBMD (ours) distillation is also not a diffusion bridge model. Due to that we do not consider one-step distillation of multi-step distilled models.

(3) It would be better to use the same notation for 'single-step' and 'multi-step' in line 99.

Thank you for this suggestion. We will change it in the final version.

(4) How much time does it take to distill each model each model I2SB and DDBM?

We present the training time of each model below:

About 75% of this training time is used to get the last 10-20% decrease of FID (e.g., drop from 3.6 to 2.5 FID in pooling SR setup or from 4.3 to 3.8 FID in JPEG with $QF=5$), while training for the first 25% of time already provides a good-quality model. On Sketch-to-image and Normal-to-image in multistep regime with 2 NFEs, convergence appears faster than in the corresponding single-step version. We will add the approximate time used for training to Table 7 of Appendix B (Table with all hyperparameters).

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. We are also open to discussing any other questions you may have.

Dear Reviewer AN8r, Thank you for your comments. Here are the answers to your questions and comments.

(1) In my understanding, the second term with \phi here acts as an "expander," which means, G can become trivial (or can collapse) without this term ... In other words, the proposed method trains another DBM just to help train $G$. I'd like to ask whether this interpretation is right, and I'd like to see a deeper discussion in the paper regarding the role of $\phi$.

Yes, this interpretation is correct. To show it more formally, note that the minimal value of the inner problem is the averaged variance of $x_0 \sim p_{\theta}(x_0|x_t, x_T)$:

$$\min_{\phi} \mathbb{E}_{p\_{\theta}(x\_t, t, x\_0, \textcolor{MyRed}{x_T})} \big[\lambda(t) \| \widehat{x}_0^{\phi}(x_t, t, \textcolor{MyRed}{x_T}) - x_0 \|^2 \big]\Big] = \mathbb{E}\_{p\_{\theta}(x\_t, t, x\_0, \textcolor{MyRed}{x\_T})} \big[\lambda(t) \| \mathbb{E}\_{p\_{\theta}(x\_0|x\_t, x\_T)}[x_0] - x_0 \|^2 \big]\Big] =$$ = \mathbb{E}\_{p\_{\theta}(t, x\_t, x\_T)} \Big( \lambda(t) \Big $ \underbrace{\mathbb{E}\_{p\_{\theta}(x\_0|x\_t, x\_T)} \big[ \| \mathbb{E}\_{p\_{\theta}(x\_0|x\_t, x\_T)}[x_0] - x_0 \|^2 \big]\Big]}\_{\text{Variance of } p\_{\theta}(x\_0|x\_t, x\_T)}\Big $ \Big).

For $t=T$, this is directly the variance of the generator $x_0 \sim p\_{\theta}(x_0|x_T)$. Since we are maximizing this part over $\theta$, it enforces the generator to produce more diverse outputs and avoid collapsing.

Following your advice, we will add this discussion on the interpretation of how the auxiliary model $\phi$ helps to train the Generator $\theta$ to the final version of the paper.

(2) How long does it take for training?

We present the training time of each model below:

About 75% of this training time is used to get the last 10-20% decrease of FID (e.g., drop from 3.6 to 2.5 FID in pooling SR setup or from 4.3 to 3.8 FID in JPEG with $QF=5$), while training for the first 25% of time already provides a good-quality model. On Sketch-to-image and Normal-to-image in multistep regime with 2 NFEs, convergence appears faster than in the corresponding single-step version. We will add the approximate time used for training to Table 7 of Appendix B (Table with all hyperparameters).

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. We are also open to discussing any other questions you may have.