Generalized Consistency Trajectory Models for Image Manipulation

## Abstract

Abstract Diffusion models (DMs) excel in unconditional generation, as well as on applications such as image editing and restoration. The success of DMs lies in the iterative nature of diffusion: diffusion breaks down the complex process of mapping noise to data into a sequence of simple denoising tasks. Moreover, we are able to exert fine-grained control over the generation process by injecting guidance terms into each denoising step. However, the iterative process is also computationally intensive, often taking from tens up to thousands of function evaluations. Although consistency trajectory models (CTMs) enable traversal between any time points along the probability flow ODE (PFODE) and score inference with a single function evaluation, CTMs only allow translation from Gaussian noise to data. This work aims to unlock the full potential of CTMs by proposing generalized CTMs (GCTMs), which translate between arbitrary distributions via ODEs. We discuss the design space of GCTMs and demonstrate their efficacy in various image manipulation tasks such as image-to-image translation, restoration, and editing.

## 1 Introduction

Diffusion-based generative models (DMs) learn the scores of noise-perturbed data distributions, which can be used to translate samples between two distributions by numerically integrating an SDE or a probability flow ODE (PFODE) (Ho et al., ; Dhariwal & Nichol, ; Song et al., ). They have achieved remarkable progress over recent years, even surpassing well-known generative models such as Generative Adversarial Networks (GANs) (Goodfellow et al., ) or Variational Autoencoders (VAEs) (Kingma & Welling, ) in terms of sample quality. Moreover, diffusion models have found wide application in areas such as image-to-image translation (Saharia et al., ), image restoration (Chung et al., ; ), image editing (Meng et al., ), etc.

The success of DMs can largely be attributed to the iterative nature of diffusion, arising from its foundation on differential equations – multi-step generation grants high-quality image synthesis by breaking down the complex process of mapping noise to data into a composition of simple denoising steps. We are also able to exert fine-grained control over the generation process by injecting minute guidance terms into each step (Chung et al., ; Ho & Salimans, ). Indeed, guidance is an underlying principle behind numerous diffusion-based image editing and restoration algorithms.

However, its iterative nature is also a curse, as diffusion inference often demands from tens to thousands of number of neural function evaluations (NFEs) per sample, rendering practical usage difficult. Consequently, there is now a large body of works on improving the inference speed of DMs. Among them, distillation refers to methods which train a neural network to translate samples along PFODE trajectories generated by a pre-trained teacher DM in one or two NFEs. Representative distillation methods include progressive distillation (PD) (Salimans & Ho, ), consistency models (CMs) (Song et al., ), and consistency trajectory models (CTMs) (Kim et al., ).

In contrast to PD or CMs which only allow traversal to the terminal point of the PFODE, CTMs enable traversal between any pair of time points along the PFODE as well as score inference, all in a single inference step. Thus, in theory, CTMs are more amenable to guidance, and are applicable to a wider variety of downstream image manipulation tasks. Yet, there is a lack of works exploring the effectiveness of CTMs in such context.

In this work, we take a step towards unlocking the full potential of CTMs. To this end, we first propose generalized CTMs (GCTMs) which generalize the theoretical framework behind CTMs with Flow Matching (Lipman et al., ) to enable translation between two arbitrary distributions. Next, we discuss the design space of GCTMs, and how each design choice influences the downstream task performance. Finally, we demonstrate the power of GCTMs on a variety of image manipulation tasks. Specifically, our contributions can be summarized as follows.

- •
Generalization of theory. We propose GCTMs, which uses conditional flow matching theory to enable one-step translation between two arbitrary distributions (Theorem 1). This stands in contrast to CTMs, which is only able to learn PFODEs from Gaussian to data. In fact, we prove CTM is a special case of GCTM when one side is Gaussian (Theorem 2).
- •
Elucidation of design space. We clarify the design components of GCTMs, and explain how each component affects downstream task performance (Section 4.1). In particular, flexible choice of couplings enable GCTM training in both unsupervised and supervised settings, allowing us to accelerate zero-shot and supervised image manipulation algorithms.
- •
Empirical verification. We demonstrate the potential of GCTMs on unconditional generation, image-to-image translation, image restoration, image editing, and latent manipulation. We show that GCTMs achieve competitive performance even with NFE = 1.

Figure: Figure 1: An illustration of GCTM and its applications – solid arrows can be implemented by a single forward pass of the GCTM network. GCTMs learn to traverse the Flow Matching ODE which is capable of interpolating two arbitrary distributions $q(\bm{x}_{0})$ and $q(\bm{x}_{1})$. GCTMs allow (a) one-step inference of ODE velocity, (b) one-step traversal between arbitrary time intervals of the ODE, (c) improved gradient-guidance by using exact posterior sample instead of posterior mean, and (d) one-step generation of varying outputs $\bm{x}_{0}$ given a fixed input $\bm{x}_{1}$.
Refer to caption: extracted/5915301/figures/main.png

## 2 Related Work

Diffusion model distillation. Despite the success of diffusion models (DMs) in generation tasks, DMs require large number of function evaluations (NFEs). As a way to improve the inference speed, the distillation method is proposed to predict the previously trained teacher DM’s output, *e.g.*, score function. Progressive distillation (PD) (Salimans & Ho, ) progressively reduces the NFEs by training the student model to learn predictions corresponding to two-steps of the teacher model’s deterministic sampling path. Consistency models (CMs) (Song et al., ) perform distillation by reducing the self-consistency function over the generative ODE. The above methodologies only consider the output of the ODE path. In contrast, consistency trajectory model (CTM) (Kim et al., ) simultaneously learns the integral and infinitesimal changes of the PFODE trajectory. Our paper extends CTM to learn the PFODE trajectory between two arbitrary distributions.

Zero-shot image restoration via diffusion.
Image restoration such as super-resolution, deblurring, and inpainting can be formulated as inverse problems, which obtain true signals from given observations. With the advancements in DMs serving as powerful priors, diffusion based inverse solvers have been explored actively. As a pioneering work, DDRM (Kawar et al., ) performs denoising steps on the spectral space of a linear corrupting matrix. DPS (Chung et al., ) and $\mathrm{\Pi}$GDM (Song et al., ) propose posterior sampling by estimating the likelihood distribution through Jensen’s approximation and Gaussian assumption, respectively. While diffusion-based inverse solvers facilitate zero-shot image restoration, they often necessitate prolonged sampling times.

Image translation via diffusion. Conditional GAN-based Pix2Pix (Isola et al., ) specifies the task of translating one image into another image as image-to-image translation. SDEdit (Meng et al., ) avoids mode collapse and learning instabilities with GANs by utilizing DMs to translate edited images along SDEs. Palette (Saharia et al., ) proposed conditional DMs for image-to-image translation tasks. To address the Gaussian prior constraint with DMs, Schrödinger bridge (SB) or direct diffusion bridge (DDB) methods have been proposed to learn SDEs between arbitrary two distributions (Liu et al., ; Kim et al., ; Delbracio & Milanfar, ; Chung et al., ). However, models that follow SDEs often require large NFEs. In contrast, our model learns ODE paths between two arbitrary distributions and demonstrates competitive results with NFE = 1.

## 3 Background

### 3.1 Diffusion Models

Diffusion models learn to reverse the process of corrupting data into Gaussian noise. Formally, the corruption process can be described by a forward SDE

$$ $\displaystyle d\bm{x}_{t}=\sqrt{2t}\,d\bm{w}_{t}$ (1) $$

defined on the time interval $t\in(0,T)$. Given $\bm{x}_{0}$ distributed according to a data distribution $p(\bm{x}_{0})$, (1) sends $\bm{x}_{0}$ to Gaussian noise as $t$ increases from 0 0 to $T$. The reverse of the corruption process can be described by the reverse SDE

$$ $\displaystyle d\bm{x}_{t}=-2t\nabla\log p(\bm{x}_{t})\,dt+\sqrt{2t}\,d\bar{\bm {w}}_{t}$ (2) $$

or its deterministic counterpart, the probability flow ODE (PFODE)

$$ $\displaystyle d\bm{x}_{t}=-t\nabla\log p(\bm{x}_{t})\,dt=t^{-1}(\bm{x}_{t}- \mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}])\,dt$ (3) $$

where $p(\bm{x}_{t})$ is the distribution of $\bm{x}_{t}$ following (1), and $\bar{\bm{w}}_{t}$ is the standard Wiener process in reverse-time. Given a noise sample $\bm{x}_{T}\sim p(\bm{x}_{T})$, $\bm{x}_{t}$ following (2) or (3) is distributed according to $p(\bm{x}_{t})$ as $t$ decreases from $T$ to 0 0. Thus, diffusion models are able to generate data from noise by approximating the scores $\nabla\log p(\bm{x}_{t})$ via score matching, and then numerically integrating (2) or (3).

