# Global Context with Discrete Diffusion in Vector Quantised Modelling
for Image Generation

Minghui Hu1 Yujie Wang2 Tat-Jen Cham 1 Jianfei Yang 1 P.N.Suganthan1 1Nanyang
Technological University  2Sensetime Research<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The integration of Vector Quantised Variational AutoEncoder (VQ-VAE) with
autoregressive models as generation part has yielded high-quality results on
image generation. However, the autoregressive models will strictly follow the
progressive scanning order during the sampling phase. This leads the existing
VQ series models to hardly escape the trap of lacking global information.
Denoising Diffusion Probabilistic Models (DDPM) in the continuous domain have
shown a capability to capture the global context, while generating high-
quality images. In the discrete state space, some works have demonstrated the
potential to perform text generation and low resolution image generation. We
show that with the help of a content-rich discrete visual codebook from VQ-
VAE, the discrete diffusion model can also generate high fidelity images with
global context, which compensates for the deficiency of the classical
autoregressive model along pixel space. Meanwhile, the integration of the
discrete VAE with the diffusion model resolves the drawback of conventional
autoregressive models being oversized, and the diffusion model which demands
excessive time in the sampling process when generating images. It is found
that the quality of the generated images is heavily dependent on the discrete
visual codebook. Extensive experiments demonstrate that the proposed Vector
Quantised Discrete Diffusion Model (VQ-DDM) is able to achieve comparable
performance to top-tier methods with low complexity. It also demonstrates
outstanding advantages over other vectors quantised with autoregressive models
in terms of image inpainting tasks without additional training.

##  1 Introduction

Vector Quantised Variational AutoEncoder (VQ-VAE) [34] is a popular method
developed to compress images into discrete representations for the generation.
Typically, after the compression and discretization representation by the
convolutional network, an autoregressive model is used to model and sample in
the discrete latent space, including PixelCNN family [23, 35, 5], transformers
family [25, 4], etc. However, in addition to the disadvantage of the huge
number of model parameters, these autoregressive models can only make
predictions based on the observed pixels (left upper part of the target pixel)
due to the inductive bias caused by the strict adherence to the progressive
scan order [16, 3]. If the conditional information is located at the end of
the autoregressive sequence, it is difficult for the model to obtain relevant
information.

Figure 1: FID v.s. Operations and Parameters. The size of the blobs is
proportional to the number of network parameters, the X-axis indicates FLOPs
on a log scale and the Y-axis is the FID score.

A recent alternative generative model is the Denoising Diffusion Model, which
can effectively mitigate the lack of global information [30, 10], also
achieving comparable or state-of-the-art performance in text [12, 1], image
[6] and speech generation [20] tasks. Diffusion models are parameterized
Markov chains trained to translate simple distributions to more sophisticated
target data distributions in a finite set of steps. Typically the Markov chain
begins with an isotropic Gaussian distribution in continuous state space, with
the transitions of the chain for reversing a diffusion process that gradually
adds Gaussian noise to source images. In the inverse process, as the current
step is based on the global information of the previous step in the chain,
this endows the diffusion model with the ability to capture the global
information.

However, the diffusion model has a non-negligible disadvantage in that the
time and computational effort involved in generating the images are enormous.
The main reason is that the reverse process typically contains thousands of
steps. Although we do not need to iterate through all the steps when training,
all these steps are still required when generating a sample, which is much
slower compared to GANs and even autoregressive models.

Some recent works [31, 22] have attempted addressing these issues by
decreasing the sampling steps, but the computation cost is still high as each
step of the reverse process generates a full-resolution image.

In this work, we propose the Vector Quantized Discrete Diffusion Model (VQ-
DDM), a versatile framework for image generation consisting of a discrete
variational autoencoder and a discrete diffusion model. VQ-DDM consists of two
stages: (1) learning an abundant and efficient discrete representation of
images, (2) fitting the prior distribution of such latent visual codes via
discrete diffusion model.

VQ-DDM substantially reduces the computational resources and required time to
generate high-resolution images by using a discrete scheme. Then the common
problem of the lack of global content and overly large number of parameters of
the autoregressive model is solved by fitting a latent variable prior using
the discrete diffusion model. Finally, since a bias of codebook will limit
generation quality, while model size is also dependent on the number of
categories, we propose a re-build and fine-tune(ReFiT) strategy to construct a
codebook with higher utilization, which will also reduce the number of
parameters in our model.

In summary, our key contributions include the following:

  * •

VQ-DDM fits the prior over discrete latent codes with a discrete diffusion
model. The use of diffusion model allows the generative models consider the
global information instead of only focusing on partially seen context to avoid
sequential bias.

  * •

We propose a ReFiT approach to improve the utilisation of latent
representations in the visual codebook, which can increase the code usage of
VQ-GAN from $31.85\%$ to $97.07\%$, while the FID between reconstruction image
and original training image is reduced from $10.18$ to $5.64$ on CelebA-HQ
$256\times 256$.

  * •

VQ-DDM is highly efficient for the both number of parameters and generation
speed. As shown in Figure 1, using only 120M parameters, it outperforms VQ-
VAE-2 with around 10B parameters and is comparable with VQ-GAN with 1B
parameters in image generation tasks in terms of image quality. It is also 10
$\sim$ 100 times faster than other diffusion models for image generation [31,
10].

Figure 2: The proposed VQ-DDM pipeline contains 2 stages: (1) Compress the
image into discrete variables via discrete VAE. (2) Fit a prior distribution
over discrete coding by a diffusion model. Black squares in the diffusion
diagram illustrate states when the underlying distributions are uninformative,
but which become progressively more specific during the reverse process. The
bar chart at the bottom of the image represents the probability of a
particular discrete variable being sampled.

##  2 Preliminaries

###  2.1 Diffusion Models in continuous state space

Given data $\mathbf{x}_{0}$ from a data distribution $q(\mathbf{x}_{0})$, the
diffusion model consists of two processes: the diffusion process and the
reverse process [30, 10].

The diffusion process progressively destroys the data $\mathbf{x}_{0}$ into
$\mathbf{x}_{T}$ over $T$ steps, via a fixed Markov chain that gradually
introduces Gaussian noise to the data according to a variance schedule
$\beta_{1:T}\in(0,1]^{T}$ as follows:

$q(\mathbf{x}_{1:T}|\mathbf{x}_{0})=\prod_{t=1}^{T}q(\mathbf{x}_{t}|\mathbf{x}_{t-1}),$
(1)
$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}).$
(2)

With an adequate number of steps $T$ and a suitable variance schedule $\beta$,
$p(\mathbf{x}_{T})$ becomes an isotropic Gaussian distribution.

The reverse process is defined as a Markov chain parameterized by $\theta$,
which is used to restore the data from the noise:

$p_{\theta}(\mathbf{x}_{0:T})=p(\mathbf{x}_{T})\prod_{t=1}^{T}p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}),$
(3)
$p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})=\mathcal{N}(\mathbf{x}_{t-1};\mu_{\theta}(\mathbf{x}_{t},t),\Sigma_{\theta}(\mathbf{x}_{t},t)).$
(4)

The objective of training is to find the best $\theta$ to fit the data
distribution $q(\mathbf{x}_{0})$ by optimizing the variational lower bound
(VLB) [19]

$\begin{split}\mathbb{E}_{q(\mathbf{x}_{0})}&[\log
p_{\theta}(\mathbf{x}_{0})]\\\
=&\mathbb{E}_{q(\mathbf{x}_{0})}\log\mathbb{E}_{q(\mathbf{x}_{1:T}|\mathbf{x}_{0})}\left[\frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_{0})}\right]\\\
\geq&\mathbb{E}_{q(\mathbf{x}_{0:T})}\left[\log\frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_{0})}\right]=:L_{\mathrm{vlb}}.\end{split}$
(5)

Ho _et al_. [10] revealed that the variational lower bound in Eq. 5 can be
calculated with closed form expressions instead of Monte Carlo estimates as
the diffusion process posteriors and marginals are Gaussian, which allows
sampling $\mathbf{x}_{t}$ at an arbitrary step $t$ with
$\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{s=0}^{t}\alpha_{s}$ and
$\tilde{\beta_{t}}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}$:

$q(\mathbf{x}_{t}|\mathbf{x}_{0})=\mathcal{N}(\mathbf{x}_{t}|\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0},(1-\bar{\alpha}_{t})\mathbf{I}),$
(6)

$\begin{split}L_{\mathrm{vlb}}=\mathbb{E}_{q(\mathbf{x}_{0})}&[D_{\mathrm{KL}}(q(\mathbf{x}_{T}|\mathbf{x}_{0})||p(\mathbf{x}_{T}))-\log
p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})\\\
&+\sum_{t=2}^{T}D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0})||p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}))].\end{split}$
(7)

Thus the reverse process can be parameterized by neural networks
$\epsilon_{\theta}$ and $\upsilon_{\theta}$, which can be defined as:

$\mu_{\theta}(\mathbf{x}_{t},t)=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(\mathbf{x}_{t},t)\right),$
(8)
$\begin{split}\Sigma_{\theta}(\mathbf{x}_{t},t)=\exp(\upsilon_{\theta}&(\mathbf{x}_{t},t)\log\beta_{t}\\\
&+(1-\upsilon_{\theta}(\mathbf{x}_{t},t))\log\tilde{\beta_{t}}).\end{split}$
(9)

