SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Consistency loss. In addition to the losses (43) and (44), we also consider an additional consistency loss. A similar idea was explored in Song (2022). In DSB (De Bortoli et al., 2021; Chen et al., 2022) and DSBM, see Algorithm 1, the processes parameterized by  $v_{\theta}$  (forward) and  $v_{\phi}$  backward are identical only at equilibrium. Thus imposing the forward and the backward processes match at each step of DSB or DSBM would lead to some bias. However, this is not the case in the joint training setting. Indeed, in that case, we have  $\mathbb{M}_f^{2n + 1} = \mathbb{M}_b^{2n + 1}$  if the families are rich enough. Therefore, we get that
$$
\mathrm {d} \mathbf {Y} _ {t} = \left\{- f _ {T - t} \left(\mathbf {Y} _ {t}\right) + v _ {\phi} (T - t, \mathbf {Y} _ {t}) \right\} \mathrm {d} t + \sigma_ {T - t} \mathrm {d} \mathbf {B} _ {t}, \quad \mathbf {Y} _ {0} \sim \pi_ {T}, \tag {45}
$$
is the time reversal of
$$
\mathrm {d} \mathbf {X} _ {t} = \left\{f _ {t} \left(\mathbf {X} _ {t}\right) + v _ {\theta} (t, \mathbf {X} _ {t}) \right\} \mathrm {d} t + \sigma_ {t} \mathrm {d} \mathbf {B} _ {t}, \quad \mathbf {X} _ {0} \sim \pi_ {0}. \tag {46}
$$
Computing the time-reversal of (45), we have
$$
\mathrm {d} \mathbf {X} _ {t} = \left\{f _ {t} \left(\mathbf {X} _ {t}\right) - v _ {\phi} (t, \mathbf {X} _ {t}) + \sigma_ {t} ^ {2} \nabla \log \Pi_ {t} ^ {2 n} \left(\mathbf {X} _ {t}\right) \right\} \mathrm {d} t + \sigma_ {t} \mathrm {d} \mathbf {B} _ {t}, \quad \mathbf {X} _ {0} \sim \pi_ {0}. \tag {47}
$$
Identifying (47) and (46), we get that for any  $t \in [0, T]$  and  $x_t \in \mathbb{R}^d$
$$
v _ {\theta} (t, x _ {t}) = - v _ {\phi} (t, x _ {t}) + \sigma_ {t} ^ {2} \nabla \log \Pi_ {t} ^ {2 n} (x _ {t}). \tag {48}
$$
We highlight that letting  $\sigma_t\to 0$  for any  $t\in [0,T]$ , we get that  $v_{\theta} = -v_{\phi}$ , which confirms that the time-reversal of an ODE is simply given by flipping the sign of the velocity. Therefore, we propose the following loss which links the parameters  $\theta$  and  $\phi$
$$
\mathcal {L} _ {\mathrm {c o n s}} (\theta , \phi) = \int_ {0} ^ {T} \mathbb {E} _ {\Pi_ {t} ^ {2 n}} [ \| v _ {\theta} (t, \mathbf {X} _ {t}) + v _ {\phi} (t, \mathbf {X} _ {t}) - \sigma_ {t} ^ {2} \nabla \log \Pi_ {t} ^ {2 n} (\mathbf {X} _ {t}) \| ^ {2} ] / \sigma_ {t} ^ {2} \mathrm {d} t.
$$
Leveraging tools from implicit score matching (Hyvarinen, 2005) and the divergence theorem, we get that
$$
\mathcal {L} _ {\mathrm {c o n s}} (\theta , \phi) = \int_ {0} ^ {T} \mathbb {E} _ {\Pi_ {t} ^ {2 n}} [ \| v _ {\theta} (t, \mathbf {X} _ {t}) + v _ {\phi} (t, \mathbf {X} _ {t}) \| ^ {2} / \sigma_ {t} ^ {2} + 2 \mathrm {d i v} (v _ {\theta} (t, \mathbf {X} _ {t}) + v _ {\phi} (t, \mathbf {X} _ {t})) ] \mathrm {d} t + C,
$$
where  $C\geq 0$  is a constant which does not depend on  $\theta$  and  $\phi$ . Alternatively, (48) shows that
$$
\nabla \log \Pi_ {t} ^ {2 n} (x _ {t}) = \mathbb {E} _ {\Pi_ {T | t} ^ {2 n}} [ \nabla \log \mathbb {Q} _ {T | t} | \mathbf {X} _ {t} = x _ {t} ] + \mathbb {E} _ {\Pi_ {0 | t} ^ {2 n}} [ \nabla \log \mathbb {Q} _ {t | 0} | \mathbf {X} _ {t} = x _ {t} ].
$$
We thus also have the following denoising score matching consistency loss
$$
\mathcal {L} _ {\mathrm {c o n s}} (\theta , \phi) = \int_ {0} ^ {T} \mathbb {E} _ {\Pi_ {0, t, T} ^ {2 n}} [ \| v _ {\theta} (t, \mathbf {X} _ {t}) + v _ {\phi} (t, \mathbf {X} _ {t}) - \sigma_ {t} ^ {2} (\nabla \log \mathbb {Q} _ {T | t} (\mathbf {X} _ {T} | \mathbf {X} _ {t}) + \nabla \log \mathbb {Q} _ {t | 0} (\mathbf {X} _ {t} | \mathbf {X} _ {0})) \| ^ {2} ] / \sigma_ {t} ^ {2} \mathrm {d} t,
$$
The advantage of this DSM loss is that it does not rely on any divergence computation.
Below, we recall the two losses used to estimate the Markovian projection (43) and (44)
$$
\mathcal {L} (\theta) = \int_ {0} ^ {T} \mathbb {E} _ {\Pi_ {t, T}} [ \| \sigma_ {t} ^ {2} \nabla \log \mathbb {Q} _ {T | t} (\mathbf {X} _ {T} | \mathbf {X} _ {t}) - v _ {\theta} (t, \mathbf {X} _ {t}) \| ^ {2} ] / \sigma_ {t} ^ {2} d t,
$$
$$
\mathcal {L} (\phi) = \int_ {0} ^ {T} \mathbb {E} _ {\Pi_ {t, 0}} [ \| \sigma_ {t} ^ {2} \nabla \log \mathbb {Q} _ {t | 0} (\mathbf {X} _ {t} | \mathbf {X} _ {0}) - v _ {\phi} (t, \mathbf {X} _ {t}) \| ^ {2} ] / \sigma_ {t} ^ {2} \mathrm {d} t.
$$
The complete loss we consider in the joint training of the algorithm is of the form
$$
\mathcal {L} _ {\lambda} (\theta , \phi) = \mathcal {L} (\theta) + \mathcal {L} (\phi) + \lambda \mathcal {L} _ {\text {c o n s}} (\theta , \phi), \tag {49}
$$
where  $\lambda > 0$  is an additional regularization parameter. We now state a version of DSBM which performs joint training in Algorithm 3.
Algorithm 3 Diffusion Schrödinger Bridge Matching (Joint Training)
1: Input: Coupling  $\Pi_{0,T}^{0}$ , tractable bridge  $\mathbb{Q}_{|0,T}$ ,  $N \in \mathbb{N}$
2: Let  $\Pi^0 = \Pi_{0,T}^0\mathbb{Q}_{|0,T}$
3: for  $n \in \{0, \dots, N - 1\}$  do
4: Learn  $v_{\phi^{\star}}, v_{\theta^{\star}}$  using (49) with  $\Pi = \Pi^n$ .
5: Let  $\mathbb{M}_f^{n + 1}$  be given by (9).
6: Let  $\mathbb{M}_b^{n + 1}$  be given by (13).
7: Let  $\mathbb{M}^{n + 1} = \frac{1}{2} (\mathbb{M}_f^{n + 1} + \mathbb{M}_b^{n + 1})$ .
8: Let  $\Pi^{n + 1} = \mathbb{M}_{0,T}^{n + 1}\mathbb{Q}_{|0,T}$ .
9: end for
10: Output:  $v_{\theta \star}, v_{\phi \star}$
H Loss Scaling
Similar to the loss weighting in standard diffusion models (Song et al., 2021b; Ho et al., 2020), we derive a similar weighting to reduce the variance of our objective. We focus on the forward direction of Markovian projection in this case, and the backward case can be derived similarly. Our forward loss in the DSBM framework is given by (10), where the inner expectation is given by
$$
\mathbb {E} _ {\Pi_ {t, T}} [ \| \sigma_ {t} ^ {2} \nabla \log \mathbb {Q} _ {T | t} (\mathbf {X} _ {T} | \mathbf {X} _ {t}) - v _ {\theta} (t, \mathbf {X} _ {t}) \| ^ {2} ].
$$
Letting  $\mathbb{Q}_{|0,T}$  be a Brownian bridge with diffusion parameter  $\sigma$  and assuming  $T = 1$ , this becomes
$$
\mathbb {E} _ {\left(\mathbf {X} _ {0}, \mathbf {X} _ {1}\right) \sim \Pi_ {0, 1}, \mathbf {Z} \sim \mathcal {N} (0, \mathrm {I d})} \left[ \left\| \mathbf {X} _ {1} - \mathbf {X} _ {0} - \sigma \sqrt {t / (1 - t)} \mathbf {Z} - v _ {\theta} \left(t, \mathbf {X} _ {t} ^ {0, T}\right) \right\| ^ {2} \right]
$$
with  $\mathbf{X}_t^{0,T} = t\mathbf{X}_1 + (1 - t)\mathbf{X}_0 + \sigma \sqrt{t(1 - t)}\mathbf{Z}$ . When  $t \approx 1$ , we see that the regression target is dominated by the noise term  $\sigma \sqrt{t / (1 - t)}\mathbf{Z}$  which needs to be predicted based on information contained within  $\mathbf{X}_t^{0,T}$ . The loss will have an approximate scale of  $\sigma^2 t / (1 - t)$  when  $t \approx 1$  which will be very large. To avoid these large values affecting gradient descent, we can downweight the loss by  $1 + \sigma^2 t / (1 - t)$  (we can add 1 to effectively cause no loss scaling when  $t$  is close to 0)
$$
\left(1 + \sigma^ {2} t / (1 - t)\right) ^ {- 1} \mathbb {E} _ {\left(\mathbf {X} _ {0}, \mathbf {X} _ {1}\right) \sim \Pi_ {0, 1}, \mathbf {Z} \sim \mathcal {N} (0, \mathrm {I d})} \left[ \| \mathbf {X} _ {1} - \mathbf {X} _ {0} - \sigma \sqrt {t / (1 - t)} \mathbf {Z} - v _ {\theta} (t, \mathbf {X} _ {t} ^ {0, T}) \| ^ {2} \right].
$$
Similar arguments can be applied to the backward loss (14)
$$
\begin{array}{l} \mathbb {E} _ {\Pi_ {t, 0}} [ \| \sigma_ {t} ^ {2} \nabla \log \mathbb {Q} _ {t | 0} (\mathbf {X} _ {t} | \mathbf {X} _ {0}) - v _ {\phi} (t, \mathbf {X} _ {t}) \| ^ {2} ] \\ = \mathbb {E} _ {(\mathbf {X} _ {0}, \mathbf {X} _ {1}) \sim \Pi_ {0, 1}, \mathbf {Z} \sim \mathcal {N} (0, \mathrm {I d})} [ \| \mathbf {X} _ {0} - \mathbf {X} _ {1} - \sigma \sqrt {(1 - t) / t} \mathbf {Z} - v _ {\phi} (t, \mathbf {X} _ {t} ^ {0, T}) \| ^ {2} ], \\ \end{array}
$$
which we then downweight by  $1 + \sigma^2 (1 - t) / t$
$$
(1 + \sigma^ {2} (1 - t) / t) ^ {- 1} \mathbb {E} _ {(\mathbf {X} _ {0}, \mathbf {X} _ {1}) \sim \Pi_ {0, 1}, \mathbf {Z} \sim \mathcal {N} (0, \mathrm {I d})} [ \| \mathbf {X} _ {0} - \mathbf {X} _ {1} - \sigma \sqrt {(1 - t) / t} \mathbf {Z} - v _ {\phi} (t, \mathbf {X} _ {t} ^ {0, T}) \| ^ {2} ].
$$
I Experiments
In this section, we present further details of the experiment setups as well as further experiment results. In all experiments, we use Brownian motion for the reference measure  $\mathbb{Q}$  with corresponding Brownian bridge (4) and  $T = 1$ . We use the Adam optimizer with learning rate  $10^{-4}$  and SiLU activations unless specified otherwise. The experiments are run on computing clusters with a mixture of both CPU and GPU resources.
I.1 2D Experiments
For the 2D experiments, we closely follow Tong et al. (2023) and the released code<sup>7</sup>, and use the same synthetic datasets and the 2-Wasserstein distance between the test set and samples simulated using probability flow ODE as the evaluation metric. However, we use 10000 samples in the test set since we find the 2-Wasserstein distance can vary greatly with only 1000 samples (which can be as high as 0.3 even between two sets of samples both drawn from the ground truth distribution). We use a simple MLP with 3 hidden layers and 256 hidden units to parameterize the forward and backward drift networks. We use batch size 128 and 20 diffusion steps with uniform schedule at sampling time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. As the initial coupling is more optimal in DSBM-IMF+, we reduce the number of outer iterations to 4 and train for 50000 steps in each outer iteration. For flow methods, we train for 20000 steps in total. For Table 2 we use  $\sigma_t = 1$  in all cases, except the moons-8gaussians dataset where we use  $\sigma_t = 5$ . Note that FM cannot be used for the moons-8gaussians task since it requires a Gaussian source, but CFM is applicable. The experiments are run using 1 CPU and take approximately 200 minutes (for both training and testing).
I.1.1 Variance of the reference measure  $\mathbb{Q}$
We comment further on the effect of  $\sigma_{t}$  in the reference path measure  $\mathbb{Q}$ . We assume a time-homogeneous  $\sigma_{t} = \sigma$  for simplicity. In Figure 2, we vary  $\sigma$  and visualize the learned transport for a 3gaussians problem of transporting between two Gaussian mixtures. In Table 4 we show the 2-Wasserstein distance between the test set and generated samples for this 3gaussians problem as well as the moons-8gaussians problem. We find that large values of  $\sigma$  result in increasingly curved transport paths, and correspondingly reduced performance when  $\sigma$  is excessively large. Conversely, we also find reduced performance when  $\sigma$  is excessively small. We conjecture this is due to increased optimization difficulty and bias accumulation. Firstly, the EOT problem becomes more difficult to solve as  $\sigma$  is taken to 0, which would require a higher number of outer iterations. Further, the introduction of noise also decreases optimization difficulty by smoothing the intermediate marginals between the two terminal distributions. The benefit of setting  $\sigma > 0$  and using a stochastic sampler was also observed in Albergo et al. (2023); Delbracio and Milanfar (2023). Finally, we conjecture setting  $\sigma > 0$  could also increase the diversity of sampled couplings and may alleviate some bias accumulation issues in the outer iterations. When  $\sigma = 0$ , these issues result in the artifacts observed in the transferred samples (marked using yellow points) in Figure 2. The appropriate value for  $\sigma$  depends on the spatial scaling of the problem as shown in Table 4, where the optimum  $\sigma$  is larger for the larger scale moons-8gaussians problem.
3gaussians	moons-8gaussians
σ	2-Wasserstein	σ	2-Wasserstein
σ = 0	0.646±0.028	σ = 0	1.459±0.008
σ = 0.1	0.724±0.039	σ = 1.0	1.285±0.346
σ = 0.3	0.546±0.169	σ = 2.0	0.916±0.292
σ = 1.0	0.439±0.072	σ = 4.0	0.818±0.249
σ = 3.0	0.543±0.078	σ = 8.0	0.989±0.179
Table 4: 2-Wasserstein distance for varying value of  $\sigma$  used in the DSBM-IMF method.
<sup>7</sup>https://github.com/atong01/conditional-flow-matching (code released under MIT license)
I.2 Gaussian Experiment
Similar to the 2D experiments, we use a simple MLP with 2 hidden layers and 256 hidden units to parameterize the forward and backward drift networks. This is a smaller network compared to the "large" network in De Bortoli et al. (2021). We use batch size 128 and 20 diffusion sampling steps with uniform time schedule at inference time. Each outer iteration is trained for 10000 steps and we train for 20 outer iterations. The experiments are also run using 1 CPU, and for the case  $d = 50$  finish in approximately 200 minutes. We note that DSB, IMF-b, and DSBM methods all solve for the Schrödinger Bridge transport, whereas Rectified Flow does not and so is not plotted on the covariance convergence plot in Figure 3 since it is not comparable.
For Table 3, we assume the marginals of the learned process  $\mathbb{P}_t$  are also independently Gaussian distributed in each dimension. We thus estimate the KL divergence from the sample mean and variance of each dimension of  $\mathbb{P}_t$  using the analytic KL formula between Gaussian distributions.
I.3 MNIST Transfer Experiment
We follow De Bortoli et al. (2021) closely for the setup of this experiment. We use the set of first 5 letters of EMNIST (A-E, a-e) such that both domains have 10 classes in total. We use the same U-Net architecture, batch size 128 and 30 diffusion sampling steps. The network size is approximately 6.6 million parameters. Each outer iteration is trained for 5000 steps. We refresh the cache datalloader every 1000 steps with 10000 new samples. Contrary to De Bortoli et al. (2021), in our experiments we find that we obtain better sampling quality for both DSB and DSBM using a uniform noisng schedule. We simply choose  $\sigma_t = 1$  for all  $t$  and  $T = 1$  in our experiments.
For DSBM-IPF, we train for at most 50 outer iterations (i.e. 250000 total number of steps). For DSBM-IMF and Rectified Flow, since their first iteration corresponds to Bridge Matching and Flow Matching respectively, we first pretrain the forward and backward networks for 100000 steps using the Bridge Matching or Flow Matching losses. DSBM-IMF then switches to iterative training with 5000 steps per outer iteration, the same as DSBM-IPF. For Rectified Flow, we train for 50000 steps per outer iteration. The experiments are performed using 2 GPUs and take approximately one day.
We provide further experiment results in Figures 13, 14 and 15. In Figure 13, we show samples generated using different algorithms and at different points of convergence. Samples generated using CFM and Rectified Flow with 2 rectified steps in (a) and (b) appear to be less clear and identifiable. OT-CFM improves upon CFM slightly in (c), but many samples still appear to be unclear. For DSB, the algorithm has not converged after 10 iterations, and many samples in (d) still appear to be letters. After 20 iterations, there are still letter-like samples in Figure 4b, and the digit classes also appear to be unbalanced with many instances of digit '0'. After 30 iterations, however, the sample quality of DSB becomes very poor in (e). On the other hand, as shown in (f)-(j), we observe that DSBM-IPF converges faster than DSB, with more accurate samples even in iterations 10 and 20, and the sample quality continues to improve until the end of training after 50 iterations.
We present some additional trajectory samples at the end of training in Figures 14 and 15 in both the forward and backward directions. We observe DSBM is able to transfer samples across the two domains faithfully, and the output samples preserve interesting similarities compared to the input, whereas Bridge Matching preserves much less similarity. The Mean Squared Distance (MSD) between the initial and final samples also confirm that DSBM methods transfer more closely to the original inputs, and DSBM-IMF further achieves the best FID score out of the three methods.
I.4 CelebA Transfer Experiment
For CelebA (Liu et al., 2015), we test DSBM on resized  $64 \times 64$  and  $128 \times 128$  resolution images. We evaluate the dependency of the results with respect to  $\sigma > 0$  on CelebA  $64 \times 64$ . We also showcase the scalability of our method by training DSBM on CelebA  $128 \times 128$ . In both cases, the dataset is split between male/old and female/young. We gather 20000 samples of each class on which we perform classical data augmentation such as horizontal flipping.
For our ablation study w.r.t. the value of  $\sigma > 0$ , we run DSBM for 20 iterations (note that the loss and the quality was still improving after 20 DSBM iterations but by stopping the runs early we were allowed to draw comparisons with more values of  $\sigma$ ). We use a U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The batch size is fixed to 64 and the EMA rate is
(a) CFM
(b) Rectified Flow @ 3
(c) OT-CFM
(d) DSB @ 10
(e) DSB @ 30
(f)DSBM-IPF @ 10
(g)DSBM-IPF @ 20
(h) DSBM-IPF @ 30
(i) DSBM-IPF @ 40
Figure 13: Samples of MNIST digits transferred from the EMNIST letters using different methods. @ indicates the progressed number of outer iterations  $n + 1$ .
(j)DSBM-IPF @ 50
fixed to 0.999. We refresh the cache dataloader every 500 steps with 10000 new samples. The SDE sampler is chosen to be the modified Euler-Maruyama sampler, see Heng et al. (2021) for instance, with a constant schedule for the step sizes. We use 100 sampling steps at inference time and to refresh the cache. For each outer DSBM iteration we train the model for 20000 iterations.
We provide additional transfer results in resolution  $128 \times 128$  in Figure 16. We do not change the training setting for this experiment.
I.5 AFHQ Transfer Experiment
For AFHQ (Choi et al., 2020), we test DSBM between classes cat and wild with  $512 \times 512$  resolution images. Each class contains approximately 5000 samples. We first pretrain the networks using Bridge Matching for 100000 steps, then run DSBM for 20 iterations with 25000 steps per outer iteration. We follow Liu et al. (2023b) and use the same U-Net architecture<sup>8</sup>. The batch size is 4 and the EMA rate is 0.999. We choose  $\sigma^2 = 5$  and again we use 100 sampling steps with constant step sizes.
I.6 CIFAR-10 Generative Modeling Experiment
We also test our method in the standard generative modeling framework on the CIFAR-10 dataset. We again use a U-Net architecture with 4 resolution levels and 2 residual blocks per resolution level. The network size is approximately 39.6 million parameters. The batch size is fixed to 128 and the EMA rate is fixed to 0.9999. The AdamW optimizer is used for this task. For DSBM-IMF and Rectified Flow, again we first pretrain the networks using the Bridge Matching or Flow Matching losses, then switch to DSBM-IMF or RF training with 100000 steps per outer iteration. We find that 1 or 2 additional outer iterations appear sufficiently effective on this task, and additional outer iterations can cause sample quality to drop. For our main experiment, the pretraining stage ran for approximately 6 days, and DSBM-IMF ran for approximately 4 additional days using 4 V100 GPUs.
The best results during training for each method are reported in Table 5 and Figure 17, where we compute the FID score between 50000 samples in the CIFAR-10 training set and 50000 generated samples following standard practice for this task. We observe that DSBM-IMF can clearly improve upon Bridge Matching at the same value of  $\sigma$ , which suggests that further outer iterations in DSBM is beneficial for improving sample quality. This is contrary to Rectified Flow, which causes the FID score to worsen compared to Flow Matching after only 1 rectified iteration. However, as  $\sigma$  increases,
<sup>8</sup>https://github.com/gnobitab/RectifiedFlow/blob/main/ImageGeneration/configs/rectified_flow/afhq_cat_pytorch_rf_gaussian.py (code released under Apache-2.0 license)
Figure 14: Left: EMNIST to MNIST sample trajectory with 30 diffusion steps at  $t = 0, 1/3, 2/3, 1$ . Right: FID score of final samples, and Mean Squared Distance between initial and final samples.
(c) DSBM-IMF
Figure 15: Left: MNIST to EMNIST sample trajectory with 30 diffusion steps at  $t = 0, 1/3, 2/3, 1$ . Right: FID score of final samples, and Mean Squared Distance between initial and final samples.
Figure 16: Transfer results between images given by the tokens female/young and male/old. Top row: original images (left) and generated images (right). Bottom row: original images (left) and generated images (right).
we observe the FID score worsens for both Bridge Matching and DSBM as more stochasticity is introduced in the sampler. The best result of DSBM-IMF is obtained using  $\sigma^2 = 0.2$  and is slightly better than FM (i.e. with  $\sigma^2 = 0$ ) using 100 Euler steps. On the other hand, using the dopri5 ODE solver, FM achieves a FID of 4.055 with on average 148 integration steps. In Figure 17, we observe that both RF and DSBM-IMF are very effective in improving sampling quality at low number of diffusion steps, i.e. low number of function evaluations (NFEs), compared to Bridge and Flow Matching as well as OT-CFM which improves upon CFM slightly. DSBM-IMF also achieves lower FID score than RF as the NFEs are taken higher. Additional strategies such as distillation and fast SDE solvers can also be useful for improving few-step sampling quality further.
I.7 Fluid Flows Experiment
We use the fluid flows dataset from Bischoff and Deck (2023). The dataset consists of unpaired low  $(64\times 64)$  and high  $(512\times 512)$  resolution fields, as well as a context field with local information for the high resolution field. The data fields consist of two channels representing supersaturation and vorticity. The context field is dependent on the wavenumber  $k_{x} = k_{y}\in \{1,2,4,8,16\}$ , which
<sup>9</sup>https://github.com/CliMA/CliIMADatasets.j1 (code released under MIT license)
σ2	FM	RF
0	4.931	6.010
σ2	BM	DSBM-IMF
0.2	5.427	4.511
0.5	8.274	6.896
1	12.749	9.881
Table 5: FID results on the CIFAR-10 train set using 100 Euler(-Maruyama) steps.
Figure 17: FID vs number of diffusion steps (NFE) with NFE between 1 and 1000.
Figure 18: Zoomed-in version of Figure 17 with NFE between 100 and 1000.
(a) Low-res
(b) Diffusion-fb
(c) Bridge Matching
Matching (d) DSBM-IPF
Figure 19: (a) Source low resolution sample; (b)(c)(d)(e) intermediate state and final reconstruction of each algorithm, for wavenumber  $k_{x} = k_{y} = 2$ .
(d) DSBM-IPF
SBM-IPF
(e) DSBM-IMF
specifies the frequency of the saturation specific humidity modulation. We follow Bischoff and Deck (2023) for data processing and network architecture, which is given by a U-Net with an additional spatial mean bypass network given by an MLP. The network size is approximately 11.3 million parameters in total. We train using batch size 4, learning rate  $2 \times 10^{-4}$ , for 5000 steps per outer iteration and  $N = 12$  outer iterations. We refresh the cache dataloader every 2500 steps with 1250 new samples. The training takes approximately 20 hours on a single RTX GPU. We used  $\sigma^2 = 0.3$  and sample using 30 diffusion steps without finetuning these parameters. For the Diffusion-fb method in Bischoff and Deck (2023), we use the released code $^{10}$  without modifying any parameters.