# DSTEA: Dialogue State Tracking with Entity Adaptive Pre-training

Yukyung Lee  Takyoung Kim  Hoonsang Yoon  Pilsung Kang  Junseong Bang  Misuk
Kim

###### Abstract

Dialogue state tracking (DST) is a core sub-module of a dialogue system, which
aims to extract the appropriate belief state (domain-slot-value) from a system
and user utterances. Most previous studies have attempted to improve
performance by increasing the size of the pre-trained model or using
additional features such as graph relations. In this study, we propose
dialogue state tracking with entity adaptive pre-training (DSTEA), a system in
which key entities in a sentence are more intensively trained by the encoder
of the DST model. DSTEA extracts important entities from input dialogues in
four ways, and then applies selective knowledge masking to train the model
effectively. Although DSTEA conducts only pre-training without directly
infusing additional knowledge to the DST model, it achieved better performance
than the best-known benchmark models on MultiWOZ 2.0, 2.1, and 2.2. The
effectiveness of DSTEA was verified through various comparative experiments
with regard to the entity type and different adaptive settings.

Dialogue State Tracking, Task-oriented dialogue, ernie, adaptive pretraining

##  1 Introduction

The main purpose of dialogue state tracking (DST) is to track the structured
belief state (domain-slot-value) accurately within a dialogue. The performance
of a dialogue system depends heavily on that of its DST system. Moreover,
since the dialogue system mainly focused on the current dialogue turn,
generating the correct multi-turn belief state is a challenging task (Kim et
al., 2020).

To address this difficulty, recent studies have attempted to improve DST
performance by applying a pre-trained model (Lin et al., 2020; Zhao et al.,
2021) trained with a large number of unlabeled corpora, such as BART (Lewis et
al., 2020) and T5 (Raffel et al., 2020), or by using a graph structure to
inject additional information (Zhu et al., 2020; Wu et al., 2020).
Alternatively, performance improvements were sought by conducting additional
training. Mehri et al. (2020) using an open external dialogue dataset.
However, when a large pre-trained model is used, an increase in inference time
is inevitable, making this strategy difficult to use in a real-world dialogue
system. The same issue arises when additional training data are used, or when
a new graph structure is embedded.

To solve the limitations of previous DST models, the ideal strategy is to
apply effective training with only a given dataset, rather than accepting the
increased model complexity or training time caused by the use of additional
data (Guu et al., 2020). To this end, we noticed that the dialogue utterance
already contains key information, i.e., an entity that appears in one
sentence, such as identifying people, places, and dates. Therefore, we propose
DSTEA (Dialogue State Tracking with Entity Adaptive pre-training), a
methodology in which the representation of the DST model trains entities in
depth after extracting important information from a given utterance.

As shown in Figure 1, previous DST model used a BERT-based model trained by
masking tokens selected at random, without considering the characteristics of
such tokens. To improve upon this, we identified entities appearing in a
dialogue and proposed the use of selective knowledge masking (SKM), which is a
method for efficient pre-training by distinguishing the identified entities
from relatively less important factors. This strategy means that DSTEA learns
important knowledge at a higher masking rate because we assume that word and
phrase entities contain indispensable knowledge and they are worth being
trained more frequently than non-entity tokens. This can boost DST performance
without increasing the parameter size of the model. Through this approach, the
encoder, which plays a key role in the DST model, can learn the inductive bias
suitable for dialogue state tracking as part of the additional pre-training
process.

In this study, we conducted comparative experiments on various versions of the
representative DST dataset, MultiWOZ (Budzianowski et al., 2018), to evaluate
the performance of adaptive pre-training. Moreover, we verified DSTEA through
an additional qualitative analysis of the domain, slot-meta information, and
values.

Figure 1: Comparison between original and proposed masking strategies. We used
selective knowledge for effective pre-training on the DST. We give the masking
probability by distinguishing between knowledge selected as important
information($p_{select}$) or relatively insignificant information
($p_{origin}$).

##  2 Related Work

###  2.1 DST after BERT

BERT (Devlin et al., 2019) is a versatile backbone in the field of natural
language processing that achieves excellent performance in various downstream
tasks. Accordingly, there have been multiple attempts to make use of the
knowledge of pre-trained BERT in DST. SUMBT (Lee et al., 2019) is an ontology-
based model that encodes system/user utterances and slot-meta information with
separate BERT encoders, and then generates values through multi-head attention
and an RNN layer. In addition, in SOM-DST (Kim et al., 2020), instead of
generating the values for all dialogue turns, a method of selectively
generating values only for slots requiring updates through the addition of a
slot gate was proposed. DS-DST (Zhang et al., 2020a) takes advantage of both
an ontology and an open vocabulary approach and attempts a hybrid approach to
find a value in a candidate list or text span using two BERT encoders. CSFN-
DST (Zhu et al., 2020) and SST (Chen et al., 2020) improve the performance by
adding schema graph information, which is prior knowledge indicating the
relationship between the domain and slot. These models perform tasks depending
on the parametric knowledge of the pre-trained encoder with respect to a large
corpus of the general domain.

In this study, previously proposed DST model achieved excellent performance
through an improvement in the pre-trained encoder. We optimized the encoder
for the dialogue domain using the proposed method based on SOM-DST, following
the settings of Zhu et al. (2020).

Figure 2: Architecture of DSTEA.

###  2.2 Pre-training Methods

####  2.2.1 Adaptive pre-training

Adaptive pre-training refers to a process in which a language model trained in
the general domain is pre-trained to learn the knowledge suitable for a
specific domain or task. Previous approaches reported enhanced performance for
a specific domain or task by adding an adaptation phase between pre-training
and fine-tuning (Gururangan et al., 2020; Beltagy et al., 2019).

In the area of DST, some studies, such as DialoGLUE (Mehri et al., 2020) and
(Zhao et al., 2021), have attempted to inject adaptive knowledge into their
models. DialoGLUE performed additional pre-training to adapt the model to the
conversation domain using various open dialogue datasets. During adaptive pre-
training, DialoGLUE used the original masked language model (MLM). In the case
of Zhao et al. (2021), various pre-training objectives were applied to T5, a
pre-training model larger than BERT, and achieved good performance without any
pre/post-processing. However, the use of a large amount of external dialogue
data to improve the performance of DialoGLUE, or an encoder-decoder model that
contains more knowledge in Zhao et al. (2021), has an intrinsic limitation
imposed by the increased training or inference time.

Therefore, we focused on adaptive pre-training for entity recognition without
significantly increasing the time complexity. In this process, external
dialogue data are not used, and instead of large models such as BART or T5,
ERNIE (Sun et al., 2020) was used with the same parameter size as BERT to
maintain the training and inference time.

####  2.2.2 Masking strategy

Since BERT first learned general-purpose knowledge through masked language
modeling, Transformer-based (Vaswani et al., 2017) language models have
improved the training ability of such models through various masking
strategies. Although BERT applies masking based on a single token, in
subsequent studies, masking in units of tokens, words, and phrases was
proposed to improve the language modeling performance (Sun et al., 2020).

ERNIE is a language model that contains entity information. It defines
meaningful tokens, entities, and phrases as knowledge, and was the first model
to perform continual learning of high-level knowledge during the pre-training
process. SpanBERT (Joshi et al., 2020) learns a continuous random span.
Because the span boundary representation is learned without relying on
individual tokens, it is more advanced than the original BERT. REALM (Guu et
al., 2020) proposed salient-span masking to learn only the entity information
for question answering. Pre-training was conducted after defining named-entity
and date information as masking targets with a BERT-based tagger. Pegasus
(Zhang et al., 2020b) proposed a gap sentence generation method to predict
masked sentences based on the sentences that remain after masking in units of
sentences from the input document. Approximately 30% of the sentences in the
given documents could be masked.

In this study, we found that the knowledge definition of ERNIE and the entity
masking of REALM are well suited to the nature of dialogues. Therefore, based
on these techniques, we propose SKM to focus on the important entities.

##  3 Proposed Method

In this study, we propose DSTEA, which enhances pre-training for DST by
focusing on important knowledge. DSTEA applies entity adaptive pre-training
using system and user utterances to learn representations that are effective
for dialogue state tracking. Subsequently, dialogue state tracking is
performed using a well-trained encoder.

###  3.1 Overview

The overall architecture of the DSTEA is shown in Figure 2. The previous model
applied DST using a pre-trained language model to conduct masked language
modeling using a fixed random probability value without considering the
characteristics of the token; this is the original masking strategy shown in
Figure 1\- (a). On the other hand, DSTEA assumes that the entity information
appearing in the utterance is significant knowledge for DST, and applies
entity adaptive pre-training using a SKM strategy shown in Figure 1\- (b).