Using a modified variant of the VLB loss as a simple loss function will offer
better results in the case of fixed $\Sigma_{\theta}$ [10]:

$L_{\mathrm{simple}}=\mathbb{E}_{t,\mathbf{x}_{0},\epsilon}\left[||\epsilon-\epsilon_{\theta}(\mathbf{x}_{t},t)||^{2}\right],$
(10)

which is a reweighted version resembling denoising score matching over
multiple noise scales indexed by $t$ [32].

Nichol _et al_. [22] used an additional $L_{\mathrm{vlb}}$ to the simple loss
for guiding a learned $\Sigma_{\theta}(\mathbf{x}_{t},t)$, while keeping the
$\mu_{\theta}(\mathbf{x}_{t},t)$ still the dominant component of the total
loss:

$L_{\mathrm{hybrid}}=L_{\mathrm{simple}}+\lambda L_{\mathrm{vlb}}.$ (11)

###  2.2 Discrete Representation of Images

van den Oord _et al_. [34] presented a discrete variational autoencoder with a
categorical distribution as the latent prior, which is able to map the images
into a sequence of discrete latent variables by an encoder and reconstruct the
image according to those variables with a decoder. Formally, given a codebook
$\mathbb{Z}\in\mathbb{R}^{K\times d}$, where $K$ represents the capacity of
latent variables in the codebook and $d$ is the dimension of each latent
variable, after compressing the high dimension input data
$\textbf{x}\in\mathbb{R}^{c\times H\times W}$ into latent vectors
$\textbf{h}\in\mathbb{R}^{h\times w\times d}$ by an encoder $E$, z is the
quantised h, which substitutes the vectors $h_{i,j}\in\textbf{h}$ by the
nearest neighbor $z_{k}\in\mathbb{Z}$. The decoder $D$ is trained to
reconstruct the data from the quantised encoding $\textbf{z}_{q}$:

$\textbf{z}=\mathrm{Quantize}(\textbf{h}):=\mathrm{arg\
min}_{k}||h_{i,j}-z_{k}||,$ (12)
$\hat{\textbf{x}}=D(\textbf{z})=D(\mathrm{Quantize}(E(\textbf{x}))).$ (13)

As $\mathrm{Quantize}(\cdot)$ has a non-differentiable operation $\mathrm{arg\
min}$, the straight-through gradient estimator is used for back-propagating
the reconstruction error from decoder to encoder. The whole model can be
trained in an end-to-end manner by minimizing the following function:

$L=||\textbf{x}-\hat{\textbf{x}}||^{2}+||sg[E(\textbf{x})]-\textbf{z}||+\beta||sg[\textbf{z}]-E(\textbf{x})||,$
(14)

where $sg[\cdot]$ denotes stop gradient and broadly the three terms are
reconstruction loss, codebook loss and commitment loss, respectively. VQ-GAN
[8] extends VQ-VAE [34] in multiple ways. It substitutes the L1 or L2 loss of
the original VQ-VAE with a perceptual loss [41], and adds an additional
discriminator to distinguish between real and generated patches [43].

The codebook update of the discrete variational autoencoder is intrinsically a
dictionary learning process. Its objective uses L2 loss to narrow the gap
between the codes $\mathbb{Z}_{t}\in\mathbb{R}^{K_{t}\times d}$ and the
encoder output $\textbf{h}\in\mathbb{R}^{h\times w\times d}$ [34]. In other
words, the codebook training is like $k$-means clustering, where cluster
centers are the discrete latent codes. However, since the volume of the
codebook space is dimensionless and h is updated each iteration, the discrete
codes $\mathbb{Z}$ typically do not follow the encoder training quickly
enough. Only a few codes get updated during training, with most unused after
initialization.

##  3 Methods

Our goal is to leverage the powerful generative capability of the diffusion
model to perform high fidelity image generation tasks with a low number of
parameters.

Our proposed method, VQ-DDM, is capable of generating high fidelity images
with a relatively small number of parameters and FLOPs, as summarized in
Figure 2. Our solution starts by compressing the image into discrete variables
via the discrete VAE and then constructs a powerful model to fit the joint
distribution over the discrete codes by a diffusion model. During diffusion
training, the darker coloured parts in Figure 2 represent noise introduced by
uniform resampling. When the last moment is reached, the latent codes have
been completely corrupted into noise. In the sampling phase, the latent codes
are drawn from an uniform categorical distribution at first, and then
resampled by performing reverse process $T$ steps to get the target latent
codes. Eventually, target latent codes are pushed into the decoder to generate
the image.

###  3.1 Discrete Diffusion Model

Assume the discretization is done with $K$ categories, i.e.
$z_{t}\in\\{1,\dots,K\\}$, with the one-hot vector representation given by
$\textbf{z}_{t}\in\\{0,1\\}^{K}$. The corresponding probability distribution
is expressed by $\textbf{z}_{t}^{\mathrm{logits}}$ in logits. We formulate the
discrete diffusion process as

$q(\textbf{z}_{t}|\textbf{z}_{t-1})=\mathrm{Cat}(\textbf{z}_{t};\textbf{z}_{t-1}^{\mathrm{logits}}\mathbf{Q}_{t}),$
(15)

where $\mathrm{Cat}(\textbf{x}|\textbf{p})$ is the categorical distribution
parameterized by p, while $\mathbf{Q}_{t}$ is the process transition matrix.
In our method, $\mathbf{Q}_{t}=(1-\beta_{t})\textbf{I}+\beta_{t}/K$, which
means $\textbf{z}_{t}$ has $1-\beta_{t}$ probability to keep the state from
last timestep and $\beta_{t}$ chance to resample from a uniform categorical
distribution. Formally, it can be written as

$q(\textbf{z}_{t}|\textbf{z}_{t-1})=\mathrm{Cat}(\textbf{z}_{t};(1-\beta_{t})\textbf{z}_{t-1}^{\mathrm{logits}}+\beta_{t}/K).$
(16)

It is straightforward to get $\textbf{z}_{t}$ from $\textbf{z}_{0}$ under the
schedule $\beta_{t}$ with $\alpha_{t}=1-\beta_{t}$,
$\bar{\alpha}_{t}=\prod_{s=0}^{t}\alpha_{s}$:

$q(\textbf{z}_{t}|\textbf{z}_{0})=\mathrm{Cat}(\textbf{z}_{t};\bar{\alpha}_{t}\textbf{z}_{0}+(1-\bar{\alpha}_{t})/K)$
(17) $or\quad
q(\textbf{z}_{t}|\textbf{z}_{0})=\mathrm{Cat}(\textbf{z}_{t};\textbf{z}_{0}\bar{\mathbf{Q}}_{t});\
\bar{\mathbf{Q}}_{t}=\prod_{s=0}^{t}\mathbf{Q}_{s}.$ (18)

We use the same cosine noise schedule as [22, 12] because our discrete model
is also established on the latent codes with a small $16\times 16$ resolution.
Mathematically, it can be expressed in the case of $\bar{\alpha}$ by

$\bar{\alpha}=\frac{f(t)}{f(0)},\quad
f(t)=\mathrm{cos}\left(\frac{t/T+s}{1+s}\times\frac{\pi}{2}\right)^{2}.$ (19)

By applying Bayes’ rule, we can compute the posterior
$q(\textbf{z}_{t-1}|\textbf{z}_{t},\textbf{z}_{0})$ as:

$\begin{split}q(\textbf{z}_{t-1}|\textbf{z}_{t},\textbf{z}_{0})&=\mathrm{Cat}\left(\textbf{z}_{t};\frac{\textbf{z}_{t}^{\mathrm{logits}}\mathbf{Q}_{t}^{\top}\odot\textbf{z}_{0}\bar{\mathbf{Q}}_{t-1}}{\textbf{z}_{0}\bar{\mathbf{Q}}_{t}{\textbf{z}_{t}^{\mathrm{logits}}}^{\top}}\right)\\\
=\mathrm{Cat}&(\textbf{z}_{t};\
\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})/\sum_{k=1}^{K}\theta_{k}(z_{t,k},z_{0,k})),\\\
\end{split}$ (20)
$\begin{split}\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})=[\alpha_{t}\textbf{z}_{t}^{\mathrm{logits}}+&(1-\alpha_{t})/K]\\\
&\odot[\bar{\alpha}_{t-1}\textbf{z}_{0}+(1-\bar{\alpha}_{t-1})/K].\end{split}$
(21)

It is worth noting that
$\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})/\sum_{k=1}^{K}\theta_{k}(z_{t,k},z_{0,k})$
is the normalized version of
$\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})$, and we use
$\mathrm{N}[\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})]$ to denote
$\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})/\sum_{k=1}^{K}\theta_{k}(z_{t,k},z_{0,k})$
below.

Hoogeboom _et al_. [12] predicted $\hat{\textbf{z}}_{0}$ from $\textbf{z}_{t}$
with a neural network $\mu(\textbf{z}_{t},t)$, instead of directly predicting
$p_{\theta}(\textbf{z}_{t-1}|\textbf{z}_{t})$. Thus the reverse process can be
parameterized by the probability vector from
$q(\textbf{z}_{t-1}|\textbf{z}_{t},\hat{\textbf{z}}_{0})$. Generally, the
reverse process $p_{\theta}(\textbf{z}_{t-1}|\textbf{z}_{t})$ can be expressed
by

