Optical Coherence Tomography Harmonization via Dual Diffusion Implicit Bridges

Abstract

Image harmonization (IH) makes images from different domains consistent, enabling multi-domain comparison with reliable quantitative measurements. Optical coherence tomography (OCT) could benefit from IH, as variations in OCT systems cause significant discrepancies in image quality due to speckle and noise. However, IH approaches typically rely on paired images during the training phase, and such images are not available in OCT data due to sparse image acquisition. A Schrödinger bridge (SB) finds an optimal coupling between arbitrary probability spaces with respect to a reference path measure and can be applied to unpaired IH. SBs are the core of many generative diffusion models, when one of the probability spaces is easy to sample from (i.e., Gaussian). SBs have been applied to natural images for image-to-image translation; however, they have not been applied to OCT IH because a SB does not guarantee anatomical consistency. In this paper, we use a dual diffusion implicit bridge (DDIB) for OCT IH, which finds independent SBs between arbitrary domains and a common Gaussian probability space. We generalize the assumption of DDIB and use segmentation masks as a condition to improve anatomical consistency. We conduct DDIB experiments with and without segmentation masks as a condition and analyze performance in terms of anatomical consistency and harmonization quality. The method does not require any paired training data and can in principle be quickly adapted to new domains.

1. INTRODUCTION

Images from different domains can have different appearances of the same scene. Image harmonization (IH) is a technique to standardize or transform images from different domains so that they are consistent and comparable. This process is essential in medical image analysis, as images from different scanners, settings, or acquisition protocols can have variations in contrast, noise, and other technical aspects, which affect downstream processing tasks such as segmentation, classification, registration, and thus affects measurement accuracy, which can impact diagnosis. IH helps mitigate these image variations, leading to more accurate analysis and diagnosis, especially when comparing images from different devices, or when tracking changes over time in longitudinal studies.

One imaging modality that would benefit from IH is optical coherence tomography (OCT),¹ which is widely used in medical fields like ophthalmology. OCT uses low-coherence interferometry to visualize micro structures in biological tissue, such as the retina and cornea. However, OCT imaging systems can introduce significant inconsistencies, including variations in contrast, speckle, and noise, as a result of different hardware configurations and scanning settings. For example, Zeiss Cirrus OCT provides fast scanning of the retina but suffers from low contrast and signal-to-noise ratio (SNR) due to speckle and noise. Heidelberg Spectralis OCT acquires multiple images around the same cross section and averages them to reduce speckle and noise, and thus has superior contrast and SNR, but fewer acquired scans in a volume. See Fig. 1(a) and Fig. 1(b) for comparison of the two scanners. These inconsistencies make it challenging to analyze multi-domain OCT datasets, as they obscure meaningful comparisons and hinder accurate interpretation of results.² Moreover, due to the sparsity of OCT data along the slow axis and the difficulty of aligning two 3D OCT volumes, it is not feasible to obtain paired images of the same anatomy from both scanners for supervised harmonization. While some studies have shown that deformable registration and generative models can enhance OCT volume interpolation along the slow axis,³ their potential for generating paired data has yet to be fully explored.

Figure 1.

(a) A Cirrus OCT B-scan. (b) A Spectralis OCT B-scan. (c) A dual diffusion implicit bridge (DDIB) for OCT IH. $C$ stands for Cirrus, $S$ stands for Spectralis, and $G$ stands for a common Gaussian probability space with zero mean and unit variance. The arrow between C and G and the arrow between S and G indicate independent SBs between arbitrary domains $(C\&S)$ and the common Gaussian probability space $G$. The shorter arrows with time $t_c$ and $t_s$ indicate a deterministic and invertible mapping between $C$ and $S$ via DDIB.

