# Exploring Quantization for Efficient Pre-Training of Transformer Language
Models

Kamran Chitsaz1,2 Quentin Fournier1,2 Gonçalo Mordido1,2,3 Sarath
Chandar1,2,3,4

1Chandar Research Lab 2Mila – Quebec AI Institute 3Polytechnique Montréal
4Canada CIFAR AI Chair
<EMAIL_ADDRESS>
###### Abstract

The increasing scale of Transformer models has led to an increase in their
pre-training computational requirements. While quantization has proven to be
effective after pre-training and during fine-tuning, applying quantization in
Transformers during pre-training has remained largely unexplored at scale for
language modeling. This study aims to explore the impact of quantization for
efficient pre-training of Transformers, with a focus on linear layer
components. By systematically applying straightforward linear quantization to
weights, activations, gradients, and optimizer states, we assess its effects
on model efficiency, stability, and performance during training. By offering a
comprehensive recipe of effective quantization strategies to be applied during
the pre-training of Transformers, we promote high training efficiency from
scratch while retaining language modeling ability. Code is available at
https://github.com/chandar-lab/EfficientLLMs.

Exploring Quantization for Efficient Pre-Training of Transformer Language
Models

Kamran Chitsaz1,2 Quentin Fournier1,2 Gonçalo Mordido1,2,3 Sarath
Chandar1,2,3,4 1Chandar Research Lab 2Mila – Quebec AI Institute
3Polytechnique Montréal 4Canada CIFAR AI Chair
<EMAIL_ADDRESS>

##  1 Introduction

Transformers (Vaswani et al., 2017) have become the dominant model for natural
language processing, with the GPT family of models (Radford et al., 2019)
showcasing their effectiveness across various tasks, from language
understanding to code generation. As the performance of Transformers scales
nicely with the number of parameters and the data size, the current state-of-
the-art models have reached unprecedented computational requirements both
during training and inference (Tay et al., 2020). For instance, pre-training a
175B parameter GPT-3 model requires a staggering number of 10,000 V100 GPUs
for 14.8 days (Patterson et al., 2021). As a result, pre-training has become
extremely expensive, beyond the reach of most research groups, and has raised
concerns over sustainability due to CO2 emissions from extensive GPU usage
(Luccioni et al., 2023).

To improve the efficiency of Transformers, quantization has gained significant
traction due to its recent successes in both post-training (Ashkboos et al.,
2024; Dettmers et al., 2022; Frantar et al., 2022) and during fine-tuning (Li
et al., 2023). Quantized pre-training, where certain parts of the
computational graph and model parameters are quantized from the beginning of
training, remains a challenging problem. In such scenarios, the training
instabilities caused by substantial changes in model parameters and emerging
model behaviors do not pair well with the added noise introduced by
quantization (Nagel et al., 2022). Additionally, quantizing model components
without compromising performance becomes increasingly difficult at larger
scales (Dettmers and Zettlemoyer, 2023). Despite its importance, quantized
pre-training of Transformer language models remains largely unexplored at
scale.

In this paper, we present the first in-depth study on the effects of
quantizing Transformer language models during pre-training and at scale. Our
primary aim is to provide a recipe for quantized pre-training by conducting a
controlled study that investigates the impact of quantization on weights,
activations, gradients, and optimizer states on model efficiency, stability,
and performance, using a simple linear quantization with 4 and 8 bits. Our
findings demonstrate that 8-bit quantization for weights and activations can
be effectively combined to provide significant memory savings and potential
speedup, achieving performance comparable to the baseline model. However,
extending quantization to gradients to utilize computational speedup in
backward matrix multiplications or reducing precision to 4 bits results in
notable training instability.

Specifically, 4-bit quantization introduces a sharper loss landscape for
weights (§4.1) and persistent outliers in the channel dimension of activations
(§4.2), significantly degrading performance despite attempts to manage them
through per-channel quantization. Additionally, gradient quantization is
particularly problematic due to spikes in gradient norms during early training
phases and the unstructured and sparse nature of gradients throughout training
(§4.3). While 8-bit gradient quantization does not hurt model convergence,
transitioning to 4 bits results in non-convergence. Additionally, quantizing
the first-order moments of Adam to even 4 bits is feasible without significant
performance loss, but the second-order moments require a more complex
quantization scheme to avoid instabilities in the Adam update, even when using
8-bits quantization (§4.4). Lastly, we present our recommended pre-training
quantization recipe for the different model components (§4.5).

##  2 Related Work

In recent years, numerous methods have been studied to improve the efficiency
of neural networks. Among these, Quantization-aware training (QAT) emerges as
an acceleration technique for inference, as parameters are stored and
operations are conducted with higher precision during training. The induced
quantization error during training serves as a regularizer, as demonstrated in
Gholami et al. (2022), ultimately facilitating the development of a more
quantization-friendly model.

In contrast, Fully Quantized Training (FQT) harnesses the accelerated gains
derived from higher throughput in INT8 or INT4 operations supported by modern
GPUs during training. Additionally, FQT capitalizes on memory savings by
storing parameters in lower precisions. As exemplified by Li et al. (2024) and
Dettmers et al. (2021), states of the Adam optimizer are stored in 4 and 8
bits, respectively, to minimize memory footprint. Another strategy introduced
by Markov et al. (2023) involves quantizing both weights and gradients to
reduce bandwidth usage in distributed training. Furthermore, Wortsman et al.
(2024) and Kim et al. (2021) advocate the replacement of linear operations
with INT8 matrix multiplications to achieve substantial speedups.

Despite the considerable advantages in accelerating the training process, FQT
poses challenges attributable to numerical stability and optimization issues
inherent in training quantized networks. Many existing methods predominantly
focus on fine-tuning Large Language Models (LLMs) using FQT, leveraging the
inherent stability of pre-trained models in each gradient update. The
evolution from 16-bit to FP8 data formats, as evidenced by remarkable results
in mixed precision training on LLMs (Peng et al., 2023), showcases the
potential of FQT. However, the scarcity and difficulty in obtaining hardware
that supports FP8 formats pose significant challenges to its widespread
adoption.

Wortsman et al. (2024) employ 8-bit quantization for linear operations in both
forward and backward passes, achieving remarkable results in pre-training
large-scale vision language models. Their approach incorporates row-wise
quantization for activations and gradients, mitigating the impact of
quantization errors on other parameters. However, vision language models
significantly differ from large textual language models. Xi et al. (2024b)
proposed 4-bit quantization in FQT with Hadamard transformation to address
outliers in activations. Acknowledging the sparsity of gradients during
training, they propose bit splitting for quantizing gradients in 4-bit
precision. However, they predominantly focus on fine-tuning, and their
experiments on pre-training only consider small models (60M) on a limited
dataset (WMT 14 En-De).

The generalization of these findings to large-scale language models (LLMs)
remains a challenge, especially considering that training such models involves
unique complexities. Additionally, their reliance on Hadamard transformations
imposes restrictions on activation dimensions, limiting applicability to
power-of-two dimensions, a constraint not easily met by recent LLM
architectures such as LLaMa.

In contrast to the aforementioned complex quantization methods, such as
quantile quantization and learnable quantization parameters, this study
focuses on a more straightforward implementation. The objective is to
investigate the effects of quantization and explore the feasibility of
training LLMs with full integer operations. The primary goal of this paper is
to offer detailed insights into quantizing different components of the model.
While more sophisticated quantization methods have the potential to enhance
the final performance, our work serves as a foundational investigation,
providing valuable insights and paving the way for further exploration in the
realm of quantization for LLMs.

##  3 Quantization Methodology

In this work, we explore various quantization schemes (§3.1) and granularities
(§3.2) on model components. This controlled approach during pre-training
allows us to examine the impacts on language modeling and downstream task
performance.

###  3.1 Quantization Scheme

We start by introducing the quantization procedure used in our study, which is
applied to all linear layers of Transformers. We perform fake quantization,
where all values and computations are stored with higher precision, and every
quantization operation is followed by de-quantization to introduce
quantization error. Simulating low-precision pre-training fits the purpose of
this study since we aim to analyze the effects of quantizing different model
components without focusing on actual training speedups that can be obtained
by implementing custom GPU kernels.

Figure 1: Overview of the quantization process in forward and backward passes.

Figure 1 illustrates how quantization error is introduced during both the
forward and backward passes. Specifically, the quantization error is injected
into either the weights or input activations during the forward pass. In the
backward pass, quantization is applied to the output gradient for computing
weight updates due to only the gradient of the weight is stored. Nonetheless,
the real-valued output gradient is used to compute the quantized input
gradient since we observed an increase in training instability when
propagating the quantization error through the entire backward path.
Additional details about gradient quantization are presented in Section 4.3.

