SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Importantly, the score matching term  $\mathcal{J}_{\mathrm{SM}}(\theta; g(t)^2)$  is stable and almost fixed in the Eq. (10), as we fine-tune the pre-trained DMs that the score network  $s_{\theta}(x_t) \approx \nabla \log p_t(x)$  is matched for almost all  $x \in \mathbb{R}^D$  and  $t \sim U[0, T]$ . In this context, optimizing the Eq. (10) equals to minimize  $\Gamma(x_t^{\mathrm{SDE}}, x_j^{\mathrm{SDE}}; \beta(t))$ , which is the upper bound of  $D_{\mathrm{KL}}(p_t \| p_t^{\mathrm{SDE}})$  [50]. Therefore, minimizing  $D_{\mathrm{KL}}(p_t \| p_t^{\mathrm{SDE}})$  results in a smaller gap between  $x_t$  and  $x_t^{\mathrm{SDE}}$ , indicating that the approximate solution obtained through numerical solvers is closer to the exact solution with the help of contrastive loss. On the other hand, each minimized  $D_{\mathrm{KL}}(p_t \| p_t^{\mathrm{SDE}})$  will decrease the corresponding  $\Delta(x_t^{\mathrm{SDE}}, x_T^{\mathrm{SDE}})$ , contributing to a smaller discretization error between step  $T$  and any step  $t$ . From this perspective, utilizing a contrastive loss to increase the similarity between  $x_j^{\mathrm{SDE}}$  and each  $x_t^{\mathrm{SDE}}$  in a sampling chain is actually minimizing  $D_{\mathrm{KL}}(p_t \| p_t^{\mathrm{SDE}})_{t \in [0, T]}$ . Hence, by minimizing the Eq. (10), one can achieve a better model distribution  $p_0^{\mathrm{SDE}}(x_0^{\mathrm{SDE}})$  with a reduction in all the discretization error present in the sampling chain. Moreover, the decreased discretization error also benefits the faster sampling speed [44, 2, 75, 63].
Contrastive learning has recently achieved remarkable performance [4, 12, 21] and has made significant waves in deep learning for computer vision tasks [35, 51]. These influential works leverage
the contrastive loss to bring similar images closer together in high-dimensional space, resulting in notable improvements in downstream tasks. Building on the insights provided in Section 4, we employ the InfoNCE loss [49] to formulate our objective function. By doing so, we effectively reduce the discretization error between  $x_{t}^{\mathrm{SDE}}$  and  $x_{t - 1}^{\mathrm{SDE}}$  by optimizing the upper bound of  $D_{\mathrm{KL}}(p_t\| p_t^{\mathrm{SDE}})$ . Consequently, numerical solvers are capable of providing more precise solutions when solving Eq. (5), and a fine-tuned DM naturally enhances the quality of generated samples. In the following sections, we will present our method, focusing on constructing contrastive sampling chains and optimizing DMs.
5.1 Contrastive Sampling Chain
As previously mentioned, our objective is to enhance the quality of samples by reducing the discretization error. Additionally, we illustrate that the presence of this discretization error creates a gap between  $x_{t}^{\mathrm{SDE}}$  and  $x_{t}$ . Analogously, optimizing this gap is equivalent to minimizing the corresponding discretization error. While directly decreasing this gap is infeasible, minimizing its upper bound provides an alternative approach to achieve the same outcome. Leveraging Theorem 1 and Eq. (10), we propose a contrastive sampling chain for fine-tuning pre-trained DMs using the contrastive loss. To accomplish this, we construct the contrastive loss and combine it with the score matching loss  $\mathcal{I}_{\mathrm{SM}}(\theta; g(t)^2)$  to jointly update the parameter  $\theta$ .
To construct the contrastive loss, we randomly select an image  $x_{t}^{\mathrm{SDE}}$  and another image  $x_{j}^{\mathrm{SDE}}$  from the defined sampling chain to form a positive pair. Meanwhile, negative instances are sampled from the training images. Next, we extract 128-dimensional latent representations of these images using the pre-trained MoCov2 [11]. The contrastive loss, known as the InfoNCE loss [49], is then computed in the subsequent steps. When applying the pre-trained MoCov2  $E$ , the InfoNCE loss can be expressed as follows:
$$
I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right) = \log \frac {\exp \left(\boldsymbol {E} \left(x _ {t} ^ {\text {S D E}}\right) \cdot \boldsymbol {E} \left(x _ {j} ^ {\text {S D E}}\right) / \tau\right)}{\exp \left(\boldsymbol {E} \left(x _ {j} ^ {\text {S D E}}\right) \cdot \boldsymbol {E} \left(x _ {t} ^ {\text {S D E}}\right) / \tau\right) + \sum_ {k -} \exp \left(\boldsymbol {E} \left(x _ {j} ^ {\text {S D E}}\right) \cdot \boldsymbol {E} \left(x ^ {-}\right) / \tau\right)}, \tag {11}
$$
where  $x_{j}^{\mathrm{SDE}}$  and  $x_{t}^{\mathrm{SDE}}$  form a positive pair which all generated via iteratively calculating Eq. (7) during the fine-tuning process. By comparison,  $x^{-}$  are negative instances which sampled from training images and  $\tau$  is a temperature hyper-parameter. Conceptually, the InfoNCE loss is crafted to bring similar features closer, thereby reducing the distance between  $x_{t}^{\mathrm{SDE}}$  and  $x_{j}^{\mathrm{SDE}}$ . In this manner, the sampling chain will become tighter and the discretization error decreased accordingly because the integration interval between sampling steps decreased.
Though we show that  $I_{\mathrm{InfoNCE}}(x_t^{\mathrm{SDE}}, x_j^{\mathrm{SDE}})$  is an upper bound of  $D_{KL}(p_t \parallel p_t^{\mathrm{SDE}})$  in Section 4, directly using Eq. (11) to fine-tune DMs will destroy the optimal result of the previous score matching in practice. To circumvent this problem, we combine the generative loss Eq. (3) and contrastive loss Eq. (11) to optimize pre-trained DMs
$$
\mathcal {L} _ {\theta} = \mathcal {J} _ {\mathrm {S M}} (\theta ; g (t) ^ {2}) + \beta (t) * I _ {\text {I n f o N C E}} \left(x _ {t} ^ {\text {S D E}}, x _ {j} ^ {\text {S D E}}\right), \tag {12}
$$
where  $\beta (t)$  is the weighting schedule to balance score matching term and contrastive loss term, which will be analyzed in detail in the following subsection. In Eq. (12), we apply the contrastive loss via calculating  $x_{t}^{\mathrm{SDE}}$  and all  $x_{j}^{\mathrm{SDE}}$  in the same sampling chain with Eq. (11). Though we amplify the  $I_{\mathrm{InfoNCE}}$  with  $\beta (t)$ , the combination of score matching term  $\mathcal{J}_{\mathrm{SM}}\left(\theta ;g(t)^{2}\right)$  is still an upper bound of  $D_{\mathrm{KL}}\left(p_0\| p_0^{\mathrm{SDE}}\right)$ .
5.2 Optimization
Considering the previous discussion, it is crucial to strike a balance between the InfoNCE loss and the generative loss. Placing excessive emphasis on the InfoNCE loss may disrupt the stability of the pre-trained generative task, while assigning a high weight to the generative loss may have a negligible impact. To address this, we draw inspiration from the fact that the InfoNCE loss quantifies the similarity between  $x_{t}^{\mathrm{SDE}}$  and  $x_{j}^{\mathrm{SDE}}$ , where the similarity increases as  $x_{j}^{\mathrm{SDE}}$  approaches  $x_{t}^{\mathrm{SDE}}$ . Consequently, it is reasonable to implement a dynamic weighting schedule based on the time step  $t$ . This schedule assigns higher weights to the loss when  $x_{j}^{\mathrm{SDE}}$  is closer to  $x_{t}^{\mathrm{SDE}}$ , and lower weights otherwise. To strike an appropriate balance between the two losses during the refinement of pretrained DMs, we have devised two dynamic weighting schedules: the linear and nonlinear weighting schedules.
Figure 1: Optimizing Diffusion Models via Back Propagation Through Time.
Linear weighting schedule: The weights in this schedule increase progressively with the discretization step as it gets closer to  $x_{j}^{\mathrm{SDE}}$ . For instance, larger weights are assigned to  $x_{t}^{\mathrm{SDE}}$  if it shares a higher similarity with  $x_{j}^{\mathrm{SDE}}$ . In contrast, smaller weights are assigned to the earlier step of the sampling chain. At this extreme, a linear weighting schedule can be expressed as
$$
\beta (t) = \alpha * (T - t), \tag {13}
$$
where  $\alpha$  is a hyper-parameter to scale the weighting schedule.
Nonlinear weighting schedule: The nonlinear weighting schedule works in a similar fashion to the linear weighting schedule, with the only difference being that the weights increase at different rates. Specifically, we calculate the noise ratio in  $x_{t}^{\mathrm{SDE}}$  compared to  $x_{j}^{\mathrm{SDE}}$  using the peak signal-to-noise ratio (PSNR) [76, 14] and subsequently apply the PSNR as the elements of this weighting schedule. Therefore, the nonlinear weighting schedule can be expressed as:
$$
\beta (t) = \alpha * \operatorname {P S N R} \left(x _ {j} ^ {\mathrm {S D E}}, x _ {t} ^ {\mathrm {S D E}}\right). \tag {14}
$$
In general, the nonlinear weighting schedule is more consistent with the principle of the sampling chain in DMs, as the weights gradually remove the noise in the samples. However, in practice, both weighting schedules demonstrate similar effectiveness.
Computational graph: As mentioned earlier, we construct a contrastive sampling chain to fine-tune a pre-trained DM. Though we utilize Eq. (6) to directly obtain the sample  $x_{t}^{\mathrm{SDE}}$ , in practice,  $x_{t}^{\mathrm{SDE}}$  is iteratively solved by Eq. (4). Hence, we need to consider the accumulated discretization error such as  $\Delta(x_{T-1}^{\mathrm{SDE}}, x_{T}^{\mathrm{SDE}})$  and  $\Delta(x_{T-2}^{\mathrm{SDE}}, x_{T-1}^{\mathrm{SDE}})$ , leads to a larger gap between  $x_{t}^{\mathrm{SDE}}$  and  $x_{T}^{\mathrm{SDE}}$ . To address this issue and influence the parameters of previous steps while optimizing the KL divergence  $D_{\mathrm{KL}}(p_{t} \| p_{t}^{\mathrm{SDE}})$  for the current step  $t$ , we propose to fine-tune pre-trained DMs using a gradient propagation mechanism similar to BPTT [73], shown in Figure 1.
Specifically, for a given sampling chain, we randomly select one sample  $x_{t}^{\mathrm{SDE}}$  and another sample  $x_{j}^{\mathrm{SDE}}$ , and calculate the InfoNCE loss with the previously mentioned weighting function. By combining it with the generative loss (Eq. (3)), we employ an optimizer to propagate the gradients of these two losses in reverse order through the chain from the current step  $t$  to the final step  $T$ . In this manner, the gradients influence the entire sampling chain and update the parameters of DMs at each time step  $t$ . Consequently, the training objective can be reformulated as follows:
$$
\mathcal {L} _ {\theta} = \mathcal {J} _ {\mathrm {S M}} (\theta ; \omega (t)) + \Gamma \left(x _ {t} ^ {\mathrm {S D E}}, x _ {j} ^ {\mathrm {S D E}}; \beta (t)\right); t \sim U [ 0, T ], \tag {15}
$$
where  $\Gamma(x_{t}^{\mathrm{SDE}}, x_{j}^{\mathrm{SDE}}; \beta(t))$  represents the InfoNCE loss Eq. (11) with weighting function  $\beta(t)$ . By following this design philosophy, the InfoNCE loss does not solely update parameters independently; rather, it operates as a unified entity. Consequently, we enable to optimize the current KL divergence  $D_{\mathrm{KL}}(p_{t} \| p_{t}^{\mathrm{SDE}})$  and simultaneously influence KL divergences in previous time steps. This approach ensures a cohesive and synchronized optimization process throughout the entire contrastive sampling chain. In summary, our method effectively mitigates the cumulative discretization errors, as each discretization error is correspondingly reduced. Below, we demonstrate experiment results to further prove our analysis.
Table 1: Performance on CIFAR-10. Our method, denoted as  $\mathrm{C} +  +$  ,has better quality than baselines with the same NFEs,and fewer NFEs for the same quality.
Method	Space	NFE↓	NLL↓	FID↓
Unconditional
VDM [36]	Data	1000	2.49	7.41
DDPM [22]	Data	1000	3.75	3.17
iDDPM [47]	Data	1000	3.37	2.90
STDDPM [34]	Data	2000	2.91	2.47
INDM [33]	Latent	2000	3.09	2.28
CLD-SGM [20]	Data	312	3.31	2.25
NCSN++ [61]	Data	2000	3.45	2.20
LSGM [67]	Latent	138	3.43	2.10
LSGM-C++ (Ours)	Latent	100	3.40	2.10
LSGM-C++ (Ours)	Latent	138	3.40	1.99
EDM[29]	Data	35	2.60	2.04
EDM-C++ (Ours)	Data	25	2.55	2.04
EDM-C++ (Ours)	Data	35	2.55	1.88
Conditional
NCSN++G [5]	Data	2000	-	2.25
EDM	Data	35	2.60	1.82
EDM-C++ (Ours)	Data	27	2.55	1.82
EDM-C++ (Ours)	Data	35	2.55	1.73
Table 2: Performance on fast samplers on CIFAR-10 and ImageNet 64x64 with FID reported. DDIM [56], DPM-Solver [44], and DEIS [75] are classical training-free fast samplers. "†" means the actual NFE is smaller than the NFE [44] given in the table. "-" represents that the FID for this NFE is not shown in original papers.
NFE
Method	10	12	15	20	50
DDPM (CIFAR-10) [22]
DDIM [56]	13.36	-	-	6.84	4.67
DDIM-C++ (Ours)	11.52	-	-	6.09	4.12
DPM-Solver-3 [44]	†24.37	8.20	5.73	†5.43	†5.29
DPM-Solver-3-C++ (Ours)	†21.13	7.14	5.22	†5.06	†4.91
SDE VP (CIFAR-10) [61]
DPM-Solver-3	†54.56	6.03	3.55	†2.90	†2.65
DPM-Solver-3-C++ (Ours)	†43.81	5.41	3.25	†2.84	†2.60
DEIS-tAB3 [75]	4.17	-	3.37	2.86	2.57
DEIS-tAB3-C++ (Ours)	4.02	-	3.20	2.77	2.48
IDDPM (ImageNet) [47]
DPM-Solver-3	†57.48	24.62	19.76	†18.95	†17.52
DPM-Solver-3-C++ (Ours)	†50.63	22.65	19.45	†18.73	†17.49
IDDPM (ImageNet Conditional) [47]
DEIS-tAB3	6.65	3.99	†3.67	3.10	2.69
DEIS-tAB3-C++ (Ours)	6.59	3.91	†3.60	3.07	2.67
DPM-Solver-2 [44]	7.93	5.36	†4.46	3.42	2.82
DPM-Solver-2-C++ (Ours)	7.78	5.22	†4.38	3.36	2.80
In this section, we demonstrate effectiveness of the contrastive sampling chain by experimental results, namely higher image quality, better log-likelihood, or slightly faster sampling speed. Comprehensive experiments are conducted on various datasets, including CIFAR-10, CelebA/FFHQ 64x64, and ImageNet 64x64. We first utilize contrastive sampling chain to improve pre-trained DMs, for which we select EDM [29] and LSGM [67] on CIFAR-10, EDM [29] on FFHQ, and STDDPM [34] on CelebA. Additionally, we further verify the performance of our method with sampling chains defined by training-free fast sampling methods, such as DDIM [56], DPM-Solver [44], and DEIS [75], when they are combined with various pre-trained DMs. Specifically, we combine these fast sampling methods with IDDPM [47] on ImageNet 64x64, DDPM [22] and SDE VP [61] on CIFAR-10. Our objective is not to focus on achieving state-of-the-art metrics in generative models, but rather to demonstrate the significant performance of our method in enhancing pre-trained DMs, either combined with or without training-free fast-sampling methods.
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 ensuring fair comparison and demonstrating the flexibility of our method. Similarly, to showcase the applicability of our method to sampling chains defined by fast sampling methods, we replace the original chains provided by pre-trained DMs with new chains for fast sampling, while retaining all other settings the same as fine-tuning pre-trained DMs. Once the fine-tuning process is completed, we test the performance of the refined model by drawing 50,000 samples from it and measuring the widely adopted Fréchet Inception Distance (FID) score, Negative Log-Likelihoods (NLL), and Neural Function Evaluations (NFEs), where lower values indicate better performance. Moreover, we also present the generated images for qualitative comparison, shown in Figure 2.
6.1 Performance on Pre-trained DMs
We first showcase the performance of our method in refining the original pre-trained DMs. Tables 1, 4 and 5 present the results of our method on CIFAR-10, FFHQ and CelebA respectively. Our method can achieve better generation quality than baselines when using the same NFEs. On the other hand, our method requires less NFEs to achieve the same quality as baselines. To evaluate on CIFAR-10, we apply our method to LSGM [67] and EDM [29] under unconditional or conditional settings. For a fair comparison, we report performances on EDM [29] under the random seed from [32, 63]. Specifically, our method improves LSGM from 2.10 FID to 1.99 FID and achieves slightly better NLL with 138 NFEs, while requiring only 100 NFEs to attain the same 2.10 FID. Moreover, we enhance EDM from 2.04 FID to 1.88 FID and reduce the NLL from 2.60 to 2.55 with 35 NFEs, while requiring only 25 NFEs for the same 2.04 FID. Furthermore, we improve the performance on conditional EDM, with 1.82 FID reduced to 1.73 FID using 35 NFEs or maintaining the same 1.82 FID with only 27 NFEs. For the evaluation on CelebA and FFHQ, we apply our method to refine STDDPM [34] and EDM
Figure 2: Randomly selected ImageNet Samples from IDDPM Improved by Our Method.
[29] respectively. Specifically, our method improves the performance of STDDPM on CelebA 64x64 from 1.90 FID to 1.73 FID with the same 131 NFEs, and achieves the same 1.90 FID with fewer NFEs (100). Similarly, we significantly enhance the performance of EDM on FFHQ 64x64 from 2.39 FID to 2.07 FID with the same 79 NFEs, and achieve the same 2.39 FID with fewer NFEs (63).
6.2 Performance on Fast Samplers
To demonstrate the applicability of our approach to sampling chains defined by training-free fast sampling methods, we carry out a sequence of experiments on CIFAR-10 and ImageNet 64x64, with detailed results in Table 2. We test our method on different fast samplers when combined with various pre-trained DMs. We utilize DDIM [56], DPM-Solver [44], and DEIS [75] to replace the original chains given by DDPM [22], SDE VP [61] and IDDPM [47]. Subsequently, for a pre-trained model with sampling chain defined by these fast samplers, we refine it with Eq. (15). In this manner, our method remarkably improves the results as illustrated in Table 2. For example, our method obtains better sample quality when applied to refine the SDE VP with DEIS. Concretely, our method improves the FID from 4.17 to 4.02 in 10 NFEs and increase the 2.57 FID to 2.48 FID in 50 NFEs. In comparison, the improvement on ImageNet is not significant as our main objective is to demonstrate the effectiveness of our method rather than achieving state-of-the-art performance metrics. For instance, when we apply our method to IDDPM with a conditional setting and replace the original chain with DEIS, we observe improvements in FID. With 10 NFEs, the FID decreases from 6.65 to 6.59, and with 14 NFEs, the FID decreases from 3.67 to 3.60. It is worth noting that the symbol "†" indicates that the actual NFE is smaller than the NFE reported in the table provided by Lu et al. [44]. In summary, our method seamlessly integrates with training-free fast sampling algorithms and enables us to achieve better overall performance.
6.3 Ablation Study
We conduct ablation studies in Table 3 to assess the impact of different techniques. Table 3 compares four different settings, i.e., the proposed contrastive loss with a single pair  $(x_{t}, x_{j})$ , a variant of contrastive loss with all samples  $\sum_{i=1, i \neq t}^{T}(x_{t}, x_{j})$  in the chain, the proposed contrastive loss with BPTT, and naive fine-tuning without contrastive loss for the same number of epochs as the previous
Table 3: FID Comparison of Ablation Study with Different Weighting Schedules on CIFAR-10. To prove the performance of BPTT and contrastive guided technique respectively, we conduct ablation study with three different settings. "EDM Baseline" [29] is the pre-trained model, "Naive Fine-Tuning" means just train more epochs as other settings. For our techniques: "Contrastive Loss" means fine-tune the model with only contrastive loss and no BPTT, "Contrastive Loss (All Steps)" is calculating the contrastive loss with all the discretization step in contrastive sampling chain and no BPTT, "Contrastive Loss + BPTT" is our setting with best performance.
Weighting Schedule
Method	Fixed	Linear	Nonlinear
EDM Baseline [29]	2.04	-	-
+ Naive Fine-Tuning	2.04	-	-
+ Contrastive Loss	2.04	1.98	1.99
+ Contrastive Loss (All Steps)	2.18	2.32	2.39
+ Contrastive Loss + BPTT	2.00	1.88	1.88
Table 4: FID comparison on FFHQ 64x64.
Method	NFE↓	FID↓
EDM [29]	79	2.39
EDM-C++ (Ours)	63	2.39
EDM-C++ (Ours)	79	2.07
Table 5:FID comparison on CelebA 64x64.
Method	NFE↓	FID↓
DDPM++ [61]	131	2.32
Soft Diffusion [16]	300	1.85
INDM [33]	132	1.75
Diffusion StyleGAN2 [71]	1	1.69
STDDPM [34]	131	1.90
STDDPM-C++ (Ours)	100	1.90
STDDPM-C++ (Ours)	131	1.73
settings. We also test three weighting schedules, including two previously mentioned dynamic weighting schedules and no schedule with a fixed weight during the entire process. As we mentioned in Section 5.2, calculating the contrastive loss with all samples in the chain disregards the consistency of the entire chain, which can destabilize pre-trained DMs on generative tasks. Therefore, it is reasonable that we obtained a worse FID when fine-tuning DMs with this setting. Moreover, the two dynamic weighting schedules show equivalent performance, both of which are much better than the fixed weight. Our results demonstrate that our contrastive method trained with BPTT for updating parameters can yield the best performance.
In this paper, we demonstrate the effectiveness of optimizing the KL divergence between the true sampling chain and the simulated chain at each time step in reducing the discretization error associated with numerical solvers used for solving SDEs. Our theoretical analysis supports the use of our objective function as an upper bound of the KL divergence between the data distribution and the model distribution. Notably, optimizing our objective function is equivalent to minimizing the KL divergence between the true sampling chain and the simulated chain at each time step. To address this, we propose a contrastive sampling chain that leverages the derived upper bound to reduce the discretization error. Additionally, we introduce the use of backpropagation through time (BPTT) to propagate gradients in the reverse direction of the sampling chain, and we design dynamic weighting schedules to enhance the stability of the refinement process. Our empirical results demonstrate that our approach significantly improves both the sample quality and the log-likelihood, while slightly accelerating pre-trained DMs without compromising image quality.
Limitations and broader impact: Although our method has shown significant improvements, there remains potential for further optimization of our method. For instance, implementing the BPTT computational graph demands a significant amount of GPU memory. Additionally, obtaining analytical solution of our weighting function will undoubtedly tighten the upper bound of the KL divergence between the data distribution and the model distribution. On the other hand, the issue of inefficient sampling remains a major obstacle to the practical application of DMs. It is reasonable to expect that the sampling speed can be greatly improved by a method that effectively optimizes the discretization error. However, it is important to acknowledge that the generation of deepfake images using our method also carries the potential risk of negative misuse of this technology.
This work was supported by the Natural Science Foundation of China under grant 61836016. Chang Xu was supported in part by the Australian Research Council under Projects DP210101859 and FT230100549. The AI training platform supporting this work were provided by High-Flyer AI (Hangzhou High-Flyer AI Fundamental Research Co., Ltd.). This work was also supported in part by the High Performance Computing Center of Central South University. This work was supported in part by the China Scholarship Council.