A Schrödinger bridge (SB)⁴ is an optimal path measure between arbitrary source and target probability spaces with respect to a reference path measure. It has been used in image restoration⁵ and image-to-image translation⁶ for natural images. It is the core of many generative diffusion models, such as denoising diffusion probabilistic models (DDPM),⁷ score-based generative models (SGM),⁸ and flow matching.⁹ These generative diffusion models are special cases of the SB approach, when one of the probability spaces is easy to sample from (i.e., Gaussian). However, traditional SB solvers^10–12 are computationally expensive, because of the high dimensional nature of the path measure. The diffusion SB (DSB)¹³ decomposes a Markov chain into smaller components to reduce computational complexity, and the simplified diffusion SB (S-DSB)¹⁴ further improves on this idea. However, both methods are slow to converge and do not generalize quickly to new domains. A dual diffusion implicit bridge (DDIB)⁶ finds two independent SBs that map two arbitrary probability spaces to a common Gaussian probability space, which is easier to train, avoids training on paired data, and adapts quickly to new domains.

SBs have seen no application in OCT IH because they do not guarantee anatomical consistency between original and harmonized images, which is critical in medical image analysis. In this paper, we use a DDIB for OCT IH. There are three main contributions in this work: 1) we generalize the assumptions of DDIB and apply it to OCT IH; 2) we conduct DDIB experiments with and without segmentation masks as a condition, and analyze the DDIB performance in terms of both anatomical consistency and IH quality versus the diffusion steps; and 3) we analyze the diffusion step discrepancies in two SBs, and find ways which may further improve the harmonization quality.

2. BACKGROUND

A SB is an optimal path measure between two arbitrary probability spaces with respect to a reference path measure (see Appendix A). Many different generative diffusion models, such as DDPM,⁷ SGM,⁸ and flow matching,⁹ are special cases of the SB approach when one of the probability spaces is Gaussian (see Appendix B). Like many generative diffusion models described by a stochastic differential equation (SDE), in SB, the corresponding probability flow ordinary differential equation (ODE) gives a deterministic and invertible mapping between the two probability spaces by the forward or backward integration (see Appendix C),

$$x_1=x_0+{\int }_0^1{v}_t\left(x_t\right)dt,$$ (1)

where $x_t\in {\mathbb{R}}^n,v:[0,1]\times {\mathbb{R}}^n\to {\mathbb{R}}^n$ is a time-dependent velocity field, $t\in [0,1]$ is time, and $n$ is the data dimension. The velocity field $v$ can be obtained by training a neural network that approximates the score function using different generative diffusion models (see Appendix D).

3. METHOD

Generalized assumption of DDIB.

A DDIB⁶ finds two independent SBs that map two domains $C$ and $S$ to a common Gaussian probability space $G$ as shown by the blue arrows in Fig. 1(c). It achieves image-to-image translation by using the velocity fields $v^c$ and $v^s$ and integration in Eq. 1. The DDIB assumes that,

$$x_0^c+{\int }_0^1{v}_t^c\left(x_t^c\right)dt=x_0^s+{\int }_0^1{v}_t^s\left(x_t^s\right)dt,$$ (2)

i.e., images $x_0^c$ and $x_0^s$ from $C$ and $S$ map to the same latent element in the common Gaussian probability space $G$. This is a strong assumption that may cause anatomical inconsistency if applied directly to medical IH. This is because forward integration resembles adding noise to images. Both noise and $G$ are zero mean and white, and the source distribution always converge to $G$ after adding enough noise. Integrating over the entire path (i.e. [0,1]) introduces significant noise and essentially decimates any meaningful anatomical information. To address this problem, we assume that there is an optimal time step along the path, where the anatomical information is mostly preserved, and the contrast difference is indistinguishable between Cirrus and Spectralis B-scans due to the addition of the noise. Therefore, we generalize the DDIB assumption to,

$$x_0^c+{\int }_0^{t_c}{v}_t^c\left(x_t^c\right)dt=x_0^s+{\int }_0^{t_s}{v}_t^s\left(x_t^s\right)dt,$$ (3)

i.e., images $x_0^c$ and $x_0^s$ map to the same latent point along the diffusion trajectory for some $t_c,t_s\in [0,1]$ as shown by the red arrows in Fig. 1(c). Our formulation degenerates into Eq. 2 when $t_c=t_s=1$. In general, we do not assume $t_c=t_s$, and we need to find an appropriate $t_c$ and $t_s$ to map between the two domains at the population level. This is because Cirrus B-scans are noisier than Spectralis B-scans, and the amount of noise added to Cirrus and Spectralis B-scans to make them indistinguishable should be different.