After applying this pre-training, DSTEA learns dialogue state tracking by
utilizing an operation-based model, called the SOM-DST, which achieves good
performance and consists of encoder and decoder structures, where the
operation is the value of the slot-meta comprising the domain and slot. An
operator can decide whether to create a value. Entity-adaptive-trained ERNIE
is used as the encoder of the model proposed in this study, and the decoder
uses a gated recurrent unit (GRU) network (Cho et al., 2014).

###  3.2 Entity Adaptive Pre-training

####  3.2.1 Why Entities?

The key to DST involves extracting information from the dialogue utterance
between the system and the user and assigning it to an appropriate domain and
slot. Most of the values that DST ultimately seeks pertain to an entity.
Therefore, DSTEA attempts to improve performance by enriching the meaning of
dialogue utterance through entity information.

####  3.2.2 Why ERNIE?

ERNIE is a framework for improving the performance of BERT, in which pre-
training is conducted by injecting entity information into BERT as prior
knowledge. The masking method in ERNIE differs from that of BERT in that it
can mask entities or phrases rather than words. As such, ERNIE is a language
model that learns using entity information, and is therefore suitable to the
direction pursued by DSTEA. Thus, we determined that it would perform well for
DST.

###  3.3 Entity Adaptive Pre-training Method

The encoder architecture is an essential part of the DST model. The purpose of
the encoder model in the proposed method is to learn an inductive bias
suitable for DST during the pre-training process so that the representation of
the pre-training model can be used to learn the dialogue information more
accurately. The previous DST model used an encoder that trained the MLM with a
random probability without considering the characteristics of the token. Using
it without an modifications results in limited performance improvement, which
we attempted to overcome by devising a pre-training method specialized for the
DST.

The adaptive pre-training method proposed for DSTEA is shown in Figure 3. Pre-
training comprises three steps: entity construction, selective knowledge
masking, and adaptive pre-training. After extracting an entity from the
utterance, a higher masking rate was assigned by defining the information as
important knowledge, and the original masking rate was assigned to the
remaining tokens. In this study, training was conducted by considering both
the word and phrase entities.

####  3.3.1 Entity Construction

One of the most important parts of entity adaptive pre-training is entity
construction. In this study, the entities were collected in four ways. First,
an entity was selected using ontology information. This is because an ontology
is the most readily available form of information from the dataset, and it
specifies the most important words. Second, after establishing a named entity
recognition (NER) model, inference was conducted on the MultiWOZ dataset to
collect the entities. An ERNIE-based BIO tagging model was used, and the model
was trained using the CONLL 2003 dataset (Tjong Kim Sang & De Meulder, 2003).
Third, entities were extracted using the spaCy (Honnibal & Montani, 2017)
library. The spaCy entity recognizer extracts entities in span units,
including entity types such as location, language, person, and product.
Finally, entities were extracted using the flair library (Akbik et al., 2019).
The flair entity tagger is a model that is trained based on CONLL 2003 and
extracts entities in span units.

The entities extracted through these methods are a mixture of words and
phrases. We attempted to learn these by distinguishing between word and phrase
entities. The ontology was used during the pre-training process, but not
during model training.

Table 1: JGA on the test sets of MultiWOZ 2.0, 2.1, and 2.2. * indicates a
model that was re-implemented. We also describe the pre-trained model used by
each methodology. Both BERT and ERNIE are the results of the experimentation
with the base model. (Parameter size : BERT - 110M, ERNIE - 110M, BART (base)
- 140M, T5 (base) - 220M)

Model | PLM | MultiWOZ 2.0 | MultiWOZ 2.1 | MultiWOZ 2.2
---|---|---|---|---
SUMBT* (Lee et al., 2019) | BERT | 42.07 | 46.99 | -
TRADE* (Wu et al., 2019) | - | 47.80 | 49.38 | 49.97
DSDST (Zhang et al., 2020a) | BERT | 52.24 | 51.21 | -
GCDST (Wu et al., 2020) | - | 50.68 | 46.09 | -
SOM-DST* (Kim et al., 2020) | BERT | 51.38 | 52.75 | 53.70
CSFN-DST (Zhu et al., 2020) | BERT | 51.57 | 52.88 | -
SOM-DST + Schema Graph (Zhu et al., 2020) | BERT | 52.23 | 53.19 | -
SST-1 (Chen et al., 2020) | BERT | 47.59 | 52.55 | -
MinTL (Lin et al., 2020) | T5 | 51.24 | 52.52 | -
MinTL (Lin et al., 2020) | BART | 52.07 | 53.62 | -
Pegasus (Zhao et al., 2021) | T5 | - | 54.40 | 56.60
DSTEA(ours) | ERNIE | 54.11 | 55.03 | 55.23

Table 2: JGA, SA and RSA scores by domain for MultiWOZ 2.0, 2.1, and 2.2. The
performance was compared with the re-implemented TRADE (Wu et al., 2019),
SUMBT (Lee et al., 2019), and SOM-DST (Kim et al., 2020) models.(* indicates
the re-implemented model.)

Domain | Model | MultiWOZ 2.0 | MultiWOZ 2.1 |  MultiWOZ 2.2
---|---|---|---|---
JGA | SA | RSA | JGA | SA | RSA | JGA | SA | RSA
Attraction | SUMBT* | 59.21 | 98.42 | 80.29 | 63.19 | 98.59 | 81.95 | - | - | -
TRADE* | 59.22 | 98.42 | 79.32 | 63.84 | 98.60 | 81.34 | 64.25 | 98.61 | 81.91
SOM-DST* | 65.80 | 98.68 | 83.21 | 68.11 | 98.78 | 84.21 | 68.02 | 98.78 | 84.48
DSTEA(Ours) | 69.21 | 98.86 | 85.31 | 70.65 | 98.92 | 86.17 | 69.54 | 98.86 | 85.39
Hotel | SUMBT* | 35.29 | 96.19 | 77.84 | 44.05 | 96.95 | 82.61 | - | - | -
TRADE* | 47.08 | 97.20 | 83.19 | 49.86 | 97.32 | 83.92 | 48.55 | 97.27 | 83.52
SOM-DST* | 46.45 | 97.11 | 83.92 | 49.22 | 97.31 | 84.93 | 46.33 | 97.04 | 83.55
DSTEA(Ours) | 48.76 | 97.26 | 84.69 | 48.61 | 97.33 | 84.93 | 49.23 | 97.37 | 85.24
Restaurant | SUMBT* | 57.95 | 98.03 | 85.10 | 58.14 | 98.16 | 86.83 | - | - | -
TRADE* | 61.59 | 98.38 | 88.34 | 58.59 | 98.24 | 87.01 | 59.25 | 98.24 | 86.93
SOM-DST* | 65.06 | 98.55 | 89.68 | 65.89 | 98.58 | 89.77 | 66.17 | 98.58 | 89.91
DSTEA(Ours) | 67.45 | 98.66 | 90.36 | 69.88 | 98.70 | 91.07 | 68.67 | 98.67 | 90.46
Taxi | SUMBT* | 25.08 | 95.51 | 52.95 | 31.96 | 95.88 | 56.18 | - | - | -
TRADE* | 36.60 | 96.64 | 64.27 | 33.18 | 96.59 | 63.16 | 36.29 | 96.65 | 64.02
SOM-DST* | 57.01 | 97.96 | 78.96 | 57.01 | 97.98 | 78.28 | 58.26 | 98.03 | 79.29
DSTEA(Ours) | 61.21 | 98.14 | 80.68 | 58.57 | 98.06 | 80.48 | 56.23 | 97.81 | 77.41
Train | SUMBT* | 57.68 | 97.85 | 84.99 | 68.46 | 98.59 | 90.27 | - | - | -
TRADE* | 66.90 | 98.52 | 89.88 | 67.31 | 98.56 | 90.38 | 68.40 | 98.62 | 90.72
SOM-DST* | 71.89 | 98.74 | 91.94 | 71.62 | 98.74 | 91.76 | 71.54 | 98.74 | 91.92
DSTEA(Ours) | 72.19 | 98.74 | 91.88 | 73.38 | 98.81 | 92.32 | 73.65 | 98.80 | 92.28

Word Entity First, each extracted entity was split into word units to compose
a word entity. To prevent overfitting during this process, information about
time and numbers was excluded from the entity. Because random times and
numbers are used when constructing dialogue datasets (Budzianowski et al.,
2018), the appearance of unseen information during a dialogue may prevent the
DST model from responding correctly to unseen slots or values. Therefore, if
we train entities regarding time and numbers with SKM, the biased model is
highly likely to generate incorrect values; therefore, we did not take these
values into account. Additionally, stopwords from the NLTK (Bird et al., 2009)
library were used to exclude words that were not helpful for training.
Moreover, filtering was conducted when the punctuation mark was extracted as
an entity.

