# Learning Fast Samplers for Diffusion Models by Differentiating Through
Sample Quality

Daniel Watson , William Chan, Jonathan Ho & Mohammad Norouzi
Google Research, Brain Team
<EMAIL_ADDRESS>
Work done as part of the Google AI Residency.

###### Abstract

Diffusion models have emerged as an expressive family of generative models
rivaling GANs in sample quality and autoregressive models in likelihood
scores. Standard diffusion models typically require hundreds of forward passes
through the model to generate a single high-fidelity sample. We introduce
Differentiable Diffusion Sampler Search (DDSS): a method that optimizes fast
samplers for any pre-trained diffusion model by differentiating through sample
quality scores. We present Generalized Gaussian Diffusion Models (GGDM), a
family of flexible non-Markovian samplers for diffusion models. We show that
optimizing the degrees of freedom of GGDM samplers by maximizing sample
quality scores via gradient descent leads to improved sample quality. Our
optimization procedure backpropagates through the sampling process using the
reparametrization trick and gradient rematerialization. DDSS achieves strong
results on unconditional image generation across various datasets (e.g., FID
scores on LSUN church 128x128 of 11.6 with only 10 inference steps, and 4.82
with 20 steps, compared to 51.1 and 14.9 with strongest DDPM/DDIM baselines).
Our method is compatible with any pre-trained diffusion model without fine-
tuning or re-training required.

##  1 Introduction

|  |  | 5 steps | 10 steps | 20 steps
---|---|---|---|---|---
1000 steps |  |  DDSS (ours) |  |  |
Real |  |  DDIM |  |  |
Figure 1: Non-cherrypicked samples from DDSS (ours) and strongest
DDIM($\eta=0$) baseline for unconditional DDPMs trained on LSUN churches
128$\\!\times\\!$128\. All samples are generated with the same random seed.
Original DDPM samples (1000 steps) and training images are shown on the left.

Denoising Diffusion Probabilistic Models (DDPM) (Sohl-Dickstein et al., 2015;
Song & Ermon, 2019; Ho et al., 2020) have emerged as a powerful family of
generative models, capable of synthesizing high-quality images, audio, and 3D
shapes (Ho et al., 2020; 2021; Chen et al., 2021a; b; Cai et al., 2020; Luo &
Hu, 2021). Recent work (Dhariwal & Nichol, 2021; Ho et al., 2021) shows that
DDPMs can outperform Generative Adversarial Networks (GAN) (Goodfellow et al.,
2014; Brock et al., 2018) in generation quality, but unlike GANs, DDPMs admit
likelihood computation and much more stable training dynamics (Arjovsky et
al., 2017; Gulrajani et al., 2017).

However, GANs are typically much more efficient than DDPMs at generation time,
often requiring a single forward pass through the generator network, whereas
DDPMs require hundreds of forward passes through a U-Net model. Instead of
learning a generator directly, DDPMs learn to convert noisy data to less noisy
data starting from pure noise, which leads to a wide variety of feasible
strategies for sampling (Song et al., 2021b). In particular, at inference
time, DDPMs allow controlling the number of forward passes (a.k.a. inference
steps) through the denoising network (Song et al., 2020; Nichol & Dhariwal,
2021).

It has been shown both empirically and mathematically that, for any
sufficiently good DDPM, more inference steps leads to better log-likelihood
and sample quality (Nichol & Dhariwal, 2021; Kingma et al., 2021). In
practice, the minimum number of inference steps to achieve competitive sample
quality is highly problem-dependent, e.g., depends on the complexity of the
dataset, and the strength of the conditioning signal if the task is
conditional. Given the importance of generation speed, recent work (Song et
al., 2020; Chen et al., 2021a; Watson et al., 2021) has explored reducing the
number of steps required for high quality sampling with pretrained diffusion
models. See Section 7 for a more complete review of prior work on few-step
sampling.

This paper treats the design of fast samplers for diffusion models as a
differentiable optimization problem, and proposes Differentiable Diffusion
Sampler Search (DDSS). Our key observation is that one can unroll the sampling
chain of a diffusion model and use reparametrization trick (Kingma & Welling,
2013) and gradient rematerialization (Kumar et al., 2019a) to optimize over a
class of parametric few-step samplers with respect to a global objective
function. Our class of parameteric samplers, which we call Generalized
Gaussian Diffusion Model (GGDM), includes Denoising Diffusion Implicit Models
(DDIM) (Song et al., 2020) as a special case and is motivated by the success
of DDIM on fast sampling of diffusion models.

An important challenge for fast DDPM sampling is the mismatch between the
training objective (e.g., ELBO or weighted ELBO) and sample quality. Prior
work (Watson et al., 2021; Song et al., 2021a) finds that samplers that are
optimal with respect to ELBO often lead to worse sample quality and Fréchet
Inception Distance (FID) scores (Heusel et al., 2017), especially with few
inference steps. We propose the use of a perceptual loss within the DDSS
framework to find high-fidelity diffusion samplers, motivated by prior work
showing that their optimization leads to solutions that correlate better with
human perception of quality. We empirically find that using DDSS with the
Kernel Inception Distance (KID) (Bińkowski et al., 2018) as the perceptual
loss indeed leads to fast samplers with significantly better image quality
than prior work (see Figure 1). Moreover, our method is robust to different
choices of kernels for KID.

Our main contributions are as follows:

  1. 1.

We propose Differentiable Diffusion Sampler Search (DDSS), which uses the
reparametrization trick and gradient rematerialization to optimize over a
parametric family of fast samplers for diffusion models.

  2. 2.

We identify a parametric family of Generalized Gaussian Diffusion Model (GGDM)
that admits high-fidelity fast samplers for diffusion models.

  3. 3.

We show that using DDSS to optimize samplers by minimizing the Kernel
Inception Distance leads to fast diffusion model samplers with state-of-the-
art sample quality scores.

##  2 Background on Denoising Diffusion Implicit Models

We start with a brief review on DDPM (Ho et al., 2020) and DDIM (Song et al.,
2020). DDPMs pre-specify a Markovian forward diffusion process, which
gradually adds noise to data in $T$ steps. Following the notation of Ho et al.
(2020),

$\displaystyle q({\bm{x}}_{0},...,{\bm{x}}_{T})$ $\displaystyle=$
$\displaystyle q({\bm{x}}_{0})\prod_{t=1}^{T}q({\bm{x}}_{t}|{\bm{x}}_{t-1})$
(1) $\displaystyle q({\bm{x}}_{t}|{\bm{x}}_{t-1})$ $\displaystyle=$
$\displaystyle\mathcal{N}({\bm{x}}_{t}|\sqrt{1-\beta_{t}}{\bm{x}}_{s},\beta_{t}{\bm{I}}),$
(2)

where $q({\bm{x}}_{0})$ is the data distribution, and $\beta_{t}$ is the
variance of Gaussian noise added at step $t$. To be able to gradually convert
noise to data, DDPMs learn to invert (1) with a model
$p_{\theta}({\bm{x}}_{t-1}|{\bm{x}}_{t})$, which is trained by maximizing a
(possibly reweighted) evidence lower bound (ELBO):

$\mathbb{E}_{q}\left[D_{\mathrm{KL}}[q({\bm{x}}_{T}|{\bm{x}}_{0})\|p({\bm{x}}_{T})]+\sum_{t=2}^{T}D_{\mathrm{KL}}[q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})\|p_{\theta}({\bm{x}}_{t-1}|{\bm{x}}_{t})]-\log
p_{\theta}({\bm{x}}_{0}|{\bm{x}}_{1})\right]$ (3)

DDPMs specifically choose the model to be parametrized as

$\displaystyle\begin{split}p_{\theta}({\bm{x}}_{t-1}|{\bm{x}}_{t})&=q\left({\bm{x}}_{t-1}\middle|{\bm{x}}_{t},\hat{{\bm{x}}}_{0}=\frac{1}{\sqrt{\bar{\alpha}_{t}}}({\bm{x}}_{t}-\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{\theta}({\bm{x}}_{t},t))\right)\\\
&=\mathcal{N}\left({\bm{x}}_{t-1}\middle|\frac{1}{\sqrt{1-\beta_{t}}}\left({\bm{x}}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)\right),\frac{1-\bar{\alpha}_{t}}{1-\bar{\alpha}_{t-1}}\beta_{t}{\bm{I}}_{d}\right)\end{split}$
(4)

where $\bar{\alpha_{t}}=\prod_{s=1}^{t}(1-\beta_{t})$ for each $t$. With this
parametrization, maximizing the ELBO is equivalent to minimizing a weighted
sum of denoising score matching objectives (Vincent, 2011).

The seminal work of Song et al. (2020) presents Denoising Diffusion Implicit
Models (DDIM): a family of evidence lower bounds (ELBOs) with corresponding
forward diffusion processes and samplers. All of these ELBOs share the same
marginals as DDPM, but allow arbitrary choices of posterior variances.
Specifically, Song et al. (2020) note that is it possible to construct
alternative ELBOs with only a subsequence of the original timesteps
$S\subset\\{1,...,T\\}$ that shares the same marginals as the construction
above (i.e., $q_{S}({\bm{x}}_{t}|{\bm{x}}_{0})=q({\bm{x}}_{t}|{\bm{x}}_{0})$
for every $t\in S$, so $q_{S}$ defines a faster sampler compatible with the
pre-trained model) by simply using the new contiguous timesteps in the
equations above. They also show it is also possible to construct an infinite
family of non-Markovian processes $\\{q_{\sigma}:\sigma\in[0,1]^{T-1}\\}$
where each $q_{\sigma}$ also shares the same marginals as the original forward
process with:

$q_{\sigma}({\bm{x}}_{0},...,{\bm{x}}_{t})=q({\bm{x}}_{0})q({\bm{x}}_{T}|{\bm{x}}_{0})\prod_{t=1}^{T-1}q_{\sigma}({\bm{x}}_{t}|{\bm{x}}_{t+1},{\bm{x}}_{0})$
(5)

and where the posteriors are defined as

$q_{\sigma}({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t-1}\middle|\sqrt{\bar{\alpha}_{t-1}}{\bm{x}}_{0}+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_{t}^{2}}\cdot\frac{{\bm{x}}_{t}-\sqrt{\bar{\alpha}_{t}}{\bm{x}}_{0}}{\sqrt{1-\bar{\alpha}_{t}}},\sigma_{t}^{2}{\bm{I}}_{d}\right)$
(6)

Song et al. (2020) empirically find that the extreme case of using all-zero
variances (a.k.a. DDIM($\eta=0$)) consistently helps with sample quality in
the few-step regime. Combined with a good selection of timesteps to evaluate
the modeled score function (a.k.a. strides), DDIM($\eta=0$) establishes the
current state-of-the-art for few-step diffusion model sampling with the
smallest inference step budgets. Our key contribution that allows improving
sample quality significantly by optimizing sampler families is constructing a
family that generalizes DDIM. See Section 4 for a more complete treatment of
our novel GGDM family.

##  3 Differentiable Diffusion Sampler Search (DDSS)

We now describe DDSS, our approach to search for fast high-fidelity samplers
with a limited budget of $K<T$ steps. Our key observation is that one can
backpropagate through the sampling process of a diffusion model via the
reparamterization trick (Kingma & Welling, 2013). Equipped with this, we can
now use stochastic gradient descent to learn fast samplers by optimizing any
given differentiable loss function over a minibatch of model samples.