We employ linear quantization for our experiments since this is a popular
approach compatible with existing hardware, and using more complex methods
could potentially hinder the practical relevance of our study. Specifically,
we map real-valued vectors $\bm{X}$ to a discrete grid of integers as follows:

$\begin{split}X_{\mathrm{int}}&=\operatorname{clip}\left(\left\lfloor\frac{X}{s}\right\rceil-z;N,P\right),\\\
\widehat{\bm{X}}&=s(X_{\mathrm{int}}+z),\end{split}$ (1)

where $\lfloor\cdot\rceil$ is the round-to-nearest integer operator, $N$ and
$P$ represent the quantization range, with $N=-2^{b-1}$, $P=2^{b-1}-1$, and
$b$ is the bit width, since we deal with signed data in our experiments. The
scaling factor $s$ is set to the maximum absolute value of $\bm{X}/P$. Unless
specified, we perform symmetric quantization by setting the offset $z$ to $0$,
which has less overhead than asymmetric quantization where $z$ is set to
$\left\lfloor\operatorname{min}(X)/s\right\rceil$ (Nagel et al., 2021). During
backpropagation, we employ the well-known straight-through estimator (STE)
(Bengio et al., 2013) mechanism to update the weights.

###  3.2 Quantization Granularity

We can choose scaling factors with different granularity: per-tensor, per-
channel, and per-token quantization, where each quantization granularity leads
to a specific trade-off between efficiency and performance. Specifically, per-
tensor quantization offers the highest efficiency since it performs a single
element-wise floating-point multiplication for the de-quantization step.
However, since only a single value is used to rescale the entire tensor,
performance degradation is likely to occur due to such uniform scaling across
the tensor elements.

On the other hand, per-channel and per-token quantization offer a finer-
grained scaling, where different scaling factors are tailored to specific
tensor element groups (i.e. channels and tokens, respectively). Even though
such approaches help in terms of performance, they introduce an overhead
during the de-quantization step. It is worth noting that, in certain
instances, these quantization granularities cannot be efficiently implemented
by hardware-accelerated GEMM kernels. For example, using per-channel
quantization for both weights and activations can not be efficiently
implemented.

###  3.3 Quantization Efficiency

Figure 2: Distribution of peak memory usage across different model sizes
(GPT-2 Small, Medium, and Large) for a constant context length of 1024, with
varying batch sizes.

To show the potential memory saving in quantized pre-training, we explore the
memory consumption of various components within GPT-2 models during training
using the PyTorch Memory Profiler. We analyze peak memory usage, as shown in
Figure 2, which illustrates memory usage for different batch sizes with a
fixed context length of 1024 across various model sizes. We observe that when
a model can fit within the GPU memory, the majority of the memory at peak
times is consumed by activations, particularly with large batch sizes and
sequence lengths. Under these conditions, gradients do not contribute to peak
memory usage. More details are in Appendix B.

Figure 3: Proportion of total execution time consumed by linear layers in the
attention block of GPT-2 models (Small, Medium, Large, and X-Large) across
different sequence lengths.

We also profile the execution time of kernels using the Nvidia Nsight Profiler
to assess the potential speedup from quantizing linear layers. Figure 3 shows
the proportion of total execution time consumed by linear layers in the
attention block of GPT-2 models of varying sizes across different sequence
lengths. This profiling includes both the forward and backward passes. We
observe that for small sequence lengths, linear layers consume a significant
portion (more than 80%) of the execution time. As the model size increases,
this proportion typically rises, but as sequence lengths increase, the
proportion of time spent in linear layers decreases, suggesting that self-
attention, due to its quadratic computational complexity, becomes the dominant
factor in execution time. This indicates that while quantizing linear layers
can offer substantial speedup, the potential gains are more pronounced with
smaller sequence lengths.

##  4 Experimental Results

We used GPT-2 small (Radford et al., 2019) (124M) with FlashAttention2 (Dao et
al., 2022) for our experiments due to its popularity and remarkable
performance-to-compute ratio. For the following experiments, we pre-trained 30
models from scratch on OpenWebText (Gokaslan and Cohen, 2019) for 300k
gradient steps with a global batch size of 512 samples and a context length of
1024 tokens. For our evaluation setup, we evaluate the performance of the
models on a range of language tasks, including ARC-Easy (Yadav et al., 2019),
ARC-Challenge (Yadav et al., 2019), Hellaswag (Zellers et al., 2019), LAMBADA
(Paperno et al., 2016), and GLUE score (Wang et al., 2018). Additional details
about our training and evaluation setups are provided in Appendix A.

###  4.1 Weight Quantization

Figure 4: Comparison of different Weight Quantization schemes. (Down)
Validation loss across training iterations for 4-bit and 8-bit quantization,
both per-tensor and per-channel, alongside the baseline. (Top) PFew-shot
accuracy on downstream tasks for the corresponding quantization approaches,
demonstrating the efficacy of 8-bit per-channel weight quantization.

The validation loss curves of applying per-tensor and per-channel quantization
to weights with 4 and 8 bits are presented in Figure 4 (down). We observe that
per-channel weight quantization with 8 bits outperforms the floating-point
baseline since the beginning of training in terms of validation loss, while
per-tensor weight quantization with 8 bits shows competitive performance. When
quantizing to 4 bits, there is a substantial difference between the different
granularities, with per-channel weight quantization significantly
outperforming per-tensor quantization, as previously discussed in § 3.2.

We also evaluate the downstream task performance of the quantized pre-trained
models in Figure 4 (top). We observe that 8-bit weight quantization
outperforms 4-bit and achieves competitive performance compared to the
floating point baseline, independently of the granularity used. Once again,
per-channel weight quantization with 8 bits achieves the best performance
among the tested quantization schemes. Overall, we observe similar findings
when comparing the performance of the different methods during the pre-
training and downstream phases. Despite the success in quantizing weights to 8
bits from scratch, we note that only performing 8-bit quantization post-
training also works well, as shown in Appendix C (Table 10). However, when it
comes to 4-bit quantization, applying quantization from scratch leads to
significantly better performance.

Figure 5: Sharpness comparison between baseline model and 4-bit weight
quantization. (Top) $m$-sharpness. (Down) Loss surfaces.

We note a pronounced drop in validation loss at the end of training with the
4-bit weight quantization schemes, as seen in the final training iterations in
Figure 4 (down). We hypothesize that this is related to reducing the learning
rate below $1e-6$ in the final steps of training in our setup. However, this
is not observed in all the schemes. Nagel et al. (2022) suggested that the
oscillations in weights originated from performing quantization with STE
during training may result in weight movements around decision thresholds.
Hence, in our use case, we hypothesize that the drop in validation loss stems
from the presence of sharp minima when quantizing weights to lower bit-widths
and that the lower learning rate regime at the end of training helps
convergence to minima with lower loss.

To further investigate this, we compare the sharpness of the different models
at the end of pre-training using $m$-sharpness (Foret et al., 2021) with
varying radii in Figure 5 (down). We observe that all quantized models
converge to sharper minima compared to the floating-point, unquantized
baseline. Moreover, there is a direct relation between the sharpness of each
quantized model and the relative drop in validation loss observed in Figure 4
(down). Specifically, per-tensor weight quantization with 4 bits shows the
highest sharpness and also the highest drop in validation loss. This
correlation is also observed on a smaller scale for the per-tensor weight
quantization model to 4 bits and the per-tensor weight quantization model to 8
bits. To visualize the loss surfaces, we employ the visualization method
introduced by Li et al. (2018). The loss surfaces of the baseline and the per-
tensor weight quantization model to 4-bits are shown in Figure 5 (right),
further illustrating the impact that quantization during pre-training has on
the sharpness of the final pre-trained model.

###  4.2 Activation Quantization

Figure 6: Training progression of activation distributions across selected
iterations, showing persistent channel-specific outliers

We present the validation loss when quantizing activations during pre-training
in Figure 7 (down). We note that quantizing activations to 4 bits is more
challenging than training with 4-bit weights, as noticed by previous work
(Xiao et al., 2023) and as seen by the divergence behavior of per-token and
per-tensor activation quantization. On the other hand, quantizing activations
to 8 bits works well, especially if performed per-token, which achieves lower
validation loss than the floating-point baseline. The performance of
downstream tasks of the baseline and quantized models is presented in Figure 7
(top). We notice a similar performance trend across the models, with 8-bit
per-token activation reaching competitive performance with the baseline model.

Figure 7: Pre-training Activation Quantization effects: (Down) Validation loss
curves for various quantization schemes; (Top) Few-shot accuracy on downstream
tasks, showing 8-bit quantization closely aligns with baseline performance.