$\begin{split}p_{\theta}(\textbf{z}_{0}|\textbf{z}_{1})&=\mathrm{Cat}(\textbf{z}_{0}|\hat{\textbf{z}}_{0}),\\\
p_{\theta}(\textbf{z}_{t-1}|\textbf{z}_{t})&=\mathrm{Cat}(\textbf{z}_{t}|\
\mathrm{N}[\boldsymbol{\theta}(\textbf{z}_{t},\hat{\textbf{z}}_{0})]).\end{split}$
(22)

Inspired by [13, 21], we use a neural network $\mu(\mathbf{Z}_{t},t)$ to learn
and predict the a noise $n_{t}$ and obtain the logits of
$\hat{\mathbf{z}}_{0}$ from

$\hat{\mathbf{z}}_{0}=\mu(\mathbf{Z}_{t},t)+\mathbf{Z}_{t}.$ (23)

It is worth noting that the neural network $\mu(\cdot)$ is based on the
$\mathbf{Z}_{t}\in\mathbb{N}^{h\times w}$, where all the discrete
representation $\mathbf{z}_{t}$ of the image are combined. The final noise
prior $\mathbf{Z}_{T}$ is uninformative, and it is possible to separably
sample from each axis during inference. However, the reverse process is
jointly informed and evolves towards a highly coupled $\mathbf{Z}_{0}$. We do
not define a specific joint prior for $\mathbf{z}_{t}$, but encode the joint
relationship into the learned reverse process. This is implicitly done in the
continuous domain diffusion. As $\mathbf{z}_{t-1}$ is based on the whole
previous representation $\mathbf{z}_{t}$, the reverse process can sample the
whole discrete code map directly while capturing the global information.

The loss function used is the VLB from Eq. 7, where the summed KL divergence
for $T>2$ is given by

$\begin{split}\mathrm{KL}(q(\textbf{z}_{t-1}|\textbf{z}_{t},\textbf{z}_{0})||p_{\theta}(\textbf{z}_{t-1}|\textbf{z}_{t}))&=\\\
\sum_{k}\mathrm{N}[\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})]&\times\log\frac{\mathrm{N}[\boldsymbol{\theta}(\textbf{z}_{t},\textbf{z}_{0})]}{\mathrm{N}[\boldsymbol{\theta}(\textbf{z}_{t},\hat{\textbf{z}}_{0})]}.\end{split}$
(24)

###  3.2 Re-build and Fine-tune Strategy

Our discrete diffusion model is based on the latent representation of the
discrete VAE codebook $\mathbb{Z}$. However, the codebooks with rich content
are normally large, with some even reaching $K=16384$. This makes it highly
unwieldy for our discrete diffusion model, as the transition matrices of
discrete diffusion models have a quadratic level of growth to the number of
classes $K$, _e.g_. $O(K^{2}T)$ [1].

To reduce the categories used for our diffusion model, we proposed a Re-build
and Fine-tune (ReFit) strategy to decrease the size $K$ of codebook
$\mathbb{Z}$ and boost the reconstruction performance based on a well-trained
discrete VAEs trained by the straight-through method.

From Eq. 14, we can find the second term and the third term are related to the
codebook, but only the second term is involved in the update of the codebook.
$||sg[E(\textbf{x})]-\textbf{z}||$ reveals that only a few selected codes, the
same number as the features from $E(\textbf{x})$, are engaged in the update
per iteration. Most of the codes are not updated or used after initialization,
and the update of the codebook can lapse into a local optimum.

We introduce a re-build and fine-tune strategy to avoid the waste of codebook
capacity. With the trained encoder, we reconstruct the codebook so that all
codes in the codebook have the opportunity to be selected. This will greatly
increase the usage of the codebook. Suppose we desire to obtain a discrete VAE
having a codebook with $\mathbb{Z}_{t}$ based on a trained discrete VAE with
an encoder $E_{s}$ and a decoder $D_{s}$. We first encode each image
$\textbf{x}\in\mathbb{R}^{c\times H\times W}$ to latent features h, or loosely
speaking, each image gives us $h\times w$ features with $d$ dimension. Next we
sample $P$ features uniformly from the entire set of features found in
training images, where $P$ is the sampling number and far larger than the
desired codebook capacity $K_{t}$. This ensures that the re-build codebook is
composed of valid latent codes. Since the process of codebook training is
basically the process of finding cluster centres, we directly employ k-means
with AFK-MC2 [2] on the sampled $P$ features and utilize the centres to re-
build the codebook $\mathbb{Z}_{t}$. We then replace the original codebook
with the re-build $\mathbb{Z}_{t}$ and fine-tune it on top of the well-trained
discrete VAE.

##  4 Experiments and Analysis

###  4.1 Datasets and Implementation Details

We show the effectiveness of the proposed VQ-DDM on CelebA-HQ [14] and LSUN-
Church [40] datasets and verify the proposed Re-build and Fine-tune strategy
on CelebA-HQ and ImageNet datasets. The details of the dataset are given in
the Appendix.

The discrete VAE follows the same training strategy as VQ-GAN[8]. All training
images are processed to $256\times 256$, and the compress ratio is set to
$16$, which means the latent vector $\textbf{z}\in\mathbb{R}^{1\times 16\times
16}$. When conducting Rebuild and Fine-tune, the sampling number $P$ is set to
$20k$ for LSUN and CelebA. For the more content-rich case, we tried a larger P
value $50k$ for ImageNet. In practical experiments, we sample $P$ images with
replacement uniformly from the whole training data and obtained corresponding
latent features. For each feature map, we make another uniform sampling over
the feature map size $16\times 16$ to get the desired features. In the fine-
tuning phase, we freeze the encoder and set the learning rate of the decoder
to $1e$-$6$ and the learning rate of the discriminator to $2e$-$6$ with 8
instances per batch.

With regard to the diffusion model, the network for estimating $n_{t}$ has the
same structure as [10], which is a U-Net [27] with self-attention [36]. The
detailed settings of hyperparameters are provided in the Appendix. We set
timestep $T=4000$ in our experiments and the noise schedule is the same as
[22]

###  4.2 Codebook Quality

A large codebook dramatically increases the cost of DDM. To reduce the cost to
an acceptable scale, we proposed a resample and fine-tune strategy to compress
the size of the codebook, while maintaining quality. To demonstrate the
effectiveness of the proposed strategy, we compare the codebook usage and FID
of reconstructed images of our method to VQ-GAN[8], VQ-VAE-2[26] and
DALL-E[25].

In this experiment, we compressed the images from $3\times 256\times 256$ to
$1\times 16\times 16$ with two different codebook capacities
$K=\\{512,1024\\}$. We also proposed an indicator to measure the usage rate of
the codebook, which is the number of discrete features that have appeared in
the test set or training set divided by the codebook capacity.

The quantitative comparison results are shown in Table 1 while the reconstruct
images are demonstrated in Figs. 3 & 4. Reducing the codebook capacity from
1024 to 512 only brings $\sim 0.1$ decline in CelebA and $\sim 1$ in ImageNet.
As seen in Figure 4, the reconstructed images (c,d) after ReFiT strategy are
richer in colour and more realistic in expression than the reconstructions
from VQ-GAN (b). The codebook usage of our method has improved significantly
compared to other methods, nearly 3x high than the second best. Our method
also achieves the equivalent reconstruction quality at the same compression
rate and with 32$\times$ lower capacity $K$ of codebook $\mathbb{Z}$.

For VQ-GAN with capacity $16384$, although it only has $976$ effective codes,
which is smaller than $1024$ in our ReFiT method when $P=20k$, it achieves a
lower FID in reconstructed images vs validation images. One possible reason is
that the value of $P$ is not large enough to cover some infrequent
combinations of features during the re-build phase. As the results in Table 1,
after we increase the sampling number $P$ from $20k$ to $100k$, we observe
that increasing the value of $P$ achieved higher performance.

Model | Latent Size | Capacity | Usage of $\mathbb{Z}$ | FID $\downarrow$
---|---|---|---|---
|  |  | CelebA | ImageNet | CelebA | ImageNet
VQ-VAE-2 | Cascade | 512 |  $\sim$65% | - | - |  $\sim$10
DALL-E | 32x32 | 8192 | - | - | - | 32.01
VQ-GAN | 16x16 | 16384 | - | 5.96% | - | 4.98
VQ-GAN | 16x16 | 1024 | 31.85% | 33.67% | 10.18 | 7.94
ours ($P=100k$) | 16x16 | 1024 | - | 100% | - | 4.98
ours ($P=20k$) | 16x16 | 1024 | 97.07% | 100% | 5.59 | 5.99
ours ($P=20k$) | 16x16 | 512 | 93.06% | 100% | 5.64 | 6.95

  * 1

All methods are trained straight-through, except DALL-E with Gumbel-Softmax
[25].

  * 2

CelebA-HQ at $256$$\times$$256$. Reported FID is between 30$k$ reconstructed
data vs training data.

  * 3

Reported FID is between 50$k$ reconstructed data vs validation data

