SDDM: Score-Decomposed Diffusion Models on Manifolds
for Unpaired Image-to-Image Translation

SBDMs (Song et al., 2021; Ho et al., 2020; Dhariwal & Nichol, 2021; Zhao et al., 2022) first progressively perturb the training data via a forward diffusion process and then learn to reverse this process to form a generative model of the unknown data distribution. Denoting $q(\mathbf{x}_{0})$ the training set with i.i.d. samples on $\mathbb{R}^{d}$ and $q(\mathbf{x}_{t})$ the intermediate distribution at time $t$, the forward diffusion process $\{\mathbf{x}_{t}\}_{t\in[0,T]}$ follows the stochastic differential equation (SDE):

where $\mathbf{f}(\cdot,t):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is the drift coefficient, $\mathrm{d}t$ denotes an infinitesimal positive timestep, $g(t)\in\mathbb{R}$ is the diffusion coefficient, and $\mathbf{w}\sim\mathcal{N}(\mathbf{0},t\mathbf{I}_{d})$ is a standard Wiener process. Denote $q_{t|0}(\mathbf{x}_{t}|\mathbf{x}_{0})$ the transition kernel from time $0$ to $t$, which is decided by $\mathbf{f}(\mathbf{x},t)$ and $g(t)$. In practice, $\mathbf{f}(\mathbf{x},t)$ is usually an affine transformation w.r.t. $\mathbf{x}$ so that the $q_{t|0}(\mathbf{x}_{t}|\mathbf{x}_{0})$ is a linear Gaussian distribution and $\mathbf{x}_{t}$ can be sampled in one step (Zhao et al., 2022). In practice, the following VP-SDE is mostly used:

and DDPM (Ho et al., 2020; Dhariwal & Nichol, 2021) use the following discrete form of the above SDE:

Normally an SDE is not time-reversible because the forward process loses information on the initial data distribution and converges to a terminal state distribution $q_{T}(\mathbf{x}_{T})$. However,  Song et al. (2021) find that the reverse process satisfies the following reverse-time SDE:

where $\mathrm{d}t$ is an infinitesimal negative timestep and $\overline{\mathbf{w}}$ is a reverse-time standard Wiener process.  Song et al. (2021) adopt a score-based model $\mathbf{s}(\mathbf{x},t)$ to approximate $\nabla_{\mathbf{x}}\log q_{t}(\mathbf{x})$, i.e. $\mathbf{s}(\mathbf{x},t)\dot{=}\nabla_{\mathbf{x}}\log q_{t}(\mathbf{x})$, obtaining the following reverse-time SDE:

In VP-SDE, $q_{T}(\mathbf{x}_{T})$ is also a standard Gaussian distribution.

For a controllable generation, it is convenient to add some guidance function $\varepsilon(\mathbf{x},t)$ to the score function and then get a new time-reverse SDE:

Unpaired I2I aims to transfer an image from a source domain $\mathcal{Y}\subset\mathbb{R}^{d}$ to a different target domain $\mathcal{X}\subset\mathbb{R}^{d}$ as the training data. This translation process can be achieved by designing a distribution $p(\mathbf{x}_{0}|\mathbf{y}_{0})$ on the target domain $\mathcal{X}$ conditioned on an image $\mathbf{y}_{0}\in\mathcal{Y}$ to transfer.

In ILVR (Choi et al., 2021), given a reference image $\mathbf{y}_{0}$, they refine $\mathbf{x}_{t}$ after each denoising step with a low-pass filter $\mathbf{\Phi}$ for the faithfulness to the reference image:

In EGSDE (Zhao et al., 2022), they carefully designed two energy-based guidance functions and follow the conditional generation method in Song et al. (2021):

Notably, energy-based methods do not avoid the intermediate distribution being overly or negatively disturbed, and they both do not fully make use of the statistics of the reference image; thus the generation results may be suboptimal.

Figure 1 overviews the main process of the proposed SDDM model. The second equation at the bottom is the equivalent SDE formulation from Eqns. (8) and (9). Starting with this equation, we have the first SDE in Figure 1 to indicate such a generation process. The illustration explains the two-stage optimization at time step $t$.