Phrase Entity To learn a phrase entity, cases that included an unknown token
(i.e., [UNK]) in the phrase were excluded. Next, the phrase entities defined
for each utterance were extracted in advance, and then pre-training was
performed using randomly selected phrase entities. Even if a phrase entity
included stopwords, filtering was not conducted for masking in span units, and
information about numbers and times was excluded, as for word entities.
However, phrase entities extracted by the NER model and flair were of poor
quality; therefore, entities were extracted using only the ontology and spaCy.

Figure 3: Selective Knowledge Masking Procedure.

####  3.3.2 Selective Knowledge Masking Strategy

SKM is a method for learning important knowledge after selecting essential
information from user and system utterances, to inject an inductive bias
suitable for DST. The pre-training model was trained as an appropriate model
for DST by masking and predicting entities that can help track the dialogue
state, rather than arbitrary mask tokens. SKM is shown in step 2 of Figure 3.
The one-turn utterance token sequence is $U_{t}\ =\
(tok_{1},tok_{2},\ldots,tok_{L})$, where the defined word entity is
$Ent_{word}\ =\ (w_{1},w_{2},\ldots,w_{N})$, the phrase entity is
$Ent_{phrase}\ =\ (P_{1},P_{2},\ldots,P_{M})$, the total length of the token
sequence is $L$, the total number of word entities is $N$, and the total
number of phrase entities is $M$.

Tokens appearing in $U_{t}$ have a masking probability. In addition, a random
token is selected according to the corresponding probability value, and masked
language modeling is applied. The SKM proposed in this study first initializes
the entire token to $prob_{origin}$, and then selectively applies it when the
tokens of $Ent_{word}$, $Ent_{phrase}$, and the utterance we defined are
matched. After changing the masking probability of the token to
$prob_{select}$, masking is conducted in units of words and phrases. Through
this process, an effective representation can be learned by training
information essential to the dialogue domain. Unlike the existing method, this
masking method has the advantage of learning knowledge that is judged to be
important.

##  4 Experimental Results

###  4.1 Datasets

We conducted experiments using MultiWOZ 2.0, 2.1, and 2.2 (Budzianowski et
al., 2018; Eric et al., 2019; Zang et al., 2020), which are the most widely
used datasets for DST. Similarly, as in previous DST model, we conducted an
experiment using five domains (restaurant, train, hotel, taxi, and
attraction). Most of the preprocessing procedures followed the script provided
in TRADE (Wu et al., 2019).

###  4.2 Experimental Settings

We compared the performance of DSTEA with that of a number of previous models.
We re-implemented those models for which either performance on the MultiWOZ
datasets has not been published, or we were unable to reproduce their
performance with the published code. In the case of SUMBT in MultiWOZ 2.2, the
performance fluctuated greatly due to ontology update issues, so the
experiment of MultiWOZ 2.2 was not reported. DSTEA only improves encoder
models using dialogue utterance during adaptive training. Therefore, we
conducted a performance comparison with models that include an encoder and do
not use additional labels. For this reason, some well-known models were not
considered in our experiments. First, Trippy (Heck et al., 2020) did not only
use the belief state but also used an additional dialogue action label to
improve the DST performance. In addition, SimpleTOD (Hosseini-Asl et al.,
2020) is a model composed of only a decoder. Since DSTEA is a model that
performed adaptive training based on an encoder, including a model with only a
decoder in the benchmark models is not appropriate.

We trained the ’pre-trained ERNIE-2.0’ on dialogue and used the huggingface
transformers (Wolf et al., 2020), with ’nghuyong/ernie-2.0-en’ as the ERNIE
model. During the experiment, the masking probability $prob_{origin}$ was set
to 0.2, and $prob_{select}$ was set to 0.4. Detailed parameters for adaptive
training and DST training can be found in Appendix A.1.

###  4.3 Joint Goal Accuracy

We compared and analyzed the performance of DSTEA using joint goal accuracy
(JGA), an evaluation metric that verifies whether the predicted belief states
exactly match the gold label. Table 1 reports JGA values obtained with DSTEA
and the baseline models for MultiWOZ 2.0, 2.1, and 2.2.

DSTEA achieved better performance than previous BERT-based DST model and also
outperformed CSFN-DST and SOM-DST + Schema Graph, in which schema graph
information is injected into SOM-DST. Furthermore, DSTEA achieved higher
accuracy than MinTL using large pre-trained models such as T5 and BART.

In particular, through the performance difference between SOM-DST and DSTEA,
we demonstrated experimentally that the SKM strategy configures the encoder as
ERNIE in the pre-training process and masks information selectively. We also
confirmed that more important entities are influential. DSTEA demonstrated
this by outperforming the benchmark model for three datasets, with the only
exception being when pre-training was conducted for MultiWOZ 2.2 on T5 based
on the Pegasus pre-training objective. Although the number of parameters of
the model proposed in Zhao et al. (2021) is approximately twice that of our
model, DSTEA outperforms it for MultiWOZ 2.1. These results show that solving
the DST task using the large pre-trained model is not optimal for performance
enhancement and that training an inductive bias appropriate for the task helps
improve performance.

Table 3:  Performance comparison table for adaptive pre-training of MultiWOZ
2.1. The experiment result is a value derived from the random seed 42.

Model | MultiWOZ 2.1
---|---
JGA | SA
SOM-DST (BERT) | 52.75 | 97.35
SOM-DST (ERNIE) | 53.97 | 97.40
| DSTEA (ERNIE)
---
\+ Random Masking ($P_{origin}=0.2$)
52.70 | 97.27
| DSTEA (ERNIE)
---
\+ Random Masking ($P_{origin}=0.4$)
54.24 | 97.42
| DSTEA (ERNIE)
---
\+ Selective Masking ($P_{select}=0.4$, $P_{origin}=0.2$)
55.03 | 97.45

Table 4: Performance according to entity type of MultiWOZ 2.1. Comparison of
individual performance of word and phrase entities according to entity module.
The experiment result is a value derived from the random seed 42.

Model | Type | Entities | MultiWOZ 2.1
---|---|---|---
JGA
DSTEA | Word Level | Ontology Only | 54.54
Entity (spaCy) Only | 54.32
Entity (NER) Only | 54.42
Entity (flair) Only | 53.90
Combine words | 54.79
Phrase Level | Ontology Only | 54.47
Entity (spaCy) Only | 54.31
Combine phrase | 54.83
Combine All | Combine All | 55.03

###  4.4 Domain Results

For each dataset, domain-specific performance comparisons between the
benchmark models and DSTEA are presented in Table 2. In addition to JGA, the
performance was evaluated using slot accuracy (SA) and relative slot accuracy
(RSA) (Kim et al., 2022). SA is an evaluation metric that identifies the
accuracy of slots among the predicted dialogue states. RSA is a recently
proposed metric that complements JGA and SA. In contrast to SA, RSA only focus
on the gold reference or predicted slots of current dialogue, rather than all
pre-defined slots in SA, it has a better discriminative power for candidate
models. The corresponding result was recorded as the performance of the
experiment at the same random seed. The results show that DSTEA achieved the
best performance in all metrics except Hotel in MultiWOZ 2.1 and Taxi in
MultiWOZ 2.2.

###  4.5 Effectiveness of Adaptive Pre-training

Table 3 shows how the performance is affected by the adaptive pre-training
setting. To demonstrate the effectiveness of SKM, the performance of the pre-
trained model and SKM as proposed in this study were compared by changing the
masking probability of the original MLM. SOM-DST (ERNIE) indicates that the
encoder of SOM-DST is changed to ERNIE, and DSTEA (ERNIE) + Random Masking
($P_{origin}=\alpha$) indicates a case in which random masking probability
$\alpha$ is used during adaptive pre-training. The experiment involved
adjusting the masking probability without considering the characteristics of
the token. The results show that performance is improved by applying ERNIE, a
language model learned by focusing the encoder on the entity. Moreover, when
the masking probability was set to 0.2, performance was lower than that of
SOM-DST (ERNIE); however, DSTEA outperformed SOM-DST (ERNIE) with a masking
probability of 0.4. However, the best overall results were achieved with DSTEA
(ERNIE) + Selective Masking ($P_{select}=\beta$, $P_{origin}=\alpha$), which
is the proposed method with SKM. Consequently, the experiment confirms that
the proposed SKM method is more effective than simply increasing the masking
probability.

Figure 5 shows the performance of DSTEA according to the change of SKM rate.
The red dashed line represents the pre-training model that does not consider
entities, which can be understood as the lower bound performance. The blue
line indicates the performance of DSTEA according to each $P_{select}$ ratio.
SKM clearly enhanced the DST performance irrespective of $P_{select}$ ratio.
Although the best JGA was reported when $P_{select}=0.4$, the worst case still
yielded a significantly enhanced JGA compared to that without SKM.

Table 5: Sample outputs of domain, slot-meta, and value mismatch in MultiWOZ
2.1 (test set). Red indicates a mismatch.

