SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In other words, given that the gradient descent update step in  $\mathrm{CDenoise}(\boldsymbol{x}_{\mathrm{t}})$  does not leave the clean data manifold, it is guaranteed that the intermediate samples generated by (13) will stay on the correct noisy manifold [6]. In this regard, CDDB can be thought of as the DDB-generalized version of DDS. Similar to DDS, CDDB does not require the computation of heavy U-Net Jacobians and hence introduces negligible computation cost to the inference procedure, while being robust in the choice of step size.
CDDB-deep As shown in DPS and IIGDM, taking deeper gradients by considering U-Net Jacobians is often beneficial for reconstruction performance. Moreover, it even provides way to impose data consistency for non-linear inverse problems, where standard gradient methods are not feasible. In order to devise an analogous method, we take inspiration from DPS, and propose to augment the solver with a gradient step that maximizes the time-dependent likelihood (w.r.t. the measurement)  $p(\pmb{y}|\pmb{x}_t)$ . Specifically, we use the Jensen approximation from [5]
$$
\begin{array}{l} p (\boldsymbol {y} | \boldsymbol {x} _ {t}) = \int p (\boldsymbol {y} | \boldsymbol {x} _ {0}) p (\boldsymbol {x} _ {0} | \boldsymbol {x} _ {t}) d \boldsymbol {x} _ {0} \tag {15} \\ = \mathbb {E} _ {p (\boldsymbol {x} _ {0} | \boldsymbol {x} _ {t})} [ p (\boldsymbol {y} | \boldsymbol {x} _ {0}) ] \simeq p (\boldsymbol {y} | \mathbb {E} [ \boldsymbol {x} _ {0} | \boldsymbol {x} _ {t} ]) = p (\boldsymbol {y} | \hat {\boldsymbol {x}} _ {0 | t}), \\ \end{array}
$$
where the last equality is naturally satisfied from the training objective (9). Using the approximation used in (15), the correcting step under the Gaussian measurement model yields
$$
\nabla_ {\boldsymbol {x} _ {t}} \log p (\boldsymbol {y} | \boldsymbol {x} _ {t}) \simeq \nabla_ {\boldsymbol {x} _ {t}} \| \boldsymbol {y} - \boldsymbol {A} \hat {\boldsymbol {x}} _ {0 | t} \| _ {2} ^ {2}. \tag {16}
$$
Implementing (16) in the place of the shallow gradient update step of Algorithm 1, we achieve CDDB-deep (see Algorithm 2). From our initial experiments, we find that preconditioning with  $A^{\dagger}$  as in IIGDM improves performance by a small margin, and hence use this setting as default.
Model, Dataset For a representative DDB, we choose  $\mathrm{I}^2\mathrm{SB}$  [26] along with the pre-trained model weights for the following reasons: 1) it is open-sourced<sup>4</sup>, 2) it stands as the current state-of-the-art, 3) the model architecture is based on ADM [11], which induces fair comparison against other DIS methods. All experiments are based on ImageNet  $256\times 256$  [10], a benchmark that is considered to be much more challenging for inverse problem solving based on generative models [5], compared to more focused datasets such as FFHQ [20]. We follow the standards of [26] and test our method on the following degradations: sr4x-{bicubic, pool}, deblur-{uniform, gauss}, and JPEG restoration with 1k validation images.
<sup>4</sup>https://github.com/NVlabs/I2SB
Figure 2:  $\mathrm{SR}(\times 4)$ -bicubic reconstruction results. CDDB/CDDB-deep corrects details wrongly captured in I $^2$ SB: line texture in  $1^{\text{st}}$  row, color/texture errors in  $2^{\text{nd}}$  row, topology of wings in  $3^{\text{rd}}$  row (see arrows).
Method	SR(x4)	Deblur
bicubic	pool	gauss	uniform
PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓	PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓
CDDB (ours)	26.41	0.860	0.198	19.88	26.36	0.855	0.184	17.79	37.02	0.978	0.059	5.007	31.26	0.927	0.193	23.15
I2SB [26]	25.22	0.802	0.260	24.13	25.08	0.800	0.258	23.53	36.01	0.973	0.067	5.800	30.75	0.919	0.198	23.01
DPS [5]	19.89	0.498	0.384	63.37	21.01	0.562	0.326	49.34	27.21	0.766	0.244	34.58	22.51	0.565	0.357	60.00
IIGDM [36]	26.20	0.850	0.252	29.36	26.07	0.849	0.256	26.97	-	-	-	-	-	-	-	-
DDRM [21]	26.05	0.838	0.270	46.49	25.54	0.848	0.257	40.40	36.73	0.975	0.071	4.346	29.21	0.901	0.210	19.97
DDNM [40]	26.41	0.801	0.230	38.63	26.04	0.792	0.218	33.15	-	-	-	-	-	-	-	-
DDS [6]	26.41	0.801	0.230	38.64	26.04	0.792	0.218	33.15	33.27	0.945	0.057	6.442	27.88	0.829	0.193	26.07
PnP-ADMM [4]	26.16	0.788	0.350	74.06	25.85	0.733	0.372	72.63	28.18	0.800	0.325	60.27	25.47	0.701	0.416	83.76
ADMM-TV	22.55	0.595	0.493	122.7	22.31	0.574	0.512	119.4	24.67	0.773	0.324	50.74	21.72	0.600	0.491	98.15
Table 2: Quantitative evaluation of SR, deblur task on ImageNet  ${256} \times  {256} - 1\mathrm{k}$  . Bold: Best,under: second best. Colored: DDB methods.
Baselines, Evaluation Along with the most important comparison against  $\mathrm{I}^2\mathrm{SB}$ , we also include comparisons with state-of-the-art DIS methods including DDRM [21], DPS [5], IIGDM [36], DDNM [40], and DDS [6]. For choosing the NFE and the hyper-parameters for each method, we closely abide to the original advised implementation: DDRM (20 NFE), DPS (1000 NFE), IIGDM (100 NFE), DDNM (100 NFE), DDS (100 NFE). We find that increasing the NFE for methods other than DPS does not induce performance gain. We also note that we exclude IIGDM and DDNM for the baseline comparison in the deblurring problem, as directly leveraging the pseudo-inverse matrix may result in unfair boost in performance [29]. We additionally perform comparisons against PnP-ADMM [4] and ADMM-TV, which are methods tailored towards higher SNR. For  $\mathrm{I}^2\mathrm{SB}$  along with the proposed method, we choose 100 NFE for JPEG restoration as we found it to be the most stable, and choose 1000 NFE for all other tasks. Further details on the experimental setup can be found in Appendix C.
Comparison against baselines Throughout the experiments, we thoroughly analyze both distortion and perception of the reconstructions obtained through our method against other DDB, DIS, and iterative optimization methods. Note that for the recent diffusion-based methods, analysis has been
Figure 3: Results on inpainting (Left) and deblurring (Right). For inpainting, boundaries are corrected (row 1) and artifacts are corrected/removed (row 2,3). For deblurring, halo artifacts are corrected, and grid patterns from the background are alleviated.
mainly focused on perceptual metrics [5, 36, 26], mainly because these methods excel on these metrics, but often compromising distortion metrics. DDB methods often take this to the extreme, where one can achieve the best PSNR with very little NFE, and the PSNR consistently degrades as one increases the NFE [9, 26] (See Fig. 1c). Despite this fact, the standard setting in DDB methods is to set a high NFE as one can achieve much improved perceptual quality. On the other hand, conventional iterative methods are often highly optimized for less distortion, albeit with low perceptual quality. While this trade-off may seem imperative, we show that CDDB can improve both aspects, putting it in the place of the state-of-the-art on most experiments (See Tab. 2, 3). A similar trend can be observed in Fig. 2, where we see that CDDB greatly improves the performance of DDB, while also outperforming DIS and iterative optimization methods.
CDDB pushes forward the Pareto-frontier It is widely known that there exists an inevitable trade-off of distortion when aiming for higher perceptual quality [3, 24]. This phenomenon has been reconfirmed consistently in the recent DIS [5, 36] and DDS [9, 26] methods. For DDB, one can flexibly control this trade-off by simply choosing different NFE values, creating a Pareto-frontier with higher NFE tailored towards perceptual quality. While this property is intriguing, the gain we achieve when increasing the
NFE decreases exponentially, and eventually reaches a bound when  $\mathrm{NFE} > 1000$ . In contrast, we show in Fig. 1 that CDDB pushes the bound further towards the optima. Specifically, 20 NFE CDDB outperforms 1000 NFE DDB in PSNR by  $>2$  db, while having lower FID (i.e. better perceptual quality). To this point, CDDB induces dramatic acceleration ( $>50\times$ ) to DDB.
Method	PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓
CDDB (ours)	26.34	0.837	0.263	19.48
I²SB [26]	26.12	0.832	0.266	20.35
IIGDM [36]	26.09	0.842	0.282	30.27
DDRM [21]	26.33	0.829	0.330	47.02
Table 3: Quantitative evaluation of the JPEG restoration (QF = 10) task.
CDDB vs. CDDB-deep The two algorithms presented in this work share the same spirit but have different advantages. CDDB generally has higher speed and stability, possibly due to guaranteed convergence. As a result, it robustly increases the performance of SR and deblurring. In contrast, considering the case of inpainting and JPEG restoration, CDDB cannot improve the performance of DDB. For inpainting, the default setting of  $\mathrm{I}^2\mathrm{SB}$  ensures consistency by iteratively applying replacement, as implemented in [40]. As the measurement stays in the pixel space, the gradients cannot impose any constraint on the missing pixel values. CDDB-deep is useful in such a situation, as the U-Net Jacobian has a global effect on all the pixels, improv
Figure 4: Results on JPEG restoration (QF=10). CDDB recovers texture details (row 1,3), and color details (row 2).
ing the performance by inducing coherence. CDDB-deep also enables the extension to nonlinear inverse problems where one cannot take standard gradient steps. This is illustrated for the case of JPEG restoration in Tab. 3 and Fig. 4, where we see overall improvement in performance compared to  $\mathrm{I}^2\mathrm{SB}$ .
Noise robustness DIS methods are often designed such that they are robust to measurement noise [21, 5]. In contrast, this is not the case for DDB as they are trained in a supervised fashion: If not explicitly trained with synthetic adding of noise, the method does not generalize well to noisy measurements, as can be seen in Fig. 5. On the other hand, note that with CDDB, we are essentially incorporating a Gaussian likelihood model, which naturally enhances the robustness to noise. As a result, while  $\mathrm{I}^2\mathrm{SB}$  tends to propagate noise (best seen in the background), we do not observe such artifacts when using CDDB.
Figure 5: Results on noisy  $\mathrm{SR} \times 4$  reconstruction.  $\mathrm{I}^2\mathrm{SB}$  propagates noise to the reconstruction. CDDB effectively removes noise.
Extension to other related works Going beyond the paired inverse problem setting and considering the Schrödinger Bridge (SB) problem [25, 8], or more generally transport mapping problems [27] between the two unmatched distributions, it is often desirable to control the deviation from the start of sampling. A concrete example would be the case of image-to-image translation [43] where one does not want to alter the content of the image. As CDDB can be thought of as a regularization method that penalizes the deviation from the starting point, the application is general and can be extended to such SB problems at inference time by using the gradients that minimize the distance from the start point. We leave this direction for future work.
Data consistency in supervised learning frameworks The first applications of supervised deep learning to solve inverse problems in medical imaging (e.g. CT [17], MRI reconstruction [39]) mostly involved directly inverting the measurement signal without considering the measurement constraints. The works that followed [1, 14] naturally extended the algorithms by incorporating measurement consistency steps in between the forward passes through the neural network. Analogously, CDDB is a natural extension of DDB but with high flexibility, as we do not have to pre-determine the number of forward passes [1] or modify the training algorithm [14].
In this work, we unify the seemingly different algorithms under the class of direct diffusion bridges (DDB) and identify the crucial missing part of the current methods: data consistency. Our train-free modified inference procedure named consistent DDB (CDDB) fixes this problem by incorporating consistency-imposing gradient steps in between the reverse diffusion steps, analogous to the recent DIS methods. We show that CDDB can be seen as a generalization of representative DIS methods (DDS, DPS) in the DDB framework. We validate the superiority of our method with extensive experiments on diverse inverse problems, achieving state-of-the-art sample quality in both distortion and perception. Consequently, we show that CDDB can push the Pareto-frontier of the reconstruction toward the desired optimum.
Limitations and societal impact The proposed method assumes prior knowledge of the forward operator. While we limit our scope to non-blind inverse problems, the extension of CDDB to blind inverse problems [7, 30] will be a possible direction of research. Moreover, for certain inverse problems (e.g. inpainting), even when do observe improvements in qualitative results, the quantitative metrics tend to slightly decrease overall. Finally, inheriting from DDS/DIS methods, our method relies on strong priors that are learned from the training data distribution. This may potentially lead to reconstructions that intensify social bias and should be considered in practice.
This research was supported by the KAIST Key Research Institute (Interdisciplinary Research Group) Project, by the National Research Foundation of Korea under Grant NRF-2020R1A2B5B03001980, by the Korea Medical Device Development Fund grant funded by the Korea government (the Ministry of Science and ICT, the Ministry of Trade, Industry and Energy, the Ministry of Health & Welfare, the Ministry of Food and Drug Safety) (Project Number: 1711137899, KMDF_PR_20200901_0015), by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST)), and by the Field-oriented Technology Development Project for Customs Administration through National Research Foundation of Korea(NRF) funded by the Ministry of Science & ICT and Korea Customs Service(NRF-2021M3I1A1097938).
[1] Hemant K Aggarwal, Merry P Mani, and Mathews Jacob. MoDL: Model-based deep learning architecture for inverse problems. IEEE transactions on medical imaging, 38(2):394-405, 2018.
[2] Thilo Balke, Fernando Davis, Cristina Garcia-Cardona, Michael McCann, Luke Pfister, and Brendt Wohlberg. Scientific Computational Imaging COde (SCICO). Software library available from https://github.com/lanl/scico, 2022.
[3] Yochai Blau and Tomer Michaeli. The perception-distortion tradeoff. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6228-6237, 2018.
[4] Stanley H Chan, Xiran Wang, and Omar A Elgendy. Plug-and-play ADMM for image restoration: Fixed-point convergence and applications. IEEE Transactions on Computational Imaging, 3(1):84-98, 2016.
[5] Hyungjin Chung, Jeongsol Kim, Michael Thompson McCann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023.
[6] Hyungjin Chung, Suhyeon Lee, and Jong Chul Ye. Fast diffusion sampler for inverse problems by geometric decomposition. arXiv preprint arXiv:2303.05754, 2023.
[7] Hyungjin Chung, Dohoon Ryu, Michael T McCann, Marc L Klasky, and Jong Chul Ye. Solving 3d inverse problems using pre-trained 2d diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[8] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
[9] Maurizio Delbrazio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303.11435, 2023.
[10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. IEEE, 2009.
[11] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[12] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc. In International conference on machine learning. PMLR, 2023.
[13] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(496):1602-1614, 2011.
[31] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
[32] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[33] Chitwan Sahara, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022.
[34] Tim Salimans and Jonathan Ho. Should ebms model the energy or the score? In Energy Based Models Workshop-ICLR 2021, 2021.
[35] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In 9th International Conference on Learning Representations, ICLR, 2021.
[36] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023.
[37] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, volume 32, 2019.
[38] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR, 2021.
[39] Shanshan Wang, Zhenghang Su, Leslie Ying, Xi Peng, Shun Zhu, Feng Liang, Dagan Feng, and Dong Liang. Accelerating magnetic resonance imaging via deep learning. In 2016 IEEE 13th international symposium on biomedical imaging (ISBI), pages 514-517. IEEE, 2016.
[40] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference on Learning Representations, 2023.
[41] Jong Chul Ye. Geometry of Deep Learning. Springer, 2022.
[42] Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE transactions on image processing, 26(7):3142-3155, 2017.
[43] Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2223-2232, 2017.