### 3.2 Consistency Trajectory Models (CTMs)

CTMs learn to translate samples between arbitrary time points of PFODE trajectories, *i.e.*, the goal of CTMs is to learn the integral of the PFODE

$$ $\displaystyle\textstyle G(\bm{x}_{t},t,s)\coloneqq\bm{x}_{t}+\int_{t}^{s}u^{-1 }(\bm{x}_{u}-\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{u})}[\bm{x}_{0}])\,du$ (4) $$

for $s\leq t$, where the terminal distribution $p(\bm{x}_{T})$ is assumed to be Gaussian. The parametrization

$$ $\displaystyle\textstyle G(\bm{x}_{t},t,s)=\frac{s}{t}\bm{x}_{t}+\left(1-\frac{ s}{t}\right)g(\bm{x}_{t},t,s)$ (5) $$

where

$$ $\displaystyle\textstyle g(\bm{x}_{t},t,s)=\bm{x}_{t}+\frac{t}{t-s}\int_{t}^{s} u^{-1}(\bm{x}_{u}-\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{u})}[\bm{x}_{0}])\,du$ (6) $$

enables both traversal along the PFODE as well as score inference, since

$$ $\displaystyle\lim_{s\rightarrow t}g(\bm{x}_{t},t,s)=\mathbb{E}_{p(\bm{x}_{0}| \bm{x}_{t})}[\bm{x}_{0}]$ (7) $$

so we may define $g(\bm{x}_{t},t,t)\coloneqq\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}]$.

Given a pre-trained DM, CTMs approximate $g$ with a neural net $g_{\theta}$ by simultaneously minimizing a distillation loss and a denoising score matching (DSM) loss. The distillation loss is

$$ $\displaystyle\mathcal{L}_{\operatorname{CTM}}(\theta)\coloneqq\mathbb{E}_{0 \leq s\leq u<t\leq T}\mathbb{E}_{p(\bm{x}_{t})}\left[d\left(G_{\theta}(\bm{x}_ {t},t,s),G_{\operatorname{sg}(\theta)}(\bm{x}_{t\rightarrow u},u,s)\right)\right]$ (8) $$

where $G_{\theta}$ is (5) with $g_{\theta}$ in place of $g$, $d(\cdot,\cdot)$ is a measure of similarity between inputs, $\operatorname{sg}$ is the stop-gradient operation, and $\bm{x}_{t\rightarrow u}$ is defined to be the integral of PFODE from time $t$ to $u$ starting from $\bm{x}_{t}$ using score estimates from the pre-trained diffusion model. Minimization of (8) causes $G_{\theta}$ to adhere to PFODE trajectories generated by the pre-trained diffusion model. The DSM loss is

$$ $\displaystyle\mathcal{L}_{\operatorname{DSM}}(\theta)\coloneqq\mathbb{E}_{0 \leq\hat{t}\leq T}\mathbb{E}_{p(\bm{x}_{0})p(\bm{x}_{1})}\mathbb{E}_{p(\bm{x}_ {\hat{t}}|\bm{x}_{0},\bm{x}_{T})}\left[\|\bm{x}_{0}-g_{\theta}(\bm{x}_{\hat{t} },\hat{t},\hat{t})\|_{2}^{2}\right]$ (9) $$

where $p(\bm{x}_{t}|\bm{x}_{0},\bm{x}_{T})=\delta_{\bm{x}_{0}+t\bm{x}_{T}}(\bm{x}_{t})$, and minimization of (9) causes $g_{\theta}$ to satisfy (7). This loss acts as a regularization which improves score accuracy, and is crucial for sampling with large NFEs. Thus, the final objective is

$$ $\displaystyle\mathcal{L}_{\operatorname{CTM}}(\theta)+\lambda_{\operatorname{ DSM}}\mathcal{L}_{\operatorname{DSM}}(\theta),$ (10) $$

and it is possible to further improve sample quality by adding a GAN loss.

### 3.3 Flow Matching (FM)

Flow Matching is another technique for learning PFODEs between two distributions $q(\bm{x}_{0})$ and $q(\bm{x}_{1})$. Specifically, let $q(\bm{x}_{0},\bm{x}_{1})$ be a joint distribution of $\bm{x}_{0}$ and $\bm{x}_{1}$. Define

$$ $\displaystyle q(\bm{x}_{t}|\bm{x}_{0},\bm{x}_{1})\coloneqq\delta_{(1-t)\bm{x}_ {0}+t\bm{x}_{1}}(\bm{x}_{t}),\quad q(\bm{x}_{t})\coloneqq\mathbb{E}_{q(\bm{x}_ {0},\bm{x}_{1})}[q(\bm{x}_{t}|\bm{x}_{0},\bm{x}_{1})]$ (11) $$

where $t\in(0,1)$ and $\delta_{\bm{y}}(\cdot)$ is a Dirac delta at $\bm{y}$. Then the ODE given by

$$ $\displaystyle d\bm{x}_{t}=\mathbb{E}_{q(\bm{x}_{0},\bm{x}_{1}|\bm{x}_{t})}[\bm {x}_{1}-\bm{x}_{0}]\,dt$ (12) $$

generates the probability path $q(\bm{x}_{t})$, *i.e.*, with terminal condition $\bm{x}_{1}\sim q(\bm{x}_{1})$, $\bm{x}_{t}$ following (12) is distributed according to $q(\bm{x}_{t})$. Analogous to denoising score matching, the velocity term in (12) can be approximated by a neural network ${\bm{v}}_{\phi}$ which solves a regression problem

$$ $\displaystyle\min_{\phi}\mathbb{E}_{q(\bm{x}_{0},\bm{x}_{1},\bm{x}_{t})}\left[ \|(\bm{x}_{1}-\bm{x}_{0})-{\bm{v}}_{\phi}(\bm{x}_{t},t)\|_{2}^{2}\right].$ (13) $$

Unlike diffusion whose terminal distribution $p(\bm{x}_{T})$ is Gaussian, $q(\bm{x}_{1})$ can be arbitrary. Also, we remark that the theory presented here is only a particular instance of FM called conditional FM.

## 4 Generalized Consistency Trajectory Models (GCTMs)

We now present GCTMs, which generalize CTMs to enable translation between arbitrary distributions. We begin with a crucial theorem which proves we can parametrize the solution to the FM ODE (12) in a form analogous to CTMs. The proof is deferred to Appendix C.1.

###### Theorem 1 .

The ODE (12) is equivalent to

$$ $\displaystyle d\bm{x}_{t}=t^{-1}(\bm{x}_{t}-\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{t })}[\bm{x}_{0}])\,dt$ (14) $$

defined on $t\in(0,1)$. Hence, we can express the solution to (12) as

$$ $\displaystyle\textstyle G(\bm{x}_{t},t,s)=\frac{s}{t}\bm{x}_{t}+\left(1-\frac{ s}{t}\right)g(\bm{x}_{t},t,s)$ (15) $\displaystyle\textstyle\textit{where}\quad g(\bm{x}_{t},t,s)=\bm{x}_{t}+\frac{ t}{t-s}\int_{t}^{s}u^{-1}(\bm{x}_{u}-\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{u})}[\bm {x}_{0}])\,du.$ (16) $$

There are two differences between (5) and (15). First, the time variables $t$ and $s$ now lie in the unit interval $(0,1)$ instead of $(0,T)$, and second, $p(\bm{x}_{0}|\bm{x}_{u})$ is replaced with $q(\bm{x}_{0}|\bm{x}_{u})$.
The second difference is what enables translation between arbitrary distributions, as $q(\bm{x}_{0}|\bm{x}_{u})$ recovers clean images $\bm{x}_{0}$ given images $\bm{x}_{u}$ perturbed by arbitrary type of vectors (e.g., Gaussian noise, images, etc.), while $p(\bm{x}_{0}|\bm{x}_{u})$ recovers clean images $\bm{x}_{0}$ only for Gaussian-perturbed samples $\bm{x}_{u}$. We call a neural network $g_{\theta}$ which approximates (16) a GCTM, and we can train such a network by optimizing the FM counterparts of $\mathcal{L}_{\operatorname{CTM}}$ and $\mathcal{L}_{\operatorname{DSM}}$:

$$ $\displaystyle\mathcal{L}_{\operatorname{GCTM}}(\theta)\coloneqq\mathbb{E}_{0 \leq s\leq u<t\leq 1}\mathbb{E}_{q(\bm{x}_{t})}\left[d\left(G_{\theta}(\bm{x}_ {t},t,s),G_{\operatorname{sg}(\theta)}(\bm{x}_{t\rightarrow u},u,s)\right)\right]$ (17) $$

where $G_{\theta}$ is (15) with $g$ replaced by $g_{\theta}$, and

$$ $\displaystyle\mathcal{L}_{\operatorname{FM}}(\theta)\coloneqq\mathbb{E}_{0\leq \hat{t}\leq 1}\mathbb{E}_{q(\bm{x}_{0},\bm{x}_{1})}\mathbb{E}_{q(\bm{x}_{\hat{ t}}|\bm{x}_{0},\bm{x}_{1})}\left[\|\bm{x}_{0}-g_{\theta}(\bm{x}_{\hat{t}},\hat {t},\hat{t})\|_{2}^{2}\right].$ (18) $$

The next theorem shows that the PFODE (3) learned by CTMs is a special case of the ODE (14) learned by GCTMs, so GCTMs indeed generalize CTMs. The proof is deferred to Appendix C.2.

###### Theorem 2 .

Consider the choice of $q(\bm{x}_{0},\bm{x}_{1})=p(\bm{x}_{0})\cdot\mathcal{N}(\bm{x}_{1}|\bm{0},\bm{I})$. Let

$$ $\displaystyle t^{\prime}\coloneqq t/(1+t),\qquad\bar{\bm{x}}_{t^{\prime}} \coloneqq\bm{x}_{t}/(1+t)$ (19) $$

where $t\in(0,\infty)$ and $\bm{x}_{t}$ follows the PFODE (3).
Then

$$ $\displaystyle\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}]=\mathbb{E}_{q( \bm{x}_{0}|\bar{\bm{x}}_{t^{\prime}})}[\bm{x}_{0}]$ (20) $$

and $\bar{\bm{x}}_{t^{\prime}}$ follows the ODE

$$ $\displaystyle d\bar{\bm{x}}_{t^{\prime}}={t^{\prime}}^{-1}(\bar{\bm{x}}_{t^{ \prime}}-\mathbb{E}_{q(\bm{x}_{0}|\bar{\bm{x}}_{t^{\prime}})}[\bm{x}_{0}])\,dt ^{\prime}$ (21) $$

on $t^{\prime}\in(0,1)$.
Furthermore, let $G_{\operatorname{CTM}}$, $g_{\operatorname{CTM}}$ denote CTM solutions and let $G_{\operatorname{GCTM}}$, $g_{\operatorname{GCTM}}$ denote GCTM solutions. Then with $s^{\prime}=s/(1+s)$,

$$ $\displaystyle\begin{cases}G_{\operatorname{CTM}}(\bm{x}_{t},t,s)=G_{ \operatorname{GCTM}}(\bar{\bm{x}}_{t^{\prime}},t^{\prime},s^{\prime})\cdot(1+s )\\ g_{\operatorname{CTM}}(\bm{x}_{t},t,t)=g_{\operatorname{GCTM}}(\bar{\bm{x}}_{t ^{\prime}},t^{\prime},t^{\prime})\end{cases}$ (22) $$

In short, (20) shows the equivalence of scores, and (21) shows the equivalence of ODEs. Thus, given $g_{\theta}$ trained with $\mathcal{L}_{\operatorname{FM}}$ and $\mathcal{L}_{\operatorname{GCTM}}$ with the setting of Thm. 1, we are able to evaluate diffusion scores and simulate diffusion PFODE trajectories with a simple change of variables (19), as shown in (22).

Given GCTM’s capability to replicate CTM, we will now outline the key components of GCTM that enable its significant extension for various downstream tasks. This flexibility offers a notable advantage of GCTM over CTM.

###### Theorem 1 .

###### Theorem 2 .

### 4.1 The Design Space of GCTMs

Coupling $q(\bm{x}_{0},\bm{x}_{1})$. In contrast to diffusion which only uses the trivial coupling $q(\bm{x}_{0},\bm{x}_{1})=q(\bm{x}_{0})q(\bm{x}_{1})$ in $\mathcal{L}_{\operatorname{DSM}}(\theta)$, FM allows us to use arbitrary joint distributions of $q(\bm{x}_{0})$ and $q(\bm{x}_{1})$ in $\mathcal{L}_{\operatorname{FM}}(\theta)$. Intuitively, $q(\bm{x}_{0},\bm{x}_{1})$ encodes our inductive bias for what kind of pairs $(\bm{x}_{0},\bm{x}_{1})$ we wish the model to learn, since FM ODE is distributed $q(\bm{x}_{t})$ at each time $t$, and $q(\bm{x}_{t})$ is the distribution of $(1-t)\bm{x}_{0}+t\bm{x}_{1}$ for $(\bm{x}_{0},\bm{x}_{1})\sim q(\bm{x}_{0},\bm{x}_{1})$. Here, we list three valid couplings of GCTM as examples (see Alg. 1 for code). In contrast,
CTM
only use a special case of independence coupling.

- •
Independent coupling:
$\displaystyle q(\bm{x}_{0},\bm{x}_{1})=q(\bm{x}_{0})q(\bm{x}_{1})$
(23)
This coupling reflects no prior assumption about the relation between $\bm{x}_{0}$ and $\bm{x}_{1}$. As shown earlier, diffusion models use this type of coupling.
- •
(Entropy-regularized) Optimal transport coupling:
$\displaystyle q=\operatorname*{arg\,min}_{\tilde{q}}\mathbb{E}_{\tilde{q}(\bm{
x}_{0},\bm{x}_{1})}\left[\|\bm{x}_{0}-\bm{x}_{1}\|_{2}^{2}\right]-\tau H(
\tilde{q})$
(24)
where $\operatorname*{arg\,min}$ is over all joint distributions $\tilde{q}$ of $q(\bm{x}_{0})$ and $q(\bm{x}_{1})$, $H$ denotes entropy, and $\tau$ is the regularization coefficient. This coupling reflects the inductive bias that $\bm{x}_{0}$ and $\bm{x}_{1}$ must be close together under the Euclidean distance. In practice, we use the Sinkhorn-Knopp (SK) algorithm (Cuturi, ) to sample OT pairs. A pseudo-code for SK is given as Alg. 3 in Appendix B.
- •
Supervised coupling:
$\displaystyle q(\bm{x}_{0},\bm{x}_{1})=\int q(\bm{x}_{0})q(\bm{H}|\bm{x}_{0})
\delta_{\bm{H}\bm{x}_{0}}(\bm{x}_{1})\,d\bm{H}$
(25)
where $\bm{H}\sim q(\bm{H}|\bm{x}_{0})$ is a random operator, possibly dependent on $\bm{x}_{0}$, which maps ground-truth data $\bm{x}_{0}$ to observations $\bm{x}_{1}$, *i.e.*, $\bm{x}_{1}=\bm{H}\bm{x}_{0}$. For instance, in the context of learning an inpainting model, $\bm{H}$ is could be a random masking operator. For a fixed $\bm{H}$, $q(\bm{H}|\bm{x}_{0})$ reduces to a Dirac delta. With this coupling, the ODE (14) tends to map observed samples $\bm{x}_{1}$ to ground-truth data $\bm{x}_{0}$ as $t\rightarrow 0$.

Table: Algorithm 1 $q(\bm{x}_{0},\bm{x}_{1})$ Sampling

Gaussian perturbation. The cardinality of the support of $q(\bm{x}_{1})$ must be larger than or equal to the cardinality of the support of $q(\bm{x}_{0})$ for there to be a well-defined ODE from $q(\bm{x}_{1})$ to $q(\bm{x}_{0})$. This is because the ODE trajectory given an initial condition is unique, so a single sample $\bm{x}_{1}\sim q(\bm{x}_{1})$ cannot be transported to multiple points in the support of $q(\bm{x}_{0})$. A simple way to address this problem is to add small Gaussian noise to $q(\bm{x}_{1})$ samples such that $q(\bm{x}_{1})$ is supported everywhere.