SUMBT |  TRADE |  SOM-DST |  DSTEA
---|---|---|---
Domain mismatch (PMUL4648 - turn 2)
|  …
---
"hotel-name-nusha"
…
|  "attraction-name-nusha",
---
"restaurant-name-nandos"
"restaurant-name-nusha", |  "attraction-name-nusha",
Slot-meta mismatch (MUL0148 - turn 2)
|  …
---
"hotel-pricerange-expensive"
…
|  …
---
"hotel-pricerange-expensive"
"hotel-type-hotel"
…
|  …
---
"hotel-pricerange-expensive"
"hotel-type-hotel"
…
|  …
---
"hotel-pricerange-expensive"
…
Value mismatch (PMUL1106 - turn 2)
|  …
---
"train-departure-none",
…
|  …
---
"train-departure-stansted
airport",
…
|  …
---
"train-departure-stan"
…
|  …
---
"train-departure-cambridge"
…

Figure 4: Domain, slot-meta, value mismatch between DST model (SUMBT, TRADE,
SOM-DST, DSTEA) and ground truth in MultiWOZ 2.1 (Total turn 7368)

###  4.6 Effectiveness of Entity Types

We confirmed in the previous section that entities provide important
information regarding the DST task, and compared and analyzed the performance
of DSTEA for various entity types, as shown in Table 4. We extracted the
entities through four modules, and both word and phrase units were considered.
First, results show that an ontology is best for individual modules at the
word level, whereas the performance is poor when word entities are extracted
from flair only. However, when information of all entities was combined,
better performance was achieved than when using individual modules. Phrase
level entities also improve performance when training by combining the phrase
entities. Finally, when both words and phrases were mixed, the best JGA of
55.03 was achieved. This is key knowledge for both word and phrase entities
for DST, and it can be confirmed that configuring an appropriate entity set
significantly improves performance.

Figure 5: Performance on DSTEA of MultiWOZ 2.1 test set with various selective
knowledge masking rate.

###  4.7 Qualitative Result

We verified the superiority of our model through the qualitative results.
Finally, the DST model generates prediction values consisting of domain-slot
values for each turn. We compared the degree of mismatch ratio between the
benchmark models and DSTEA with respect to the domain, slot-meta, and value.
Figure 4 shows that DSTEA effectively reduces mismatching compared to the
benchmark models. This means that the domain and slot information of the
current utterance turn can be understood accurately, and a value suitable for
each slot can be generated. These results suggest that the model trained using
DSTEA can learn the appropriate inductive bias and improve its performance.

Table 5 shows a sample output for each model. The domain mismatch example
shows that DSTEA can accurately classify domains in a dataset for a multi-
domain dialogue. In addition, the slot-meta and value mismatch example
confirms that DSTEA can match an appropriate slot corresponding to the domain
and generate a relevant value.

##  5 Conclusion

In this study, we propose DSTEA, which applies entity adaptive pre-training by
focusing on entities and essential knowledge in dialogue state tracking. DSTEA
uses ERNIE, a model specialized for entities. In addition, SKM was proposed to
learn word and phrase entities effectively to focus on important information.
Through this, DSTEA outperforms the benchmark models by learning an inductive
bias suitable for dialogue state tracking during pre-training, and generating
a suitable representation. DSTEA achieved satisfactory performance for
MultiWOZ 2.0, 2.1, and 2.2 without using a large pre-trained model or directly
injecting additional information into the model. We expect that this pre-
training method can be used effectively for both dialogue state tracking under
insufficient input data and entity-related tasks.

## References

  * Akbik et al. (2019) Akbik, A., Bergmann, T., Blythe, D., Rasul, K., Schweter, S., and Vollgraf, R.  Flair: An easy-to-use framework for state-of-the-art nlp.  In _NAACL 2019, 2019 Annual Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations)_ , pp. 54–59, 2019.
  * Beltagy et al. (2019) Beltagy, I., Lo, K., and Cohan, A.  SciBERT: A pretrained language model for scientific text.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pp. 3615–3620, Hong Kong, China, November 2019. Association for Computational Linguistics.  doi: 10.18653/v1/D19-1371.
  * Bird et al. (2009) Bird, S., Klein, E., and Loper, E.  _Natural language processing with Python: analyzing text with the natural language toolkit_.  " O’Reilly Media, Inc.", 2009.
  * Budzianowski et al. (2018) Budzianowski, P., Wen, T.-H., Tseng, B.-H., Casanueva, I., Ultes, S., Ramadan, O., and Gašić, M.  MultiWOZ - a large-scale multi-domain Wizard-of-Oz dataset for task-oriented dialogue modelling.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pp. 5016–5026, Brussels, Belgium, October-November 2018. Association for Computational Linguistics.  doi: 10.18653/v1/D18-1547.
  * Chen et al. (2020) Chen, L., Lv, B., Wang, C., Zhu, S., Tan, B., and Yu, K.  Schema-guided multi-domain dialogue state tracking with graph attention neural networks.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 34(05):7521–7528, Apr. 2020.  doi: 10.1609/aaai.v34i05.6250.
  * Cho et al. (2014) Cho, K., van Merriënboer, B., Bahdanau, D., and Bengio, Y.  On the properties of neural machine translation: Encoder–decoder approaches.  In _Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation_ , pp. 103–111, Doha, Qatar, October 2014. Association for Computational Linguistics.  doi: 10.3115/v1/W14-4012.
  * Devlin et al. (2019) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  BERT: Pre-training of deep bidirectional transformers for language understanding.  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)_ , pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.
  * Eric et al. (2019) Eric, M., Goel, R., Paul, S., Sethi, A., Agarwal, S., Gao, S., and Hakkani-Tur, D.  Multiwoz 2.1: Multi-domain dialogue state corrections and state tracking baselines.  _arXiv preprint arXiv:1907.01669_ , 2019.
  * Gururangan et al. (2020) Gururangan, S., Marasović, A., Swayamdipta, S., Lo, K., Beltagy, I., Downey, D., and Smith, N. A.  Don’t stop pretraining: Adapt language models to domains and tasks.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 8342–8360, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.740.
  * Guu et al. (2020) Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M.-W.  Realm: Retrieval-augmented language model pre-training.  _ArXiv_ , abs/2002.08909, 2020.
  * Heck et al. (2020) Heck, M., van Niekerk, C., Lubis, N., Geishauser, C., Lin, H.-C., Moresi, M., and Gašić, M.  TripPy: A triple copy strategy for value independent neural dialog state tracking.  In _Proceedings of the 21st Annual Meeting of the Special Interest Group on Discourse and Dialogue_ , pp. 35–44, 1st virtual meeting, July 2020. Association for Computational Linguistics.
  * Honnibal & Montani (2017) Honnibal, M. and Montani, I.  spaCy 2: Natural language understanding with Bloom embeddings, convolutional neural networks and incremental parsing.  To appear, 2017.
  * Hosseini-Asl et al. (2020) Hosseini-Asl, E., McCann, B., Wu, C.-S., Yavuz, S., and Socher, R.  A simple language model for task-oriented dialogue.  _arXiv preprint arXiv:2005.00796_ , 2020.
  * Joshi et al. (2020) Joshi, M., Chen, D., Liu, Y., Weld, D. S., Zettlemoyer, L., and Levy, O.  SpanBERT: Improving pre-training by representing and predicting spans.  _Transactions of the Association for Computational Linguistics_ , 8:64–77, 2020.
  * Kim et al. (2020) Kim, S., Yang, S., Kim, G., and Lee, S.-W.  Efficient dialogue state tracking by selectively overwriting memory.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 567–582, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.53.
  * Kim et al. (2022) Kim, T., Yoon, H., Lee, Y., Kang, P., and Kim, M.  Mismatch between multi-turn dialogue and its evaluation metric in dialogue state tracking, 2022.
  * Lee et al. (2019) Lee, H., Lee, J., and Kim, T.-Y.  SUMBT: Slot-utterance matching for universal and scalable belief tracking.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pp. 5478–5483, Florence, Italy, July 2019. Association for Computational Linguistics.  doi: 10.18653/v1/P19-1546.
  * Lewis et al. (2020) Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O., Stoyanov, V., and Zettlemoyer, L.  BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 7871–7880, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.703.
  * Lin et al. (2020) Lin, Z., Madotto, A., Winata, G. I., and Fung, P.  MinTL: Minimalist transfer learning for task-oriented dialogue systems.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pp. 3391–3405, Online, November 2020\. Association for Computational Linguistics.  doi: 10.18653/v1/2020.emnlp-main.273.
  * Mehri et al. (2020) Mehri, S., Eric, M., and Hakkani-Tur, D.  Dialoglue: A natural language understanding benchmark for task-oriented dialogue.  _ArXiv_ , abs/2009.13570, 2020.
  * Raffel et al. (2020) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of Machine Learning Research_ , 21(140):1–67, 2020.
  * Sun et al. (2020) Sun, Y., Wang, S., Li, Y., Feng, S., Tian, H., Wu, H., and Wang, H.  Ernie 2.0: A continual pre-training framework for language understanding.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 34(05):8968–8975, Apr. 2020.  doi: 10.1609/aaai.v34i05.6428.
  * Tjong Kim Sang & De Meulder (2003) Tjong Kim Sang, E. F. and De Meulder, F.  Introduction to the CoNLL-2003 shared task: Language-independent named entity recognition.  In _Proceedings of the Seventh Conference on Natural Language Learning at HLT-NAACL 2003_ , pp. 142–147, 2003.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I.  Attention is all you need.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.
  * Wolf et al. (2020) Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Le Scao, T., Gugger, S., Drame, M., Lhoest, Q., and Rush, A.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pp. 38–45, Online, October 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.emnlp-demos.6.
  * Wu et al. (2019) Wu, C.-S., Madotto, A., Hosseini-Asl, E., Xiong, C., Socher, R., and Fung, P.  Transferable multi-domain state generator for task-oriented dialogue systems.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pp. 808–819, Florence, Italy, July 2019. Association for Computational Linguistics.  doi: 10.18653/v1/P19-1078.
  * Wu et al. (2020) Wu, P., Zou, B., Jiang, R., and Aw, A.  GCDST: A graph-based and copy-augmented multi-domain dialogue state tracking.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pp. 1063–1073, Online, November 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.findings-emnlp.95.
  * Zang et al. (2020) Zang, X., Rastogi, A., Sunkara, S., Gupta, R., Zhang, J., and Chen, J.  Multiwoz 2.2: A dialogue dataset with additional annotation corrections and state tracking baselines.  In _Proceedings of the 2nd Workshop on Natural Language Processing for Conversational AI, ACL 2020_ , pp. 109–117, 2020.
  * Zhang et al. (2020a) Zhang, J., Hashimoto, K., Wu, C.-S., Wang, Y., Yu, P., Socher, R., and Xiong, C.  Find or classify? dual strategy for slot-value predictions on multi-domain dialog state tracking.  In _Proceedings of the Ninth Joint Conference on Lexical and Computational Semantics_ , pp. 154–167, Barcelona, Spain (Online), December 2020a. Association for Computational Linguistics.
  * Zhang et al. (2020b) Zhang, J., Zhao, Y., Saleh, M., and Liu, P.  PEGASUS: Pre-training with extracted gap-sentences for abstractive summarization.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 11328–11339. PMLR, 13–18 Jul 2020b.
  * Zhao et al. (2021) Zhao, J., Mahdieh, M., Zhang, Y., Cao, Y., and Wu, Y.  Effective sequence-to-sequence dialogue state tracking.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pp. 7486–7493, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.emnlp-main.593.
  * Zhu et al. (2020) Zhu, S., Li, J., Chen, L., and Yu, K.  Efficient context and schema fusion networks for multi-domain dialogue state tracking.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pp. 766–781, Online, November 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.findings-emnlp.68.