To investigate 4-bit activation quantization further, we also tried applying
an asymmetric scheme. The intuition is that, while most activations exhibit
symmetry around zero, this is not the case for activation after GELU
activation functions (Hendrycks and Gimpel, 2016). Hence, having an asymmetric
scheme can lead to a better utilization of the available bits for
representation. However, we observe that while the asymmetric scheme provides
an improvement over 4-bit per-token symmetric quantization, the model still
diverges. We analyze the activation distributions of the output projection
layer within the attention block of layer 7 in Figure 6. We see that outliers
predominantly reside within specific channels and persistently affect the same
channels throughout training. Given our use of per-tensor and per-token
quantization for activations, it is evident that such outliers can influence
all tokens, given their consistent pattern across the channel dimension.

Figure 8: Validation loss and activation outliers for 4-bit per-channel
quantization. (Left) Validation loss shows convergence but underperforms
compared to the baseline. (Right) Histogram of activation with massive
outliers.

Since activation outliers are mostly predominant in particular channels during
training, we explore the efficacy of per-channel 4-bit activation quantization
in Figure 8 (left). We observe that this variant does converge, even though it
fails to be competitive with the floating-point baseline. Such degradation in
validation loss can be attributed to the presence of massive outlier
activations in specific layers. Despite being important for the model’s
performance Sun et al. (2024), these big activations pose a challenge for both
per-token and per-channel quantization. An example of these is presented in
Figure 8 (right), showing the presence of large activations in the FC2 layer
in the final attention block.

###  4.3 Gradient Quantization

We perform gradient quantization and present the validation losses obtained
during pre-training in Figure 9 (down). We observe that, with 4-bit gradient
quantization, training becomes highly unstable or completely fails to
converge. With 8 bits, only per-token quantization converges despite showing
worse performance compared to our baseline model. The observations are similar
when measuring the performance of the different models on downstream tasks, as
shown in Figure 9 (top).

Figure 9: Gradient Quantization. (Down) Validation loss showing non-
convergence for 4-bit and 8-bit per-tensor quantization. (Top) Few-shot
accuracy on downstream tasks, with only 8-bit per-token approaching baseline
performance.

As previously discussed in §3.1, we quantize the output gradients only for the
weight updates, avoiding the instability in training that can result from
propagating quantization errors when quantizing activation gradients. This is
illustrated in Figure 10 (top), where quantizing activation gradients in the
initial training stages leads to an explosion in the validation loss, followed
by divergence. Moreover, we observe an increase in the L2 norm between the
floating-point gradients and the quantized counterparts when quantizing
activation gradients compared to weight gradients.

Figure 10: (Top) Validation loss spike illustrating instability from
quantizing activation gradients. (Down) Gradient magnitude histogram for a
linear layer, highlighting sparsity and potential for quantization error.

To further analyze the subpar performance of 8-bit per-token gradient
quantization, even when only applied to weight gradients, we analyze the
gradients for the QKV projection at the first layer of the model early on in
training in Figure 10 (down). We observe that gradients are mostly sparse
during training and are prone to induce high quantization errors, rendering
instabilities.

###  4.4 Optimizer States Quantization

Figure 11: Quantization of Adam optimizer’s first-order moments. (Down)
Comparison of validation loss for different quantization approaches; (Top)
Few-shot accuracy on downstream tasks, with 8-bit per-channel closely matching
baseline.

We quantize optimizer states, particularly the first and second moments in the
Adam optimizer. Specifically, the quantized values of each state are stored
until the next training iteration, which are then dequantized and used for
Adam’s update. To better assess the effect of quantizing each state, we
quantize them separately and individually. The validation losses when
quantizing Adam’s first state are presented in Figure 11 (down). We observe
that per-channel quantization to 8 bits works well, achieving performance
similar to that of the baseline model. Notably, only per-tensor quantization
to 4 bits failed to converge out of all tested configurations. Similar
findings are found when evaluating the downstream performance of the different
quantization schemes in Figure 11 (right).

Figure 12: Second-order moments of Adam quantization. (Top) Validation loss
quickly diverges. (Down) Histogram showing the concentration of quantized
values in the zero bin, which contributes to instability in weight updates.

The results of quantizing Adam’s second state are presented in Figure 12
(left). We observe that the quantized model failed to converge smoothly
throughout training, even when applying per-channel quantization with 8 bits.
This can be explained by the usage of a linear symmetric quantization function
around zero in our scheme. This causes all small values to be set to zero
after quantization, hurting performance, as presented in Figure 12 (right).
Given that the second state plays a pivotal role in the denominator of Adam’s
update, such clustering to zero leads to excessively large weight updates,
causing training to diverge from the onset, as observed in Figure 12 (left).

###  4.5 Multiple Components Quantization

Figure 13: (Down) Validation loss across training iterations for weight,
activation, and gradient quantization together; (Top) Few-shot accuracy on
downstream tasks.

As our last studied components, we trained models using 8-bit quantization for
weights, activations, and gradients, as well as isolating quantization to only
weights and activations with per-channel granularity for weights and per-token
granularity for activations and gradients. Our findings, as depicted in Figure
13, demonstrate that quantizing both weights and activations to 8-bit allows
performance to closely align with the baseline model. However, extending
quantization to include gradients results in a notable decrease in
performance. This observation is consistent with earlier results from our
independent quantization of gradients, highlighting the significant challenges
this introduces.

##  5 Conclusion

This study presents an extensive analysis of the impact of quantizing specific
Transformer components in 4 and 8 bits during pre-training, in contrast to
concurrent work that focuses on individual methods without comprehensive
ablations Xi et al. (2024a). Our study reveals that quantizing weights to 8
bits from the beginning of pre-training is generally successful. However,
4-bit weight quantization can significantly affect model convergence due to a
sharper loss landscape. We also found that carefully managing activation
outliers is crucial to avoid performance drops with lower-precision
quantization. Additionally, we explain the sensitive nature of gradient
quantization and its potential to fail at lower bit-widths. We observed that
while the first-order moments of the Adam optimizer can be effectively
quantized to 4 bits, the second-order moments pose a greater challenge even
for 8-bit quantization. Overall, our work establishes an important foundation
for future developments in quantization approaches tailored to Transformers,
opening the door to efficiently train large-scale models from scratch for
improved accessibility.

## Limitation

We acknowledge the limitations of our work:

  * •

We employed linear quantization for our experiments. While this approach is
widely used and allows for a controlled study, it may not capture the full
potential of more sophisticated quantization methods.

  * •

Due to the cost of pre-training and the number of experiments, we limited our
study to GPT-2 small, and our findings may not generalize to larger models.

  * •

The efficiency gains discussed are estimated using profiling data, and
implementing these improvements in practice is challenging due to the
complexity of kernel optimizations.

## Acknowledgements

Sarath Chandar is supported by the Canada CIFAR AI Chairs program, the Canada
Research Chair in Lifelong Machine Learning, and the NSERC Discovery Grant.
Gonçalo Mordido was supported by an FRQNT postdoctoral scholarship (PBEEE)
during part of this work. The authors acknowledge the computational resources
provided by the Digital Research Alliance of Canada.