To explicitly optimize the tangled distributions during image generation, we use moments of the perturbed reference image $\mathbf{y}_{0}$ as constraints for constructing separable manifolds, thus disentangling the distributions of adjacent time steps. As shown in Figure 1, the manifolds of adjacent time steps $t$ and $t-1,$ $\mathcal{M}_{t}$ and $\mathcal{M}_{t-1}$ are separable, which indicates the conditional distributions of adjacent time steps $\mathbf{x}_{t}$ and $\mathbf{x}_{t-1}$ are also separable. Furthermore, at time step $t$, the manifold $\mathcal{M}_{t}$ decompose the score function $\mathbf{s}(\mathbf{x},t)$ into the content refinement part $\mathbf{s}_{r}(\mathbf{x},t)$ and the image denoising part $\mathbf{s}_{d}(\mathbf{x},t)$, and also separate out the content refinement parts $\nabla_{\mathbf{x}}\varepsilon_{1r}(\mathbf{x},\mathbf{y}_{0},t),\nabla_{\mathbf{x}}\varepsilon_{2r}(\mathbf{x},\mathbf{y}_{0},t)$ of $\nabla_{\mathbf{x}}\varepsilon_{1}(\mathbf{x},\mathbf{y}_{0},t),\nabla_{\mathbf{x}}\varepsilon_{2}(\mathbf{x},\mathbf{y}_{0},t)$ on the tangent space $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$. Therefore, The entire optimization process at each time step is divided into two stages: one is to optimize on the manifold $\mathcal{M}_{t}$, and the other stage is to map to the next manifold $\mathcal{M}_{t-1}$ properly.

In the first stage, we optimize on the manifold $\mathcal{M}_{t}$. We apply a multi-objective optimization algorithm to get the red vector MOO, which is the optimal direction considering the score function and energy guidance on the tangent space $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$. Then at the second stage, we use the rest of the first equation in Figure 1, which contains $\left[\mathbf{f}(\mathbf{x},t)-g(t)^{2}\mathbf{s}_{d}(\mathbf{x},t)+\boldsymbol{\delta}\right]\mathrm{d}t+g(t)\mathrm{d}\overline{\mathbf{w}}$ to map the $\mathbf{x}_{t}+\mathrm{MOO}$ to the next manifold $\mathcal{M}_{t-1}$ properly. Note that here we use $\boldsymbol{\delta}$ to indicate the restriction on $\mathcal{M}_{t-1}$ for the consistency of form.

Given a score function $\mathbf{s}(\mathbf{x})=\nabla_{\mathbf{x}}\log p(\mathbf{x})$ on $\mathbb{R}^{d}$, suppose $\mathcal{M}$ is a smooth, compact submanifold of $\mathbb{R}^{d}$. We let $p_{\mathcal{M}}(\mathbf{x})$ is the corresponding probability distribution restricted on $\mathcal{M}$. Then we have the following definitions:

which is the score function on the manifold. If there is a series of manifolds $\{\mathcal{M}_{t}\}$, and the original score function is denoted $\mathbf{s}(\mathbf{x},t)$, we denote $\mathbf{s}_{r}(\mathbf{x},t)$ the tangent score function on $\mathcal{M}_{t}$.

which is the score function on the normal space of the manifold. We also denote $\mathbf{s}_{d}(\mathbf{x},t)$ the normal score function on the manifold $\mathcal{M}_{t}$.

Then we have the following score function decomposition:

which can be derived when knowing $\mathbf{s}_{r}(\mathbf{x})=\mathbf{s}(\mathbf{x})|_{T_{\mathbf{x}}\mathcal{M}}$.

Normally this division is meaningless because the manifolds of adjacent time steps are coupled with each other. Previous researchers usually treat the entire $\cup_{t}\mathbf{x}_{t}$ as an entire manifold (Liu et al., 2022), or use strong assumptions (Chung et al., 2022). However, in some conditional generation tasks, for example, the image-to-image transition task, a given reference image $\mathbf{y}_{0}$ can provide compact manifolds at different time steps, and manifolds of adjacent time steps can be well separated. In this situation, the tangent score function can be treated as a refinement part on the manifold. The normal score function is part of the mapping function between manifolds of adjacent time steps.