We emphasize that Gaussian perturbation allows GCTMs to achieve one-to-many generation when we use the supervised coupling. Concretely, consider the scenario where there are multiple labels $\bm{x}_{0}\sim q(\bm{x}_{0}|\bm{x}_{1})$ which correspond to an observed $\bm{x}_{1}$. Then, the perturbation $\bm{\epsilon}$ added to $\bm{x}_{1}$ acts as a source of randomness, allowing the GCTM network to map $\bm{x}_{1}+\bm{\epsilon}$ to distinct labels $\bm{x}_{0}$ for distinct $\bm{\epsilon}$. This stands in contrast to simply regressing the neural network output of $\bm{x}_{1}$ to corresponding labels $\bm{x}_{0}\sim q(\bm{x}_{0}|\bm{x}_{1})$ with $\ell_{2}$ loss, as this will cause the network to map $\bm{x}_{1}$ to the blurry posterior mean $\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{1})}[\bm{x}_{0}]$ instead of a sharp image $\bm{x}_{0}$. Indeed, in Section 5.2, we observe blurry outputs if we use regression instead of GCTMs.

## 5 Experiments

We now explore the possibilities of GCTMs on unconditional generation, image-to-image translation, image restoration, image editing, and latent manipulation. In particular, GCTM admits NFE = 1 sampling via $\bm{x}_{t}\mapsto G_{\theta}(\bm{x}_{t},t,0)$. Due to the similarities between CTMs and GCTMs as detailed in Thm. 1, GCTMs can be trained using CTM training methods. In fact, we run Alg. 2 with the method in Section 5.2 of (Kim et al., ) to train all GCTMs without pre-trained teacher models. A complete description of training settings are deferred to Appendix A.

Figure: (a) Indep. $N=4$, FID = 24.7
Refer to caption: extracted/5915301/figures/in_4.png

### 5.1 Fast Unconditional Generation

In the scenario where we do not have access to data pairs, we must resort to either the independent coupling or the OT coupling. Here, we show that the optimal transport coupling can significantly accelerate the convergence speed of GCTMs during training, especially when we use a smaller number of timesteps $N$.

**Figure 3: Training acceleration.**
| Method | Teacher | FID $\downarrow$ |
| --- | --- | --- |
| CTM | ✓ | 5.28 |
|  | ✗ | 9.00 |
| CM | ✓ | 3.55 |
|  | ✗ | 8.70 |
| iCM | ✗ | 2.51 |
| GCTM (OT) | ✗ | 5.32 |

Using small $N$ may be of interest when we wish to trade-off training speed for performance, since per-iteration training cost of GCTMs increases linearly with $N$. For instance, when $t=1$ and $u=s=0$ in the GCTM loss (17), we need to integrate along the entire time interval $(0,1)$, which requires $N$ steps of ODE integration.

In Figure 3, we observe up to $\times$2.5 acceleration in terms of training iterations when we use OT coupling instead of independent coupling. Indeed, in Figure 2, OT coupling samples are visually sharper than independent coupling samples. We postulate this is because (1) OT coupling leads to straighter ODE trajectories, so we can accurately integrate ODEs with smaller $N$, and (2) lower variance from OT pairs leads to smaller variance in loss gradients, as discussed in (Pooladian et al., ).

In Table 1, we compare the Fréchet Inception Distance (FID) (Heusel et al., ) of GCTM and relevant baselines on CIFAR10 with NFE = 1. In the setting where we do not use a pre-trained teacher diffusion model, GCTM with OT coupling outperforms all methods with the exception of iCM (Song & Dhariwal, ), which is an improved variant of CM. Moreover, GCTM is on par with CTM trained with a teacher. We speculate that further fine-tuning of hyper-parameters could push the performance of GCTMs to match that of iCMs, and we leave this for future work.

### 5.2 Fast Image-to-Image Translation

**Table 2: Quantitative evaluation of I2I translation. Best is in bold, second best is underlined.**
| Method | NFE | Time<br>(ms) | Edges$\rightarrow$Shoes | Night$\rightarrow$Day | Facades |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| FID $\downarrow$ | IS $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | IS $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | IS $\uparrow$ | LPIPS $\downarrow$ |  |  |  |
| Regression | 1 | 87 | 54.3 | 3.41 | 0.100 | 189.2 | 1.85 | 0.373 | 121.8 | 3.28 | 0.274 |
| Pix2Pix (Isola et al., ) | 1 | 33 | 77.0 | 3.17 | 0.208 | 158.0 | 1.68 | 0.418 | 134.1 | 2.74 | 0.288 |
| Palette (Saharia et al., ) | 5 | 166 | 334.1 | 1.90 | 0.861 | 350.2 | 1.16 | 0.707 | 259.3 | 2.47 | 0.394 |
| I^2SB(Liu et al., ) | 5 | 284 | 53.9 | 3.23 | 0.154 | 145.8 | 1.79 | 0.376 | 135.2 | 2.51 | 0.269 |
| GCTM | 1 | 87 | 40.3 | 3.54 | 0.097 | 148.8 | 2.00 | 0.317 | 111.3 | 2.99 | 0.230 |

Figure: Figure 4: Qualitative evaluation of image-to-image translation on Edges$\rightarrow$Shoes (top), Night$\rightarrow$Day (middle) and Facades (bottom). NFE = 5 for I^2SB and Palette.
Refer to caption: extracted/5915301/figures/fig1_revised.png

Unlike previous distillation methods such as CM or CTM, GCTM can learn ODEs between arbitrary distributions, enabling image-to-image translation. To numerically validate this theoretical improvement, we train GCTMs on three translation tasks Edges$\rightarrow$Shoes, Night$\rightarrow$Day, and Facades (Isola et al., ), scaled to 64$\times$64, with the supervised coupling. We consider three baseline methods: $\ell_{2}$-regression, Pix2Pix (Isola et al., ), Palette (Saharia et al., ) and I^2SB(Liu et al., ). To evaluate translation performance, we use FID and Inception Score (IS) (Barratt & Sharma, ) to rate translation quality and LPIPS (Zhang et al., ) to assess faithfulness to input. We control NFEs such that all methods have similar inference times, and we calculate all metrics on validation samples.

In Table 2, we see GCTM shows strong performance on all tasks. In particular, GCTM is good at preserving input structure, as supported by low LPIPS values.
SDE-based methods I^2SB and Palette show poor performance at low NFEs, even when trained with pairs.
Qualitative results in Figure 4 are in line with the metrics. Baselines produce blurry or nonsensical samples, while GCTM produces sharp and realistic images that are faithful to the input.

**Table 3: Quantitative evaluation of image restoration on FFHQ.**
|  | Method | NFE | Time<br>(ms) | SR2 - Bicubic | Deblur - Gaussian | Inpaint - Center |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |  |  |  |
|  | DPS | 32 | 1079 | 31.19 | 0.935 | 0.015 | 27.88 | 0.878 | 0.041 | 24.69 | 0.876 | 0.042 |
|  | CM | 32 | 1074 | 30.80 | 0.930 | 0.010 | 27.85 | 0.871 | 0.027 | 23.02 | 0.857 | 0.050 |
| 0-Shot | GCTM | 32 | 1382 | 31.61 | 0.939 | 0.015 | 28.19 | 0.885 | 0.037 | 24.47 | 0.876 | 0.042 |
|  | Regression | 1 | 87 | 33.46 | 0.964 | 0.015 | 31.19 | 0.942 | 0.015 | 28.76 | 0.922 | 0.028 |
|  | Palette | 5 | 166 | 17.88 | 0.556 | 0.234 | 17.81 | 0.571 | 0.234 | 16.12 | 0.489 | 0.357 |
|  | I^2SB | 5 | 284 | 26.74 | 0.869 | 0.033 | 26.20 | 0.853 | 0.038 | 26.01 | 0.874 | 0.038 |
| Superv. | GCTM | 1 | 87 | 32.37 | 0.954 | 0.009 | 30.56 | 0.935 | 0.009 | 27.37 | 0.896 | 0.027 |

### 5.3 Fast Image Restoration

Figure: Figure 5: Reg. vs. GCTM.
Refer to caption: extracted/5915301/figures/inv_sample.png

We consider two settings on the FFHQ $64\times 64$ dataset, where we either know or do not know the corruption operator. In the former case, we train an unconditional GCTM with the independent coupling, with which we implement three zero-shot image restoration algorithms: DPS, CM-based image restoration, and the guided generation algorithm illustrated in Figure 1, where the loss is given as inconsistency between observations (see Append. B.2 for pseudo-codes and a detailed discussion of the differences). In the latter case, we train a GCTM with the supervised coupling and $\ell_{2}$-regression, I^2SB  and Palette for comparison. Notably, GCTM is the only model applicable to both situations, thanks to the flexible choice of couplings. We again control NFEs such that all methods have similar inference speed.