## References

  * Ashkboos et al. (2024) Saleh Ashkboos, Amirkeivan Mohtashami, Maximilian L Croci, Bo Li, Martin Jaggi, Dan Alistarh, Torsten Hoefler, and James Hensman. 2024.  Quarot: Outlier-free 4-bit inference in rotated llms.  _arXiv preprint arXiv:2404.00456_.
  * Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _arXiv preprint arXiv:1308.3432_.
  * Chelba et al. (2014) Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. 2014.  One billion word benchmark for measuring progress in statistical language modeling.  _Preprint_ , arXiv:1312.3005.
  * Dao et al. (2022) Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  _Advances in Neural Information Processing Systems_ , 35:16344–16359.
  * Dettmers et al. (2022) Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022.  GPT3.int8(): 8-bit matrix multiplication for transformers at scale.  In _Advances in Neural Information Processing Systems_.
  * Dettmers et al. (2021) Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 2021.  8-bit optimizers via block-wise quantization.  _arXiv preprint arXiv:2110.02861_.
  * Dettmers and Zettlemoyer (2023) Tim Dettmers and Luke Zettlemoyer. 2023.  The case for 4-bit precision: k-bit inference scaling laws.  In _Proceedings of the 40th International Conference on Machine Learning_ , ICML’23. JMLR.org.
  * Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. 2021.  Sharpness-aware minimization for efficiently improving generalization.  In _International Conference on Learning Representations_.
  * Frantar et al. (2022) Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022.  GPTQ: Accurate post-training compression for generative pretrained transformers.  _arXiv preprint arXiv:2210.17323_.
  * Gholami et al. (2022) Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. 2022.  A survey of quantization methods for efficient neural network inference.  In _Low-Power Computer Vision_ , pages 291–326. Chapman and Hall/CRC.
  * Gokaslan and Cohen (2019) Aaron Gokaslan and Vanya Cohen. 2019.  Openwebtext corpus.  http://Skylion007.github.io/OpenWebTextCorpus.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel. 2016.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_.
  * Kim et al. (2021) Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. 2021.  I-bert: Integer-only bert quantization.  In _International conference on machine learning_ , pages 5506–5518. PMLR.
  * Li et al. (2024) Bingrui Li, Jianfei Chen, and Jun Zhu. 2024.  Memory efficient optimizers with 4-bit states.  _Advances in Neural Information Processing Systems_ , 36.
  * Li et al. (2018) Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. 2018.  Visualizing the loss landscape of neural nets.  _Advances in neural information processing systems_ , 31.
  * Li et al. (2023) Zhikai Li, Xiaoxuan Liu, Banghua Zhu, Zhen Dong, Qingyi Gu, and Kurt Keutzer. 2023.  Qft: Quantized full-parameter tuning of llms with affordable resources.  _arXiv preprint arXiv:2310.07147_.
  * Luccioni et al. (2023) Alexandra Sasha Luccioni, Sylvain Viguier, and Anne-Laure Ligozat. 2023.  Estimating the carbon footprint of bloom, a 176b parameter language model.  _Journal of Machine Learning Research_ , 24(253):1–15.
  * Marcus et al. (1993) Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. 1993.  Building a large annotated corpus of English: The Penn Treebank.  _Computational Linguistics_ , 19(2):313–330.
  * Markov et al. (2023) Ilia Markov, Adrian Vladu, Qi Guo, and Dan Alistarh. 2023.  Quantized distributed training of large models with convergence guarantees.  _arXiv preprint arXiv:2302.02390_.
  * Merity et al. (2016) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016.  Pointer sentinel mixture models.  _Preprint_ , arXiv:1609.07843.
  * Nagel et al. (2021) Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart Van Baalen, and Tijmen Blankevoort. 2021.  A white paper on neural network quantization.  _arXiv preprint arXiv:2106.08295_.
  * Nagel et al. (2022) Markus Nagel, Marios Fournarakis, Yelysei Bondarenko, and Tijmen Blankevoort. 2022.  Overcoming oscillations in quantization-aware training.  In _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 16318–16330. PMLR.
  * Paperno et al. (2016) Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. 2016.  The LAMBADA dataset: Word prediction requiring a broad discourse context.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1525–1534, Berlin, Germany. Association for Computational Linguistics.
  * Patterson et al. (2021) David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. 2021.  Carbon emissions and large neural network training.  _arXiv preprint arXiv:2104.10350_.
  * Peng et al. (2023) Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, et al. 2023.  Fp8-lm: Training fp8 large language models.  _arXiv preprint arXiv:2310.18313_.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9.
  * Sun et al. (2024) Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. 2024.  Massive activations in large language models.  _arXiv preprint arXiv:2402.17762_.
  * Tay et al. (2020) Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2020.  Efficient transformers: A survey.(2020).  _arXiv preprint cs.LG/2009.06732_.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2018.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  _arXiv preprint arXiv:1804.07461_.
  * Wortsman et al. (2024) Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. 2024.  Stable and low-precision training for large-scale vision-language models.  _Advances in Neural Information Processing Systems_ , 36.
  * Xi et al. (2024a) Haocheng Xi, Yuxiang Chen, Kang Zhao, Kaijun Zheng, Jianfei Chen, and Jun Zhu. 2024a.  Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization.  _arXiv preprint arXiv:2403.12422_.
  * Xi et al. (2024b) Haocheng Xi, Changhao Li, Jianfei Chen, and Jun Zhu. 2024b.  Training transformers with 4-bit integers.  _Advances in Neural Information Processing Systems_ , 36.
  * Xiao et al. (2023) Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. 2023.  Smoothquant: Accurate and efficient post-training quantization for large language models.  In _International Conference on Machine Learning_ , pages 38087–38099. PMLR.
  * Yadav et al. (2019) Vikas Yadav, Steven Bethard, and Mihai Surdeanu. 2019.  Quick and (not so) dirty: Unsupervised selection of justification sentences for multi-hop question answering.  _arXiv preprint arXiv:1911.07176_.
  * Zellers et al. (2019) Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019.  Hellaswag: Can a machine really finish your sentence?  _arXiv preprint arXiv:1905.07830_.

##  Appendix A Experimental Setup

###  A.1 Model and Training Configuration

Our experiments leverage the computational enhancements of the FlashAttention
library (Dao et al., 2022), utilizing its GPT-2 implementation within the
HuggingFace Trainer framework
111https://huggingface.co/docs/transformers/en/main_classes/trainer for our
training processes.

Training was conducted on the OpenWebText corpus (Gokaslan and Cohen, 2019),
adopting a set of training configurations similar to (Dao et al., 2022) and
nanoGPT 222https://github.com/karpathy/nanoGPT, without hyperparameter tuning
due to computational constraints. We use AdamW optimizer with a learning rate
of $6e-4$, combined with a cosine learning rate scheduler set to a half cycle.
We adopt mixed precision training in bfloat16. The experiments are conducted
with a fixed batch size of 512, employing gradient accumulation as necessary
to accommodate the computational constraints of our setup across 4xA100 80G
GPUs. This training configuration remains consistent across all experiments,
culminating in a training duration averaging 4.3 days for completion of 300k
steps.

###  A.2 Data and Evaluation Metrics

The OpenWebText corpus was randomly split into training and validation sets,
with 0.5% of the data reserved for validation. Following the methodology of
Radford et al. (2019), the performance of our pre-trained models was evaluated
on several benchmark datasets. These evaluations focused on perplexity
measurements on well-known datasets including WikiText (Merity et al., 2016),
PTB (Marcus et al., 1993), and 1BW (Chelba et al., 2014). Our models were
further evaluated using a range of downstream tasks:

  * •

GLUE (Wang et al., 2018): A collection of natural language understanding tasks
including question answering, sentiment analysis, and textual entailment,
designed to benchmark the generalization capabilities of models across a
diverse range of linguistic challenges.

  * •

ARC (Yadav et al., 2019): Comprising the ARC-Easy and ARC-Challenge, these
datasets test the model’s reasoning ability through science-based question
answering. ARC-Easy contains simpler questions, while ARC-Challenge includes
more complex queries demanding deeper reasoning.

  * •

Hellaswag (Zellers et al., 2019): This dataset challenges models to predict
the most plausible continuation of a narrative from a large corpus of everyday
contexts and movie scripts, testing the commonsense reasoning ability of the
models.

  * •

LAMBADA (Gokaslan and Cohen, 2019): Evaluates the model’s capability to
predict the final word of a textual passage, focusing on the contextual
understanding of the language.

For evaluation on downstream tasks, we adopted a few-shot approach, utilizing
a 5-shot prompting method with greedy decoding and we report the average and
variance of accuracy over 5 different seeds. Each prompt was structured with a
task instruction followed by the five examples of training data and we
computed accuracy on the validation set. We implemented our few-shot
evaluation protocol following the guidelines provided by the
lm_evaluation_harness library
333https://github.com/zphang/lm_evaluation_harness. For each configuration
reported in the tables, the average GLUE score is computed first by taking the
mean across all GLUE tasks. Subsequently, an overall average is calculated by
averaging this GLUE score with the scores from ARC Easy, ARC Challenge,
Hellaswag, and LAMBADA tasks.

###  A.3 Comparison of Baseline and Pre-trained Model

Table 1: Comparison of Baseline and Pre-trained Models: The latter trained for
at least twice the duration.

| WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---
baseline | 39.94 | 34.32 | 35.13 | 44.03
pre-trained | 29.17 | 24.67 | 35.86 | 45.87

To establish a solid baseline for our experiments, we benchmark our trained
model against the pre-trained GPT-2 weights provided by OpenAI across several
downstream tasks. The comparison, detailed in Table 1, reveals that our model
achieves results closely aligned with the original, validating the efficacy of
our training approach. In the following sections, we will delve into the
individual components of the model, discussing the impact and outcomes of our
quantization experiments on each.

##  Appendix B Memory Analysis

Figure 14: Distribution of peak memory usage across different model sizes
(GPT-2 Small, Medium, and Large) for a constant context length of 1024, with
varying batch sizes. This figure demonstrates how memory dedicated to
activations increases as batch size increases, highlighting the impact of
batch size on memory allocation dynamics in large-scale models. Figure 15:
Peak memory usage profile for different sequence lengths while maintaining a
constant batch size of 4, across different model sizes (GPT-2 Small, Medium,
and Large). This figure illustrates the shift in peak memory usage from
gradients to activations as sequence length increases, and the significant
impact of sequence length on memory dynamics during training.

In this section, we explore the memory consumption patterns of various
components within GPT-2 models during training, using the PyTorch Memory
Profiler. This profiling tool allows for precise monitoring of memory usage
throughout the lifecycle of a model’s operation, particularly during the
forward and backward passes, and during optimization steps. "Peak memory" in
this context refers to the maximum memory usage observed at any point in these
stages, providing insights into how different model configurations impact
overall memory requirements. We profiled all attention blocks as well as lm
head of the Transformer model.