We begin with a pre-trained DDPM and a family of $K$-step samplers that we
wish to optimize for the given DDPM. We parametrize this family’s degrees of
freedom as simple transformations of trainable variables. We experiment with
the following families in this paper, but emphasize that DDSS is applicable to
any other family where model samples are differentiable with respect to the
trainable variables:

  * •

DDIM: we parametrize the posterior variances with
$\sigma_{t}=\sqrt{1-\bar{\alpha}_{t-1}}$ $\mathrm{sigmoid}(v_{t})$, where
$v_{1},...,v_{K}$ are trainable variables (the $\sqrt{1-\bar{\alpha}_{t-1}}$
constant ensures real-valued mean coefficients; see the square root in
Equation 6).

  * •

VARS: we parametrize the marginal variances of a DDPM as
$\mathrm{cumsum}(\mathrm{softmax}([{\bm{v}};1]))_{t}$ instead of fixing them
to $1-\bar{\alpha}_{t}$. This ensures they are monotonically increasing with
respect to $t$ (appending a one to ensure $K$ degrees of freedom).

  * •

GGDM: a new family of non-Markovian samplers for diffusion models with more
degrees of freedom illustrated in Figure 2 and defined in Section 4. We
parametrize $\mu_{tu}$ and $\sigma_{t}$ of a GGDM for all $t$ as sigmoid
functions of trainable variables.

  * •

GGDM +PRED: we parametrize all the $\mu_{tu}$ and $\sigma_{t}$ identically to
GGDM, but also learn the marginal coefficients with a
$\mathrm{cumsum}\circ\mathrm{softmax}$ parametrization (identical to VARS)
instead of computing them via Theorem 1, as well as the coefficients that
predict ${\bm{x}}_{0}$ from $a_{t}{\bm{x}}_{t}-b_{t}\bm{\epsilon}$ with
$1+\mathrm{softplus}$ and $\mathrm{softplus}$ parametrizations.

  * •

[family]+TIME: for any sampler family, we additionally parametrize the
timesteps used to query the score model with a
$\mathrm{cumsum}\circ\mathrm{softmax}$ parametrization (identical to VARS).

Figure 2: Illustration of GGDM. To improve sample quality, our novel family of
samplers combines information from all previous (noisier) images at every
denoising step.

As we will show in the experiments, despite the fact that our pre-trained
DDPMs are trained with discrete timesteps, learning the timesteps is still
helpful. In principle, this should only be possible for DDPMs trained with
continuous time sampling (Chen et al., 2021a; Song et al., 2021b; Kingma et
al., 2021), but in practice we find that DDPMs trained with continuously
embedded discrete timesteps are still well-behaved when applied at timesteps
not present during training. We think this is due to the regularity of the
sinusoidal positional encodings Vaswani et al. (2017) used in these model
architectures and training with a sufficiently large number of timesteps $T$.

###  3.1 Differentiable sample quality scores

We can differentiate through a stochastic sampler using the reparameterization
trick, but the question of which objective to optimize still remains. Prior
work has shown that optimizing log-likelihood can actually worsen sample
quality and FID scores in the few-step regime (Watson et al., 2021; Song et
al., 2021a). Thus, we instead design a perceptual loss which simply compares
mean statistics between model samples and real samples in the neural network
feature space. These types of objectives have been shown in prior work to
better correlate with human perception of sample quality (Johnson et al.,
2016; Heusel et al., 2017), which we also confirm in our experiments.

We rely on the representations of the penultimate layer of a pre-trained
InceptionV3 classifier (Szegedy et al., 2016) and optimize the Kernel
Inception Distance (KID) (Bińkowski et al., 2018). Let $\phi({\bm{x}})$ denote
the inception features of an image ${\bm{x}}$ and $p_{\psi}$ represent a
diffusion sampler with trainable parameters $\psi$. For a linear kernel, which
works best in our experiments, the objective is:

$\mathcal{L}(\psi)=\mathop{\mathbb{E}}_{{\bm{x}}_{p}\sim
p_{\psi}}\mathop{\mathbb{E}}_{{\bm{x}}^{\prime}_{p}\sim
p_{\psi}}\phi({\bm{x}}_{p})^{\top}\phi({\bm{x}}^{\prime}_{p})-\mathop{\mathbb{E}}_{{\bm{x}}_{p}\sim
p_{\psi}}\mathop{\mathbb{E}}_{{\bm{x}}_{q}\sim
q}\phi({\bm{x}}_{p})^{\top}\phi({\bm{x}}_{q})$ (7)

More generally, KID can be expressed as:

$\mathcal{L}_{\textrm{KID}}(\psi)~{}=~{}\left\|\mathop{\mathbb{E}}_{{\bm{x}}_{p}\sim
p_{\psi}}f^{*}({\bm{x}}_{p})-\mathop{\mathbb{E}}_{{\bm{x}}_{q}\sim
q}f^{*}({\bm{x}}_{q})\right\|_{2}^{2}~{},$ (8)

where $f^{*}({\bm{x}})=\mathbb{E}_{{\bm{x}}^{\prime}_{p}\sim
p_{\psi}}k_{\phi}({\bm{x}},{\bm{x}}^{\prime}_{p})-\mathbb{E}_{{\bm{x}}^{\prime}_{q}\sim
q({\bm{x}}_{0})}k_{\phi}({\bm{x}},{\bm{x}}^{\prime}_{q})$ is the witness
function for any differentiable, positive definite kernel $k$, and
$k_{\phi}({\bm{x}},{\bm{y}})=k(\phi({\bm{x}}),\phi({\bm{y}}))$. Note that
$f^{*}$ attains the supremum of the MMD. To enable stochastic gradient
descent, we use an unbiased estimator of KID using a minibatch of $n$ model
samples ${\bm{x}}_{p}^{(1)}\ldots{\bm{x}}_{p}^{(n)}\sim p_{\psi}$ and $n$ real
samples ${\bm{x}}_{q}^{(1)}\ldots{\bm{x}}_{q}^{(n)}\sim q$:

$\displaystyle\frac{1}{n(n-1)}\sum_{i\neq
j}^{n}k_{\phi}({\bm{x}}_{p}^{(i)},{\bm{x}}_{p}^{(j)})-\frac{2}{n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}k_{\phi}({\bm{x}}_{p}^{(i)},{\bm{x}}_{q}^{(j)})+c~{},$
(9)

where $c$ is constant in $\psi$. Since the sampling chain of any Gaussian
diffusion process admits using the reparametrization trick, our loss function
is fully differentiable with respect to the trainable variables $\psi$. We
empirically find that using the perceptual features is crucial; i.e., by
trying $\phi({\bm{x}})={\bm{x}}$ to compare images directly on pixel space
rather than neural network feature space (as above), we observe that our
method makes samples consistently worsen in apparent quality as training
progresses.

###  3.2 Gradient Rematerialization

In order for backpropagation to be feasible under reasonable memory
constraints, one final problem must be addressed: since we are taking
gradients with respect to model samples, the cost in memory to maintain the
state of the forward pass scales linearly with the number of inference steps,
which can quickly become unfeasible considering the large size of typical DDPM
architectures. To address this issue, we use gradient rematerialization (Kumar
et al., 2019b). Instead of storing a particular computation’s output from the
forward pass required by the backward pass computation, we recompute it on the
fly. To trade $\mathcal{O}(K)$ memory cost for $\mathcal{O}(K)$ computation
time, we simply rematerialize calls to the pre-trained DDPM (i.e., the
estimated score function), but keep in memory all the progressively denoised
images from the sampling chain. In JAX (Bradbury et al., 2018), this is
trivial to implement by simply wrapping the score function calls with
`jax.remat`.

##  4 Generalized Gaussian Diffusion Models

We now present Generalized Gaussian Diffusion Model (GGDM), our novel family
of Gaussian diffusion processes that includes DDIM as a special case mentioned
in section 3. We define a joint distribution with no independence assumptions

$q_{\mu,\sigma}({\bm{x}}_{0},...,{\bm{x}}_{T})=q({\bm{x}}_{0})q({\bm{x}}_{T}|{\bm{x}}_{0})\prod_{t=1}^{T-1}q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t},{\bm{x}}_{0})$
(10)

where the new factors are defined as

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t},{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|\sum_{u\in
S_{t}}\mu_{tu}{\bm{x}}_{u},\sigma_{t}^{2}{\bm{I}}_{d}\right)$ (11)

(letting $S_{t}=\\{0,...,T\\}\setminus\\{1,...,t\\}$ for notation
compactness), with $\sigma_{t}$ and $\mu_{tu}$ free parameters $\forall
t\in\\{1,...,T\\},u\in S_{t}$. In other words, when predicting the next, less
noisy image, the sampler can take into account all the previous, noisier
images in the sampling chain, and similarly to DDIM, we can also control the
sampler’s variances. As we prove in the appendix (A.2), this construction
admits Gaussian marginals, and we can differentiably compute the marginal
coefficients from arbitrary choices of $\mu$ and $\sigma$:

Theorem 1. Given some $t\in\\{1,...,T\\}$, let $a_{tu}^{(1)}=\mu_{tu}$
$\forall u\in S_{t}$ and $v_{t}^{(1)}=\sigma_{t}^{2}$. For each
$i\in\\{1,...,T-t\\}$, recursively define

$a_{tu}^{(i+1)}=a_{t,t+i}^{(i)}\mu_{t+i,u}+a_{tu}^{(i)}\;\forall u\in
S_{t+i}\quad\textrm{and}\quad
v_{t}^{(i+1)}=v_{t}^{(i)}+\left(a_{t,t+i}^{(i)}\sigma_{t+i}\right)^{2}.$

Then, it follows that

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t+i},{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|\sum_{u\in
S_{t+i}}a^{(i+1)}_{tu}{\bm{x}}_{u},v_{t}^{(i+1)}{\bm{I}}_{d}\right).$ (12)

In other words, instead of letting the $\beta_{t}$ (or equivalently, the
$\bar{\alpha}_{t}$) define the forward process as done by a usual DDPM, the
GGDM family lets the $\mu_{tu}$ and $\sigma_{t}$ define the process. In
particular, an immediate corollary of Theorem 1 is that the marginal
coefficients are given by

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|a^{(T-t+1)}_{t0}{\bm{x}}_{0},v_{t}^{(T-t+1)}{\bm{I}}_{d}\right)$
(13)

The reverse process is thus defined as
$p({\bm{x}}_{T})\prod_{t=1}^{T}p({\bm{x}}_{t-1}|{\bm{x}}_{t})$ with
$p({\bm{x}}_{T})\sim\mathcal{N}(\bm{0},{\bm{I}}_{d}))$ and

$p_{\theta}({\bm{x}}_{t}|{\bm{x}}_{>t})=q_{\mu,\sigma}\left({\bm{x}}_{t}|{\bm{x}}_{>t},\hat{{\bm{x}}}_{0}=\frac{1}{a^{(T-t+1)}_{t0}}\left({\bm{x}}_{t}-\sqrt{v_{t}^{(T-t+1)}}\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)\right)\right).$
(14)

###  4.1 Ignoring the matching marginals condition

Unlike DDIM, the GGDM family does not guarantee that the marginals of the new
forward process match that of the original DDPM. We empirically find, however,
that this condition can often be too restrictive and better samplers exist
where the marginals
$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|a_{t0}^{(T-t+1)}{\bm{x}}_{0},v_{t}^{(T-t+1)}{\bm{I}}_{d}\right)$
of the new forward process differ from the original DDPM’s marginals. We
verify this empirically by applying DDSS to both the family of DDIM sigmas and
DDPM variances (“VARS” in Section 3): both sampler families have the same
number of parameters (the reverse process variances), but the latter does not
adjust the mean coefficients like DDIM to ensure matching marginals and still
achieves similar or better scores than the former across sample quality
metrics (and even outperforms the DDIM($\eta=0$) baseline); see Section 5.2.

