# Diffusion Probabilistic Model Made Slim

Xingyi Yang1 Daquan Zhou2 Jiashi Feng2 Xinchao Wang1
National University of Singapore1 ByteDance Inc.2
<EMAIL_ADDRESS>{daquanzhou<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Despite the recent visually-pleasing results achieved, the massive
computational cost has been a long-standing flaw for diffusion probabilistic
models (DPMs), which, in turn, greatly limits their applications on resource-
limited platforms. Prior methods towards efficient DPM, however, have largely
focused on accelerating the testing yet overlooked their huge complexity and
sizes. In this paper, we make a dedicated attempt to lighten DPM while
striving to preserve its favourable performance. We start by training a small-
sized latent diffusion model (LDM) from scratch, but observe a significant
fidelity drop in the synthetic images. Through a thorough assessment, we find
that DPM is intrinsically biased against high-frequency generation, and learns
to recover different frequency components at different time-steps. These
properties make compact networks unable to represent frequency dynamics with
accurate high-frequency estimation. Towards this end, we introduce a
customized design for slim DPM, which we term as Spectral Diffusion (SD), for
light-weight image synthesis. SD incorporates wavelet gating in its
architecture to enable frequency dynamic feature extraction at every reverse
steps, and conducts spectrum-aware distillation to promote high-frequency
recovery by inverse weighting the objective based on spectrum magni-
tudes.Experimental results demonstrate that, SD achieves 8-18$\times$
computational complexity reduction as compared to the latent diffusion models
on a series of conditional and unconditional image generation tasks while
retaining competitive image fidelity.

##  1 Introduction

Diffusion Probabilistic Models (DPMs) [16, 55, 53] have recently emerged as a
power tool for generative modeling, and have demonstrated impressive results
in image synthesis [42, 8, 43], video generation [19, 15, 68] and 3D editing
[39]. Nevertheless, the gratifying results come with a price: DPMs suffer from
massive model sizes. In fact, state-of-the-art DPMs requires billions of
parameters, with hundreds or even thousands of inference steps per image. For
example, _DALL $\cdot$ E 2_ [42], which is composed of 4 separate diffusion
models, requires 5.5B parameters and 356 sampling steps in total. Such
enormous model size, in turn, makes DPMs extremely cumbersome to be employed
in resource-limited platforms.

However, existing efforts towards efficient DPMs have focused on model
acceleration but largely overlooked model lightening. For examples, the
approaches of [36, 47, 34, 51, 1, 30, 33] strive for faster sampling, while
those of [17, 13, 57, 43] rely on reducing the input size. Admittedly, all of
these methods give rise to shortened training or inference time, yet still,
the large sizes preclude them from many real-world application scenarios.

Figure 1: (1) Visualizing the frequency domain gap among generated images with
the full DPM [43], Lite DPM and our SD on FFHQ [26] dataset. Lite-DPM is
unable to recover fine-grained textures, while SD can produce sharp edges and
realistic patterns. (2) Model size, Multiply-Add cumulation (MACs) and FID
score on class-conditioned ImageNet [7]. Our model achieves compelling visual
quality, with minimal parameters and computational cost. ∗ indicates our re-
implemented version.

In this paper, we make a dedicated attempt towards building compact DPMs. To
start with, we train a lite version of the popular latent diffusion model
(LDM) [43] by reducing the channel size. We show the image generated by the
original and and lite DPM in Figure 1. Although the lite LDM indeed sketches
the overall structure of the faces, the high-frequency components, such as the
skin and hair textures, are unfortunately poorly recovered. This phenomenon
can be in fact revealed by the Discrete Fourier Transform (DFT) coefficient
shown on the right column, indicating that the conventional design for DPMs
leads to high-frequency deficiency when the model is made slim.

We then take an in-depth inspection on the DPMs through the lens of frequency,
which results in two key observations. (1) Frequency Evolution. Under mild
assumptions, we mathematically prove that DPMs learn different functionalities
at different time-steps. Specifically, we show that the optimal denoiser in
fact boils down to a cascade of wiener filters [61] with growing bandwidths.
After recovering the low-frequency components, high-frequency features are
added gradually in the later denoising stages. This evolution property, as a
consequence, small DPMs fails to learn dynamic bandwidths with limited
parameters. (2) Frequency Bias. DPM is biased towards dominant frequency
components of the data distribution. It is most obvious when the noise
amplitude is small, leading to inaccurate noise prediction at the end of the
reverse process. As such, small DPMs struggle to recover the high-frequency
band and image details.

Motivated by these observations, we propose a novel Spectral Diffusion (SD)
model, tailored for light-weight image synthesis. Our core idea is to
introduce the frequency dynamics and priors into the architecture design and
training objective of the small DPM, so as to explicitly preserve the high-
frequency details. The proposed solution consists of two parts, each
accounting for one aforementioned observations. For the frequency evolution,
we propose a wavelet gating operation, which enables the network to
dynamically adapt to the spectrum response at different time-steps. In the
upsample and downsample stage, the input feature is first decomposed through
wavelet transforms and the coefficients are re-weighted through a learnable
gating function. It significantly lowers the parameter requirements to
represent the frequency evolution in the reverse process.

To compensate for the frequency bias for small DPMs, we distill the high-
frequency knowledge from teacher DPM to a compact network. This is implemented
by inverse weighting the distillation loss based on spectrum magnitudes.
Specifically, high-frequency recovery is strengthened by over-weighting
frequency bands of small magnitudes. Students thereby focus on the textual
recovery for image generation. By seamlessly integrating both designs, we are
able to build a slim latent diffusion model, SD, which largely preserve the
performance of LDM. Notably, SD by nature inherits the merits of DPMs,
including superior sample diversity, training stability and tractable
parameterization. As shown in Figure 1, our model is $8\sim 18\times$ smaller
and runs $2\sim 5\times$ faster than the original LDM, while achieving
competitive image fidelity.

The contributions of this study are threefold:

  1. 1.

This study investigates the task of diffusion model slimming, which remains
largely unexplored before.

  2. 2.

We identify that the key challenge lies in its unrealistic recovery for the
high-frequency components. By probing DPMs from a frequency perspective, we
show that there exists a spectrum evolution over different denoising steps,
and the rare frequencies cannot be accurately estimated by small models.

  3. 3.

We propose SD, a slim DPM that effectively restores imagery textures by
enhancing high-frequency generation performance. SD achieves gratifying
performance on image generation tasks at a low cost.

##  2 Related Work

Diffusion Probabilistic Models. DPMs [50, 16] have achieved state-of-the-art
results in terms of both log-likelihood estimation [52] and sample quality
[8], compared to Generative adversarial Network (GAN)-based [26, 12, 25]
approaches. It has been pointed out that DPM, in its essence, is a score-based
model [60, 55, 54] with annealed noise scheduling [53]. The reverse process is
considered as solving reverse stochastic differential equations (SDE) [55].
Current best-performed DPMs are implemented as a time-conditioned UNet [45, 8,
55] armed with self-attention [59] and cross-attention [43, 21]. Parameter
moving average [38], re-weighted objective [16] and advanced scheduling [38]
significantly improves the visual quality. In this work, we focus on small
diffusion designed for image generation, which has rarely been studied before.

Efficient Diffusion. The efficient diffusion model for low-resource inferences
has recently become a popular research topic. One approach is through reducing
the sampling steps, which is either done by distilling multiple steps into a
single step [36, 47, 34], or shortening the reverse steps while maintaining
the image fidelity [51, 1, 30, 33]. Another possible solution explores the
idea of diffusing in a lower dimensional space , and then scaling it up, with
a cascade structure [17] or in the latent space [57, 43]. In distinction from
them, we build an efficient diffusion model using light-weight architecture
and knowledge distillation.

Frequency Analysis for Generative Model. Neural networks tend to fit low-
frequency signals first and shift to the high-frequency components, which is
referred to as _frequency principle_ of deep neural network [63, 64, 2]. The
frequency bias is also observed when training deep generative models like GANs
[10, 5, 27, 49], where the generator struggles to build up natural high-
frequency details.

In this paper, we examine the frequency behavior of DPMs. Taking advantage of
its frequency properties, our SD achieves realistic image generation at a low
cost.

##  3 Background

###  3.1 Denoising Diffusion Probabilistic Models

Diffusion model reverses a progressive noise process based on latent
variables. Given data $\mathbf{x}_{0}\sim q(\mathbf{x}_{0})$ sampled from the
real distribution, we consider perturbing data with Gaussian noise with zero
mean and $\beta_{t}$ variance for $T$ steps

$\displaystyle
q(\mathbf{x}_{t}|\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1},\beta_{t}\mathbf{I})$
(1)

where $t\in[1,T]$ and $0<\beta_{1:T}<1$ denote the noise scale scheduling. At
the end of day, $\mathbf{x}_{T}\to\mathcal{N}(0,\mathbf{I})$ converge to a
Gaussian white noise. Although sampling from noise-perturbed distribution
$q(\mathbf{x}_{t})=\int q(\mathbf{x}_{1:t}|\mathbf{x}_{0})d\mathbf{x}_{1:t-1}$
requires a tedious numerical integration over steps, the choice of Gaussian
noise provides a close-form solution to generate arbitrary time-step
$\mathbf{x}_{t}$ through

$\displaystyle\mathbf{x}_{t}=\sqrt{\bar{\alpha}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}}\bm{\epsilon},\quad\text{where}\quad\epsilon\sim\mathcal{N}(0,\mathbf{I})$
(2)