Table 3 presents the numerical results in both settings. In the zero-shot setting, we see GCTM outperforming both DPS and CM. In particular, CM is slightly worse than DPS. Sample quality degradation due to error accumulation for CMs at large NFEs have already been observed in unconditional generation (*e.g.*, see Fig. 9 in (Kim et al., )), and we speculate a similar problem occurs for CMs in image restoration as well.
On the other hand, GCTMs avoid this problem, as they are able to traverse to a smaller time using the ODE velocity approximated via $g_{\theta}$.
In the supervised setting, we see regression attains the best PSNR and SSIM. This is a natural consequence of perception-distortion trade-off. Specifically, regression minimizes the MSE loss, so it leads to best distortion metrics (Delbracio & Milanfar, ) while producing blurry results. GCTM, which provides best results if we exclude regression on distortion metrics (PSNR and SSIM) and best results on perception metrics (LPIPS), strikes the best balance between perception and distortion. For instance, in Fig. 5 inpainting results, regression sample lacks detail (e.g., wrinkles) while GCTM sample is sharp. We show more samples in Appendix E.
In particular, in Table 5, we demonstrate image restoration task of GCTM on ImageNet with higher resolution (256 $\times$ 256 resolution) to demonstrate it scalability.

Figure: Figure 6: Image editing with GCTM, NFE = 1.
Refer to caption: extracted/5915301/figures/edit.png

### 5.4 Fast Image Editing

In this section, we demonstrate that GCTM can perform realistic and faithful image editing without any special purpose training. Figure 6 shows image editing with an Edges$\rightarrow$Shoes model and an unconditional FFHQ model. On Edges$\rightarrow$Shoes, to edit an image, a user creates an edited input, which is an edge image painted to have a desired color and / or modified to have a desired outline. We then interpolate the edited input and the original edge image to a certain time point $t=s$ and send it to time $t=0$ with GCTM to produce the output. On FFHQ, analogous to SDEdit (Meng et al., ), we interpolate an edited image with Gaussian noise and send it to time $t=0$ with GCTM to generate the output. In contrast to previous image editing models such as SDEdit, GCTM requires only a single step to edit an image. Moreover, we observe that GCTM faithfully preserves source image structure while making the desired changes to the image.

Figure: (a) Controlling latent vector strength $\gamma$
Refer to caption: extracted/5915301/figures/interp_strength.png

### 5.5 Fast Latent Manipulation

In this section, we demonstrate that GCTMs have a highly controllable latent space. Since there are plenty of works on latent manipulation with unconditional diffusion models, we focus on latent manipulation with GCTMs trained for image-to-image translation. For an image-to-image translation GCTM trained with Gaussian perturbation in Section 4.1, we assert that the perturbation added to $\bm{x}_{1}$ can be manipulated to produce desired outputs $\bm{x}_{0}$. In other words, the perturbation acts as a “latent vector” which controls the factors of variation in $\bm{x}_{0}$. To test this hypothesis, in Figure 7, we display outputs $G_{\theta}(\bm{x}_{1}+\gamma\bm{\epsilon},1,0)$ for particular choices of $\bm{\epsilon}$. In the left panel, we observe generated outputs increasingly adhere to the texture of latent $\bm{\epsilon}$ as we increase guidance strength $\gamma$. Interestingly, GCTM generalizes well to latent vectors unseen during training, such as leopard spots or the color red. In the right panel, we explore the effect of linearly combining red, green, and blue latent vectors. We see that the desired color change is reflected faithfully in the outputs. These observations validate our hypothesis that image-to-image GCTMs have an interpretable latent space.

### 5.6 Ablation Study

Figure: Figure 8: Ablation study of GCTM.
Refer to caption: extracted/5915301/figures/ablation.png

We now perform an ablation study on the design choices of Section 4.1. We have already illustrated the power of using appropriate couplings in previous sections, so we explore the importance of $\sigma_{\max}$. A robust choice for $\sigma_{\max}$ for unconditional generation is well-known to be $\sigma_{\max}=80$ (Karras et al., ; Kim et al., ), and we found using this choice to perform sufficiently well for GCTMs when learning to translate noise to data with independent or OT couplings. So, we restrict our attention to image-to-image translation.

In Figure 8, we display the learning curves on Edges$\rightarrow$Shoes for GCTMs trained without and with Gaussian perturbation, and $\sigma_{\max}\in\{80,500\}$. We observe that GCTM trained without perturbation and $\sigma_{\max}=80$ exhibits unstable dynamics, and is unable to minimize the FID below $30$. On other hand, GCTM trained with perturbation and $\sigma_{\max}=80$ surpasses the model trained without perturbation. This demonstrates Gaussian perturbation is indeed crucial for one-to-many generation, as noted in the last paragraph of Section 4.1. Finally, GCTM with both perturbation and $\sigma_{\max}=500$ minimizes FID the fastest. This shows high-curvature regions for image-to-image ODEs lie near $\bm{x}_{1}$, so we need to use a large $\sigma_{\max}$ which places more discretization points near $t=1$.

## 6 Conclusion

Our work marks a significant advancement in the realm of ODE-based generative models, particularly on the transformative capabilities of Consistency Trajectory Models (CTMs). While the iterative nature of diffusion has proven to be a powerful foundation for high-quality image synthesis and nuanced control, the computational demands associated with numerous neural function evaluations (NFEs) per sample have posed challenges for practical implementation. Our proposal of Generalized CTMs (GCTMs) extends the reach of CTMs by enabling one-step translation between arbitrary distributions, surpassing the limitations of traditional CTMs confined to Gaussian noise to data transformations. Through an insightful exploration of the design space, we elucidate the impact of various components on downstream task performance, providing a comprehensive understanding that contributes to a broadly applicable and stable training scheme. Empirical validation across diverse image manipulation tasks demonstrates the potency of GCTMs, showcasing their ability to accelerate and enhance diffusion-based algorithms. In summary, our work not only contributes to theoretical advancements but also delivers tangible benefits, showcasing GCTMs as a key element in unlocking the full potential of diffusion models for practical, real-world applications in image synthesis, translation, restoration, and editing.

## Appendix A Full Experiment Settings

### A.1 Training

In this section, we introduce training choices which provided reliable performance across all experiments in our paper.

Bootstrapping scores. In all our experiments, we train GCTMs without a pre-trained score model. So, analogous to CTMs, we use velocity estimates given by an exponential moving average $\theta_{\operatorname{EMA}}$ of $\theta$ to solve ODEs. We use exponential moving average decay rate $0.999$.

Time discretization. In practice, we discretize the unit interval into a finite number of timesteps $\{t_{n}\}_{n=0}^{N}$ where

$$ $\displaystyle t_{0}=0<t_{1}<\cdots<t_{N}=1$ (26) $$

and learn ODE trajectories integrated with respect to the discretization schedule. EDM (Karras et al., ), which has shown robust performance on a variety of generation tasks, solves the PFODE on the time interval $(\sigma_{\min},\sigma_{\max})$ for $0<\sigma_{\min}<\sigma_{\max}$ according to the discretization schedule

$$ $\displaystyle\sigma_{n}=(\sigma_{\min}^{1/\rho}+(n/N)(\sigma_{\max}^{1/\rho}- \sigma_{\min}^{1/\rho}))^{\rho}$ (27) $$

for $n=0,\ldots,N$ and $\rho=7$. Thus, using the change of time variable (19) derived in Theorem 1, we convert PFODE EDM schedule to FM ODE discretization

$$ $\displaystyle t_{0}=0,\quad t_{n}=\sigma_{n}/(1+\sigma_{n})\quad\text{for} \quad n=1,\ldots,N-1,\quad t_{N}=1.$ (28) $$

In our experiments, we fix $\sigma_{\min}=0.002$ and control $\sigma_{\max}$. We note that $\sigma_{\max}$ controls the amount of emphasis on time near $t=1$, i.e., larger $\sigma_{\max}$ places more time discretization points near $t=1$.

Number of discretization steps $N$. CTMs use fixed $N=18$. In contrast, analogous to iCMs, we double $N$ every $100k$ iterations, starting from $N=4$.

Time $\hat{t}$ distribution. For unconditional generation, we sample

$$ $\displaystyle\hat{t}=\sigma/(1+\sigma),\qquad\log\sigma\sim\mathcal{N}(-1.2,1. 2^{2})$ (29) $$

in accordance with EDM. For image-to-image translation, we sample

$$ $\displaystyle\hat{t}\sim\text{beta}(3,1).$ (30) $$

Network conditioning. We use the EDM conditioning, following CTMs.

Distance $d$. CTMs use $d$ defined as