Generative diffusion models.

For both Cirrus and Spectralis B-scans, we trained two independent generative diffusion models using diffusion model U-Nets,¹⁵ one without segmentation masks as a condition, and another with segmentation masks as a condition. Therefore, there are four independent generative diffusion models that are trained. For the generative diffusion model, we use a variance perserving SDE and a linear beta scheduler, as proposed in DDPM,⁷ to learn the velocity field $v^c$ and $v^s$. After training, we use a DDIM¹⁶ sampler to solve the probability flow ODE of DDIB in Eq. 3. For simplicity, DDIB-o refers to the DDIB without segmentation masks as a condition, and DDIB-w refers to the DDIB with segmentation masks as a condition. The segmentation mask is required and obtained by segmenting the source B-scans during DDIB-w inference. The probability flow ODE is invertible by construction and thus cycle-consistent. However, in practice, it accumulates discretization errors from the DDIM solver. Therefore, to reduce errors, we do not skip time steps in the DDIM.

4. EXPERIMENTS

Dataset.

Our in-house OCT dataset consists of 18,000 Cirrus B-scans sampled from 422 Cirrus OCT volumes, and 18,000 Spectralis B-scans sampled from 368 Spectralis OCT volumes. The 422 Cirrus OCT volumes come from 212 subjects, representing 422 unique eyes. The 368 Spectralis OCT volumes come from 174 subjects, representing 288 unique eyes. We divided the data into a training set, a validation set, and a testing set. The training set contains 15,000 Cirrus B-scans and 15,000 Spectralis OCT B-scans. Both the validation set and the testing set contain 1,500 Cirrus B-scans and 1,500 Spectralis B-scans. For preprocessing, we cropped both Cirrus and Spectralis B-scans to reduce the background areas and centered the retina in the image, and resized them to 128 × 512 pixels, to ensure that the cropped and resized Cirrus and Spectralis B-scans have the same physical resolutions.

Time parameter scanning.

Finding the proper time parameter $t_c$ and $t_s$ in Fig. 1(c) requires a two-dimensional grid search and is computationally expensive. To speed this up, we first restrict $t_c=t_s=t$ and find the proper $t$, because we expect $t_c$ and $t_s$ to be relatively close. We scan $t$ from 0 to 1 with a step size of 0.025. For each $t$, we run both Cirrus to Spectralis IH and Spectralis to Cirrus IH using DDIB-w and DDIB-o in Eq. 3. We segment the original and harmonized OCT images using a deep learning-based retinal OCT segmentation algorithm,^17,18 and calculate the mean absolute distance (MAD) between nine segmented boundaries. A smaller MAD means the retinal layers are better aligned before and after harmonization. We also calculate the Fréchet inception distance (FID)^19,20 between the harmonized and real images in the validation set. We use the MAD and FID values for different $t$’s to find out the proper time parameter when $t_c=t_s$.

Time classifier models.

To find the correct relationship between $t_c$ and $t_s$, we trained two independent time classifier models using vision transformers (ViT)²¹ on Cirrus and Spectralis B-scans, respectively. For time classifier prediction, we divided the time $t$ in the range [0, 1] into 1,000 discrete classes. During training, we randomly sampled $t$ from these 1,000 classes assuming uniform distribution and use them as ground truth time labels. We then simulated the noisy Cirrus or Spectralis B-scans at corresponding time $t$ in the diffusion trajectory. The ViTs predict time logits from simulated noisy Cirrus and Spectralis B-scans, and the loss function is set to be the cross entropy between the predicted time logits and the ground truth time labels. After training both time classifiers on the training set, we first ran the Cirrus time classifier on the Cirrus testing set and the Spectralis time classifer on the Spectralis testing set to check their performance. We then ran the Cirrus time classifier on the Spectralis testing set and the Spectralis time classifer on the Cirrus testing set to check if there is a bias between $t_c$ and $t_s$. We check this through a linear regression between the ground truth time and the predicted time. If there is no bias, the linear fit shows no offset. If the bias is present, we adjust $t_c$ and $t_s$ accordingly to compensate the bias.