where $\alpha_{t}=1-\beta_{t}$ and
$\bar{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$. A variational Markov chain in
the reverse process is parameterized as a time-conditioned denoising neural
network $\mathbf{s}(\mathbf{x},t;\bm{\theta})$ with
$p_{\bm{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_{t})=\mathcal{N}(\mathbf{x}_{t-1};\frac{1}{\sqrt{1-\beta_{t}}}(\mathbf{x}_{t}+\beta_{t}\mathbf{s}(\mathbf{x}_{t},t;\bm{\theta})),\beta_{t}\mathbf{I})$.
The denoiser is trained to minimize a re-weighted evidence lower bound (ELBO)
that fits the noise

$\displaystyle\mathcal{L}_{\text{DDPM}}$
$\displaystyle=\mathbb{E}_{t,\mathbf{x}_{0},\bm{\epsilon}}\Big{[}||\bm{\epsilon}-\mathbf{s}(\mathbf{x}_{t},t;\bm{\theta})||_{2}^{2}\Big{]}$
(3)
$\displaystyle=\mathbb{E}_{t,\mathbf{x}_{0},\bm{\epsilon}}\Big{[}||\nabla_{\mathbf{x}_{t}}\log
p(\mathbf{x}_{t}|\mathbf{x}_{0})-\mathbf{s}(\mathbf{x}_{t},t;\bm{\theta})||_{2}^{2}\Big{]}$
(4)

where the $\nabla_{\mathbf{x}_{t}}\log p(\mathbf{x}_{t}|\mathbf{x}_{0})$ are
also called the score function [53]. Thus, the denoiser equivalently learns to
recover the derivative that maximize the data log-likelihood [22, 60]. With a
trained
$\mathbf{s}(\mathbf{x},t;\bm{\theta}^{*})\approx\nabla_{\mathbf{x}_{t}}\log
p(\mathbf{x}_{t}|\mathbf{x}_{0})$, we generate the data by reversing the
Markov chain

$\displaystyle\mathbf{x}_{t-1}\leftarrow\frac{1}{\sqrt{1-\beta_{t}}}(\mathbf{x}_{t}+\beta_{t}\mathbf{s}(\mathbf{x}_{t},t;\bm{\theta}))+\sqrt{\beta_{t}}\bm{\epsilon}_{t}$
(5)

The reverse process could be understood as going along
$\nabla_{\mathbf{x}_{t}}\log p(\mathbf{x}_{t}|\mathbf{x}_{0})$ from
$\mathbf{x}_{T}$ to maximize the data likelihood.

###  3.2 Frequency Domain Representation of Images

Frequency domain analysis decomposes a image according to a sets of basis
functions. We focus on two discrete transformations: _Fourier_ and _Wavelet_
Transform.

Given a $H\times W$ input signal111For simplicity, we only introduce the
formulation for gray-image, while it is extendable to multi-channel inputs.
$\mathbf{x}\in\mathbb{R}^{H\times W}$, Discrete Fourier Transform (DFT)
$\mathcal{F}$ projects it onto a collections of sine and cosine waves of
different frequencies and phases

$\displaystyle\mathcal{X}(u,v)=\mathcal{F}[\mathbf{x}]=\sum_{x=1}^{H}\sum_{y=1}^{W}\mathbf{x}(x,y)e^{-j2\pi(\frac{u}{H}x+\frac{v}{W}y)}$

$\mathbf{x}(x,y)$ is the pixel value at $(x,y)$; $\mathcal{X}(u,v)$ represents
complex value at frequency $(u,v)$; $e$ and $j$ are Euler’s number and the
imaginary unit.

On the other hand, Discrete Wavelet Transform (DWT) projects it onto multi-
resolution wavelets functions. In a singlescale case, $\mathbf{x}$ is
decomposed into 4 wavelet coefficients
$\mathbf{x}_{\textsf{LL}},\mathbf{x}_{\textsf{LH}},\mathbf{x}_{\textsf{HL}},\mathbf{x}_{\textsf{HH}}=\textsf{DWT}(\mathbf{X})$
with halving the scale, where
$\mathbf{x}_{\\{\textsf{LL},\textsf{LH},\textsf{HL},\textsf{HH}\\}}\in\mathbb{R}^{\frac{H}{2}\times\frac{W}{2}}$.
$\mathbf{x}_{\textsf{LL}}$ stands for low-frequency component and
$\mathbf{x}_{\\{\textsf{LH},\textsf{HL},\textsf{HH}\\}}$ are high-frequency
components that contains the textural details. The coefficients could then be
inverted and up-sampled back to the original input
$\mathbf{x}=\textsf{IDWT}(\mathbf{x}_{\textsf{LL}},\mathbf{x}_{\textsf{LH}},\mathbf{x}_{\textsf{HL}},\mathbf{x}_{\textsf{HH}})$.

##  4 Frequency Perspective for Diffusion

In general signal processing, denoising is often performed in frequency space.
Similar to Figure 1, Table 1 compares Low-freq and High-freq error222The error
computed as the
$\mathbb{E}_{f}[\mathbb{E}[|\mathcal{F}_{real}|]-\mathbb{E}[|\mathcal{F}_{gen}|]]$
over 300 real and generated samples, with the low-high cut-off frequency of
28Hz. for different DPMs on FFHQ dataset. Lite-LDM performs poorly due to its
lack of high-frequency generation.

Method | #Param |  FID$\downarrow$ |  Low-freq Error$\downarrow$ |  High-freq Error$\downarrow$
---|---|---|---|---
LDM | 274.1M | 5.0 | 0.11 | 0.75
Lite-LDM | 22.4M | 17.3 |  0.28(+0.17) |  3.35(+2.17)
Table 1: Low-freq and High-freq error for different model size.

Thus, we examine DPM’s behavior in the frequency domain. As illustrated in
Figure 2, we make two findings: (1) _Frequency Evolution._ Diffusion model
learns to recover the low-frequency components at first, and gradually adds in
photo-realistic and high-frequency details. (2) _Frequency Bias._ Diffusion
model makes biased recovery for the minority frequency band.

Figure 2: Illustration of the Frequency Evolution and Bias for Diffusion
Models. In the reverse process, the optimal filters recover low-frequency
components first and add on the details at the end. The predicted score
functions may be incorrect for rare patterns, thus failing to recover complex
and fine-grained textures.

###  4.1 Spectrum Evolution over Time

DPM optimizes a time-conditioned network to fit noise at multiple scales,
which gives rise to a denoising trajectory over time-steps. We take a close
look at this trajectory through the lens of frequency. When assuming the
network is a linear filter, we give the optimal filter in terms of its
spectrum response at every timestep. This filter is commonly known as Wiener
filter [61].

Proposition 1. Assume $\mathbf{x}_{0}$ is a wide-sense stationary signal and
$\bm{\epsilon}$ is white noise of variance $\sigma^{2}=1$. For
$\mathbf{x}_{t}=\sqrt{\bar{\alpha}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}}\bm{\epsilon}$,
the optimal linear denoising filter $h_{t}$ at time $t$ that minimize
$J_{t}=\|h_{t}\ast\mathbf{x}_{t}-\bm{\epsilon}\|^{2}$ has a closed-form
solution

$\displaystyle\mathcal{H}_{t}^{*}(f)=\frac{1}{\bar{\alpha}|\mathcal{X}_{0}(f)|^{2}+1-\bar{\alpha}}$
(6)

where $|\mathcal{X}_{0}(f)|^{2}$ is the power spectrum of $\mathbf{x}_{0}$ and
$\mathcal{H}^{*}_{t}(f)$ is the frequency response of $h_{t}^{*}$.

Although the linear assumption poses a strong restriction on the model
architecture, we believe it provides valuable insights into what has been done
in the reverse process.

DPM goes from structure to details. In this study, we make a widely accepted
assumption about the power spectra of natural images
$\mathbb{E}[|X_{0}(f)|^{2}]=A_{s}(\theta)/f^{\alpha_{S}(\theta)}$that follows
a power law [58, 3, 9, 56]. $A_{s}(\theta)$ is called an amplitude scaling
factor and $\alpha_{S}(\theta)$ is the frequency exponent. If we set
$A_{s}(\theta)=1$ and $\alpha_{S}(\theta)=2$, the frequency response of the
signal reconstruction filter $1-\sqrt{1-\bar{\alpha}}h$ is in Figure 3.

In the reverse process, $t$ goes from $T\to 0$, and $\bar{\alpha}$ increases
from $0\to 1$. Therefore, DPM displays a spectrum-varying behavior over time.
In the beginning, we have a narrow-banded filter ($\bar{\alpha}=0.1$ and
$\bar{\alpha}=0.01$) that only restores the low-frequency components that
control the rough structures. $t$ goes down and $\bar{\alpha}$ gradually
increases, with more details and high-frequency components restored in the
images, like the human hairs, wrinkles, and pores.

We plot the denoised predictions $\hat{\mathbf{x}}_{0}$ at different steps
using pre-trained LDM [43] in Figure 2, which shows that DPM generates low-
frequency first and transits into high-frequency. The same empirical
observation that DPM goes from rough to details has been shown in [16, 35, 6,
43], while we are the first to give its numerical solutions.

Figure 3: $1-(1-\bar{\alpha})|H^{*}(f)|^{2}$ of the optimal linear denoising
filter with different $\bar{\alpha}$. Figure 4: Toy example for 1D signal
fitting. Small DPM is unable to recover minority frequency components.

###  4.2 Frequency Bias in Diffusion Model

Another challenge in diffusion-based model is the inaccurate denoising
estimation in low-density regions [53]. It results from the expectation over
$p(\mathbf{x}_{0})$ in the loss function

$\displaystyle\mathcal{L}_{\text{DDPM}}=\int
p(\mathbf{x}_{0})\mathbb{E}_{t,\bm{\epsilon}}\Big{[}||\bm{\epsilon}-\mathbf{s}(\mathbf{x}_{t},t;\bm{\theta})||_{2}^{2}\Big{]}\text{d}\mathbf{x}_{0}$
(7)

Since the denoising objective is weighted by $p(\mathbf{x}_{0})$, the trained
diffusion will be biased towards the high-density region, while ignoring the
long-tail patterns.

For image generation tasks, one long-tail pattern is the frequency bias. While
the low-frequency images are dominant with large $p(\mathbf{x}_{0})$, very few
samples contain high-frequency components. Training small DPMs on the biased
data distribution makes it difficult to generate samples with complex textures
and realistic high-frequency patterns.

Example 1. We fit a toy diffusion model to 1D functions $f(x)=cos(\alpha 2\pi
x)$, where $P(\alpha=3)=0.2$ and $P(\alpha=5)=0.8$. We adopt a two-layer feed-
forward neural network, with 1000 denoising steps and hidden units
$M=\\{64,1024\\}$. More details in in Supplementary.

We plot the 300 generated signals in Figure 4 (Top), their DFT magnitudes in
(Button Right), and the mean frequency histogram in (Button Left). Small model
($M=64$) faces difficulty recovering the minority frequencies other than
$\alpha=3$, while large model ($M=1024$) achieves smooth denoised results over
all freq bands, especially when $\alpha=5$.

It provides concrete evidence that small DPMs have intrinsic defects in
recovering the high frequencies.

##  5 Spectral Diffusion Model

As explained above, our goal is to slim down the DPMs by introducing the
frequency dynamics and priors into the architecture design and training
objectives. Taking the LDM [43] as our baseline, we design a wavelet-gating
module to enable time-dynamic inference for the network with a limited model
size. A spectrum-aware distillation is applied to enhance the high-frequency
generation performance. Both modifications allow us to achieve photo-realistic
image generation with minimal model size and computational effort.

###  5.1 Dynamic Wavelet Gating

As depicted in Section 4.1, the reverse process requires a cascade of filters
with dynamic frequency response. Vanilla UNet [45], while being effective in
reconstructing image details, is incapable to incorporate dynamic spectrum
into a single set of parameters. As a result, the small-size DPM is incapable
to compensate for the changing bandwidth.

In response to such frequency evolution, we propose to insert the Wavelet
Gating (WG) module into the network to automatically adapt it to varing
frequency response. WG decomposes the feature map into wavelet bands and
selectively attends to the proper frequency at different reverse steps, which
is uniquely tailored for the diffusion model.

Gating over the Wavelet Coefficients. We replace all down-sample and up-sample
in UNet with DWT and IDWT [11, 65], and pose a soft gating operation on
wavelet coefficients to facilitate step-adaptive image denoising. We call them
WG-Down and WG-Up, as shown in Figure 5.

Following the channel attention operation [62, 20, 40], information from input
feature $\mathbf{X}$ is aggregated to produce a soft gating mask

$\displaystyle g_{\\{\textsf{LL},\textsf{LH},\textsf{HL},\textsf{HH}\\}}$
$\displaystyle=\text{Sigmoid}(\text{FFN}(\text{Avgpool}(\mathbf{X})))$ (8)

where $g_{i}$ is the gating score of each wavelet band; FFN is a 2 layer where
feed-forward network and Avgpool stands for the average pooling. The
coefficients are then gated with $g_{i}$ to produce the output
$\mathbf{X}^{\prime}$.

Figure 5: WG-Down and WG-Up with wavelet gating.

In the WG-Down, we apply WG after the DWT operation to fuse the sub-band
coefficients with weighted summation
$\mathbf{X}^{\prime}=\sum_{i\in\\{\textsf{LL},\textsf{LH},\textsf{HL},\textsf{HH}\\}}g_{i}\odot\mathbf{X}_{i}$,
where $\odot$ is the element-wise multiplication. In the WG-Up, the input
feature is splitted into 4 chunks as the wavelet coefficients. Then, WG is
carried out to re-weight each sub-band before
$\mathbf{X}^{\prime}=\textsf{IDWT}(g_{\textsf{LL}}\odot\mathbf{X}_{\textsf{LL}},g_{\textsf{LH}}\odot\mathbf{X}_{\textsf{LH}},g_{\textsf{HL}}\odot\mathbf{X}_{\textsf{HL}},g_{\textsf{HH}}\odot\mathbf{X}_{\textsf{HH}})$.
In this paper, we apply Haar wavelet by default.

###  5.2 Spectrum-Aware Knowledge Distillation

Diffusion model has difficulty in modelling the high-frequency components (in
Section 4.2), especially for efficient requirements. In combat with spectrum
deficiency in image generation, we distill the prediction of a large pre-
trained teacher model to a compact WG-Unet student. Beyond output matching
with a L2 loss, a Spectrum-Aware Distillation is applied to guide the student
to synthesize naturalistic image details. Our intuition is to re-weight the
distillation loss according to the spectrum magnitude. For components with low
magnitudes, such as high-frequency bands, we increase the error penalty; while
the weight for the low-frequency elements is reduced.

Suppose a teacher diffusion model $\bm{s}_{T}(\cdot;\bm{\theta}_{T})$, we
would like to distill a student $\bm{s}_{T}(\cdot;\bm{\theta}_{T})$ by
mimicking the outputs and features . At time-step $t$, the perturbed image
$\textbf{x}_{t}$ is fed into both networks to produce the outputs and
features. A L2 loss [44, 31] is use to quantify their spatial distance

$\displaystyle\mathcal{L}_{\text{spatial}}=\sum_{i}\|\mathbf{X}^{(i)}_{T}-\mathbf{X}^{(i)}_{S}\|_{2}^{2}$
(9)

where $\mathbf{X}^{(i)}_{T}$ and $\mathbf{X}^{(i)}_{S}$ stand for the pair of
teacher/student’s output features or outputs of the same scale. A single
1$\times$1 Conv layer is used to align the dimensions between a prediction
pair.

In addition to the spatial distillation, inspired by the imbalanced learning
[28, 4, 23] and long-tail learning [67, 24], we design a distillation loss to
encourage the model for minority frequency recovery. Given a pair of model
predictions and the clean image $\mathbf{x}_{0}$, we first interpolate
$\mathbf{x}_{0}$ to the same size of the feature map, then take their 2D DFT

$\displaystyle\mathcal{X}^{(i)}_{T}=\mathcal{F}[\mathbf{X}^{(i)}_{T}],\mathcal{X}^{(i)}_{S}=\mathcal{F}[\mathbf{X}^{(i)}_{S}],\mathcal{X}^{(i)}=\mathcal{F}[\text{Resize}(\mathbf{x}_{0})]$
(10)

The $\mathcal{X}_{0}$ is then applied to modulate the difference between
$\mathcal{X}^{(i)}_{T}$ and $\mathcal{X}^{(j)}_{S}$

$\displaystyle\mathcal{L}_{\text{freq}}=\frac{}{}\sum_{i}\omega_{i}\|\mathcal{X}^{(i)}_{T}-\mathcal{X}^{(j)}_{S}\|_{2}^{2},\text{where
}\omega=|\mathcal{X}^{(i)}|^{\alpha}$ (11)

with a scaling factor $\alpha<0$ ($\alpha=-1$ in our experiment),
$\mathcal{L}_{\text{freq}}$ pushes the student towards learning the minority
frequencies yet down-weights the majority components. Together with the DDPM
objective in Eq. 3, our training objective becomes
$\mathcal{L}=\mathcal{L}_{\text{DDPM}}+\lambda_{s}\mathcal{L}_{\text{spatial}}+\lambda_{f}\mathcal{L}_{\text{freq}}$
with weighting factors $\lambda_{s}=0.1$ and $\lambda_{f}=0.1$.

Note that our method aims to learn accurate score prediction at each denoising
step, which is orthogonal to existing distillation on sampling step reduction
[47, 36].

##  6 Experiments

This section demonstrates the ability of our approach SD on high-resolution
image synthesis (Section 6.1) with limited computation, and validates the
significance of each proposed module via ablation study in Section 6.2.

FFHQ $256\times 256$
---
Model | #Param | MACs |  FID$\downarrow$
DDPM [16] | 113.7M | 248.7G | 8.4
P2 [6] | 113.7M | 248.7G | 7.0
LDM [43] | 274.1M | 96.1G | 5.0
Lite-LDM |  22.4M($12.2\times$) |  7.9G($12.2\times$) |  17.3($-12.3$)
Ours |  21.1M($13.0\times$) |  6.7G($14.3\times$) |  10.5($-5.5$)

CelebA-HQ $256\times 256$
---
Model | #Param | MACs |  FID$\downarrow$
Score SDE [55] | 65.57M | 266.4G | 7.2
DDGAN [57] | 39.73M | 69.9G | 7.6
LDM [43] | 274.1M | 96.1G | 5.1
Lite-LDM |  22.4M($12.2\times$) |  7.9G($12.2\times$) |  14.3($-9.2$)
Ours |  21.1M($13.0\times$) |  6.7G($14.3\times$) |  9.3($-4.2$)

LSUN-Bedroom $256\times 256$
---
Model | #Param | MACs |  FID$\downarrow$
DDPM [16] | 113.7M | 248.7G | 4.9
IDDPM [38] | 113.7M | 248.6G | 4.2
ADM [8] | 552.8M | 1114.2G | 1.9
LDM [43] | 274.1M | 96.1G | 3.0
Lite-LDM |  22.4M($12.2\times$) |  7.9G($12.2\times$) |  10.9($-7.9$)
Ours |  21.1M($13.0\times$) |  6.7G($14.3\times$) |  5.2($-2.2$)

LSUN-Church $256\times 256$
---
Model | #Param | MACs |  FID$\downarrow$
DDPM [16] | 113.7M | 248.7G | 4.9
IDDPM [38] | 113.7M | 248.6G | 4.3
ADM [8] | 552.8M | 1114.2G | 1.9
LDM [43] | 295.0M | 18.7G | 4.0
Lite-LDM |  32.8M($9.0\times$) |  2.1G($8.9\times$) |  13.6($-9.6$)
Ours |  33.8M($8.7\times$) |  2.1G($8.9\times$) |  8.4($-4.4$)

Table 2: Unconditional generation results comparison to prior DPMs. The
results are taken from the original paper, except that DDPM is take from the
[6].

Figure 6: Throughput for unconditional image generation.

Datasets and Evaluation. We evaluate our model on 4 unconditional generation
datasets and 2 conditional benckmarks. Specially, we train our unconditional
SSD models on LSUN-Churches/Bedrooms [66], FFHQ [26], and CelebA-HQ [25]. We
also validate the model on class-conditioned ImageNet [7] and MS-COCO [29]
text-to-image generation. For the text-to-image task, we first train on
LAION-400M [48] and test on MS-COCO directly.

Training and Evaluation Details. We build our model on the LDM [43]
frameworks. All pre-trained teachers and auto-encoders are downloaded from the
official repository333https://github.com/CompVis/latent-diffusion. For fair
comparison444[Generative models from other families (e.g. GAN, VAE, and Flow)
are excluded intentionally for fair computation comparison.], we implement a
lite-version of LDM, with a channel dimension of $64$ as our baseline model.
We call it Lite-LDM.

On 4 unconditional benchmarks, we train our spectral diffusion for 150k
iterations with a mini-batch size of $512$. We use AadmW [32] optimizer with
initial learning rate $1.024\text{\times}{10}^{-3}$ and linear lr decay. For
the class- and text-conditioned generation, the initial learning rate is set
to $5.12\text{\times}{10}^{-4}$ instead, with other parameters unchanged.
Classifier-free guidance [18] is applied. The synthesized image quality is
measured by the FID score [14] with 50k generated samples at the resolution of
$256$. We use a 200-step DDIM [51] sampling by default. We also compare the
model size and computational cost in terms of parameter number and Multiply-
Add cumulation (MACs)555https://github.com/sovrasov/flops-counter.pytorch.
Throughput is reported as our measurement of running speed. All experiments
are run on 8 NVIDIA Tesla V100 GPUs. More details are specified in the
Supplementary Material.

Figure 7: Randomly sampled $256\times 256$ images generated by our models
trained on CelebA-HQ [25], FFHQ [26], LSUN-Bedroom and LSUN-Church [66],
ImageNet [7]. All images are sampled with 200 DDIM steps.

###  6.1 Image Generation Results

Unconditional Image Generation. We train our SD on LSUN-Churches/Bedrooms [66]
FFHQ [26], and CelebA-HQ [25], and evaluate the sample quality. As shown in
Table 2, directly training small-sized diffusion models largely deteriorates
the model performance, such that Lite-LDM achieves an FID drop of $12.3$ on
FFHQ and $13.2$ on CelebA-HQ. Our proposed SD achieves $8\sim 14$ times
parameter and computation reduction compared to official LDM while being
competitive in image fidelity. For example, with a 21.1M Unet model and 6.7G
MACs, our SD gets an FID score of 5.2, which is very close to the 4.9 FID in
DDPM, but with only $\frac{1}{37}$ of its computation cost.

Throughput is reported in Figure 6. It refers to the number of time steps that
model runs per second. We measure its value with a batch-size of 64 by
averaging over 30 runs. We see that, Lite-LDM, while being fast, suffer
greatly from low visual quality. In comparison, our SD is $4.6\times$ faster
on CPU and $3.6\times$ on GPU compared to LDM on 3 of the 4 datasets.

We inspect the visual quality of the synthesized sample in Figure 7, row 1-4.
With much less parameters and complexity, our SD still produces realistic
samples with decent high-frequency details and sample diversity.

Class-conditional Image Generation. We validate our performance for class-
conditioned image generation on ImageNet. The results are demonstrated in
Table 3. With super-mini architecture and classifier-free guidance of $w=3.0$,
our SD reaches an FID score of 10.6. As the comparison, the ADM [8] only gets
FID=10.9, but with 553.8M parameters and 1114.2 MACs. Lite-LDM, though being
comparably fast, suffers from its inability for high-frequency generation,
gets a high FID score of 20.1.

Generated results are visualized in Figure 7 row 5-10. Our SD is able to
produce diverse images of different categories, particularly good at animal
generation like corgi and bear. However, we still observe failure cases with
distorted faces and shapes. For example, our models suffer in crowded instance
generation such as on banana.

Method | #Param | MACs |  FID$\downarrow$
---|---|---|---
IDDPM [38] | 273.1M | 1416.3G | 12.3
ADM [8] | 553.8M | 1114.2G | 10.9
LDM [43] | 400.9M | 99.8G | 10.6
ADM-G [8] | 553.8+54.1M | 1114.2+72.2G | 4.6
LDM-CFG [43] | 400.9M | 99.8G | 3.6
Lite-LDM-CFG |  47.0M($8.5\times$) |  11.1G ($9.0\times$) |  20.1($-16.5$)
Ours-CFG |  45.4M($8.8\times$) |  9.9G ($10.1\times$) |  10.6($-7.0$)
Table 3: Comparison of class-conditional image generation methods on ImageNet [7] with recent state-of-the-art methods. “G” stands for the classifier guidance and “CFG” refers to the classifer-free guidance for conditional image generation. Method | #Param |  FID$\downarrow$
---|---|---
GLIDE [37] | 5.0B | 12.24
DALLE2 [42] | 5.5B | 10.39
Imagen [46] | 3.0B | 7.27
LDM [43] | 1.45B | 12.63
Ours |  77.6M($18.7\times$) | 18.43
Table 4: Zero-Shot evaluation on MS-COCO text-to-image generation. We only
count the model size of diffusion part but exclude language encoder.  Figure
8: Selected samples from Spectral Diffusion using classifier-free guidance
$w=5.0$ for text-to-image generation.

Text-to-Image Generation. Following prior work [43], we train our text-
conditioned SD with a fixed CLIP encoder [41] on LAION-400M [48], and then do
zero-shot inference on MS-COCO [29] with $w=2.0$. Since each MS-COCO images
contains multiple captions, during evaluations, we randomly select 50k
descriptions from the train set, with one caption corresponding to a unique
image.

The evaluation results are provided in Table 4. Again, with a 77.6M model, we
gets to a FID score of 18.43, while being $18.7\times$ smaller than LDM. We
also provide qualitative analysis for text-to-image generation with new
prompts, in Figure 8. Although the image quality is not as perfect as in those
large-sized diffusion models, our model learns to compose vivid drawing
according to the descriptions, with minimal computational cost and portable
model size. Our SD is good at abstract or carton style paintings. However, it
is still challenging to generate human body and faces, as in the “basketball
player” example.

###  6.2 Ablation Study and Analysis

In this section, we validate the effectiveness of wavelet gating and spectrum-
aware distillation, on whether and how they help to improve the image
fidelity.

Effectiveness of Wavelet Gating. We validate the effectiveness of the Wavelet
Gating by replacing our WG upsample and downsample with the nearest neighbor
resizer in LDM [43] and train on the FFHQ dataset. As shown in Table 5,
removing WG significantly increases the FID from $10.5\to 12.4$. Besides, WG
alone improves Lite-LDM’s FID score by $2.6$. Both results indicate that WG

effectively promote the sample quality of the small DPMs.

In addition, we plot the values of the gating functions at different denoising
steps for a pre-trained text-to-image SD model in Figure 9. Each curve is
calculated by averaging the gating coefficient for 100 generated images. The
trends of the downsample and upsample operations diverge. In the end of
denoising (large $t$), high-frequency details emerged in
$\hat{\mathbf{x}}_{t}$. The WG-Down thus enhances the high-frequency signals
with increased $g_{\\{\textsf{HL},\textsf{LH},\textsf{HH}\\}}$ while keeping
the low-frequency part constant. In contrast, the WG-Up (Right) promotes
$g_{\textsf{LL}}$ in the late stage of denoising. Predicted noises boost its
low-frequency components, resulting in high-frequency element recovery in the
$\hat{\mathbf{x}}_{0}=\frac{\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}\bm{\epsilon}}}{\sqrt{\bar{\alpha}}}$.