##  5 Experiments

Table 1: FID and IS scores for DDSS against baseline methods for a DDPM trained on CIFAR10 with the $L_{\textrm{simple}}$ objective proposed by (Ho et al., 2020). FID scores (lower is better) are the numbers at the left of each entry, and IS scores (higher is better) are at the right. Sampler $\setminus$ $K$ | 5 | 10 | 15 | 20 | 25
---|---|---|---|---|---
DDPM (linear stride) | 84.27 / 5.396 | 43.39 / 7.034 | 31.40 / 7.609 | 25.94 / 7.879 | 22.60 / 8.043
DDPM (quadratic stride) | 76.25 / 5.435 | 42.03 / 6.965 | 27.78 / 7.714 | 20.225 / 8.128 | 16.17 / 8.350
DDIM (linear stride) | 44.41 / 6.750 | 19.11 / 7.965 | 14.06 / 8.190 | 11.82 / 8.420 | 10.52 / 8.512
DDIM (quadratic stride) | 32.66 / 7.090 | 13.62 / 8.190 | 9.318 / 8.495 | 7.500 / 8.641 | 6.560 / 8.759
GGDM +PRED+TIME |  13.77 / 8.520 |  8.227 / 8.903 |  6.115 / 9.050 |  4.722 / 9.261 |  4.250 / 9.186
Table 2: FID / IS scores for DDSS against baseline methods for a DDPM trained on ImageNet 64x64 with the $L_{\textrm{hybrid}}$ objective proposed by Nichol & Dhariwal (2021). Sampler $\setminus$ $K$ | 5 | 10 | 15 | 20 | 25
---|---|---|---|---|---
DDPM (linear stride) | 122.0 / 5.878 | 58.78 / 10.67 | 39.30 / 13.22 | 31.36 / 14.72 | 26.36 / 15.71
DDPM (quadratic stride) | 394.8 / 1.351 | 129.5 / 5.997 | 80.10 / 9.595 | 61.34 / 11.60 | 49.60 / 13.01
DDIM (linear stride) | 135.4 / 5.898 | 40.70 / 12.225 | 28.54 / 13.99 | 24.225 / 14.75 | 22.13 / 15.16
DDIM (quadratic stride) | 409.1 / 1.380 | 148.6 / 5.533 | 67.65 / 9.842 | 45.60 / 11.99 | 36.11 / 13.225
GGDM +PRED+TIME |  55.14 / 12.90 |  37.32 / 14.76 |  24.69 / 17.225 |  20.69 /17.92 |  18.40 / 18.12

In order to emphasize that our method is compatible with any pre-trained DDPM,
we apply our method on pre-trained DDPM checkpoints from prior work.
Specifically, we experiment with the DDPM trained by Ho et al. (2020) with
$L_{\textrm{simple}}$ on CIFAR10, as well as a DDPM following the exact
configuration of Nichol & Dhariwal (2021) trained on ImageNet 64x64 (Deng et
al., 2009) with their $L_{\textrm{hybrid}}$ objective (with the only
difference being that we trained the latter ourselves for 3M rather than 1.5M
steps). Both of these models utilize adaptations of the UNet architecture
(Ronneberger et al., 2015) that incorporate self-attention layers (Vaswani et
al., 2017).

We evaluate all of our models on both FID and Inception Score (IS) (Salimans
et al., 2016), comparing the samplers discovered by DDSS against DDPM and DDIM
baselines with linear and quadratic strides. As previously mentioned, more
recent methods for fast sampling are outperformed by DDIM when the budget of
inference steps is as small as those we utilize in this work (5, 10, 15, 20,
25). All reported results on both of these approximate sample quality metrics
were computed by comparing 50K model and training data samples, as is standard
in the literature. Also as is standard, IS scores are computed 10 times, each
time on 5K samples, and then averaged.

In all of our experiments, we optimize the DDSS objective presented in Section
3.1 with the following hyperparameters:

  1. 1.

For every family of models we search over, we initialize the degrees of
freedom such that training begins with a sampler matching DDPM with $K$
substeps following Song et al. (2020); Nichol & Dhariwal (2021).

  2. 2.

We apply gradient updates using the Adam optimizer (Kingma & Ba, 2015). We
sweeped over the learning rate and used $\lambda=0.0005$. We did not sweep
over other Adam hyperparameters and kept $\beta_{1}=0.9$, $\beta_{2}=0.999$,
and $\epsilon=1\times 10^{-8}$.

  3. 3.

We tried batch sizes of 128 and 512 and opted for the latter, finding that it
leads to better sample quality upon inspection. Since the loss depends on
averages over examples as our experiments are on unconditional generation,
this choice was expected.

  4. 4.

We run all of our experiments for 50K training steps and evaluate the
discovered samplers at this exact number of training steps. We did not sweep
over this value.

We include our main results in Table 1 for CIFAR10 and Table 2 for ImageNet
64x64, comparing DDSS applied to GGDM +PRED+TIME against DDPM and DDIM
baselines with linear and quadratic strides. All models use a linear kernel,
i.e., $k_{\phi}({\bm{x}},{\bm{y}})=\phi({\bm{x}})^{\top}\phi({\bm{y}})$, which
we found to perform slightly better than the cubic kernel used by Bińkowski et
al. (2018) (we ablate this in section 5.1). We omit the use of the learned
variances of the ImageNet 64x64 model (i.e., following Nichol & Dhariwal
(2021)), as we search for the variances ourselves via DDSS. We include samples
for 5, 10 and 25 steps comparing the strongest DDIM baselines to DDSS + GGDM
with a learned stride; see Figures 1 and 3. We include additional ImageNet
64x64 samples (A.1) and results for larger resolution datasets (A.4) in the
appendix.

###  5.1 Ablations for KID kernel and GGDM variants

As our approach is compatible with any choice of KID kernel, we experiment
with different choices of kernels. Namely, we try the simplest possible linear
kernel, $k_{\phi}({\bm{x}},{\bm{y}})=\phi({\bm{x}})^{\top}\phi({\bm{y}})$, as
well as the cubic kernel
$k_{\phi}({\bm{x}},{\bm{y}})=\left(\frac{1}{d}\phi({\bm{x}})^{\top}\phi({\bm{y}})+1\right)^{3}$
used by Bińkowski et al. (2018). We compare the performance of these two
kernels, as well as different variations of GGDM (i.e., with and without TIME
and PRED as defined in Section 3). Results are included for CIFAR10 across all
budgets in Table 3. We also include a smaller version of this ablation for
ImageNet 64x64 in the appendix (A.3).

Table 3: FID / IS scores for the KID kernel ablation on CIFAR10. When not learning the timesteps, we fix them to a quadratic stride, as Table 1 shows this performs best on CIFAR10. Sampler $\setminus$ $K$ | 5 | 10 | 15 | 20 | 25
---|---|---|---|---|---
DDSS (linear kernel) |  |  |  |  |
GGDM +PRED+TIME | 13.77 / 8.520 | 8.227 / 8.903 | 6.115 / 9.050 |  4.722 / 9.261 |  4.250 / 9.186
GGDM +PRED | 14.26 / 8.406 | 8.617 / 8.842 | 5.939 / 9.035 | 4.893 / 9.153 | 4.574 / 9.145
GGDM +TIME | 12.85 / 8.383 |  7.858 / 8.895 | 6.265 / 9.075 | 5.367 / 9.136 | 4.887 / 9.229
GGDM) | 14.45 / 8.281 | 8.154 / 8.892 | 7.045 / 8.939 | 5.477 / 9.183 | 4.815 / 9.189
DDSS (cubic kernel) |  |  |  |  |
GGDM +PRED+TIME | 14.41 / 8.527 | 8.2257 / 9.007 |  5.895 / 9.036 | 4.932 / 9.092 | 4.278/ 9.286
GGDM +PRED | 14.39 / 8.401 | 8.977 / 8.870 | 6.517 / 8.970 | 4.915 / 9.132 | 4.471 / 9.247
GGDM +TIME |  12.35 / 8.406 | 7.879 / 8.852 | 6.682 / 8.999 | 5.639 / 9.058 | 4.631 / 9.189
GGDM | 14.57 / 8.297 | 8.2252 / 8.836 | 6.727 / 8.904 | 5.569 / 9.177 | 4.668 / 9.192

The results in this ablation show that the contributions of the linear kernel,
timestep learning, and the empirical choice of learning the coefficients that
predict ${\bm{x}}_{0}$ all slightly contribute to better FID and IS scores.
Importantly, however, removing any of these additions still allows us to
comfortably outperform the strongest baselines. See also the results on LSUN
in the appendix A.4, which also do not include these additional trainable
variables.

###  5.2 Search space ablation

Now, in order to further demonstrate the key importance of optimizing our GGDM
family to find high-fidelity samplers, we also apply DDSS to the less general
DDIM and VARS families. We show that, while we still attain better scores than
a regular DDPM, searching these less flexible families of samplers does not
yield improvements as significant as with out novel GGDM family. In
particular, optimizing the DDIM sigma coefficients does not outperform the
corresponding DDIM($\eta=0$) baseline on CIFAR10, which is not a surprising
result as Song et al. (2020) show empirically that most choices of the
$\sigma_{t}$ degrees of freedom lead to worse FID scores than setting them all
to 0. These results also show that optimizing the VARS can outperform DDSS
applied to the DDIM family, and even the strongest DDIM($\eta=0$) baselines
for certain budgets, justifying our choice of not enforcing the marginals to
match (as discussed in Section 4.1).

Table 4: FID / IS scores for the DDSS search space ablation on CIFAR10. All runs fix the timesteps to a quadratic stride and use a linear kernel except for the last row (we only include the GGDM results for ease of comparison). Sampler $\setminus$ $K$ | 5 | 10 | 15 | 20 | 25
---|---|---|---|---|---
DDIM($\eta=0$) | 32.66 / 7.090 |  13.62 / 8.190 |  9.318 / 8.495 | 7.500 / 8.641 | 6.560 / 8.759
DDSS |  |  |  |  |
VARS | 33.08 / 7.096 | 15.33 / 8.559 | 9.693 / 8.845 |  7.297 / 8.924 |  6.172 / 9.057
DDIM |  32.61 / 7.084 | 16.29 / 7.966 | 11.31 / 8.372 | 9.120 / 8.563 | 7.853 / 8.644
GGDM | 14.45 / 8.281 | 8.154 / 8.892 | 7.045 / 8.939 | 5.477 / 9.183 | 4.815 / 9.189
GGDM +PRED+TIME | 13.77 / 8.520 | 8.227 / 8.903 | 6.115 / 9.050 | 4.722 / 9.261 | 4.250 / 9.186

##  6 Discussion