$$ $\displaystyle d(\bm{x}_{t},\hat{\bm{x}}_{t})=\operatorname{LPIPS}(G_{\theta_{ \operatorname{EMA}}}(\bm{x}_{t},t,0),G_{\theta_{\operatorname{EMA}}}(\hat{\bm{ x}}_{t},t,0))$ (31) $$

which compares the perceptual distance of samples projected to time $t=0$. In contrast, following iCMs, we use the pseudo-huber loss

$$ $\displaystyle d(\bm{x}_{t},\hat{\bm{x}}_{t})=\sqrt{\|\bm{x}_{t}-\hat{\bm{x}}_{ t}\|_{2}^{2}+c^{2}}-c$ (32) $$

where $c=0.00054\sqrt{d}$, where $d$ is the dimension of $\bm{x}_{t}$.

Batch size. We use batch size $128$ for $32\times 32$ resolution images and batch size $64$ for $64\times 64$ resolution images.

Optimizer. We use the Adam optimizer (Kingma & Ba, ) with learning rate

$$ $\displaystyle\eta=0.0002/(128/\text{{batch_size}})$ (33) $$

and default $(\beta_{1},\beta_{2})=(0.9,0.999)$.

Coefficient for $\mathcal{L}_{\operatorname{FM}}(\theta)$. We use $\lambda_{\operatorname{FM}}=0.1$ for all experiments.

Network. We modify SongUNet provided at [https://github.com/NVlabs/edm](https://github.com/NVlabs/edm) to accept two time conditions $t$ and $s$ by using two time embedding layers.

ODE Solver. We use the second order Heun solver to calculate $\mathcal{L}_{\operatorname{GCTM}}(\theta)$.

Gaussian perturbation. We apply a Gaussian perturbation from a normal distribution multiplied by 0.05 to sample $\bm{x}_{1}$, excluding inpainting task.

### A.2 Evaluation

In this section, we describe the details of the evaluation to ensure reproducibility of our experiments.

Datasets. In unconditional generation task, we compare our GCTM generation performance using CIFAR10 training dataset. In image-to-image translation task, we evaluate the performance of models using test sets of Edges$\rightarrow$Shoes, Night$\rightarrow$Day, Facades from Pix2Pix. In image restoration task, we use FFHQ and apply following corruption operators $\bm{H}$ from I^2SB to obtain measurement: bicubic super-resolution with a factor of 2, Gaussian deblurring with $\sigma=0.8$, and center inpainting with Gaussian. We then assess model performance using test dataset.

Baselines. In image-to-image translation task, we compare three baselines: Pix2Pix from [https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix), Palette model from [https://github.com/Janspiry/Palette-Image-to-Image-Diffusion-Models](https://github.com/Janspiry/Palette-Image-to-Image-Diffusion-Models), and I^2SB from [https://github.com/NVlabs/I2SB](https://github.com/NVlabs/I2SB). We modify the image resolution to 64 $\times$ 64 and keep the hyperparameters as described in their code bases, except that Pix2Pix due to the input size constraints of the discriminator. Same configuration is used in supervised image restoration task.

Metrics details. We calculate FID using [https://github.com/mseitzer/pytorch-fid](https://github.com/mseitzer/pytorch-fid) and IS from [https://github.com/pytorch/vision/blob/main/torchvision/models/inception.py](https://github.com/pytorch/vision/blob/main/torchvision/models/inception.py). We assess LPIPS from [https://github.com/richzhang/PerceptualSimilarity](https://github.com/richzhang/PerceptualSimilarity) with AlexNet version 0.1. In generation task, we employ the entire training dataset to obtain FID scores, and in the other task, we sample 5,000 test datasets. To obtain PSNR and SSIM, we convert the data type of model output to uint8 and normalize it. We use [https://github.com/scikit-image/scikit-image](https://github.com/scikit-image/scikit-image) for PSNR and SSIM.

Sampling time. To compare inference speed, we measure the average time between the model taking in one batch size as input and outputting it.

## Appendix B Algorithms

### B.1 Optimal Transport

Figure: Algorithm 3 Sinkhorn-Knopp (SK)

### B.2 Image Restoration

Figure: Algorithm 4 Zero-shot Image Restoration

In Alg. 4, we describe three zero-shot image restoration algorithms, DPS, CM, and GCTM. DPS uses the posterior mean $\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{t_{i}}^{\prime})}[\bm{x}_{0}]$ to both traverse to a smaller time $t_{i-1}$ and to approximate measurement inconsistency. As the posterior mean generally do not lie in the data domain, using it to calculate measurement inconsistency can be problematic. Indeed, approximation error in DPS is closely related to the discrepancy between the posterior mean and $\bm{x}_{t_{i}\rightarrow 0}^{\prime}$ (see Theorem 1 in (Chung et al., ) for a formal statement). On the other hand, CM uses the ODE terminal point $\bm{x}_{t_{i}\rightarrow 0}^{\prime}$ to traverse to a smaller time $t_{i-1}$ and to approximate measurement inconsistency. While CM can have better guidance gradients as $\bm{x}_{t_{i}\rightarrow 0}^{\prime}$ lie within the data domain, using $\bm{x}_{t_{i}\rightarrow 0}^{\prime}$ to traverse to $t_{i-1}$ can accumulate truncation error and degrade sample quality. For instance, see Figure 9 (a) in (Kim et al., ). GCTM mitigates both problems by enabling parallel evaluation of posterior mean and ODE endpoint, as shown in Line 12-13 of Alg. 4.

### B.3 Image Editing

Figure: Algorithm 5 Image Editing

## Appendix C Proofs

### C.1 Proof of Theorem 1

###### Proof.

We observe that the velocity term in (12) may be expressed as

$$ $\displaystyle\mathbb{E}_{q(\bm{x}_{0},\bm{x}_{1}|\bm{x}_{t})}[\bm{x}_{1}-\bm{x }_{0}]$ $\displaystyle=\mathbb{E}_{q(\bm{x}_{0},\bm{x}_{1}|\bm{x}_{t})}[(\bm{x}_{t}-\bm {x}_{0})/t]$ (34) $\displaystyle=\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{t})}[(\bm{x}_{t}-\bm{x}_{0})/t]$ (35) $\displaystyle=(\bm{x}_{t}-\mathbb{E}_{q(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}])/t$ (36) $$

since $\bm{x}_{1}$ is determined given $\bm{x}_{0}$ and $\bm{x}_{t}$. This shows the equivalence between (12) and (14). Eqs. (15) and (16) are straightforward consequences of the equivalence between ODEs.
∎

###### Proof.

### C.2 Proof of Theorem 2

#### C.2.1 Proof of part (i)

###### Proof.

We first show equivalence of scores. We note that

$$ $\displaystyle\bm{x}_{t}\mapsto\bar{\bm{x}}_{t^{\prime}}$ (37) $$

is a bijective transformation, so by change of variables,

$$ $\displaystyle q(\bar{\bm{x}}_{t^{\prime}}|\bm{x}_{0})=(1+t)\cdot\mathcal{N}( \bm{x}_{t}|\bm{x}_{0},t\bm{I})=(1+t)\cdot p(\bm{x}_{t}|\bm{x}_{0})$ (38) $$

and marginalizing out $\bm{x}_{0}$, we get

$$ $\displaystyle q(\bar{\bm{x}}_{t^{\prime}})=(1+t)\cdot p(\bm{x}_{t}).$ (39) $$

It follows by Bayes’ rule that

$$ $\displaystyle p(\bm{x}_{0}|\bm{x}_{t})$ $\displaystyle=\frac{p(\bm{x}_{t}|\bm{x}_{0})p(\bm{x}_{0})}{p(\bm{x}_{t})}$ (40) $\displaystyle=\frac{(1+t)^{-1}q(\bar{\bm{x}}_{t^{\prime}}|\bm{x}_{0})q(\bm{x}_ {0})}{(1+t)^{-1}q(\bar{\bm{x}}_{t^{\prime}})}$ (41) $\displaystyle=\frac{q(\bar{\bm{x}}_{t^{\prime}}|\bm{x}_{0})q(\bm{x}_{0})}{q( \bar{\bm{x}}_{t^{\prime}})}$ (42) $\displaystyle=q(\bm{x}_{0}|\bar{\bm{x}}_{t^{\prime}})$ (43) $$

and thus

$$ $\displaystyle\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}]=\mathbb{E}_{q( \bm{x}_{0}|\bar{\bm{x}}_{t^{\prime}})}[\bm{x}_{0}].$ (44) $$