Effectiveness of Spectrum-Aware Distillation. To understand the value of the
proposed SA-Distillation, we sequentially remove each loss term. Figure 5
shows that, while the spatial term only accounts for 0.9 FID, the frequency
term takes up 1.8 FID improvement, highlighting its importance in high-quality
image generation.

We also visualize the images generated by trained models with (W) or without
(W/O) the frequency term in Figure 10, with their DFT difference. The model
without $\mathcal{L}_{freq}$ makes smoother predictions, while our method
recovers the details like hair or architectural textures. By penalizing high-
frequency distillation, our proposed SA-Distillation resulted in large
differences and improvements in high-frequency components in
$|\mathcal{F}_{f}-\mathcal{F}_{nof}|$.

Method | FFHQ $256\times 256$ |
---|---|---
\+ Wavelet Gating |  | ✓ |  |  | ✓ |  | ✓ | ✓
\+ Spatial Distill |  |  | ✓ |  | ✓ | ✓ |  | ✓
\+ Freq Distill |  |  |  | ✓ |  | ✓ | ✓ | ✓
FID$\downarrow$ | 17.3 | 14.7 | 16.6 | 15.3 | 12.3 | 12.4 | 11.4 | 10.5
Table 5: Ablation study on FFHQ dataset. Figure 9: Wavelet gating function
values at different $t$. We plot the mean$\pm$std for 100 generated images.
Figure 10: Generated images W or W/O the freq term, as well as their DFT
difference $|\mathcal{F}_{\text{f}}-\mathcal{F}_{\text{nof}}|$.

