# Momentum2 Teacher: Momentum Teacher with Momentum Statistics
for Self-Supervised Learning

Zeming Li  Songtao Liu  Jian Sun
MEGVII Technology
{lizeming, liusongtao<EMAIL_ADDRESS>
###### Abstract

In this paper, we present a novel approach, Momentum2 Teacher, for student-
teacher based self-supervised learning. The approach performs momentum update
on both network weights and batch normalization (BN) statistics. The teacher’s
weight is a momentum update of the student, and the teacher’s BN statistics is
a momentum update of those in history. The Momentum2 Teacher is simple and
efficient. It can achieve the state of the art results (74.5%) under ImageNet
linear evaluation protocol using small-batch size(, 128), without requiring
large-batch training on special hardware like TPU or inefficient across GPU
operation (, shuffling BN, synced BN). Our implementation and pre-trained
models will be given on
GitHub111https://github.com/zengarden/momentum2-teacher.

##  1 Introduction

Student-teacher framework is one of the key ingredients in current state-of-
the-art self-supervised visual representation learning, (_e.g._ , MoCo [21,
9], BYOL [19]). The framework learns two networks with the same architecture,
where the student network is trained with gradient back-propagation and the
teacher network is a momentum version of the student network.

In the student-teacher framework, as illustrated in Figure 1 (a)(b), one image
is augmented into two different views for a student and a teacher. The student
is trained to predict the representation of the teacher, and the teacher is
updated with a “momentum update” (exponential moving average) of the student.
The success of MoCo and BYOL has proven the effectiveness of the student-
teacher framework. With the momentum update, the teacher obtains more stable
parameters in a temporal ensembling manner [45].

Besides network parameters, stable statistics (i.e. batch normalization) is
another critical factor for training modern deep networks, especially for
self-supervised learning. MoCo [21] adopts a “shuffling batch normalization
(shuffling BN)” on the teacher, which shuffles the sample order in the current
batch before distributing it among multiple GPUs and rolls back after
encoding. This operation prevents the information leaking of the intra-batch
communication and still allows training to benefit from BN [21]. However,
shuffling BN only operates on the samples independently for each GPU with
small sizes (typical 32 images per GPU), which limits the benefit from larger
batch normalization with more stable statistics.

By leveraging specially designed TPU, BYOL [19] uses larger batch (, 4096) on
both student and teacher to attain more stable statistics, which is equivalent
to perform a “synchronized batch normalization (synced BN)” on multiple GPUs.
The leading performances of BYOL indicates that stable statistics is
important. When training with small batch size (, 128 for BN), its performance
significantly downgrades [19]. The demand of large batching is unfriendly in
many scenarios. Accessing TPU is not common to the research community, and
synced BN is very inefficient across multiple GPUs on multiple machines.

Figure 1: Pipeline of student-teacher based self-supervised approaches. The
positive pairs ($v,v^{\prime}$), the different augment views of the same
image, are fed into student and teacher networks respectively to learn the
representations ($z,z^{\prime}$). In MoCo [21], the student is trained with
(small-batch) BN and the teacher uses shuffling BN.
($v^{\prime}_{0},v^{\prime}_{1},v^{\prime}_{2},...$) are negative pairs from
other images. In BYOL [19], both student and teacher are trained with large
batch and BN statistics with synchronized batch-normalization (Synced BN
[40]). In our Momentum2 Teacher, the student use small-batch statistics to
keep efficiency, and the teacher is equipped with a simple and efficient
“Momentum BN” to obtain more stable statistics.

In this paper, we propose _Momentum 2 Teacher_, which applies momentum update
on both network parameters and BN statistics to obtain a more stable teacher.
As shown in Figure 1 (c), we replace all the batch normalization layers of the
teacher with a simple and efficient operation, “Momentum BN”, to obtain more
stable statistics. Momentum BN conducts momentum update on the historical
batch statistics (_i.e._ , mean and standard deviation) to normalize features,
instead of just using current mini-batch calculations. This operation enables
small batch size training (128) to generate more stable statistics while
keeping the efficiency. Furthermore, as the teacher does not need back-
propagation, Momentum BN does not have a gradient issue of the
indifferentiable history statistics. Our method can achieve the state-of-the-
art result, without resorting to large batch training on TPU or slow synced BN
operation.

The main contributions of this work are:

  1. 1

We propose a novel approach, named _Momentum 2 Teacher_, which keeps
efficiency and hardware-friendly of small batch training while achieves
competitive performance against large batch training.

  2. 2

The core of our method, Momentum BN, which conducts momentum update on the
batch statistics, can benefit all student-teacher based self-supervised
methods. It can improves both MoCo and BYOL.

  3. 3

We obtain the state-of-the-art result, 74.5% Top-1 accuracy, on ImageNet under
the linear evaluation protocol.

##  2 Related Works

##### Self-supervised learning:

Self-supervised approaches have largely reduced the performance gap between
supervised models and even achieved superior results on down-stream vision
tasks. Contrastive learning measures the (dis)similarities of the sample pairs
in a latent space, such as [21, 9, 8, 47, 37, 24, 3, 25, 46, 33]. Pretext
tasks are also heavily researched topic. Some mainly generate pseudo labels
by, e.g., clustering features [5, 6, 7], augmenting different views of the
single image (“exemplar”) [17], relative and ordered patches [35, 14, 15], or
consecutiveness in videos [51, 38]. Others propose to recover the input from
corruption, e.g., inpainting [39], denoising [50] and colorization [60, 61]
with auto-encoders and GAN.

##### Student-teacher framework:

Mean-teacher [45] introduces a student and teacher (moving-averaged version of
the student) network to learn with each other. MoCo [21, 9] combines the
contrastive mechanism and the student-teacher framework [45], with a memory
bank to maintain a large number of negative samples. BYOL [19] further removes
the negative samples, using an additional predictor to avoid collapsing. It
relies on the large number of positive samples, typically 4096 for achieving
the best performance.

##### Normalization:

BN [27] has widely proven effectively and efficiently in most of the vision
tasks. It normalizes the internal feature maps using channel-wise statistics
along batch dimension. In practice, BN relies on sufficient batch-size, which
is not easily satisfied, especially when large-batch training is acquired in
self-supervised learning. There are many techniques proposed to maintain or
approximate large batch statistics. Synced BN [40] increases the batch-size by
computing the mean and variance across multiple devices (GPUs), however,
introducing a lot of overhead. Batch Renormalization [26] and EvalNorm [43]
correct batch statistics during training and inference, while compared to
Synced BN they generally perform worse.

Moving averaged batch normalization [56] and online normalization [11] adopt
similar momentum updating of statistics like our momentum BN during the
forward pass. However, they need further correct backpropagation for valid SGD
optimization, which requires additional computation and memory resources. As
there is no back pass within the teacher network, our momentum BN do not need
this gradient revision, making it more efficient.

Another family of normalization is functionally based. Instead of normalizing
across samples, layer-norm [2] performs normalizes across features which makes
it irrelevant to batch-size. Group normalization [52] and instance
normalization [49] further extend this by partitioning features into groups.
Normalization which is applied to network weights are also proposed, such as
weight normalization [42] and normalization propagation [1]. The very recent
study [41] shows that BYOL even works using functional-based normalization.
But the result is still slightly worse than its counterpart like Synced BN
[40], which is consistent with the observation under supervised training.

##  3 Momentum2 Teacher

We first give our motivation by studying the impact of BN statistics in a
series of controlled experiments. Then, we present our method to increase the
stability of BN statistics.

###  3.1 Importance of Stable Statistics

To analysis the role of BN statistics under the student-teacher framework, we
design a set of exploratory experiments on STL10 [12] dataset, based on BYOL
baseline [19].

##### Experiment setup:

STL10 contains 96$\times$96 pixel RGB images belonging to ten classes. For
each class, it provides 500 images for training and 800 images for testing. We
use ResNet18 as the basic network for fast ablation. More training details
including image augmentation can be referred in Sec.3.4.

Table 1: Comparison on STL10 with different BN statistics in student and teacher. BN and Synced BN indicate the BN statistics are calculated within a single GPU and cross GPUs. All experiments are conducted with 400 epochs. We also report the training speed (seconds/iteration). Student | Teacher | Top1 | Sec./Iter
---|---|---|---
Synced BN | Synced BN | 88.06 | 1.1s
BN | Synced BN | 87.80 | 0.39s
Synced BN | BN | 87.12 | 0.49s
BN | BN | 84.16 | 0.24s
BN | Momentum BN | 88.18 | 0.25s
Synced BN | Momentum BN | 88.25 | 0.50s

##### Observations.

BYOL baseline applies “Synced BN” (performing BN cross all GPUs) on both
student and teacher. Since the teacher does not backprogpagate the gradient,
the role of BN in the teacher more replies on the statistics collected in the
forward pass. To exam if Synced BN matters, we replace Synced BN with “BN”
(performing BN on single GPU, independently) in student or teacher or both. As
showed in Table 1, we have four observations.

  1. 1.

Synced BN is critical. Without Synced BN, the accuracy significantly drops
from 88.06 to 84.16. Similar result was also reported in [19] when using small
batch size. This verified the importance of stable statistics in BN.

  2. 2.

Synced BN slows down the training speed. Because Sycned BN performs cross GPU
oprations, the communication overhead is very significant. For example, the
speed of Synced BN/Synced BN combination is more than four times slower than
BN/BN.

  3. 3.

It is not essential to apply Synced BN on both student and teacher. We can get
decent results (87.80 or 87.12) with a stable teacher or a stable student.
This will enable us to decouple design of BN in student and teacher.

  4. 4.

A stable teacher is essential. At the third row in the table, although we did
not directly apply Synced BN in the teacher, the teacher still benefit from
stable statistics by copying BN parameters in a moving averaged way. Directly
applying Synced BN on the teacher is better than on the student (87.80 v.s.
87.12).

Based on the above observations, we conclude that the stable BN statistics
brought by large batch training is crucial for student-teacher self-supervised
learning. Although stable statistics can be obtained from large-batched
samples, it hurts efficiency. Next, we will introduce a new, simple method for
higher accuracy and better efficiency.

###  3.2 Momentum BN

Synced BN simply uses large batch to obtain stable statistics. We note that in
the student-teacher framework, we do not propagate gradient in the teacher. If
we can leverage this characteristic, we may obtain stable statistics using
small batch.

We follow the annotation of the batch-normalization. Considering a mini-batch
$\mathcal{B}$ of size $m$, then we have $m$ values $x_{1..m}$ of feature. BN
operation first calculate two important statistics mean $\mu$ and variance
$\sigma$:

$\displaystyle\mu$ $\displaystyle=\frac{1}{m}\Sigma_{i=1}^{m}x_{i},$ (1)
$\displaystyle\sigma$
$\displaystyle=\frac{1}{m}\Sigma_{i=1}^{m}(x_{i}-\mu)^{2}.$

Then, BN performs a linear transformation to get final output $y_{1..m}$:

$\displaystyle\widehat{x_{i}}=\frac{x_{i}-{\mu}}{\sqrt{\sigma+\epsilon}},$ (2)
$\displaystyle y_{i}=\gamma\widehat{x}_{i}+\beta,$

where $\gamma$ and $\beta$ are two learnable parameters.

In the student-teacher framework, two statistics mean $\mu$ and variance
$\sigma$ are calculated from samples (encoded by the teacher) in current
batch, while $\gamma$ and $\beta$ are momentum version of the student.

Since we do not have to estimate $\gamma$ and $\beta$, and motivated by that
the teacher is a kind of temporal ensemble of the student, we perform a
momentum update of the BN statistics:

$\displaystyle\mu_{t}$ $\displaystyle=(1-\alpha)\mu_{t}+\alpha\mu_{t-1},$ (3)
$\displaystyle\sigma_{t}$
$\displaystyle=(1-\alpha)\sigma_{t}+\alpha\sigma_{t-1},$

where $\alpha$ is a momentum coefficient. We simply use an exponential moving
average of historical BN statistics to make the teacher more stable. We call
the BN with the above momentum update “ _Momentum BN_ ”.

Using momentum statistics is not new in the deep learning. For example,
running-mean and running variance have been provided in the deep learning
framework like Tensor-Flow for better validation and inference after the
training. Here momentum statistics are used for training in the student-
teacher framework.

##### Lazy update.

Information leaking is one of the main issues in self-supervised learning. For
two samples $v,v^{\prime}$ of the same image, BYOL _sequentially_ calculates
symmetrized losses as follows:

$\displaystyle\mathcal{L}_{1}$
$\displaystyle=||\text{student}(v),\text{teacher}(v^{\prime})||,$ (4)
$\displaystyle\mathcal{L}_{2}$
$\displaystyle=||\text{student}(v^{\prime}),\text{teacher}(v)||.$

After calculating the loss $\mathcal{L}_{1}$, the statistics of $v^{\prime}$
will be fed into the teacher model. When we calculate the loss
$\mathcal{L}_{2}$, if we use Momentum BN straightforward, we will include the
statistics of $v^{\prime}$ (as historical statistics) into the teacher. This
will make the learning more trivial and hurt the performance.

We address this issue by lazily updating Momentum BN statistics. We first
perform Momentum BN for L1 and L2 independently, using the statistics of
previous batch $\\{\mu_{t-1},\sigma_{t-1}\\}$:

$\displaystyle\\{\mu_{t}^{1},\sigma_{t}^{1}\\}$
$\displaystyle=(1-\alpha)\\{\mu_{t}^{v^{\prime}},\sigma_{t}^{v^{\prime}}\\}+\alpha\\{\mu_{t-1},\sigma_{t-1}\\},$
(5) $\displaystyle\\{\mu_{t}^{2},\sigma_{t}^{2}\\}$
$\displaystyle=(1-\alpha)\\{\mu_{t}^{v},\sigma_{t}^{v}\\}+\alpha\\{\mu_{t-1},\sigma_{t-1}\\},$

Then, we conduct BN transformation via $\\{\mu_{t}^{1},\sigma_{t}^{1}\\}$ and
$\\{\mu_{t}^{2},\sigma_{t}^{2}\\}$. Last, we update BN statistics:

$\displaystyle\\{\mu_{t},\sigma_{t}\\}$
$\displaystyle=(1-\alpha)\\{\frac{\mu_{t}^{v^{\prime}}+\mu_{t}^{v}}{2},\frac{\sigma_{t}^{v^{\prime}},\sigma_{t}^{v}}{2}\\}+\alpha\\{\mu_{t-1},\sigma_{t-1}\\}$
(6)

##### Results.

We replace Synced BN with Momentum BN in the teacher. As shown in Table 1,
Momentum BN achieves better performance (88.18 vs 87.8) without requiring
cross machine communication. It is as fast as we use small batch BN.

We also apply Momentum BN on the student. The improvement is marginal,
verifying that the stable statistics in teacher is essential.

###  3.3 Momentum2 Teacher

From Table 1, we can see that student with a small batch (32) statistics
already performs closely to that using large-batch (2048) Synced BN, which
demonstrates that student can be designed with much smaller batch-size for BN
statistics compared to teacher. Therefore, in this paper, we recommend the
combination of “ _student with small batch_ \+ _teacher with (small batch)
momentum BN_ ” as our main method. This method has the best trade-off between
accuracy and efficiency.

Because the teacher in our method uses momentum mechanism twice, one for
weights update from student, the other for calculating BN statistic, we call
our method “ _Momentum 2 Teacher_”.

###  3.4 Implementation Details

Figure 2: Momentum2 Teacher based on BYOL. Teacher maintains stable BN
statistics to improve the performance, and student uses plain BN to keep the
efficiency.

_Baseline:_ Figure 2 shows our method, which uses BYOL [19] as the baseline.
BYOL learns representations by maximizing the similarity between two different
augmented views $v$ and $v^{\prime}$ from the same image $X$. $v$ is passed
into the student which consists of a basic encoder $f_{\theta}$, a MLP
$g_{\theta}$, and a predictor $q_{\theta}$; and $v^{\prime}$ is fed into the
teacher has only a basic encoder $f_{\xi}$ and a MLP $g_{\xi}$. The parameters
of $f_{\xi}$ and $g_{\xi}$ are momentum update of $f_{\theta}$ and
$g_{\theta}$:

$\displaystyle f_{\xi}\leftarrow(1-m)f_{\xi}+f_{\theta},$ (7) $\displaystyle
g_{\xi}\leftarrow(1-m)g_{\xi}+g_{\theta}.$

where $m$ is a momentum coefficient.

_Image augmentations:_ We use the same set of image augmentations as in BYOL
and SimCLR [8]. Specifically, a crop of fixed size is taken from a random
resized and horizontal flipped image, followed by a commonly used color
distortion. Then, Gaussian blur and an optional gray-scale are applied to
these patches. Finally, solarization is adopted whose probability is set to 0
for student and 0.2 for teacher.

_Architecture:_ On STL10, we use ResNet-18 [23] as the basic encoder
($f_{\theta},f_{\xi}$), which produces a feature with 512 dimension
($y,y^{\prime}$) by average pooling. The dimension is set to 512 for first
linear layer in the MLP and 128 for the second linear layer.

_Training:_ We train STL10 with 64 2080TI GPUs in order to simulate the cross
machine (8 machines) communication in Synced BN. SGD with momentum of 0.9 is
adopted without LARS [57]. All experiments involve 32 image crops for each GPU
unless special statements. At the end of the training, we dump the model from
teacher. Learning rate is decay by cosine strategy [32]. Basic learning rate
is set to 0.1. It warm-ups 10 epochs with 0.001 factor and is scaled linearly
[18] with batch-size. Weight-decay is set to 1e-4. The momentum coefficient of
$m$ starts from $m_{base}=0.032$ and is decreased to zero at the end of the
training. In our Momentum BN, the momentum coefficient $\alpha$ starts from
$\alpha_{base}=1$ and is decreased to 0 with cosine schedule:

$\displaystyle\begin{split}m\leftarrow m_{base}\times(cos(\pi k/K)+1)/2,\\\
\alpha\leftarrow\alpha_{base}\times(cos(\pi k/K)+1)/2,\end{split}$ (8)

where, $k$ is the current iteration and $K$ is the total number of iterations.

_Evaluation:_ We follow the linear classification protocol. The features are
frozen and attached by a new linear classifier normalized by BN to fine-tune
with given classes. On STL10, we train 80 epochs using learning rate starting
from 0.5 (for 256 batch-size) and decayed with cosine schedule. We train on 8
2080 GPUs using SGD with a momentum of 0.9, without weight-decay.

##  4 Experiments

In this section, we perform training on ImageNet and evaluate the model with
the linear evaluation of classification task and some downstream tasks on
COCO. Our setup is briefed as follows:

ImageNet [13] ILSVRC-2012 dataset has about 1.28 million images belonging to
1000 different classes. The class labels are ignored in the self-supervised
learning. The image augmentation is same as STL10’s except that image crop
uses 224$\times$224.

We use ResNet-50 [23] as the encoder. After average pooling, the feature
($y,y^{\prime}$) has 2048 dimensions. MLP uses a 4098 layer and a 256 layer to
generate output ($z,z^{\prime}$).

We use 128 2080 GPUs with LARS [57] optimizer, but bias and BN weight are
excluded from the LARS adaptation and weight-decay. Basic learning rate is set
to 0.3, and weight-decay is set to 1.5e-6. The momentum coefficient of
parameter $m$ starts from $m_{base}=0.01$ and is gradually decreased to zero.
Fine-tuning uses learning rate 0.2.

###  4.1 Effectiveness of Momentum BN

We first validate the effectiveness of Momentum BN on both BYOL and MoCo
frameworks. Table 2 shows BYOL without Synced BN decreases from 72.5 to 61.5,
when we use batch size 128. Our Momentum2 Teacher significantly boosts the
performances to 72.9. It is worth noting that Momentum2 Teacher trains as fast
as BYOL w/o Synced BN, demonstrating the efficiency (0.6s v.s. 5.25s).

Table 2: Comparison of accuracies between BYOL [9] and Momentum2 Teacher. Following BYOL, we train with 300 epochs. “BN” indicates the number of samples to calculate BN. Model | Top1 | Top5 | BN | Sec./Iter
---|---|---|---|---
BYOL w/ Synced BN | 72.5 | 87.6 | 4096 | 5.25s
BYOL w/o Synced BN | 61.5 | 84.6 | 32 | 0.6s
Momentum2 Teacher | 72.0 | 90.6 | 32 | 0.6s
Momentum2 Teacher | 72.9 | 90.6 | 128 | -

Then, we replace Shuffling BN in MoCoV2 [9] with our Momentum BN. Following
the practice of MoCoV2, we fix the momentum coefficient $\alpha$ in Momentum
BN to 0.064. For fast experiment, we train on 32 2080TI GPUs222Training on 4
machines with 32 GPUs, slightly reduces the MoCo baseline (trained on 8 GPUs)
from 67.5 to 66.8.. Table 3 shows that more stable statistics also benefits
MoCo, which validates the generality of Momentum BN in the student-teacher
framework. Moreover, Momentum BN is near twice faster than shuffling BN.

Table 3: Accuracies of MoCoV2 [9] with Shuffling BN and Momentum BN on ImageNet. Following MoCo, we train with 200 epochs. | Top1 | Top5 | Sec./Iter
---|---|---|---
MoCo w/ Shuffling BN | 66.8 | 87.6 | 0.65s
MoCo w/ Momentum BN | 67.8 | 88.0 | 0.35s

###  4.2 Small Batch-Size

Next, we compare BYOL with Momentum2 Teacher at different batch sizes and the
same training schedule of 300 epochs. We keep using 128 GPUs and changes the
number of samples within each GPU. Thus the total batch-sizes will be changed.
To achieve stable performances, we extend the linear scaling rule [18] by
introducing a new equivalence rule for the parameter’s momentum coefficient:
When the mini-batch size is multiplied by $k$, multiply the learning rates by
$k$ and multiply the basic momentum coefficient $m_{base}$ of parameters by
$k$ simultaneously.

Figure 3: Comparison of using different batch-size to calculate batch-norm in student-network. We plot the ‘number” of batch-size to calculate the batch-normalization. BYOL results are referred from its paper. Momentum2 teacher significantly outperforms BYOL when the batch size is reduced. All results are pre-trained with 300 epochs. Table 4: Comparison among different BN batch-sizes (i.e. the number of samples per GPU) on ImageNet with 300 epochs. BYOL results are referred from its paper conducted by accumulating _N_ -steps gradients in parallel.  Batch |  |  |  |  |  |  |
---|---|---|---|---|---|---|---
size | 16 | 32 | 48 | 64 | 128 | 256 | 512
BYOL | - | - | - | 59.7 | 69.6 | 71.8 | 72.2
Ours | 68.3 | 72.0 | 72.5 | 72.6 | 72.9 | - | -

From Table 4 and Fig. 3, we can see that the number of samples to calculate BN
statistics in the student can be very small in our Momentum2 Teacher. Using a
batch-size of 32 can still get a decent result, which is on par with BYOL at
512 batch-size. On the contrary, the performance of BYOL rapidly deteriorates
when reducing batch size.

###  4.3 Training on a Single Machine

Thanks to small batch-size training, Momentum2 teacher is also validated on a
single machine with 8 GPUs. We use a learning rate of $lr\times$BatchSize/256
(linear scaling rule [18]) with a basic $lr=0.05$, and then adopt a cosine
decay strategy. The weight decay is 0.0001 and the momentum of SGD is 0.9.
Large-batch optimizers such as LARS are not involved.

Results are shown in Tab. 5. Compared with large batch-size training, our
method achieve even better performances on a single machine with small batch-
sizes. We also compare Momentum2 teacher with recent work SimSiam [10].
Momentum2 teacher consistently outperforms other counterparts at different
settings. It is worth noting that our superior results is attained without
synced BN operation, which makes it more efficient for training.

Table 5: Top 1 accuracies of small batch training on ImageNet linear classification. Our result is obtained on a single machine with 8 V100 GPUs, using batch size of 1024. Method |  | batch
---
size
100e | 200e | 300e
Ours | 256 | 70.4 | - | -
Ours | 1024 | 70.7 | 72.7 | 73.8
SimSiam [10] | 256 | 68.1 | 70.0 | -
BYOL [19] | 4096 | - | - | 72.5

###  4.4 Parameter Sensitivity

Momentum2 teacher introduces an additional coefficient $\alpha$ for updating
Momentum BN in the teacher. In the previous experiments, $\alpha$ was
_dynamically_ set from 1 to 0: at the beginning of the training, BN mainly
relies on the statistics of the current mini-batch for normalizing; as the
training stabilizes, it will rely more on the statistics attained by momentum
of historical iterations.

Table 6 compares dynamically adjusted $\alpha$ and fixed $\alpha$. All of
fixed $\alpha$ get inferior results. In dynamic adjustment, $\alpha$ will
eventually drop to 0, which means it can calculate statistics over the entire
data set, making more robust results. Table 7 comares the choices of different
$\alpha_{base}$ (in Eqn. 8). Large $\alpha_{base}$ consistently improves the
performance, further validating the early stage of the training should rely
more on current batch samples.

Table 6: Varying $\alpha$ on ImageNet. All results are pre-trained with 300 epochs. $\alpha$ | 1 $\to$ 0 | 1 | 0.5 | 0.2 | 0.1 | 0.01
---|---|---|---|---|---|---
Top1 | 72.0 | 61.6 | 71.2 | 69.1 | 69.8 | 68.9
Top5 | 90.6 | 84.6 | 90.0 | 89.1 | 89.2 | 88.4
Table 7: Varying $\alpha_{base}$ on ImageNet. All results are pre-trained with 300 epochs. $\alpha_{base}$ | 1 | 0.75 | 0.5 | 0.2 | 0.1 | 0.05
---|---|---|---|---|---|---
Top1 | 72.0 | 71.6 | 71.2 | 69.7 | 69.5 | 69.6
Top5 | 90.6 | 90.4 | 90.1 | 89.0 | 89.0 | 88.8
Table 8: Comparisons under the linear classification protocol on ImageNet. All methods are instantiated with ResNet-50 basic extractor. We report top-1 and top-5 accuracies in % on the test set.  Method | Epoch | Top-1 | Top-5
---|---|---|---
Jigsaw [36] | 90 | 45.7 | -
InstDis [54] | 200 | 56.5 | -
BigBiGAN [16] | - | 56.6 | -
Local Agg. [62] | - | 60.2 | -
CPC v2 [24] | 200 | 63.8 | 85.3
CMC [46] | - | 66.2 | 87.0
SimCLR [8] | 200 | 66.6 | -
MOCOv2 [9] | 200 | 67.5 | -
PCL-v2 [28] | 200 | 67.6 | -
InfoMin Aug [47] | 200 | 70.1 | 89.4
BYOL [19] | 300 | 72.5 | 90.8
Ours | 300 | 72.9 | 91.2
PIRL [33] | 800 | 63.6 | -
SimCLR [8] | 1000 | 69.3 | 89.0
MOCOv2 [9] | 800 | 71.1 | -
InfoMin Aug [47] | 800 | 73.0 | 91.1
BYOL [41] | 1000 | 74.3 | 91.6
Ours | 1000 | 74.5 | 91.7
Table 9: Fine-tuning of proposal based object detection and instance
segmentation on COCO-train2017. We validate on FPN and Cascade R-CNN extended
with mask branch. We use ResNet-50-FPN extractor, and report the bounding box
AP and mask AP on val2017. 200e and 300e indicate the pre-training epochs.

pre-train | APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$ | APmk | AP${}^{mk}_{50}$ | AP${}^{mk}_{75}$
---|---|---|---|---|---|---
random init. | 31.0 | 49.5 | 33.2 | 28.5 | 46.8 | 30.4
supervised | 38.9 | 59.6 | 42.7 | 35.4 | 56.5 | 38.1
MoCo 200e | 38.5 | 58.9 | 42.0 | 35.1 | 55.9 | 37.7
BYOL 300e | 39.6 | 60.9 | 43.3 | 36.7 | 58.0 | 39.3
Ours 300e | 39.7 | 61.2 | 43.2 | 36.8 | 58.0 | 39.6
(a) Mask R-CNN, 1$\times$ schedule

APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$ | APmk | AP${}^{mk}_{50}$ | AP${}^{mk}_{75}$
---|---|---|---|---|---
36.7 | 56.7 | 40.0 | 33.7 | 53.8 | 35.9
40.6 | 61.3 | 44.4 | 36.8 | 58.1 | 39.5
40.8 | 61.6 | 44.7 | 36.9 | 58.4 | 39.7
41.6 | 62.9 | 45.8 | 38.2 | 59.9 | 41.1
41.8 | 62.8 | 45.9 | 38.4 | 60.1 | 41.2
(b) Mask R-CNN, 2$\times$ schedule

pre-train | APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$ | APmk | AP${}^{mk}_{50}$ | AP${}^{mk}_{75}$
---|---|---|---|---|---|---
random init. | 35.3 | 51.0 | 38.3 | 31.0 | 48.8 | 33.3
supervised | 42.1 | 59.8 | 45.9 | 36.4 | 57.1 | 39.3
MoCo 200e | 42.4 | 59.8 | 46.1 | 37.0 | 57.2 | 40.1
BYOL 300e | 43.6 | 61.5 | 47.5 | 38.3 | 59.0 | 41.6
Ours 300e | 43.6 | 61.6 | 47.3 | 38.2 | 58.8 | 41.2
(c) Cascade R-CNN, 1$\times$ schedule

APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$ | APmk | AP${}^{mk}_{50}$ | AP${}^{mk}_{75}$
---|---|---|---|---|---
40.7 | 57.7 | 44.2 | 35.9 | 55.5 | 39.1
43.5 | 61.3 | 47.4 | 37.7 | 58.7 | 40.8
44.3 | 61.9 | 48.1 | 38.7 | 59.5 | 42.0
45.0 | 62.9 | 48.8 | 39.3 | 60.5 | 42.6
45.1 | 63.0 | 49.0 | 39.4 | 60.7 | 42.8
(d) Cascade R-CNN, 2$\times$ schedule

###  4.5 Comparison with State-of-the-art

In this section, we first compare performances of Momentum2 Teacher’s
representation with recent state-of-the-art self-supervised approaches on
ImageNet. As the main merit of self-supervised learning is to learn
_transferrable_ feature, we then measure the transfer capabilities on COCO
[31] and LVIS [20] dataset which includes both annotations for object
detection and segmentation. To achieve better results, the pre-training uses
128 samples within each GPU.

####  4.5.1 Linear Evaluation on ImageNet

As shown in Table 8, our method (ResNet-50 based) obtains 72.9% and 74.5 top-1
accuracies under 300 and 1000 epochs, which outperform the previous state of
the arts. It is worth noting that we only adopt 128 samples for batch-
normalization within student, while BYOL requires 4096. Our superior result is
also achieved without any additional architecture requirement and stronger
augmentation, keeping it simple and effective for practice.

Recent SwAV [7] method can achieve higher accuracy 75.3% by using additional 6
crops within each training iteration to, at the cost of more forward
computations.

Table 10: Fine-tuning of dense object detection on COCO-train2017. We validate
on RetinaNet and FCOS with ResNet-50-FPN and report the AP of bounding boxes.

pre-train | APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$
---|---|---|---
random init. | 24.4 | 38.8 | 25.8
supervised | 37.3 | 56.6 | 39.8
MoCo 200e | 37.2 | 56.4 | 40.0
BYOL 300e | 36.4 | 55.9 | 39.2
Ours 300e | 35.9 | 55.3 | 38.4
(a) RetinaNet, 1$\times$ schedule

APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$
---|---|---
31.2 | 48.1 | 33.3
38.7 | 58.2 | 41.4
39.0 | 58.3 | 41.6
39.5 | 59.3 | 42.6
39.4 | 59.3 | 42.2
(b) RetinaNet, 2$\times$ schedule

APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$
---|---|---
25.0 | 39.2 | 26.4
38.7 | 57.6 | 41.9
39.0 | 57.5 | 42.0
37.6 | 56.0 | 41.0
37.7 | 56.8 | 40.8
(c) FCOS, 1$\times$ schedule

APbb | AP${}^{bb}_{50}$ | AP${}^{bb}_{75}$
---|---|---
31.8 | 48.2 | 33.6
38.5 | 57.0 | 41.3
38.8 | 57.1 | 41.7
39.5 | 58.3 | 42.5
39.9 | 59.0 | 43.2
(d) FCOS, 2$\times$ schedule

####  4.5.2 Transferring Features

For COCO dataset, we fine-tune on the _train2017_ (about 118k images) set and
tested on _val2017_ (5k). The typical FPN [29] backbone is adopted and further
cooperated with Synced BN. The short-edge of training image is in [640, 800]
sampled with 32 pixels intervals, while fixed at 800 during testing. We follow
the typical 1$\times$ or 2$\times$ training strategy provided by Detectron2
[53] repository. For LVIS 0.5 dataset, we uses 56K images over 1230 categories
for training, and 5k images for validation.

##### Proposal based Detector and Segmentor

We adopt FPN and Cascade R-CNN extended with mask branch [29, 22, 4] to
validate the transfer capabilities on COCO object detection and segmentation.
Following common practice, we stack 4 $3\times 3$ convolutions in the mask
branch and utilize 2 linear layers in the detection branch.

Table 9 shows the bounding box AP and mask AP on COCO _val2017_. For both two-
stage and multi-stage instance segmentor, our method significantly outperforms
the supervised counterpart, achieving comparable results as BYOL with more
efficient implementation. Meanwhile, we surpass another student-teacher based
self-supervised method, MoCo, which uses a small-batch of 32 by a large
margin.

##### Dense Object Detector

Next we compare the representation by fine-tuning dense object detectors
(a.k.a single-stage object detector). We adopt two typical dense detector,
namely anchor-based RetinaNet [30] and anchor-free FCOS [48]. We follow the
baseline provided by Detectron2. Particularly, FCOS utilizes group
normalization [52] for convolutions in detection head, while RetinaNet does
not. Moreover, RetinaNet adopt the standard multi-scale training setting while
FCOS trains with a single-scale of 800 pixels.

Table 10 shows the bounding box AP of the detectors initialized with different
pre-training weights. Our method shows great potential when training with a
long schedule (2$\times$). We achieve comparable performance as BYOL on
RetinaNet while yield much better results on FCOS, outperforming the
supervised counterpart. We observe that both BYOL and our method get inferior
results when training with the 1$\times$ schedule, which we leave its
exploration for future work.

##### LVIS Instance Segmentation

We further transfer the instance segmentation on LVIS which contains about
1000 long-tailed distributed categories. Our method significantly outperforms
the supervised pre-training by an mAP of 1.7% in Tab. 11, validating the
generalization of our approach. We hypothesis that self-supervised pre-
training can benefit more for the task with fewer annotated samples.

Table 11: Fine-tuning of Mask R-CNN on LVIS 0.5. All experiments are instantiated with ResNet-50-FPN and trained with 2$\times$ schedule. pre-train | APmk | AP${}^{mk}_{50}$ | AP${}^{mk}_{75}$
---|---|---|---
random init. | 22.5 | 34.8 | 23.8
supervised | 24.4 | 37.8 | 25.8
MoCo 200e | 24.1 | 37.4 | 25.5
Ours 300e | 26.1 (+1.7) | 39.8 (+2.0) | 28.1 (+2.3)
Table 12: Accuracy of models trained with few image labels. We report top-1 and top-5 accuracies in % on the test set.  Method | epoch | Top-1 | Top-5
---|---|---|---
1% | 10% | 1% | 10%
Supervised[58] | - | 25.4 | 56.4 | 48.4 | 80.4
_Semi-supervised:_ |  |  |  |  |
Pseudolabels [59] | - | - | - | 51.6 | 82.4
VAT [34] | - | - | - | 47.0 | 83.4
UDA [55] | - | - | 68.8 | - | 88.5
FixMatch [44] | - | - | 71.5 | - | 89.1
_Self-supervised:_ |  |  |  |  |
InstDisc [54] | 200 | - | - | 39.2 | 77.4
PIRL [33] | 800 | - | - | 57.2 | 83.8
MoCov2 [9] | 800 | 42.3 | 63.8 | 70.1 | 86.2
PCL [28] | 200 | - | - | 75.3 | 85.6
SimCLR [8] | 1000 | 48.3 | 65.6 | 75.5 | 87.8
SWAV(B=256)[7] | 200 | 51.3 | 67.8 | 76.6 | 88.6
SWAV(B=4096)[7] | 200 | 52.6 | 68.5 | 77.7 | 89.2
SWAV(B=4096)[7] | 800 | 53.9 | 70.2 | 78.5 | 89.9
BYOL [19] | 1000 | 53.2 | 68.8 | 78.4 | 89.0
Ours | 300 | 57.7 | 70.2 | 80.8 | 89.3
Ours | 1000 | 62.3 | 72.2 | 84.1 | 90.1

###  4.6 Semi-Supervised Training on ImageNet

Last, we evaluate the classification performance obtained when fine-tuning
ours representation using a small subset of the ImageNet’s train set with
label information. We use the same fixed splits of 1% and 10 % provided by
SimCLR. The training most follows common semi-supervised protocol in [8, 19,
24], we perform training on 8 V100 GPUs with total batch-size of 1024.
Learning rate is set to 0.08 and decayed by cosine strategy. We fine-tune
20/15 epochs for splits of 1% /10% respectively. After training, the
statistics of BN are recounted for the later testing. We report the top-1 and
top-5 accuracies on the test-set in Tab. 12. Our method consistently
outperforms the previous approaches.

##  5 Conclusion

We have presented Momentum2 Teacher, for self-supervised learning, by
introducing a simple and efficient Momentum BN operation on the teacher. With
a more stable teacher, we are able to use fast small-batch training to obtain
the leading results, which is more friendly to the majority of the
researchers.

## References

  * [1] Devansh Arpit, Yingbo Zhou, Bhargava U Kota, and Venu Govindaraju.  Normalization propagation: A parametric technique for removing internal covariate shift in deep networks.  arXiv preprint arXiv:1603.01431, 2016.
  * [2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  arXiv preprint arXiv:1607.06450, 2016.
  * [3] Philip Bachman, R Devon Hjelm, and William Buchwalter.  Learning representations by maximizing mutual information across views.  In Advances in Neural Information Processing Systems, pages 15535–15545, 2019.
  * [4] Zhaowei Cai and Nuno Vasconcelos.  Cascade r-cnn: Delving into high quality object detection.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6154–6162, 2018.
  * [5] Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze.  Deep clustering for unsupervised learning of visual features.  In Proceedings of the European Conference on Computer Vision (ECCV), pages 132–149, 2018.
  * [6] Mathilde Caron, Piotr Bojanowski, Julien Mairal, and Armand Joulin.  Unsupervised pre-training of image features on non-curated data.  In Proceedings of the IEEE International Conference on Computer Vision, pages 2959–2968, 2019.
  * [7] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  Advances in Neural Information Processing Systems, 33, 2020.
  * [8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  arXiv preprint arXiv:2002.05709, 2020.
  * [9] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He.  Improved baselines with momentum contrastive learning.  arXiv preprint arXiv:2003.04297, 2020.
  * [10] Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning.  arXiv preprint arXiv:2011.10566, 2020.
  * [11] Vitaliy Chiley, Ilya Sharapov, Atli Kosson, Urs Koster, Ryan Reece, Sofia Samaniego de la Fuente, Vishal Subbiah, and Michael James.  Online normalization for training neural networks.  Advances in Neural Information Processing Systems, 32:8433–8443, 2019.
  * [12] Adam Coates, Andrew Ng, and Honglak Lee.  An analysis of single-layer networks in unsupervised feature learning.  In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215–223, 2011.
  * [13] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [14] Carl Doersch, Abhinav Gupta, and Alexei A Efros.  Unsupervised visual representation learning by context prediction.  In Proceedings of the IEEE international conference on computer vision, pages 1422–1430, 2015.
  * [15] Carl Doersch and Andrew Zisserman.  Multi-task self-supervised visual learning.  In Proceedings of the IEEE International Conference on Computer Vision, pages 2051–2060, 2017.
  * [16] Jeff Donahue and Karen Simonyan.  Large scale adversarial representation learning.  In Advances in Neural Information Processing Systems, pages 10542–10552, 2019.
  * [17] Alexey Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox.  Discriminative unsupervised feature learning with convolutional neural networks.  In Advances in neural information processing systems, pages 766–774, 2014.
  * [18] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  arXiv preprint arXiv:1706.02677, 2017.
  * [19] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  arXiv preprint arXiv:2006.07733, 2020.
  * [20] Agrim Gupta, Piotr Dollar, and Ross Girshick.  Lvis: A dataset for large vocabulary instance segmentation.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 5356–5364, 2019.
  * [21] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729–9738, 2020.
  * [22] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In Computer Vision (ICCV), 2017 IEEE International Conference on, pages 2980–2988. IEEE, 2017.
  * [23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [24] Olivier J Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, SM Eslami, and Aaron van den Oord.  Data-efficient image recognition with contrastive predictive coding.  arXiv preprint arXiv:1905.09272, 2019.
  * [25] R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio.  Learning deep representations by mutual information estimation and maximization.  arXiv preprint arXiv:1808.06670, 2018.
  * [26] Sergey Ioffe.  Batch renormalization: Towards reducing minibatch dependence in batch-normalized models.  In Advances in neural information processing systems, pages 1945–1953, 2017.
  * [27] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  arXiv preprint arXiv:1502.03167, 2015.
  * [28] Junnan Li, Pan Zhou, Caiming Xiong, Richard Socher, and Steven CH Hoi.  Prototypical contrastive learning of unsupervised representations.  arXiv preprint arXiv:2005.04966, 2020.
  * [29] Tsung-Yi Lin, Piotr Dollár, Ross B Girshick, Kaiming He, Bharath Hariharan, and Serge J Belongie.  Feature pyramid networks for object detection.  In CVPR, volume 1, page 4, 2017.
  * [30] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  In Proceedings of the IEEE international conference on computer vision, pages 2980–2988, 2017.
  * [31] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In European Conference on Computer Vision, pages 740–755. Springer, 2014.
  * [32] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  arXiv preprint arXiv:1608.03983, 2016.
  * [33] Ishan Misra and Laurens van der Maaten.  Self-supervised learning of pretext-invariant representations.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6707–6717, 2020.
  * [34] Takeru Miyato, Shin-ichi Maeda, Masanori Koyama, and Shin Ishii.  Virtual adversarial training: a regularization method for supervised and semi-supervised learning.  IEEE transactions on pattern analysis and machine intelligence, 41(8):1979–1993, 2018.
  * [35] Mehdi Noroozi and Paolo Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In European Conference on Computer Vision, pages 69–84. Springer, 2016.
  * [36] Mehdi Noroozi and Paolo Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In European Conference on Computer Vision, pages 69–84. Springer, 2016.
  * [37] Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  arXiv preprint arXiv:1807.03748, 2018.
  * [38] Deepak Pathak, Ross Girshick, Piotr Dollár, Trevor Darrell, and Bharath Hariharan.  Learning features by watching objects move.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 2701–2710, 2017.
  * [39] Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros.  Context encoders: Feature learning by inpainting.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2536–2544, 2016.
  * [40] Chao Peng, Tete Xiao, Zeming Li, Yuning Jiang, Xiangyu Zhang, Kai Jia, Gang Yu, and Jian Sun.  Megdet: A large mini-batch object detector.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6181–6189, 2018.
  * [41] Florent Altché Corentin Tallec Florian Strub Andrew Brock Samuel Smith Soham De Razvan Pascanu Bilal Piot Michal Valko Pierre H. Richemond, Jean-Bastien Grill.  Byol works even without batch statistics.  2020\.
  * [42] Tim Salimans and Durk P Kingma.  Weight normalization: A simple reparameterization to accelerate training of deep neural networks.  In Advances in neural information processing systems, pages 901–909, 2016.
  * [43] Saurabh Singh and Abhinav Shrivastava.  Evalnorm: Estimating batch normalization statistics for evaluation.  In Proceedings of the IEEE International Conference on Computer Vision, pages 3633–3641, 2019.
  * [44] Kihyuk Sohn, David Berthelot, Chun-Liang Li, Zizhao Zhang, Nicholas Carlini, Ekin D Cubuk, Alex Kurakin, Han Zhang, and Colin Raffel.  Fixmatch: Simplifying semi-supervised learning with consistency and confidence.  arXiv preprint arXiv:2001.07685, 2020.
  * [45] Antti Tarvainen and Harri Valpola.  Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results.  In Advances in neural information processing systems, pages 1195–1204, 2017.
  * [46] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive multiview coding.  arXiv preprint arXiv:1906.05849, 2019.
  * [47] Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola.  What makes for good views for contrastive learning.  arXiv preprint arXiv:2005.10243, 2020.
  * [48] Zhi Tian, Chunhua Shen, Hao Chen, and Tong He.  Fcos: Fully convolutional one-stage object detection.  In Proceedings of the IEEE international conference on computer vision, pages 9627–9636, 2019.
  * [49] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky.  Instance normalization: The missing ingredient for fast stylization.  arXiv preprint arXiv:1607.08022, 2016.
  * [50] Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol.  Extracting and composing robust features with denoising autoencoders.  In Proceedings of the 25th international conference on Machine learning, pages 1096–1103, 2008.
  * [51] Xiaolong Wang and Abhinav Gupta.  Unsupervised learning of visual representations using videos.  In Proceedings of the IEEE international conference on computer vision, pages 2794–2802, 2015.
  * [52] Yuxin Wu and Kaiming He.  Group normalization.  In Proceedings of the European conference on computer vision (ECCV), pages 3–19, 2018.
  * [53] Yuxin Wu, Alexander Kirillov, Francisco Massa, Wan-Yen Lo, and Ross Girshick.  Detectron2.  https://github.com/facebookresearch/detectron2, 2019.
  * [54] Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin.  Unsupervised feature learning via non-parametric instance discrimination.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3733–3742, 2018.
  * [55] Qizhe Xie, Zihang Dai, Eduard Hovy, Minh-Thang Luong, and Quoc V Le.  Unsupervised data augmentation for consistency training.  arXiv preprint arXiv:1904.12848, 2019.
  * [56] Junjie Yan, Ruosi Wan, Xiangyu Zhang, Wei Zhang, Yichen Wei, and Jian Sun.  Towards stabilizing batch statistics in backward propagation of batch normalization.  arXiv preprint arXiv:2001.06838, 2020.
  * [57] Yang You, Igor Gitman, and Boris Ginsburg.  Large batch training of convolutional networks.  arXiv preprint arXiv:1708.03888, 2017.
  * [58] Xiaohua Zhai, Avital Oliver, Alexander Kolesnikov, and Lucas Beyer.  S4l: Self-supervised semi-supervised learning.  In Proceedings of the IEEE international conference on computer vision, pages 1476–1485, 2019.
  * [59] Xiaohua Zhai, Avital Oliver, Alexander Kolesnikov, and Lucas Beyer.  S4l: Self-supervised semi-supervised learning.  In Proceedings of the IEEE international conference on computer vision, pages 1476–1485, 2019.
  * [60] Richard Zhang, Phillip Isola, and Alexei A Efros.  Colorful image colorization.  In European conference on computer vision, pages 649–666. Springer, 2016.
  * [61] Richard Zhang, Phillip Isola, and Alexei A Efros.  Split-brain autoencoders: Unsupervised learning by cross-channel prediction.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1058–1067, 2017.
  * [62] Chengxu Zhuang, Alex Lin Zhai, and Daniel Yamins.  Local aggregation for unsupervised learning of visual embeddings.  In Proceedings of the IEEE International Conference on Computer Vision, pages 6002–6012, 2019.


# Efficient and Effective Vocabulary Expansion
Towards Multilingual Large Language Models

Seungduk Kim Seungtaek Choi11footnotemark: 1 Myeongho Jeong
Yanolja, South Korea
{seungduk.kim, seungtaek.choi<EMAIL_ADDRESS>
Equal Contribution.

###### Abstract

This report introduces EEVE-Korean-v1.0, a Korean adaptation of large language
models that exhibit remarkable capabilities across English and Korean text
understanding. Building on recent highly capable but English-centric LLMs,
such as SOLAR-10.7B and Phi-2, where non-English texts are inefficiently
processed with English-centric tokenizers, we present an efficient and
effective vocabulary expansion (EEVE) method, which encompasses parameter
freezing and subword initialization. In contrast to previous efforts that
believe new embeddings require trillions of training tokens, we show that our
method can significantly boost non-English proficiency within just 2 billion
tokens. Surpassing most instruction-tuned LLMs on the Open Ko-LLM Leaderboard,
as of January 2024, our model EEVE-Korean-10.8B-v1.0 ranks as the leading
Korean pre-trained model in the open-source community, according to Hugging
Face’s leaderboard. We open-source our models on Huggingface to empower the
open research community in various languages.

Efficient and Effective Vocabulary Expansion
Towards Multilingual Large Language Models

Seungduk Kim††thanks:  Equal Contribution. Seungtaek Choi11footnotemark: 1
Myeongho Jeong Yanolja, South Korea {seungduk.kim, seungtaek.choi,
<EMAIL_ADDRESS>

##  1 Introduction

Recent advancements in the field of large language models (LLMs), such as
GPT-4 OpenAI (2023), Gemini Team et al. (2023a), and Claude Anthropic (2023),
have demonstrated remarkable capabilities in processing and understanding
multiple languages. On the other hand, though notable models in open source
community, such as LLaMA Touvron et al. (2023a, b), MPT Team et al. (2023b),
Falcon Almazrouei et al. (2023), Mistral Jiang et al. (2023), Mixtral Jiang et
al. (2024), SOLAR Kim et al. (2023), and Phi-1.5 Li et al. (2023b) have set
benchmarks in English tasks, these developments have predominantly favored
English, leading to a performance gap in non-English languages.

Such disparity can be found not only in their language proficiency but also in
computational efficiency, where non-English languages like Korean require
significantly more tokens than English even for equivalent semantic content
(Figure 1). And, of course, this negatively affects the user experiences, such
as longer response times, shorter context lengths, and higher API costs Petrov
et al. (2023). Expanding the tokenizer vocabulary, which introduces some
frequently used yet long words as additional tokens, is thus indispensable for
non-English users, but vocabulary expansion is a very challenging task because
new embeddings require trillions of training tokens Zhao et al. (2024).

To this end, this technical report presents a novel approach for efficient and
effective vocabulary expansion, namely EEVE, which can better train the
embeddings of newly added tokens. For ease of adaptation, we utilize subword-
based embedding initialization and design seven training stages with parameter
freezing, which elaborately adjust the order and amount of parameters to be
trained. We meticulously transfer the advanced capabilities of foundational
models from English to Korean by initially focusing on the training of only
input embeddings and progressively expanding to encompass the full parameters
in the final stage.

Using EEVE, we officially release a family of Korean LLMs, EEVE-
Korean-10.8B-v1.0111https://huggingface.co/yanolja/EEVE-Korean-10.8B-v1.0 and
EEVE-Korean-2.8B-v1.0222https://huggingface.co/yanolja/EEVE-Korean-2.8B-v1.0,
which are built on recent English-centric LLMs, specifically SOLAR-10.7B Kim
et al. (2023) and Phi-2 Li et al. (2023b), with further Korean-centric pre-
training. We evaluate our models on lm-evaluation-
harness333https://github.com/EleutherAI/lm-evaluation-harness Gao et al.
(2023) for both English and Korean language tasks, such as boolean question
answering (BoolQ; Clark et al. 2019), commonsense causal reasoning (COPA;
Roemmele et al. 2011, context-sensitive word understanding (WiC; Pilehvar and
Camacho-Collados 2019), commonsense reasoning (HellaSwag; Zellers et al.
2019), and sentiment negation recognition (SentiNeg). From the evaluation, we
observe that our models outperform the recent open Korean pre-trained LLMs
like OPEN-SOLAR-KO-10.7B L. Junbum (2024), Polyglot-Ko Ko et al. (2023), and
KoGPT Kim et al. (2021), while preserving the strong English capability of the
base English-centric LLMs in terms of benchmark performance, being ranked as
the leading Korean pre-trained model in Open Ko-LLM Leaderboard Park et al.
(2023).

##  2 Efficient and Effective Vocabulary Expansion

To address the challenge of efficiently extending English-centric Language
Models (LLMs) to include non-English languages, we introduce a novel
methodology for vocabulary expansion. This method combines parameter freezing
with subword-based embedding initialization to effectively incorporate and
adapt to new linguistic tokens from languages beyond its initial training
scope, thereby enhancing its applicability across various linguistic contexts.
Our approach outlines a structured seven-stage training process, as
illustrated in Figure 1, meticulously designed to effectively integrate new
tokens into the model’s vocabulary. During pre-training, our objective is
causal language modeling.

Our core assumption is that foundational models, having been extensively
trained in English texts, possess a substantial level of understanding and
reasoning capabilities. Transferring these capabilities from English to
another language, such as Korean, could be more efficient than developing
performance from standalone Korean pre-training.

###  2.1 Preliminary 1: Tokenizer Training

---
English (8 tokens)
“Hello, the weather is nice today.”
[‘_Hello’, ‘,’, ‘_the’, ‘_weather’, ‘_is’, ‘_nice’, ‘_today’, ‘.’]
Korean (26 tokens)
“안녕하세요, 오늘은 날씨가 좋네요.”
[‘_’, ‘안’, ‘$<$0xEB$>$’, ‘$<$0x85$>$’, ‘$<$0x95$>$’, ‘하’, ‘세’, ‘요’, ‘,’, ‘_’,
‘오’, ‘$<$0xEB$>$’, ‘$<$0x8A$>$’, ‘$<$0x98$>$’, ‘은’, ‘_’, ‘날’, ‘$<$0xEC$>$’,
‘$<$0x94$>$’, ‘$<$0xA8$>$’, ‘가’, ‘_’, ‘좋’, ‘네’, ‘요’, ‘.’]
Expanded Tokenizer (9 tokens)
[‘_안’, ‘녕’, ‘하세요’, ‘,’, ‘_오늘은’, ‘_날씨가’, ‘_좋’, ‘네요’, ‘.’]

Table 1: A comparison of token consumption between English and Korean. We used
the tokenizers of SOLAR Kim et al. (2023) and our EEVE-Korean-10.8B-v1.0.

We trained a new tokenizer on our Korean corpus. Since our goal is to maximize
the leverage of the base model’s performance, we maintained the base model’s
vocabulary and added 8,960 tokens from the corpus that appeared at least 6,000
times, prioritizing those with the highest frequency. Ultimately, the
tokenizer’s vocabulary expanded to 40,960 tokens for EEVE-Korean-10.8B-v1.0.
This process entailed several rounds of tokenizer training and a manual
curation of tokens, based on an analysis of token frequency, ensuring a
comprehensive and relevant vocabulary for our model. As shown in Table 1, the
overall token consumption for Korean texts is significantly improved, almost
three-fold, contributing to the reduction of computational costs during the
entire training process.

(a) Stage 0

(b) Stage 1

(c) Stage 2

(d) Stage 3

(e) Stage 4

(f) Stage 5

(g) Stage 6

(h) Stage 7

Figure 1: Training stages with parameter freezing. The fire and snowflake
emojis indicate the trainable and frozen parameters respectively.

###  2.2 Preliminary 2: Subword-based Embeddings Initialization

The integration process starts before actual training, introducing new input
and output embeddings, called embed_tokens and lm_head, to the model’s
parameters. This preliminary step is crucial for preparing for the
sophisticated learning process that follows.

For the input embeddings of the newly added tokens, we adopt the approach of
using the average embeddings of the subword tokens that make up these new
tokens as in Hewitt (2021); Welch et al. (2020). This method utilizes the
semantic richness of the model’s existing subword embeddings to offer a
meaningful starting point for the new tokens’ representations.

Conversely, the output embeddings for the newly added tokens are initialized
with the embeddings of the first subword token that comprises the new token.
This strategy aims to align the new tokens’ output representations closely
with the semantic characteristics of their constituent subwords, enabling a
smoother integration into the model’s predictive framework. The significance
of such initialization will be further discussed.

###  2.3 Multi-stage Training

Here we describe the nuanced approach of our seven-stage training methodology
for efficient vocabulary expansion, emphasizing the meticulous process of
integrating new tokens derived from languages beyond the initial English-
centric training scope.

Figure 2: An illustrative example of showing how our subword-based embedding
initialization enables harmonize the old and new tokens at Stage 1. In Stage
1, the output embeddings of newly added tokens are initialized with the output
embeddings of their first subword tokens that make up these new tokens, such
that the last hidden representation for predicting “하세요” yields the same
logits for the newly added token “하세요” with its first subword token “하”. Even
if we give the new token “하세요” as a gold token, the gradients are eventually
computed based on its subword token “하”, so the model takes the input
embeddings of “하세요” to predict its subword “하”.

Stage 1 (new input embeddings): Initially, our focus is narrow yet critical:
to learn the input embeddings of the newly added tokens while freezing all
other model parameters. This stage is foundational, allowing the model to
adjust its recognition and processing of these tokens from the beginning. The
pre-initialized embeddings serve as a starting point, guiding the model to
better utilize these new tokens in its existing framework. Our principal
hypothesis here is that if the input and output token sequences in causal
language modeling can be differentiated, by utilizing both the old and new
tokenizers at the same time, the model can more efficiently and effectively
learn new vocabulary embeddings, as it could leverage its established
knowledge in the embedding spaces from old tokens. However, employing distinct
tokenizers for input and output sequences at once poses implementation
challenges, such as the difficulty of applying teacher forcing due to
mismatched input/output sequences. Here, the subword-based embedding
initialization (Sec 2.2) provides a proxy for using the old tokenizer for
output sequences, such that the model is tasked to generate the subword token
(old) given the whole word token (new). In other words, the model could learn
to align their representations for generating the new token with that for
generating its first subword token, by optimizing only the input embeddings
without any modification of input/output token sequences as described in
Figure 2. However, at this stage, the model is not yet able to distinguish
between tokens sharing the same hidden state.

Stage 2 (new output embeddings): Our goal is to enhance the model’s
proficiency in accurately generating new tokens across various contexts by
solely adjusting the output embeddings (lm_head). The decision to freeze all
other parameters stems from the model’s current unstable state. Allowing both
input and output embeddings to be trained simultaneously would complicate
achieving convergence, thus hindering the model’s progress toward optimal
performance. By freezing most of the parameters, we achieve more stable
convergence. Moreover, this approach significantly reduces the training time,
as it eliminates the necessity for backpropagation through the other layers.

Stage 3 (new input and output embeddings): At this stage, the input embeddings
(embed_tokens) still remain optimized based on the initial embeddings of the
output embeddings. This stage allows the updates of both input and output
embeddings of the newly added tokens simultaneously. By aligning between input
and output embeddings, the model learns to use the new tokens in both
understanding and prediction.

Stage 4 (all output embeddings): As all the original parameters of the base
model were frozen until this stage, we assumed the logits between old and new
tokenizers were differently scaled, or less optimized to be used as a whole
vocabulary. To this end, we begin to allow the update of the old parameters,
specifically the output embeddings of old tokens here, making the model better
generate the new tokens. In our preliminary experiments, we found this stage
is critical for improving the model’s generative capabilities.

Stage 5 (new input and all output embeddings): At this stage, the training
extends to fine-tuning all output embeddings across the model’s vocabulary
while continuing to refine the input embeddings for the newly added tokens.
The goal is to ensure that the model can accurately predict any token within
its expanded vocabulary. This phase emphasizes the integration of new tokens
within the broader context of the model’s linguistic understanding, ensuring
that they are both well-represented as inputs and accurately generated as
outputs. This dual focus aids in harmonizing the model’s overall performance,
ensuring that the expanded vocabulary is seamlessly woven into its language
generation processes.

Stage 6 (all layers): Contrary to being the final phase, this stage represents
an advanced step in the vocabulary expansion process, where all model
parameters are subject to optimization, including both newly introduced and
pre-existing ones. The focus here is on integrating the enhancements made to
the embedding layers within the model’s overall parameters. Techniques such as
QLoRA are utilized not just for efficiency but to ensure the preservation of
the model’s strong capabilities as much as possible, while allowing effective
integration of the expanded vocabulary.

Stage 7 (internal layers): Following the extensive integration and
optimization efforts, this stage serves as a “cool down” phase, focusing on
updating the model’s internal layers, which includes all the layers except the
input and output embedding layers. The objective is to ensure that the
enhancements made during the vocabulary expansion are deeply embedded within
the model’s core processing capabilities. This phase prepares the model for
robust performance, ensuring that it not only recognizes and generates the new
tokens but does so with a nuanced understanding of their use in varied
linguistic contexts.

##  3 Implementation Details

###  3.1 Datasets

For pre-training, we curated publicly available Korean corpora from diverse
sources, such as Korean web content, English vocabulary, and parallel corpus
in Korean AI Hub444https://aihub.or.kr/, etc. To construct a high-quality pre-
training corpus, we applied a set of preprocessing rules: 1) perplexity-based
filtering, 2) n-gram repetition Li et al. (2023a)-based filtering, and 3)
stopword-based filtering. For the efficient training of newly added Korean
tokens, we intentionally filtered out documents that do not contain many of
these tokens. Subsequently, we acquired a pre-training corpus totaling 3.2M
documents (or, 6.7GB).

Model |  | Total
---
(tokens)
| Average
---
(tokens)
SOLAR-10.7B | 3.1B | 964
EEVE-Korean-10.8B-v1.0 | 1.6B | 500
Phi-2 | 5.6B | 1748
EEVE-Korean-2.8B-v1.0 | 1.6B | 484

Table 2: Comparison of tokenizers for our 6.7GB pre-training corpus of a total
3.2M documents.

As can be seen in Table 2, for the entire corpus, the SOLAR tokenizer needed
to use 3.1B tokens to represent them, but our new tokenizer can do so with
almost half, using only 1.6B tokens. This difference becomes even more
pronounced in the case of Phi-2 and EEVE-Korean-2.8B models, where they
require 5.6B tokens and 1.6B tokens respectively. Considering that
transformers have a quadratic-increasing computation complexity with respect
to token length, this can be interpreted in two significant ways. First, it
allows for processing sequences more than 4 times longer on the same GPU. Or
second, it means our model can be trained nearly 4 times more computationally
efficiently on the same dataset. This difference becomes even more pronounced
in the case of the Phi-2 and EEVE-Korean-2.8B tokenizers.

For fine-tuning of EEVE-Korean models, we employed the Direct Preference
Optimization (DPO; Rafailov et al. 2023) based on LLaMA-Factory
implementation. To further enhance the models’ capabilities of following
Korean instructions, we translated the publicly available instruction
datasets, specifically Orca555https://huggingface.co/datasets/Open-
Orca/SlimOrca-Dedup Mukherjee et al. (2023); Lian et al. (2023) and
UltraFeedback666https://huggingface.co/datasets/argilla/ultrafeedback-
binarized-preferences-cleaned Cui et al. (2023) into Korean. In the process of
translating these datasets into Korean, ensuring the integrity of programming
code formats and correcting translation errors, such as instances where both
the source and target languages were inadvertently translated into Korean, was
crucial for maintaining the quality and effectiveness of our fine-tuned
models. We named the fine-tuned models as EEVE-Korean-Instruct.

###  3.2 Training

As foundational architectures, we opt for SOLAR-10.7B Kim et al. (2023) and
Phi-2 Li et al. (2023b), because both have shown outstanding performances
among similar sizes of LLMs. This choice of foundational architectures aligns
with our strategic training objectives, leveraging their proven strengths to
ensure our new models achieve similar levels of language understanding and
reasoning capabilities in Korean.

For the training of the model variants, we utilized two distinct codebases:
Axolotl777https://github.com/OpenAccess-AI-Collective/axolotl for the initial
pre-training phase and LLaMA-Factory888https://github.com/hiyouga/LLaMA-
Factory hiyouga (2023) for subsequent fine-tuning. These codebases provided a
strong and reliable base for our training process.

Specifically, we train our models with a setup of 8 x NVIDIA H100 GPUs with
80GB memory each, utilizing 64 CPU cores. For EEVE-Korean-10.8B-v1.0, under
bf16 precision, the training process is configured with a sequence of length
4096, gradient accumulation steps set to 4, and a micro-batch size of 8,
whereas EEVE-Korean-2.8B-v1.0 adopts a sequence length of 2048, gradient
accumulation of 16, and a micro-batch size of 16. We employ the AdamW
Loshchilov and Hutter (2018) optimizer, paired with a cosine learning rate
scheduler that includes a warmup phase of 10 steps. The learning rate for the
10.8B variant is set to 4e-5, while we used 2e-4 for the small model. We
continued training at each stage until the loss converged, observing the loss
converged before reaching 400 global steps, which signifies the efficiency of
our training strategy. Though our training strategy involves 7 different
stages, it is noteworthy that, for our 2.8B variant, the overall pre-training
can be done in less than two days as optimizing only the output embeddings
doesn’t incur much computation.

##  4 Evaluations

Model | Types | English | Korean | Avg.
---|---|---|---|---
BQ (0) |  CP (0) |  HS (0) |  BQ (0) |  CP (0) |  HS (0) |  SN (0) |  WIC (0)
meta-llama/Llama-2-7b-hf | PT | 0.7774 | 0.8700 | 0.5714 | 0.5242 | 0.5700 | 0.4420 | 0.4610 | 0.4881 | 0.5880
meta-llama/Llama-2-13b-hf | PT | 0.8055 | 0.9100 | 0.6006 | 0.5214 | 0.6010 | 0.4380 | 0.5038 | 0.4881 | 0.6086
mistralai/Mistral-7B-v0.1 | PT | 0.8379 | 0.9200 | 0.6129 | 0.6282 | 0.5880 | 0.4300 | 0.5365 | 0.4881 | 0.6302
meta-llama/Llama-2-7b-chat-hf | FT | 0.7976 | 0.8700 | 0.5779 | 0.5157 | 0.5530 | 0.4160 | 0.4987 | 0.4881 | 0.5896
meta-llama/Llama-2-13b-chat-hf | FT | 0.8165 | 0.8800 | 0.6072 | 0.5057 | 0.5760 | 0.4040 | 0.4685 | 0.4881 | 0.5933
upstage/SOLAR-10.7B-v1.0 (base) | PT | 0.8257 | 0.8700 | 0.6393 | 0.5057 | 0.5750 | 0.4320 | 0.6146 | 0.4881 | 0.6188
upstage/SOLAR-10.7B-Instruct-v1.0 | FT | 0.8853 | 0.9400 | 0.6866 | 0.8184 | 0.6370 | 0.4560 | 0.5668 | 0.4921 | 0.6853
beomi/OPEN-SOLAR-KO-10.7B∗ | PT | 0.8187 | 0.8800 | 0.5570 | 0.8355 | 0.8010 | 0.5040 | 0.6952 | 0.4897 | 0.6976
yanolja/EEVE-Korean-10.8B-v1.0∗ | PT | 0.8492 | 0.9000 | 0.6203 | 0.8568 | 0.7530 | 0.4900 | 0.6675 | 0.4992 | 0.7045
yanolja/EEVE-Korean-Instruct-10.8B-v1.0∗ | FT | 0.8810 | 0.9300 | 0.6502 | 0.8860 | 0.7610 | 0.4700 | 0.9521 | 0.4937 | 0.7530
microsoft/Phi-2 (base) | PT | 0.8336 | 0.9000 | 0.5583 | 0.5021 | 0.4770 | 0.3280 | 0.5063 | 0.4881 | 0.5742
daekeun-ml/phi-2-ko-v0.1∗ | PT | 0.6141 | 0.5800 | 0.3257 | 0.5164 | 0.6100 | 0.3860 | 0.4484 | 0.4881 | 0.4961
yanolja/EEVE-Korean-2.8B-v1.0∗ | PT | 0.7404 | 0.8900 | 0.5247 | 0.5299 | 0.5820 | 0.3800 | 0.5164 | 0.4881 | 0.5814
yanolja/EEVE-Korean-Instruct-2.8B-v1.0∗ | FT | 0.8248 | 0.8700 | 0.5392 | 0.7066 | 0.5640 | 0.3660 | 0.5290 | 0.5230 | 0.6153

Table 3: Main evaluation results based on lm-evaluation-harness. The dataset
names are abbreviated for brevity: BQ for BoolQ, CP for COPA, HS for
HellaSwag, and SN for SentiNeg. Korean tasks are from Jang et al. (2022).
Accuracy (acc) is used as the evaluation metric for all tasks. In the ‘Types’
column, the models are categorized into two groups: pre-trained (PT) and fine-
tuned (FT). We denote the models trained in Korean datasets with ∗. For ease
of reproduction, we adopt their official names at HuggingFace.

We evaluate our models on both Korean and English LLM benchmarks, to highlight
the advantages of our vocabulary expansion method, which could efficiently
leverage the strong multilingual capabilities of base foundational models.
Desirably, we expect a model to show improved performance in Korean tasks and
comparable performance in English tasks.

###  4.1 Benchmarks

For Korean tasks, we adopt the KoBEST benchmark Jang et al. (2022), whose
tasks are designed to evaluate the various aspects of language understanding
and reasoning. Specifically, this benchmark provides a Korean-translated
version of language understanding tasks: boolean question answering (BoolQ;
Clark et al. 2019), commonsense causal reasoning (COPA; Roemmele et al. 2011,
context-sensitive word understanding (WiC; Pilehvar and Camacho-Collados
2019), commonsense reasoning (HellaSwag; Zellers et al. 2019), and sentiment
negation recognition (SentiNeg). For English tasks, we employ the following
original tasks of KoBEST, BoolQ, COPA, and HellaSwag, which can better
highlight the alignment between the English and Korean capabilities of LLMs.
To ensure consistent comparisons, we employ an open-source LLM evaluation
framework, lm-evaluation-harness999https://github.com/EleutherAI/lm-
evaluation-harness Gao et al. (2023).

###  4.2 Results

We now present evaluation results for both our EEVE-Korean and EEVE-Korean-
Instruct variants with other top-performing models in Table 3. EEVE-
Korean-10.8B-v1.0 outperforms other pre-trained models of similar sizes in the
average performance. It is noteworthy that, EEVE-Korean is the only case where
the performance in Korean is improved without compromising the performance in
English. For example, though OPEN-SOLAR-KO-10.7B, which is built on the same
base model as ours, performs slightly better than our EEVE-Korean-
Instruct-10.8B-v1.0, it fails to preserve the English capabilities, showing
lower performance in English tasks than its base model, SOLAR-10.7B-v1.0. We
observe similar trends even for our smaller model, EEVE-Korean-2.8B-v1.0 in
comparison with the phi-2-ko-v0.1 model, sharing Phi-2 as its base model. This
demonstrates the effectiveness of our training strategy, especially
considering that we used even fewer training tokens than our competitors.

Notably, but not surprisingly, preference tuning on English datasets even
makes the models underperform in Korean tasks. For example, LLaMA-2-chat
variants, which are the preference-tuned version of LLaMA-2 checkpoints, show
improved performances in English tasks (Llama-2-7b 0.7774 $\rightarrow$
Llama-2-7b-chat 0.7976 in English BoolQ), while underperforming in Korean
tasks (Llama-2-7b 0.5242 $\rightarrow$ Llama-2-7b-chat 0.5157 in Korean
BoolQ), which highlights the importance of Korean-specific training for LLMs.
On the other hand, we observe that preference tuning our models on Korean
instruction datasets doesn’t hurt the model performance in English tasks,
rather even improving it. We posit that it is because the embedding spaces are
already well-aligned between Korean and English tokens, thus fine-tuning on a
specific language doesn’t incur a significant change in model parameters.

##  5 Conclusion & Future Work

The report introduces EEVE-Korean-v1.0, a Korean adaptation of large language
models that utilizes an Efficient and Effective Vocabulary Expansion (EEVE)
method to enhance Korean text processing capabilities significantly. The
method, based on parameter freezing and subword initialization, enables the
EEVE-Korean-10.8B-v1.0 model to excel in Korean language tasks while
maintaining strong English capabilities. Achieved with a corpus of just 2
billion tokens, this approach represents a notable advancement in language
model training efficiency and effectiveness. By making these models available
to the research community, the project aims to contribute to the development
of more inclusive and efficient language processing technologies.

Expanding our vision, future efforts will explore the application of our
vocabulary expansion methodology to additional languages, assessing its
generalizability and effectiveness. We aim to not only extend the EEVE-Korean
model’s linguistic range but also to delve deeper into evaluating its
reasoning and generative capabilities through diverse tasks, including complex
mathematical reasoning tests like GSM8K Cobbe et al. (2021), and human
evaluations in interactive settings like chatbots Zheng et al. (2023).
Moreover, efforts to enhance pre-training data quality, and to analyze
performance in code-switching scenarios Zhang et al. (2023) will underpin our
commitment to refining the model’s robustness and versatility. These
initiatives are designed to broaden the model’s applicability and efficacy,
pushing the boundaries of what is achievable with advanced language models.

## References

  * Almazrouei et al. (2023) Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxandra Cojocaru, Mérouane Debbah, Étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, et al. 2023.  The falcon series of open language models.  _arXiv preprint arXiv:2311.16867_.
  * Anthropic (2023) Anthropic. 2023.  Model card and evaluations for claude models.  _Anthropic technical Report_.
  * Clark et al. (2019) Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019.  Boolq: Exploring the surprising difficulty of natural yes/no questions.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 2924–2936.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021.  Training verifiers to solve math word problems.  _arXiv preprint arXiv:2110.14168_.
  * Cui et al. (2023) Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. 2023.  Ultrafeedback: Boosting language models with high-quality feedback.  _arXiv preprint arXiv:2310.01377_.
  * Gao et al. (2023) Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2023.  A framework for few-shot language model evaluation.
  * Hewitt (2021) John Hewitt. 2021.  Initializing new word embeddings for pretrained language models.
  * hiyouga (2023) hiyouga. 2023.  Llama factory.  https://github.com/hiyouga/LLaMA-Factory.
  * Jang et al. (2022) Myeongjun Jang, Dohyung Kim, Deuk Sin Kwon, and Eric Davis. 2022.  Kobest: Korean balanced evaluation of significant tasks.  In _Proceedings of the 29th International Conference on Computational Linguistics_ , pages 3697–3708.
  * Jiang et al. (2023) Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. 2023.  Mistral 7b.  _arXiv preprint arXiv:2310.06825_.
  * Jiang et al. (2024) Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024.  Mixtral of experts.  _arXiv preprint arXiv:2401.04088_.
  * Kim et al. (2023) Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, et al. 2023.  Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling.  _arXiv preprint arXiv:2312.15166_.
  * Kim et al. (2021) Ildoo Kim, Gunsoo Han, Jiyeon Ham, and Woonhyuk Baek. 2021.  Kogpt: Kakaobrain korean(hangul) generative pretrained transformer.
  * Ko et al. (2023) Hyunwoong Ko, Kichang Yang, Minho Ryu, Taekyoon Choi, Seungmu Yang, Sungho Park, et al. 2023.  A technical report for polyglot-ko: Open-source large-scale korean language models.  _arXiv preprint arXiv:2306.02254_.
  * L. Junbum (2024) L. Junbum. 2024.  Solar-ko-10.7b.
  * Li et al. (2023a) Huayang Li, Tian Lan, Zihao Fu, Deng Cai, Lemao Liu, Nigel Collier, Taro Watanabe, and Yixuan Su. 2023a.  Repetition in repetition out: Towards understanding neural text degeneration from the data perspective.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Li et al. (2023b) Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. 2023b.  Textbooks are all you need ii: phi-1.5 technical report.  _arXiv preprint arXiv:2309.05463_.
  * Lian et al. (2023) Wing Lian, Guan Wang, Bleys Goodson, Eugene Pentland, Austin Cook, Chanvichet Vong, "Teknium", and Nathan Hoos. 2023.  Slimorca dedup: A deduplicated subset of slimorca.
  * Loshchilov and Hutter (2018) Ilya Loshchilov and Frank Hutter. 2018.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_.
  * Mukherjee et al. (2023) Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar, Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah. 2023.  Orca: Progressive learning from complex explanation traces of gpt-4.  _arXiv preprint arXiv:2306.02707_.
  * OpenAI (2023) OpenAI. 2023.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_.
  * Park et al. (2023) Chanjun Park, Hwalsuk Lee, Hyunbyung Park, Hyeonwoo Kim, Sanghoon Kim, Seonghwan Cho, Sunghun Kim, and Sukyung Lee. 2023.  Open ko-llm leaderboard.
  * Petrov et al. (2023) Aleksandar Petrov, Emanuele La Malfa, Philip HS Torr, and Adel Bibi. 2023.  Language model tokenizers introduce unfairness between languages.  _arXiv preprint arXiv:2305.15425_.
  * Pilehvar and Camacho-Collados (2019) Mohammad Taher Pilehvar and Jose Camacho-Collados. 2019.  Wic: the word-in-context dataset for evaluating context-sensitive meaning representations.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 1267–1273.
  * Rafailov et al. (2023) Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. 2023.  Direct preference optimization: Your language model is secretly a reward model.  _arXiv preprint arXiv:2305.18290_.
  * Roemmele et al. (2011) Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. 2011.  Choice of plausible alternatives: An evaluation of commonsense causal reasoning.  In _2011 AAAI Spring Symposium Series_.
  * Team et al. (2023a) Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. 2023a.  Gemini: a family of highly capable multimodal models.  _arXiv preprint arXiv:2312.11805_.
  * Team et al. (2023b) MosaicML NLP Team et al. 2023b.  Introducing mpt-30b: Raising the bar for open-source foundation models.
  * Touvron et al. (2023a) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023a.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_.
  * Touvron et al. (2023b) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023b.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_.
  * Welch et al. (2020) Charles Welch, Rada Mihalcea, and Jonathan K Kummerfeld. 2020.  Improving low compute language modeling with in-domain embedding initialisation.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8625–8634.
  * Zellers et al. (2019) Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019.  Hellaswag: Can a machine really finish your sentence?  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 4791–4800.
  * Zhang et al. (2023) Ruochen Zhang, Samuel Cahyawijaya, Jan Christian Blaise Cruz, and Alham Fikri Aji. 2023.  Multilingual large language models are not (yet) code-switchers.  _arXiv preprint arXiv:2305.14235_.
  * Zhao et al. (2024) Jun Zhao, Zhihao Zhang, Qi Zhang, Tao Gui, and Xuanjing Huang. 2024.  Llama beyond english: An empirical study on language capability transfer.  _arXiv preprint arXiv:2401.01055_.
  * Zheng et al. (2023) Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. 2023.  Judging llm-as-a-judge with mt-bench and chatbot arena.  _arXiv preprint arXiv:2306.05685_.