Table 1: FID between reconstructed images and original images on CelebA-HQ and
ImageNet  Figure 3: Reconstruction images $384\times 384$ from ImageNet based
VQ-GAN and ReFiT

(a) Source

(b) VQ-GAN

(c) ReFiT K=1024

(d) ReFiT K=512

Figure 4: Reconstruction images of CelebA HQ $256\times 256$ from VQ-GAN and
ReFiT.

###  4.3 Generation Quality

We evaluate the performance of VQ-DDM for the unconditional image generation
on CelebA-HQ $256\times 256$. Specifically, we evaluated the performance of
our approach in terms of FID and compared it with various likelihood-based
based methods including GLOW [17], NVAE [33], VAEBM [39], DC-VAE [24], VQ-GAN
[8] and likelihood-free method, e.g., PGGAN [14]. We also conducted an
experiment on LSUN-Church.

In CelebA-HQ experiments, the discrete diffusion model was trained with
$K=512$ and $K=1024$ codebooks respectively. We also report the different FID
from $T=2$ to $T=4000$ with corresponding time consumption in Figure 6.
Regarding the generation speed, it took about 1000 hours to generate $50k$
$256\times 256$ images using DDPM with 1000 steps on a NVIDIA 2080Ti GPU, 100
hours for DDIM with 100 steps [31], and around 10 hours for our VQ-DDM with
1000 steps.

Figure 5: Steps and corresponding FID during the sampling. The text
annotations are hours to sample 50k latent feature maps on 1 NVIDIA 2080Ti GPU
Figure 6: Hours to sampling 50k latent codes by VQ-DDM and generating 50k
images with VQ-DDM and DDPM

(a) Samples $(256\times 256)$ from a VQ-DDM model trained on CelebA HQ.
FID=$13.2$ (b) Samples $(256\times 256)$ from a VQ-DDM model trained on LSUN-
Church. FID=$16.9$

Figure 7: Samples from VQ-DDM models.

Table 2 shows the main results on VQ-DDM along with other established models.
Although VQ-DDM is also a likelihood-based method, the training phase relies
on the negative log-likehood (NLL) of discrete hidden variables, so we do not
compare the NLL between our method and the other methods. The training NLL is
around $1.258$ and test NLL is $1.286$ while the FID is $13.2$. Fig. 7(a)
shows the generated samples from VQ-DDM trained on the CelebA-HQ.

For LSUN-Church, the codebook capacity $K$ is set to $1024$, while the other
parameters are set exactly the same. The training NLL is $1.803$ and the test
NLL is $1.756$ while the FID between the generated images and the training set
is $16.9$. Some samples are shown in Fig. 7(b).

After utilizing ReFiT, the generation quality of the model is significantly
improved, which implies a decent codebook can have a significant impact on the
subsequent generative phase. Within a certain range, the larger the codebook
capacity leads to a better performance. However, excessive number of codebook
entries will cause the model collapse [12].

###  4.4 Image Inpainting

Autoregressive models have recently demonstrated superior performance in the
image inpainting tasks [4, 8]. However, one limitation of this approach is
that if the important context is found at the end of the autoregressive
series, the models will not be able to correctly complete the images. As
mentioned in Sec. 3.1, the diffusion model will directly sample the full
latent code map, with sampling steps based on the _full_ discrete map of the
previous step. Hence it can significantly improve inpainting as it does not
depend on context sequencing.

We perform the mask diffusion and reverse process in the discrete latent
space. After encoding the masked image $x_{0}\sim q(\textbf{x}_{0})$ to
discrete representations $z_{0}\sim q(\textbf{z}_{0})$, we diffuse
$\textbf{z}_{0}$ with $t$ steps to $\tilde{\textbf{z}}_{t}\sim
q(\textbf{z}_{t}|\textbf{z}_{0})$. Thus the last step with mask
$\tilde{\textbf{z}}_{T}^{m}$ can be demonstrated as
$\tilde{\textbf{z}}_{T}^{m}=(1-m)\times\tilde{\textbf{z}}_{T}+m\times\mathbb{C}$,
where $\mathbb{C}\sim\mathrm{Cat}(K,1/K)$ is the sample from a uniform
categorical distribution and $m\in\\{0,1\\}^{K}$ is the mask, $m=0$ means the
context there is masked and $m=1$ means that given the information there. In
the reverse process, $\textbf{z}_{T-1}$ can be sampled from
$p_{\theta}(\mathbf{z}_{T-1}|\tilde{\textbf{z}}_{T}^{m})$ at $t=T$, otherwise,
$\textbf{z}_{t-1}\sim p_{\theta}(\mathbf{z}_{t-1}|\textbf{z}_{t}^{m})$, and
the masked
$\textbf{z}_{t-1}^{m}=(1-m)\times\textbf{z}_{t-1}+m\times\tilde{\textbf{z}}_{t-1}$.

We compare our approach and another that exploits a transformer with a sliding
attention window as an autoregressive generative model [8]. The completions
are shown in Fig. 8, in the first row, the upper 62.5% (160 out of 256 in
latent space) of the input image is masked and the lower 37.5% (96 out of 256)
is retained, and in the second row, only a quarter of the image information in
the lower right corner is retained as input. We also tried masking in an
arbitrary position. In the third row, we masked the perimeter, leaving only a
quarter part in the middle. Since the reverse diffusion process captures the
global relationships, the image completions of our model performs much better.
Our method can make a consistent completions based on arbitrary contexts,
whereas the inpainting parts from transformer lack consistency. It is also
worth noting that our model requires no additional training in solving the
task of image inpainting.

Method | FID $\downarrow$ | Params | FLOPs
---|---|---|---
Likelihood-based |  |
GLOW [17] | 60.9 | 220 M | 540 G
NVAE [33] | 40.3 | 1.26 G | 185 G
ours ($K=1024$ w/o ReFiT) | 22.6 | 117 M | 1.06 G
VAEBM [39] | 20.4 | 127 M | 8.22 G
ours ($K=512$ w/ ReFiT) | 18.8 | 117 M | 1.04 G
DC-VAE [24] | 15.8 | - | -
ours ($K=1024$ w/ ReFiT) | 13.2 | 117 M | 1.06 G
DDIM(T=100) [31] | 10.9 | 114 M | 124 G
VQ-GAN + Transformer [8] | 10.2 | 802 M | 102 Ga
Likelihood-free |  |
PG-GAN [14] | 8.0 | 46.1 M | 14.1 G

  * a

VQ-GAN is an autoregressive model, and the number in the table is the
computation needed to generate the full size latent feature map. The FLOPs
needed to generate one discrete index out of 256 is 0.399 G.

Table 2: FID on CelebA HQ $256\times 256$ dataset. All the FLOPs in the table
only consider the generation stage or inference phase for one $256\times 256$
images.  Figure 8: Completions with the arbitrary masks.

##  5 Related Work

###  5.1 Vector Quantised Variational Autoencoders

VQ-VAE [34] leads a trend of discrete representation of images. The common
practice is to model the discrete representations using an autoregressive
model, e.g. PixelCNN [35, 5], transformers [8, 25, 25], etc. Some works had
attempted to fit the prior distribution of discrete latent variables using a
light non-autoregressive approach, like EM approach [28] and Markov chain with
self-organizing map [9], but yet they are struggling to fit a large scale of
data. Ho _et al_. [10] have also shown that the diffusion models can be
regarded as an autoregressive model along the time dimension, but in reality,
it is non-autoregressive along the pixel dimension.

A concurrent work [7] follow a similar pipeline which uses a diffusion model
on discrete latent variables, but the work uses parallel modeling of multiple
short Markov chains to achieve denoising.

###  5.2 Diffusion Models

Sohl-Dickstein _et al_. [30] presented a simple discrete diffusion model,
which diffused the target distribution into the independent binomial
distribution. Recently, Hoogeboom _et al_. [12] have extended the discrete
model from binomial to multinomial. Further, Austin _et al_. [1] proposed a
generalized discrete diffusion structure, which provides several choices for
the diffusion transition process.

In the continuous state space, there are some recent diffusion models that
surpassed the state-of-the-art in the image generation area. With the guidance
from the classifiers, Dhariwal _et al_. [6] enabled diffusion models called
ADM to generate images beyond BigGAN, which was previously one of the most
powerful generative models. In CDM [11], the authors performed the cascade
pipeline on the diffusion model to generate the image with ultra-high fidelity
and reach state-of-the-art on conditional ImageNet generation. In addition,
there have been several recent works that have attempted to use diffusion
models to modelling the latent variables of VAE [18, 37], while revealed the
connection among several diffusion models mentioned above.

##  6 Conclusion

In this paper, we introduce VQ-DDM, a high-fidelity image generation model
with a two-stage pipeline. In the first stage, we train a discrete VAE with a
well-utilized content-rich codebook. With the help of such an efficient
codebook, it is possible to generate high-quality images by a discrete
diffusion model with relatively tiny parameters in the second stage.
Simultaneously, benefiting from the discrete diffusion model, the sampling
process captures the global information and the image inpainting is no longer
affected by the location of the given context and mask. Meanwhile, in
comparison with other diffusion models, our approach further reduces the gap
in generation speed with respect to GAN. We believe that VQ-DDM can also be
utilized for audio, video and multimodal generation.

### Limitations