##  7 Conclusion

In the study, we focus on reducing the computation cost for diffusion models.
The primary obstacle to training small DPMs is their inability to provide
realistic high-frequency, which results from the frequency evolution and bias
of diffusion process. In order to resolve these problems, we propose Spectral
Diffusion (SD) for efficient image generation. It performs spectrum dynamic
denoising by using a wavelet gating operation, which automatically enhances
different frequency bands at different reverse steps. A large pre-trained
network helps to improve the performance of high-frequency generation by
knowledge distillation. By seamlessly integrating both modifications, our
model is 8-18 $\times$ slimer and runs 2-5$\times$ faster than the latent
diffusion model, with negligible performance drop.

## References

  * [1] Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang.  Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models.  arXiv preprint arXiv:2201.06503, 2022.
  * [2] Ronen Basri, Meirav Galun, Amnon Geifman, David Jacobs, Yoni Kasten, and Shira Kritchman.  Frequency bias in neural networks for input of non-uniform density.  In International Conference on Machine Learning, pages 685–694. PMLR, 2020.
  * [3] Geoffrey J Burton and Ian R Moorhead.  Color and spatial structure in natural scenes.  Applied optics, 26(1):157–170, 1987.
  * [4] Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma.  Learning imbalanced datasets with label-distribution-aware margin loss.  Advances in neural information processing systems, 32, 2019.
  * [5] Yuanqi Chen, Ge Li, Cece Jin, Shan Liu, and Thomas Li.  Ssd-gan: Measuring the realness in the spatial and spectral domains.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 1105–1112, 2021.
  * [6] Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon.  Perception prioritized training of diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11472–11481, 2022.
  * [7] 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, pages 248–255. Ieee, 2009.
  * [8] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  Advances in Neural Information Processing Systems, 34:8780–8794, 2021.
  * [9] David J Field.  Relations between the statistics of natural images and the response properties of cortical cells.  Josa a, 4(12):2379–2394, 1987.
  * [10] Joel Frank, Thorsten Eisenhofer, Lea Schönherr, Asja Fischer, Dorothea Kolossa, and Thorsten Holz.  Leveraging frequency analysis for deep fake image recognition.  In International conference on machine learning, pages 3247–3258. PMLR, 2020.
  * [11] Minghan Fu, Huan Liu, Yankun Yu, Jun Chen, and Keyan Wang.  Dw-gan: A discrete wavelet transform gan for nonhomogeneous dehazing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 203–212, 2021.
  * [12] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial networks.  Communications of the ACM, 63(11):139–144, 2020.
  * [13] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo.  Vector quantized diffusion model for text-to-image synthesis.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10696–10706, 2022.
  * [14] 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.  Advances in neural information processing systems, 30, 2017.
  * [15] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al.  Imagen video: High definition video generation with diffusion models.  arXiv preprint arXiv:2210.02303, 2022.
  * [16] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  * [17] Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans.  Cascaded diffusion models for high fidelity image generation.  J. Mach. Learn. Res., 23:47–1, 2022.
  * [18] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  arXiv preprint arXiv:2207.12598, 2022.
  * [19] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet.  Video diffusion models.  arXiv:2204.03458, 2022.
  * [20] Jie Hu, Li Shen, and Gang Sun.  Squeeze-and-excitation networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7132–7141, 2018.
  * [21] Zilong Huang, Xinggang Wang, Lichao Huang, Chang Huang, Yunchao Wei, and Wenyu Liu.  Ccnet: Criss-cross attention for semantic segmentation.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 603–612, 2019.
  * [22] Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  Journal of Machine Learning Research, 6(4), 2005.
  * [23] Liming Jiang, Bo Dai, Wayne Wu, and Chen Change Loy.  Focal frequency loss for image reconstruction and synthesis.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 13919–13929, 2021.
  * [24] Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis.  Decoupling representation and classifier for long-tailed recognition.  In International Conference on Learning Representations, 2020.
  * [25] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.  Progressive growing of gans for improved quality, stability, and variation.  arXiv preprint arXiv:1710.10196, 2017.
  * [26] Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401–4410, 2019.
  * [27] Mahyar Khayatkhoei and Ahmed Elgammal.  Spatial frequency bias in convolutional generative adversarial networks.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7152–7159, 2022.
  * [28] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  In Proceedings of the IEEE international conference on computer vision, pages 2980–2988, 2017.
  * [29] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In European conference on computer vision, pages 740–755. Springer, 2014.
  * [30] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao.  Pseudo numerical methods for diffusion models on manifolds.  In International Conference on Learning Representations, 2022.
  * [31] Yifan Liu, Ke Chen, Chris Liu, Zengchang Qin, Zhenbo Luo, and Jingdong Wang.  Structured knowledge distillation for semantic segmentation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2604–2613, 2019.
  * [32] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [33] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.  Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.  arXiv preprint arXiv:2206.00927, 2022.
  * [34] Eric Luhman and Troy Luhman.  Knowledge distillation in iterative generative models for improved sampling speed.  arXiv preprint arXiv:2101.02388, 2021.
  * [35] Hengyuan Ma, Li Zhang, Xiatian Zhu, and Jianfeng Feng.  Accelerating score-based generative models with preconditioned diffusion sampling.  In European Conference on Computer Vision, 2022.
  * [36] Chenlin Meng, Ruiqi Gao, Diederik P Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans.  On distillation of guided diffusion models.  arXiv preprint arXiv:2210.03142, 2022.
  * [37] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen.  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  arXiv preprint arXiv:2112.10741, 2021.
  * [38] Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In International Conference on Machine Learning, pages 8162–8171. PMLR, 2021.
  * [39] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall.  Dreamfusion: Text-to-3d using 2d diffusion.  arXiv, 2022.
  * [40] Zequn Qin, Pengyi Zhang, Fei Wu, and Xi Li.  Fcanet: Frequency channel attention networks.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 783–792, 2021.
  * [41] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In International Conference on Machine Learning, pages 8748–8763. PMLR, 2021.
  * [42] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.  Hierarchical text-conditional image generation with clip latents.  arXiv preprint arXiv:2204.06125, 2022.
  * [43] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684–10695, 2022.
  * [44] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets.  arXiv preprint arXiv:1412.6550, 2014.
  * [45] 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, pages 234–241. Springer, 2015.
  * [46] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al.  Photorealistic text-to-image diffusion models with deep language understanding.  arXiv preprint arXiv:2205.11487, 2022.
  * [47] Tim Salimans and Jonathan Ho.  Progressive distillation for fast sampling of diffusion models.  In International Conference on Learning Representations, 2022.
  * [48] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki.  LAION-400M: open dataset of clip-filtered 400 million image-text pairs.  CoRR, abs/2111.02114, 2021.
  * [49] Katja Schwarz, Yiyi Liao, and Andreas Geiger.  On the frequency bias of generative models.  Advances in Neural Information Processing Systems, 34:18126–18136, 2021.
  * [50] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In International Conference on Machine Learning, pages 2256–2265. PMLR, 2015.
  * [51] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502, 2020.
  * [52] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon.  Maximum likelihood training of score-based diffusion models.  Advances in Neural Information Processing Systems, 34:1415–1428, 2021.
  * [53] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  Advances in Neural Information Processing Systems, 32, 2019.
  * [54] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon.  Sliced score matching: A scalable approach to density and score estimation.  In Uncertainty in Artificial Intelligence, pages 574–584. PMLR, 2020.
  * [55] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In International Conference on Learning Representations, 2021.
  * [56] David J Tolhurst, Yoav Tadmor, and Tang Chao.  Amplitude spectra of natural images.  Ophthalmic and Physiological Optics, 12(2):229–232, 1992.
  * [57] Arash Vahdat, Karsten Kreis, and Jan Kautz.  Score-based generative modeling in latent space.  Advances in Neural Information Processing Systems, 34:11287–11302, 2021.
  * [58] van A Van der Schaaf and JH van van Hateren.  Modelling the power spectra of natural images: statistics and information.  Vision research, 36(17):2759–2770, 1996.
  * [59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [60] Pascal Vincent.  A connection between score matching and denoising autoencoders.  Neural computation, 23(7):1661–1674, 2011.
  * [61] Norbert Wiener, Norbert Wiener, Cyberneticist Mathematician, Norbert Wiener, Norbert Wiener, and Cybernéticien Mathématicien.  Extrapolation, interpolation, and smoothing of stationary time series: with engineering applications, volume 113.  MIT press Cambridge, MA, 1949.
  * [62] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In So Kweon.  Cbam: Convolutional block attention module.  In Proceedings of the European conference on computer vision (ECCV), pages 3–19, 2018.
  * [63] Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma.  Frequency principle: Fourier analysis sheds light on deep neural networks.  arXiv preprint arXiv:1901.06523, 2019.
  * [64] Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao.  Training behavior of deep neural network in frequency domain.  In International Conference on Neural Information Processing, pages 264–274. Springer, 2019.
  * [65] Mengping Yang, Zhe Wang, Ziqiu Chi, and Wenyi Feng.  Wavegan: Frequency-aware gan for high-fidelity few-shot image generation.  arXiv preprint arXiv:2207.07288, 2022.
  * [66] 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.
  * [67] Yifan Zhang, Bingyi Kang, Bryan Hooi, Shuicheng Yan, and Jiashi Feng.  Deep long-tailed learning: A survey.  arXiv preprint arXiv:2110.04596, 2021.
  * [68] Daquan Zhou, Weimin Wang, Hanshu Yan, Weiwei Lv, Yizhe Zhu, and Jiashi Feng.  Magicvideo: Efficient video generation with latent diffusion models.  arXiv preprint arXiv:2211.11018, 2022.


Nancy Xu, Giannis Nikolentzos, Michalis Vazirgiannis, and Henrik Boström
11institutetext: KTH Royal Institute of Technology, Stockholm, Sweden
11email<EMAIL_ADDRESS>22institutetext: Athens University of
Economics and Business, Athens, Greece
22email<EMAIL_ADDRESS>33institutetext: École Polytechnique, Palaiseau,
France

# Image Keypoint Matching using Graph Neural Networks

Nancy Xu 11 Giannis Nikolentzos 22 Michalis Vazirgiannis 1133 Henrik Boström
11

###### Abstract

Image matching is a key component of many tasks in computer vision and its
main objective is to find correspondences between features extracted from
different natural images. When images are represented as graphs, image
matching boils down to the problem of graph matching which has been studied
intensively in the past. In recent years, graph neural networks have shown
great potential in the graph matching task, and have also been applied to
image matching. In this paper, we propose a graph neural network for the
problem of image matching. The proposed method first generates initial soft
correspondences between keypoints using localized node embeddings and then
iteratively refines the initial correspondences using a series of graph neural
network layers. We evaluate our method on natural image datasets with keypoint
annotations and show that, in comparison to a state-of-the-art model, our
method speeds up inference times without sacrificing prediction accuracy.

###### keywords:

keypoint matching, graph neural networks, graph matching

##  1 Introduction

Image matching, the task of finding correspondences between the key features
extracted from one image and those extracted from another image of the same
object, is a fundamental task in computer vision. Indeed, image matching lies
at the heart of several applications in computer vision including 3D structure
recovery, simultaneous localization and mapping (SLAM), and change detection,
just to name a few. Over the past decades, a large amount of research was
devoted to image matching and led to the development of a diverse set of
approaches [23].

The pipeline that is typically followed by most image matching algorithms
consists of the extraction of a set of key features (or keypoints) from each
image and the generation of descriptors for those features. The above steps
are followed by some matching approach which is responsible for establishing
correspondences between the keypoints of the two images. Recently, the last
part of the pipeline, i. e., the task of keypoint matching in natural images,
has been formulated as a graph matching problem and has been addressed using
graph neural network architectures [28, 25, 9]. Images are represented as
graphs where nodes correspond to keypoints and edges capture proximity or
other types of relations between keypoints. In fact, the problem of graph
matching has been studied a lot in the past, but mainly from an algorithmic
viewpoint [5]. Unfortunately, most derived formulations correspond to NP-hard
problems, and therefore, are not useful for most real-world settings. With the
advent of deep learning and the success of graph neural networks, the problem
of learning how to compute correspondences of nodes (or of other structures)
between two graphs attracted increasing attention [28, 15, 26, 1]. Though not
exact, approaches that utilize graph neural networks to find correspondences
between nodes are very efficient, while they also learn representations that
are best suited for the task at hand.

In this paper, we capitalize on recent advances in graph representation
learning, and we propose a graph neural network for establishing meaningful
structural correspondences of keypoints between two natural images. More
specifically, we take as starting point the Deep Graph Matching Consensus
(DGMC) method proposed in [9]. DGMC operates in two stages, a local feature
matching stage which generates an initial soft correspondence matrix, and a
neighborhood consensus stage which iteratively refines the soft correspondence
matrix by repeatedly applying a single graph neural network to the scores. The
final correspondence matrix is a combination of the initial soft matrix and
the final result of the iterations.

We develop a new approach for obtaining and combining soft correspondence
matrices that is more intuitive and more attractive in terms of running time.
The proposed method first generates initial soft correspondences between
keypoints using localized node embeddings, and then iteratively refines the
initial correspondences using a series of graph neural network layers. The
final correspondence matrix is formed of all the matricies produced during the
consensus stage. This change increases the complexity of the neighborhood
consensus stage which reduces the number of iterations needed, thus decreasing
inference time. Experiments on standard keypoint matching datasets demonstrate
that in comparison to the DGMC, the proposed method is in most cases more
efficient, while in terms of performance, it matches or outperforms the
baseline model.

The rest of this paper is organized as follows. Section 2 provides an overview
of the related work. Section 3 provides a detailed description of the proposed
model. Section 4 evaluates the proposed model in the task of keypoint matching
in natural images. Finally, section 5 concludes.

##  2 Related Work

Traditional methods for identifying correspondences between images typically
deal with instance matching tasks where correspondences are found between two
images of the same scene or object. These methods use representations such as
image pyramids which are formed by repeatedly filtering and subsampling a
given image. Image features are computed densely over image pyramids and can
be sampled densely [7] or used to identify keypoints [22]. Notably, SIFT-Flow
combines the image pyramid with optical flow methods for instance matching
[19]. Unfortunately, most traditional methods depend on handcrafted features
which must be manually defined.

In comparison, deep learning methods automatically learn important image
features and are widely used in computer vision. For instance, Long et al.
applied in [21] a straightforward convolutional neural network to image
matching and achieved similar and in some cases, superior results to those of
SIFT-Flow. The introduction of datasets specialized for image matching such as
Proposal Flow (PASCAL-PF) [14], WILLOW-ObjectClass [3], and the PASCAL-VOC
Berkeley annotations [2], greatly improved the ease of developing new deep
learning methods. Examples include FCSS, which uses a convolutional self-
similarity layer to encode local self-similarity patterns [16] and SFNet which
outputs a semantic flow field using a CNN [18].

Finding correspondences between keypoints in images can also be seen as a
graph matching problem. Graph matching can be formulated as a quadratic
assignment problem, which turns out to be a notoriously difficult problem
[20]. A more computationally tractable formulation is to treat the problem as
an optimal transport problem. Xu et al. demonstrated in [27] the success of
using this formulation along with the Sinkhorn algorithm for graph matching
[17]. Graph neural networks can effectively and automatically generate node or
graph representations using message passing schemes [13]. Several recent
methods use a combination of graph neural networks and the Sinkhorn algorithm
for graph matching [28, 26, 25, 9]. By using a graph neural network to
generate similarity scores followed by the application of the Sinkhorn
normalization, we can build an end-to-end trainable framework for semantic
matching between keypoints extracted from images.

##  3 Graph Neural Networks for Image Keypoint Matching

Before presenting the proposed method, we begin by introducing some key
notation for graphs which will be used later. Let $G=(V,E)$ be a graph where
$V$ denotes the set of nodes and $E$ the set of edges. We will denote by $n$
the number of vertices and by $m$ the number of edges. Let also
$\mathbf{A}\in\mathbb{R}^{n\times n}$ denote the adjacency matrix of $G$.
Following previous work [9], we represent images as graphs where nodes
correspond to keypoints and edges are obtained via the Delaunay triangulation
of the keypoints. Furthermore, every node in the graph is associated with a
feature vector (i. e., the input features of keypoints), and we use
$\mathbf{X}\in\mathbb{R}^{n\times d}$ to denote those features where $d$ is
the feature dimensionality. The edges of the graph can also be potentially
associated with a feature vector (e. g., distance of two endpoints), and we
use $\mathbf{E}\in\mathbb{R}^{n\times n\times d}$ to denote those features.

Given a source and a target image, we fist map these images to graphs $G_{s}$
and $G_{t}$, respectively. Let $n_{s}$ and $n_{t}$ denote the number of nodes
of the two graphs. Without loss of generality, we assume that $n_{s}\leq
n_{t}$. The goal of the method is to establish correspondences between $G_{s}$
and $G_{t}$ and match two pairs of keypoints.

Similar to DGMC [9], our method consists of a local feature matching stage
followed by a neighborhood consensus stage, which are both presented in the
following subsections. The local feature matching stage finds initial
correspondences between images based on the similarity of their node
embeddings and is implemented as a graph neural network layer
$\text{GNN}^{(0)}$. The neighborhood consensus stage then refines these
initial correspondences which we implement as a series of graph neural network
layers $\text{GNN}^{(1)},\ldots,\text{GNN}^{(K)}$.

###  3.1 Local Feature Matching

Starting with $G_{s}$ and $G_{t}$, the representations of the nodes are
updated by feeding both graphs into a graph neural network, i. e.,
$\mathbf{H}=\text{GNN}^{(0)}(\mathbf{A},\mathbf{X},\mathbf{E})$. Since we use
a graph neural network to update the features of the keypoints, we obtain
localized and permutation equivariant representations. Then, the similarity
between keypoints extracted from the two images is determined by the
similarity of their node representations $\mathbf{H}_{s}$ and
$\mathbf{H}_{t}$. Specifically, we generate matrix $\mathbf{K}^{(0)}$ as
follows $\mathbf{K}^{(0)}=\mathbf{H}_{s}\,\mathbf{H}_{t}^{\top}$, and then, we
can obtain initial soft correspondences between the keypoints of the two
images via the Sinkhorn algorithm [6]:

$\mathbf{S}^{(0)}=\text{sinkhorn}\big{(}\mathbf{K}^{(0)}\big{)}\in[0,1]^{n_{s}\times
n_{t}}$

The Sinkhorn normalization produces a rectangular doubly-stochastic
correspondence matrix such that
$\sum_{j=1}^{n_{t}}\mathbf{S}_{i,j}^{(0)}=1,\,\forall
i\in\\{1,\ldots,n_{s}\\}$ and $\sum_{i=1}^{n_{s}}\mathbf{S}_{i,j}^{(0)}\leq
1,\,\forall j\in\\{1,\ldots,n_{t}\\}$.

As discussed in [9], the Sinkhorn algorithm is expensive in terms of running
time, while it also might lead to vanishing gradients. Furthermore, it may
also converge to inconsistent solutions which the next stage of the algorithm
might find difficult to refine. Therefore, instead of applying Sinkhorn, the
problem is partially relaxed by dropping the following constraint
$\sum_{i=1}^{n_{s}}\mathbf{S}_{i,j}^{(0)}\leq 1,\,\forall
j\in\\{1,\ldots,n_{t}\\}$, and the initial correspondence matrix is computed
by applying row-wise softmax normalization on matrix $\mathbf{K}^{(0)}$ as
follows:

$\mathbf{S}^{(0)}=\text{softmax}\big{(}\mathbf{K}^{(0)}\big{)}$

###  3.2 Neighborhood Consensus

The initial correspondences only take local features into account, which makes
it more prone to falsely matching nodes which are only superficially similar.
The neighborhood consensus step aims to detect these false correspondences and
iteratively resolve them. The iteration is done by applying a series of graph
neural network layers which we refer to as
$\text{GNN}^{(1)},\ldots,\text{GNN}^{(K)}$. In contrast to DGMC proposed in
[9], which uses a shared graph neural network for all iterations, in our case
no weight sharing is performed and each graph neural network consists of its
own trainable parameters. The idea behind this change is to increase the
complexity of the neighborhood consensus stage while increasing efficiency by
reducing the number of iterations needed to refine the correspondence
matrices.

At each iteration, a graph neural network layer produces a new correspondence
matrix by updating the previous correspondence matrix. This process is applied
$k$ times resulting in $\mathbf{S}^{(1)},\ldots,\mathbf{S}^{(K)}$.

For every iteration $k$, where $k\in\\{1,\ldots,K\\}$, each consensus loop
begins by assigning an arbitrary coloring to the nodes of $G_{s}$ using
randomly generated vectors. Let $\mathbf{R}_{s}$ be a matrix whose $i$-th row
contains the vector of the $i$-th node of $G_{s}$. Then, the model assigns a
coloring to the nodes of $G_{t}$ based on their relationship to the nodes of
$G_{s}$ (as captured by matrix $\mathbf{S}^{(k-1)}$). Thus, matrix
$\mathbf{R}_{t}$ is produced as follows:
$\mathbf{R}_{t}=\big{[}\mathbf{S}^{(k-1)}\big{]}^{\top}\,\mathbf{R}_{s}$.
These colors are distributed across their respective graphs by using
synchronous message passing, i. e., by applying $\text{GNN}^{(k)}$ to each
graph. That is, the model computes new matrices of features
$\mathbf{Z}_{s}^{(k)}=\text{GNN}^{(k)}(\mathbf{A}_{s},\mathbf{R}_{s},\mathbf{E}_{s})$
and
$\mathbf{Z}_{t}^{(k)}=\text{GNN}^{(k)}(\mathbf{A}_{t},\mathbf{R}_{t},\mathbf{E}_{t})$.
By comparing the emerging representations of two nodes, a matrix
$\mathbf{D}\in\mathbb{R}^{n_{s}\times n_{t}}$ is produced which measures the
neighborhood consensus between all pairs of nodes. For example, for the $i$-th
node of $G_{s}$ and the $j$-th node of $G_{t}$, we have that:

$\mathbf{D}_{i,j}=\big{[}\mathbf{Z}_{s}^{(k)}\big{]}_{i,:}-\big{[}\mathbf{Z}_{t}^{(k)}\big{]}_{j,:}$

where $\big{[}\mathbf{Z}\big{]}_{i,:}$ denotes the $i$-th row of matrix
$\mathbf{Z}$. Matrix $\mathbf{D}$ is then transformed using a multi-layer
perceptron and the new correspondence matrix is computed as:

$\mathbf{S}^{(k)}=\text{softmax}\big{(}\mathbf{S}^{(k-1)}+\text{MLP}(\mathbf{D})\big{)}$

Instead of considering only two correspondence matrices as in [9], in our
setting, the final correspondence matrix $\mathbf{S}$ is computed as the
weighted sum of all the soft correspondence matrices
$\mathbf{S}^{(0)},\ldots,\mathbf{S}^{(K)}$, where the weights correspond to
trainable parameters followed by a final softmax function (applied in a row-
wise manner).

$\mathbf{S}=\text{softmax}\Big{(}\mathbf{w}_{0}\,\mathbf{S}^{(0)}+\mathbf{w}_{1}\,\mathbf{S}^{(1)}+\ldots+\mathbf{w}_{k}\,\mathbf{S}^{(K)}\Big{)}$

where $\mathbf{w}\in\mathbb{R}^{k}$ is a trainable vector. To train the model,
we employ the categorical cross-entropy:

$\mathcal{L}=-\sum_{i=1}^{n_{s}}\log\big{(}\mathbf{S}_{i,\pi(i)}\big{)}$

where $\pi(i)$ gives the node of $G_{t}$ that matches to node $i$ of $G_{s}$.
Since the loss function is differentiable, we can train the model in an end-
to-end fashion using stochastic gradient descent.

##  4 Experiments

In this section, we report on the experimental evaluation of the proposed
method and compare with state-of-the-art methods.

###  4.1 Datasets

We evaluate the proposed method in the task of keypoint matching in natural
images and in the task of geometric feature matching, where only the
coordinates of the keypoints are available and not their visual features. For
the first task, we experiment with the following two datasets: (i) PASCAL-VOC
[8] with Berkeley annotations of keypoints [2]; and (ii) WILLOW-ObjectClass
[3]. The PASCAL-VOC dataset is an extension of PASCAL-VOC 2011 and contains
annotations of body parts for 20 semantic classes. The Willow ObjectClass
dataset consists of 5 categories from Caltech-256 and Pascal VOC 2007. For the
second task, we experiment with the PASCAL-PF dataset [14]. The PASCAL-PF
dataset consists of annotations for 20 semantic classes. The annotations for
each image pair consist of a set of keypoint coordinate pairs where each pair
shares the same arbitrary numeric label.

###  4.2 Baselines

We compare the proposed model against recently proposed neural network models
that have achieved stat-of-the-art performance. More specifically, on the
PASCAL-VOC and WILLOW ObjectClass datasets, we compare the proposed model
against the following three models: (i) GMN [28]; (ii) PCA-GM [26]; and (iii)
DGMC [9]. On the PASCAL-PF dataset, we compare our model against two models:
(i) GMN [29]; and (ii) DGMC [9]. For the DGMC model, $L$ refers to the number
of neighborhood consensus iterations used. For all baselines, we use the
results reported in [9].

###  4.3 Experimental Setup

On all three considered datasets, we follow the experimental setups of
previous studies [28, 26, 9, 4], and we use the same training and test splits
as in those studies.

For the proposed model, we set the number of epochs to $30$, and the batch
size to $512$. We use the Adam optimizer with a learning rate of $10^{-3}$. We
set the hidden-dimension size of the graph neural networks to $128$. We tune
the following two hyperparameters: (i) the number of graph neural networks
used in the neighborhood consensus step $K$; and (ii) the number of message
passing iterations of each graph neural network $r$. For the proposed model
and DGMC, two types of edge features are evaluated: (i) isotropic edge
features that use normalized relative distances which result in the neighbors
of a node being equally considered during message passing; and (ii)
anisotropic edge features that use 2D cartesian coordinates which result in
some neighbors having a greater weight during message passing.
$\text{GNN}^{(0)},\ldots,\text{GNN}^{(K)}$ all correspond to SplineCNN models
[11]. The message passing scheme employed by the SplineCNN model is defined
as:

$\mathbf{h}_{i}^{(t+1)}=\sigma\Big{(}\mathbf{W}^{(t+1)}\mathbf{h}_{i}^{(t)}+\sum_{j\in\mathcal{N}(i)}\mathbf{\Phi}_{\theta}^{(t+1)}(\mathbf{e}_{i,j})\mathbf{h}_{j}^{(t)}\Big{)}$

where $\mathcal{N}(i)$ denotes the set of neighbors of the $i$-th node of the
graph, and the trainable B-spline based kernel function
$\mathbf{\Phi}_{\theta}(\cdot)$ is conditioned on edge features.

The proposed model is implemented in PyTorch [24] using the PyTorch Geometric
library [10] and KeOPs [12]. All experiments were performed on a single
machine equipped with NVidia Tesla P100-PCIE GPU.

Table 1: Hits@1 (%) on the PASCAL-PF and PASCAL-VOC dataset

PASCAL-VOC | Aero | Bike | Bird | Boat | Bottle | Bus | Car | Cat | Chair | Cow | Table | Dog | Horse | M-Bike | Person | Plant | Sheep | Sofa | Train | TV | Mean
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
GMN | 31.1 | 46.2 | 58.2 | 45.9 | 70.6 | 76.5 | 61.2 | 61.7 | 35.5 | 53.7 | 58.9 | 57.5 | 56.9 | 49.3 | 34.1 | 77.5 | 57.1 | 53.6 | 83.2 | 88.6 | 57.9
PCA-GM | 40.9 | 55.0 | 65.8 | 47.9 | 76.9 | 77.9 | 63.5 | 67.4 | 33.7 | 66.5 | 63.6 | 61.3 | 58.9 | 62.8 | 44.9 | 77.5 | 67.4 | 57.5 | 86.7 | 90.9 | 63.8
DGMC $L=20$ Iso | 50.1 | 65.4 | 55.7 | 65.3 | 80.0 | 83.5 | 78.3 | 69.7 | 34.7 | 60.7 | 70.4 | 59.9 | 70.0 | 62.2 | 56.1 | 80.2 | 70.3 | 88.8 | 81.1 | 84.3 | 68.3
Ours $K=1$, $r=4$ Iso | 47.8 | 66.5 | 57.5 | 56.9 | 74.4 | 85.1 | 75.7 | 69.4 | 33.8 | 55.9 | 80.9 | 63.1 | 69.5 | 58.0 | 49.8 | 78.6 | 61.5 | 90.2 | 77.6 | 83.1 | 66.8
DGMC $L=20$ Aniso | 47.0 | 65.7 | 56.8 | 67.6 | 86.9 | 87.7 | 85.3 | 72.6 | 42.9 | 69.1 | 84.5 | 63.8 | 78.1 | 55.6 | 58.4 | 98.0 | 68.4 | 92.2 | 94.5 | 85.5 | 73.0
Ours $K=1$, $r=5$ Aniso | 52.0 | 71.5 | 56.0 | 72.1 | 84.6 | 85.5 | 83.9 | 75.7 | 48.0 | 70.5 | 88.1 | 66.0 | 85.1 | 64.5 | 58.2 | 96.2 | 67.5 | 93.5 | 93.6 | 88.7 | 75.0

PASCAL-PF | Aero | Bike | Bird | Boat | Bottle | Bus | Car | Cat | Chair | Cow | Table | Dog | Horse | M-Bike | Person | Plant | Sheep | Sofa | Train | TV | Mean
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
(Zhang & Lee, 2019) | 76.1 | 89.8 | 93.4 | 96.4 | 96.2 | 97.1 | 94.6 | 82.8 | 89.3 | 96.7 | 89.7 | 79.5 | 82.6 | 83.5 | 72.8 | 76.7 | 77.1 | 97.3 | 98.2 | 99.5 | 88.5
DGMC L = 20 | 81.1 | 92.0 | 94.7 | 100.0 | 99.3 | 99.3 | 98.9 | 97.3 | 99.4 | 93.4 | 100.0 | 99.1 | 86.3 | 86.2 | 87.7 | 100.0 | 100.0 | 100.0 | 100.0 | 99.3 | 95.7
Ours $K$ = 5, $r=4$ | 82.2 | 92.2 | 94.5 | 99.4 | 99.3 | 98.4 | 98.6 | 98.2 | 98.6 | 94.9 | 100.0 | 99.0 | 85.6 | 87.2 | 87.8 | 99.5 | 100.0 | 99.4 | 99.5 | 98.9 | 95.7

###  4.4 Results

Following previous work, we use Hits@$1$ as our evaluation metric, which
measures the proportion of correct matches in the top result. We also measure
the running time of the proposed model.

#### PASCAL-VOC.

Table 1 (top) shows the Hits@1 scores achieved by the different methods. We
observe that the anisotropic variants of the proposed model and of DGMC
outperform all the other models by wide margins. Interestingly, the
anisotropic variant of the proposed model is the best performing method since
it yields a relative increase of $2\%$ on the mean Hits@1 over DGMC. In terms
of running time, as shown in Table 2 (top), the variants of the proposed model
are faster than those of DGMC both during training and during inference. This
can be of high importance for real-time scenarios where generating predictions
in a small amount of time is of a critical nature.

Table 2: Average inference and training times on the Pascal-PF and PASCAL-VOC
datasets.

PASCAL-VOC | Inference Time | Training Time (per epoch)
---|---|---
DGMC $L=20$ Iso | 89ms | 29s
Ours $K=1$, $r=4$ Iso | 36ms | 17s
DGMC $L=20$ Aniso | 158ms | 60s
Ours $K=1$, $r=5$ Aniso | 94ms | 47s

PASCAL-PF | Inference Time | Training Time (per epoch)
---|---|---
DGMC $L=20$ | 40ms | 34s
Ours $K=5$, $r=4\qquad\,\,\,$ | 19ms | 36s

#### WILLOW-ObjectClass.

Tables 3 and 4 illustrate the performance (Hits@1 score) and running time of
the different models, respectively. The anisotropic variant of DGMC
outperforms all the other methods, while the anisotropic variant of the
proposed model is the second best method and is slightly outperformed by DGMC.
Surprisingly, the isotropic instance of our model achieves much smaller Hits@1
score than that of DGMC, while it is even outperformed by PCA-GM. In both
cases, the proposed model is trained in a smaller amount of time than DGMC,
while it is also more efficient at inference time. In Figure 1, we also
provide some examples where the proposed model has successfully find
correspondences between keypoints extracted from pairs of images.

Table 3: Hits@1 (%) on the WILLOW-ObjectClass dataset WILLOW | face | motorbike | car | duck | winebottle
---|---|---|---|---|---
GMN | 99.3 | 71.4 | 74.3 | 82.8 | 76.7
PCA-GM | 100.0 | 76.7 | 84.0 | 93.5 | 96.9
DGMC $L=20$ Iso | 100.00 ± 0.00 | 92.05 ± 3.24 | 90.28 ± 4.67 | 88.97 ± 3.49 | 97.14 ± 1.83
Ours $K=3$, $r=3$ Iso | 99.40 ± 0.46 | 87.78 ± 3.76 | 83.82 ± 4.09 | 85.20 ± 2.86 | 92.53 ± 1.69
DGMC $L=10$ Aniso | 100.00 ± 0.00 | 98.80 ± 1.58 | 96.53 ± 1.55 | 93.22 ± 3.77 | 99.87 ± 0.31
Ours $K=3$, $r=4$ Aniso | 99.85 ± 0.24 | 99.35 ± 0.71 | 97.88 ± 1.09 | 93.17 ± 2.67 | 99.15 ± 0.73
Table 4: Average inference and training times on the WILLOW-ObjectClass dataset. WILLOW | Inference Time | Training Time (per Epoch)
---|---|---
DGMC L = 20 Iso | 57ms | 13s
Ours $K=3$, $r=3$ Iso | 37ms | 8s
DGMC L = 10 Aniso | 64ms | 23s
Ours $K=3$, $r=4$ Aniso | 51ms | 16s

#### PASCAL-PF.

Table 1 (bottom) reports the Hits@1 score achieved by the proposed model and
the baselines. We observe that the proposed model matches the performance of
DGMC, while both models outperform the model proposed in [29]. This indicates
that our model is very effective even if visual features are not available.
Table 2 (bottom) illustrates the running time of the proposed model and that
of DGMC. The proposed model’s average inference time per image is much smaller
than that of DGMC, while training times of the two models are similar to each
other.

Figure 1: Examples of matching keypoints from the WILLOW-ObjectClass dataset.

##  5 Conclusion

In this paper, we presented a graph neural network for the problem of image
matching. Our starting point is a previous work on graph neural networks for
graph matching [9], which we modify to use not only two correspondence
matrices, but all the matrices that are produced during the consensus stage.
Furthermore, we also replace the iteration in the neighborhood consensus step
with a series of graph neural networks. Experiments on standard image matching
datasets showed that the proposed model is more efficient than the baseline
model while it maintains and in some cases improves prediction performance.

##  6 Acknowledgements

This research is co-financed by Greece and the European Union (European Social
Fund- ESF) through the Operational Programme «Human Resources Development,
Education and Lifelong Learning»in the context of the project “Reinforcement
of Postdoctoral Researchers - 2nd Cycle” (MIS-5033021), implemented by the
State Scholarships Foundation (IKY). This work was also supported by the
Wallenberg AI, Autonomous Systems and Software Program (WASP).

## References

  * [1] Bai, Y., Ding, H., Bian, S., Chen, T., Sun, Y., Wang, W.: Simgnn: A neural network approach to fast graph similarity computation. In: Proceedings of the 12th ACM International Conference on Web Search and Data Mining. pp. 384–392 (2019)
  * [2] Bourdev, L., Malik, J.: Poselets: Body part detectors trained using 3d human pose annotations. In: International Conference on Computer Vision (ICCV) (2009), http://www.eecs.berkeley.edu/~lbourdev/poselets
  * [3] Cho, M., Alahari, K., Ponce, J.: Learning graphs to match. In: 2013 IEEE International Conference on Computer Vision. pp. 25–32 (2013)
  * [4] Choy, C.B., Gwak, J.Y., Savarese, S., Chandraker, M.: Universal correspondence network. In: Proceedings of the 30th International Conference on Neural Information Processing Systems. pp. 2414–2422 (2016)
  * [5] Conte, D., Foggia, P., Sansone, C., Vento, M.: Thirty years of graph matching in pattern recognition. International Journal of Pattern Recognition and Artificial Intelligence 18(03), 265–298 (2004)
  * [6] Cuturi, M.: Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems 26, 2292–2300 (2013)
  * [7] Dalal, N., Triggs, B.: Histograms of oriented gradients for human detection. In: 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05). vol. 1, pp. 886–893 vol. 1 (2005)
  * [8] Everingham, M., Gool, L., Williams, C.K., Winn, J., Zisserman, A.: The pascal visual object classes (voc) challenge. Int. J. Comput. Vision 88(2), 303–338 (Jun 2010), https://doi.org/10.1007/s11263-009-0275-4
  * [9] Fey, M., Lenssen, J.E., Morris, C., Masci, J., Kriege, N.M.: Deep graph matching consensus. In: International Conference on Learning Representations (ICLR) (2020)
  * [10] Fey, M., Lenssen, J.E.: Fast graph representation learning with PyTorch Geometric. In: ICLR Workshop on Representation Learning on Graphs and Manifolds (2019)
  * [11] Fey, M., Lenssen, J.E., Weichert, F., Müller, H.: SplineCNN: Fast geometric deep learning with continuous B-spline kernels. In: IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2018)
  * [12] Feydy, J., Glaunès, J., Charlier, B., Bronstein, M.: Fast geometric learning with symbolic matrices. Advances in Neural Information Processing Systems 33 (2020)
  * [13] Gilmer, J., Schoenholz, S.S., Riley, P.F., Vinyals, O., Dahl, G.E.: Neural message passing for quantum chemistry. In: Precup, D., Teh, Y.W. (eds.) Proceedings of the 34th International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 70, pp. 1263–1272. PMLR (06–11 Aug 2017)
  * [14] Ham, B., Cho, M., Schmid, C., Ponce, J.: Proposal flow. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (2016)
  * [15] Heimann, M., Shen, H., Safavi, T., Koutra, D.: Regal: Representation learning-based graph alignment. In: Proceedings of the 27th ACM international conference on information and knowledge management. pp. 117–126 (2018)
  * [16] Kim, S., Min, D., Ham, B., Jeon, S., Lin, S., Sohn, K.: Fcss: Fully convolutional self-similarity for dense semantic correspondence. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), IEEE (2017)
  * [17] Knopp, P., Sinkhorn, R.: Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics 21(2), 343 – 348 (1967)
  * [18] Lee, J., Kim, D., Ponce, J., Ham, B.: Sfnet: Learning object-aware semantic correspondence. In: 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 2273–2282 (2019)
  * [19] Liu, C., Yuen, J., Torralba, A.: Sift flow: Dense correspondence across scenes and its applications. IEEE Transactions on Pattern Analysis and Machine Intelligence 33(5), 978–994 (2011)
  * [20] Loiola, E.M., de Abreu, N.M.M., Boaventura-Netto, P.O., Hahn, P., Querido, T.: A survey for the quadratic assignment problem. European Journal of Operational Research 176(2), 657–690 (2007), https://www.sciencedirect.com/science/article/pii/S0377221705008337
  * [21] Long, J., Zhang, N., Darrell, T.: Do convnets learn correspondence? In: Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 1. p. 1601–1609. NIPS’14, MIT Press, Cambridge, MA, USA (2014)
  * [22] Lowe, D.: Object recognition from local scale-invariant features. In: Proceedings of the Seventh IEEE International Conference on Computer Vision. vol. 2, pp. 1150–1157 vol.2 (1999)
  * [23] Ma, J., Jiang, X., Fan, A., Jiang, J., Yan, J.: Image matching from handcrafted to deep features: A survey. International Journal of Computer Vision 129(1), 23–79 (2021)
  * [24] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., Chintala, S.: Pytorch: An imperative style, high-performance deep learning library. In: Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems 32, pp. 8024–8035. Curran Associates, Inc. (2019), http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * [25] Sarlin, P.E., DeTone, D., Malisiewicz, T., Rabinovich, A.: SuperGlue: Learning feature matching with graph neural networks. In: CVPR (2020), https://arxiv.org/abs/1911.11763
  * [26] Wang, R., Yan, J., Yang, X.: Learning combinatorial embedding networks for deep graph matching. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 3056–3065 (2019)
  * [27] Xu, H., Luo, D., Zha, H., Duke, L.C.: Gromov-Wasserstein learning for graph matching and node embedding. In: Chaudhuri, K., Salakhutdinov, R. (eds.) Proceedings of the 36th International Conference on Machine Learning. Proceedings of Machine Learning Research, vol. 97, pp. 6932–6941. PMLR (09–15 Jun 2019)
  * [28] Zanfir, A., Sminchisescu, C.: Deep learning of graph matching. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 2684–2693 (2018)
  * [29] Zhang, Z., Lee, W.S.: Deep graphical feature learning for the feature matching problem. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 5087–5096 (2019)