# QuEST: Low-bit Diffusion Model Quantization via
Efficient Selective Finetuning

Haoxuan Wang  Yuzhang Shang  Zhihang Yuan  Junyi Wu  Yan Yan

###### Abstract

Diffusion models have achieved remarkable success in image generation tasks,
yet their practical deployment is restrained by the high memory and time
consumption. While quantization paves a way for diffusion model compression
and acceleration, existing methods totally fail when the models are quantized
to low-bits. In this paper, we unravel three properties in quantized diffusion
models that compromise the efficacy of current methods: imbalanced activation
distributions, imprecise temporal information, and vulnerability to
perturbations of specific modules. To alleviate the intensified low-bit
quantization difficulty stemming from the distribution imbalance, we propose
finetuning the quantized model to better adapt to the activation distribution.
Building on this idea, we identify two critical types of quantized layers:
those holding vital temporal information and those sensitive to reduced bit-
width, and finetune them to mitigate performance degradation with efficiency.
We empirically verify that our approach modifies the activation distribution
and provides meaningful temporal information, facilitating easier and more
accurate quantization. Our method is evaluated over three high-resolution
image generation tasks and achieves state-of-the-art performance under various
bit-width settings, as well as being the first method to generate readable
images on full 4-bit (i.e. W4A4) Stable Diffusion. Code is been made publicly
available at URL .

Machine Learning, ICML

##  1 Introduction

Diffusion models (Ho et al., 2020; Dhariwal & Nichol, 2021; Rombach et al.,
2021) have achieved great success in the field of image generation recently.
Given the high capacity and flexibility of diffusion models, there still
remains two challenges that limit their efficiency (Croitoru et al., 2023).
The first obstacle is the time-consuming denoising process which requires
hundreds or thousands of time steps that slows down the inference speed
drastically. The other one is the increasing model size driven by the need of
better image fidelity and higher image resolutions. Both factors contribute to
considerable latency and increased computational requirements, impeding the
application of diffusion models to real-world settings where both time and
computational power are restricted.

Neural network quantization offers a feasible solution for accelerating
inference speed and reducing memory consumption simultaneously (Gholami et
al., 2022). It aims to compress high-bit model parameters into low-bit
approximations with negligible performance degradation. For example, using
4-bit weight and 4-bit activation quantization can achieve up to 8$\times$
inference time speedup and 8$\times$ memory reduction theoretically (Liang et
al., 2021). The benefits will increase consistently when the bit-widths are
further lowered. Hence, low-bit quantization of diffusion models emerges as a
viable approach for enhancing their efficiency.

However, existing diffusion model quantization approaches fail under low-bit
settings. These methods focus either on timestep-aware calibration data
construction (Shang et al., 2023; Li et al., 2023a) or quantization noise
correction (He et al., 2023b), targeting at tailoring existing quantization
techniques to the properties of diffusion models that are distinct from other
model types such as CNNs (Pilipović et al., 2018) and ViTs (Chitty-Venkata et
al., 2023). Such approaches overlook the intrinsic diffusion model mechanisms
that are related with quantization, leading to inconsistencies between method
deployment and model characteristics. Therefore, we intend to identify unique
features in diffusion models that are associated with quantization.

In this paper, we reveal three properties of quantized diffusion models that
hinder effective quantization:  ❶ The activation distributions are likely to
be imbalanced, where the majority of values are close to 0 but the other
values are large and appear inconsistently. Conventional quantization methods
tend to either approximate either the large and sparse values, inadequately
estimating the numerous small values, or focus on small values while
overlooking the large ones, thereby impeding the reduction of quantization
error.  ❷ Diffusion models exhibit varying functions at distinct time steps
(Choi et al., 2022), such as forming object outlines at early time steps and
constructing details latter. Thus, preserving accurate temporal information is
important during quantization. However, current methods struggle to represent
extensive time steps with limited bit-width under low-bit quantization,
leading to similar time embeddings across different time steps and results in
error accumulation.  ❸ Diffusion models typically have complex network
architectures, incorporating convolutional layers and linear layers that
construct various kinds of attention blocks (Dosovitskiy et al., 2021) and
residual blocks (He et al., 2015). Whereas previous works consider each model
component as equally important and apply quantization uniformly, we reveal
that certain layers are particularly sensitive to perturbations from
quantization, while others are more resilient.

Addressing the revealed properties, we propose a novel quantization approach
for diffusion models, termed QuEST (Quantization via Efficient Selective
FineTuning). Confronting property  ❶ , we first theoretically justify that
weight finetuning can enhance model robustness towards large perturbations
under low-bit settings for reducing quantization error, whereas previous
methods struggle to find appropriate quantization parameters. Then we
empirically show that by finetuning the model, the activation distributions
are actually modified to be more amenable to quantization. Following the idea
of weight finetuning, we identify two types of layers that are responsible for
performance degradation: Time embedding layers exhibiting property  ❷ and
attention-related layers corresponding to property  ❸ . Thus, we selectively
and progressively finetune these small layer subsets in conjunction with all
activation quantization parameters, as illustrated in Figure 1(b). With less
than 10% of the total parameters involved, QuEST not only substantially
enhances low-bit quantized model performance, but is also notably time-
efficient and requires only supervision from the full-precision model in a
data-free manner. To preserve the essential time information, we further
introduce an efficient Time-aware Activation Quantizer (shown in Figure 1(a))
that can flexibly adapt to the varying activation distributions over the
numerous time steps. Our contributions are summarized as follows:

  * •

We uncover and validate three properties that present challenges to low-bit
diffusion model quantization. We also theoretically discuss why previous
methods are expected to fail and justify the sufficiency of weight finetuning
under low-bit settings.

  * •

Addressing the identified properties, we introduce QuEST, an innovative
finetuning strategy that trains the model selectively and progressively,
achieving low-bit quantization capability with high time and memory
efficiency. An efficient activation quantizer is also proposed for effective
time information estimation. Practically, our method can be executed in a
data-free and parameter-efficient manner.

  * •

Experiments on three high-resolution image generation tasks over four models
demonstrate the superiority of our method, achieving state-of-the-art
performance under various bit-width settings. We are also the first method to
generate reasonable images for full 4-bit Stable Diffusion.

Figure 1: Framework of the proposed method at time step $t$. $s_{l}^{t}$
represents the quantization parameters of layer $l$ at time $t$,
$Q(\textbf{w}_{l})$ represents the quantized weight of layer $l$. (a) In a
quantized layer, the scaling factor $s_{l}^{t}$ is updated while the others
are frozen at time $t$. (b) $e(t)$ and $x_{t}$ are the input time embedding
and input image representation respectively. The model layers’ weights are
optimized selectively and the quantization parameters are all updated. While
different selected layers are learned progressively, the corresponding weight
and scaling factors are updated simultaneously using intermediate activation
alignment and task loss guidance.

##  2 Related Work

###  2.1 Diffusion Model Inference

Diffusion models (Ho et al., 2020; Rombach et al., 2021) generate samples via
a learned iterative denoising process. During inference, the initial input is
sampled from a Gaussian distribution:
$x_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, and the final output $x_{0}$ is
obtained by removing noises from $x_{T}$ subsequently using:

$\displaystyle
p_{\theta}(x_{t-1}|x_{t})=\mathcal{N}(x_{t-1};\tilde{\boldsymbol{\mu}}_{\theta,t}(x_{t}),\tilde{\beta_{t}}\mathbf{I}).$
(1)

In this Markov Chain, $\tilde{\boldsymbol{\mu}}_{\theta,t}$ and
$\tilde{\beta_{t}}$ are calculated from the model’s output. For a typical
diffusion model, the denoising stage requires hundreds to thousands of
iterations. This laborious sequential process makes efficient inference and
high-resolution generation extremely challenging. Latent Diffusion Models
(LDMs) (Rombach et al., 2021) were developed to overcome the challenges. They
encode the data into a latent space and then apply the diffusion process
within this space, enabling high-resolution image generation with fewer time
steps. The output images are produced by decoding the latent representations.
LDM’s concept and network quantization complement each other, mutually
enhancing the efficiency of diffusion models.

Practically, diffusion models typically adopt a UNet architecture (Ronneberger
et al., 2015), with some modules substituted by attention blocks. In a latent
diffusion model, an encoder and a decoder are incorporated for projecting
images into latent representations and subsequently generating images from the
latent space, respectively. Usually encoders and decoders are lightweight and
computationally inexpensive, so our focus is on quantizing the UNets in LDMs,
in alignment with other works (as described below).

Table 1: Comparison with Different Quantization Frameworks. Method | Data | Time & Memory | Low-bit
---|---|---|---
| Free | Efficient | Compatible
PTQ |  ✔ |  ✔ |  ✗
QAT |  ✗ |  ✗ |  ✔
Ours |  ✔ |  ✔ |  ✔

###  2.2 Diffusion Model Quantization

Model quantization is a dominant technique for optimizing the inference memory
and speed of deep learning models by reducing the precision of the tensors
used in computation. Few researches were done for quantizing diffusion models
and they fall into two mainstream frameworks: Quantization-Aware Training
(QAT) (Jacob et al., 2017; Li et al., 2022; Xu et al., 2023; Li et al., 2023b)
and Post-Training Quantization (PTQ) (Wang et al., 2020; Nahshan et al., 2021;
Li et al., 2021; Wei et al., 2023a; Liu et al., 2023). QAT methods such as
Q-DM (Li et al., 2023b) consider value precision as an additional optimization
constraint, training all parameters from scratch with the target task
objective. While it is effective for low-bit quantization, it is even more
resource-intensive than training a full-precision model. Efficient-DM (He et
al., 2023a) further uses Low-rank adapters (LoRA) (Hu et al., 2021) to reduce
training cost. However, it introduces extra weight parameters and still
requires substantial training iterations.

Post-Training Quantization methods aim to calibrate the quantization
parameters based on a small calibration set. Pioneering works such as PTQ4DM
(Shang et al., 2023) and Q-Diffusion (Li et al., 2023a) focus on sampling the
outputs of the full-precision model under different time steps with certain
probabilities to construct the calibration set. ADP-DM (Wang et al., 2023)
partition time steps into different groups and optimize each group with
differentiable search algorithms. PTQ-D (He et al., 2023b) decomposes the
quantization noise and correct them individually with full-precision model
statistics. However, these methods generally fail under 4-bit or lower bit-
width. To enable low-bit compatibility with high efficiency, our method adopts
the general framework of PTQ but further introduces an efficient weight fine-
tuning strategy. Table 1 shows the unique capabilities of our method, where we
are able to realize QAT ability with PTQ efficiency.

##  3 Methodology

###  3.1 Preliminaries

The quantization process for a single value $x$ in a vector can be formulated
as:

$\displaystyle\hat{x}$
$\displaystyle=\text{clamp}(\text{round}(\frac{x}{s})+Z;q_{min},q_{max}),$ (2)