For a complete diffusion, we need a large number of steps, which will result
in a very fluctuating training process and limit the image generation quality.
Hence, our model may suffer from underperformance when exposed to the large
scale and complex datasets.

## References

  * [1] Jacob Austin, Daniel Johnson, Jonathan Ho, Danny Tarlow, and Rianne van den Berg.  Structured denoising diffusion models in discrete state-spaces.  arXiv preprint arXiv:2107.03006, 2021.
  * [2] Olivier Bachem, Mario Lucic, Hamed Hassani, and Andreas Krause.  Fast and provably good seedings for k-means.  Advances in neural information processing systems, 29:55–63, 2016\.
  * [3] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer.  Scheduled sampling for sequence prediction with recurrent neural networks.  arXiv preprint arXiv:1506.03099, 2015.
  * [4] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In International Conference on Machine Learning, pages 1691–1703. PMLR, 2020.
  * [5] Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel.  Pixelsnail: An improved autoregressive generative model.  In International Conference on Machine Learning, pages 864–872. PMLR, 2018.
  * [6] Prafulla Dhariwal and Alex Nichol.  Diffusion models beat gans on image synthesis.  arXiv e-prints, pages arXiv–2105, 2021.
  * [7] Patrick Esser, Robin Rombach, Andreas Blattmann, and Björn Ommer.  Imagebart: Bidirectional context with multinomial diffusion for autoregressive image synthesis.  arXiv preprint arXiv:2108.08827, 2021.
  * [8] Patrick Esser, Robin Rombach, and Bjorn Ommer.  Taming transformers for high-resolution image synthesis.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12873–12883, 2021.
  * [9] Vincent Fortuin, Matthias Hüser, Francesco Locatello, Heiko Strathmann, and Gunnar Rätsch.  Som-vae: Interpretable discrete representation learning on time series.  arXiv preprint arXiv:1806.02199, 2018.
  * [10] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  arXiv preprint arxiv:2006.11239, 2020.
  * [11] 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.
  * [12] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling.  Argmax flows and multinomial diffusion: Towards non-autoregressive language models.  arXiv preprint arXiv:2102.05379, 2021.
  * [13] Eric Jang, Shixiang Gu, and Ben Poole.  Categorical reparameterization with gumbel-softmax.  arXiv preprint arXiv:1611.01144, 2016.
  * [14] 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.
  * [15] 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.
  * [16] Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah.  Transformers in vision: A survey.  arXiv preprint arXiv:2101.01169, 2021.
  * [17] Diederik P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  arXiv preprint arXiv:1807.03039, 2018.
  * [18] Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  arXiv preprint arXiv:2107.00630, 2021.
  * [19] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * [20] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro.  Diffwave: A versatile diffusion model for audio synthesis.  In International Conference on Learning Representations, 2020.
  * [21] Chris J Maddison, Andriy Mnih, and Yee Whye Teh.  The concrete distribution: A continuous relaxation of discrete random variables.  arXiv preprint arXiv:1611.00712, 2016.
  * [22] Alex Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  arXiv preprint arXiv:2102.09672, 2021.
  * [23] Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu.  Conditional image generation with pixelcnn decoders.  arXiv preprint arXiv:1606.05328, 2016.
  * [24] Gaurav Parmar, Dacheng Li, Kwonjoon Lee, and Zhuowen Tu.  Dual contradistinctive generative autoencoder.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 823–832, 2021.
  * [25] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  arXiv preprint arXiv:2102.12092, 2021.
  * [26] Ali Razavi, Aaron van den Oord, and Oriol Vinyals.  Generating diverse high-fidelity images with vq-vae-2.  In Advances in neural information processing systems, pages 14866–14876, 2019.
  * [27] 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.
  * [28] Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar.  Theory and experiments on vector quantized autoencoders.  arXiv preprint arXiv:1805.11063, 2018.
  * [29] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma.  Pixelcnn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications.  arXiv preprint arXiv:1701.05517, 2017.
  * [30] 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.
  * [31] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In International Conference on Learning Representations, 2020.
  * [32] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  In Proceedings of the 33rd Annual Conference on Neural Information Processing Systems, 2019.
  * [33] Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  arXiv preprint arXiv:2007.03898, 2020.
  * [34] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu.  Neural discrete representation learning.  In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 6309–6318, 2017.
  * [35] Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu.  Pixel recurrent neural networks.  In International Conference on Machine Learning, pages 1747–1756. PMLR, 2016.
  * [36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Advances in neural information processing systems, pages 5998–6008, 2017.
  * [37] Antoine Wehenkel and Gilles Louppe.  Diffusion priors in variational autoencoders.  In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.
  * [38] Yuxin Wu and Kaiming He.  Group normalization.  In Proceedings of the European conference on computer vision (ECCV), pages 3–19, 2018.
  * [39] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat.  Vaebm: A symbiosis between variational autoencoders and energy-based models.  In International Conference on Learning Representations, 2020.
  * [40] 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.
  * [41] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.  The unreasonable effectiveness of deep features as a perceptual metric.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586–595, 2018.
  * [42] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.  The unreasonable effectiveness of deep features as a perceptual metric.  In CVPR, 2018.
  * [43] Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros.  Unpaired image-to-image translation using cycle-consistent adversarial networkss.  In Computer Vision (ICCV), 2017 IEEE International Conference on, 2017.

## Appendix

### Datasets

CelebA-HQ is a high-quality version of the CelebA dataset, consisting of 30000
images generated by PG-GAN. We followed [14] instructions to obtain the
dataset.

LSUN [40] includes ten scenes and twenty object categories, totally about one
million images with label. We mainly use the Church, which contains about
126,000 images. The image pre-processing method follows StyleGAN [15].

### Discrete VAEs

Our architecture for discrete image representation follows that in [8]. For
completeness, a brief description is as follows:

Encoder | Decoder
---|---
Conv2D | Conv2D
4$\times${ResDown} | Middle Block
Middle Block | 4$\times${ResDown}
GN, Swish, Conv2D | GN, Swish, Conv2D

  * 1

ResDown is the combination of a Residual Block and Downsample Block, if the
feature map size matches the preset value, there will be an addition non-local
self-attention block.

  * 2

Middle Block is the cascade of one Residual Block, one Self-attention Block
and one more Residual Block.

  * 3

GN means the group normalization [38]

Table 3: Brief Architecture of the VQ-GAN encoder and decoder

For CelebA-HQ and ImageNet, we obtain the pre-trained checkpoints from the
official release, for LSUN-Church, we trained a model from scratch under the
same configurations for ImageNet in [8]. Specifically, the embedding dimension
is 256 and the number of embedded tokens is 1024. The channel numbers of the
encoder-decoders is 128, the self-attention block is introduced when the
feature map size meets $16\times 16$. We set the learning rate is 4.5e-6 for
each instance and the learning rate is fixed.

### Discrete Diffusion Models

The network structures and hyperparameter settings of discrete diffusion
models follow [10]. In detail, the model architecture is based on the backbone
of PixelCNN++ [29], which is a U-Net [27] with group normalization. Instead of
only adding a self-attention block at $16\times 16$ feature map resolution
level, we increase two more self-attention blocks on $8\times 8$ and $4\times
4$ separately. We have 117M parameters for the diffusion models.

For the logits of
$\tilde{p_{\theta}}(\tilde{z_{0}}|z_{t})=\mathrm{Cat}(\tilde{z_{0}}|p_{\theta})$,
we predict a noise using the neural network and add it to $z_{t}$ instead of
predicting the $\tilde{z_{0}}$ directly. As shown in Eq. 23, the desired
logits is obtained by superimposing the predicted noise
$\mathrm{nn_{\theta}}(z_{t})$ on a calculated $z_{t}$

The noise schedule $\alpha_{t}$ is the same as [22]. The difference is that
their parameter $\sqrt{\hat{\alpha_{t}}}$ is assigned to the mean of the
Gaussian distribution, while our factor $\hat{\alpha_{t}}$ is the parameter of
the categorical distribution. The definition is given in Eq. 19 and $s=0.008$.
We also sample $t$ with $q(t)\propto\sqrt{\mathbb{E}[L^{2}_{t}]}$ instead of
uniform sampling [22].

The batch size is $180$ per GPU and the learning rate is $0.0001$ with Adam
optimizer with standard settings. The learning rate scheduler is the cosine
annealing scheduler with 1 million steps. We have not employed any dropout in
the model.

### Additional Results

In Figures 9 & 10, we show additional generation results based on CelebA and
on LSUN-Church. We also provide additional results for image inpainting in
Fig. 11.

### Risk of overfitting

As described in [8], FID scores cannot detect an overfitting, while early-
stopping based on validation NLL can prevent overfitting. In Fig 12, we show
top-$10$ nearest neighbors based on LPIPS distance [42] for the training
image. We can find that the nearest neighboring generated image is not the
reproduced original image and we can infer that there is no overfitting in
such model.

### Societal Impact

Our work is an extension of the diffusion model, which also belongs to the
family of generative models. It can be used to generate fake images or videos
to disseminate disinformation, however, as our adopted datasets are collected
from the Internet, which will contain the biases, the generated images from
our model are also difficult to escape from the bias caused by training data.

Figure 9: Additional samples on LSUN-Church. Figure 10: Additional samples on
CelebA-HQ. Figure 11: Additional samples on image inpainting for CelebA-HQ.
Figure 12: Nearest Neighbours for CelebA-HQ $256\times 256$ model. The left
column are the images generated by our model, and the remaining images are the
nearest neighbors(with minimum LPIPS distance) from the training set.


# SHMC-NET: A Mask-guided Feature Fusion Network for Sperm Head Morphology
Classification

###### Abstract

Male infertility accounts for about one-third of global infertility cases.
Manual assessment of sperm abnormalities through head morphology analysis
encounters issues of observer variability and diagnostic discrepancies among
experts. Its alternative, Computer-Assisted Semen Analysis (CASA), suffers
from low-quality sperm images, small datasets, and noisy class labels. We
propose a new approach for sperm head morphology classification, called SHMC-
Net, which uses segmentation masks of sperm heads to guide morphology
classification of sperm images. SHMC-Net generates reliable segmentation masks
using image priors, refines object boundaries with an efficient graph-based
method, and trains an image network with sperm head crops and a mask network
with the corresponding masks. In the intermediate stages of the networks,
image and mask features are fused with a fusion scheme to better learn
morphological features. To handle noisy class labels and regularize training
on small datasets, SHMC-Net applies Soft Mixup to combine mixup augmentation
and a loss function. We achieve state-of-the-art results on SCIAN and HuSHeM
datasets, outperforming methods that use additional pre-training or costly
ensembling techniques.

Index Terms—  Human Sperm Analysis, Sperm Head Morphology Classification,
Feature Fusion, Intra-class Mixup

##  1 Introduction

Approximately 1 in 6 adults worldwide experience infertility [16], and nearly
one-third of these cases are attributed to male infertility. Human sperm head
morphology classification is clinically significant because different types of
sperm head morphology abnormalities are indicative of genetic or environmental
factors that impact treatment decisions [13]. However, manual classification
of these abnormalities is subjective, time-consuming, and heavily dependent on
human experience, which leads to observer variability [4] and diagnostic
discrepancies (noisy class labels) even among experts. Away from subjective
manual assessments, computational approaches such as Computer Assisted Semen
Analysis (CASA) have advanced morphology analysis toward a more objective and
reproducible process [15]. Traditional CASA methods for sperm morphology
classification used hand-crafted shape descriptors and yielded competitive
results compared to manual classification [3, 2, 19, 1]. Recent advances in
Computer Vision (CV) and Deep Learning (DL) have developed methods [11, 24]
that no longer require hand-crafted features and generally outperform
traditional methods.

Existing sperm datasets often contain a limited quantity of images, which may
cause overfitting or over-focusing of DL models on clinically irrelevant image
regions [24]. The labels of such images may be noisy or inconsistent, due to
the subjective nature of expert annotations. (e.g., see row 2 of Fig. 1).
Sperm images are often of low quality and have various artifacts in the
background, varying textures in the foreground, and irregular structures
(e.g., in the $2^{nd}$ or $5^{th}$ crop of row 2 of Fig. 1, the thick sperm
mid-section could be taken as part of the head), which can affect DL models to
learn and discriminate effectively, especially with a limited number of
samples and noisy labels.

Fig. 1: Row 1: Sperm head crops with different labels (HuSHeM). Row 2: Class
label discrepancies (SCIAN). Fig. 2: An overview of our SHMC-Net.

To address these challenges we propose to incorporate additional information
from sperm head masks to help morphology classification. Sperm head masks can
be reliably and efficiently extracted from images, as shown in previous work
[17, 3, 24, 8]. Such masks in general have well-delineated boundaries with
fewer background artifacts. The objects of interest in the masks are
represented as the foreground with morphologically relevant shape, size, and
structure information that is easy to extract while containing little
distracted structures that may inhibit shape analysis. Motivated by these
observations, we propose a new approach for effective Sperm Head Morphology
Classification, called SHMC-Net. Our major contributions are (1) a new sperm
head morphology-aware classification approach that explicitly utilizes
morphology information from sperm head segmentation masks, (2) a
classification network with a novel architecture, Fusion Encoder, that applies
an effective feature fusion scheme to utilize information from both the input
images and masks (3) a novel combination of augmentation and loss function to
handle noisy labels and regularize training on small datasets, (4) an
efficient graph-based sperm head mask boundary refinement method and (5) SOTA
performances on 2 public sperm datasets outperforming even the best-known
methods that use additional pre-training or costly ensembling of multiple
convolution neural network (CNN) models.

##  2 Methodology

Fig. 2 gives an overview of our proposed Sperm Head Morphology Classification
Network, SHMC-Net. Given a raw sperm image, we first apply our mask generation
and refinement module to obtain sperm head crops and their corresponding
boundary-refined masks. The head crops and their corresponding masks are fed
into image network and mask network of the Fusion Encoder, respectively. In
the deeper stages of the two networks, image and mask features are fused and
passed into the subsequent stages of the image network. The final features
obtained from both networks are fused and passed through a linear classifier
to obtain the class predictions. SHMC-Net also uses Soft Mixup, which consists
of intra-class Mixup augmentation for both images and masks and a loss
function, to address the observer variability and class imbalance problems and
regularize the training on small datasets. The three major components of SHMC-
Net— mask generation and refinement, Fusion Encoder, and Soft Mixup
regularization — are detailed below.

###  2.1 Mask Generation and Refinement

Reasonable head masks can be extracted from sperm images [17, 3, 24, 8] in an
unsupervised fashion but these masks may have imperfect boundaries. Hence,
given a raw sperm image, we first generate segmentation masks for sperm head
crops using anatomical and image priors and further refine mask boundaries
using a graph-based method. Using the HPM method [24], we obtain the sperm-
head-only crops from the sperm image, right-align them, and produce their
initial pseudo-masks. We propose a sperm head shape-aware Graph-based Boundary
Refinement (GrBR) method. GrBR generates the boundary-refined masks
instantaneously ($<7$ ms per image). The GrBR module refines the boundary
contour $C$ of the HPM pseudo-mask in a head crop to obtain a closed curve
$C^{\prime}$ that captures the accurate head boundary of the mask. We
formulate the optimal contour refinement problem as computing the shortest
path in a directed graph $G$ defined on $C$.

We first uniformly sample $n$ points clockwise on $C$, and for each sample
point $a_{j}\in C$, a line segment $l(a_{j})$ is drawn orthogonal to the
curvature of $C$ at $a_{j}$, with $a_{j}$ being the center point of
$l(a_{j})$. We then sample $m$ points uniformly on each line segment
$l(a_{j})$. Let $p_{i}^{j}:=(x(p_{i}^{j}),y(p_{i}^{j}))$ represent the $i$-th
sample point on the line segment $l(a_{j})$, $i=1,2,\ldots,m$. All the sample
points on $C$ and on the line segments are vertices of the graph $G$ and have
weights equal to the negatives of their gradients. Since sperm heads have
almost convex shapes, we impose a smoothness constraint and a near convex
constraint for computing $C^{\prime}$. The smoothness constraint is enforced
by using a smoothness parameter $s$ such that each $p_{i}^{j}$ can have a
directed edge to any point $p_{i^{\prime}}^{j+1-\lfloor j/n\rfloor\times n}$
only, for each $i^{\prime}$ with $|i^{\prime}-i|\leq s$, where
$i=1,2,\dots,m$, $i^{\prime}=1,2,\dots,m$, and $j=1,2,\dots,n$. The shape
constraint is enforced by penalizing any concave edge-to-edge connection along
$C^{\prime}$ with a larger connection cost, i.e., edge weight =
$c\times(\theta-\pi)$ if $\theta>\pi$, and else edge weight = 0, where $c$ =
cost penalty = initial edge weight and $\theta$ is the angle between the two
adjacent edges. We make $G$ into a directed acyclic graph by making a copy of
all the vertices on the first line segment $l(a_{1})$. These vertex copies are
now on the line segment copy $l^{\prime}(a_{1})$ and their edges are from the
vertices of $l(a_{n})$.

An optimal closed curve $C^{\prime}$ in the head crop subject to the
smoothness and shape constraints based on $C$ is obtained by computing the
shortest path in the graph $G$ in polynomial time using dynamic programming.
The obtained $C^{\prime}$ corresponds to the boundary-refined sperm head mask.

Table 1:  Results on HuSHeM datasets. In the column of Pre-training, IN
denotes ImageNet pre-trained weights, US denotes unsupervised pre-training,
and SMIDS denotes pre-training on the SMIDS dataset. ‘–’ marks that
corresponding results are not available.

2* Method |  2*Pre-training | SCIAN Partial Agreement (PA) | HuSHeM
---|---|---|---
|  | Accuracy | Precision |  Recall |  F1 | Accuracy | Precision |  Recall | F1
CE-SVM [2] | ✕ | – | – |  57.6 |  – | 78.7 | 80.6 |  78.6 | 79.6
MSER [9] | ✕ | – | – |  – |  – | 86.6 | 88.4 |  86.6 | 86.1
ADPL [19] | ✕ | – | – |  62.4 |  – | 92.6 | 92.7 |  92.7 | 92.7
LCLM-MDL[14] | ✕ | – | – |  – |  – | 95.6 | 95.7 |  95.7 | 95.7
MC-HSH [10] | ✕ | 63.0 | 56.0 |  68.0 |  61.0 | 95.7 | 96.1 |  95.5 | 95.5
Yüzkat et.al [22] | ✕ | 71.9 | 67.0 |  52.8 |  59.1 | 85.2 | 85.2 |  85.3 | 85.3
FT-VGG [18] | IN | 61.7 | 61.9 |  61.7 |  61.8 | 94.1 | 94.3 |  94.1 | 94.2
TL [11] | IN | – | – |  62.0 |  – | 96.4 | 96.4 |  96.4 | 96.4
SHMCNN [20] | IN | – | – |  – |  – | 95.4 | 95.6 |  95.3 | 95.4
Ilhan et.al[7] | IN+SMIDS | 73.2 | 66.2 |  57.6 |  61.6 | 92.1 | 92.3 |  92.1 | 92.2
Zhang et.al[24] | IN+US | 65.9 | 58.7 |  68.9 |  63.3 | 96.5 | 96.8 |  96.6 | 96.6
SHMC-Net (ours) | IN |  73.6 $\pm$ 0.4 | 66.5  $\pm$ 0.8 |  63.6 $\pm$ 0.9 |  65.0 $\pm$ 0.3 |  98.2 $\pm$ 0.3 |  98.3 $\pm$ 0.3 |  98.1 $\pm$ 0.3 |  98.2 $\pm$ 0.3

###  2.2 Fusion Encoder

We propose a novel Fusion Encoder, to leverage the mask features and
facilitate the morphology classification. Using the masks directly to train a
classification network is not ideal (see Table 4) because the masks lack
texture details and are polarized to all-or-nothing errors due to their coarse
discretized representation. Hence, we use a one-directional feature fusion
scheme (from mask network to image), at the deeper network layers (evidence
for this design in Section 3.5) to exploit the learned semantics from both the
image and mask networks. The Fusion Encoder consists of an image network, a
mask network, fusion blocks, and a classifier.

####  2.2.1 Image and Mask Network

We use two 5-staged lightweight ShuffleNet [12] backboned encoders as our
image and mask network. In the first three stages, both networks learn the
features independent from each other. After the third stage, features from the
image and mask networks are fused and passed to the fourth stage of the image
network. The features from the two networks are fused again before the fifth
stage and passed to the image network to obtain the final features from the
image network.

####  2.2.2 Fusion Blocks and Classifier

Fusion blocks are placed at the beginning of the $4^{th}$ and $5^{th}$ stages
of the networks and consist of a pixel-wise summation of the image and mask
features, followed by a convolution, batch normalization, and ReLU operations.
A pixel-wise summation is performed between the final features obtained at the
end of the image and mask networks, followed by average pooling and fed
through the classifier. We use a randomly initialized fully connected linear
layer as a classifier that outputs (after the softmax and argmax operations)
the morphology class predictions of the input sperm head crop.

###  2.3 Soft Mixup Regularization

Accurately determining the sperm head morphology classes is difficult even for
experts, which may lead to disagreement on expert-labeled class IDs of samples
(i.e., for the SCIAN dataset a sample may have a majority class label with
consensus among at least 2/3 expert labels and a minority class label with no
consensus).

####  2.3.1 Soft Mixup Augmentation

We apply an intra-class Mixup augmentation [23] scheme to handle samples with
disagreed labels and reduce class conditional input variance. Given a sample
image crop $x_{i}$ and its corresponding mask $m_{i}$ with a majority label
$y^{a}_{i}$ and a minority label $y^{b}_{i}$, we randomly pick another sample
$x_{j}$ (with a mask $m_{j}$, a majority label $y^{a}_{j}$, and a minority
label $y^{b}_{j}$) if $y^{a}_{i}=y^{a}_{j}$. Prior to this step, we perform
random oversampling of all the under-represented classes to match the number
of samples in the most represented class. We then apply separate image
augmentation ($\Phi$) to $x_{i}$ and $x_{j}$. The input, $\ddot{x}_{i}$, to
the Fusion Encoder is a linear combination of the augmented mixup samples
$\Phi_{i}(x_{i})$ and $\Phi_{j}(x_{j})$, and its target label $\ddot{y}_{i}$
is equivalent to the majority label ($y_{i}^{a}$). Formally:

$\displaystyle\ddot{x}_{i}$
$\displaystyle:=\lambda\;\Phi_{i}(x_{i})+(1-\lambda)\;\Phi_{j}(x_{j}),\;\;\ddot{y}_{i}:=y^{a}_{i}=y^{a}_{j},$
(1) $\displaystyle\ddot{m}_{i}$
$\displaystyle:=\lambda\;\Phi_{i}(m_{i})+(1-\lambda)\;\Phi_{j}(m_{j}).$

where $\lambda$ is the mixup strength. Soft Mixup augmentation yields a mixed
crop $\ddot{x}_{i}$, mixed mask $\ddot{m}_{i}$, target label $\ddot{y}_{i}$,
and using these as input, the model outputs class prediction
$\widehat{y}_{i}$.

####  2.3.2 Soft Mixup Loss

We modify the Soft Loss in [24] to adapt intra-class mixup to handle partial
agreements among the morphology class labels. Soft Loss is defined as:

$\mathcal{L}_{s}(\widehat{y},y^{a},y^{b})=\gamma\mathcal{L}_{ce}(\widehat{y},y^{a})+(1-\gamma)\mathcal{L}_{ce}(\widehat{y},y^{b}),\vspace{-0.1cm}$
(2)

where $\mathcal{L}_{ce}$ is the cross entropy loss and $\gamma$ is a hyper-
parameter that weighs the majority label. Then, we combine our soft loss with
the intra-class mixup to build the Soft Mixup Loss, as:

$\mathcal{L}_{smx}=\lambda\mathcal{L}_{s}(\widehat{y}_{i},y^{a}_{i},y^{b}_{i})+(1-\lambda)\mathcal{L}_{s}(\widehat{y}_{i},y^{a}_{j},y^{b}_{j}),\vspace{-0.1cm}$
(3)

where $\lambda$ is the same as used in Eq. (1). When minority labels are not
available we retain our formulation of $\mathcal{L}_{smx}$ but with
$y_{i}^{a}=y_{i}^{b}$, reducing our loss to that in the original Mixup [23].

##  3 Dataset, Experiments, and Results

###  3.1 Datasets

We demonstrate the performance of SHMC-Net on two public sperm morphology
datasets. SCIAN-Morpho Sperm GS (SCIAN) [1] contains 1854 gray-scale sperm
images of size 35 $\times$ 35 (height x width), each belonging to one of five
morphological classes: Normal (N), Tapered (T), Pyriform (P), Small (S), or
Amorphous (A) as labeled by 3 experts. Samples where all 3 experts agree on
the same class label form the SCIAN total Agreement (SCIAN-TA) dataset (384
images: 35 N, 69 T, 7 P, 11 S, 262 A), and the samples where at least 2
experts agree on class labels form the SCIAN Partial Agreement (SCIAN-PA)
dataset (1132 images: 100 N, 228 T, 76 P, 72 S, 656 A)

Similarly, the Human Sperm Head Morphology (HuSHeM) dataset [19] contains 216
RGB sperm images of size 131 $\times$ 131 belonging to one of the 4
morphological classes: Normal (54), Tapered (53), Pyriform (57), and Amorphous
(52).

###  3.2 Implementation Details

ImageNet pre-trained ShuffleNet V2 [12] is used as the backbone for our image
and mask networks and trained using the AdamW optimizer with an initial
learning rate of 0.00015 and the Cosine Annealing learning rate scheduler. The
model is trained on a single NVIDIA TITAN-Xp GPU for 300 epochs (batch size of
128 for SCIAN-PA and 64 for SCIAN-TA and HuSHeM). Dropout with 40% probability
and weight decay of 0.01 was applied as additional regularization. The
majority class weight $\gamma$ is set to 0.85 and mixup strength, $\lambda$,
is sampled randomly from the beta distribution with $\alpha=0.5$. All the
image crops are first resized to 64$\times$64 and randomly square-cropped
between sizes 45 to 64 and resized again to 64*64 along with 10-degree
rotation, vertical flip, and horizontal/vertical shift by up to 10% of their
width or height. We evaluate the performance in terms of Overall Accuracy,
macro-averaged class-wise Precision and Recall, and the F1 score attained as a
harmonic mean of the two. We apply 5-fold cross-validation and report the
average of 5 runs along with the standard deviation.

###  3.3 Performance Comparison on SCIAN

As seen in Table 1, our SHMC-Net yields the highest Accuracy and F1 on SCIAN-
PA. Known SOTA methods traded Precision and Recall, while SHMC-Net balances
these two and gives the best F1. Similarly, for SCIAN-TA, we achieve the best
Accuracy, Precision, and F1 outperforming even the ensemble methods (Table 4).

###  3.4 Performance Comparison on HuSHeM

On HuSHeM dataset SHMC-Net outperforms the previous SOTA methods on all the
metrics despite the diminished margins for performance improvement (Table 1).
We also implement an ensemble version of SHMC-Net (SHMC-Net-E) using three
backbones (ShuffleNet V2, ResNet34 [5], and DenseNet121 [6]), and apply
majority voting. With 6 times fewer parameters ($\sim$50 millions) SHMC-Net-E
outperforms [21] ($\sim$300 millions) by about 0.65% on all metrices.

###  3.5 Component Studies

We study different fusion directions: image network to mask (I2M), mask to
image (M2I), and bidirectional (BD) where both networks mutually help each
other and no fusion (✗). At the classifier level, we explore the options to
use features from either the image network (I) or both networks (B). In Table
4, we observe that avoiding fusion in the early stages, using mask features
for fusion (M2I) in the later stages, and fusing both the image and mask
features before the classifier offers the best performance.

Table 2:  Additional Results on SCIAN-TA and HuSHeM.

Method | Ensemble |  Accuracy | Precision | Recall | F1
---|---|---|---|---|---
SCIAN Total Agreement (TA)
FT-VGG [18] | ✕ |  – | – | 72.0 | –
CE-SVM [2] | ✓ |  – | – | 73.2 | –
MC-HSH [10] | ✕ |  77.0 | 64.0 | 88.0 | 74.0
Spencer et al. [21] | ✓ |  86.3 | 77.0 | 91.8 | 81.2
SHMC-Net (ours) | ✕ |  86.9 $\pm$ 0.8 |  86.7 $\pm$ 1.4 | 83.4  $\pm$ 1.6 |  85.0 $\pm$ 1.5
Ensemble for HuSHeM
Spencer et al. [21] | ✓ |  98.52 | 98.53 | 98.52 | 98.52
SHMC-Net-E (ours) | ✓ |  99.17 $\pm$ 0.20 |  99.19 $\pm$ 0.21 |  99.16 $\pm$ 0.20 |  99.17 $\pm$ 0.20

Table 3:  Study of key components.

2*Fusion |  2*Mixup | SCIAN-TA | HuSHeM
---|---|---|---
|  | Acc | F1 |  Accuracy |  F1
Fusion Options
M2I M2I M2I M2I B | Intra-Class | 62.9 | 58.5 |  90.7 |  91.2
BD BD BD BD B | Intra-Class | 64.0 | 58.4 |  90.7 |  91.1
I2M I2M M2I M2I B | Intra-Class | 70.4 | 62.9 |  97.2 |  97.3
✗ ✗ M2I M2I I | Intra-Class | 72.3 | 64.5 |  97.7 |  97.8

Table 4: Ablation study results. C.L. = classifier level, I.S. = intermediate
stage fusion.

2*  Image |  2*  Masks |  2*  C.L. Fusion |  2*  I.S. Fusion |  2*  Soft Mixup | SCIAN-PA | HuSHeM
---|---|---|---|---|---|---
|  |  |  |  | Accuracy | F1 | Accuracy | F1
✓ |  |  |  |  | 70.0 | 57.5 | 94.5 | 94.7
| ✓ |  |  |  | 58.8 | 43.9 | 82.0 | 82.2
✓ | ✓ |  ✓ |  |  | 71.6 | 60.9 | 96.3 | 96.5
✓ | ✓ |  ✓ | ✓ |  | 72.5 | 62.8 | 97.7 | 97.8
✓ | ✓ |  ✓ | ✓ | ✓ | 73.6 | 65.0 | 98.2 | 98.2

###  3.6 Ablation Study

Using only masks to train the model yields poor performance (Table 4) since
their coarse discretized representation lacks texture details and is polarized
to all-or-nothing errors. However, combining mask and image features at the
classifier level has improvements (3.4% F1 on SCIAN-PA and 1.8% on HuSHeM)
over using image-only. Intermediate stage fusion of mask and image features
further strengthens the model’s ability to learn morphology features (further
1.9% on SCIAN-PA and 1.3% on HuSHeM). Finally, handling class imbalance and
label discrepancies through Soft Mixup adds 2.2% further F1 improvement on
SCIAN-PA, and the regularization through Soft Mixup contributes to an
additional 0.4% F1 on the small-size HuSHeM dataset.

##  4 Conclusions

We proposed a new sperm morphology classification framework, SHMC-Net, which
employs morphological information of sperm head masks. SHMC-Net generates
boundary-refined masks using image priors and incorporates an effective
feature fusion method that leverages features from the raw images and their
masks. The Soft Mixup component was introduced to handle noisy labels and
regularize training on small datasets. SOTA performance was obtained on two
sperm morphology datasets, without additional pre-training or costly
ensembling of multiple CNNs. Ablation and component studies demonstrated the
effectiveness of SHMC-Net.

##  5 COMPLIANCE WITH ETHICAL STANDARDS

This research study involved a retrospective analysis of human subject data
obtained from two publicly accessible open-access datasets. Ethical approval
was not required as confirmed by the licenses attached within the open access
of the two datasets [1, 19] used in the experiments of our method.

## References

  * [1] V. Chang, A. Garcia, N. Hitschfeld, and S. Härtel.  Gold-standard for computer-assisted morphological sperm analysis.  CBM, 83:143–150, 2017.
  * [2] V. Chang, L. Heutte, C. Petitjean, S. Härtel, and N. Hitschfeld.  Automatic classification of human sperm head morphology.  CBM, 84:205–216, 2017.
  * [3] V. Chang, J.M. Saavedra, V. Castañeda, L. Sarabia, N. Hitschfeld, and S. Härtel.  Gold-standard and improved framework for sperm head segmentation.  Computer Methods and Programs in Biomedicine, 117(2):225–237, 2014.
  * [4] M. Freund.  Standards for the rating of human sperm morphology. a cooperative study.  International Journal of Fertility, 11(1):97–180, 1966.
  * [5] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In CVPR, pages 770–778, 2016.
  * [6] G. Huang, Z. Liu, L. Van Der Maaten, and K.Q. Weinberger.  Densely connected convolutional networks.  In CVPR, pages 4700–4708, 2017.
  * [7] H.O. Ilhan and G. Serbes.  Sperm morphology analysis by using the fusion of two-stage fine-tuned deep networks.  Biomedical Signal Processing and Control, 71:103246, 2022.
  * [8] H.O. Ilhan, G. Serbes, and N. Aydin.  Automatic directional masking technique for better sperm morphology segmentation and classification analysis.  Electronics Letters, 55(5):256–258, 2019.
  * [9] H.O. Ilhan, G. Serbes, and N. Aydin.  Automated sperm morphology analysis approach using a directional masking technique.  CBM, 122:103845, 2020.
  * [10] I. Iqbal, G. Mustafa, and J. Ma.  Deep learning-based morphological classification of human sperm heads.  Diagnostics, 10(5):325, 2020.
  * [11] R. Liu, M. Wang, M. Wang, J. Yin, Y. Yuan, and J. Liu.  Automatic microscopy analysis with transfer learning for classification of human sperm.  Applied Sciences, 11(12):5369, 2021.
  * [12] N. Ma, X. Zhang, H. Zheng, and J. Sun.  ShuffleNet V2: Practical guidelines for efficient CNN architecture design.  In ECCV, pages 116–131, 2018.
  * [13] R. Menkveld, Cas A. Holleboom, and J.P Rhemrev.  Measurement and significance of sperm morphology.  Asian Journal of Andrology, 13(1):59, 2011.
  * [14] T. Ni, Y. Ding, J. Xue, K. Xia, X. Gu, and Y. Jiang.  Local constraint and label embedding multi-layer dictionary learning for sperm head classification.  ACM Transactions on Multimedia Computing, Communications, and Applications (TOMM), 17(3s):1–16, 2021.
  * [15] World Health Organization et al.  WHO laboratory manual for the examination and processing of human semen.  2010\.
  * [16] World Health Organization et al.  Infertility prevalence estimates: 1990–2021.  2023\.
  * [17] K.S. Park, W.J. Yi, and J.S. Paick.  Segmentation of sperms using the strategic Hough transform.  Annals of Biomedical Engineering, 25:294–302, 1997.
  * [18] J. Riordon, C. McCallum, and D. Sinton.  Deep learning for the classification of human sperm.  CBM, 111:103342, 2019.
  * [19] F. Shaker, S.A. Monadjemi, J. Alirezaie, and A.R. Naghsh-Nilchi.  A dictionary learning approach for human sperm heads classification.  CBM, 91:181–190, 2017.
  * [20] M.A.C. Soares, D.H.M. Falci, M.F.A. Farnezi, H.C.M. Farnezi, F.S. Parreiras, and J.V.B. Gomide.  Automated sperm head morphology classification with deep convolutional neural networks.  In 2022 35th SIBGRAPI Conference on Graphics, Patterns and Images (SIBGRAPI), volume 1, pages 270–275. IEEE, 2022.
  * [21] L. Spencer, J. Fernando, F. Akbaridoust, K. Ackermann, and R. Nosrati.  Ensembled deep learning for the classification of human sperm head morphology.  Advanced Intelligent Systems, 4(10):2200111, 2022.
  * [22] M. Yüzkat, H.O. Ilhan, and N. Aydin.  Multi-model CNN fusion for sperm morphology analysis.  CBM, 137:104790, 2021.
  * [23] H. Zhang, M. Cisse, Y.N. Dauphin, and D. Lopez-Paz.  mixup: Beyond empirical risk minimization.  In ICLR, 2018.
  * [24] Y. Zhang, J. Zhang, X. Zha, Y. Zhou, Y. Cao, and D. Chen.  Improving human sperm head morphology classification with unsupervised anatomical feature distillation.  In ISBI, pages 1–5. IEEE, 2022.