When applied to a sufficiently flexible family (such as the GGDM family
proposed in this work), DDSS consistently finds samplers that achieve better
image generation quality than the strongest baselines in the literature for
very few steps. This is qualitatively apparent in non-cherrypicked samples
(e.g., DDIM($\eta=0$) tends to generate blurrier images and with less
background details as the budget decreases), and multiple quantitative sample
quality metrics (FID and IS) also reflect these results. Still, we observe
limitations to our method. Finding samplers with inference step budgets as
small as $K<10$ that have little apparent loss in quality remains challenging
with our proposed search family. And, while on CIFAR10 the metrics indicate
significant relative improvement over sample quality metrics, the relative
improvement on ImageNet 64x64 is less pronounced. We hypothesize that this is
an inherent difficulty of ImageNet due to its high diversity of samples, and
that in order to retain sample quality and diversity, it might be impossible
to escape some minimum number of inference steps with score-based models as
they might be crucial to mode-breaking.

Beyond the empirical gains of applying our procedure, our findings shed
further light into properties of pre-trained score-based generative models.
First, we show that without fine-tuning a DDPM’s parameters, these models are
already capable of producing high-quality samples with very few inference
steps, though the default DDPM sampler in this regime is usually suboptimal
when using a few-step sampler. We further show that better sampling paths
exist, and interestingly, these are determined by alternative variational
lower bounds to the data distribution that make use of the score-based model
but do not necessarily share the same marginals as the original DDPM forward
process. Our findings thus suggest that enforcing this marginal-sharing
constraint is unnecessary and can be too restrictive in practice.

##  7 Other Related Work

| 5 steps | 10 steps | 20 steps | Reference
---|---|---|---|---
DDIM |  |  |  |
Ours |  |  |  |
Figure 3: Non-cherrypicked samples for a DDPM trained on CIFAR10, comparing
the strongest DDIM($\eta=0$) baseline and our approach. All samples were
generated with the same random seeds. For reference, we include DDPM samples
using all 1,000 steps (top right) and real images (bottom right).

Besides DDIM (Song et al., 2020), there have been more recent attempts at
reducing the number of inference steps for DDPMs. Jolicoeur-Martineau et al.
(2021) proposed a dynamic step size SDE solver that can reduce the number of
calls to the modeled score function to $\sim 150$ on CIFAR10 (Krizhevsky et
al., 2009) with minimal cost in FID scores, but quickly falls behind
DDIM($\eta=0$) with as many as 50 steps. Watson et al. (2021) proposed a
dynamic programming algorithm that chooses log-likelihood optimal strides, but
find that log-likelihood reduction has a mismatch with FID scores,
particularly with in the very few step regime, also falling behind
DDIM($\eta=0$) in this front. Other methods that have been shown to help
sample quality in the few-step regime include non-Gaussian variants of
diffusion models (Nachmani et al., 2021) and adaptively adjusting the sampling
path by introducing a noise level estimating network (San-Roman et al., 2021),
but more thorough evaluation of sample quality achieved by these approaches is
needed with budgets as small as those considered in this work.

Other approaches to sampling DDPMs have also been recently proposed, though
not for the explicit purpose of efficient sampling. Song et al. (2021b) derive
a reverse SDE that, when discretized, uses different coefficients than the
ancestral samplers considered in this work. The same authors also derive
“corrector” steps, which introduce additional calls to the pre-trained DDPM as
a form of gradient ascent (Langevin dynamics) that help with quality but
introduce computation cost, as well as an alternative sampling procedure using
a probability flow ODE that shares the same marginals as the DDPM’s original
forward process. Huang et al. (2021) generalize this family of samplers to a
“plug-in reverse SDE” that interpolates between a probability flow ODE and the
reverse SDE, similarly to how the DDIM $\eta$ interpolates between an implicit
probabilistic model and a stochastic reverse process. Our proposed search
family includes discretizations of most of these cases for Gaussian processes,
notably missing corrector steps, where reusing a single timestep is
considered.

##  8 Conclusion and Future Work

We propose Differentiable Diffusion Sampler Search (DDSS), a method for
finding few-step samplers for Denoising Diffusion Probabilistic Models. We
show how to optimize a perceptual loss over a space of diffusion processes
that makes use of a pre-trained DDPM’s samples by leveraging the
reparametrization trick and gradient rematerialization. Our results
qualitatively and quantitatively show that DDSS is able to significantly
improve sample quality for unconditional image generation over prior methods
on efficient DDPM sampling. The success of our method hinges on searching a
novel, wider family of Generalized Gaussian Diffusion Model (GGDM) than those
identified in prior work (Song et al., 2020). DDSS does not fine-tune the pre-
trained DDPM, only needs to be applied once, has few hyperparameters, and does
not require re-training the DDPM.

Our findings suggest future directions to further reduce the number of
inference steps while retaining high fidelity in generated samples. For
instance, it is plausible to use different representations for the perceptual
loss instead of those of a classifier, e.g., use representations from an
unsupervised model such as SimCLR (Chen et al., 2020), to using internal
representations learned by the pre-trained DDPM itself, which would eliminate
the burden of additional computation. Moreover, considering the demonstrated
benefits of applying DDSS to our proposed GGDM family of samplers (as opposed
to narrower families like DDIM), we motivate future work on identifying more
general families of samplers and investigating whether they help uncover even
better samplers or lead to overfitting. Finally, identifying other variants of
perceptual losses (e.g., that do not sample from the model), or alternative
optimization strategies (e.g., gradient-free methods) that lead to similar
results is important future work. This would make DDSS itself a more efficient
procedure, as gradient-based optimization of our proposed loss requires
extensive memory or computation requirements to backpropagate through the
whole sampling chain.

## References

  * Arjovsky et al. (2017) Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein GAN.  In _arXiv_ , 2017.
  * Bińkowski et al. (2018) Mikołaj Bińkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton.  Demystifying mmd gans.  _arXiv preprint arXiv:1801.01401_ , 2018.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Brock et al. (2018) Andrew Brock, Jeff Donahue, and Karen Simonyan.  Large scale GAN training for high fidelity natural image synthesis.  _arXiv preprint arXiv:1809.11096_ , 2018.
  * Cai et al. (2020) Ruojin Cai, Guandao Yang, Hadar Averbuch-Elor, Zekun Hao, Serge Belongie, Noah Snavely, and Bharath Hariharan.  Learning Gradient Fields for Shape Generation.  In _ECCV_ , 2020.
  * Chen et al. (2021a) Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan.  WaveGrad: Estimating Gradients for Waveform Generation.  In _ICLR_ , 2021a.
  * Chen et al. (2021b) Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, Najim Dehak, and William Chan.  WaveGrad 2: Iterative Refinement for Text-to-Speech Synthesis .  In _INTERSPEECH_ , 2021b.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton.  Big Self-Supervised Models are Strong Semi-Supervised Learners.  In _NeurIPS_ , 2020.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Dhariwal & Nichol (2021) Prafulla Dhariwal and Alex Nichol.  Diffusion models beat gans on image synthesis.  _arXiv preprint arXiv:2105.05233_ , 2021.
  * Goodfellow et al. (2014) Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial networks.  _arXiv preprint arXiv:1406.2661_ , 2014.
  * Gulrajani et al. (2017) Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville.  Improved training of wasserstein gans.  _arXiv preprint arXiv:1704.00028_ , 2017.
  * Heusel et al. (2017) Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _arXiv preprint arXiv:1706.08500_ , 2017.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising Diffusion Probabilistic Models.  _NeurIPS_ , 2020.
  * Ho et al. (2021) Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans.  Cascaded diffusion models for high fidelity image generation.  _arXiv preprint arXiv:2106.15282_ , 2021.
  * Huang et al. (2021) Chin-Wei Huang, Jae Hyun Lim, and Aaron Courville.  A variational perspective on diffusion-based generative models and score matching.  _arXiv preprint arXiv:2106.02808_ , 2021.
  * Johnson et al. (2016) Justin Johnson, Alexandre Alahi, and Li Fei-Fei.  Perceptual losses for real-time style transfer and super-resolution.  In _European conference on computer vision_ , pp. 694–711. Springer, 2016.
  * Jolicoeur-Martineau et al. (2021) Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas.  Gotta go fast when generating data with score-based models.  _arXiv preprint arXiv:2105.14080_ , 2021.
  * Kingma & Ba (2015) Diederik Kingma and Jimmy Ba.  Adam: A Method for Stochastic Optimization.  In _ICLR_ , 2015.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In _ICLR_ , 2013.
  * Kingma et al. (2021) Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  _arXiv preprint arXiv:2107.00630_ , 2021.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  _Technical Report_ , 2009.
  * Kumar et al. (2019a) Kundan Kumar, Rithesh Kumar, Thibault de Boissiere, Lucas Gestin, Wei Zhen Teoh, Jose Sotelo, Alexandre de Brebisson, Yoshua Bengio, and Aaron Courville.  MelGAN: Generative Adversarial Networks for Conditional Waveform Synthesis.  In _NeurIPS_ , 2019a.
  * Kumar et al. (2019b) Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua R Wang.  Efficient rematerialization for deep networks.  In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ , pp. 15172–15181, 2019b.
  * Luo & Hu (2021) Shitong Luo and Wei Hu.  Diffusion probabilistic models for 3d point cloud generation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 2837–2845, 2021.
  * Nachmani et al. (2021) Eliya Nachmani, Robin San Roman, and Lior Wolf.  Non gaussian denoising diffusion models.  _arXiv preprint arXiv:2106.07582_ , 2021.
  * Nichol & Dhariwal (2021) Alex Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  _arXiv preprint arXiv:2102.09672_ , 2021.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _International Conference on Medical image computing and computer-assisted intervention_ , pp. 234–241. Springer, 2015.
  * Salimans et al. (2016) Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training gans.  _Advances in neural information processing systems_ , 29:2234–2242, 2016.
  * San-Roman et al. (2021) Robin San-Roman, Eliya Nachmani, and Lior Wolf.  Noise estimation for generative diffusion models.  _arXiv preprint arXiv:2104.02600_ , 2021.
  * Sohl-Dickstein et al. (2015) Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pp. 2256–2265. PMLR, 2015.
  * Song et al. (2020) Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_ , 2020.
  * Song & Ermon (2019) Yang Song and Stefano Ermon.  Generative Modeling by Estimating Gradients of the Data Distribution.  _NeurIPS_ , 2019.
  * Song et al. (2021a) Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon.  Maximum likelihood training of score-based diffusion models.  _arXiv e-prints_ , pp. arXiv–2101, 2021a.
  * Song et al. (2021b) Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-Based Generative Modeling through Stochastic Differential Equations.  In _ICLR_ , 2021b.
  * Svensén & Bishop (2007) Markus Svensén and Christopher M Bishop.  Pattern recognition and machine learning, 2007.
  * Szegedy et al. (2016) Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna.  Rethinking the Inception Architecture for Computer Vision.  In _CVPR_ , 2016.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention Is All You Need.  In _NIPS_ , 2017.
  * Vincent (2011) Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural Computation_ , 23(7):1661–1674, 2011\.
  * Watson et al. (2021) Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan.  Learning to efficiently sample from diffusion probabilistic models.  _arXiv preprint arXiv:2106.03802_ , 2021.
  * Yu et al. (2015) Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop.  _arXiv preprint arXiv:1506.03365_ , 2015.

##  Appendix A Appendix

###  A.1 Additional ImageNet 64x64 samples

We provide additional samples for our results on ImageNet 64x64. The DDPM and
DDIM($\eta=0$) samples (left and middle, respectively) use a linear stride,
while our DDSS + GGDM samples (right) use a learned stride.

| DDPM | DDIM | Ours
---|---|---|---
5 steps |  |  |
10 steps |  |  |
20 steps |  |  |
Reference |  |  |
Figure A.1: Additional samples on ImageNet 64x64. For reference, we include
DDPM samples with all 4,000 steps (bottom left) and real samples (bottom
middle).

