SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[67] A. Vahdat, K. Kreis, and J. Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:11287-11302, 2021.
[68] P. Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7): 1661-1674, 2011.
[69] G. Wang, Y. Jiao, Q. Xu, Y. Wang, and C. Yang. Deep generative learning via Schrödinger bridge. In International Conference on Machine Learning, pages 10794–10804. PMLR, 2021.
[70] Y. Wang, X. Wang, A.-D. Dinh, B. Du, and C. Xu. Learning to schedule in diffusion probabilistic models. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 2478-2488, 2023.
[71] Z. Wang, H. Zheng, P. He, W. Chen, and M. Zhou. Diffusion-GAN: Training GANs with diffusion. In International Conference on Learning Representations, 2023.
[72] D. Watson, W. Chan, J. Ho, and M. Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. In International Conference on Learning Representations, 2022.
[73] Y. Yu, X. Si, C. Hu, and J. Zhang. A review of recurrent neural networks: LSTM cells and network architectures. Neural computation, 31(7):1235-1270, 2019.
[74] Q. Zhang and Y. Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34:16280-16291, 2021.
[75] Q. Zhang and Y. Chen. Fast sampling of diffusion models with exponential integrator. International Conference on Learning Representations, 2023.
[76] R. Zhang, P. Isola, A. A. Efros, E. Shechtman, and O. Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586-595, 2018.
A Related Works
In this paper, our focus lies in reducing the discretization error by optimizing the Kullback-Leibler (KL) divergence at each time step between the true sampling chain and the simulated chain. To this end, our method involves fine-tuning pre-trained Diffusion Models (DMs) by minimizing a combination of the contrastive loss and the score matching loss. Thus, our objective function aims to optimize the upper bound of each KL divergence between the true sampling chain and the simulated chain. In line with this design philosophy, our method can be regarded as solving the Schrödinger Bridges [39, 40, 13]. Within this context, extensive research has been conducted to enhance DMs [57, 22, 61, 15]. In [74], a diffusion normalizing flow is proposed, which jointly trains the two neural stochastic differential equations (SDEs) [64] to minimize a common cost function that quantifies the difference between the two. By comparison, the diffusion Schrödinger Bridge optimizes an entropy-regularized transport problem on path spaces, yields diffusions which generate samples from the data distribution in finite time [18, 10]. In [69], they propose to learn a generative model via entropy interpolation [24] with a Schrödinger Bridge. This however comes with a heavy training overhead since they train the model from scratch. By contrast, our method fine-tune pre-trained DMs via reducing the gap between the true sampling chain and the simulated chain.
On the other hand, certain fast sampling methods aim to essentially reduce the discretization error in order to enhance both the sampling speed and the quality of the results. In [56] the authors use a non-Markovian forward noising. The resulted algorithm, DDIM, achieves significant acceleration than DDPMs. More recently, the authors of [44] optimize the backward Markovian process to approximate the non-Markovian forward process and get an analytic expression of optimal variance in denoising process. Concurrent to this work, Zhang and Chen [75] propose diffusion exponential integrator sampler, which leverages a semilinear structure of the learned diffusion process to reduce the discretization error. Orthogonal to those direction, our focus lies in minimizing the discretization error by optimizing the upper bound of the Kullback-Leibler (KL) divergence between the true data distribution and the model distribution.
We follow the regularity assumptions in [60, 43] to prove the Theorem 1. For completeness, we list all these assumptions in this section. We use  $\mathcal{C}$  to denote all continuous functions, and let  $\mathcal{C}^k$  denote the family of functions with continuous  $k$ -th order derivatives.
Assumptions We make the following assumptions throughout the paper:
(i)  $p(\mathbf{x})\in \mathcal{C}^2$  and  $\mathbb{E}_{\mathbf{x}\sim p}\left[\| \mathbf{x}\| _2^2\right] <   \infty .$
(ii)  $\pi (\mathbf{x})\in \mathcal{C}^2$  and  $\mathbb{E}_{\mathbf{x}\sim \pi}\left[\| \mathbf{x}\| _2^2\right] <   \infty .$
(iii)  $\forall t\in [0,T]:\pmb {f}(\cdot ,t)\in \mathcal{C}^1,\exists C > 0\forall \mathbf{x}\in \mathbb{R}^D,t\in [0,T]:\| \pmb {f}(\mathbf{x},t)\| _2\leqslant C\left(1 + \| \mathbf{x}\| _2\right).$
(iv)  $\exists C > 0, \forall \mathbf{x}, \mathbf{y} \in \mathbb{R}^D: \| \pmb{f}(\mathbf{x}, t) - \pmb{f}(\mathbf{y}, t) \|_2 \leqslant C \| \mathbf{x} - \mathbf{y} \|_2.$
(v)  $g\in \mathcal{C}$  and  $\forall t\in [0,T],|g(t)| > 0$
(vi) For any open bounded set  $\mathcal{O}$ ,  $\int_0^T\int_{\mathcal{O}}\| p_t(\mathbf{x})\| _2^2 +Dg(t)^2\| \nabla_{\mathbf{x}}p_t(\mathbf{x})\| _2^2$ $\mathrm{d}\mathbf{x}\mathrm{d}t < \infty$
(vii)  $\exists C > 0\forall \mathbf{x}\in \mathbb{R}^D,t\in [0,T]:\| \nabla_{\mathbf{x}}\log p_t(\mathbf{x})\| _2\leqslant C\left(1 + \| \mathbf{x}\| _2\right)$
(viii)  $\exists C > 0, \forall \mathbf{x}, \mathbf{y} \in \mathbb{R}^D: \| \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) - \nabla_{\mathbf{y}} \log p_t(\mathbf{y}) \|_2 \leqslant C \| \mathbf{x} - \mathbf{y} \|_2.$
$(\mathbf{ix})\exists C > 0\forall \mathbf{x}\in \mathbb{R}^{D},t\in [0,T]:\| s_{\pmb{\theta}}(\mathbf{x},t)\|_{2}\leqslant C\left(1 + \| \mathbf{x}\|_{2}\right).$
(x)  $\exists C > 0, \forall \mathbf{x}, \mathbf{y} \in \mathbb{R}^D: \| s_{\pmb{\theta}}(\mathbf{x}, t) - s_{\pmb{\theta}}(\mathbf{y}, t) \|_2 \leqslant C \| \mathbf{x} - \mathbf{y} \|_2.$
(xi) Novikov's condition:  $\mathbb{E}\left[\exp \left(\frac{1}{2}\int_0^T\|\nabla_{\mathbf{x}}\log p_t(\mathbf{x}) - s_\pmb{\theta}(\mathbf{x},t)\|_2^2\mathrm{d}t\right)\right] < \infty$
(xii)  $\forall t\in [0,T]\exists k > 0:p_t(\mathbf{x}) = O\left(e^{-\| \mathbf{x}\| _2^k}\right)as\| \mathbf{x}\| _2\to \infty .$
Below we provide all proofs for our theorems.
Proof. Let  $\mu$  and  $\nu$  in the Theorem 1 denote the path measures [40, 38] of SDEs in Eq. (1) and Eq. (4) with  $\lambda = 1$  respectively. Due to assumptions (i) (ii) (iii) (iv) (v) (ix) and (x), both  $\mu$  and  $\nu$  are uniquely given by the corresponding SDEs. Consider a Markov kernel  $K(\{x_{t}\}_{t\in}[0,T],y)\coloneqq \delta(z_0=y)$ . Since  $x_0\sim p_0(x_0)$ , and  $x_0^{\mathrm{SDE}}\sim p_0^{\mathrm{SDE}}(x_0^{\mathrm{SDE}})$ , we have the following result
$$
\int K (\{x _ {t} \} _ {t \in [ 0, T ]}, x) d \pmb {\mu} (\{x _ {t} \} _ {t \in [ 0, T ]}) = p _ {0} (x _ {0}),
$$
$$
\int K (\{x _ {t} ^ {\mathrm {S D E}} \} _ {t \in [ 0, T ]}, x) d \pmb {\nu} (\{x _ {t} ^ {\mathrm {S D E}} \} _ {t \in [ 0, T ]}) = p _ {0} ^ {\mathrm {S D E}} (x _ {0} ^ {\mathrm {S D E}}).
$$
Here the Markov kernel  $K$  essentially performs marginalization of path measures to obtain "sliced" distributions at  $t = 0$ . We can use the data processing inequality with this Markov kernel to obtain
$$
\begin{array}{l} D _ {\mathrm {K L}} \left(p _ {0} \| p _ {0} ^ {\mathrm {S D E}}\right) = D _ {\mathrm {K L}} \left(\int K \left(\left\{x _ {t} \right\} _ {t \in [ 0, T ]}, x\right) \mathrm {d} \boldsymbol {\mu} \| \int K \left(\left\{x _ {t} ^ {\mathrm {S D E}} \right\} _ {t \in [ 0, T ]}, x ^ {\mathrm {S D E}}\right) \mathrm {d} \boldsymbol {\nu}\right) \tag {16} \\ \leqslant D _ {\mathrm {K L}} (\boldsymbol {\mu} \| \boldsymbol {\nu}) = \int_ {0} ^ {T} D _ {\mathrm {K L}} (p _ {t} \| p _ {t} ^ {\mathrm {S D E}}) d p _ {t}. \\ \end{array}
$$
From the perspective of the Schrödinger Bridge, the KL divergence between the true data distribution and the model distribution is bounded by the integral of KL at each time step  $t \in [0,T]$  between  $\mu$  and  $\nu$ . On the other hand, the value of any KL is greater than or equal to zero. Hence, we enable also to obtain
$$
D _ {\mathrm {K L}} \left(p _ {0} \| p _ {0} ^ {\mathrm {S D E}}\right) \leqslant D _ {\mathrm {K L}} \left(\boldsymbol {\mu} _ {t} \| \boldsymbol {\nu} _ {t}\right) = \int_ {0} ^ {t} D _ {\mathrm {K L}} \left(p _ {i} \| p _ {i} ^ {\mathrm {S D E}}\right) d p _ {i}. \tag {17}
$$
Recall that by definition  $x_{t} \sim p_{t}(x_{t})$  and  $x_{t}^{\mathrm{SDE}} \sim p_{t}^{\mathrm{SDE}}(x_{t}^{\mathrm{SDE}})$ , and now we assume  $p_{t}^{\mathrm{SDE}}$  is a prior distribution. Leveraging the chain rule of KL divergences [40], we have
$$
\left. D _ {\mathrm {K L}} \left(\boldsymbol {\mu} _ {t} \| \boldsymbol {\nu} _ {t}\right) = D _ {\mathrm {K L}} \left(p _ {t} \| p _ {t} ^ {\mathrm {S D E}}\right) + \mathbb {E} _ {\mathbf {z} \sim p _ {t}} \left[ D _ {\mathrm {K L}} \left(\boldsymbol {\mu} (\cdot \mid \mathbf {x} _ {t} = \mathbf {z}) \| \boldsymbol {\nu} (\cdot \mid \mathbf {x} _ {t} ^ {\mathrm {S D E}}) = \mathbf {z})\right) \right]. \right. \tag {18}
$$
Under assumptions (i) (iii) (iv) (v) (vi) (vii) (viii), the SDE in Eq. (1) has a corresponding reverse-time SDE given by
$$
d x = \left[ \boldsymbol {F} _ {t} x d t - \boldsymbol {G} _ {t} \boldsymbol {G} _ {t} ^ {T} \nabla \log p _ {t} \left(x _ {t}\right) \right] d t + \boldsymbol {G} _ {t} d \omega . \tag {19}
$$
Since Eq. (19) is the time reversal of Eq. (1), it induces the same path measure  $\pmb{\mu}$ . As a result,  $D_{\mathrm{KL}}\left(\pmb {\mu}(\cdot \mid \mathbf{x}_t = \mathbf{z})\| \pmb {\nu}\left(\cdot \mid \mathbf{x}_t^{\mathrm{SDE}}\right) = \mathbf{z}\right)$  can be viewed as the KL divergence between the path measures induced by the following two (reverse-time) SDEs:
$$
d x = \left[ \boldsymbol {F} _ {t} x d t - \boldsymbol {G} _ {t} \boldsymbol {G} _ {t} ^ {T} \nabla \log p _ {t} (x _ {t}) \right] d t + \boldsymbol {G} _ {t} d \omega , x _ {t} = x,
$$
$$
d x = \left[ \boldsymbol {F} _ {t} x d t - \boldsymbol {G} _ {t} \boldsymbol {G} _ {t} ^ {T} s _ {\theta} \left(x _ {t} ^ {\text {S D E}}\right) \right] d t + \boldsymbol {G} _ {t} d \omega , x _ {t} ^ {\text {S D E}} = x.
$$
The KL divergence between two SDEs with shared diffusion coefficients and starting points exists under assumptions (vii) (viii) (ix) (x) (xi) (see in [41, 66]), and can be computed via the Girsanov theorem [48]
$$
\begin{array}{l} D _ {\mathrm {K L}} \left(\boldsymbol {\mu} _ {t} (\cdot | x _ {t} = \mathbf {z}) \| \boldsymbol {\nu} _ {t} (\cdot | x _ {t} ^ {\mathrm {S D E}} = \mathbf {z})\right) \\ = - \mathbb {E} _ {\boldsymbol {\mu} _ {t}} \left[ \log \frac {\mathrm {d} \boldsymbol {\nu} _ {t}}{\mathrm {d} \boldsymbol {\mu} _ {t}} \right] \\ \stackrel {(j)} {=} \mathbb {E} _ {\boldsymbol {\mu}} \left[ \int_ {0} ^ {t} g (i) \left(\nabla_ {\mathbf {x}} \log p _ {i} (\mathbf {x}) - s _ {\boldsymbol {\theta}} (\mathbf {x}, i)\right) \mathrm {d} \overline {{\mathbf {w}}} _ {i} + \frac {1}{2} \int_ {0} ^ {t} g (i) ^ {2} \| \nabla_ {\mathbf {x}} \log p _ {i} (\mathbf {x}) - s _ {\boldsymbol {\theta}} (\mathbf {x}, i) \| _ {2} ^ {2} \mathrm {d} i \right] \\ \stackrel {(j j)} {=} \mathbb {E} _ {\boldsymbol {\mu}} \left[ \frac {1}{2} \int_ {0} ^ {t} g (i) ^ {2} \| \nabla_ {\mathbf {x}} \log p _ {i} (\mathbf {x}) - \boldsymbol {s} _ {\boldsymbol {\theta}} (\mathbf {x}, i) \| _ {2} ^ {2} d i \right] \\ = \frac {1}{2} \int_ {0} ^ {t} \bar {\mathbb {E}} _ {p _ {i} (\mathbf {x})} \left[ g (i) ^ {2} \| \nabla_ {\mathbf {x}} \log p _ {i} (\mathbf {x}) - s _ {\boldsymbol {\theta}} (\mathbf {x}, i) \| _ {2} ^ {2} \right] \mathrm {d} i \\ \stackrel {(j j j)} {\leq} \frac {1}{2} \int_ {0} ^ {T} \mathbb {E} _ {p _ {i} (\mathbf {x})} \left[ g (i) ^ {2} \| \nabla_ {\mathbf {x}} \log p _ {i} (\mathbf {x}) - s _ {\boldsymbol {\theta}} (\mathbf {x}, i) \| _ {2} ^ {2} \right] \mathrm {d} i \\ = \mathcal {J} _ {\mathrm {S M}} (\boldsymbol {\theta}; g (\cdot) ^ {2}), \tag {20} \\ \end{array}
$$
where (i) is due to Girsanov Theorem II [[48], Theorem 8.6.6], and (ii) is due to the martingale property of Ito integrals. Combining Eqs. (16), (18) and (20) completes the proof of Theorem 1.
Though we demonstrate that the combination of score matching losses  $\mathcal{I}_{\mathrm{SM}}(\pmb{\theta};g(\cdot)^2)$  and a KL term  $D_{\mathrm{KL}}(p_t\| p_t^{\mathrm{SDE}})$  is an upper bound of  $D_{\mathrm{KL}}(p_0\| p_0^{\mathrm{SDE}})$ , it is intractable to optimize Eq. (9) due to the unknown function form of  $p_t$ . To circumvent this problem, we combined with the mutual information (MI) theory [50] that MI between  $p_t(x_t)$  and  $p_t^{\mathrm{SDE}}(x_t^{\mathrm{SDE}})$  can be expressed as  $I(p_t^{\mathrm{SDE}}(x_t^{\mathrm{SDE}}),p_j^{\mathrm{SDE}}(x_j^{\mathrm{SDE}}))\leq D_{KL}(p_t^{\mathrm{SDE}}\parallel p_t)$ , where  $j\in [0,T]$  and  $j\neq t$ . The detailed process can be shown as
$$
\begin{array}{l} I (p _ {t} ^ {\mathrm {S D E}} (x _ {t} ^ {\mathrm {S D E}}), p _ {j} ^ {\mathrm {S D E}} (x _ {j} ^ {\mathrm {S D E}})) \equiv \mathbb {E} _ {p (x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}})} \left[ \log \frac {p (x _ {t} ^ {\mathrm {S D E}} \mid p _ {j} ^ {\mathrm {S D E}})}{p _ {t} ^ {\mathrm {S D E}} (x _ {t} ^ {\mathrm {S D E}})} \right] \\ = \mathbb {E} _ {p \left(x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}\right)} \left[ \log \frac {p \left(x _ {t} ^ {\mathrm {S D E}} \mid p _ {j} ^ {\mathrm {S D E}}\right) p _ {t} \left(x _ {t}\right)}{p _ {t} \left(x _ {t}\right) p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right)} \right] \tag {21} \\ = \mathbb {E} _ {p \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right)} \left[ \log \frac {p \left(x _ {t} ^ {\text {S D E}} \mid p _ {j} ^ {\text {S D E}}\right)}{p _ {t} \left(x _ {t}\right)} \right] - D _ {\mathrm {K L}} \left(p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right) \| p _ {t} \left(x _ {t}\right)\right) \\ \leq \mathbb {E} _ {p _ {j} ^ {\mathrm {S D E}} (x _ {j} ^ {\mathrm {S D E}})} \left[ D _ {\mathrm {K L}} \left(p \left(x _ {t} ^ {\mathrm {S D E}} \mid p _ {j} ^ {\mathrm {S D E}}\right) \parallel p _ {t} (x _ {t})\right) \right] \\ = \mathbb {E} _ {p _ {j} ^ {\text {S D E}} (x _ {j} ^ {\text {S D E}})} \left[ D _ {\text {K L}} \left(p _ {t} ^ {\text {S D E}} \| p _ {t}\right) \right]. \\ \end{array}
$$
Therefore, we obtain the upper bound of  $I(p_{t}^{\mathrm{SDE}}(x_{t}^{\mathrm{SDE}}), p_{j}^{\mathrm{SDE}}(x_{j}^{\mathrm{SDE}}))$  given by  $D_{KL}(p_{t}^{\mathrm{SDE}} \parallel p_{t})$ . However, it still intractable since the marginal distribution  $p_{t}(x_{t})$  is inaccessible. To remedy this, when applied the Jensen's inequality to this term, we enable to obtain a upper bound of  $D_{KL}(p_{t} \parallel p_{t}^{\mathrm{SDE}})$  by a InfoNCE loss:  $I_{\mathrm{InfoNCE}}(x_{t}^{\mathrm{SDE}}, x_{j}^{\mathrm{SDE}}) \geq D_{KL}(p_{t} \parallel p_{t}^{\mathrm{SDE}})$  [49], where  $j \in [0, T]$  and  $j \neq t$ . To prove this upper bound in detail, we first review the original lower bound of MI [49]
$$
I \left(x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}\right) \geq \log N - I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}\right), \tag {22}
$$
where  $\log N$  means the total numbers of positive pair  $(x_{i}^{\mathrm{SDE}}, x_{j}^{\mathrm{SDE}})$  when calculating the InfoNCE loss. To show clearly the connection between MI and the InfoNCE loss, we demonstrate the derivation process as follows
$$
\begin{array}{l} I _ {\mathrm {I n f o N C E}} (x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}) = - \underset {X} {\mathbb {E}} \log \left[ \frac {\frac {p (x _ {t} ^ {\mathrm {S D E}} | x _ {j} ^ {\mathrm {S D E}})}{p (x _ {t} ^ {\mathrm {S D E}})}}{\frac {p (x _ {t} ^ {\mathrm {S D E}} | x _ {j} ^ {\mathrm {S D E}})}{p (x _ {t} ^ {\mathrm {S D E}})} + \sum_ {x ^ {-} \in X _ {\mathrm {n e g}}} \frac {p (x ^ {-} | x _ {j} ^ {\mathrm {S D E}})}{p (x ^ {-})}} \right] \\ = \mathbb {E} _ {X} \log \left[ 1 + \frac {p \left(x _ {t} ^ {\text {S D E}}\right)}{p \left(x _ {t} ^ {\text {S D E}} \mid x _ {j} ^ {\text {S D E}}\right)} \sum_ {x ^ {-} \in X _ {\text {n e g}}} \frac {p \left(x ^ {-} \mid x _ {j} ^ {\text {S D E}}\right)}{p \left(x ^ {-}\right)} \right] \\ \approx \underset {X} {\mathbb {E}} \log \left[ 1 + \frac {p \left(x _ {t} ^ {\mathrm {S D E}}\right)}{p \left(x _ {t} ^ {\mathrm {S D E}} \mid x _ {j} ^ {\mathrm {S D E}}\right)} (N - 1) \underset {x ^ {-}} {\mathbb {E}} \frac {p \left(x ^ {-} \mid x _ {j} ^ {\mathrm {S D E}}\right)}{p \left(x ^ {-}\right)} \right] \tag {23} \\ = \underset {X} {\mathbb {E}} \log \left[ 1 + \frac {p \left(x _ {t} ^ {\mathrm {S D E}}\right)}{p \left(x _ {t} ^ {\mathrm {S D E}} \mid x _ {j} ^ {\mathrm {S D E}}\right)} (N - 1) \right] \\ \geq \mathbb {E} _ {X} \log \left[ \frac {p \left(x _ {t} ^ {\text {S D E}}\right)}{p \left(x _ {t} ^ {\text {S D E}} \mid x _ {j} ^ {\text {S D E}}\right)} N \right] \\ = - I \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) + \log (N). \\ \end{array}
$$
Here,  $x^{-}$  are negative instances sampled from training images. Eqs. (23) quickly becomes more accurate as  $N$  increases. At the same time  $\log(N) - I(x_{t}^{\mathrm{SDE}}, x_{j}^{\mathrm{SDE}})$  also increases, so it's useful to use large values of  $N$ . In practice, we choose a bigger batch size for the large  $N$ . Concretely, we obtain much more samples from the given sampling chain to form the positive pairs in each training iteration. In consequence, when combined with Eq. (21), we have
$$
I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) \geq \log N - I \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) = \log N - D _ {\mathrm {K L}} \left(p _ {t} ^ {\text {S D E}} \| p _ {t}\right). \tag {24}
$$
Subsequently, we apply the Jensen's inequality to further scale Eq. (24)
$$
\begin{array}{l} I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) \geq \log N - D _ {\mathrm {K L}} \left(p _ {t} ^ {\text {S D E}} \| p _ {t}\right) \\ = \log N - p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right) \log \frac {p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right)}{p _ {t} \left(x _ {t}\right)} \tag {25} \\ \geq \log N - \log \frac {p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right) p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right)}{p _ {t} \left(x _ {t}\right)} \\ = \log N - 2 \log p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right) + \log p _ {t} \left(x _ {t}\right). \\ \end{array}
$$
However, the KL in Eq. (9) is  $D_{\mathrm{KL}}(p_t \parallel p_t^{\mathrm{SDE}})$  which is not equivalent to  $D_{\mathrm{KL}}(p_t^{\mathrm{SDE}} \parallel p_t)$ . To fill the gap, we scale the former as follows
$$
D _ {\mathrm {K L}} \left(p _ {t} \| p _ {t} ^ {\mathrm {S D E}}\right) = p _ {t} \left(x _ {t}\right) \log \frac {p _ {t} \left(x _ {t}\right)}{p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right)} \tag {26}
$$
$$
\leq \log \frac {p _ {t} (x _ {t}) p _ {t} (x _ {t})}{p _ {t} ^ {\mathrm {S D E}} (x _ {t} ^ {\mathrm {S D E}})} = 2 \log p _ {t} (x _ {t}) - \log p _ {t} ^ {\mathrm {S D E}} (x _ {t} ^ {\mathrm {S D E}}).
$$
Therefore, when combined Eq. (25) and Eq. (26), we have
$$
\begin{array}{l} I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) - D _ {\mathrm {K L}} \left(p _ {t} \| p _ {t} ^ {\text {S D E}}\right) \\ = \log N - 2 \log p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right) + \log p _ {t} \left(x _ {t}\right) - \left[ 2 \log p _ {t} \left(x _ {t}\right) - \log p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right) \right] \\ = \log N - \left[ \log p _ {t} ^ {\mathrm {S D E}} \left(x _ {t} ^ {\mathrm {S D E}}\right) - \log p _ {t} \left(x _ {t}\right) \right] \tag {27} \\ \stackrel {(a)} {=} \log N - \log p _ {t} ^ {\text {S D E}} \left(x _ {t} ^ {\text {S D E}}\right) p _ {t} \left(x _ {t}\right) \\ \geq l o g N, \\ \end{array}
$$
where  $(a)$  is due to both  $p_t^{\mathrm{SDE}}(x_t^{\mathrm{SDE}})\leq 1$  and  $p_t(x_t)\leq 1$  because of the attribute of marginal distribution. Hence, it is reasonable that  $p_t^{\mathrm{SDE}}(x_t^{\mathrm{SDE}})p_t(x_t)\leq 1$  and naturally there is  $\log p_t^{\mathrm{SDE}}(x_t^{\mathrm{SDE}})p_t(x_t)\leq 0$ . Accordingly, we have  $I_{\mathrm{InfoNCE}}(x_t^{\mathrm{SDE}},x_j^{\mathrm{SDE}})\geq D_{\mathrm{KL}}(p_t\parallel p_t^{\mathrm{SDE}})$ . Under such a perspective, we replace  $D_{\mathrm{KL}}(p_t\parallel p_t^{\mathrm{SDE}})$  in Eq. (9) by  $I_{\mathrm{InfoNCE}}(x_t^{\mathrm{SDE}},x_j^{\mathrm{SDE}})$  and we have
$$
D _ {\mathrm {K L}} \left(p _ {0} \| p _ {0} ^ {\mathrm {S D E}}\right) \leq I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}\right) + \mathcal {J} _ {\mathrm {S M}} \left(\theta ; g (t) ^ {2}\right). \tag {28}
$$
In practice, we enable to fine-tune pre-trained DMs with Eq. (28). Moreover,  $\mathcal{J}_{\mathrm{SM}}(\theta ;g(t)^2)$  is almost fixed in each pre-trained DM. Similarly, optimizing Eq. (28) is equivalent to minimizing  $D_{\mathrm{KL}}(p_t\parallel p_t^{\mathrm{SDE}})$ . Therefore, our objective function is essentially solving the Schrödinger Bridges.
C More Experiment Details
In this section, we will provide a more comprehensive account of our experiments, focusing on the implementation of DMs. To conduct our experiments, we utilized NVIDIA A100 GPUs for conducting experiments on CIFAR-10, CelebA, FFHQ, and ImageNet datasets. Additionally, we employed Nvidia V100 GPUs for experiments involving the combination of fast sampling methods. Based on our theoretical analysis, to optimize the KL divergence  $D_{\mathrm{KL}}(p_t \parallel p_t^{\mathrm{SDE}})$ , it is sufficient to add an InfoNCE loss to mitigate the discretization error. As a result, this flexible method can be applied to fine-tuning various pre-trained DMs, both with and without training-free fast sampling algorithms, thereby enhancing sample quality or achieving marginally faster sampling speeds.
To achieve this, one simply needs to include sampling chains given by pre-trained DMs or defined by fast sampling algorithms in the training process. This involves first selecting one sample  $x_{t}^{\mathrm{SDE}}$  and then randomly choose another sample  $x_{j}^{\mathrm{SDE}}$  from the chain to calculate the contrastive loss and form the total loss using Eq. (15). Afterwards, back-propagation through time (BPTT) is utilized to update the parameters of DMs, with the gradient flow directed in the opposite direction of this chain. It is worth noting that we maintain all the training settings of the pre-trained DMs and only modify the part that constructs the contrastive loss, thereby demonstrating the flexibility of our method. Similarly, to showcase the scalability of our method on chains defined by fast sampling methods, we replace the original chains provided by pre-trained DMs with new chains defined by fast sampling methods, while retaining all other settings the same as fine-tuning pre-trained DMs. To fine-tune a pre-trained
DM with a sampling chain of default  $N$  steps in practice, we extend it to a chain of  $2N$ ,  $3N$ , or even  $4N$  steps to construct the contrastive loss. Once the fine-tuning process is completed, we test the performance of the  $N$ -step chain by drawing 50,000 samples from it and measure the widely adopted FID score, negative log-likelihood (NLL), and NFEs, where lower values indicate better performance. Ideally, a default chain in pre-trained DM with more steps will demonstrate better performance when fine-tuned by our method. However, a very long sampling chain suffers from gradient disappearance when BPTT is executed. Hence, we utilize different steps to fine-tune different DMs for different target datasets and demonstrate only the best results in this paper.