5. RESULTS

Time parameter scanning.

We scanned the time parameter $t$ assuming that $t_c=t_s=t$, and the MAD and FID results for different $t$’s are shown in Fig. 2. From Fig. 2(a) and Fig. 2(b), we see that the anatomical structure slowly deviates as $t$ increases for both DDIB-o and DDIB-w. The DDIB-w preserves the anatomy much better than DDIB-o. From Fig. 2(c) and Fig. 2(d), we see that FID generally improves as $t$ increases, but degrades as $t$ gets close to 1. The DDIB-w achieves better FID than DDIB-o when $t<0.600$ but has worse FID when $t>0.750$. Both the blue dashed line and the red solid line are above the green dashed line for all $t$’s, which indicates that the IH quality does not match the original distributions for both Cirrus to Spectralis and Spectralis to Cirrus IH.

Figure 2.

(a) Mean absolute distance (MAD) between the segmentation boundaries of harmonized images and original images for Cirrus to Spectralis IH. (b) MAD between the segmentation boundaries of harmonized images and original images for Spectralis to Cirrus IH. The blue dashed line is the median MAD value for DDIB-o. The red solid line is the median MAD value for DDIB-w. The green area indicates the interquartile range between the 25% (Q1) and 75% (Q3) quartiles. (c) Fréchet inception distance (FID) between harmonized images and the Spectralis validation set for Cirrus to Spectralis IH. (d) FID between harmonized images and the Cirrus validation set for Spectralis to Cirrus IH. The blue dashed line is the FID value for DDIB-o. The red solid line is the FID value for DDIB-w. The green dashed line are the baseline FID values between the testing set and the training set for Spectralis and Cirrus, respectively.

Time classifier models.

We test both the Cirrus and Spectralis time classifiers on the Cirrus and Spectralis testing sets, and the results are shown in Fig. 3. From Fig. 3(a) and Fig. 3(b), we see that both time classifiers achieve good prediction performance with negligible bias when used on the correct data domains. However, from Fig. 3(c) and Fig. 3(d), we see that there is a systematic although small bias around 0.014 between the predicted time and the ground truth time when using the Cirrus time classifier on Spectralis B-scans and when using the Spectralis time classifier on Cirrus B-scans. This is because original Cirrus B-scans are noisier than original Spectralis B-scans. We take into account this bias and set $t_c=0.493$ and $t_s=0.507$, and compare it with $t_c=t_s=0.500$ for Spectralis to Cirrus IH. The results show that for DDIB-o, the FID improves from 60.97 to 50.74; and for DDIB-w, the FID improves from 47.76 to 40.81. This suggests that finding the optimal relationship between $t_c$ and $t_s$ can further improve the DDIB IH quality.

Figure 3.

Time classifiers. (a) Cirrus time classifier on Cirrus B-scans. (b) Spectralis time classifier on Spectralis B-scans. (c) Spectralis time classifier on Cirrus B-scans. (d) Cirrus time classifier on Spectralis B-scans. Orange dots are data samples $(N=20,000)$, and blue line is the linear regression.

Qualitative comparison.

We show some qualitative results for both Cirrus to Spectralis IH in Fig. 4 and Spectralis to Cirrus IH in Fig. 5. From Fig. 4 and Fig. 5, we can see that when $t_c=t_s=0.250$, the results have good anatomical consistency, but poor IH quality because the harmonized images look closer to the original domain than the target domain. When $t_c=t_s=0.750$ and when $t_c=t_s=1.000$, the results have good IH quality, but poor anatomical consistency. This can be seen from the harmonized results using DDIB-o, where the anatomy deviates a lot from the original images as shown by the blue arrows. This can also be seen from the harmonized images using DDIB-w, where some vessels and shadows are created which do not exist in the original images as shown by the yellow arrows. This is because the segmentation mask used in the condition contain only the eight retinal layers, but no vessels or shadows. When $t_c=t_s=0.500$ and when $t_c=0.493,t_s=0.507$ (highlighted in red), the results achieve a good trade-off between anatomical consistency and IH quality.

Figure 4.