##  Appendix A Appendix

###  A.1 Training details

Our models are built with the huggingface transformers. And we trained our
model with NVIDIA GeForce RTX 3090.

Table 6: Hyperparameters for entity adaptive pre-training

Hyperparameter | Assignment
---|---
number of epoch | 10
batch size | 16
learning rate | 6e-5
optimizer | AdamW
weight decay | 0.01
warmup ratio | 0.06
selective masking ratio | 0.4
original masking ratio | 0.2

Table 7: Hyperparameters for DSTEA

Hyper parameter | Assignment
---|---
number of epochs | 30
batch size | 64
encoder learning rate | 4e-5
decoder learning rate | 2e-4
drop out | 0.1
warm up ratio | 0.1
max sequence length | 256

###  A.2 Data statistics

Table 8: Data statistics on MultiWOZ 2.0, 2.1 and 2.2

domain |  MultiWOZ 2.0, 2.1, 2.2
---|---
train | dev | test
attraction | 2717 | 401 | 395
hotel | 3381 | 416 | 394
restaurant | 3813 | 438 | 437
taxi | 1654 | 207 | 195
train | 3103 | 484 | 494

###  A.3 Entity extraction details

We use spaCy, flair library to extract entities. The settings we used are
below.

Table 9: Settings for spaCy entity extraction Setting | Assignment
---|---
model | "en-core-web-sm"
version | 3.0.2
Table 10: Settings for flair entity extraction Setting | Assignment
---|---
model | "flair/ner-english-large"
version | 0.10


# SMILE: Scaling Mixture-of-Experts with Efficient Bi-level Routing

Chaoyang He1, Shuai Zheng2, Aston Zhang2, George Karypis2, Trishul Chilimbi2
Mahdi Soltanolkotabi1, Salman Avestimehr1
1University of Southern California
2AWS AI
Email<EMAIL_ADDRESS>

###### Abstract

The mixture of Expert (MoE) parallelism is a recent advancement that scales up
the model size with constant computational cost. MoE selects different sets of
parameters (i.e., experts) for each incoming token, resulting in a sparsely-
activated model. Despite several successful applications of MoE, its training
efficiency degrades significantly as the number of experts increases. The
routing stage in MoE relies on the efficiency of the All2All communication
collective, which suffers from network congestion and has poor scalability. To
mitigate these issues, we introduce SMILE, which exploits heterogeneous
network bandwidth and splits a single-step routing into bi-level routing. Our
experimental results show that the proposed method obtains a 2.5$\times$
speedup over Switch Transformer in terms of pretraining throughput on the
“Colossal Clean Crawled Corpus” without losing any convergence speed.

##  1 Introduction

Gigantic models have recently gained significant attention due to their
remarkable performance on natural language processing [1], computer vision
[2], and cross-modal learning [3]. However, the training of gigantic models
requires significant computational resources and data. As the model size is
scaled up, such large-scale training becomes computationally intensive and
environmentally unfriendly [4]. In perspective, the total emissions of net
tCO2e for GPT-3 are around 552 tCO2e [4], while a direct round trip of a
single passenger jet between San Francisco and New York emits about 180 tCO2e
[4]. Recent studies have started seeking alternative approaches to enable
greater computational efficiency.

Mixture of Experts (MoEs) [5] have emerged as the foundational neural network
to scale up model capacity using a massive number of parameters while
maintaining a constant computational cost by routing the input to a small
subset of experts with a router. While MoEs are promising in terms of model
performance and inference efficiency, they require careful design and tuning
of the router. A practical router should either enable a balanced workload for
experts to avoid downgrading the model performance or reduce the communication
overhead to guarantee that training finishes in a reasonable time [6, 7, 8].
For example, Switch Transformer [7] was introduced to train a model with 1.6
trillion parameters by simplifying MoE to route each token to only a single
expert and using an auxiliary loss to improve the workload balance during the
training.

Despite the success of the aforementioned literature in pre-training giant
models, they have two main disadvantages. First, they rely on an All2All
communication collective for both intra-node and inter-node data exchange,
which have heterogeneous bandwidths that differ by a large gap (e.g., in AWS
EC2 P4d, the peak bandwidth of EFA is 400Gbps, while the aggregated bandwidth
of NVSwitch inside a node is 600GB/s), hence downgrading the communication
efficiency. Second, computing a balanced router is increasingly more expensive
as more experts are used. In Switch Transformer, a large number of experts
cause congestion, generate network hotspots, and adversely affect performance.

Figure 1: Overview of SMILE $(N=m\times n)$. Bi-level routing divides token
dispatch into two stages, where in the first stage each token is routed to a
node via an inter-node router and then gets assigned to a GPU by an intra-node
router. This allows us to better utilize heterogeneous bandwidths to achieve
greater communication efficiency.

To tackle the above bottlenecks, we introduce bi-level routing to scale up MoE
with more efficient routing, dubbed as SMILE. The system overview is
illustrated in Figure 1. We divide the experts into two levels based on the
mesh topology. All the experts within a node are considered a group. Each
token is first routed to a node and then gets dispatched to a particular GPU
within the node. In this way, inter-node network congestion is dramatically
reduced. Moreover, the launch overhead on a node for All2All communication is
reduced from $\mathcal{O}(mn)$ to $\mathcal{O}(m+n)$, where $m$ and $n$ denote
the number of GPUs per node and the total number of nodes, respectively.
Similarly, the time complexity of routing is reduced from $\mathcal{O}(mnTd)$
to $\mathcal{O}(max(n,m)Td)$, where $T$ is the total number of tokens and $d$
is the model hidden size.

In the experiments, we demonstrate that the proposed SMILE improves throughput
over Switch Transformer by 2.5$\times$ for training a 3.7B parameter model on
128 GPUs while being able to maintain the same convergence speed. The
scalability analyses show that SMILE achieves significantly better weak and
strong scaling efficiencies than Switch Transformer, and maintains a good
performance advantage when we increase the model size from 3.7B to 48B. The
profiling of a single MoE layer confirms our motivation that All2All
communication is the major bottleneck in MoE and the proposed bi-level routing
significantly mitigates the communication overhead.