We have Proposition 1 to describe the manifolds.

Proposition 1 shows that we can use statistical constraints to define concentrated manifolds with lower dimensions than $\mathbb{R}^{d}$. We can also use the chunking trick to lower the manifold dimensions, which will be introduced in Section 4. Therefore, we can use such manifolds to represent the maintenance of the statistics, which indicates that the tangent space $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$ can separate the “refinement” part well.

We also have Lemma 2 to show that perturbed distributions of adjacent time steps, $\mathbf{y}_{t}$ and $\mathbf{y}_{t-1}$ can be well separated.

Therefore, we can use $\mathcal{M}_{t}$ to decompose $\mathbf{s}$ into $\mathbf{s}_{r}$ and $\mathbf{s}_{d}$ approximately. More generally, we can decouple the optimization space with the tangent space $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$. With $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$, we can operate the score function of SBDM and energy more elaborately. We can also split the “refinement” part out, thus preventing the “denoising” part of the score function from being overly disturbed.

Firstly, we will give some main definitions about manifold optimization and muti-objective optimization in our task. We use restriction $\mathbf{R}_{\mathbf{x}_{t}}$ represent the function that maps the points on $T_{\mathbf{x}_{t}}\mathcal{M}_{t}$ near $\mathbf{x}_{t}$ to the manifold $\mathcal{M}_{t}$, which is normally an orthogonal projection onto $\mathcal{M}_{t}$.

Because that given $t$, the score function $\mathbf{s}(\mathbf{x},t)$ is an estimation of $\nabla_{\mathbf{x}}\log q_{t}(\mathbf{x})$, and we can use $\log q_{t}(\mathbf{x})$ as the potential energy of $\mathbf{s}(\mathbf{x},t)$, so are the guidance of energy funcitons. Then Stage 1 is a manifold optimization.

Then, the goal of multi-objective optimization is to find the Pareto optimal solution. The local Pareto optimality can also be reached via gradient descent like single-objective optimization. We just follow the multiple gradient descent algorithm (MGDA) (Désidéri, 2012). MGDA also leverages the Karush-Kuhn-Tucker (KKT) conditions for the multi-objective optimization, which in our task is that:

All points that satisfy the above conditions are called Pareto stationary points. Every Pareto optimal point is Pareto stationary point, while the reverse is not true. Désidéri (2012) showed that the optimization solution for the problem :

gives a descent direction that improves all tasks or gives a Pareto stationary point. For a balanced result, we normalize all gradients first.

However, In our task, we can search Pareto stationary points in $\mathbf{B}_{\epsilon}(\mathbf{x}_{t})\cap\mathcal{M}_{t}$ for a small $\epsilon$ because we have many time steps of different manifolds. $\mathbf{B}_{\epsilon}(\mathbf{x}_{t})$ is an open ball with center $\mathbf{x}_{t}$, radius $\epsilon$.

We have the following algorithm:

After the optimization on the manifold $\mathcal{M}_{t}$, we get $\mathbf{x}_{t}^{*}$ that dominates $\mathbf{x}_{t}$, then we use $\mathbf{f}(\mathbf{x}_{t}^{*},t)$, the “denoising” part score function $\mathbf{s}_{d}(\mathbf{x}^{*},t)$, reverse-time noise and restriction function on $\mathcal{M}_{t-1}$ to map to the adjacent manifold $\mathcal{M}_{t-1}$.

Firstly, we have the following proposition to describe the properties of the adjacent map.

However, if we use $\mathbf{v_{\mathbf{x}_{t}}}$ as the adjacent map, we will lose the impact of $\mathbf{s}_{d}$ and the reverse-time noise. Therefore, we follow the reverse SDE, using the extra part of which on the normal space $N_{\mathbf{x}_{t}}\mathcal{M}_{t}$ and a restriction function on $\mathcal{M}_{t-1}$ as the adjacent map, we denote this part as $\mathbf{v_{\mathbf{x}_{t}}^{*}}$.