###  A.2 Proof for Theorem 1

Theorem 1. Given some $t\in\\{1,...,T\\}$, let $a_{tu}^{(1)}=\mu_{tu}$
$\forall u\in S_{t}$ and $v_{t}^{(1)}=\sigma_{t}^{2}$. For each
$i\in\\{1,...,T-t\\}$, recursively define

  * •

$a_{tu}^{(i+1)}=a_{t,t+i}^{(i)}\mu_{t+i,u}+a_{tu}^{(i)}$ $\forall u\in
S_{t+i}$

  * •

$v_{t}^{(i+1)}=v_{t}^{(i)}+\left(a_{t,t+i}^{(i)}\sigma_{t+i}\right)^{2}$

Then, it follows that

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t+i},{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|\sum_{u\in
S_{t+i}}a^{(i+1)}_{tu}{\bm{x}}_{u},v_{t}^{(i+1)}{\bm{I}}_{d}\right)$

Proof. Let us prove this result with mathematical induction. Note that, for
each such $t$, we have by definition that

$\displaystyle q_{\mu,\sigma}({\bm{x}}_{t+1}|{\bm{x}}_{>t+1},{\bm{x}}_{0})$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+1}}\mu_{t+1,u}{\bm{x}}_{u},\sigma_{t+1}^{2}{\bm{I}}_{d}\right)$ and
$\displaystyle
q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{t+1},{\bm{x}}_{>t+1},{\bm{x}}_{0})$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t}}\mu_{tu}{\bm{x}}_{u},\sigma_{t}^{2}{\bm{I}}_{d}\right)=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t}}a_{tu}^{(1)}{\bm{x}}_{u},v_{t}^{(1)}{\bm{I}}_{d}\right)$

Therefore, following Svensén & Bishop (2007) (2.115), by prior conjugacy it
follows that

$\displaystyle q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t+1},{\bm{x}}_{0})$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|a_{t,t+1}^{(1)}\textstyle\sum_{u\in
S_{t+1}}\mu_{t+1,u}{\bm{x}}_{u}+\textstyle\sum_{u\in
S_{t+1}}a_{tu}^{(1)},\left(v_{t}^{(1)}+a_{t,t+1}^{(1)}\sigma_{t+1}^{2}a_{t,t+1}^{(1)}\right){\bm{I}}_{d}\right)$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+1}}\left(a_{t,t+1}^{(1)}\mu_{t+1,u}+a_{tu}^{(1)}\right){\bm{x}}_{u},v_{t}^{(2)}{\bm{I}}_{d}\right)$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+1}}a_{tu}^{(2)}{\bm{x}}_{u},v_{t}^{(2)}{\bm{I}}_{d}\right).$

This proves the base case for our induction argument. Now, let us prove the
inductive step. Suppose there exists some integer $j\in\\{1,...,T-t+1\\}$ such
that

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t+j},{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+j}}a_{tu}^{(j+1)}{\bm{x}}_{u},v_{t}^{(j+1)}{\bm{I}}_{d}\right).$

By definition, we already know
$q({\bm{x}}_{t+j+1}|{\bm{x}}_{>t+j+1},{\bm{x}}_{0})$, so we have

$\displaystyle
q_{\mu,\sigma}({\bm{x}}_{t+j+1}|{\bm{x}}_{>t+j+1},{\bm{x}}_{0})$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t+j+1}\middle|\textstyle\sum_{u\in
S_{t+j+1}}\mu_{t+j+1,u}{\bm{x}}_{u},\sigma_{t+j+1}^{2}{\bm{I}}_{d}\right)$ and
(rewriting the inductive hypothesis) $\displaystyle
q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{t+j+1},{\bm{x}}_{>t+j+1},{\bm{x}}_{0})$
$\displaystyle=\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+j}}a_{tu}^{(j+1)}{\bm{x}}_{u},v_{t}^{(j+1)}{\bm{I}}_{d}\right).$

Therefore, by prior conjugacy again, it follows that

$\displaystyle q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{>t+j+1},{\bm{x}}_{0})$
$\displaystyle=$
$\displaystyle\mathcal{N}\left({\bm{x}}_{t}\middle|a_{t,t+j+1}^{(j+1)}\textstyle\sum_{u\in
S_{t+j}}\mu_{t+j+1,u}{\bm{x}}_{u}+\textstyle\sum_{u\in
S_{t+j}}a_{tu}^{(j+1)},\left(v_{t}^{(j+1)}+a_{t,t+j+1}^{(j+1)}\sigma_{t+j+1}^{2}a_{t,t+j+1}^{(j+1)}\right){\bm{I}}_{d}\right)$
$\displaystyle=$
$\displaystyle\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+j}}\left(a_{t,t+j+1}^{(j+1)}\mu_{t+j+1,u}+a_{tu}^{(j+1)}\right){\bm{x}}_{u},v_{t}^{(j+2)}{\bm{I}}_{d}\right)$
$\displaystyle=$
$\displaystyle\mathcal{N}\left({\bm{x}}_{t}\middle|\textstyle\sum_{u\in
S_{t+j+1}}a_{tu}^{(j+2)}{\bm{x}}_{u},v_{t}^{(j+2)}{\bm{I}}_{d}\right).$

This concludes the proof of the inductive step. Hence, we have proven the
result for any $i\in\\{1,...,T-t\\}$. In particular,

$q_{\mu,\sigma}({\bm{x}}_{t}|{\bm{x}}_{0})=\mathcal{N}\left({\bm{x}}_{t}\middle|a_{t0}^{(T-t+1)}{\bm{x}}_{0},v_{t}^{(T-t+1)}{\bm{I}}_{d}\right).\quad\square$

###  A.3 Additional ablation of KID kernel and GGDM variants for ImageNet
64x64

We also ran a smaller version of the ablation results presented in Section
5.1, but for ImageNet 64x64 instead of CIFAR10, as these are more
computationally intensive to do a full grid search. Results for a step budget
$K=15$ are included below. When not learning the timesteps, we fix them to a
linear stride, as Table 2 shows this performs best on ImageNet 64x64.

Sampler $\setminus$ $K$ | 15
---|---
DDSS (linear kernel) |
GGDM +PRED+TIME |  24.69 / 17.225
GGDM +PRED | 27.08 / 16.44
GGDM +TIME | 25.73 / 17.27
GGDM | 28.34 / 16.63
DDSS (cubic kernel) |
GGDM +PRED+TIME | 26.52 / 16.29
GGDM +PRED | 27.82 / 16.3
GGDM +TIME | 26.87 / 16.99
GGDM | 28.83 / 16.32

###  A.4 Results on larger resolution datasets

We include results for LSUN (Yu et al., 2015) bedrooms and churches at the
128x128 resolution. We trained the models for 400K and 200K steps
(respectively), and all other hyperparameters are identical: we use the Adam
optimizer with learning rate 0.0003 (linearly warmed up for the first 1000
training steps), batch size 2048, gradient clipping at norms over 1.0, dropout
of 0.1, and EMA over the weights with decay rate 0.9999. We train the models
using a linear stride of 1000 evenly-spaced timesteps, fixing the log-signal-
to-noise-ratio schedule to a cosine function monotonically decreasing from 20
to -20. The ELBO is reweighted with $L_{\textrm{simple}}$ following Ho et al.
(2020), but we additionally reweight each term by $\max(1,\textrm{SNR})$ which
we found to be slightly helpful in resulting FID scores (note this is
equivalent to minimizing the worst mean squared error between either the
${\bm{x}}_{0}$ or $\bm{\epsilon}$). The UNet employs five down/up-sampling
resolutions with $768\times(1,2,4,6,8)$ respective channels, 3 ResNet blocks
per resolution, and spatial self-attention at the 3 smallest resolutions,
i.e., 8, 16, and 32.

After training the models, we run DDSS using just the GGDM model family for
simplicity (i.e., we don’t use the +PRED and +TIME we experiment with in the
paper) at 5, 10 and 20 evenly-spaced inference steps. DDSS training occurs for
50K steps, using the Adam optimizer with learning rate of 0.0005 and batch
size 512, optimizing the linear kernel for the KID loss. We compare against
the usual DDPM and DDIM($\eta=0$) baselines at the same inference budgets and
include the FID scores with all 1000 steps for reference. Results and samples
are included below.

Sampler $\setminus$ $K$ | 5 | 10 | 20 | 1000
---|---|---|---|---
LSUN Bedroom |  |  |  |
DDPM | 95.38 | 44.84 | 16.88 | 2.547
DDIM($\eta=0$) | 168.7 | 56.33 | 9.527 | –
DDSS (GGDM) | 29.15 | 11.01 | 4.817 | –
LSUN Church |  |  |  |
DDPM | 96.67 | 51.05 | 16.53 | 2.718
DDIM($\eta=0$) | 133.1 | 54.39 | 14.96 | –
DDSS (GGDM) | 30.24 | 11.59 | 6.736 | –
| DDPM | DDIM | Ours
---|---|---|---
5 steps |  |  |
10 steps |  |  |
20 steps |  |  |
Reference |  |  |
Figure A.2: Non-cherrypicked samples for a DDPM trained on LSUN bedroom 128x128, comparing DDPM and DDIM($\eta=0$) to our approach. All samples were generated with the same random seeds and a linear stride. For reference, we include DDPM samples using all 1,000 steps (bottom left) and real images (bottom middle). | DDPM | DDIM | Ours
---|---|---|---
5 steps |  |  |
10 steps |  |  |
20 steps |  |  |
Reference |  |  |
Figure A.3: Non-cherrypicked samples for a DDPM trained on LSUN church
128x128, comparing DDPM and DDIM($\eta=0$) to our approach. All samples were
generated with the same random seeds and a linear stride. For reference, we
include DDPM samples using all 1,000 steps (bottom left) and real images
(bottom middle).


#  Optical frequency tweezers

Rikizo Ikuta  Graduate School of Engineering Science, Osaka University,
Toyonaka, Osaka 560-8531, Japan  Center for Quantum Information and Quantum
Biology, Osaka University, Osaka 560-8531, Japan  Masayo Yokota  Graduate
School of Engineering Science, Osaka University, Toyonaka, Osaka 560-8531,
Japan  Toshiki Kobayashi  Graduate School of Engineering Science, Osaka
University, Toyonaka, Osaka 560-8531, Japan  Center for Quantum Information
and Quantum Biology, Osaka University, Osaka 560-8531, Japan  Nobuyuki Imoto
Center for Quantum Information and Quantum Biology, Osaka University, Osaka
560-8531, Japan  Takashi Yamamoto  Graduate School of Engineering Science,
Osaka University, Toyonaka, Osaka 560-8531, Japan  Center for Quantum
Information and Quantum Biology, Osaka University, Osaka 560-8531, Japan

###### Abstract

We show a concept of optical frequency tweezers for tweezing light in the
optical frequency domain with a high resolution, which is the frequency
version of the optical tweezers for the spatial manipulation of microscopic
objects. We report the proof-of-principle experiment via frequency conversion
inside a cavity only for the converted light. Owing to the atypical
configuration, the experimental results successfully achieve the tweezing
operation in the frequency domain, which picks a light at a target frequency
from the frequency-multiplexed input light and converts it to a different
frequency, without touching any other light sitting in different frequency
positions and shaking the frequency by the pump light.

##  I introduction