for all $t\in(0,\infty)$ and $\bm{x}_{t}$. We now show equivalence of ODEs. Let us first re-state the diffusion PFODE below.

$$ $\displaystyle d\bm{x}_{t}=\frac{\bm{x}_{t}-\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t} )}[\bm{x}_{0}]}{t}\,dt.$ (45) $$

With the change of variable

$$ $\displaystyle\bar{\bm{x}}_{t}=\bm{x}_{t}/(1+t),$ (46) $$

we have

$$ $\displaystyle d\bar{\bm{x}}_{t}$ $\displaystyle=-\frac{\bm{x}_{t}}{(1+t)^{2}}\,dt+\frac{1}{1+t}\,d\bm{x}_{t}$ (47) $\displaystyle=-\frac{\bm{x}_{t}}{(1+t)^{2}}\,dt+\frac{\bm{x}_{t}-\mathbb{E}_{p (\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}]}{t(1+t)}\,dt$ (48) $\displaystyle=-\frac{\bar{\bm{x}}_{t}}{1+t}\,dt+\frac{(1+t)\bar{\bm{x}}_{t}- \mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm{x}_{0}]}{t(1+t)}\,dt$ (49) $\displaystyle=\frac{\bar{\bm{x}}_{t}-\mathbb{E}_{p(\bm{x}_{0}|\bm{x}_{t})}[\bm {x}_{0}]}{t(1+t)}\,dt$ (50) $\displaystyle=\frac{\bar{\bm{x}}_{t}-\mathbb{E}_{q(\bm{x}_{0}|\bar{\bm{x}}_{t^ {\prime}})}[\bm{x}_{0}]}{t(1+t)}\,dt$ (51) $$

where we have used equivalence of scores at the last line. We then make the change of time variable

$$ $\displaystyle t^{\prime}=t/(1+t)\implies dt^{\prime}=\frac{1}{(1+t)^{2}}\,dt$ (52) $$

which gives us

$$ $\displaystyle d\bar{\bm{x}}_{t^{\prime}}$ $\displaystyle=\frac{\bar{\bm{x}}_{t^{\prime}}-\mathbb{E}_{q(\bm{x}_{0}|\bar{ \bm{x}}_{t^{\prime}})}[\bm{x}_{0}]}{t/(1+t)}\,dt^{\prime}$ (53) $\displaystyle=\frac{\bar{\bm{x}}_{t^{\prime}}-\mathbb{E}_{q(\bm{x}_{0}|\bar{ \bm{x}}_{t^{\prime}})}[\bm{x}_{0}]}{t^{\prime}}\,dt^{\prime}.$ (54) $$

This concludes the proof.
∎

###### Proof.

#### C.2.2 Proof of part (ii)

###### Proof.

For the first equality in (22), transform PFODE variables $(\bm{x}_{t},t)$ into FM ODE variables $(\bar{\bm{x}}_{t^{\prime}},t^{\prime})$ with (19), transport $\bar{\bm{x}}_{t^{\prime}}$ to $\bar{\bm{x}}_{s^{\prime}}$ with $G_{\operatorname{GCTM}}$, and then transform FM ODE variables $(\bar{\bm{x}}_{s^{\prime}},s^{\prime})$ into PFODE variables $(\bm{x}_{s},s)$ with the inverse of (19). Second equality in (22) follows directly from (20).
∎

###### Proof.

## Appendix D Limitation, Social Impacts, and Reproducibility

Limitations. GCTMs are yet unable to reach state-of-the-art unconditional generative performance. We speculate further tuning of hyper-parameters in the manner of iCMs could improve the performance, and leave this for future work.

Social impacts. GCTM generalizes CTM to achieve fast translation between any two distributions. Hence, GCTM may be used for beneficial purposes, such as fast medical image restoration. However, GCTM may also be used for malicious purposes, such as generation of malicious images, and this must be regulated.

Reproducibility. We will open-source our code upon acceptance. Moreover, we have submitted experiment codes as a supplementary material.

## Appendix E Additional Experiments

### E.1 Comparing I2I performance with other baseline models

We compare the image-to-image (I2I) performance of our model with two baseline approaches: EGSDE (Zhao et al., ) and BBDM (Li et al., ). Since BBDM, an I2I framework based on the Brownian Bridge process, leverages a latent diffusion model, we train it with a pixel-space diffusion model for a fair comparison. Both BBDM and EGSDE are trained on the Edges$\rightarrow$Shoes dataset. As shown in Table LABEL:tab:edit_base, our GCTM outperforms all baselines across various metrics, even when evaluated with fewer sampling steps.

In addition, we visualize the image editing results in Fig. 9. While EGSDE generates realistic images, it fails to faithfully preserve the given conditions. BBDM, on the other hand, struggles to perform robustly on (unseen) conditional images. In contrast, GCTM produces realistic images while accurately maintaining the original conditions.

**Table 4: Evaluation of I2I translation on Edges$\rightarrow$Shoes with other baslines.**
| Method | NFE | Time (ms) | FID $\downarrow$ | IS $\uparrow$ | LPIPS $\downarrow$ |
| --- | --- | --- | --- | --- | --- |
| BBDM (Li et al., ) | 5 | 75 | 43.7 | 3.43 | 0.099 |
| EGSDE (Zhao et al., ) | 500 | 2590 | 198.1 | 2.87 | 0.476 |
| GCTM | 1 | 87 | 40.3 | 3.54 | 0.097 |

Figure: Figure 9: Comparison on image editing with GCTM and other baselines
Refer to caption: extracted/5915301/figures/edit_base.png

### E.2 Controllable Image Editing

In this section, we demonstrate that effectiveness of image editing can be controlled. In Algorithm 5, we control the time point $t$ to determine how much of the edited image to reflect. In Fig. 10, the results visualize how $t$ effect the output of model output. We observe that the larger $t$, the more realistic the image, and the smaller $t$, the more faithful the edit feature. We set $t=0.95$ and $t=0.4$ at supervised coupling and independent coupling, respectively.

Figure: Figure 10: Controllability of image editing by $t$.
Refer to caption: extracted/5915301/figures/fig4.png

### E.3 High-resolution Image restoration

In Table 5, we demonstrate image restoration task of GCTM on ImageNet with higher resolution.

**Table 5: GCTM evaluation of image restoration on ImageNet with 256 $\times$ 256 resolution.**
| Method | NFE | SR4 - Bicubic | Deblur - Gaussian | Inpaint - Center |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |  |  |
| Corrupt | - | 24.48 | 0.708 | 0.340 | 25.26 | 0.830 | 0.223 | 12.92 | 0.708 | 0.598 |
| GCTM | 1 | 26.70 | 0.771 | 0.223 | 34.65 | 0.948 | 0.032 | 21.56 | 0.808 | 0.229 |

### E.4 Additional Image-to-Image Translation Samples

Figure: Figure 11: Additional results on image-to-image translation task on Edges$\rightarrow$Shoes (top), Night$\rightarrow$Day (middle) and Facades (bottom).
Refer to caption: extracted/5915301/figures/app_i2i.png

### E.5 Additional Image Restoration Samples

Figure: Figure 12: Additional results of supervised image restoration task on FFHQ 64$\times$64.
Refer to caption: extracted/5915301/figures/app_i2i1.png

Figure: Figure 13: Additional results of zero-shot image restoration task on FFHQ 64$\times$64.
Refer to caption: extracted/5915301/figures/zeroshot_gctm.png

Figure: Figure 14: Qualitative comparison of zero-shot algorithms on FFHQ 64$\times$64.
Refer to caption: extracted/5915301/figures/zeroshot_compare.png

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: Can you provide the details of computational costs, and training time required for GCTM, and compare this with  CTM, and the  teacher model?


In CTMs [Kim et al., 2024b], adversarial training can be incorporated to further enhance the quality of samples. Can adversarial training also be incorporated in the proposed GCTMs?
soundness: 4
strengths: This is overall a nice submission with significant novel technical contributions.

The proposed  generalized CTM is formulated in a principled manner and is explained well.

The paper nicely extends consistency trajectory models to allow translation between arbitrary distributions via flow matching.

The design space is systematically examined, studying the effects of different couplings,  Gaussian perturbation and $\sigma_{max}$ for  stable training.