Figure 14 shows how memory usage changes with different batch sizes for a
fixed context length of 1024, across various model sizes. Notably, as the
batch size increases, the memory allocated for activations becomes more
dominant, especially in larger models. This trend is primarily due to the need
to store activations for the computation of gradients during the backward
pass, which increases with larger batch sizes.

Figure 15 examines the memory usage across various sequence lengths while
keeping the batch size constant at 4. When both batch size and sequence length
are small, peak memory typically occurs towards the end of the backward
propagation phase. At this stage, memory includes the parameters, optimizer
states, gradients from all layers, and activations from the initial layers.
However, as the sequence length and batch size increase, peak memory usage
shifts to the beginning of the backward propagation. At this point, the memory
comprises parameters, optimizer states, all activations, and notably, the
output gradient of the final layers, which matches the size of the logits
(proportional to batch size * sequence length * vocabulary size).

In conclusion, the analysis reveals that when a model can fit within the GPU
memory, the majority of the memory at peak times is consumed by activations,
especially when working with sufficiently large batch sizes and sequence
lengths. Under these conditions, gradients do not contribute to peak memory
usage. Consequently, quantizing gradients will not lead to significant memory
savings.

##  Appendix C Quantization Results

Table 2: Weight Quantization: evaluation of perplexity across multiple
datasets.

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | 55.50 | 46.70 | 52.38 | 59.14
per-channel | 56.43 | 47.32 | 38.18 | 46.30
8 bit | per-tensor | 48.52 | 40.01 | 37.02 | 45.04
per-channel | 42.43 | 35.94 | 34.81 | 43.47

Table 3: Activation Quantization: evaluation of perplexity across multiple
datasets.

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | 418.63 | 264.07 | 261.64 | 310.43
per-token | 69.82 | 54.53 | 56.58 | 72.06
8 bit | per-tensor | 64.77 | 51.06 | 37.96 | 45.26
per-token | 42.86 | 37.17 | 35.43 | 43.38

Table 4: Gradient Quantization: evaluation of perplexity across multiple
datasets.

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | 17990.70 | 15560.03 | 6632.20 | 6393.07
per-token | 128.71 | 92.74 | 106.73 | 110.14
8 bit | per-tensor | 123.08 | 87.81 | 104.90 | 111.51
per-token | 59.24 | 47.50 | 42.28 | 51.89

Table 5: Adam Optimizer’s First Moments: evaluation of perplexity across
multiple datasets.

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | 78.78 | 62.08 | 66.50 | 85.60
per-channel | 43.02 | 36.70 | 38.57 | 47.90
8 bit | per-tensor | 42.93 | 36.91 | 39.72 | 46.63
per-channel | 39.84 | 33.78 | 35.67 | 44.29

Table 6: Weight Quantization: Few-shot accuracy on downstream tasks.

| GLUE Score | ARC |  |
---|---|---|---|---
# of bit | granularity | MNLI | MRPC | RTE | QNLI | SST | WNLI | Easy | Challenge | Hellaswag | LAMBADA | Average |
baseline | $33.3_{\pm 0.4}$ | $61.7_{\pm 1.3}$ | $49.7_{\pm 2.6}$ | $49.2_{\pm 0.3}$ | $53.8_{\pm 1.9}$ | $46.8_{\pm 3.9}$ | $45.7_{\pm 0.3}$ | $22.5_{\pm 1.1}$ | $28.9_{\pm 0.1}$ | $36.17$ | $36.46$ |
4 bit | per-tensor | $32.1_{\pm 0.2}$ | $46.1_{\pm 1.2}$ | $49.0_{\pm 3.1}$ | $49.1_{\pm 0.2}$ | $54.1_{\pm 1.6}$ | $35.5_{\pm 4.2}$ | $39.9_{\pm 0.7}$ | $19.3_{\pm 0.5}$ | $27.2_{\pm 0.1}$ | $27.03$ | $31.54$ |
per-channel | $33.3_{\pm 0.1}$ | $53.1_{\pm 2.2}$ | $49.7_{\pm 2.8}$ | $49.9_{\pm 0.5}$ | $54.6_{\pm 0.8}$ | $45.9_{\pm 4.6}$ | $44.6_{\pm 0.9}$ | $22.5_{\pm 0.7}$ | $28.7_{\pm 0.1}$ | $34.21$ | $35.55$ |
8 bit | per-tensor | $33.2_{\pm 0.3}$ | $59.1_{\pm 1.3}$ | $49.7_{\pm 2.6}$ | $49.1_{\pm 0.1}$ | $56.9_{\pm 2.6}$ | $41.1_{\pm 3.5}$ | $45.5_{\pm 0.8}$ | $20.4_{\pm 1.4}$ | $28.7_{\pm 0.2}$ | $34.81$ | $35.53$ |
per-channel | $34.9_{\pm 0.3}$ | $62.7_{\pm 1.5}$ | $53.2_{\pm 1.3}$ | $49.5_{\pm 0.1}$ | $56.9_{\pm 1.1}$ | $53.5_{\pm 3.6}$ | $44.6_{\pm 0.5}$ | $21.1_{\pm 0.9}$ | $29.1_{\pm 0.2}$ | $36.43$ | $36.59$ |

Table 7: Activation Quantization: Few-shot accuracy on downstream tasks.

| GLUE Score | ARC |  |
---|---|---|---|---
# of bit | granularity | MNLI | MRPC | RTE | QNLI | SST | WNLI | Easy | Challenge | Hellaswag | LAMBADA | Average |
baseline | $33.3_{\pm 0.4}$ | $61.7_{\pm 1.3}$ | $49.7_{\pm 2.6}$ | $49.2_{\pm 0.3}$ | $53.8_{\pm 1.9}$ | $46.8_{\pm 3.9}$ | $45.7_{\pm 0.3}$ | $22.5_{\pm 1.1}$ | $28.9_{\pm 0.1}$ | $36.17$ | $36.46$ |
4 bit | per-tensor | $34.6_{\pm 0.2}$ | $31.8_{\pm 0.2}$ | $50.5_{\pm 2.0}$ | $49.5_{\pm 0.0}$ | $50.2_{\pm 1.0}$ | $52.1_{\pm 2.7}$ | $30.4_{\pm 1.3}$ | $20.6_{\pm 2.6}$ | $26.1_{\pm 0.1}$ | $1.53$ | $24.67$ |
per-token | $34.2_{\pm 0.4}$ | $58.3_{\pm 1.4}$ | $51.2_{\pm 2.3}$ | $49.3_{\pm 0.4}$ | $50.8_{\pm 0.3}$ | $52.4_{\pm 3.1}$ | $37.1_{\pm 1.6}$ | $21.0_{\pm 1.8}$ | $26.7_{\pm 0.1}$ | $18.88$ | $30.62$ |
| per-token asymmetric | $34.3_{\pm 0.2}$ | $54.1_{\pm 2.0}$ | $50.7_{\pm 2.9}$ | $49.1_{\pm 0.5}$ | $53.0_{\pm 1.4}$ | $55.2_{\pm 4.6}$ | $38.6_{\pm 0.9}$ | $18.9_{\pm 0.3}$ | $26.9_{\pm 0.2}$ | $22.30$ | $31.22$ |
8 bit | per-tensor | $32.7_{\pm 0.2}$ | $48.3_{\pm 1.4}$ | $47.1_{\pm 2.1}$ | $49.4_{\pm 0.3}$ | $54.6_{\pm 1.4}$ | $39.7_{\pm 3.0}$ | $44.1_{\pm 0.8}$ | $22.3_{\pm 1.5}$ | $28.8_{\pm 0.1}$ | $34.33$ | $34.97$ |
per-token | $33.8_{\pm 0.4}$ | $56.2_{\pm 0.7}$ | $50.5_{\pm 3.4}$ | $49.8_{\pm 0.3}$ | $56.4_{\pm 1.5}$ | $54.1_{\pm 3.0}$ | $44.5_{\pm 0.6}$ | $20.0_{\pm 1.7}$ | $29.3_{\pm 0.1}$ | $36.46$ | $36.08$ |

Table 8: Gradient Quantization: Few-shot accuracy on downstream tasks.