where $\hat{x}$ is the quantized integer result, $\text{round}(\cdot)$
represents rounding algorithms such as the round-to-nearest operator (Li et
al., 2021) and AdaRound (Nagel et al., 2020), $s$ is referred to as the
scaling factor and $Z$ is the zero-point. clamp is the function that clamps
values into the range of $[q_{min},q_{max}]$, which is determined by the bit-
width. Reversely, transforming the quantized values back into the full-
precision form is formulated as:

$\displaystyle\tilde{x}=(\hat{x}-Z)*s.$ (3)

This is denoted as the dequantization process. The quantization and
dequantization processes are performed on both model weights and layer outputs
(also termed as ’activation’). Equation 2 indicates that the quantization
error is composed of two factors: The clipping error produced by range
clamping and the rounding error caused by the rounding function, where they
exhibit a trade-off relationship (Li et al., 2021). While previous approaches
strive for an optimal balance between the two errors, they neglect the
intrinsic mechanisms in quantized diffusion models. In the next section, we
investigate three properties of diffusion models that relate to quantization
performance and show why typical quantization methods tend to fail.

###  3.2 Quantization-aware Properties of Diffusion Models

Property  ❶ : The activation outputs are likely to have a majority of values
close to zero, with the other values being numerically large and important.
Previous works (Shang et al., 2023; Li et al., 2023a) paid attention to the
varying activation distribution over different time steps, but none of them
focused on the activation distribution itself. As shown in Figure 2, we find
that in some layers, though the majority of values are close to 0, there
exists values that are relatively large and diverse. Take the bin plot on the
left for example, the original activation values (blue line) ranges from [-10,
34] but with over 99% values between [-0.6, 1.7], posing difficulties in
minimizing the clipping error. This phenomenon is different from the outlier
distribution observed in LLMs (Lin et al., 2023; Xiao et al., 2023), where the
outliers are always positive and show up at certain positions consistently
(Wei et al., 2023b). Moreover, we observe that the larger values are important
for generation performance. With the large values being important and the
small values appearing frequently, neither of them are negligible and need to
be carefully quantized at the same time. Unfortunately, typical quantization
methods fall short of this ability under low-bit settings, where the rounding
error often outweighs the clipping error during optimization and results in
over-clipped large values, generating corrupted images. This phenomenon
inspires us to reshape the activation distributions via model finetuning to
attain more quantization-friendly distributions.

Figure 2: Illustrations of imbalanced activation distributions. We cluster
activation values into uniformly distributed bins and manipulate the frequency
for better visualization. In the activation distribution of the full-precision
model, the majority of values cluster near zero with sporadic large values,
presenting challenges for low-bit quantization. Our method reduces the large
values to enable easier quantization for improved performance.

Property  ❷ : Accurate temporal information is critical for quantized
diffusion model performance. As shown in Equation 1, diffusion model
predictions are highly correlated with time step information. Specifically,
time steps are transformed and projected into time embeddings and then added
onto the image features, guiding the model to denoise correctly. Table 2 shows
the performance comparison of quantizing time embeddings or preserving their
full-precision values in the LDM-4 of LSUN-Bedrooms (Yu et al., 2015). Under
W8A8 and W4A8, quantizing the time embeddings leads to an increase of 0.81 and
1.04 (relatively 15%) in FID, respectively. This indicates that inaccurate
time embedding quantization can lead to large degradation in generation
performance, yielding even larger corruptions under lower bit-widths. Due to
the different focus diffusion models have on image generation at different
time steps, we infer that inaccurate time embeddings can cause mismatched
input image and model functionality, resulting in possible oscillations in the
sequence of noise removal. Thus, preserving diffusion model performance under
quantization requires accurate temporal information.

Property  ❸ : Different activations have different sensitivity towards
decreasing bit-width. Low-bit activation quantization is often more difficult
than weight quantization and remains the key challenge for diffusion model
quantization. Compared with weights than can be categorized into linear layers
and convolutional layers, activation types are more diverse and complicated.
We hence examine the effects on model performance when quantizing different
activations. Figure 3 shows how sensitive different activations are to
quantization in Stable Diffusion. We quantize three different types of
activations to lower bits while maintaining the others’ bit-width (8-bit). We
observe that the FeedForward layer (Feng et al., 2023) activations cause
generation failure at 6 bits, whereas the activations of all other linear
layers (containing 5 times more layers) barely fail at 4 bits and all
convolution layers (containing 3 times more layers) only fail at 4 bits. This
phenomenon indicates that a small amount of activations are especially
sensitive to low-bit quantization, making them essential to be specially dealt
with.

Table 2: FP & Quantized Time Embedding Comparisons. Settings |  FID $\downarrow$ |  sFID $\downarrow$
---|---|---
W8A8 (Quantized Time Embedding) | 7.58 | 22.07
W8A8 (FP Time Embedding) | 6.77 | 22.03
W8A8 (TAQuant+$\text{SLA}_{\text{TE}}$) | 5.61 | 21.22
W4A8 (Quantized Time Embedding) | 8.59 | 22.74
W4A8 (FP Time Embedding) | 7.55 | 21.69
W4A8 (TAQuant+$\text{SLA}_{\text{TE}}$) | 6.95 | 23.17

###  3.3 Quantization via Efficient Selective Finetuning

In this section, we first theoretically discuss the potential reasons for the
failure of previous methods and the underlying rationale for weight fine-
tuning. Then we introduce a finetuning method QuEST for diffusion models that
can significantly boost low-bit performance with less time and memory usage.
Figure 1 illustrates the proposed framework, which is a distillation-based
finetuning strategy consisting of selective weight optimization and network-
wise scaling factor optimization. To further address the loss of temporal
information due to quantization, we propose an efficient time-aware activation
quantizer.

####  3.3.1 Sufficiency of Finetuning

We first review the underlying theoretical foundation underpinning previous
methods, which typically employ the reconstruction-based post-training
quantization approach. Denote the diffusion model’s activations at time $t$ as
$\mathbf{z}_{t}=[z_{1,t},z_{2,t},...,z_{n,t}]$, the task loss as
$L(z_{n,t},\overline{z}_{n,t})$, where $n$ is the number of layers and
$\overline{z}_{n,t}$ is the ground-truth. $L$ can be any loss function and
here we use the mean squared error (MSE). We treat quantization as a type of
perturbation and formulate the influence of activation quantization using
Taylor expansion, assuming model weight $\mathbf{w}$ is frozen:

$\displaystyle\mathbb{E}[L(z_{n,t}$
$\displaystyle+\Delta;\mathbf{w})]-\mathbb{E}[L(z_{n,t};\mathbf{w})]$
$\displaystyle\approx\Delta^{T}\overline{\mathbf{g}}^{(z_{n,t})}+\frac{1}{2}\Delta^{T}\overline{\mathbf{H}}^{(z_{n,t})}\Delta,$
(4)

where $\Delta$ is the activation perturbation,
$\overline{\mathbf{g}}^{(\mathbf{z})}$ is the gradient and
$\overline{\mathbf{H}}^{(z_{n,t})}$ is the Hessian matrix. According to (Li et
al., 2021; Yuan et al., 2022), for a well-trained model,
$\overline{\mathbf{g}}^{(z_{n,t})}=\nabla_{z_{n,t}}L$ approaches $0$. Thus the
above equation can be simplified to:

$\displaystyle\mkern-18.0mu\frac{1}{2}\Delta^{T}\overline{\mathbf{H}}^{(z_{n,t})}\Delta=\frac{1}{2}(\tilde{z}_{n,t}-z_{n,t})^{T}\overline{\mathbf{H}}^{(z_{n,t})}(\tilde{z}_{n,t}-z_{n,t}).\mkern-18.0mu$
(5)

This leads to a rather straightforward conclusion: Quantization error can be
reduced by reconstructing the model output. However, under low-bit settings,
the reasoning from Equation 3.3.1 to Equation 5 is inaccurate, where the
activation perturbation $\Delta$ is too large for meaningful Taylor expansion.
Thus we have the following proposition:

######  Proposition 3.1.

Reconstruction-based post-training quantization methods may lose their
theoretical guarantee due to the large value perturbations under low-bit
quantization.

To resolve the above problem, we can use the following theorem by transforming
$\Delta$ into a smaller perturbation $\epsilon$:

######  Theorem 3.2.

Given an $n$ layer diffusion model at time $t$ with quantized activations as
$\tilde{\mathbf{z}}_{t}=[\tilde{z}_{1,t},\tilde{z}_{2,t},...,\tilde{z}_{n,t}]$
and $\tilde{z}_{n,t}=z_{n,t}+\Delta$, where $z_{n,t}$ is the ground-truth and
$\Delta$ is the large perturbation caused by low-bit quantization. Denote the
target task MSE loss as $L(\tilde{z}_{n,t},\overline{z}_{n,t})$, the
quantization error can be transformed into:

$\displaystyle\mathbb{E}[L(z_{n,t}+\Delta;\mathbf{w})]-\mathbb{E}[L(z_{n,t};\mathbf{w})]$
$\displaystyle\approx
2\epsilon^{T}\sum_{i=1}^{K}(\tilde{z}_{n-1,t}^{i}\cdot\mathbf{w}_{n}-z_{\text{FP}})$
$\displaystyle+\frac{1}{2}\sum_{i=1}^{K}(\tilde{z}_{n,t}^{i}-z_{\text{FP}})^{T}\overline{\mathbf{H}}^{(z_{n,t}+(i-1)\epsilon)}(\tilde{z}_{n,t}^{i}-z_{\text{FP}})$
(6)

where $\mathbf{w}_{n}$ is the weight for layer $n$ and
$\tilde{z}_{n,t}^{i}=\tilde{z}_{n-1,t}^{i}\cdot\mathbf{w}_{n}$, $\epsilon$ is
a perturbation small enough for Taylor expansion, $K$ is a constant and
$\Delta=K\epsilon$.

Theorem 3.2 indicates that: To minimize the quantization error, ideally we can
finetune $\mathbf{w}_{n}$ so that for any $i$, the weights can fit to the
data/activation $\tilde{z}_{n-1,t}^{i}+(i-1)\epsilon$ which the full-precision
counterpart may not fit, making the model more robust to varying inputs. In
other words, we are finetuning the model weights for better robustness towards
large input activation perturbations, therefore enabling easier quantization.
We also find that the second term in Equation 3.2 can be minimized
simultaneously with the first term, so that the total quantization error can
be reduced via a common goal: Activation output alignment via finetuning
$\mathbf{w}_{n}$.

Figure 3: Effect of decreasing different activations’ bit-width on the model
performance in Stable Diffusion. FeedForward layers experience generation
failures at 6 bits, while all other linear layers barely fail at 4 bits and
all convolution layers only fail at 4 bits.