Finally, we have Algorithm 2 in the following to generate images with our proposed SDDM.

We compare our experiments with other GANs-based and SBDM-based methods, as shown in Table 1.

Compared with other SBDM-based methods, our SDDM improves on both metrics FID and SSIM, which indicates the effectiveness of the two-stage generation process of our SDDM via the decomposition of score function and energy guidance with manifolds. Especially compared with EGSDE*, which has strong pre-trained energy functions, in the Cat $\rightarrow$ Dog task, our SDDM improves the FID score by 3.53 and SSIM score by 0.007 with much lower time steps, $1000\rightarrow 100$. For the comparison with EGSDE** having 200 diffusion steps, SDDM improves the FID score by 7.87 and the SSIM score by 0.011 in the Cat $\rightarrow$ Dog task and improves the FID score by 0.75 and the SSIM score by 0.014 in the Male $\rightarrow$ Female task, which suggests the advantage of our SDDM in fewer diffusion steps. The visual comparison is in Appendix F.

Observations on Score Components. While performing the Cat $\rightarrow$ Dog experiment, we report the $L_{2}$ norms of the deterministic guidance values on $T_{\mathbf{x}_{t}}\mathcal{M}$ and $N_{\mathbf{x}_{t}}\mathcal{M}$. As shown in Figure 2, the component on the normal space has one in 128 dimensions but contains the most value of the deterministic guidance of diffusion models, while the component on the tangent space $\mathbf{s}_{r}(\cdot,\cdot)$ has 127 in 128 dimensions but contains a minimal value, which indicates we have relative large optimization space on the manifold which will not excessively interfere with the intermediate distributions.

Comparison of Different Manifolds. We compare SDDM with different manifold methods and report the results in Table 2. Compared with the manifold restricted with a low-pass filter, the manifold restricted with our BAdaIN has better performance both on FID and SSIM, because our manifold separates the content refinement part and image denoising part better.

Comparison of Different Coefficients. We have two coefficients at each iteration step, the coefficient of the step size $\lambda$ of optimal multi-objection direction and the coefficient of the energy guidance $\lambda_{1}$. As in Table 3, the larger $\lambda$ is, the better FID will be because, at each optimization on the manifold, it reaches a position with higher probability $p_{\mathcal{M}}(\mathbf{x})$, but when $\lambda$ is too large, the FID score will be worse again. The $\lambda_{1}$ has a negative connection with the impact of energy guidance, which indicates that smaller $\lambda_{1}$ makes the energy guidance stronger and thus has a better SSIM score.

Comparison w./w.o. Multi-Objective Optimization. We compare SDDM with SDDM without the MOO method and report the FID, SSIM, and probability of negative impact (PNI), which indicates the situation that the total guidance including score and energy decreases the $p(\mathbf{x})$ in Table 4. The proposed SDDM method avoids such situations and reaches better performance.

$\epsilon$ Policy in The Optimization on Manifolds. We mainly compare three different $\epsilon$ policies:

• •

  Policy 1: The $\epsilon$ is very small such that in Algorithm 1, we iterate only once. With Remark 1, this method will introduce no additional calculations of scores.

• •

  Policy 2: $\epsilon_{t}=\|\mathbf{s}_{r}(\mathbf{x},t)\|$, which normally iterates twice in Algorithm 1. In practice, we iterate twice.

• •

  Policy 3: At the time step $T_{0}$ in Algorithm 2, we set $\epsilon_{t}$ larger to iterate another 4 times. and other time steps are as same as Policy 1.

We report the FID and SSIM of different policies in Table 5. Policy 3 has the best performance, which reveals that iteration a little more at $T_{0}$ time step can balance different metrics better without introducing too much cost.

The Choice of Middle-Time $T_{0}$ and Block Number. As shown in Figure 3, when we chunk more blocks or set the $T_{0}$ smaller, the generated image is more faithful to the reference image. But too many blocks will also introduce some bad details, like the mouth in the left bottom image.