| GLUE Score | ARC |  |
---|---|---|---|---
# of bit | granularity | MNLI | MRPC | RTE | QNLI | SST | WNLI | Easy | Challenge | Hellaswag | LAMBADA | Average |
baseline | $33.3_{\pm 0.4}$ | $61.7_{\pm 1.3}$ | $49.7_{\pm 2.6}$ | $49.2_{\pm 0.3}$ | $53.8_{\pm 1.9}$ | $46.8_{\pm 3.9}$ | $45.7_{\pm 0.3}$ | $22.5_{\pm 1.1}$ | $28.9_{\pm 0.1}$ | $36.17$ | $36.46$ |
4 bit | per-token | $33.4_{\pm 0.4}$ | $36.6_{\pm 0.6}$ | $47.4_{\pm 2.6}$ | $49.5_{\pm 0.0}$ | $50.5_{\pm 0.8}$ | $58.3_{\pm 2.6}$ | $33.9_{\pm 1.0}$ | $20.2_{\pm 1.0}$ | $26.1_{\pm 0.1}$ | $19.06$ | $29.04$ |
8 bit | per-tensor | $34.2_{\pm 0.6}$ | $32.7_{\pm 0.6}$ | $50.3_{\pm 1.2}$ | $49.4_{\pm 0.0}$ | $50.5_{\pm 0.7}$ | $56.3_{\pm 2.0}$ | $32.1_{\pm 0.7}$ | $21.9_{\pm 1.0}$ | $26.3_{\pm 0.1}$ | $18.22$ | $28.81$ |
per-token | $33.2_{\pm 0.1}$ | $42.6_{\pm 2.4}$ | $51.6_{\pm 1.8}$ | $49.5_{\pm 0.5}$ | $53.0_{\pm 2.0}$ | $22.3_{\pm 3.1}$ | $45.0_{\pm 1.0}$ | $19.5_{\pm 1.3}$ | $28.5_{\pm 0.1}$ | $33.81$ | $33.77$ |

Table 9: Adam Optimizer’s First-Order Moments Quantization: Few-shot accuracy
on downstream tasks.

| GLUE Score | ARC |  |
---|---|---|---|---
# of bit | granularity | MNLI | MRPC | RTE | QNLI | SST | WNLI | Easy | Challenge | Hellaswag | LAMBADA | Average |
baseline | $33.3_{\pm 0.4}$ | $61.7_{\pm 1.3}$ | $49.7_{\pm 2.6}$ | $49.2_{\pm 0.3}$ | $53.8_{\pm 1.9}$ | $46.8_{\pm 3.9}$ | $45.7_{\pm 0.3}$ | $22.5_{\pm 1.1}$ | $28.9_{\pm 0.1}$ | $36.17$ | $36.46$ |
4 bit | per-tensor | $32.2_{\pm 0.3}$ | $64.9_{\pm 1.6}$ | $47.7_{\pm 0.6}$ | $50.2_{\pm 0.5}$ | $51.6_{\pm 1.7}$ | $38.9_{\pm 1.7}$ | $35.1_{\pm 1.6}$ | $20.1_{\pm 0.8}$ | $26.5_{\pm 0.2}$ | $20.12$ | $29.87$ |
per-column | $33.8_{\pm 0.1}$ | $66.6_{\pm 1.2}$ | $50.0_{\pm 3.0}$ | $49.5_{\pm 0.0}$ | $53.5_{\pm 0.9}$ | $50.4_{\pm 2.7}$ | $44.6_{\pm 1.6}$ | $20.6_{\pm 0.9}$ | $28.3_{\pm 0.1}$ | $33.15$ | $35.44$ |
8 bit | per-tensor | $33.7_{\pm 0.4}$ | $59.2_{\pm 2.5}$ | $51.4_{\pm 2.7}$ | $50.1_{\pm 0.6}$ | $56.1_{\pm 1.4}$ | $51.8_{\pm 3.3}$ | $44.4_{\pm 0.7}$ | $18.5_{\pm 0.6}$ | $28.7_{\pm 0.1}$ | $32.91$ | $34.97$ |
per-column | $33.3_{\pm 0.2}$ | $67.1_{\pm 1.0}$ | $50.8_{\pm 1.5}$ | $49.5_{\pm 0.0}$ | $55.3_{\pm 1.1}$ | $48.2_{\pm 2.1}$ | $43.8_{\pm 1.7}$ | $21.7_{\pm 1.3}$ | $28.9_{\pm 0.1}$ | $36.52$ | $36.34$ |

This appendix presents the granular results from our experiments on weight,
activation, gradient, and optimizer states quantization. The tables detail the
performance metrics, like perplexity and accuracy, under various quantization
settings. Finally, we present post-training quantization results of our
baseline model in Tables 10,11.

Table 10: Post-training weight quantization results

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
(ppl) | (ppl) | (ppl) | (ppl)
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | 16196.10 | 17256.89 | 17471.35 | 13761.79
per-column | 98.39 | 75.56 | 81.28 | 94.40
8 bit | per-tensor | 46.45 | 39.23 | 41.18 | 52.15
per-column | 40.15 | 34.45 | 35.23 | 44.11

Table 11: Post-training activation quantization results

#bit | granularity | WikiText103 | WikiText2 | PTB | 1BW
---|---|---|---|---|---
(ppl) | (ppl) | (ppl) | (ppl)
baseline | 39.94 | 34.32 | 35.13 | 44.03
4 bit | per-tensor | - | - | - | -
per-token | 14022.78 | 17933.29 | 13392.28 | 8763.06
8 bit | per-tensor | 70.07 | 58.45 | 64.99 | 149.35
per-token | 40.09 | 34.44 | 35.43 | 44.37


11institutetext: Ho Chi Minh University of Technology, Vietnam.
22institutetext: National Cheng Kung University, Taiwan. 33institutetext:
Vietnamese-German University, Vietnam

# Semi-supervised learning for White Blood Cell Semantic Segmentation

Quoc-Vinh Luu 11 Khanh-Duy Le 22 Thanh Huy Nguyen 33

###### Abstract

AI’s application in healthcare, especially in white blood cell cancer
diagnosis, is hindered by two primary challenges: the lack of large-scale
labeled datasets for white blood cell (WBC) segmentation and outdated
segmentation methods. To address the first challenge, we must have a semi-
supervised learning framework so that we can efficiently annotate that large
dataset. In this work, we address this issue and propose a novel self-training
pipeline with the incorporation of FixMatch. We conduct experiments on white
blood cell datasets and evaluate the strengths and weaknesses of our pipeline.

###### Keywords:

Semi-supervised learning Semantic segmentation White blood cell

##  1 Introduction

Cancer stands among the deadliest diseases in the world. In Vietnam and the
U.S., cancer is only second to heart disease in terms of mortality rate. In
2020, it is estimated that nearly 19 million cases of cancer arise worldwide,
with the number of deaths standing at approximately 10 million [1]. Cancer of
the white blood cells (WBC), namely leukemia and lymphoma, are two of the most
common types of blood cancers [2]. However, certain screening and diagnosis of
blood cancers, particularly cancer of the white blood cells (WBC), still
employ a manual process, such as in peripheral blood smear analysis. These
tests often serve the purpose of classifying and counting the number of WBCs
in the blood, and if the number is abnormally high or low, further examination
is needed to decide if the patient is inflicted with cancer or not.

White blood cell segmentation is a rarely touched area among AI experts. In
our experience, it is hard to find a dataset that both provides images and
their ground truth masks to perform semantic segmentation. Most datasets
concern themselves with white blood cell classification and do not contain the
masks needed for segmentation. And when we found datasets that suited our
purposes, the number of data was very insignificant.

Recognizing that there is a shortage of meaningful datasets in the field of
white blood cell segmentation, we aim to provide a semi-supervised learning
technique so that in the presence of a large-scale dataset, we can annotate
that dataset most efficiently.

Our contributions to this paper are as follows:

  * -

We aim to provide a semi-supervised learning approach to annotate a large
dataset of white blood cells.

  * -

From that semi-supervised learning approach, we should devise a more accurate,
fine-tuned way to segment white blood cells.

##  2 Related works

####  2.0.1 Semi-supervised learning.

is a paradigm in which the dataset used for training is not fully annotated.
It sits between supervised learning, where the whole dataset is labeled, and
unsupervised learning, where no data is labeled. In the early days, semi-
supervised learning was confined to traditional machine-learning approaches
[5], such as semi-supervised support vector machines, graph-based, and co-
training. Semi-supervised learning can be categorized into 5 groups:
generative, consistency regularization, graph-based, pseudo-labeling, and
hybrid methods.

####  2.0.2 Semi-supervised semantic segmentation

addresses the challenges of applying semi-supervised learning principles to
the task of semantic image segmentation. It is a more realistic application of
semi-supervised learning since semantic segmentation involves pixel-level
labeling which is a more intensive labeling requirement than classification or
detection. The techniques in this field can also be categorized into 5
categories: Adversarial, contrastive learning, consistency regularization,
pseudo-labeling, and hybrid methods.

####  2.0.3 Consistency regularization

The consistency regularization method commonly used in semi-supervised
semantic segmentation can be categorized into 4 main groups: Input
perturbations, feature perturbations, network perturbations, and the three
combined methods. The common goal for all of these sub-groups is to make the
model perform consistently across all transformations applied to the
unlabelled data.