Qualitative comparison of Cirrus to Spectralis IH. Columns #1 to #3 show original images, harmonized images using DDIB-o, and harmonized images using DDIB-w, respectively. Rows #1 to #5 show harmonization results for different $t_c$’s and $t_s$’s, respectively. Red rectangles highlight the results when $t_c\ne t_s$. Red solid lines and yellow dashed lines are the retinal layer segmentation on the original and harmonized images, respectively. Blue and yellow arrows indicate the anatomy shifts.

Figure 5.

Qualitative comparison of Spectralis to Cirrus IH. Columns #1 to #3 show original images, harmonized images using DDIB-o, and harmonized images using DDIB-w, respectively. Rows #1 to #5 show harmonization results for different $t_c$’s and $t_s$’s, respectively. Red rectangles highlight the results when $t_c\ne t_s$. Red solid lines and yellow dashed lines are the retinal layer segmentation on the original and harmonized images, respectively. Blue and yellow arrows indicate the anatomy shifts.

6. CONCLUSION

In this paper, we use a DDIB for OCT IH, which finds independent SBs between arbitrary domains and a common Gaussian probability space. We generalize the assumption of DDIB and use segmentation masks as a condition to improve anatomical consistency. We conduct DDIB experiments with and without segmentation masks as a condition and analyze its performance in terms of anatomical consistency and harmonization quality, and find ways to further improve the harmonization quality. A unique element of this work is the consideration of the difference between $t_c$ and $t_s$. The assumption that these values are the same is not correct and may be important in medical imaging scenarios in which the images under consideration for harmonization are more disparate.

There are several limitations of this work. First, we only used a linear beta scheduler for our DDPM. Other schedulers such as a cosine scheduler or other generative diffusion models need to be explored. Second, the relationship between $t_c$ and $t_s$ is obtained on the population level and by a simple linear regression model; their relationship can be nonlinear and adapted to each individual data to be more effective. Third, the anatomical structures such as shadows and vessels still deviate from the original using DDIB-w for large $t$ because the segmentation mask does not contain these structures. This suggests that additional strategies should be employed to improve anatomical consistency.

APPENDIX A. SCHRÖDINGER BRIDGE

Let $\Omega =C\left([0,1],{\mathbb{R}}^n\right)$ be the set of all continuous ${\mathbb{R}}^n$-valued paths on the unit time interval [0, 1], we construct the probability space $(\Omega ,ℱ,P)$ by adding the $\sigma$-algebra $ℱ\subseteq 2^{\Omega }$ and the path measure $P:ℱ\to [0,1]$. The Schrödinger bridge problem finds the optimal path measure $p$ with respect to a reference measure $q$ by minimizing their Kullback–Leibler (KL) divergence,

$$p=\underset{\pi \in \Pi \left({\mu }_0,{\mu }_1\right)}{\text{arg}\text{min}}{D}_{KL}(\pi \Vert q),$$ (4)

where $\Pi \left({\mu }_0,{\mu }_1\right)$ is the subset of path measures with its marginal densities at $t=0$ and $t=1$ being ${\mu }_0$ and ${\mu }_1$, respectively, and $q$ is restricted to the subset of path measures of random process $Q_t$ described by the forward stochastic differential equation (SDE),

$$d{Q}_t=f_t\left(Q_t\right)dt+g_{t}d{W}_t,Q_0~{\mu }_0,$$ (5)

where $f:[0,1]\times {\mathbb{R}}^n\to {\mathbb{R}}^n$ and $g:[0,1]\to \mathbb{R}$ are the drift and diffusion coefficients, and $W_t$ is the standard Wiener process. Let $\Psi ,\hat{\Psi }:[0,1]\times {\mathbb{R}}^n\to \mathbb{R}$ be the solutions to the following partial differential equations (PDE),

$$\frac{\partial {\Psi }_t(x)}{\partial t}=-{\left({\nabla }_x{\Psi }_t\left(x\right)\right)}^{\top }\cdot f_t\left(x\right)-\frac{1}{2}{g}_t^2{\Delta }_x{\Psi }_t\left(x\right),$$ (6)