##  2 Related Works

Mixture of Experts (MoE), in the context of modern deep learning
architectures, was proven effective in [5]. In this work, a MoE layer was
applied between stacked LSTM [9] layers, and tokens were separately routed to
different subsets of experts. In our work, we consider a hybrid of data +
expert parallelism, where each worker holds a single expert for each MoE
layer, and the number of experts scales with the number of workers, i.e.,
$N=nm$ with $m$ and $n$ denoting the number of GPUs per node and total number
of nodes. For the token assignment, the router is equipped with a variable
$W_{r}\in\mathcal{R}^{N\times d}$ and produces logits $r(x)=W_{r}x$, where
$x\in\mathcal{R}^{d}$ is the token hidden vector. The logits are normalized
via a softmax to construct the probability of selecting each expert. The
routing probability for expert $e$ is given by

$\displaystyle p_{e}(x)=\frac{r_{e}(x)}{\sum_{i=1}^{N}r_{i}(x)}.$ (1)

The top-k experts are then selected for processing the given token. Denote the
set of chosen experts by $\mathcal{I}$. The output of the top-k experts is
given by

$\displaystyle y(x)=\sum_{e\in\mathcal{I}}p_{e}(x)E_{e}(x),$ (2)

where $E_{e}(\cdot)$ is the sub-model (e.g., multi-layer perceptron) for
expert $e$. In this way, the number of model parameters increases linearly
with respect to the number of experts only with a small amount of extra
computational cost for routing. MoE offers state-of-the-art performance in
language modeling and machine translation benchmarks. The routing step has a
total complexity of $\mathcal{O}(kmnTd)$, where $T$ is the total number of
tokens. The MoE layer was reintroduced into the Transformer architecture by
the Mesh Tensorflow library [10] where MoE layers were introduced as a
substitute for the Feed-forward Network (FFN) layers in Transformers [11].
However, there were no accompanying NLP results. With recent advances in
machine learning infrastructure, GShard [6], which extended the XLA compiler,
uses MoE to dramatically improve machine translation across 100 languages. In
[12], a different deterministic MoE strategy is adopted to split the model
parameters into non-overlapping groups of languages. Switch Transformer [7]
simplifies the routing process and only selects a single expert for each
token. BASE [8] is another MoE variant that stacks multiple FFN layers as a
single expert and inserts them into a standard Transformer architecture. This
significantly increases the inference time compared with a vanilla
Transformer.

Our proposed method introduces bi-level routing to better leverage
heterogeneous communication bandwidth, and utilize two additive losses for
load balancing, achieving a large speedup. A concurrent work [13]
independently considers intra-node all-to-all first followed by an inter-node
all-to-all. We would like to emphasize that this work does not change routing
mechanisms and only exploits hierarchical all-to-all for inference.

##  3 Method

###  3.1 Background: Bottleneck of MoE in Scaling to a Large Number of GPUs

for (int r = 0; r < numranks; r++) if (count != 0) ncclSend(sendbuff +
r*rankdiff, count, type, r, comm, stream);

ncclRecv(recvbuff + r*rankdiff, count, type, r, comm, stream);

Figure 2: The implementation of All2All Communication in NCCL
(torch/csrc/cuda/nccl.cpp)

Figure 3: Throughput results when scaling Switch Transformer to a large number
of GPUs

MoE heavily relies on the performance of All2All. Depending on the network
topology, different All2All algorithms result in different communication
costs, latency, and network congestion behaviors. Suppose that there are $N$
workers in total. For a ring topology, All2All has a quadratic communication
cost and linear latency on $N$ while communication cost and latency are
reduced to $\mathcal{O}(N^{3/2})$ and $\mathcal{O}(N^{1/2})$, respectively for
mesh topology such as TPUs [14]. Regardless of the underlying topology of the
network, another trivial approach is to send all the messages asynchronously
into the network as illustrated in Figure 3.1. This naive algorithm implements
pairwise one-to-one routings
111https://github.com/pytorch/pytorch/blob/2b7943c47c8561a46103488b0fe9a592b87dc5bb/torch/csrc/cuda/nccl.cpp#L637
and suffers from network congestion because of the bisection width of the
network [15]. As in Figure 3, Switch Transformer has very poor scaling
efficiency when the number of nodes is increased from 1 to 16 (8 to 128 GPUs).
The throughput on 8 nodes are even worse than that on 4 nodes.

###  3.2 Efficient MoE Layer via Bi-level Routing

To tackle the above bottlenecks, we introduce bi-level routing to scale up MoE
with more efficient routing, dubbed as SMILE.

####  3.2.1 Model Architecture: Orchestration of Inter-node All2All and Intra-
node All2All

Figure 4: Illustration of SMILE Layer (m = 8, n = 8)

To account for the heterogeneous and hierarchical nature of the inter-
connection network, we divide the experts into two-level All2All operation.
All the experts within a node are grouped together. As shown in Figure 4, when
a token is ready for the dispatch, it is first routed to a node via an inter-
node router (blue) and is then assigned to a GPU via an intra-node router
(orange).

The proposed bi-level routing reduces the launch time on a node for All2All
from $\mathcal{O}(mn)$ to $\mathcal{O}(m+n)$, where $m$ and $n$ are the number
of GPUs per node and the total number of nodes, respectively. In terms of the
communication efficiency, bi-level routing parallelizes multiple All2All
collectives, which minimizes network interference between flows and
significantly reduces inter-node network congestion compared to the naive
algorithm implemented in NCCL (Ref. Figure 3.1).

Bi-level routing also simplifies router optimization where the size of routing
problem in Switch Transformer (Equation 1) is decreased from $mn$ to
$max(n,m)$. In particular, the complexity of routing in Switch Transformer is
reduced from $\mathcal{O}(mnTd)$ to $\mathcal{O}(\max(n,m)Td)$.

The dash curve in Figure 4 also illustrates the way we calculate the output of
the top-1 expert:

$\displaystyle h_{out}=p_{i}(h_{in})q_{j}(h_{in})E_{i,j}(h_{in}),$ (3)

where $p_{i}(h_{in})$ and $q_{j}(h_{in})$ are the top-1 routing probability
(generated by Equation 1) allocated for node $i$ and local expert $j$,
respectively, and $E_{i,j}$ is the $j$-th expert on node $i$. Both the inter-
node and intra-node routers have tied parameters across all the workers,
ensuring that the results do not change when an incoming example is processed
by a different worker. Compared to a single router, bi-level routers also
reduce the total number of router parameters from $\mathcal{O}(mn)$ to
$\mathcal{O}(m+n)$.

####  3.2.2 Additive Load Balancing Loss

Different from the one-hop load balancing loss in previous works [5, 10, 6,
7], we use an additive load balancing (LB) loss for bi-level routing. Given
$N$ experts indexed by $i=1$ to $N$, we decouple it into $n\times m$, where
$n$ is the number of GPU nodes and $m$ is the number of GPUs on a single node
(e.g., 8 GPUs is a common configuration for a GPU node). For a batch
$\mathcal{B}$ with $T$ tokens, the additive load balancing loss has two
components, either of which is computed as the scaled dot-product between
dispatch fraction and routing probability vectors:

$\text{loss}_{\text{lb}}=\underbrace{\alpha\cdot
n\cdot\sum_{i=1}^{n}f_{i}\cdot P_{i}}_{\text{inter-node LB
loss}}+\underbrace{\beta\cdot m\cdot\sum_{j=1}^{m}f_{j}\cdot
Q_{j}}_{\text{intra-node LB loss}},$ (4)

where $f_{i}$ and $f_{j}$ are the fraction of tokens dispatched to node $i$
and local expert $j$, respectively:
$f_{i}=\frac{1}{T}\sum_{x\in\mathcal{B}}1\\{\texttt{argmax}p(x)=i\\}$,
$f_{j}=\frac{1}{T}\sum_{x\in\mathcal{B}}1\\{\texttt{argmax}q(x)=j\\}$; $P_{i}$
and $P_{j}$ are the fraction of the router probability allocated for node $i$
and expert $j$ respectively;
$P_{i}=\frac{1}{T}\sum_{x\in\mathcal{B}}p_{i}(x)$,
$Q_{j}=\frac{1}{T}\sum_{x\in\mathcal{B}}q_{j}(x)$; hyper-parameter $\alpha$
and $\beta$ are multiplicative coefficients. The minimum is attained under
uniform inter-node and intra-node routing, i.e.,
$\min\text{loss}_{\text{lb}}=\alpha\cdot n\cdot\sum_{i=1}^{n}1/n\cdot
1/n+\beta\cdot m\cdot\sum_{j=1}^{m}1/m\cdot 1/m=\alpha+\beta$. In practice, we
simply use $\alpha=\beta$.