####  3.3.2 Quantization via Selective Finetuning

Motivated by the unraveled properties of quantized diffusion models and the
shortcomings in existing quantization approaches, we propose an efficient
finetuning strategy for low-bit diffusion model quantization. The components
of our method are detailed below.

Data-Free Partial-network-wise Training. To alleviate the need of substantial
training data, we propose to adopt the general setting of Post-Training
Quantization and construct the calibration set in a data-free manner. By
feeding sampled Gaussian noises $x_{T}$ into the full-precision model and
sample iteratively over different time steps using Equation 1, we can obtain
the calibration data needed for finetuning the quantized model. In practice,
the number of samples generated for each time step is 256, indicating that we
only have to inference the full-precision model a few times to obtain the
needed number of calibration samples.

Moreover, inspired by the fact that we can view the activation quantization
parameters as additional model parameters which all contribute to the final
output, we propose a partial-network-wise training strategy. Different from
quantization methods using layer-wise or block-wise reconstruction (Shang et
al., 2023; Li et al., 2023a) that binds quantization parameters with their
corresponding layers or blocks, we optimize all activation scaling factors for
each time step and only train partial weight parameters (e.g. 7% of the
parameters in LDM-4). As depicted in Figure 1, the majority of the model
weights remain unchanged at time step $t$, and only a small subset of
parameters (including layer weights and scaling factors) related to this time
step are updated using the full-precision model’s outputs as ground-truth. The
choices for the weights to be finetuned are discussed in the following
sections. Additionally, while layer/block-wise optimization methods can only
reconstruct sequentially, our approach update the required parameters at once.
In this way, we significantly save the time and memory needed for
quantization. Note that the quantization parameters for the weights are not
updated once initialized, in order to maintain memory efficiency during
deployment.

Table 3: Unconditional image generation results of LDM-4 models on LSUN-Bedrooms/Churches 256$\times$256\. ”N/A” denotes generation failure. ”FP” denotes the full precision model. Dataset | Method |  | Bit-width
---
(W/A)
| Size
---
(MB)
FID$\downarrow$
LSUN- Bedrooms (LDM-4) | FP | 32/32 | 1045.6 | 2.95
PTQ4DM | 8/8 | 279.1 | 4.75
Q-Diffusion | 8/8 | 279.1 | 4.53
PTQ-D | 8/8 | 279.1 | 3.75
Ours | 8/8 | 279.1 | 3.07
PTQ4DM | 4/8 | 148.4 | -
Q-Diffusion | 4/8 | 148.4 | 5.37
PTQ-D | 4/8 | 148.4 | 5.94
Ours | 4/8 | 148.4 | 3.26
PTQ4DM | 4/4 | 148.4 | N/A
Q-Diffusion | 4/4 | 148.4 | N/A
PTQ-D | 4/4 | 148.4 | N/A
Ours | 4/4 | 148.4 | 5.64
LSUN- Churches (LDM-8) | FP | 32/32 | 1125.4 | 4.02
PTQ4DM | 8/8 | 330.6 | 5.54
Q-Diffusion | 8/8 | 330.6 | 6.94
PTQ-D | 8/8 | 330.6 | 6.40
Ours | 8/8 | 330.6 | 6.55
PTQ4DM | 4/8 | 189.9 | -
Q-Diffusion | 4/8 | 189.9 | 7.80
PTQ-D | 4/8 | 189.9 | 7.33
Ours | 4/8 | 189.9 | 7.33
PTQ4DM | 4/4 | 189.9 | N/A
Q-Diffusion | 4/4 | 189.9 | N/A
PTQ-D | 4/4 | 189.9 | N/A
Ours | 4/4 | 189.9 | 11.76

Efficient Time-aware Activation Quantization. Temporal information is a unique
and essential factor in diffusion models. As highlighted in (Shang et al.,
2023; Li et al., 2023a), the variation of activation distribution over
different time steps is fatal for conventional PTQ methods. Also, as discussed
in Section 3.2, accurate time embedding quantization is critical. To
discriminate between different time steps, we design the Time-aware Activation
Quantizer (TAQuant) that adopts different sets of activation quantization
parameters for different time steps. While these parameters consumes
negligible memory and does not effect inference speed, learning independent
parameters for each time step leads to immense time cost. Inspired by the fact
that neighboring time steps yield similar functions in diffusion models (Choi
et al., 2022), we use the simple strategy of clustering neighboring time steps
so that they share the same calibration data and quantization parameters. Then
the set of activation quantization parameters are formulated as:

$\displaystyle\mathbf{s}=\\{\mathbf{s}^{t}(x_{t}),\mathbf{s}^{2t}(x_{2t}),...,\mathbf{s}^{T}(x_{T})\\},$
(7)

where the time steps used for optimization are sampled uniformly, and each
chosen time step’s statistics are used to represent $T/t$ adjacent ones.
Though most time steps’ quantization parameters are approximated, this
strategy saves the optimization time by $t$ times without compromising much
performance. Section 4.3 further provides an analysis on the trade-off between
time efficiency and model performance.

Selective & Progressive Layer Alignment. As investigated in Section 3.2, we
identify two crucial types of activations for image generation performance and
select them for weight finetuning: Time embeddings and attention-related
activations (e.g. FeedForward layers in Stable Diffusion). We progressively
update these components due to their distinct, non-overlapping
functionalities.

Table 4: Class-conditional image generation results of the LDM-4 model on ImageNet 256$\times$256, containing 1,000 categories. | Bit-width
---
(W/A)
Method |  | Size
---
(MB)
FID$\downarrow$ |  sFID$\downarrow$ |  IS$\uparrow$
32/32 | FP | 1529.7 | 11.28 | 7.70 | 364.73
8/8 | Q-Diffusion | 428.7 | 10.60 | 9.29 | 350.93
PTQ-D | 428.7 | 10.05 | 9.01 | 359.78
Ours | 428.7 | 10.43 | 6.07 | 365.12
4/8 | Q-Diffusion | 237.5 | 9.29 | 9.29 | 336.80
PTQ-D | 237.5 | 8.74 | 7.98 | 344.72
Ours | 237.5 | 8.48 | 6.55 | 346.02
4/4 | Q-Diffusion | 237.5 | N/A | N/A | N/A
PTQ-D | 237.5 | N/A | N/A | N/A
Ours | 237.5 | 5.98 | 7.93 | 202.45
Table 5: Text-to-image generation results (512$\times$512) of Stable Diffusion v1.4 using COCO2014 validation set prompts. Bit-width (W/A) | Method | Size (MB) |  CLIP Score$\uparrow$
---|---|---|---
32/32 | FP | 3279.1 | 31.50
8/8 | Q-Diffusion | 949.0 | 31.43
Ours | 949.0 | 31.47
4/8 | Q-Diffusion | 539.1 | 31.39
Ours | 539.1 | 31.50
4/4 | Q-Diffusion | 539.1 | N/A
Ours | 539.1 | 28.85

Figure 4: Generated images compared with Q-Diffusion under two bit-width
types: W8A8 and W4A4, using LDM-4 in LSUN-Bedrooms. Our method not only
succeed well in generating low-bit images, but also aligns more closely with
the full-precision counterpart.

During a single forward process, identical time embeddings are injected into
different parts of the model, passed through projection layers, and merged
with the latent image representations. This implies that the time information
operates independently from the primary network flow. Thus, we first refine
the time embedding layer $l$’s weight $\mathbf{w}_{l}$ along with its
activation quantization parameters $\mathbf{s}_{l}^{t}$ using:

$\displaystyle\mathcal{L}_{\text{TE}}^{t}=\sum_{l\in\mathbb{C}_{\text{TE}}}\mathbb{E}_{t}[||\overline{O}(e(t);\mathbf{w}_{l})-\tilde{O}(e(t);\mathbf{w}_{l},\mathbf{s}_{l}^{t})||^{2}],$
(8)

where $\mathbb{C}_{\text{TE}}$ represents the set of time embedding layers,
$e(t)$ is the fixed time embedding for time $t$,
$\overline{O}(e(t);\mathbf{w}_{l})$ is the activation output of the full-
precision model representing the ground-truth, and
$\tilde{O}(e(t);\mathbf{w}_{l},\mathbf{s}_{l}^{t})$ is the quantized output.
This loss function indicates that the chosen weight parameters are
consistently updated across different time steps to ensure robustness to
diverse inputs, whereas only the subset of scaling factors relevant to each
specific time step are adjusted for temporal information discrimination.

Apart from time embedding layers, attention-related layers are also vital. As
depicted in Figure 3, reducing activation bit-width of these layers degrades
model performance the most dramatically. Denote $\mathbb{C}_{\text{A}}$ as the
set containing all these layers and given the image calibration samples
$x_{t}$, we optimize $\mathbf{w}_{l}$ and $\mathbf{s}^{t}$ but exclude the
quantization parameters already optimized in Equation 8:

$\displaystyle\mkern-18.0mu\mathcal{L}_{\text{A}}^{t}=\sum_{l\in\mathbb{C}_{A}}\mathbb{E}_{t}[||\overline{O}(x_{t};\mathbf{w}_{:l})-\tilde{O}(x_{t};\mathbf{w}_{:l},\mathbf{s}^{t}\setminus
s_{\mathbb{C}_{\text{TE}}}^{t})||^{2}],\mkern-18.0mu$ (9)

where $\mathbf{w}_{:l}$ are the weight parameters of the first $l$ layers.
This strategy resembles linear probing, where only the layer directly
associated with the activation output is finetuned.

Furthermore, we aim not only for each independent layer’s quantization error
reduction, but also for the final generated images’ quality improvement.
Therefore, the target task loss $\sum_{t}\mathcal{L}_{d}^{t}$ is introduced,
where the ground-truth is approximated by the output of the full-precision
counterpart. By integrating Equation 8 and Equation 9, the final objective is
formulated as follows:

$\displaystyle\mathcal{L}$
$\displaystyle=\sum_{t}((\mathcal{L}_{\text{TE}}^{t}+\mathcal{L}_{d}^{t})+(\mathcal{L}_{\text{A}}^{t}+\mathcal{L}_{d}^{t}))$
$\displaystyle=\sum_{t}(\mathcal{L}_{\text{TE}}^{t}+\mathcal{L}_{\text{A}}^{t}+2\mathbb{E}[||\overline{O}(x_{t};\mathbf{w})-\tilde{O}(x_{t};\mathbf{w},\mathbf{s}^{t})||^{2}]).$
(10)

##  4 Experiments

###  4.1 Experiment Settings