$$\frac{\partial {\hat{\Psi }}_t(x)}{\partial t}=-{\nabla }_x\cdot \left({\hat{\Psi }}_t\left(x\right)f_t\left(x\right)\right)+\frac{1}{2}{g}_t^2{\Delta }_x{\hat{\Psi }}_t\left(x\right),$$ (7)

with boundary conditions ${\Psi }_0{\hat{\Psi }}_0={\mu }_0$ and ${\Psi }_1{\hat{\Psi }}_1={\mu }_1$, the marginal density of the optimal measure can be factorized as $p_t={\Psi }_t{\hat{\Psi }}_t$ through Hopf-Cole transformation.^22–24 The optimal solution to Eq. 4 can be described as the path measure of the random process $X_t$ to the following forward SDE,^25–27

$$d{X}_t=\left(f_t\left(X_t\right)+g_t^2{\nabla }_x\text{log}{\Psi }_t\left(X_t\right)\right)dt+g_{t}d{W}_t,X_0~{\mu }_0.$$ (8)

or equivalently the backward SDE,²⁸

$$d{X}_t=\left(f_t\left(X_t\right)-g_t^2{\nabla }_x\text{log}{\hat{\Psi }}_t\left(X_t\right)\right)dt+g_{t}d{W}_t,X_1~{\mu }_1.$$ (9)

where ${\nabla }_x\text{log}{\Psi }_t\left(X_t\right)$ and ${\nabla }_x\text{log}{\hat{\Psi }}_t\left(X_t\right)$ are the optimal forward and backward drifts for the Schrödinger bridge.

APPENDIX B. CONNECTION TO GENERATIVE DIFFUSION MODELS

Consider ${\mu }_0$ as the data distribution, and ${\mu }_1$ as the prior which is easy to sample from, for example a Gaussian, the optimal solution can be obtained by choosing ${\Psi }_t(x)=1$. In this special case, the PDE in Eq. 6 is automatically satisfied and the PDE in Eq. 7 becomes the Kolmogorov forward equation, also known as Fokker-Planck equation²⁹ by substituting $p_t(x)={\hat{\Psi }}_t(x)$,

$$\frac{\partial p_t(x)}{\partial t}=-{\nabla }_x\cdot \left(f_t(x)p_t(x)\right)+\frac{1}{2}{g}_t^2{\Delta }_x{p}_t(x).$$ (10)

The forward SDE in Eq. 8 becomes the corresponding SDE of Eq. 10,

$$d{X}_t=f_t\left(X_t\right)dt+g_{t}d{W}_t,X_0~{\mu }_0,$$ (11)

which is the same as the reference SDE when the marginal density is designed as a Gaussian at $t=1$. From Eq. 11, we observe that the Schrödinger bridge degenerates to a DDPM,⁷ i.e., variance preserving (VP) SDE, when we choose $f_t(x)=-\frac{1}{2}\beta (t)x$ and $g_t=\sqrt{\beta (t)}$; the Schrödinger bridge degenerates to a SGM,⁸ i.e., variance exploding (VE) SDE, when we choose $f_t(x)=0$ and $g_t=\sqrt{d{\sigma }^2(t)/dt}$; and the Schrödinger bridge degenerates to flow matching⁹ when we choose $g_t=0$, i.e., the SDE becomes an ordinary differential equation (ODE).

APPENDIX C. PROBABILITY FLOW ORDINARY DIFFERENTIAL EQUATION

The Fokker-Planck equation describes the temporal development of the marginal density of a path measure. There is a family of SDEs that satisfy a given Fokker-Planck equation. This can be seen by rewriting Eq. 10 as,

$$\frac{\partial p_t(x)}{\partial t}=-{\nabla }_x\cdot \left[\left(f_t\left(x\right)-\frac{1}{2}\left(g_t^2-{\lambda }_t^2\right){\nabla }_x\text{log}{p}_t(x)\right)p_t(x)\right]+\frac{1}{2}{\lambda }_t^2{\Delta }_x{p}_t(x),$$ (12)

where $\lambda :[0,1]\to \mathbb{R}$ is an arbitrary choice of the diffusion coefficients, and it corresponds to the following forward SDE,