####  2.0.4 Mean teacher.

Almost all of the methods in consistency regularization are based on the Mean
Teacher model. The Mean Teacher model consists of two main components: the
student model and the teacher model. The student model is the main model that
is being trained on the available labeled data. Like in standard supervised
learning, the student model learns to make predictions based on the labeled
examples it sees. For the Teacher model, this is a secondary model whose
weights are an exponential moving average (EMA) of the student model’s weights
over time. The teacher model is not trained directly with gradient descent.
Instead, its weights are updated as the moving average of the student model’s
weights after each training step.

##  3 Methodology

###  3.1 Problem definition

Semi-supervised semantic segmentation can be formulated as follows: Let
$X=\\{X_{L},X_{U}\\}$ be the dataset, and $X_{L}=\\{x_{i}\\}_{i=1}^{L}$ be the
labeled subset, $Y_{L}=(y_{1},y_{2},\ldots,y_{L})$ is the labels for such set
and $X_{U}=\\{(x_{i})\\}_{i=1}^{U}$ be the unlabeled subset. The assumption is
that $L\ll U$. The objective is to minimize the following loss function:

$\min_{\theta}\sum_{x\in X_{L},y\in
Y_{L}}\mathcal{L}_{s}(x,y,\theta)+\alpha\sum_{x\in
X_{U}}\mathcal{L}_{u}(x,\theta)+\beta\sum_{x\in X}\mathcal{R}(x,\theta)$

$\mathcal{L}_{s}$ denotes the supervised losses with labeled data,
$\mathcal{L}_{u}$ is the unsupervised loss and $\mathcal{R}$ is the
regularization term. Depending on the technique chosen, this regularization
term can be consistency regularization or another designed regularization
technique. $\theta$ is the model’s parameters and $\alpha$ and $\beta$ are
trade-off terms respectively.

###  3.2 Self-training (ST)

Similar to Mean teacher, self-training also incorporated the idea of the
teacher-student model. But one significant difference is that the teacher and
student are trained in two separate stages. The steps for the plain Self-
training are as follows:

  1. 1.

Train the teacher model $T$ on labeled data $X_{L}=\\{x_{i}\\}_{i=1}^{L}$.

  2. 2.

Use the model that has just been trained on labeled data to generate
predictions on the unlabeled data. These predictions in the context of
semantic segmentation can be considered as pseudo-masks. This set of pseudo-
masks can be formulated as
$\hat{D}_{U}=\left\\{(x_{i},T(x_{i}))\right\\}_{i=1}^{U}$

  3. 3.

Extend the labeled data ${D}_{L}=\left\\{(x_{i},y_{i})\right\\}_{i=1}^{L}$
with the unlabeled ones $\hat{D}_{U}$ by unioning the labeled set with the
unlabeled set, then re-train the student model S on this new set.

In [16], the authors propose a stronger baseline for this plain self-training
pipeline by incorporating strong data augmentation in the final stage on
unlabelled data. The performance upon implementing these is competitive in
comparison with other methods at the time.

Algorithm 1 ST Pseudocode

1:Labeled training set $D_{L}=\\{(x_{i},y_{i})\\}_{i=1}^{L}$, Unlabeled
training set $D_{U}=\\{u_{i}\\}_{i=1}^{U}$

2:Fully trained model $S$

3:Initialize teacher model $T$ with random weights

4:Pre-train $T$ on $D_{L}$ using supervised learning

5:Use model $T$ to infer labels on $D_{U}$ to get $\hat{D}_{U}$

6:Combine $D_{L}$ and $\hat{D}_{U}$ to create an augmented training set
$D_{aug}$

7:Re-train student model $S$ on $D_{aug}$

8:return $S$

###  3.3 Self-training++ (ST++)

Despite the simplicity and impressive results of Self-training, it doesn’t
have a mechanism to take into account the pseudo-masks that are extremely
incorrect. By re-training on pseudo-masks that are extremely wrong, called
unreliable samples, the model’s performance inevitably deteriorates. In
contrast, the pseudo-masks that are close to the actual ground truth of the
unlabeled image, called reliable samples, play a key role in improving the
model’s performance. Therefore, in [16], the authors propose a way to harness
the power of reliable images to compensate for unreliable ones. Instead of re-
training directly on labeled data and all unlabelled data in the final step,
the authors first train the model on labeled data and reliable samples. Then,
they use this model to generate pseudo-masks for the unreliable data. Finally,
they re-train the model on labeled data, reliable data, and unreliable data.

Same as the strong baseline the authors in [16] for plain self-training, the
authors apply strong data augmentation on the unlabeled images whenever re-
training on both labeled and unlabeled data.

Algorithm 2 ST++ Pseudocode

1:Labeled training set $D_{L}$, Unlabeled training set $D_{U}$

2:Fully trained student model $S$

3:Train teacher model $T$ on $D_{L}$

4:Obtain pseudo labels for $D_{U}$ using $T$

5:Select a subset of $D_{U}$ with the most reliable pseudo labels

6:Train student model $S$ on $D_{L}$ and the reliable subset of $D_{U}$

7:Re-label the remaining subset of $D_{U}$ using $S$

8:Re-train $S$ on $D_{L}$ and the entire $D_{U}$

9:return $S$

###  3.4 FixMatch

Inspired by the input perturbation method in the area of consistency-
regularization applied in semi-supervised semantic segmentation, FixMatch [17]
is a method introduced in 2020 based on the assumption that a model should
produce the same prediction for an unlabeled image, regardless of whether it
has been weakly or strongly augmented. At the time of its introduction, the
FixMatch method outperformed all other methods in datasets such as CIFAR-100,
SVHN, and STL-100 [17] in the field of image classification.

###  3.5 ST++ with FixMatch

Realizing that the ST++ approach has not yet fully utilized consistency
regularization approaches, we aspire to incorporate FixMatch among the steps
described in the pseudo-code for ST++. Realizing that the first step of
training in a fully supervised manner in ST++ can be easily adjusted to
incorporate FixMatch. We have done just that.

Algorithm 3 ST++ with FixMatch

1:Labeled dataset $D_{L}=\\{(x_{i},y_{i})\\}$, Unlabeled dataset
$D_{U}=\\{u_{j}\\}$, Model parameters $\theta$, Confidence threshold $\tau$,
Strong augmentation function $S(\cdot)$, Weak augmentation function $W(\cdot)$

2:Initialize: Train model on $L$ for a few epochs

3:for each training epoch do

4: for each mini-batch do

5: Draw a batch of labeled samples $\\{(x_{i},y_{i})\\}$ from $D_{L}$

6: Apply weak augmentation: $x^{\prime}_{i}=\alpha(x_{i})$

7: Draw a batch of unlabeled samples $\\{u_{j}\\}$ from $U$

8: Apply weak augmentation: $u^{\prime}_{j}=W(u_{j})$

9: Apply strong augmentation: $u^{\prime\prime}_{j}=S(u_{j})$

10: Obtain predictions for weakly-augmented unlabeled data:
$p_{j}=M(u^{\prime}_{j};\theta)$

11: Generate pseudo-labels: $\hat{y}_{j}=\arg\max(p_{j})$ if
$\max(p_{j})>\tau$ else ignore

12: Supervised Training Step (incorporating FixMatch):

13: Train model on $\\{(x^{\prime}_{i},y_{i})\\}$ using supervised loss

14: Train model on $\\{(u^{\prime\prime}_{j},\hat{y}_{j})\\}$ for samples with
pseudo-labels, using consistency loss

15: end for

16: Update model parameters $\theta$

17:end for

18:return Optimized model parameters $\theta$

##  4 Experiments

###  4.1 Dataset

####  4.1.1 Dataset 1.

[14] Sourced from Jiangxi Tecom Science Corporation in China, comprises 300
images of white blood cells (WBCs), each measuring 120×120 pixels with a color
depth of 24 bits. These images were captured using a Motic Moticam Pro 252A
camera attached to an N800-D motorized auto-focus microscope. The blood smear
samples in this dataset were stained using a novel hematology reagent
specifically designed for rapid WBC staining, resulting in images that
predominantly exhibit a yellow background.

Figure 1: Dataset1: A sample

####  4.1.2 Dataset 2

[14] Dataset 2 is composed of 100 color images, each with dimensions of
300×300 pixels, acquired from the CellaVision blog. These images predominantly
feature cells in a purple hue and often include numerous red blood cells
surrounding the white blood cells. This characteristic sets Dataset 2 apart
from Dataset 1.

####  4.1.3 LISC.

[15] The LISC dataset consists of hematological images taken from the
peripheral blood of healthy subjects. The dataset contains images of normal
leukocytes, including basophils, eosinophils, lymphocytes, monocytes, and
neutrophils, manually segmented and classified by experts. The images,
collected from 8 normal subjects and amounting to 400 samples from 100
microscope slides, are in BMP format with a resolution of 720×576 pixels.