To verify the effectiveness of our proposed method, we conduct experiments on
three types of generation tasks: Unconditional image generation on LSUN-
Bedrooms and LSUN-Churches datasets (Yu et al., 2015), class-conditional image
generation on ImageNet (Deng et al., 2009), and text-to-image generation. The
model architectures we quantize include LDM and Stable Diffusion (Rombach et
al., 2021), and uses ”WnAm” to represent the quantization setting: n-bit
weight quantization and m-bit activation quantization. DDIM samplers (Ho et
al., 2020) are adopted for LDMs and the PLMS sampler (Liu et al., 2022) is
used for Stable Diffusion. We generate 256 samples per time step for
constructing the calibration set, and train 20 epochs over each time step. The
Adam optimizer (Kingma & Ba, 2017) is adopted and the learning rate for weight
finetuning and scaling factor finetuning are set as $1e^{-5}$ and $1e^{-4}$
respectively.

The performance of different quantized LDMs are evaluated using the
Fr$\acute{\text{e}}$chet Inception Distance (FID) (Heusel et al., 2018),
spatial FID (sFID) (Nash et al., 2021) and Inception Score (IS) (Barratt &
Sharma, 2018). Unless specified, quantitative results are obtained by sampling
50,000 images from the models and evaluated using the official evaluation
scripts (Dhariwal & Nichol, 2021). For Stable Diffusion, we use the CLIP Score
(Hessel et al., 2022) for evaluation.All experiments are conducted on A6000
GPUs using the PyTorch framework (Paszke et al., 2019).

###  4.2 Experiment Results

Unconditional Generation: We evaluate the performance of our method over LDM-4
(LSUN-Bedrooms 256$\times$256) and LDM-8 (LSUN-Churches 256$\times$256) using
the DDIM sampler with 200 and 500 time steps, respectively. Results are shown
in table 3. We compare with methods including PTQ4DM (Shang et al., 2023),
Q-Diffusion (Li et al., 2023a) and PTQ-D (He et al., 2023b). The results are
evaluated using FID, where the Inception Score is not a reasonable metric for
datasets that have significantly different domains and categories from
ImageNet (Li et al., 2023a).

Table 6: Different component comparison and efficiency comparison on LDM-4 (LSUN-Bedrooms 256 $\times$ 256). Method |  | Bit-width
---
(W/A)
| Calibration
---
data size
| Time cost
---
(GPU hours)
| Memory cost
---
(MB)
| Model size
---
(MB)
FID$\downarrow$
FP | 32/32 | - | - | - | 1045.6 | 2.95
PTQ (Li et al., 2023a) | 4/4 | 5120 | 23.08 | 10334 | 148.4 | 428.8
TAQuant | 4/4 | 5120 | 11.52 | 9822 | 148.4 | 12.56
\+ $\text{SLA}_{\text{TE}}$ | 4/4 | 5120 | 13.13 | 12134 | 148.4 | 6.41
\+ $\text{SLA}_{\text{A}}$ | 4/4 | 5120 | 15.25 | 12450 | 148.4 | 5.78

Class-conditional Generation: We evaluate the performance using LDM-4 on
ImageNet 256$\times$256 using the DDIM sampler (20 steps). As shown in Table
4, three metrics are used for evaluation. Note that sFID uses additional
intermediate spatial features for calculation compared with FID. We can also
see that FID is not a valid metric for ImageNet LDM-4 evaluation: All methods
have lower FID when quantized to lower-bits, conflicting with human
perception. We show that our method not only succeeds in W4A4 quantization,
but also significantly improve the generation quality on higher bit settings.
Under W8A8, our method is able to achieve 2.82 improvement over sFID and 8.13
increase in Inception Score. Under W4A8, we improve the SOTA PTQ method by
1.57 sFID and 8.77 IS.

Text-to-image Generation: We use Stable Diffusion v1.4 as the model for
quantization with the PLMS sampler sampling 50 time steps. Table 5 shows the
results. Images are generated based on the 10,000 prompts sampled from the
COCO2014 (Lin et al., 2015) validation set, and CLIP Score is calculated based
on the ViT-B/16 backbone. Given the limited works done on Stable Diffusion, we
can only compare with Q-Diffusion and the full-precision baseline.

###  4.3 Experimental Analysis

Efficiency comparison and effects of different component. In Table 6, we
present a comparison of efficiency between our method and the post-training
quantization approach on the LSUN-Bedrooms dataset. Our approach, while
utilizing the same quantity of calibration data as the PTQ method, is more
time-efficient and incurs only a 20% increase in memory usage. We also show
how each proposed component improves the generation performance. The results
indicate that sequentially finetuning time embedding layers followed by
attention-related layers leads to consistent improvements in model
performance.

Table 7 presents the comparison of with and without the use of task loss
$\mathcal{L}_{d}$. The results show that using the difference between the
outputs of the quantized model and the full-precision counterpart for
supervision is essential for performance improvement, enhancing the FID by
2.58 and 5.21 for $\text{SLA}_{\text{TE}}$ and $\text{SLA}_{\text{A}}$,
respectively. However, when the learning process is only supervised by the
task loss, the performance degrades by 7.13 FID and 9.39 sFID for
$\text{SLA}_{\text{TE}}$, indicating that the task loss alone is insufficient.

Table 7: Influence of task loss supervision on model performance. Method |  | $\text{SLA}_{\text{TE}}$
---
w/o $\mathcal{L}_{d}$
| $\text{SLA}_{\text{TE}}$
---
w/ $\mathcal{L}_{d}$
| $\text{SLA}_{\text{A}}$
---
w/o $\mathcal{L}_{d}$
| $\text{SLA}_{\text{A}}$
---
w/ $\mathcal{L}_{d}$
FID$\downarrow$ | 8.99 | 6.41 | 11.41 | 6.20
sFID$\downarrow$ | 15.23 | 11.18 | 15.64 | 12.32

Time step sampling frequency analysis. While time steps are sampled uniformly,
will the performance consistently improve when we sample more frequently?
Table 8 illustrates the trade-off between the increased number of sampled time
steps (’T-Num’) and generation performance. Metrics are calculated using
10,000 samples generated from LDM-4 on LSUN-Bedrooms. The results show that
the though performance improves consistently, the amount of growth gradually
decreases. Starting from 5 time steps, using 15 more time steps leads to 25.22
sFID improvement. But using another 30 time steps (from 20 to 50) only leads
to 2.92 sFID improvement. Given the linearly increasing time cost associated
with using more time steps, we opt for only 20 out of 200 time steps to
enhance quantization efficiency.

Table 8: Effect of number of uniformly sampled time steps for W4A4 LDM-4 on LSUN-Bedrooms dataset. T-Num | 5 | 10 | 20 | 25 | 40 | 50
---|---|---|---|---|---|---
FID $\downarrow$ | 19.34 | 15.27 | 12.56 | 12.16 | 11.04 | 10.54
sFID $\downarrow$ | 54.00 | 37.47 | 28.78 | 27.99 | 26.42 | 25.86

How QuEST influences model properties. Our approach is based on the certain
characteristics observed in quantized diffusion models, hence we aim to
analyze how our fine-tuning strategy counters these identified challenges. As
shown in Table 2, finetuning time embedding layers leads to even better
performance than the full-precision counterpart. We infer that this might be
due to both the additional parameters introduced by quantization and the
additional task loss supervision. Figure 2 further implies that our method
adjusts the activation distribution so that it can be more easily quantized.
We see that the ranges of the activation value distributions shrink from [-10,
34] to [-4, 14] and [-11, 20] to [-4, 4]. The standard deviations also
decrease from 0.171 to 0.157 and from 0.073 to 0.071. This indicates a higher
concentration of values around the mean, which can effectively reduce both
rounding and clipping errors.

##  5 Conclusion

In this work, we propose QuEST, an efficient data-free finetuning framework
for low-bit diffusion model quantization. Our method is motivated by the three
underlying properties found in quantized diffusion models. We also
theoretically demonstrate the sufficiency of finetuning, interpreting it as a
way to enhance model robustness against large activation perturbations. To
alleviate the performance degradation, we propose to finetune the time
embedding layers and attention-related layers under the supervision of the
full-precision counterpart. A time-aware activation quantizer is also
introduced to handle varying time steps. Experimental results over three high-
resolution image generation tasks demonstrate the effectiveness and efficiency
of QuEST, achieving low-bit compatibility with less time and memory cost.