Optical tweezers Ashkin (1970); Ashkin et al. (1986) that pick up and move
microscopic objects using a highly focused laser are an indispensable
technology that support modern science and engineering such as the optical
trap of atoms Chu et al. (1986), manipulation of atoms in an optical lattice
Browaeys and Lahaye (2020), and biological science applications Ashkin and
Dziedzic (1987). Similar to the optical tweezers for spatially deployed
objects, to manipulate photonic lattices in frequency synthetic dimension Hu
et al. (2020); Buddhiraju et al. (2021), there is an increasing demand for
tweezing light in the optical frequency domain with a high resolution that
precisely picks up only light in a target frequency mode from broadened or
densely frequency multiplexed light and converts it to any frequency mode. In
addition, it is also required in the field of quantum information processing
for realizing photonic quantum computation based on quantum frequency combs
Kues et al. (2017); Lukens and Lougovski (2017); Niu et al. (2018); Kues et
al. (2019); Reimer et al. (2019), and quantum internet Kimble (2008) to
interconnect various quantum systems including frequency/wavelength-division
multiplexing technologies Sinclair et al. (2014); Wengerowsky et al. (2018);
Lingaraju et al. (2021). However, such optical frequency tweezers have never
been realized.

In this study, we propose an optical system for the optical frequency
tweezers, as illustrated in Fig. 1 (a), and perform its proof-of-principle
demonstration. To this end, we adopt a frequency conversion based on a
nonlinear optical interaction inside a cavity that only confines light to be
generated by the conversion process. Fig. 1 (b) presents the system design
based on a second-order nonlinearity ($\chi^{(2)}$). The input signal
interacts with the cavity mode via the frequency conversion with a strong pump
light, and then the converted light is extracted because of the cavity decay.
The system is considered as an all optical implementation of a multiplexed
$\Lambda$-type three-level system, as illustrated in Fig. 1 (c). The cavity
resonant frequency modes adopted as the discretized excited levels modulate
the bandwidth of the nonlinear optical interaction while possessing broadband
bandwidths for the signal and pump modes. Therefore, it is feasible for the
pump light to select a target frequency mode to be converted from a frequency-
multiplexed input signal. It represents the minimum of required properties for
the optical frequency tweezers. In addition, the tweezers should precisely
pick up the target without shaking and should not disturb the unconverted
signals around the target. The unconventional cavity configuration addresses
these requirements for the tweezers, which have never been reported with the
use of multiply resonant cavities such as microresonator systems. There is no
thermal fluctuation owing to no confinement of the pump light, and there is no
disturbance and cavity loss of the unconverted light due to no confinement of
the input signals.

##  II basic concept

Figure 1:  (a) Concept of optical frequency tweezers. The tweezers pick up
only light at the target frequency and deploy it to the desired frequency
mode. (b) Implementation of the optical frequency tweezers based on
$\chi^{(2)}$-based frequency conversion inside a cavity only for the converted
light. The mode coupling relationship is also illustrated. $a_{s}$, $a_{r}$,
and $a_{c}$ correspond to the signal, converted and cavity modes. $\omega_{p}$
denotes the pump frequency. (c) Diagram of the frequency tweezing system when
only one of the converted frequency of the input frequencies is resonant on
the cavity. Only the middle of the input signal can be converted.

The optical frequency tweezers introduced in this study are realized by the
photonic $\Lambda$-type three-level system based on the $\chi^{(2)}$
nonlinearity inside a single-sided cavity around the converted frequencies
(Fig. 1 (b)). We focus on the frequency conversion process around the resonant
frequency of the cavity described by the single mode $a_{c}$. The cavity mode
is coupled to two modes outside the cavity, namely $a_{r}$ and $a_{s}$, of
which the former corresponds to the cavity decay through its single side and
the latter corresponds to the input signal field obtained via the
$\chi^{(2)}$-based frequency conversion. The interaction Hamiltonians of these
coupling systems are described by
$H_{C}=i\sqrt{\gamma_{r}}a_{c}a_{r}^{\dagger}+H.c.$ and $H_{NL}=i\xi
a_{c}a_{s}^{\dagger}+H.c.$, where $H.c.$ denotes the Hermitian conjugate,
$\sqrt{\gamma_{r}}$ is a coupling constant between the internal and external
modes of the cavity, and $\xi=|\xi|e^{i\phi}$ represents an effective coupling
constant of the frequency conversion that is proportional to the complex
amplitude of the sufficiently strong pump light with its phase $\phi$. From
the system Hamiltonian $H_{C}+H_{NL}$, the time evolution of the cavity mode
$a_{c}=a_{c}(t)$ in a frame rotating at the resonant frequency $\omega_{c}$ is
described by Collett and Gardiner (1984)

$\displaystyle\frac{\mathrm{d}a_{c}}{\mathrm{d}t}=i\Delta_{c}a_{c}-\frac{\gamma_{\rm
all}+|\xi|^{2}}{2}a_{c}+\sqrt{\gamma_{r}}a_{r,IN}+\xi^{*}a_{s,IN},$ (1)

where $\gamma_{\rm all}:=\gamma_{r}+\gamma_{\rm int}$ is the total loss of the
cavity including the internal loss $\gamma_{\rm int}$.
$\Delta_{c}:=\omega-\omega_{c}$ denotes the detuning of the light in the
cavity, where $\omega$ represents the frequency of mode $a_{c}$, and is
determined by the frequencies $\omega_{s}$ and
$\omega_{p}(=\omega_{s}-\omega)$ of mode $a_{s}$ and the pump light.
$a_{s,IN}$ is the input signal field to be converted. $a_{r,IN}$ is the
external drive field, and is assumed to be the vacuum for the standard usage
of this device, which we will describe in the next paragraph. This drive field
can be used to control the frequency conversion process of $a_{s,IN}$, which
will be explained in detail based on Eq. (4). The system described by Eq. (1)
can be regarded as the $\Lambda$-type three-level quantum system, which
comprises the excited level $a_{c}$ and two ground levels, $a_{s}$ and
$a_{r}$, as illustrated in Fig. 1 (c). We note that this diagram does not
present the energy relationship among the relevant modes as conventional ones
in $\Lambda$-type atomic systems. The excited level $a_{c}$ in the diagram
exhibits the energy lower than that of the ground level $a_{s}$, while $a_{s}$
is excited to $a_{c}$ by the coupling based on the frequency conversion. In
addition, although $a_{c}$ and another ground level $a_{r}$ exhibit the same
energy, $a_{c}$ decays to $a_{r}$ due to the spatial coupling at the end
mirror.

The frequency conversion from $a_{s}$ to $a_{r}$ is the case of $a_{r,IN}=0$.
From Eq. (1) and the input-output relationships Walls and Milburn (2007)
$a_{s,IN}-a_{s,OUT}=\xi a_{c}$ and
$a_{r,IN}+a_{r,OUT}=\sqrt{\gamma_{r}}a_{c}$, we obtain

$\displaystyle t_{ss}:=\frac{a_{s,OUT}}{a_{s,IN}}$ $\displaystyle=$
$\displaystyle\frac{\frac{1}{2}(1-\widetilde{C})-i\widetilde{\Delta}_{c}}{\frac{1}{2}(1+\widetilde{C})-i\widetilde{\Delta}_{c}},$
(2) $\displaystyle r_{rs}:=\frac{a_{r,OUT}}{a_{s,IN}}$ $\displaystyle=$
$\displaystyle\sqrt{\widetilde{\gamma}_{r}}\frac{e^{-i\phi}\sqrt{\widetilde{C}}}{\frac{1}{2}(1+\widetilde{C})-i\widetilde{\Delta}_{c}},$
(3)

where $\widetilde{\gamma}_{r}:=\gamma_{r}/\gamma_{\rm all}$ and
$\widetilde{\Delta}_{c}:=\Delta_{c}/\gamma_{\rm all}$.
$\widetilde{C}:=|\xi|^{2}/\gamma_{\rm all}$ corresponds to the cooperativity
parameter. $t_{ss}$ and $r_{rs}$ are complex amplitudes for the staying
probability $T(=|t_{ss}|^{2})$ at mode $a_{s}$ and the transition probability
$R(=|r_{rs}|^{2})$ to mode $a_{r}$. These equations hold at each frequency
interval corresponding to the free spectral range (FSR) of the cavity with the
same level of the nonlinearity strength, regardless of the dispersion of the
medium. This is because the singly resonant cavity structure eliminates the
restriction in conventional doubly or triply resonant cavities, such that the
relevant three frequencies satisfying the energy conservation must be
simultaneously resonant on the cavities Ikuta et al. (2019).

From the analysis, we can observe remarkable properties in the frequency
conversion based on the $\Lambda$-type structure as follows: (i) The device
can pick up any single frequency mode from multiplexed input frequency modes,
and then convert to any frequency mode by designing the FSR and pump
frequency, such that only the converted frequency corresponding to the target
frequency to be chosen for conversion is resonant, as illustrated in Fig. 1
(c). (ii) The remaining signal light which is far from the cavity resonances
passes through the device without any disturbance and loss, which is observed
from $T=1$ for sufficiently large $\widetilde{\Delta}_{c}$ in Eq. (2). (iii)
No frequency shake occurs during the pick up process. This is due to the
absence of the thermal heating effect triggered by the pump light, which is
not confined in the cavity, as well as the stable operation without a severe
frequency locking. Properties (i) – (iii) are exactly the requirements for the
“tweezers.” Hence, we call the device optical frequency tweezers.

In addition to the aforementioned basic properties, our device also has
advantages as the cavity-based frequency converter as follows: (iv) It has no
optical impedance matching problem for the input signal, regardless of the
cavity enhancement characterized by $\widetilde{C}$. (v) The spectral shape of
the converted light is not split or distorted, and obeys the Lorentzian, which
differs from the mode splitting phenomena analogous to the Autler-Townes
splitting that appears in coupled resonator systems Peng et al. (2014); Guo et
al. (2016). (vi) The maximum conversion efficiency $\widetilde{\gamma}_{r}$
would be higher than that determined by the product of cavity losses for both
signal and converted modes in the frequency conversion based on multiply
resonant cavities Li et al. (2016); Guo et al. (2016). (vii) All optical
adjustments of $\widetilde{\Delta}_{c}$ are possible by the pump frequency
tuning.

Figure 2:  (a) Setup for picking a frequency mode around $s_{i}$ and deploying
it to another mode around $r_{i}$ using a pump light at $\omega_{p,i}$. The
$\chi^{(2)}$ medium for the frequency conversion is the PPLN/WG. (b) Cascade
use of the above setup for any-to-any tweezing, as presented in Fig. 1 (a).
(c) Experimental setup for the proof-of-principle experiment. The signal light
A at $780\text{\,}\mathrm{n}\mathrm{m}$ is used to characterize the
performance of the frequency conversion. The demonstration of the tweezing
operation is performed by using both signal lights A and B with frequency
detuning $\nu_{s}$. PD; photodetector, RFSA; radio frequency spectral
analyzer, DSO; digital sampling oscilloscope, OSA; optical spectral analyzer.

##  III experiments

The experimental design for the optical frequency tweezers is presented in
Fig. 2 (a). The scenario considered here involves tweezing one tooth from a
frequency comb or quantum frequency comb Kues et al. (2019). When all the
frequency intervals and the cavity FSR satisfy a relationship in which only
the converted frequency corresponding to the target tooth is resonant on the
cavity, the tweezers can pick up and convert the frequency of the tooth and
leave the rest of the teeth on the unconverted modes without disturbance.
Consequently, the cascade use of the tweezers presented in Fig. 2 (b) can
realize any-to-any tweezing operation, as illustrated in Fig. 1 (a).