To compute the total model loss during training, we sum up
$\text{loss}_{\text{lb}}$ in all SMILE layers as the auxiliary loss:

$\text{loss}_{\text{total}}=\text{loss}_{\text{train}}+\sum_{l=1}^{L}{\text{loss}_{\text{lb}}^{l}},$
(5)

where $L$ denotes the total number of SMILE layers and
$\text{loss}_{\text{lb}}^{l}$ is the load balancing loss in $l$-th SMILE
layer.

####  3.2.3 Bi-level Process Group Management

Figure 5: Bi-level Process Group Management and its Pseudocode

The system implementation of the SMILE layer requires two levels of
distributed process management. The first level process group handles node-
level All2All, and the second level manages All2All between intra-node GPU
processes. In addition, these two groups of processes should be connected to
complete the bi-level routing without mutual interference. The left side of
Figure 5 shows the process of cooperation between the first-level inter-node
process group and the second-level intra-node process group to complete bi-
level routing. Based on such requirements, we propose a process management
mechanism based on PyTorch dist.new_group API. As shown on the right, for each
GPU process, we create an inter-node process group and an intra-node process
group, where the process ranks in the group are shown in blue text and orange
text, respectively. Based on this process group management, when performing
the All2All operation for the BiMoE Layer, we only need to specify the
inter_node_process_group instance and intra_node_process_group instance
according to local rank. This method greatly simplifies the management of the
process group so that the MoE layer itself does not need to care about the
system implementation details. The right side of the figure also shows the
process of four sequential All2All operations. Two additional All2All
operations are required because of the reversed routing for the consecutive
attention layer.

##  4 Experiments

###  4.1 Experimental Setup

##### Task and Dataset.

We evaluate SMILE on NLP pre-training tasks with large Transformer models. We
use a masked language modeling task [16, 17, 18] where the model is trained to
predict missing tokens. We evaluate the performance of SMILE by pre-training
on the “Colossal Clean Crawled Corpus” (C4), a collection of English-language
text sourced from the public Common Crawl web scrape. It includes heuristics
to extract only natural language (as opposed to boilerplate and other
gibberish) in addition to extensive deduplication [19]. The C4 dataset is
obtained from the curated version hosted by Hugging Face Dataset
222https://huggingface.co/datasets/c4. It has 129 billion tokens (words) in
the training dataset and 129 million tokens (words) in the validation dataset.
For the parallel training on large number of GPUs, we split the training
dataset into 32768 (1024 x 24) files, and validation dataset into 256 files.
We use the same vocabulary as the original T5 (11B) model
333https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/tokenization_t5.py
(vocabulary size is 32128).

##### Model Architecture.

We compare SMILE to Switch Transformer to demonstrate the efficiency of bi-
level routing. The proposed method can also be used in conjunction with other
MoE models such as GShard [6] and BASE [8]. For fair comparison, Switch
Transformer and SMILE use the same BERT-like architectures (a stack of many
standard Transformer layers) but replaces the every other feed forward network
(FFN) layer in Transformer with a MoE (mixture of experts) layer. In each
Transformer layer, the MoE layer follows after a multi-head attention layer,
and they are enhanced with a skip connection followed by a LayerNorm operation
afterwards. The activation function in attention and FFN layer is set to GELU
with a dropout rate 0.1. In SMILE the routing architecture is defined
according to the design in Section 3.2.1. For different model sizes, we only
change the number of hidden layers, hidden size, and intermediate size
(details are introduced in Section A.2).

##### Hardware.

We run experiments on state-of-the-arts hardware in AWS with advanced supports
for computing, communication, storage: 1\. GPU accelerators: we evaluate all
baselines and SMILE on AWS P4d nodes 444https://aws.amazon.com/ec2/instance-
types/p4/. Each node is equipped with 8 NVIDIA A100 GPUs. We scale up to 16
nodes to evaluate the scalability; 2\. High bandwidth Communicator: We utilize
AWS EFA (Elastic Fabric Adapter) 555https://aws.amazon.com/hpc/efa/ for 400
Gbps high bandwidth inter-node networking. Compared to commonly used NVIDIA
InfiniBand, EFA’s custom-built operating system (OS) bypass hardware interface
enhances the performance of inter-instance communications. Our experiments
show that even in such an high bandwidth setting, the All2All communication is
still a bottleneck in MoE models (e.g., Switch Transformer); 3\. Networked
File System: To boost the performance of accessing the data files in a
distributed manner, we use AWS FSx 666https://aws.amazon.com/fsx/ with SSD
support. The total storage cost for C4 dataset and all source code is around
800G.

##### Training Hyper-parameters.

We train MoE models with the LAMB optimizer [20], where the learning rates are
tuned in the range $\\{0.0001,0.0003,0.001,0.003\\}$, the weight decay is
fixed to 0.01, and $\epsilon$ is set to 1e-6. We clip gradients if their l2
norm exceeds 1.0. As a common practice to reduce the GPU memory cost in LAMB
optimizer, we also enable half precision (fp16). We use a sequence length of
128. Unless otherwise specified, we fix overall training batch size to 16384
and micro batch size to 128, where micro batch size refers to the batch size
per GPU per micro step and total_batch_size = micro_batch_size *
num_micro_steps. Gradient accumulation is adopted when the number of micro
steps is larger than 1. We use 128 because it is the maximum size that can be
used under GPU memory constraints with our hardware configuration. For
scalability analysis, we scale the number of nodes from 1 node (8 GPUs) to 16
nodes (128 GPUs).

##### Implementation.

Our source code is well-maintained as a Python pip package. We implement our
algorithm with the integration of the PyTorch DDP and DeepSpeed frameworks.
The process group management introduced in is handled by PyTorch DDP [21]
grouping APIs. We use LAMB optimizer implemented by DeepSpeed [22]. For GPU
memory-efficient training of large dense model, we reuse a few techniques
supported by DeepSpeed, including ZERO optimization [23], activation
checkpointing, and half precision (fp16). To analyze the fine-grained time
breakdown for communication and computation in MoE layer, we use PyTorch
Profiler. Our data loader for C4 dataset is customized with the pre-fetching
mechanism for efficient distributed loading.

###  4.2 Comparison with BERT and Switch Transformer

Figure 6: The curve of iteration-to-perplexity
Figure 7: Unscaled load balancing loss.

##### Accurate and Faster Training.

In addition to Switch Transformer, we compare SMILE with BERT (110M) and BERT
(3.7B) baselines which have the same model floating point operations (FLOPs)
and number of parameters as Switch Transformer and SMILE, respectively. We use
$\alpha=0.01$ for Switch Transformer and $\alpha=\beta=0.005$ (introduced in
Equation (4)) for SMILE in our experiments, and set the capacity factor for
routing as 2.0. We replace every other shared feed-forward layer in the
Transformer architecture with a MoE (Resp. SMILE) layer.

Table 1: Throughput (samples/second) Model | Throughput
---|---
BERT (110M) | 93282
BERT (3.7B) | 5114
Switch Transformer | 8112
SMILE | 20011

From Figures 6, 7 and Table 1, we have four important observations. First,
SMILE has the same convergence behaviour as Switch Transformer, and it
converges faster than BERT (110M). Second, both SMILE and Switch Transformer
converge slower than BERT (3.7B), which is expected since the MoE models
trade-off convergence for greater computational efficiency. Third, both Switch
Transformer and SMILE are slower than BERT (110M), indicating that routing is
the major bottleneck in the MoE models. And, SMILE runs 2.5x and 3.9x faster
than Switch Transformer and BERT (3.7B), respectively. This proves that bi-
level routing is effective in reducing the overhead of standard MoE layers.
Lastly, SMILE achieves the twice unscaled balancing loss of that of Switch
Transformer, which is expected since SMILE has two additive losses. When
scaled with $\alpha$, two curves will roughly overlap with each other.

Next, we provide fine-grained performance analysis and ablation studies to
justify the necessity of bi-level routing.

###  4.3 Scalability

####  4.3.1 Scalability on High Bandwidth (400 Gbps) Inter-node Communication

(a)  Weak Scaling (b) Strong Scaling

Figure 8: Switch Transformer vs. SMILE. The number of nodes is increased from
1 to 16.

We compared the throughput (samples per second) between SMILE and Switch
Transformers when scaling the number of GPU nodes (each node has 8 GPUs) from
1 to 16 in high bandwidth. Both weak scaling and strong scaling are evaluated.
In weak scaling, the global batch size is adjusted with the number of GPUs,
while in strong scaling, both the global batch size and the micro batch size
per GPU are fixed (the number of gradient accumulation steps decreases when
the node number is scaling up). From the results in Figure 8, we have the
following observations.