###  4.2 Settings

The labeled set accounts for 1/4 of the size of the training set. This
proportion is in accordance with the setting used in the semi-supervised
learning setting, such as ST++ [16].

The experiment was conducted on the RTX3090. The experiment was run under two
settings: With FixMatch and without FixMatch. Under each of these settings, we
employ three frameworks in semi-supervised learning:

  * •

SupOnly: The model is trained on the labeled dataset only.

  * •

ST: The model is trained in a standard self-training manner.

  * •

ST++: The model is trained in a self-traning++ framework according to the
pseudo-code 2.

Without consistency, the batch size for Dataset 1, Dataset 2, and LISC was 16.
With the consistency incorporated, the batch size for the labeled and the
unlabeled data is 16. For Dataset 2 and LISC, the labeled data’s batch size is
16 and the unlabeled data’s batch size is 8. The network architectures are
DeepLabV3+, DeepLabv2, and PSPNet. Each network architecture consists of the
ResNet-50 and ResNet-101 backbone, but DeepLabv2 only has ResNet-101 backbone.

We will also run two algorithms: Self-Training (ST) and Self-Training++ (ST++)

##  5 Results

Table 1: The mIOU of the models run under different networks with different
backbones.

Model | Backbone | Backbone | Dataset
---|---|---|---
|  |  | Dataset 1 | Dataset2 | LISC
|  |  | SupOnly | ST | ST++ | SupOnly | ST | ST++ | SupOnly | ST | ST++
|  | without consistency | 89.49 | 90.71 | 90.02 | 86.19 | 86.43 | 87.34 | 75.2 | 76.67 | 77.64
DeepLabv3 | ResNet-50 | with consistency | 87.81 | 90.04 | 90.69 | 84.19 | 87.02 | 87.37 | 73.96 | 76.48 | 76.49
|  | without consistency | 88.86 | 89.83 | 89.99 | 86.03 | 86.35 | 87.22 | 75.95 | 75.87 | 76.55
DeepLabv3 | ResNet-101 | with consistency | 86.61 | 89.8 | 90.29 | 85.15 | 87.28 | 87.09 | 75.09 | 77.02 | 76.6
|  | without consistency | 82.54 | 84.21 | 85.01 | 74.46 | 79.7 | 78.74 | 72.2 | 75.27 | 74.84
DeepLabv2 | ResNet-101 | with consistency | 80.08 | 85.05 | 85.07 | 71.47 | 79.32 | 80.38 | 70.47 | 74.71 | 75.75
|  | without consistency | 80.27 | 83.08 | 83.64 | 75.76 | 80.2 | 80.92 | 71.17 | 74.69 | 75.12
PSPNet | ResNet-50 | with consistency | 79.02 | 83.44 | 83.76 | 74.95 | 80.31 | 81.79 | 71.53 | 75.5 | 75.39
|  | without consistency | 81.56 | 83.93 | 84.3 | 77.76 | 80.39 | 81.35 | 72.91 | 75.19 | 75.17
PSPNet | ResNet-101 | with consistency | 80.69 | 84.56 | 84.58 | 77.33 | 79.8 | 80.11 | 71.35 | 74.96 | 74.82

From the above results 1, there can be two conclusions:

  * •

There is a performance improvement when transitioning from Supervised-Only
(SupOnly) to Self-Training (ST) and further to ST++, with ST++ demonstrating
the highest level of performance.

  * •

With FixMatch incorporated, the performance on ST and ST++ framework improves
in certain scenarios.

##  6 Discussion

###  6.1 Performance on the supervised training stage when incorporating
FixMatch

When incorporating FixMatch into the supervised training stage, there was a
noticeable decrease in performance on the labeled images. This is a commonly
seen trade-off when we try to do two tasks in 1 stage.

Table 2: Noticeable decrease when incorporating FixMatch into the model. | Dataset1 | Dataset2 | LISC
---|---|---|---
| $\times$ |  | $\times$ |  | $\times$ |
DeepLabv3+ - ResNet-50 | 89.49 | 87.81 | 86.19 | 84.19 | 75.2 | 73.96
DeepLabv3+ - ResNet-101 | 88.86 | 86.61 | 86.03 | 85.15 | 75.95 | 75.09
DeepLabv2 - ResNet101 | 82.54 | 80.08 | 74.46 | 71.47 | 72.2 | 70.47
PSPNet - ResNet50 | 80.27 | 79.02 | 75.76 | 74.95 | 71.17 | 71.53
PSPNet - ResNet101 | 81.56 | 80.69 | 77.76 | 77.33 | 72.91 | 71.35

However, the performance in the later stage remains the same or even improves,
suggesting that the trade-off in the first stage is sometimes necessary to
improve the performance when extending the labeled images with unlabeled ones.

###  6.2 The quality of the pseudo-masks produced when incorporating FixMatch

The quality of the pseudo-masks between the two settings of FixMatch and non-
FixMatch approaches varies. On the one hand, the pseudo-masks of easy,
reliable images remain essentially the same between the two settings.

Figure 2: Sample of Ground truth mask, Non-FixMatch, and FixMatch mask
generated from the LISC dataset.

However, for more complicated, unreliable images, the masks generated can have
a large deviation from the ground truth mask.

Figure 3: Deviations from the Ground truth mask of Non-FixMatch and FixMatch
mask generated from the LISC dataset.

This suggests that incorporating FixMatch into the model has yet to be able to
take on the nuances presented in challenging images.

## References

  * [1] International Agency for Research on Cancer. Retrieved from https://gco.iarc.fr/today/data/factsheets/populations/900-world
-fact-sheets.pdf
  * [2] VINMEC. (n.d.). Cac dau hieu canh bao ung thu mau. Retrieved from https://www.vinmec.com/vi/ung-buou-xa-tri/thong-tin-suc-khoe/
cac-dau-hieu-canh-bao-ung-thu-mau/

  * [3] Jähne, B. (2005). Digital image processing. Springer Science & Business Media.
  * [4] O. Chapelle, B. Schlkopf, and A. Zien, “Semi-supervised learning,” IEEE Transactions on Neural Networks, vol. 20, 2006.
  * [5] Yang, X., Song, Z., King, I., & Xu, Z. (2022). A survey on deep semi-supervised learning. IEEE Transactions on Knowledge and Data Engineering.
  * [6] Peláez-Vegas, A., Mesejo, P., & Luengo, J. (2023). A Survey on Semi-Supervised Semantic Segmentation. arXiv preprint arXiv:2302.09899.
  * [7] Tarvainen, A., & Valpola, H. (2017). Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. Advances in neural information processing systems, 30.
  * [8] Yuan, J., Liu, Y., Shen, C., Wang, Z., & Li, H. (2021). A simple baseline for semi-supervised semantic segmentation with strong data augmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 8229-8238).
  * [9] Chen, X., Yuan, Y., Zeng, G., & Wang, J. (2021). Semi-supervised semantic segmentation with cross pseudo supervision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 2613-2622).
  * [10] Ouali, Y., Hudelot, C., & Tami, M. (2020). Semi-supervised semantic segmentation with cross-consistency training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 12674-12684).
  * [11] An, S., Zhu, H., Zhang, J., Ye, J., Wang, S., Yin, J., & Zhang, H. (2022). Deep Tri-Training for Semi-Supervised Image Segmentation. IEEE Robotics and Automation Letters, 7(4), 10097-10104.
  * [12] Liu, Y., Tian, Y., Chen, Y., Liu, F., Belagiannis, V., & Carneiro, G. (2022). Perturbed and strict mean teachers for semi-supervised semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 4258-4267).
  * [13] Wu, Y., Liu, C., Chen, L., Zhao, D., Zheng, Q., & Zhou, H. (2023). Perturbation consistency and mutual information regularization for semi-supervised semantic segmentation. Multimedia Systems, 29(2), 511-523.
  * [14] Zheng, X., Wang, Y., Wang, G., & Liu, J. (2018). Fast and robust segmentation of white blood cell images by self-supervised learning. Micron, 107, 55-71.
  * [15] Rezatofighi, S. H., & Soltanian-Zadeh, H. (2011). Automatic recognition of five types of white blood cells in peripheral blood. Computerized Medical Imaging and Graphics, 35(4), 333-343.
  * [16] Yang, L., Zhuo, W., Qi, L., Shi, Y., & Gao, Y. (2022). St++: Make self-training work better for semi-supervised semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 4268-4277).
  * [17] Sohn, K., Berthelot, D., Carlini, N., Zhang, Z., Zhang, H., Raffel, C. A., … & Li, C. L. (2020). Fixmatch: Simplifying semi-supervised learning with consistency and confidence. Advances in neural information processing systems, 33, 596-608.