The experimental setup for the proof-of-principle demonstration is shown in
Fig. 2 (c). The continuous wave signal light around
$780\text{\,}\mathrm{n}\mathrm{m}$ (corresponding to $\omega_{s}$) with $\sim$
$1\text{\,}\mathrm{m}\mathrm{W}$ and the pump light for the tweezers at
$1540\text{\,}\mathrm{n}\mathrm{m}$ ($\omega_{p}$) are combined at a dichroic
mirror (DM). Subsequently, they are focused on the periodically poled lithium
niobate waveguide resonator (PPLN/WR), which only confines the converted light
around $1581\text{\,}\mathrm{n}\mathrm{m}$ ($\omega_{r}$).

PPLN/WR satisfies the type-0 quasi-phase-matching condition for the vertically
polarized light. The length of the waveguide is
$20\text{\,}\mathrm{m}\mathrm{m}$. The bandwidth of frequency conversion for
the input light at $780\text{\,}\mathrm{n}\mathrm{m}$ is over
$150\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$ Ikuta et al. (2021). The FSR of
the resonator is $\nu_{\rm FSR}=$3.5\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$$
Ikuta et al. (2018). In addition, the dielectric multilayers on the end faces
of the waveguide achieve high reflectance of $\sim$ $94\text{\,}\%$ for light
around $1581\text{\,}\mathrm{n}\mathrm{m}$. We determined the full width at
the half maximum (FWHM) of the cavity without the internal loss by the
reflectance as $\gamma_{0}:=$71\text{\,}\mathrm{M}\mathrm{H}\mathrm{z}$$. The
reflectance for $780\text{\,}\mathrm{n}\mathrm{m}$ is a few percent and thus
the anti-reflection coating is achieved. The transmittance of
$780\text{\,}\mathrm{n}\mathrm{m}$ light to the PPLN/WR is $90\text{\,}\%$.
Although the reflectance for $1540\text{\,}\mathrm{n}\mathrm{m}$ is slightly
larger at approximately $30\text{\,}\%$, it forms a very lossy cavity.
Considering the fineness is smaller than $\pi$ Regener and Sohler (1988);
Stefszky et al. (2018), there should not be any cavity enhancement effect for
the frequency conversion.

After the PPLN/WR, the signal light and converted light are extracted using
another DM and a volume Bragg grating (VBG) with its center wavelength and
bandwidth of $1581\text{\,}\mathrm{n}\mathrm{m}$ and
$1\text{\,}\mathrm{n}\mathrm{m}$, respectively. They are coupled to single
mode fibers followed by proper experimental apparatuses.

Figure 3:  (a) Observed four transmission spectra (p1 – p4) related to the
longitudinal cavity modes separated by the FSR of the 1581-nm (upper) and
780-nm light (lower) for pump power $25\text{\,}\mathrm{m}\mathrm{W}$,
$200\text{\,}\mathrm{m}\mathrm{W}$, and $500\text{\,}\mathrm{m}\mathrm{W}$
from the top. (b) Pump power dependency of observed bandwidths for 780-nm
(red) and 1581-nm light (blue). (c) Pump power dependency of the transition
probability $R(P)$ (blue) and staying probability $T(P)$ (red). (d) Observed
spectra around $1581\text{\,}\mathrm{n}\mathrm{m}$ for
$\nu_{s}=$7.0\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$=2\,\nu_{\rm FSR}$ (light
blue) and $\nu_{s}=$9.1\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$=2.6\,\nu_{\rm
FSR}$ (dark blue). (e) Same figure as Fig. 3 (d) with the linear scale for the
vertical axis.

First, we characterize the performance of the frequency conversion from
$780\text{\,}\mathrm{n}\mathrm{m}$ to $1581\text{\,}\mathrm{n}\mathrm{m}$. We
adopt the signal light emerging from light source A. By scanning the pump
frequency in the range of approximately
$15\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$, we observed four peaks (labelled
p1 – p4) in the transmission spectra for $780\text{\,}\mathrm{n}\mathrm{m}$
and $1581\text{\,}\mathrm{n}\mathrm{m}$, which is related to the longitudinal
cavity modes separated by the FSR. The spectra are presented in Fig. 3 (a). As
an example, we focus on the rightmost peaks (p4) of the figure. The bandwidths
of the frequency conversion around the peaks are presented in Fig. 3 (b). As
expected in Eqs. (2) and (3), the bandwidths obtained by the 780-nm and
1581-nm light are almost the same and proportional to the pump power $P$
measured before the PPLN/WR. The best fit to the data with a function $\alpha
P+\gamma_{\rm all}^{{\rm(ex)}}$ gives
$\alpha=$0.68\text{\,}\mathrm{M}\mathrm{H}\mathrm{z}\mathrm{/}\mathrm{m}\mathrm{W}$$
and $\gamma_{\rm
all}^{{\rm(ex)}}=$106\text{\,}\mathrm{M}\mathrm{H}\mathrm{z}$$. $\gamma_{\rm
all}^{{\rm(ex)}}$ reflects the total loss of the resonator corresponding to
$\gamma_{\rm all}$ introduced in Eq. (1), which includes the internal loss of
the resonator, as well as the loss from the end mirrors.

For the conversion efficiency, its maximum is limited to
$\widetilde{\gamma_{r}}=\gamma_{0}/\gamma_{\rm
all}^{{\rm(ex)}}=$67\text{\,}\%$$ from Eq. (3) for the internal loss of the
resonator estimated to be $\gamma_{\rm
int}=$35\text{\,}\mathrm{M}\mathrm{H}\mathrm{z}$$. Besides, a mode mismatch of
the signal and the converted modes propagating in the waveguide should degrade
the conversion efficiency. To estimate the amount $M$ of the mode match, we
plot the bottom peaks of the observed spectra for the780-nm light normalized
by the value at $P=$0\text{\,}\mathrm{m}\mathrm{W}$$ in Fig. 3 (c). The best
fit to the data with a function $MT(P)+(1-M)$ gives $M=0.94$, where
$T(P)=|t|^{2}$ is given by Eq. (2) with $\widetilde{\Delta}_{c}=0$.
Accordingly, the maximum of the achievable internal conversion efficiency is
estimated as $\eta_{\rm max}:=M\widetilde{\gamma_{r}}=0.63$. By normalizing
the observed peaks for the 1581-nm light with the maximum, we can observe the
internal conversion efficiency, as presented in Fig. 3 (c). From the best fit
to the data with a function $\eta_{\rm max}R(P)$ given by Eq. (3) with
$\widetilde{\Delta}_{c}=0$, the unit power cooperativity is estimated to be
$\beta=$0.0062\text{\,}\mathrm{m}\mathrm{W}^{-1}$$, which agrees well with
$\alpha/\gamma_{\rm
all}^{{\rm(ex)}}=$0.0064\text{\,}\mathrm{m}\mathrm{W}^{-1}$$, as estimated
from the experimental result about the bandwidth. Although the PPLN/WR used in
this experiment is the Fabry-Pérot cavity and the converted light emerges from
its two end faces, an asymmetric mirror coating could limit the output to one
side only, or these two modes can be coherently combined by adopting an
interferometer such as the Sagnac loop Yamazaki et al. (2021).

We notice that the pump power for the maximum conversion efficiency
($\widetilde{C}=1$) is $\beta^{-1}\sim$160\text{\,}\mathrm{m}\mathrm{W}$$
while the maximum in the case without cavities was
$700\text{\,}\mathrm{m}\mathrm{W}$ Ikuta et al. (2011). This clearly indicates
the cavity enhancement effect without the pump light confinement. In addition,
the good agreement of the theoretical curves and the experimental results for
all pump powers implies the absence of unwanted nonlinear optical interaction
that facilitates noise light generation around the target wavelengths to be
converted and the pump power consumption observed in the $\chi^{(2)}$-based
frequency conversion with the pump light confinement Ikuta et al. (2021).

We performed the same analysis as described above for the other peaks in Fig.
3 (a). The estimated values are listed in Table 1. It is evident that almost
the same values are obtained for every parameter. This implies that this
device achieved equal performance on the frequency conversion around every
resonant frequency.

peak number | p1 | p2 | p3 | p4
---|---|---|---|---
$\alpha$ (MHz/mW) | $0.63\pm 0.03$ | $0.74\pm 0.02$ | $0.63\pm 0.02$ | $0.68\pm 0.02$
$\gamma_{\rm all}^{{\rm(ex)}}$ (MHz) | $117\pm 4$ | $105\pm 3$ | $106\pm 2$ | $106\pm 2$
$M$ (%) | $93\pm 0.2$ | $94\pm 0.1$ | $94\pm 0.2$ | $94\pm 0.1$
$10^{3}\beta$ (/mW) | $5.6\pm 0.1$ | $5.9\pm 0.1$ | $6.1\pm 0.1$ | $6.2\pm 0.1$
Table 1:  Estimated values on the four peaks.

Next, we perform the tweezing operation. To suppress the shaking during the
picking and maximize the conversion efficiency, we adopted part of the
converted light to lock the pump frequency. The observed power of signal A
after PPLN/WR was reduced from $0.6\text{\,}\mathrm{m}\mathrm{W}$ to
$0.04\text{\,}\mathrm{m}\mathrm{W}$ via frequency conversion with the
frequency-locked pump light. This decrease ratio of 0.93 agrees well with the
observed value $M$ in the previous section. Under this setup, we input an
additional signal light denoted by B with a frequency difference $\nu_{s}$
from the frequency of signal A. We show the observed spectra of the converted
light around $1581\text{\,}\mathrm{n}\mathrm{m}$ in Figs. 3 (d) and (e) with
logarithmic and linear scales for the vertical axes, respectively. It can be
observed that for
$\nu_{s}=$7.0\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$=2\,\nu_{\rm FSR}$, both
signals were simultaneously converted to $1581\text{\,}\mathrm{n}\mathrm{m}$.
In contrast, for
$\nu_{s}=$9.1\text{\,}\mathrm{G}\mathrm{H}\mathrm{z}$=2.6\,\nu_{\rm FSR}$, the
converted light by signal B was significantly smaller (a suppression of
$16\text{\,}\mathrm{d}\mathrm{B}$) because the $\widetilde{\Delta}_{c}$ value
for signal B is sufficiently large. Based on this result, we conclude that the
proof-of-principle experiment for the optical frequency tweezers was
successfully achieved.

Figure 4:  (a) Diagram of the three-level system with an external input light
from mode $a_{r}$. (b) Theoretical simulation of the interference effect by
using Eq. (4) with $\phi=0$ and $\widetilde{C}\widetilde{\gamma_{r}}=0.25$.
The control light is assumed to be
$a_{r,IN}/a_{s,IN}=e^{i\theta}0.5\widetilde{\Gamma}/(\widetilde{\Gamma}+i\widetilde{\delta}_{c})$
with $\widetilde{\Gamma}=0.05$. The top figure presents three curves for
$\widetilde{\delta}_{c}=\widetilde{\Delta}_{c}-0.75$ (green),
$\widetilde{\Delta}_{c}$ (red), and $\widetilde{\Delta}_{c}+0.75$ (blue) in
the case of $\theta=0$. The middle figure shows three curves for
$\widetilde{\delta}_{c}=\widetilde{\Delta}_{c}-0.75$ (blue),
$\widetilde{\Delta}_{c}$ (yellow), and $\widetilde{\Delta}_{c}+0.75$ (green)
in the case of $\theta=\pi$. The bottom figure shows the curves for the three
detunings $\widetilde{\delta}_{c}=\widetilde{\Delta}_{c}\pm
0.75,\widetilde{\Delta}_{c}$, each with the random phase of $\theta=k\pi/20$
($k=0,\ldots,40$). (c) The observed transmission spectra at
$780\text{\,}\mathrm{n}\mathrm{m}$ with the control light around
$1581\text{\,}\mathrm{n}\mathrm{m}$ with four different detunings. The figures
on the right are enlarged figures on the left.