## References

  * Barratt & Sharma (2018) Barratt, S. and Sharma, R.  A note on the inception score, 2018.
  * Chitty-Venkata et al. (2023) Chitty-Venkata, K. T., Mittal, S., Emani, M., Vishwanath, V., and Somani, A. K.  A survey of techniques for optimizing transformer inference.  _Journal of Systems Architecture_ , pp. 102990, 2023.
  * Choi et al. (2022) Choi, J., Lee, J., Shin, C., Kim, S., Kim, H., and Yoon, S.  Perception prioritized training of diffusion models, 2022.
  * Croitoru et al. (2023) Croitoru, F.-A., Hondru, V., Ionescu, R. T., and Shah, M.  Diffusion models in vision: A survey.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 45(9):10850–10869, 2023.
  * Deng et al. (2009) Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. IEEE, 2009.
  * Dhariwal & Nichol (2021) Dhariwal, P. and Nichol, A.  Diffusion models beat gans on image synthesis, 2021.
  * Dosovitskiy et al. (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.  An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
  * Feng et al. (2023) Feng, Z., Zhang, Z., Yu, X., Fang, Y., Li, L., Chen, X., Lu, Y., Liu, J., Yin, W., Feng, S., Sun, Y., Chen, L., Tian, H., Wu, H., and Wang, H.  Ernie-vilg 2.0: Improving text-to-image diffusion model with knowledge-enhanced mixture-of-denoising-experts.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 10135–10145, June 2023.
  * Gholami et al. (2022) Gholami, A., Kim, S., Dong, Z., Yao, Z., Mahoney, M. W., and Keutzer, K.  A survey of quantization methods for efficient neural network inference.  In _Low-Power Computer Vision_ , pp. 291–326. Chapman and Hall/CRC, 2022.
  * He et al. (2015) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition, 2015.
  * He et al. (2023a) He, Y., Liu, J., Wu, W., Zhou, H., and Zhuang, B.  Efficientdm: Efficient quantization-aware fine-tuning of low-bit diffusion models, 2023a.
  * He et al. (2023b) He, Y., Liu, L., Liu, J., Wu, W., Zhou, H., and Zhuang, B.  Ptqd: Accurate post-training quantization for diffusion models, 2023b.
  * Hessel et al. (2022) Hessel, J., Holtzman, A., Forbes, M., Bras, R. L., and Choi, Y.  Clipscore: A reference-free evaluation metric for image captioning, 2022.
  * Heusel et al. (2018) Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S.  Gans trained by a two time-scale update rule converge to a local nash equilibrium, 2018.
  * Ho et al. (2020) Ho, J., Jain, A., and Abbeel, P.  Denoising diffusion probabilistic models, 2020.
  * Hu et al. (2021) Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W.  Lora: Low-rank adaptation of large language models, 2021.
  * Jacob et al. (2017) Jacob, B., Kligys, S., Chen, B., Zhu, M., Tang, M., Howard, A., Adam, H., and Kalenichenko, D.  Quantization and training of neural networks for efficient integer-arithmetic-only inference, 2017.
  * Kingma & Ba (2017) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization, 2017.
  * Li et al. (2023a) Li, X., Liu, Y., Lian, L., Yang, H., Dong, Z., Kang, D., Zhang, S., and Keutzer, K.  Q-diffusion: Quantizing diffusion models.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , pp. 17535–17545, October 2023a.
  * Li et al. (2021) Li, Y., Gong, R., Tan, X., Yang, Y., Hu, P., Zhang, Q., Yu, F., Wang, W., and Gu, S.  Brecq: Pushing the limit of post-training quantization by block reconstruction.  _arXiv preprint arXiv:2102.05426_ , 2021.
  * Li et al. (2022) Li, Y., Xu, S., Zhang, B., Cao, X., Gao, P., and Guo, G.  Q-vit: Accurate and fully quantized low-bit vision transformer, 2022.
  * Li et al. (2023b) Li, Y., Xu, S., Cao, X., Zhang, B., and Sun, X.  Q-dm: An efficient low-bit quantized diffusion model.  In _NeurIPS 2023_ , October 2023b.
  * Liang et al. (2021) Liang, T., Glossner, J., Wang, L., Shi, S., and Zhang, X.  Pruning and quantization for deep neural network acceleration: A survey.  _Neurocomputing_ , 461:370–403, 2021.
  * Lin et al. (2023) Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., Gan, C., and Han, S.  Awq: Activation-aware weight quantization for llm compression and acceleration, 2023.
  * Lin et al. (2015) Lin, T.-Y., Maire, M., Belongie, S., Bourdev, L., Girshick, R., Hays, J., Perona, P., Ramanan, D., Zitnick, C. L., and Dollár, P.  Microsoft coco: Common objects in context, 2015.
  * Liu et al. (2023) Liu, J., Niu, L., Yuan, Z., Yang, D., Wang, X., and Liu, W.  Pd-quant: Post-training quantization based on prediction difference metric, 2023.
  * Liu et al. (2022) Liu, L., Ren, Y., Lin, Z., and Zhao, Z.  Pseudo numerical methods for diffusion models on manifolds, 2022.
  * Nagel et al. (2020) Nagel, M., Amjad, R. A., Van Baalen, M., Louizos, C., and Blankevoort, T.  Up or down? adaptive rounding for post-training quantization.  In _International Conference on Machine Learning_ , pp. 7197–7206. PMLR, 2020.
  * Nahshan et al. (2021) Nahshan, Y., Chmiel, B., Baskin, C., Zheltonozhskii, E., Banner, R., Bronstein, A. M., and Mendelson, A.  Loss aware post-training quantization.  _Machine Learning_ , 110(11-12):3245–3262, 2021.
  * Nash et al. (2021) Nash, C., Menick, J., Dieleman, S., and Battaglia, P. W.  Generating images with sparse representations, 2021.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library, 2019.
  * Pilipović et al. (2018) Pilipović, R., Bulić, P., and Risojević, V.  Compression of convolutional neural networks: A short survey.  In _2018 17th International Symposium INFOTEH-JAHORINA (INFOTEH)_ , pp. 1–6. IEEE, 2018.
  * Rombach et al. (2021) Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.  High-resolution image synthesis with latent diffusion models, 2021.
  * Ronneberger et al. (2015) Ronneberger, O., Fischer, P., and Brox, T.  U-net: Convolutional networks for biomedical image segmentation, 2015.
  * Shang et al. (2023) Shang, Y., Yuan, Z., Xie, B., Wu, B., and Yan, Y.  Post-training quantization on diffusion models.  In _CVPR_ , 2023.
  * Wang et al. (2023) Wang, C., Wang, Z., Xu, X., Tang, Y., Zhou, J., and Lu, J.  Towards accurate data-free quantization for diffusion models, 2023.
  * Wang et al. (2020) Wang, P., Chen, Q., He, X., and Cheng, J.  Towards accurate post-training network quantization via bit-split and stitching.  In _International Conference on Machine Learning_ , pp. 9847–9856. PMLR, 2020.
  * Wei et al. (2023a) Wei, X., Gong, R., Li, Y., Liu, X., and Yu, F.  Qdrop: Randomly dropping quantization for extremely low-bit post-training quantization, 2023a.
  * Wei et al. (2023b) Wei, X., Zhang, Y., Zhang, X., Gong, R., Zhang, S., Zhang, Q., Yu, F., and Liu, X.  Outlier suppression: Pushing the limit of low-bit transformer language models, 2023b.
  * Xiao et al. (2023) Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S.  Smoothquant: Accurate and efficient post-training quantization for large language models, 2023.
  * Xu et al. (2023) Xu, S., Li, Y., Lin, M., Gao, P., Guo, G., Lu, J., and Zhang, B.  Q-detr: An efficient low-bit quantized detection transformer, 2023.
  * Yu et al. (2015) Yu, F., Zhang, Y., Song, S., Seff, A., and Xiao, J.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop.  _arXiv preprint arXiv:1506.03365_ , 2015.
  * Yuan et al. (2022) Yuan, Z., Xue, C., Chen, Y., Wu, Q., and Sun, G.  Ptq4vit: Post-training quantization framework for vision transformers with twin uniform quantization, 2022.


11institutetext: Computer Science and Engineering, University of Minnesota –
Twin Cities 11email: {helle246, deanx252<EMAIL_ADDRESS>

# Imperfect Segmentation Labels:
How Much Do They Matter?

Nicholas Heller  Joshua Dean  Nikolaos Papanikolopoulos

###### Abstract

Labeled datasets for semantic segmentation are imperfect, especially in
medical imaging where borders are often subtle or ill-defined. Little work has
been done to analyze the effect that label errors have on the performance of
segmentation methodologies. Here we present a large-scale study of model
performance in the presence of varying types and degrees of error in training
data. We trained U-Net, SegNet, and FCN32 several times for liver segmentation
with 10 different modes of ground-truth perturbation. Our results show that
for each architecture, performance steadily declines with boundary-localized
errors, however, U-Net was significantly more robust to jagged boundary errors
than the other architectures. We also found that each architecture was very
robust to non-boundary-localized errors, suggesting that boundary-localized
errors are fundamentally different and more challenging problem than random
label errors in a classification setting.

##  1 Introduction

Automatic semantic segmentation has wide applications in medicine, including
new visualization techniques [19], surgical simulation [10], and larger
studies of morphological features [5], all of which would remain prohibitively
expensive if segmentations were provided manually.

In the past 4 years, Deep Learning (DL) has risen to the forefront of semantic
segmentation techniques with virtually all segmentation challenges currently
dominated by DL-based entries [8]. Deep learning is a subfield of machine
learning which uses labeled input, or training data to learn functions that
map unlabeled input data to its correct response. In the case of semantic
segmentation, the model learns from image and mask pairs, where the mask
assigns each pixel or voxel to one of a set number of classes. These masks are
typically provided manually by a domain expert and often contain some errors.

Typically the most challenging and expensive task in using deep learning for
semantic segmentation is curating a ground-truth dataset that is sufficiently
large for the trained model to effectively generalize to unseen data.
Practitioners are often faced with a tradeoff between the quantity of ground-
truth masks and their quality [9].

We categorize ground truth errors to be either biased or unbiased. Biased
errors stem from errors of intention, where the expert creating the labels
would repeat the error if asked to label the instance again. These errors are
pernicious because they can result in systemic inaccuracies in the dataset
that may then be imparted to the learned model. These errors can often be
mitigated by giving clear and unambiguous instructions for those performing
the labeling.

Unbiased errors are all other types of errors. For instance, if an annotator’s
hand shakes when performing labeling, this would be an unbiased error so long
as his hand is not more likely to shake on certain features than on others. We
define the gold standard ground truth to be what an unbiased annotator would
produce if he were to annotate every instance an infinite number of times and
then take plurality votes to produce the final labels. For semantic
segmentation, each pixel would be an instance in this example.

Errors can be difficult to recognize in annotated images, but in medical
imaging, 3D imaging modalities such as Computed Tomography (CT) allow us to
scrutinize the annotations from the other anatomical planes. In Fig. 1 we can
clearly see that the expert is somewhat inconsistent in his treatment of the
region boundary in the axial plane, since there are clear discontinuities in
the contour when viewed from the saggital plane. This is important in medical
image processing because often models are trained on all three anatomical
planes to produce a more robust final segmentation [15], or volumetric models
are used [13]. It’s conceivable that this jagged boundary might confuse a
learned model by suggesting that the predicted segmentations should also have
jagged boundaries.

Figure 1: A sagittal cross-section of an annotation from the Pancreas
Segmentation Dataset [18] that was performed in the axial plane (best viewed
in color).

In this work, we study how errors in ground truth masks affect the performance
of trained models for the task of semantic segmentation in medical imaging. In
particular, we simulate ground truth errors in the widely used Liver
Segmentation Dataset111https://competitions.codalab.org/competitions/17094 by
perturbing the training annotations to various degrees in a ”natural”,
”choppy”, and ”random” way. The validation and testing annotations were left
untouched. We repeatedly train three widely used DL-based segmentation models
(U-Net [17], SegNet [3], and FCN32 [20]) on the perturbed training data and
report the corresponding degradation in performance.

##  2 Related Work

In [2] Angluin and Laird analyzed mislabeled examples from the standpoint of
Probably Approximately Correct learning. They show that the learning problem
remains feasible so long as the noise affects less than half of the instances
on average, although sample complexity increases with label noise.

Considerable work has been done to characterize the effect of label noise on
classical algorithms such as Decision Trees, Support Vector Machines, and
k-Nearest Neighbors, and robust variants of these have been proposed. For a
detailed survey, see [7]. Many data-cleansing algorithms have been proposed to
reduce the incidence of mislabeled data in datasets [14] [4] [21] but
challenges arise in distinguishing mislabeled instances from instances that
are difficult but informative.

With the rise to prominence of deep learning for computer vision tasks, the
ready availability of vast quantities of noisily labeled data on the internet,
and the lack of sufficient data-cleansing algorithms, many have turned their
attention to studying the pitfalls of training Deep Neural Networks for image
recognition, attribute learning, and scene classification using noisy labels.

In [22] the authors find that transfer learning from a noisy dataset to a
smaller but clean dataset for the same task does better than fine-tuning on
the clean dataset alone. They go on to extend Convolutional Neural Networks
with a probabilistic framework to model how mislabelings occur and infer true
labels with the Expectation Maximization algorithm.

In [16] the authors utilize what they call ”perceptual consistency”. They
argue this is implicit in the network parameters and that it holds an internal
representation of the world. Thus, it can serve as a basis for the network to
”disagree” with the provided labels and relabel data during training. The
network then ”bootstraps” itself in this way, using what it learns from the
relabeling as a basis to relabel more data, and so on.

These techniques are very robust to label noise in the image recognition but
label errors in semantic segmentation present a fundamentally different
problem, since label errors overwhelmingly occur at region boundaries, and no
such concept exists for holistic image analysis. In addition, learning in
semantic segmentation is done with fixed cohorts of pixels (images) within
random batches. Therefore, a DL model may learn a general rule about feasible
region size and discourage an otherwise positive prediction for a pixel in the
absence of positive predictions for its neighbors.

##  3 Methods

###  3.1 Perturbations

We attempted to perturb ground truth masks such that they closely mimicked the
sorts of errors that human experts often make when drawing freehand contours.
In order to achieve this, we first retrieved the contours from an existing
binary mask using OpenCV’s `findContours()` function. We then sampled points
from this contour and moved them a random offset either towards or away from
the contour’s center. We used a simple fill to produce the perturbed
annotation. The offsets were produced by a normal distribution with a given
variance and zero mean. We call these offsets natural perturbations. A natural
perturbation applied to a circle can be seen in Fig. 2 (middle-left).

In addition, we wanted to mimic the sort of errors that occur when natural
errors are made in a single plane of a volume and data is viewed from an
orthogonal plane, as seen in Fig. 1. For this, we iterated over every row in
the masks, found each block of consecutive positive labels, and shifted the
block’s starting and end points by some amount that was once again sampled
from a normal distribution with zero mean and provided variance. We call these
choppy perturbations. A choppy perturbation applied to a circle can be seen in
Fig. 2 (middle-right).

Finally, in order to simulate random errors in a classification setting, we
randomly chose an equal proportion of voxels from both the negative and
positive classes and flipped their values. We call these random perturbations
(Fig. 2, right).

Three parameter settings were chosen for each perturbation mode in order to
produce perturbed ground truth with 0.95, 0.90, and 0.85 Dice-Sorensen
agreement with the original ground truth, i.e. we chose 9 total parameter
setting. Each was tuned by randomly choosing 1000 slices and using bisection
with the terminal condition that the upper and lower bounds each producing
Dice scores within 0.005 of the target.

Figure 2: From left to right: unperturbed, natural perturbations, choppy
perturbations, and random perturbations.

###  3.2 Training

We ran the experiments using the Keras [6] framework with a TensorFlow [1]
back-end. We optimized our models using the Adam algorithm [11] with the
default parameter values. We addressed the imbalance of the problem by equally
sampling from each class, and we used mini-batches of 20 slices, where each
slice is a 512x512 array of Hounsfield Units from the axial plane. For each
model, we started with 6 initial convolutional kernels and the number doubled
with each down-sampling. Each model was trained for 100 epochs with 35 steps
per epoch.

For each architecture and perturbation pair, we trained five times in order to
improve statistical power, resulting in 150 total training sessions.

##  4 Results

Our results show that the performance of each model steadily declined with the
extent of boundary-localized perturbations, but that model performance was
very robust to random perturbations. This suggests that flawed ground truth
labels, particularly in border regions, are hindering the performance of DL-
based models for semantic segmentation.

As can be seen in Fig. 3, other than the large choppy perturbations for U-Net,
the responses of each architecture to the different degrees of boundary-
localized perturbation were surprisingly uniform. This suggests that there may
be a general predictive relationship between the incidence of ground-truth
errors and the expected performance of these models. Additionally, it appears
that each of the models are very resilient to random perturbations in ground
truth, in some cases outperforming the Dice-Sorenson score of the training
data itself by more than 5%.

| U-Net | SegNet | FCN32
---|---|---|---
Control | 0.9134 | 0.8993 | 0.8870
Natural 0.95 | 0.8880 | 0.8640 | 0.8587
Natural 0.90 | 0.8193 | 0.8265 | 0.8250
Natural 0.85 | 0.7521 | 0.7717 | 0.7581
Choppy 0.95 | 0.8928 | 0.8799 | 0.8660
Choppy 0.90 | 0.8321 | 0.8268 | 0.8202
Choppy 0.85 | 0.8058 | 0.7823 | 0.7782
Random 0.95 | 0.9124 | 0.9050 | 0.8881
Random 0.90 | 0.9213 | 0.9013 | 0.8751
Random 0.85 | 0.9182 | 0.9068 | 0.8676
Table 1: Mean Dice-Sorensen score for each model-perturbation pair for Liver
Segmentation Figure 3: The results for liver segmentation for each model with
each type of mode of ground-truth perturbation (best viewed in color).

U-Net’s anomalously good performance in the presence of large choppy
perturbations is interesting. We hypothesize that this is because U-Net’s
”skip connections” allow it to very effectively preserve border information
from activation functions early on in the network. Thus, borders are likely
still emphasized because even though the contour has become jagged, the region
edges are centered on the true contour. This is not the case for the ”natural”
perturbations.

##  5 Limitations

Better performance has been reported for the liver segmentation problem [12],
but that is due to the use of ensembles and hyperparameter tuning. It would
not be feasible to engineer and train such techniques for each and every data
point. It is possible (although we believe unlikely) that these findings do
not translate to large ensemble settings, but this must be the subject for
future work.

Additionally, these experiments were all run on a single dataset with binary
labels. More work must be done to study whether these results generalize to
different problems, and problems with many class labels.

Finally, this study did not examine the effect of biased labels, which are
also likely to exist in semantic segmentation datasets. Our intuition is that
the models will tend to exhibit the same bias as the expert, but it’s unclear
what the effect on performance would be when there are multiple experts, each
with different biases. This, too, must be the subject for future work.

##  6 Conclusion

In this work we tested how three widely-used deep learning based models
responded to various modes of errors in ground-truth labels for semantic
segmentation of the liver in abdominal CT scans. We found that in general,
these models each experience relatively uniform performance degradation with
increased incidence of label errors, but that U-Net was especially robust to
large amounts of ”choppy” noise on the liver regions.

There are many opportunities to continue this work. In particular, we would
like to expand the scope of this study to look also at how the hyperparameters
of the architectures and training procedures affect its sensitivity. We also
believe it would be useful to explore the effect of dataset size on
sensitivity, since it’s possible that models will have a more difficult time
coping with noisy data when they have less data to look at. Finally, we plan
to study how deep-learning-based architectures for semantic segmentation can
be modified in order to be more robust to ground truth errors at region
boundaries.

The code for our experiments has been made available at
https://github.com/neheller/labels18.

## Acknowledgements

Research reported in this publication was supported by the National Cancer
Institute of the National Institutes of Health under Award Number R01CA225435.
The content is solely the responsibility of the authors and does not
necessarily represent the official views of the National Institutes of Health.

## References

  * [1] Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., et al.: Tensorflow: A system for large-scale machine learning. In: OSDI. vol. 16, pp. 265–283 (2016)
  * [2] Angluin, D., Laird, P.: Learning From Noisy Examples. Machine Learning 2(4), 343–370 (1988). https://doi.org/10.1023/A:1022873112823
  * [3] Badrinarayanan, V., Kendall, A., Cipolla, R.: Segnet: A deep convolutional encoder-decoder architecture for image segmentation. CoRR abs/1511.00561 (2015), http://arxiv.org/abs/1511.00561
  * [4] Brodley, C.E., Friedl, M.A.: Identifying mislabeled training data. Journal of artificial intelligence research 11, 131–167 (1999)
  * [5] Chang, R.F., Wu, W.J., Moon, W.K., Chen, D.R.: Automatic ultrasound segmentation and morphology based diagnosis of solid breast tumors. Breast Cancer Research and Treatment 89(2), 179 (Jan 2005). https://doi.org/10.1007/s10549-004-2043-z, https://doi.org/10.1007/s10549-004-2043-z
  * [6] Chollet, F., et al.: Keras. https://keras.io (2015)
  * [7] Frénay, B., Verleysen, M.: Classification in the presence of label noise: A survey. IEEE Transactions on Neural Networks and Learning Systems 25(5), 845–869 (2014). https://doi.org/10.1109/TNNLS.2013.2292894
  * [8] Guo, Y., Liu, Y., Oerlemans, A., Lao, S., Wu, S., Lew, M.S.: Deep learning for visual understanding: A review. Neurocomputing 187, 27–48 (2016)
  * [9] Heller, N., Stanitsas, P., Morellas, V., Papanikolopoulos, N.: A web-based platform for distributed annotation of computerized tomography scans. Intravascular Imaging and Computer Assisted Stenting, and Large-Scale Annotation of Biomedical Data and Expert Label Synthesis 10552, 136–145 (2017). https://doi.org/10.1007/978-3-319-67534-3, http://link.springer.com/10.1007/978-3-319-67534-3
  * [10] Huff, T.J., Ludwig, P.E., Zuniga, J.M.: The potential for machine learning algorithms to improve and reduce the cost of 3-dimensional printing for surgical planning. Expert Review of Medical Devices 15(5), 349–356 (2018). https://doi.org/10.1080/17434440.2018.1473033, https://doi.org/10.1080/17434440.2018.1473033, pMID: 29723481
  * [11] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 (2014)
  * [12] Le, T.N., Huynh, H.T., et al.: Liver tumor segmentation from mr images using 3d fast marching algorithm and single hidden layer feedforward neural network. BioMed research international 2016 (2016)
  * [13] Milletari, F., Navab, N., Ahmadi, S.A.: V-net: Fully convolutional neural networks for volumetric medical image segmentation. In: 3D Vision (3DV), 2016 Fourth International Conference on. pp. 565–571. IEEE (2016)
  * [14] Muhlenbach, F., Zighed, D.A.: Relabeling Mislabeled Instances pp. 5–15 (2002)
  * [15] Prasoon, A., Petersen, K., Igel, C., Lauze, F., Dam, E., Nielsen, M.: Deep feature learning for knee cartilage segmentation using a triplanar convolutional neural network. In: International conference on medical image computing and computer-assisted intervention. pp. 246–253. Springer (2013)
  * [16] Reed, S., Lee, H., Anguelov, D., Szegedy, C., Erhan, D., Rabinovich, A.: Training Deep Neural Networks on Noisy Labels with Bootstrapping pp. 1–11 (2014). https://doi.org/10.2200/S00196ED1V01Y200906AIM006, http://arxiv.org/abs/1412.6596
  * [17] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. CoRR abs/1505.04597 (2015), http://arxiv.org/abs/1505.04597
  * [18] Roth, H.R., Lu, L., Farag, A., Shin, H.C., Liu, J., Turkbey, E.B., Summers, R.M.: Deeporgan: Multi-level deep convolutional networks for automated pancreas segmentation. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 556–564. Springer (2015)
  * [19] Roth, H.R., Oda, H., Zhou, X., Shimizu, N., Yang, Y., Hayashi, Y., Oda, M., Fujiwara, M., Misawa, K., Mori, K.: An application of cascaded 3d fully convolutional networks for medical image segmentation. CoRR abs/1803.05431 (2018), http://arxiv.org/abs/1803.05431
  * [20] Shelhamer, E., Long, J., Darrell, T.: Fully convolutional networks for semantic segmentation. CoRR abs/1605.06211 (2016), http://arxiv.org/abs/1605.06211
  * [21] Verbaeten, S., Van Assche, A.: Ensemble methods for noise elimination in classification problems. In: International Workshop on Multiple Classifier Systems. pp. 317–325. Springer (2003)
  * [22] Xiao, T., Xia, T., Yang, Y., Huang, C., Wang, X.: Learning from massive noisy labeled data for image classification. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition 07-12-June, 2691–2699 (2015). https://doi.org/10.1109/CVPR.2015.7298885


# vieCap4H Challenge 2021: Vietnamese Image Captioning for Healthcare Domain
using Swin Transformer and Attention-based LSTM

Thanh Tin Nguyen1 Long H. Nguyen2 Nhat Truong Pham3,4, Liu Tai Nguyen2 Van
Huong Do5 Hai Nguyen6 Ngoc Duy Nguyen7 1 Human Computer Interaction Lab,
Sejong University, Seoul, South Korea
2 Faculty of Information Technology, Ton Duc Thang University, Ho Chi Minh
City, Vietnam
3 Division of Computational Mechatronics, Institute for Computational Science,
Ton Duc Thang University, Ho Chi Minh City, Vietnam
4 Faculty of Electrical and Electronics Engineering, Ton Duc Thang University,
Ho Chi Minh City, Vietnam
5 ASICLAND, Suwon, South Korea
6 Khoury College of Computer Sciences, Northeastern University, Boston, USA
7 Institute for Intelligent Systems Research and Innovation,
Deakin University, Victoria, Australia

###### Abstract

This study presents our approach on the automatic Vietnamese image captioning
for healthcare domain in text processing tasks of Vietnamese Language and
Speech Processing (VLSP) Challenge 2021, as shown in Figure 1. In recent
years, image captioning often employs a convolutional neural network-based
architecture as an encoder and a long short-term memory (LSTM) as a decoder to
generate sentences. These models perform remarkably well in different
datasets. Our proposed model also has an encoder and a decoder, but we instead
use a Swin Transformer in the encoder, and a LSTM combined with an attention
module in the decoder. The study presents our training experiments and
techniques used during the competition. Our model achieves a BLEU4 score of
0.293 on the vietCap4H dataset, and the score is ranked the 3rd place on the
private leaderboard. Our code can be found at https://git.io/JDdJm.

###### keywords:

image captioning , swin transformer , encoder-decoder , covid-19 , health ,
Vietnamese.

Figure 1: A sample image of the image captioning task. The image is described
by three Vietnamese sentences generated by our model.

##  1 Introduction

Generating meaningful captions for images is recently a challenging topic in
artificial intelligence (AI). The task involves both natural language
processing (NLP) and computer vision (CV) techniques because it requires the
machine to understand an image and translate the understanding into a
meaningful caption. Solving the problem leads to several practical
applications such as virtual assistants for blind and visually impaired
people, conducting visual content indexing and searching. Recently, Vu et al.
[1] also proposed another application of image captioning to generate online
personalized reviews via a multimodal approach, such as text, image, and
ratings. This helps to understand users’ behavior because the generated
reviews are based on users’ privacy and fairness.

Although the image captioning task has been tackled by a variety of
techniques, there are little research on Vietnamese domain. To encourage
conducting research on Vietnamese image captioning, [2] created a dataset for
Vietnamese domain, also serving as a premise for researching on Vietnamese
image captioning for healthcare domain.

The vieCap4H Challenge 2021 [3] aims to be a competition for developing
machine learning algorithms that use Vietnamese to describe the visual content
in healthcare settings, especially images that describe the COVID-19 pandemic.
Similar to this task, the most recent studies were presented in [4] and [5]
that proposed a network involving a deep Convolutional Neural Network (CNN) as
an encoder and a Recurrent Neural Network (RNN) as a decoder. While the
encoder is in charge of extracting features from images, the decoder takes
these features as input and infers a descriptive caption. The study in [5] is
an extension of the study in [4] by adding an attention mechanism into the
decoder part. Inspired by the these studies, we employ a specialized encoder
architecture called Swin Transformer [6], and keep the decoder intact as
proposed in [5].

The rest of the paper is organized as follows. Section 2 presents related
studies that we have investigated. Section 3 describes our proposed method and
the different techniques used during the competition. Analytical results are
presented in Section 4, and a detailed discussion is presented in Section 5.
Finally, Section 6 concludes the paper and outlines potential future
directions.

##  2 Related Work

Our approach was initially inspired by the early works in [4] and [5] for
image captioning task. Additionally, we also reference the code from [7].
Recently introduced ideas of using CNNs and RNNs can be found in [4] and [5].
Some proposed bottom-up methods are mentioned in [8], [9], or [10]. These
bottom-up methods all require to incorporate an object detection module to
extract object proposals, then give them to the encoder.

Although CNN has a profound impact on CV, Transformer, first introduced in the
paper “Attention is all you need" [11], has progressively replaced CNN in this
field. There are plenty of variants of Transformer, but the two most well-
known used for extracting visual contents are the Vision Transformer [12] and
the Swin Transformer [6] architecture.

Due to the time and resource limitation, our team only improved the encoder
side by using the Swin Transformer.

##  3 Methodology

###  3.1 Attempted Techniques

Our attempts in this competition consist of the following:

  * 1.

Changing the architecture of the encoder by using different CNN architecture.

  * 2.

In the decoder, RNN and Transformer were adopted to extract a caption for an
image.

  * 3.

Tuning the image size, the learning rate, and the learning rate scheduler.

  * 4.

Scrubbing the ground-truth captions, as well as attempting to use a pre-
trained embedding for text.

  * 5.

Augmenting images and text.

  * 6.

Using noise injection.

  * 7.

Experimenting beam search with different beam widths.

  * 8.

Using cross validation.

  * 9.

Attempting ensemble models.

We start from the model proposed by [4]. Moreover, we also test the idea of
predicting the English captions of images and then translating captions from
English to Vietnamese. However, this approach failed because translation
cannot capture the semantic and relational information of a sentence.

###  3.2 Pre-processing

The data in the challenge consists of two training datasets containing 8,032
images and a validation dataset containing 1,002 images. Both datasets are
used in the public phase of the competition. Additionally, there is another
testing dataset containing 1,032 images.

In the public training dataset, there are a few images with defective
captions, such as repeated or redundant characters in a word, two consecutive
words are stick together, or English instead of Vietnamese captions. Due to
these erroneous labels, the initial performance of the model is poor and
unstable. After fixing the errors, we also convert all of the captions to
lowercase and remove all existing punctuation and numbers.

Figure 2: Our proposed model has two main components: (1) The encoder uses
Swin Transformer pre-trained on the ImageNet [13] dataset; (2) The decoder
consists several layers of LSTM.

###  3.3 Model Architecture

Our proposed model as shown in Figure 2 is an end-to-end integration of a Swin
Transformer encoder and attention-based LSTM decoder. These architectures were
chosen because, firstly, the Swin Transformer presents a structure in which an
image is divided into multiple patches and each patch is divided into multiple
windows to perform self-attention within the window. Also, it allows a
hierarchical structure resulting in a good performance in several tasks such
as object detection, and image segmentation. Secondly, the LSTM is well-known
and straightforward for capturing the semantic meaning of natural languages.

The model takes in an image and outputs the corresponding caption as follows.
Firstly, the pre-processed image with the size of 224$\times$224 is fed into
the pre-trained Swin Transformer encoder to extract visual features. Then
these extracted visual features are used as the inputs of the attention-based
LSTM decoder to generate the caption. This attention step is called late
fusion because it aggregates two features extracted from two different
architectures. More concretely, the attention module calculates how much
attention to a specific hidden vector is given to an output hidden vector. In
this case, $h_{T}$ is the last hidden vector of the LSTM, $h_{output}$ is the
output of the encoder which is Swin Transformer, after computing attention, it
will output $h_{attended}$ state, this is the representation vector for
decoding step. The self-attention can be seen in Fig. 3. Finally, post-
processing is applied using beam search to improve the accuracy of the
proposed model.

Figure 3: Our attention module: The hidden state $h_{n}$ is multiplied with
the encoded state $h_{output}$, then pass the output through a softmax
function to create a weighted distribution over the $h_{output}$. Finally,
$h_{attended}$ is computed by multiplying the weighted distribution with the
$h_{output}$ to highlight salient regions.

###  3.4 Noise Injection

In the field of image captioning, the accuracy of prediction for the next
character can be relatively high if we can ensure that the previous sequence
is predicted correctly. However, if a character is mispredicted, then the
prediction failure rate will gradually increase.

To tackle the above issue, we use noise injection, which is done by randomly
replacing ground truth characters with other characters during training, and
new sentences will be assigned fake labels called $fake$. Then, using the
modified sentences as the input for the decoder will force it to correctly
predict the next character on the basis that the previous one was wrong.
Specifically, there will be two losses, the cross-entropy loss between the
ground truth and the prediction, and the same loss however between the $fake$
and their predictions. In the final loss, we combine the two losses with a
weighting $\beta=0.1$ for the second loss

$\mathcal{L}_{CE}=-\sum_{i=1}^{N}y^{(i)}*\log{\hat{y}^{(i)}}-\beta\sum_{i=1}^{N}y^{(i)}*\log{\hat{y}^{(i)}_{fake}},$
(1)

where $\mathcal{L}_{CE}$ is the Cross Entropy loss, $y^{(i)}$ and
$\hat{y}^{(i)}$ denote a label and its predicted probability, respectively.
$N$ represents the total number of classes.

###  3.5 Beam Search

Beam search is a well-known search algorithm adopted in several problems such
as Image Captioning, Sequence to Sequence. Normally, at test time, in order to
generate the output text, greedy search is employed to get the maximum token
at each time $t$. This algorithm can not assure to extract the best sequence,
because of that, beam search comes into rescue to keep a set of $k$ (beam
size) best tokens at each time step. This results in a best approximately
sentence. We tried beam search sizes from 1 to 10, and the best BLEU4 score is
achieved with a beam size of 2. Thus, we chose 2 as the beam search size.

###  3.6 Training Configuration

All experiments were trained on a single Titan Xp GPU. The batch size is 16,
the input image size is 224$\times$224, the learning rate of the encoder is
1e-4, that of the decoder is 4e-4, the Adam optimizer is used in this model
with a weight decay of 1e-6. Moreover, the Cosine Annealing Warm Restarts [14]
scheduler is used for scheduling the learning rate. In addition, the model
uses $k$-fold cross validation with $k=4$. We also use common augmentation
techniques such as HorizontalFlip, RandomCrop with probability of 0.5, and
Normalize with mean and std are (0.485, 0.456, 0.406), (0.229, 0.224, 0.225),
respectively. Finally, our model uses a cross entropy as the loss function.

##  4 Analysis

###  4.1 Results

#### BLEU4 Scores:

The model was tested on a private test set which contains 1,032 images. In
this competition, the BLEU4 score is used as the metric to evaluate models.
Our model achieves the BLEU4 score of 0.293 and is ranked the 3rd on the
private leaderboard. In terms of the public test dataset, which contains 1,002
images, the highest BLEU4 score evaluated on this dataset is 0.302. The
results during the public phase are summarized in Table 1.

Table 1:  The table shows the performance of our approach by using a CNN-based model to the Transformer model evaluated on the public dataset. By using additional techniques such as beam search, noise injection, and augmentation, we achieved the highest score on the public test dataset, which is 0.302. Encoder |  Decoder |  Additional Methods |  BLEU4
---|---|---|---
Resnet101 |  LSTM+Attention |  Beam search |  0.253
Efficientnetv2 |  LSTM+Attention |  Beam search |  0.266
Vision Transformer |  LSTM+Attention |  Beam search |  0.277
Swin Transformer |  LSTM+Attention |  Beam search |  0.292
Swin Transformer |  LSTM+Attention |  Beam search + Noise Injection + Augmentation |  0.302

#### Computational processing times and resources:

In this study, the proposed method is implemented with one GPU Titan Xp.
Besides, our proposed method takes about 4 hours to train with 4-fold cross
validation.

###  4.2 Visualization

Visualizing attention maps is to explain how the model learns to concentrate
on different parts of an image with its corresponding words. These attention
images are shown in Figure 4. For example, the caption of the image in Figure
4a is ‘Các chiếc khẩu trang được xếp chồng lên nhau’ in which the words ‘khẩu’
and ‘trang’ are attending in the middle of the image where exactly the object
‘khẩu trang’ is located. Figure 4b describes a mistaken caption of its
corresponding image. The caption of the image in Figure 4b is ‘Một người phụ
nữ đang đứng cạnh một chiếc bàn’, however, there is no ‘một người phụ nữ’ (a
woman) in the image.

(a) Attention visualization of ‘Các chiếc khẩu trang được xếp chồng lên nhau’
sample.

(b) Attention visualization of ‘Một người phụ nữ đang đứng cạnh một chiếc bàn’
sample.

Figure 4: Attention visualization of test samples. In Figure 4a, the words
‘khẩu’ and ‘trang’ are attending in the middle of the image where exactly the
object ‘khẩu trang’ is located. Figure 4b describes a mistaken caption of its
corresponding image.

###  4.3 Ablation tests

We have performed several ablation tests before perfecting the final model.
Table 2 shows our analysis of the contribution of each process in the model
including: without pre-processing data, without beam search, without noise
injection, and the final model.

Table 2: The table shows the score of our final model and when removing important parts in the model evaluated on the public dataset. Model |  BLEU4
---|---
Without pre-processing data |  0.273
Without beam search |  0.286
Without noise injection |  0.293
Our final model |  0.302

###  4.4 Discussion

We found that the provided data has some minor errors such as false captions,
redundant words, and typos, and after fixing all these errors, the accuracy
increased a little bit. Beam search is a must for efficient decoding, in fact,
it helps us boost the performance. Finally, with the noise injection in the
loss function, the model can predict the next character more properly. In
conclusion, following the encoder-decoder manner, we adopted two well-known
architectures which are Swin Transformer and LSTM, and the result has
demonstrated that our proposed model is a powerful tool for image captioning.

##  5 Lessons Learned

We have tried a variety of techniques, some perform well, and some do not. In
this section, we describe different techniques that are used during the
competition.

###  5.1 Working Techniques

Firstly, we examine the training data and apply suitable adjustments to
defective captions. During the tuning process, we use a grid search of
different learning rates and image sizes.

Swin Transformer [6] is also a key factor that outperforms most CNN-based
models, including Efficientnet [15], Efficientnet v2 [16], Restnet 50 [17],
and Resnet 101 [17] concerning the accuracy, training time, and memory
capacity. It also outperforms the Vision Transformer [12], which is another
variant of a Transformer.

Augmentation is a promising method to ensure the model capacitates with
different transformations of the images. In this technique, we use a Resize
factor of 224 for each dimension, a HorizontalFlip factor of 0.5, and a
RandomCrop factor of 0.5. Finally, a Normalization factor is used with a mean
of (0.485, 0.456, 0.406) and a standard deviation of (0.229, 0.224, 0.225).

Noise injection is a regularization technique used to randomly replace ground
truth characters with other characters during training to create new
sentences. This technique helps to improve the prediction of the next
character while the previous one is imperfect.

Beam search is a technique that can seek the best potential sentence while
decoding the features. When cross-validating, we run our model on four
different subsets of the data to achieve different outcome measures. Although
it can take a long time, we can achieve a better result.

###  5.2 Non-working Techniques

We conduct the performance evaluation by using two different image sizes:
224$\times$224 and 384$\times$384\. It is surprising that an image size of
224$\times$224 provides a better result compared with an image size of
384$\times$384.

Various CNN-based models such as Efficientnet [15], Efficientnetv2 [16],
Resnet 50 [17], Resnet 101 [17], or Vision Transformer [12] (ViT) have been
evaluated as the encoder. They consume a higher memory and have low
performance in terms of accuracy and training time.

Although pre-trained word embedding models, e.g., [18] and [19], have trained
on large corpus in Vietnamese, we failed to incorporate it into our model due
to a small embedding size of the pre-trained ones.

We also employed the Transformer as the decoder, but the achieved accuracy was
not as good as that achieved by the LSTM. Moreover, we also failed at
implementing beam search for the Transformer decoder due to its complicated
architecture.

Finally, ensemble is a regularization technique that helps boost overall
accuracy, but unfortunately, we failed to implement it because of lack of
time.

###  5.3 Unique domain/data specific insights that you have uncovered

Noise injection, beam search, and cross validation are unique techniques that
helps to improve the accuracy. Moreover, carefully screening the captions also
helps the model learn better. Although these are not new and novel to
somebody, but these still gave us surprises, and we have learnt a lot from it.

##  6 Conclusion

This paper describes our approach in the automatic Vietnamese image captioning
for healthcare domain in text processing tasks of the VLSP Challenge 2021. The
proposed model employs an encoder-decoder architecture, and experimental
results show a potentially useful network for tackling the problem.
Furthermore, techniques such as noise injection, beam search, and cross
validation help boost the algorithm’s performance during the competition.
Finally, cleaning and pre-processing the data improves the performance of the
algorithm.

One potential direction to further improve the performance of the decoder such
as using a Transformer-based decoder. Additionally, using beam search while
decoding, there are repeated words. Fixing this problem would potentially
increase the performance.

## References

  * [1] X.-S. Vu, T.-S. Nguyen, D.-T. Le, L. Jiang, Multimodal review generation with privacy and fairness awareness, in: 28th International Conference on Computational Linguistics (COLING), Barcelona, Spain (Online), December 8-13, 2020., International Committee on Computational LinguisticsInternational Committee on Computational Linguistics, 2020, pp. 414–425.
  * [2] Q. H. Lam, Q. D. Le, V. K. Nguyen, N. L.-T. Nguyen, Uit-viic: A dataset for the first evaluation on vietnamese image captioning, in: International Conference on Computational Collective Intelligence, Springer, 2020, pp. 730–742.
  * [3] T. M. Le, L. H. Dang, T.-S. Nguyen, T. M. H. Nguyen, X.-S. Vu, Vlsp 2021 - viecap4h challenge: Automatic image caption generation for healthcare domain in vietnamese, in: Proceedings of the 8th International Workshop on Vietnamese Language and Speech Processing, VNU Journal of Science: Computer Science and Communication Engineering, HCM, Vietnam, 2021.
  * [4] O. Vinyals, A. Toshev, S. Bengio, D. Erhan, Show and tell: A neural image caption generator, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2015, pp. 3156–3164.
  * [5] K. Xu, J. Ba, R. Kiros, K. Cho, A. Courville, R. Salakhudinov, R. Zemel, Y. Bengio, Show, attend and tell: Neural image caption generation with visual attention, in: International conference on machine learning, PMLR, 2015, pp. 2048–2057.
  * [6] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, B. Guo, Swin transformer: Hierarchical vision transformer using shifted windows, arXiv preprint arXiv:2103.14030.
  * [7] S. Vinodababu, https://github.com/sgrvinod/a-p-
ytorch-tutorial-to-image-captioning.

  * [8] M. Pedersoli, T. Lucas, C. Schmid, J. Verbeek, Areas of attention for image captioning, in: Proceedings of the IEEE international conference on computer vision, 2017, pp. 1242–1250.
  * [9] S. Herdade, A. Kappeler, K. Boakye, J. Soares, Image captioning: Transforming objects into words, arXiv preprint arXiv:1906.05963.
  * [10] P. Anderson, X. He, C. Buehler, D. Teney, M. Johnson, S. Gould, L. Zhang, Bottom-up and top-down attention for image captioning and visual question answering, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 6077–6086.
  * [11] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, Attention is all you need, in: Advances in neural information processing systems, 2017, pp. 5998–6008.
  * [12] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al., An image is worth 16x16 words: Transformers for image recognition at scale, arXiv preprint arXiv:2010.11929.
  * [13] A. Krizhevsky, I. Sutskever, G. E. Hinton, Imagenet classification with deep convolutional neural networks, Advances in neural information processing systems 25 (2012) 1097–1105.
  * [14] I. Loshchilov, F. Hutter, Sgdr: Stochastic gradient descent with warm restarts, arXiv preprint arXiv:1608.03983.
  * [15] M. Tan, Q. Le, Efficientnet: Rethinking model scaling for convolutional neural networks, in: International Conference on Machine Learning, PMLR, 2019, pp. 6105–6114.
  * [16] M. Tan, Q. V. Le, Efficientnetv2: Smaller models and faster training, arXiv preprint arXiv:2104.00298.
  * [17] K. He, X. Zhang, S. Ren, J. Sun, Deep residual learning for image recognition, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778.
  * [18] P. Bojanowski, E. Grave, A. Joulin, T. Mikolov, Enriching word vectors with subword information, Transactions of the Association for Computational Linguistics 5 (2017) 135–146.
  * [19] A. T. Nguyen, M. H. Dao, D. Q. Nguyen, A pilot study of text-to-sql semantic parsing for vietnamese, arXiv preprint arXiv:2010.01891.