Experiments are performed on a variety of image restoration and editing tasks in both zero-shot and  supervised manner. In image restoration, the proposed method outperforms consistency models in zero-shot setting, and provides results competitive with supervised regression with a higher perceptual quality.
summary: Consistency trajectory models (CTM) are a recent technique for accelerated sampling of diffusion models, which involves training a model to predict any intermediate point of the Probability Flow ODE (PFODE) trajectory allowing traversal between any time points along the PFODE. This work extends CTMs to develop generalized CTMs (GCTMs) using conditional flow matching to enable one-step translation between two arbitrary distributions via ODEs instead of only Gaussian noise-to-data transformation in CTMs. They utilize flow matching to allow more flexible couplings between starting and target distributions, including independent coupling used by diffusion models as a special case. This broadens the applicability of GCTMs to handle arbitrary image-to-image translation tasks in addition to unconditional generation, and reduces the computational costs associated with these tasks by reducing the need for multiple neural function evaluations.
weaknesses: I do not find any major concerns in the paper.

In Fig.6, the noise to image GCTM  editing results with NFE = 1 look a little blurred, with background details washed out.

The authors could consider citing and discussing the following parallel works which also aim to perform image-to-image tasks with few NFEs:
Mei etal. CoDi: Conditional Diffusion Distillation for Higher-Fidelity and Faster Image Generation. In CVPR 2024
Zhao etal. CoSIGN: Few-Step Guidance of ConSIstency Model to Solve General INverse Problems. In ECCV 2024
He etal. Consistency Diffusion Bridge Models. In NeurIPS 2024.
Xiao etal. CCM: Real-Time Controllable Visual Content Creation Using Text-to-Image Consistency Models. In ICML 2024
Starodubcev etal. Invertible Consistency Distillation for Text-Guided Image Editing in Around 7 Steps. Arxiv June 2024.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: Please refer to the weakness section above.
soundness: 3
strengths: * The paper is theoretically sound and inspring in that: 
(1) The reparameterization of the FM ODE to resemble the CTM form allows the proposed GCTM to perform consistent trajectory sampling between arbitrary distributions, which is highly flexible and useful.
(2) GCTM is rigorously shown to generalize CTMs. This makes it possible to train the GCTM using the same training stratagies for training CTMs. 

* The paper includes exploration of the design space for the proposed model, such as OT coupling for accelerated sampling (~2.5x as the baseline), Gaussian perturbation for more diverse generation. 

* GCTM achieves state-of-the-art results across various applications, such as image restoration and translation, without needing a pre-trained teacher model.
summary: This paper proposes Generalized Consistency Trajectory Models (GCTMs), which extend Consistency Trajectory Models (CTMs) to support sampling from arbitrary distributions rather than being restricted to Gaussian distributions. By reparameterizing the flow matching (FM) ODE in a way analogous to CTMs, the authors provide a general framework that includes CTMs as a special case when the sampled distributions is Gaussian. The paper also explores the design space of the proposed method. The proposed method achieves state-of-the-art performance in multiple image generation and manipulation tasks without relying on pre-trained teacher models.
weaknesses: *Gaussian Perturbation in Section 4.1:

The introduction of Gaussian perturbation seems contradictory to the main goal of GCTM - sampling from arbitrary distributions. Specifically, in the case of multiple labels corresponding to a single observation (L302-303), Gaussian perturbation does not apply necessarily, as these label variations are not expected to follow an iid Gaussian distribution.

*Lack of teacher model distillation training scheme:

The reason for the absence of distillation from a teacher network is not fully addressed, although the method is claimed to be a generalization of CTM method. Given that distilling from a teacher could potentially improve training efficiency, an explanation for this choice would strengthen the paper.

*Performance gap with iCM:

GCTM does not outperform improved Consistency Models (iCM), but the reason for this gap is not fully explored. A deeper analysis of this performance difference would be valuable.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: What is the value for the batch size M used for OT-based couplings ?

Table 1: The CTM paper [Kim, Lai et al 2024] reports FID of 1.98 for Diffusion Models – Distillation Sampling with NFE = 1. What is the 5.28 FID reported here?
soundness: 2
strengths: The paper is a mixed approach of CTM and Conditional Flow Matching that is rigorous.
The theoretical presentation is new as far as I know.

While this is surely of interest, in its current state, the quality of the presentation and the quality of the experiments are not high enough for a pulication at ICLR.
summary: This work presents "GENERALIZED CONSISTENCY TRAJECTORY MODELS" a generative model framework that allows consistency trajectory models for arbitrary coupling of image distribution while the original approach of CTM is restricted to the classical (image, pure noise) independent coupling.
weaknesses: The main issue of this paper is that CTM is not used on high-resolution images.
The community does not need methods for 64x64 image processing, this resolution is not high enough to assess image to image translation or image restoration.
In comparison,

Unpaired Image-to-Image Translation via Neural Schrödinger Bridge
Download PDF
Beomsu Kim, Gihyun Kwon, Kwanyoung Kim, Jong Chul Ye
ICLR 2024
uses 256x256 and 512x512 and discuss difficulty for other methods to be of higher resolutions.

As another recent example,
Extremal Domain Translation with Neural Optimal Transport
Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, Evgeny Burnaev
NeurIPS 2023
uses 64x64 and 128x128 examples (already to low in my opinion).

Also comparison with the original pix2pix GAN method is not fair using 64x64 resolution since it is based on a local patch-based discriminator (less resolution means less patches to analyze).


l. 425  "we demonstrate image restoration task of GCTM on ImageNet with higher resolution (256 × 256 resolution) to demonstrate it scalability."
There is no comparison done in the appendix.
Comparison with few NFE and many NFE methods (eg DPS, [Chung & Kim et al ICLR 2023]) is necessary to assess the performance.


All the background section is written without references.
Especially, Section 3.3 FLOW MATCHING (FM) seems far from the original paper by Lipman et al, at least in notation.
Is it just a reinterpretation of CFM using diffusion-like notation?


Another main issue is the use of Optimal Transport coupling (l. 255 (Entropy-regularized) Optimal transport coupling).
This is a chicken and egg problem. If one knew the OT coupling presented by Equation (24) this would be a very good generative model (the best one in a sense) and we would not require a CTM.
What is done in practice is Algorithm 3 Sinkhorn-Knopp (SK) (l. 760) with two batches of size M.
(on a side remark, I don't understand why using entropy regularization here instead of computing an optimal assignment (eg using the emd function from the POT library) unless the batch size M is larger than 1000. This would ensure to use all the data.)
Anyway, there is a big gap between equation (24) which is continuous and discrete algorithm restricted on two discrete batches. Expressing the coupling in continuous space resulting from this discrete procedure is probably intractable (and far from Equation (24)).
This is linked with l. 361: We postulate this is because (1) OT coupling leads to straighter ODE trajectories.


Sampling algorithm and details are missing. This sentence is hard to understand "especially when we use a smaller number of timesteps N (l. 323)". Is gamma sampling from [Kim, Lai et al 2024] used?


Table 1: The CTM paper [Kim, Lai et al 2024] reports FID of 1.98 for Diffusion Models – Distillation Sampling with NFE = 1. What is the 5.28 FID reported here?


The Theorems are more Propositions since their proofs only involves reparameterization of ODE/expectation.


Minor corrections / suggestions:
* l. 170 eq (9): There is a cohabitation of [0,T] and [0,1] convention (change x1 in xT)
* l. 245: joint distributions of q(x0) and q(x1) (speak of q(x0,x1) instead ?)
* l. 253:  add with normal distribution
* l. 269 is hard to inderstands. An ODE as t tends to 0 does not have a rigorous meaning.
* Many references point to ArXiv papers while published at conferences or journals (l. 544, l. 547, l. 553, ...)

Meta-review: This paper addresses the computational inefficiency of the multi-step denoising process in diffusion-based generative models. Its key contribution is the introduction of the Generalized Consistency Trajectory Model (GCTM), designed to flexibly translate between two arbitrary distributions. All reviewers agreed that the paper is theoretically well-grounded and that GCTM represents a highly novel approach. Following the author-review discussion phase, the authors further enhanced the manuscript, particularly by extending the experimental results. The AC noted the negative score provided by Reviewer PKWr, as well as private communications from the authors. Upon careful examination of Reviewer PKWr’s comments, the AC found that the reviewer expressed a positive stance on the theoretical aspects of the paper. Reviewer PKWr also provided valuable feedback, particularly the need for experiments on higher-resolution images. In response, the authors conducted experiments with high-resolution patches in early December. The AC believes this partially satisfies Reviewer PKWr’s recommendation for a major revision and a fairer comparison with prior work. Reviewer PKWr stated, “...a major revision and a fairer way of comparing with previous work. It will then be a strong contribution.”

Given the above considerations, the AC recommends accepting this paper.

Decision: Accept (Poster)