The tweezing operation presented above is based on the frequency conversion of
the 780-nm light. In addition to the function, the $\Lambda$-type system
enables us to adopt an external field at $1581\text{\,}\mathrm{n}\mathrm{m}$,
to control the tweezing operation (refer to Fig. 4 (a)). The external field is
described by $a_{r,IN}$ in Eq. (1), and the complex amplitude of the staying
probability becomes

$\displaystyle\frac{a_{s,OUT}}{a_{s,IN}}=\frac{e^{i\phi}\sqrt{\widetilde{C}\widetilde{\gamma}_{r}}a_{r,IN}/a_{s,IN}+\frac{1}{2}(1-\widetilde{C})-i\widetilde{\Delta}_{c}}{\frac{1}{2}(1+\widetilde{C})-i\widetilde{\Delta}_{c}}.$
(4)

The equation above demonstrates that the output of the signal includes an
interference effect between modes $a_{r}$ and $a_{s}$. Unlike the coupled-
resonator-based $\Lambda$-type optical systems Peng et al. (2014); Guo et al.
(2016), it is possible to freely tune the relative phase of the free fields.
Consequently, both destructive and constructive interferences corresponding to
the electromagnetically-induced transparency (EIT) and the
electromagnetically-induced absorption (EIA) based on the Fano interference in
atomic systems are achieved without altering the cooperativity $\widetilde{C}$
Naweed et al. (2005). The theoretical simulation of the optical
implementations of EIT and the EIA under an assumption that the frequency of
the control light is locked to that of the pump light is presented in Fig. 4
(b). The results indicate the possibilities that EIT and the EIA can,
respectively, obstruct and facilitate the tweezing operation close to the
resonant frequencies. This property enables the device to provide additional
functions. For example, the switching operation to all frequency modes by the
pump light is performed at once for the frequency multiplexed input with the
frequency interval corresponding to the FSR. However, the control light
enables us to switch one of the frequency modes. Another usage of this
property involves performing the tweezing operation of the signal sensitive to
the phase of the control beam as an additional reference light.

The observed spectra of the 780-nm light for various wavelengths of the
control light are presented in Fig. 4(c). In response to the wavelength of the
control light changes, the wavelength for observing the interference changed.
As is expected, both the destructive and constructive interferences related to
the EIT and EIA are observed at the control frequencies near the resonant
points, while the interference signatures reflect the random phase structure
similar to the bottom of Fig. 4 (b) because the signal and control lasers were
not phase locked. This result implies that the tweezing operation can be
switched by using a more stable laser for the control field.

##  IV conclusion

In this study, we presented the optical frequency tweezers with experimental
implementation based on frequency conversion with the atypical cavity
structure. It facilitated the precise frequency tweezing operation without
shaking the frequency or disturbing the light, other than the target to be
chosen for tweezing. Using the control light as the analogies of EIT and EIA
in atomic systems, we demonstrated the possibilities of suppressing and
enhancing the tweezing operation. Such optical frequency tweezers will be
indispensable as the channel exchange for a massive any-to-any channel switch
among frequency-multiplexed channels, which has the same role as a
conventional channel exchange for an any-to-any switch among spatially
assigned channels. In addition to the above tweezing operation, we verified
the advantages of our frequency converter. Importantly, the cavity enhancement
without severe impedance matching and the higher maximum conversion efficiency
will be beneficial in the frequency conversion of a single photon with quantum
information that is never amplified.

##  V acknowledgments

R.I., N.I., and T.Y. acknowledge members of Quantum Internet Task Force (QITF)
for comprehensive and interdisciplinary discussions of the quantum internet.
This work was supported by Moonshot R & D, JST JPMJMS2066; CREST, JST
JPMJCR1671; MEXT/JSPS KAKENHI Grants No. JP21H04445, JP20H01839, and Asahi
Glass Foundation.

## References

  * Ashkin (1970) A. Ashkin, Acceleration and trapping of particles by radiation pressure, Phys. Rev. Lett. 24, 156 (1970).
  * Ashkin et al. (1986) A. Ashkin, J. M. Dziedzic, J. E. Bjorkholm, and S. Chu, Observation of a single-beam gradient force optical trap for dielectric particles, Optics Letters 11, 288 (1986).
  * Chu et al. (1986) S. Chu, J. E. Bjorkholm, A. Ashkin, and A. Cable, Experimental observation of optically trapped atoms, Phys. Rev. Lett. 57, 314 (1986).
  * Browaeys and Lahaye (2020) A. Browaeys and T. Lahaye, Many-body physics with individually controlled rydberg atoms, Nature Physics 16, 132 (2020).
  * Ashkin and Dziedzic (1987) A. Ashkin and J. M. Dziedzic, Optical trapping and manipulation of viruses and bacteria, Science 235, 1517 (1987).
  * Hu et al. (2020) Y. Hu, C. Reimer, A. Shams-Ansari, M. Zhang, and M. Loncar, Realization of high-dimensional frequency crystals in electro-optic microcombs, Optica 7, 1189 (2020).
  * Buddhiraju et al. (2021) S. Buddhiraju, A. Dutt, M. Minkov, I. A. Williamson, and S. Fan, Arbitrary linear transformations for photons in the frequency synthetic dimension, Nature Communications 12, 1 (2021).
  * Kues et al. (2017) M. Kues, C. Reimer, P. Roztocki, L. R. Cortés, S. Sciara, B. Wetzel, Y. Zhang, A. Cino, S. T. Chu, B. E. Little, D. J. Moss, L. Caspani, J. Azaña, and R. Morandotti, On-chip generation of high-dimensional entangled quantum states and their coherent control, Nature 546, 622 (2017).
  * Lukens and Lougovski (2017) J. M. Lukens and P. Lougovski, Frequency-encoded photonic qubits for scalable quantum information processing, Optica 4, 8 (2017).
  * Niu et al. (2018) M. Y. Niu, I. L. Chuang, and J. H. Shapiro, Qudit-basis universal quantum computation using ${\chi}^{(2)}$ interactions, Phys. Rev. Lett. 120, 160502 (2018).
  * Kues et al. (2019) M. Kues, C. Reimer, J. M. Lukens, W. J. Munro, A. M. Weiner, D. J. Moss, and R. Morandotti, Quantum optical microcombs, Nature Photonics 13, 170 (2019).
  * Reimer et al. (2019) C. Reimer, S. Sciara, P. Roztocki, M. Islam, L. R. Cortés, Y. Zhang, B. Fischer, S. Loranger, R. Kashyap, A. Cino, S. T. Chu, B. E. Little, D. J. Moss, L. Caspani, W. J. Munro, J. Azaña, M. Kues, and R. Morandotti, High-dimensional one-way quantum processing implemented on $d$-level cluster states, Nature Physics 15, 148 (2019).
  * Kimble (2008) H. J. Kimble, The quantum internet, Nature 453, 1023 (2008).
  * Sinclair et al. (2014) N. Sinclair, E. Saglamyurek, H. Mallahzadeh, J. A. Slater, M. George, R. Ricken, M. P. Hedges, D. Oblak, C. Simon, W. Sohler, and W. Tittel, Spectral multiplexing for scalable quantum photonics using an atomic frequency comb quantum memory and feed-forward control, Phys. Rev. Lett. 113, 053603 (2014).
  * Wengerowsky et al. (2018) S. Wengerowsky, S. K. Joshi, F. Steinlechner, H. Hübel, and R. Ursin, An entanglement-based wavelength-multiplexed quantum communication network, Nature 564, 225 (2018).
  * Lingaraju et al. (2021) N. B. Lingaraju, H.-H. Lu, S. Seshadri, D. E. Leaird, A. M. Weiner, and J. M. Lukens, Adaptive bandwidth management for entanglement distribution in quantum networks, Optica 8, 329 (2021).
  * Collett and Gardiner (1984) M. Collett and C. Gardiner, Squeezing of intracavity and traveling-wave light fields produced in parametric amplification, Phys. Rev. A 30, 1386 (1984).
  * Walls and Milburn (2007) D. F. Walls and G. J. Milburn, Quantum Optics (Springer Science & Business Media, 2007).
  * Ikuta et al. (2019) R. Ikuta, R. Tani, M. Ishizaki, S. Miki, M. Yabuno, H. Terai, N. Imoto, and T. Yamamoto, Frequency-multiplexed photon pairs over 1000 modes from a quadratic nonlinear optical waveguide resonator with a singly resonant configuration, Phys. Rev. Lett. 123, 193603 (2019).
  * Peng et al. (2014) B. Peng, Ş. K. Özdemir, W. Chen, F. Nori, and L. Yang, What is and what is not electromagnetically induced transparency in whispering-gallery microcavities, Nature Communications 5, 1 (2014).
  * Guo et al. (2016) X. Guo, C.-L. Zou, H. Jung, and H. X. Tang, On-chip strong coupling and efficient frequency conversion between telecom and visible optical modes, Phys. Rev. Lett. 117, 123902 (2016).
  * Li et al. (2016) Q. Li, M. Davanço, and K. Srinivasan, Efficient and low-noise single-photon-level frequency conversion interfaces using silicon nanophotonics, Nature Photonics 10, 406 (2016).
  * Ikuta et al. (2021) R. Ikuta, T. Kobayashi, T. Yamazaki, N. Imoto, and T. Yamamoto, Cavity-enhanced broadband photonic Rabi oscillation, Phys. Rev. A 103, 033709 (2021).
  * Ikuta et al. (2018) R. Ikuta, M. Asano, R. Tani, T. Yamamoto, and N. Imoto, Frequency comb generation in a quadratic nonlinear waveguide resonator, Optics Express 26, 15551 (2018).
  * Regener and Sohler (1988) R. Regener and W. Sohler, Efficient second-harmonic generation in $\rm{Ti:LiNbO_{3}}$ channel waveguide resonators, JOSA B 5, 267 (1988).
  * Stefszky et al. (2018) M. Stefszky, R. Ricken, C. Eigner, V. Quiring, H. Herrmann, and C. Silberhorn, High-power waveguide resonator second harmonic device with external conversion efficiency up to 75%, Journal of Optics 20, 065501 (2018).
  * Yamazaki et al. (2021) T. Yamazaki, R. Ikuta, T. Kobayashi, S. Miki, F. China, H. Terai, N. Imoto, and T. Yamamoto, Massive-mode polarization entangled biphoton frequency comb, arXiv preprint arXiv:2101.04410 (2021).
  * Ikuta et al. (2011) R. Ikuta, Y. Kusaka, T. Kitano, H. Kato, T. Yamamoto, M. Koashi, and N. Imoto, Wide-band quantum interface for visible-to-telecommunication wavelength conversion, Nature Communications 2, 537 (2011).
  * Naweed et al. (2005) A. Naweed, G. Farca, S. I. Shopova, and A. Rosenberger, Induced transparency and absorption in coupled whispering-gallery microresonators, Phys. Rev. A 71, 043804 (2005).