$$d{X}_t=\left(f_t\left(X_t\right)-\frac{1}{2}\left(g_t^2-{\lambda }_t^2\right){\nabla }_x\text{log}{p}_t\left(X_t\right)\right)dt+{\lambda }_{t}d{W}_t,X_0~{\mu }_0,$$ (13)

or equivalently the backward SDE,²⁸

$$d{X}_t=\left(f_t\left(X_t\right)-\frac{1}{2}\left(g_t^2+{\lambda }_t^2\right){\nabla }_x\text{log}{p}_t\left(X_t\right)\right)dt+{\lambda }_{t}d{W}_t,X_1~{\mu }_1,$$ (14)

where different $\lambda$’s produce different stochastic processes, different path measures, but the same marginal density at every time $t$. When ${\lambda }_t=0$, both the forward and the backward SDEs degenerate to the same probability flow ODE,³⁰

$$d{x}_t=\left(f_t\left(x_t\right)-\frac{1}{2}{g}_t^2{\nabla }_{x_t}\text{log}{p}_t\left(x_t\right)\right)dt,$$ (15)

which is a deterministic trajectory. Note, we use lower case letters as these are not stochastic processes. When $f_t\left(x_t\right)$ is an affine drift coefficient, the probability flow ODE in Eq. 15 is the same as the denoising diffusion implicit model (DDIM).¹⁶ The probability flow ODE gives a deterministic and invertible mapping between the two probability spaces ${\mu }_0$ and ${\mu }_1$ by the forward or the backward integration,

$$x_1=x_0+{\int }_0^1{v}_t\left(x_t\right)dt,$$ (16)

where $v:[0,1]\times {\mathbb{R}}^n\to {\mathbb{R}}^n$ is the time-dependent velocity field defined by $v_t\left(x_t\right)=f_t\left(x_t\right)-\frac{1}{2}{g}_t^2{\nabla }_{x_t}\text{log}{p}_t\left(x_t\right)$.

APPENDIX D. TRAINING OBJECTIVE

Both the drift and diffusion coefficients $f$ and $g$ are chosen by design, and we need to determine the score function ${\nabla }_{x_t}\text{log}{p}_t\left(x_t\right)$ in the velocity field $v_t\left(x_t\right)$ in Eq. 1. Theoretically, the score function can be calculated based on,

$${\nabla }_{x_t}\text{log}{p}_t\left(x_t\right)=\frac{{\text{E}}_{x_0}\left[{\nabla }_{x_t}{p}_{t\mid 0}\left(x_t\mid x_0\right)\right]}{{\text{E}}_{x_0}\left[p_{t\mid 0}\left(x_t\mid x_0\right)\right]},$$ (17)

because the conditional probability $p_{t\mid 0}\left(x_t\mid x_0\right)$ can be obtained analytically from the forward SDE in Eq. 8. However, the calculation involves averaging over the whole data sample, which is computationally expensive and generally infeasible. Notice that Eq. 17 can be reformulated as an optimization problem,

$${\nabla }_{x_t}\text{log}{p}_t\left(x_t\right)=\underset{s}{\text{arg}\text{min}}{\text{E}}_{x_0}\left[p_{t\mid 0}\left(x_t\mid x_0\right){\left(s-{\nabla }_{x_t}\text{log}{p}_{t\mid 0}\left(x_t\mid x_0\right)\right)}^2\right],$$ (18)

for a fixed $x_t$ and $t$. Therefore, the score function can be approximated by a neural network $s_{\theta }\left(t,x_t\right)$ with parameters $\theta$ by minimizing the following objective,

$$ℒ={\text{E}}_{x_0~{\mu }_0,x_t~p_{t\mid 0}\left(x_t\mid x_0\right),t~U(0,1)}\left[{\left(s_{\theta }\left(t,x_t\right)-{\nabla }_{x_t}\text{log}{p}_{t\mid 0}\left(x_t\mid x_0\right)\right)}^2\right],$$ (19)

which is obtained by integrating the function inside the arg min in Eq. 18 over $x_t$ and $t$. The velocity field is approximated by $v_t\left(x_t\right)=f_t\left(x_t\right)-\frac{1}{2}{g}_t^2{s}_{\theta }\left(t,x_t\right)$.