1\. MoE overhead in Switch Transformer (All2All Communication) is non-trivial
even in a large bandwidth supported by advanced communication adaptor (AWS
EFA). Its scaling efficiency is far below the linear scaling, which can be
explained by the additional inter-node communication cost that drags down the
performance; what is even worse is that the final throughput on 16 nodes is
not notably better than that on a single node and 8 nodes has worse throughput
than 4 nodes.

2\. Compared to Switch Transformer, SMILE scales up much better from 1 node to
16 nodes. The throughputs on 16 nodes are 7.7x and 4x higher than those on 1
node with weak and strong scaling, respectively. Moreover, different from
Switch Transformer, when scaling from 4 nodes to 8 nodes, the throughput still
increases. We observe worse performance of SMILE on 1 node with weak scaling,
which is due to additional overhead in the implementation. On a single node,
we should directly use Switch Transformer.

Therefore, we conclude that bi-level routing is efficient in inter-node MoE
scaling, and the scalability is largely improved when SMILE is applied.

####  4.3.2 Scalability w.r.t. Different Model Sizes

Table 2: Comparison of Throughput between Switch Transformer and SMILE (16 P4d
nodes). We fix the total batch size to 16384 and vary the micro batch size
depending on the model size and GPU memory.

Model Size (128 Experts) | Model Configuration | Throughput (samples/second)
---|---|---
Switch Transformer | SMILE
3.7B |  micro_batch_size = 128 num_layers = 12 hidden_size = 768 intermediate_size = 3072 | 8112 | 20011 ($2.47\times\uparrow$)
13B |  micro_batch_size = 64 num_layers = 24 hidden_size = 1024 intermediate_size = 4096 | 4001 | 6829 ($1.71\times\uparrow$)
48B |  micro_batch_size = 64 num_layers = 36 hidden_size = 1600 intermediate_size = 6400 | 889 | 2223 ($2.50\times\uparrow$)

To understand the performance in different model sizes, we evaluate on three
model configurations as shown in Table 2. We conduct this experiment on 128
GPUs, and the number of experts is fixed to 128. The first two use BERT_base
and BERT_large as the base dense models, while the third one is constructed by
increasing the hidden size and model depth. This result demonstrates that
SMILE is not restricted to a specific model architecture, and still achieves
1.7-2.5 times faster training speed when the model size increases
significantly.

###  4.4 Understanding SMILE: Time Breakdown and Performance Analysis

To demystify the scalability benefit of SMILE, we also did fine-grained
performance analysis. It would be difficult and inaccurate by directly
dissecting the performance of a single MoE layer from an end-to-end training
pipeline, since there are other factors involved by the interaction between
data parallelism (AllReduce) and MoE layer (All2All). As such, we develop a
tiny model with only a single MoE layer and perform training with dummy data
on the same GPU cluster with AWS EFA (16 P4d nodes)777Our source code also
includes this evaluation framework. By this way, we dissect the CUDA time cost
for different phases in the MoE layer using PyTorch Profiler.

Figure 9: Dissecting the time cost in MoE layers (16 nodes; screenshotted via
PyTorch Profiler) Table 3: Breakdown of the time cost per iteration (micro-
batch FP) in MoE layers (16 P4d nodes)

| Switch Transformer | SMILE
---|---|---
Total Time | 535 ms | 146 ms
All2All Time Cost | 382 ms | Inter node | 77 ms
Intra node | 9 ms
FFN Expert and Others (e.g., operations other than All2Alls) | 153 ms | 60 ms
Ratio (All2All Time vs. Total Time) | 71% | 59%

The results are summarized in Figure 9 and Table 3. In Figure 9, we mainly
annotate the time cost for the All2All operations (due to two additional hops
in routing for the reversed order, SMILE has more All2Alls). The following
observations provide us with clear evidence to support our motivation to
design bi-level routing: (1) SMILE can largely improve the overhead of a
single MoE layer: bi-level routing running time (including EFA communication,
NVSwitch communication, and GPU computation for expert networks) is 3.7 times
less than one-hop routing across many nodes (146 ms vs. 535 ms); (2) All2All
time cost in SMILE is also 4.4 times smaller than that of Switch Transformer
(382 ms vs. 86 ms), matching the analysis we have explained in Section 3.1;
(3) Compared to the time cost on inter-node All2Alls (77 ms), the time cost on
intra-node All2Alls (9 ms) is much smaller due to higher bandwidth (600GB/s
vs. 50GB/s); (4) When applying SMILE, the ratio (All2All Time vs. Total Time)
is also reduced from 71% to 59%.

To understand more fine-grained details of the performance, we refer to the
whole screenshot for visualizing performance results from PyTorch Profiler in
Figure 10 and 11 in Appendix.

##  5 Conclusion

We propose a new routing algorithm and system for sparsely activated mixture-
of-experts (MoE) layer. Specifically, we introduce SMILE with bi-level routing
that better leverages heterogeneous communication bandwidth. The bi-level
routing significantly reduces network contention, launch overhead, and routing
complexity. Our experiments demonstrate that the proposed SMILE improves
training throughput by 2.5$\times$ compared to SwitchTransformer on 128 GPUs
without affecting convergence.

## References

  * [1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  Advances in neural information processing systems, 33:1877–1901, 2020.
  * [2] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [3] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In International Conference on Machine Learning, pages 8748–8763. PMLR, 2021.
  * [4] David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean.  Carbon emissions and large neural network training.  arXiv preprint arXiv:2104.10350, 2021.
  * [5] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  arXiv preprint arXiv:1701.06538, 2017.
  * [6] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  arXiv preprint arXiv:2006.16668, 2020.
  * [7] William Fedus, Barret Zoph, and Noam Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  arXiv preprint arXiv:2101.03961, 2021.
  * [8] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer.  Base layers: Simplifying training of large, sparse models.  In International Conference on Machine Learning, pages 6265–6274. PMLR, 2021.
  * [9] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  Neural computation, 9(8):1735–1780, 1997.
  * [10] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al.  Mesh-tensorflow: Deep learning for supercomputers.  Advances in neural information processing systems, 31, 2018.
  * [11] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [12] Angela Fan, Shruti Bhosale, Holger Schwenk, Zhiyi Ma, Ahmed El-Kishky, Siddharth Goyal, Mandeep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, et al.  Beyond english-centric multilingual machine translation.  Journal of Machine Learning Research, 22(107):1–48, 2021.
  * [13] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He.  Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale.  arXiv preprint arXiv:2201.05596, 2022.
  * [14] Vipin Kumar, Ananth Grama, Anshul Gupta, and George Karypis.  Introduction to parallel computing, volume 110.  Benjamin/Cummings Redwood City, CA, 1994.
  * [15] Susanne E Hambrusch, Farooq Hameed, and Ashfaq A Khokhar.  Communication operations on coarse-grained mesh architectures.  Parallel Computing, 21(5):731–751, 1995.
  * [16] Wilson L Taylor.  “cloze procedure”: A new tool for measuring readability.  Journalism quarterly, 30(4):415–433, 1953.
  * [17] William Fedus, Ian Goodfellow, and Andrew M Dai.  Maskgan: Better text generation via filling in the_.  arXiv preprint arXiv:1801.07736, 2018.
  * [18] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [19] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  arXiv preprint arXiv:1910.10683, 2019.
  * [20] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training bert in 76 minutes.  arXiv preprint arXiv:1904.00962, 2019.
  * [21] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al.  Pytorch distributed: Experiences on accelerating data parallel training.  arXiv preprint arXiv:2006.15704, 2020.
  * [22] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.  In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 3505–3506, 2020.
  * [23] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.  Zero: Memory optimizations toward training trillion parameter models.  In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16. IEEE, 2020.

##  Appendix A Appendix

###  A.1 Performance Profiling on Different MoE Layers

Figure 10: Switch Transformer MoE layer All2All time cost profiling (16 P4d
nodes)

Figure 11: SMILE layer All2All time cost profiling (16 P4d nodes)

###  A.2 Another Angle to Understand the Overhead of All2All Communication

Figure 12: The throughput varying with the number of chunks in pipelined
overlapping

From the result, we can observe that the inter-node-communication cost is
roughly a few times larger than the sum of intra-node communication and expert
forward propagation. Motivated by such a relationship in time cost, it is
possible to overlap communication cost and computational cost.

In order to verify this idea, we utilize the pipeline mechanism to parallelize
the execution of communication and computation on different hardware
resources, i.e., GPU and NIC. We evaluated the throughput in a varying number
of chunks. The results are shown in Figure 12. Unfortunately, no matter how we
manipulate the chunk size, the performance still cannot improve. We argue that
the performance degradation is due to the increase of more All2All operations.
As we know from Section 4.4 that the All2All operation is non-trivial.
Although communication and communication are overlapped in some degree, the
number of All2All operations are largely increased due to that the number of
All2All communication operations inside the MoE layer increases linearly with
respect to the number of chunks. This provides a new angle to understand the
overhead of the All2All communication in the MoE layer.