S PARSE M O E WITH R ANDOM R OUTING AS THE N EW D ROPOUT : T RAINING B IGGER AND S ELF -S CALABLE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Exploiting scale to revamp information absorption has recently become central to the success of deep learning and transformers have become de facto choice achieving numerous breakthrough performances on many real-world applications.
Despite their enormous success, gigantic transformers suffer not only from exorbitant computational and memory footprints during training but also from severe collapse as evidenced by a high degree of parameter redundancy.
Recently proposed
Sparsely-activated Mixture-of-Experts (SMoEs) models have shown promise to mitigate the issue of training efficiency, yet they have some critical limitations.
In particular, SMoEs models are prone to redundant experts due to representational collapse and poor scalability during inference and downstream fine-tuning primarily due to overfitting of the learned routing policy to the number of activated experts during training.
As recent research efforts are predominantly focused on improving routing policies to encourage expert specializations, our work focuses on exploring the overlooked scalability bottleneck of SMoEs , and leveraging it as an auxiliary tool to effectively benefit scaling large-scale dense transformers.
To this end, we propose a new plug-and-play training framework, SMoE-Dropout to enable scaling transformers to better accuracy in the full capacity setting without collapse.
Specifically, SMoE-Dropout consists of a randomly initialized and fixed router network to activate experts and gradually increase their number as training progresses over time.
SMoE-Dropout naturally provides a “self-slimmable” property offering consistent boosted performance for transformers with an increase in activated experts during inference and downstream fine-tuning, subjected to resource availability.
Our extensive experiments across diverse transformer architectures on a variety of tasks validate superior performance and substantial computation savings, compared to densely trained baselines with equivalent parameter counts.
More precisely, our trained BERT outperforms their densely trained counterpart with consistent improvements of { 1 . 03% , 0 .
78% , 1 . 09% } on challenging reasoning tasks { ASDiv-A , MAWPS , SVAMP } , respectively.
Codes and models will be publicly released.
I NTRODUCTION
Scaling neural networks, historically with the blessing of modern hardware, have dramatically improved the state-of-the-art on a wide array of real-world machine learning applications and leaderboards, conforming to the empirical scaling laws (Kaplan et al., 2020), where the final model quality has been found to have a power-law relationship with the amount of data, model size, and compute time.
Transformers (Vaswani et al., 2017), swiftly after their introduction have become de facto choice for many natural language processing (NLP)
(Yang et al., 2019c; Liu et al., 2019b; Talmor et al., 2018; Jaiswal et al., 2021;
Yang et al., 2019b; Wang et al., 2018; Ding et al., 2019; Chowdhery et al., 2022; Wei et al., 2022) and computer vision (Dosovitskiy et al., 2020; Han et al., 2020;
Touvron et al., 2021; Mao et al., 2022; Zheng et al., 2021; Parmar et al., 2018) applications and now their parameter counts are typically measured in billions rather than millions.
Unfortunately, this exploitation of parameters actuates a roughly quadratic blow-up in training costs , as both the model size and the number of training examples increase especially for dense advanced transformer-based models ( e.g., BERT (Devlin et al., 2018) and GPT (Brown et al., 2020)) and need thousands of GPU days for training.
Additionally, these gigantic transformers suffer from representation collapse
is- sue during vanilla training, which is affirmed by a high degree of parameter redundancy (Guo et al., 2019; Ganesh et al., 2020; McCarley et al., 2019) and observed ineffective usage of the transformer expressiveness (Michel et al., 2019; Chen et al., 2022a).
[Figure]
Sparse Mixture-of-Experts (SMoEs), which perform input-dependent conditional computing provide an effective way to scale model capacity given a fixed computational cost, and facilitate training transformers with significantly high parameter counts at moderately increased cost, compared to their dense counterpart, thereby improving the training efficiency .
For instance, with similar training FLOPS, Switch-Large (Fedus et al., 2021) (a kind of SMoE) is 35 × larger than a T5-Large dense model (Raffel et al., 2020).
Despite their benefits in mitigating computational and energy footprints, SMoEs tend to have many critical limitations.
Firstly , the current learning-based routing mechanisms in SMoEs tend to push hidden representations clustering around expert centroids (Chi et al., 2022), implying a trend toward representation collapse , which in turn lead to redundant experts, inferior expert specialization, thereby substandard performance (Mittal et al., 2022; Chen et al., 2022b).
Secondly , SMoEs suffer from poor scalability during inference and downstream fine-tuning prominently due to overfitting of the learned routing policy to the number of activated experts during training.
Adopting naive changes to mitigate the issue of the immutability of sparsely activated subnetworks, often lead to performance degradation.
As recent research efforts for SMoEs are predominantly focused on improving routing policies to encourage expert specializations, we explore the overlooked scalability bottleneck of SMoEs to benefit transformers and ask: Does there exist a principled, and neatly pluggable approach to modify SMoE training in pursuit to benefit scalability at inference and downstream fine-tuning of large-scale transformers, by dynamically adapting the number of activated experts subject to resource availability?
To this end, this paper proposes a novel plug-and-play training framework, SMoE-Dropout to enable scaling transformers to better accuracy in the full capacity setting without collapse.
More specifically, SMoE-Dropout consists of a randomly initialized and fixed router network to activate experts and gradually increase their number as training progresses over time.
Our simple, yet highly effective strategy has a multi-fold win-win for trained transformers, specifically: ❶ obtaining a “selfslimmable” property during inference and downstream fine-tuning subject to resource availability, which delivers a once-for-all in-situ trade-off between efficiency and performance; ❷ mitigating representational collapse and effectively utilizing the full model capacity, where activating more experts produces superior performance (Figure 1 (blue)); ❸ eliminating the overhead of learning routing policies for SMoE.
Note that SMoE-Dropout can be swiftly adapted for training any deep learning network (e.g. CNNs), given some splitting techniques (Zhang et al., 2021), but this work primarily focus on transformers considering their exploding computational footprints.
Our innovative contributions can be summarized as:
⋆ We propose a new plug-and-play training framework, SMoE-Dropout , to enable scaling transformers in the full capacity setting without collapse.
SMoE-Dropout facilitates the randomly and sparsely activated structure of network modules, playing an implicit regularization role similar to dropout .
Our new framework leads to an enhanced generalization and reduced training costs ( e.g. , up tp 37% running time savings) in comparison with the vanilla training of large dense transformers having equivalent parameter counts.
⋆ SMoE-Dropout naturally provides a “self-slimmable” property offering consistent boosted performance for transformers with an increase in activated experts during inference and downstream fine-tuning (Figure 1 (blue)), subjected to resource availability using a randomly initialized and fixed router network.
It enjoys a once-for-all in-situ trade-off between efficiency and performance when considering deployment.
⋆ Our extensive experiments across representative architectures on a variety of tasks validate the effectiveness of our proposed SMoE-Dropout.
Specifically, during pre-training,
our approach has { 1 . 37 , 4 . 10 } , { 2 . 53 , 12 . 44 } and { 154 . 12 , 188 . 00 } ( × 10 − 2 ) lower
BPC than { vanilla dense training (with the same parameter counts), learned SMoE } for Transformer-XL, BERT, and RoBERTa, respectively; after transferring, SMoEDropout obtains { 0 .
78% , 1 . 09% } performance improvements for BERT and {− , 5 . 88% , 0 .
07% , 5 . 04% } for RoBERTa, on { CSQA , ASDiv-A , MAWPS , SVAMP } reasoning tasks compared to its dense training counterpart.
R ELATED W ORKS
Mixture of Experts (MoE).
MoE is a special kind of neural network, where its parameters are partitioned into a series of sub-modules (or called experts), and a conditional computation is performed in an input-dependent fashion (Jacobs et al., 1991; Jordan & Jacobs, 1994;
Chen et al., 1999; Yuksel et al., 2012).
The traditional dense MoEs require intensive computational cost since they adopt all experts for each input (Eigen et al., 2013).
Fortunately, recent investigations (Shazeer et al., 2017;
Lepikhin et al., 2020; Fedus et al., 2021) have proved the effectiveness of MoEs with sparsely activated experts ( i.e. , SMoE) at both training and inference stages, which greatly trim down the cost and scale language models to enormous sizes like trillions of parameters (Fedus et al., 2021).
This efficient fashion of SMoEs gains increasing popularity in various NLP (Shazeer et al., 2017; Lepikhin et al., 2020; Zhou et al., 2022; Zhang et al., 2021; Zuo et al., 2022; Jiang et al., 2021) and vision (Riquelme et al., 2021; Eigen et al., 2013; Ahmed et al., 2016; Gross et al., 2017; Wang et al., 2020; Yang et al., 2019a; Abbas & Andreopoulos, 2020; Pavlitskaya et al., 2020) tasks.
However, its sparse-gated manner incurs several downsides of: ( 1 ) Unstable training .
Zoph et al.
(2022) points out that many methods like gradient clipping can stabilize SMoEs’ training, but usually at the expense of worse quality.
The router z -loss
(Zoph et al., 2022) is an ideal solution for the double-win of performance and stability.
( 2 ) Poor specialization .
One of the intriguing goals of SMoE is to divide-and-conquer the learning task by solving each piece of the task with adaptively selected experts (Aoki et al., 2021;
Hazimeh et al., 2021; Ma et al., 2018; Mittal et al., 2022).
To encourage the specialization and reduce the redundancy across experts (Chen et al., 2022b), Dai et al.
(2022) pre-defines the expert assignment for different input categories, and Hazimeh et al.
(2021) advocates multiple diverse router policies.
( 3 ) Representation collapse and load imbalance among experts .
As the major issue behind learning-based SMoEs, various approaches are proposed to mitigate its negative effects, which date back to 2017 .
Shazeer et al.
(2017) injects Gaussian noises into gating networks to promote the routing balance.
Later on, for the same purpose, Lepikhin et al.
(2020); Fedus et al.
(2021) apply an auxiliary loss of load balancing regularizers; Lewis et al.
(2021) performs the routing by dealing with a linear assignment problem; Clark et al.
(2022) utilizes reinforcement learners; and Zhou et al.
(2022) routes top-k inputs for each expert instead of selecting experts for each input sample.
Beyond learned routing policies, Roller et al.
(2021) and Zuo et al.
(2022) design deterministic hashing and stochastic assignments, respectively, which discard the need for router networks.
Zuo et al.
(2022), one close literature, also recommends the advantage of a stochastic assignment.
They randomly activate experts for each input during training and inference, which leads to inconsistent inference results.
To alleviate the prediction randomness, Zuo et al.
(2022) uses a consistent regularized loss to penalize the discrepancy between different experts.
However, such regularization is prone to enforce redundancy to SMoEs and sacrifice the network capacity.
In our proposal, the fixed router with random weights generates deterministic inferences.
Meanwhile, the presented “self-slimmable” attribute suggests sufficiently exploited expressiveness of the full model.
Dropout and Other Training Techniques for Transformers in NLP.
Dropout (Srivastava et al., 2014) is invented to prevent overfitting in the training phase of over-parameterized networks by randomly omitting neurons and their associated connections.
Followup studies develop plenty of dropout variants (Zhang & He, 2020; Wan et al., 2013; Ba & Frey, 2013; Kingma et al., 2015;
Gal et al., 2017; Wu & Gu, 2015; Tompson et al., 2015; DeVries & Taylor, 2017; Park & Kwak, 2016; Semeniuta et al., 2016).
In parallel, McAllester (2013); Mou et al.
(2018); Mianjy & Arora (2020); Zhang & Xu (2022); Neklyudov et al.
; Gal & Ghahramani (2016) devote themselves in deriving the theoretical foundation for dropout and explaining its implicit regularization impacts.
Other notorious bottlenecks of transformer training mainly lie in overfitting and instability, resulting from poor optimization (Zhang et al., 2019; Liu et al., 2019a; 2020a), insufficient or heterogeneous downstream data (Variˇs & Bojar, 2021; Zhang & Vaidya, 2021), etc. .
Accordingly, numerous
Normalization
[Figure]
Attention
Normalization
Multi-Layer Perceptron ...
Multi-Layer Perceptron
Expert
Expert Expert
Transformer Layer
Expert Expert
Expert Expert
Expert ...
Splitting
Modulization ......
Token Embedding ...
Random Router
Expert Expert Expert
Expert ...
Point-
[Figure]
Point
[Figure]
Point-wise Multiplication
Summation
SMoE-Dropout ......
Gradually Increased k
Along with Training remedies are developed to address the issues.
For example, data augmentations (Sun et al., 2020), improved initialization (Liu et al., 2020b;a; Xu et al., 2020; Zhu et al., 2021), upgraded normalization (Wang et al., 2022; Yang et al., 2022), enhanced optimizers (Cohen et al., 2022), weight decay (Loshchilov & Hutter, 2017), and early stopping.
M ETHODOLOGY 3. P RELIMINARY
Sparse Mixture of Experts (SMoEs).
SMoE models take advantage of conditional computation activating different subsets of a network for different inputs.
A building block of SMoEs is the expert layer including a multi-head attention block and multiple experts in parallel.
In this work, we consider SMoE for Transformers, where SMoE layers are incorporated into contiguous Transformer blocks.
SMoE expert can be normally constructed by either splitting the vanilla MLP of transformers into smaller pieces (Zhang et al., 2021) or replicating the MLP (Fedus et al., 2021).
Most existing SMoE works deal with the MLP component in transformers because MLPs constitute roughly 2 / 3 of total model parameters counts storing significant amounts of learned knowledge as memory networks (Geva et al., 2020; Dai et al., 2021).
Let {E i } N i =1 denote the experts, where i is the index of expert and N is the total number of experts.
A gating network or router R is inserted to choose the top- k experts with the largest scores R ( x ) i , and x represents the input embedding.
Usually, k ≪ N , which implies a sparsely activated setting.
Specifically, the resultant output of the expert layer can be depicted as follows:
[Equation]
if v is the top k
[Equation]
otherwise (1) where G is the critical part of a router R .
For a learnable routing, G is a neural network that can be one or a few layers MLP (Shazeer et al., 2017; Fedus et al., 2021).
E j ( x ) stands for features from the expert E j .
It will be further summed with a scaling coefficient R ( x ) j to form the final output y .
The TopK function maintains the largest k values and sets the reset elements to zero.
In practice, a load or important balancing loss (Shazeer et al., 2017) is employed to avoid the representation collapse issue, i.e. , always picking the same experts for different inputs and ignoring others.
Dropout and its variants.
Dropout is a classic training technique for mitigating potential overfitting.
The vanilla dropout is typically applied to fully connected layers with a dropping probability p .
For each training iteration, neurons will be disabled with the probability p .
In other words, the omission of neurons follows a Bernoulli ( p ) distribution.
As for the inference phase, there is no dropout and all neurons are activated.
To compensate for the additional information during training, the output logits are reweighted by 1 − p .
In this paper, we selected two representatives among diverse proposed dropout variants, concrete dropout (Gal et al., 2017) and dropblock (Ghiasi et al., 2018) as our comparison baselines.
▷ Concrete Dropout.
It replaces the discrete Bernoulli ( p ) distribution of dropout with a continuous relaxation, i.e. , Concrete distribution, and allows an automatic tuning of the dropping probability p .
For example, considering the one-dimensional case, as shown in Gal et al.
a
Concrete random variable z is described as z = sigmoid (cid:0) 1 t × (cid:0) log( p ) − log(1 − p )
+log( u ) − log(1 − u ) (cid:1)(cid:1) , where u ∼ Unif (0 , 1) is a uniform random variable and t denotes a temperature hyperparameter.
Note that parameter p is optimized in a data-driven way.
▷ DropBlock.
Instead of performing Bernoulli dropping per feature map, Ghiasi et al.
(2018) applies it in areas within feature maps.
They claim that DropBlock improves the generalization and limits overfitting by hiding certain areas of features or input samples
. 3. A N EW T RAINING P IPELINE : SM O E-D ROPOUT
Modulization.
Our first step in our SMoE-Dropout, turns a large densely connected MLP into multiple smaller MLPs with the same size, as demonstrated in Figure 2.
Without loss of generality, in Figure 2, we use a single-layer MLP f with a dimension d for illustrations.
After the modulization, it is divided into a set of MLPs {E 1 , E 2 · · · , E N } , where they have the same hidden dimension d N .
Random Routing Policy.
Few prior works have investigated some form of random routing policies, such as Roller et al.
(2021) utilizes a hash table to enforce a pre-defined deterministic random mapping from inputs to experts and Zuo et al.
(2022) adopts a fully random assignment in each training iteration.
Although they have shown some benefits from random policies, both methods suffer from inconsistent inference predictions, and can not outperform the densely trained models with the same parameter counts.
Our proposed framework, SMoE-Dropout considers a randomly initialized and fixed router network to guide the token assignment.
Different from previous works, our proposal’s assignment is ( 1 ) implicitly optimized during training, since feature embeddings are kept updated for the same input sample; ( 2 ) deterministic during inference thanks to the fixed weights in R .
Extensive results in Section 4 verify the superiority of our proposal, compared to existing random policies and the dense baseline with the same model parameters.
Additionally, another crucial design in SMoE-Dropout’s routing is the progressively enlarged number of activated experts ( k ) .
Riquelme et al.
; Jiang et al.
(2021) reveals that altering k in the inference phase incurs significant performance degradation if the SMoE is learned with a fixed k .
For example, (Riquelme et al., 2021)’s SMoE trained with k = 1 has 20% ∼ 30% accuracy drops on ImageNet, when activating k ≥ 7 experts during the evaluation.
This drawback largely limits the practical usage of SMoEs.
Because in real-world applications, different scenarios have distinctive resource budgets, requiring flexible and effective network capacity at the inference stage.
To tackle this limitation, we adopt a gradually enriched active network capacity along with the training, which is implemented through linearly increasing the number of selected experts k .
Such training approach coincided with the principle of curriculum learning brings the attractive “self-slimmable” ability offering consistent boosted performance for transformers with an increase in activated experts during inference and downstream fine-tuning, as shown in Figure 1.
SMoE-Dropout.
Our effective proposal is formed by three simple and highly effective steps, as described in Figure 2.
First, it breaks the MLP into a series of MLPs with a reduced size for modulization ( Middle Left of Figure 2).
Then, a random policy parameterized by fixed weights is introduced to route token embeddings to k experts with the largest response ( Middle Right of Figure 2).
Finally, it progressively actives more experts, preventing the overfitting to the amounts of used network capacity during training.
( Right of Figure 2).
E XPERIMENT 4. I MPLEMENTATION D ETAILS
Network Architectures and Comparison Baselines.
In our experiments, we have adopted three representative transformer-based networks, including BERT (Devlin et al., 2018), TransformerXL
(Dai et al., 2019), and RoBERTa (Liu et al., 2019b).
Specifically, we use double-size BERT base / RoBERTa base that have 12 transformer layers, 768 -dimensional encoder layers, 6144 -/ 3072 - dimensional feed-forward networks (MLPs), and 12 attention heads.
For both Transformer-XL, we choose a reduced size due to limited resources, which has 4 layers, 256 -dimensional encoder layers, 8192 -dimensional feed-forward networks, and 8 attention heads with a head size of 64 .
For sufficient comparisons with our proposal, Training w. SMoE-Dropout, we consider five baselines: ( i ) Densely Training w. Dropout, where the vanilla dropout is applied to feed-forward networks (MLPs); ( ii ) Densely Training w. Concrete Dropout (Gal et al., 2017); ( iii ) Densely Training w. DropBlock (Ghiasi et al., 2018).
Note that both Concrete Dropout and DropBlock are inserted in feed-forward networks, replacing the vanilla dropout; ( iv ) Training w. Learnable SMoE (Fedus et al., 2021); ( v ) Training w. THOR
(Zuo et al., 2022), where THOR is another random SMoE that randomly activates a pair of experts for each input sample and adopts an auxiliary consistency regularization based on Kullback-Leibler (KL) divergence.
To compute the regularization term, two forward processes are needed in each training iteration.
Pre-Training.
▷ Datasets.
Transformer-XL is pre-trained on enwik8 (Mahoney, 2011) dataset, while we use BooksCorpus (Zhu et al., 2015) for BERT and RoBERTa.
▷ Training Configurations.
For Transformer-XL, we follow the official training setups, using Adam optimizer and the learning rate starts from 2 .
5 × 10 − 4 and decreases according to a cosine annealing scheduler.
We use a batch size of 22 and optimize the network for 4 × 10 5 iterations.
As for BERT pre-training, we adopt an AdamW optimizer with an initial learning rate of 5 × 10 − 5 that linearly decays to 0 .
The batch size and total training steps are 64 and 1 × 10 5 , respectively.
RoBERTa’s pre-training configurations strictly follow the default from HuggingFace 1 , but with reduced training steps of 1 × 10 5 .
Moreover, we conduct a grid search and set the coefficient of THOR’s regularization term as 2 .
Similarly, the temperature in Concrete dropout is t = 0 .
▷ Evaluation Metrics.
Since both performance and efficiency are essential, we assess the pre-training performance via Bits-Per-Character (BPC) on the hold-out validation set, where a smaller BPC value indicates a better pre-training; and we report training time per iteration & the number of floating point operations (FLOPs) of singlesample inference, for evaluating the efficiency.
{ 1 RTX A 6000 , batch size 22 } and { 8 V 100 , batch size 64 } are adopted for time measurements of Transformer-XL and BERT/RoBERTa, respectively.
Downstream Fine-Tuning.
▷ Datasets.
Five benchmarks across three downstream tasks are examined in this paper, including text classification ( SST-2 (Socher et al., 2013)), arithmetic reasoning ( ASDiv-A (Miao et al., 2020), MAWPS (Koncel-Kedziorski et al., 2016), SVAMP Patel et al.
(2021)), and commonsense reasoning ( CSQA (Talmor et al., 2018)).
▷ Training Configurations.
We perform dense fine-tuning for all approaches.
Given a downstream parameter budget, SMoE-based methods will select the most voted experts based on their routing policies.
Detailed training setups are listed as follows.
We fine-tune the pre-trained Transformer-XL with a smaller learning rate of 1 × 10 − 4 and a batch size of 64 on SST-2 benchmark.
And for BERT and RoBERTa, we fine-tune the models on the aforementioned four reasoning datasets.
The learning rate is fixed at 2 × 10 − 5 and the batch size is 64 .
In each downstream task, the fine-tuning continues for 3 epochs, while other configurations are kept the same as the ones in pre-training.
▷ Evaluation Metrics.
At the evaluation phase, accuracy (%) and the problem solving rate (%) (Wei et al., 2022) are reported on the test set of SST-2 and other reasoning tasks, respectively.
4. S UPERIOR P ERFORMANCE OF SM O E-D ROPOUT
We adopt classical transformer-based models, i.e. , { Transformer-XL, BERT, RoBERTa } , and train them on { enwik8 , BookCorpus , BookCorpus } in a dense or SMoE-based manner.
Evaluation results are summarized in Table 1, where all models are compared under the same number of parameter counts.
The following observations can be drawn: ❶ Our SMoE-Dropout demonstrates superior performance compared to all other training algorithms.
Specifically, SMoE-Dropout with all experts selected obtains 1 .
37 ∼ 18 . 49 , 0 .
56 ∼ 12 . 44 , and 152 .
82 ∼ 188 .
( × 10 − 2 ) lower BPC for Transformer-XL, BERT, and RoBERTa, respectively.
It validates the effectiveness of our proposals.
❷ Appropriate random routing policies show consistent performance benefits across all three network backbones.
Meanwhile, our randomly weighted router surpasses the completely random allocation in THOR, which is within expectation since our assignment is implicitly “optimized” given evolved feature embeddings.
❸
In terms of training efficiency, SMoE-Dropout has up to 21% , 37% , and 25% training time savings compared to the dense training of three backbones.
If only activating half experts ( k = N2 ), ours enjoy extra 23% ∼ 34% inference FLOPs reduction with a comparable
BPC.
Although learnable SMoE reaches the best efficiency, it ends up with inferior performance.
[Table]
[Figure]
Besides, we report another group of experiments varying the expert numbers (parameter counts) during evaluation.
As shown in Figure 3, for SMoE-based approaches, we directly change the number of activated experts at the inference stage, which is an in-situ fashion from the single trained transformer.
While for dense training baselines, each dot in their curve requires a separately trained model since it does not allow modifications of network capacity without further fine-tuning.
From the results, we see: ❶ The performance of SMoE-Dropout is stably improved along with more parameters used, and it outperforms the others after 1 . 0 , 10 , and 8 ( × 10 7 ) parameter counts for three backbones.
Such “slimmable” property enables scaling transformers to the full capacity without collapse, bringing a once-for-all trade-off respected to inference resource availability.
❷
In contrast, learnable SMoE’s and THOR’s BPC are quickly saturated and deteriorated when adopting more experts, which implies the existence of expert redundancy (or representation collapse).
The potential reasons for their substandard results are ( i ) the overfitting to fixed
# experts utilized during training for learnable SMoE, ( ii ), and the consistency regularization between experts’ predictions for THOR.
4. T RANSFER S TUDY OF SM O E-D ROPOUT : S ELF -S LIMMABLE
We further investigate SMoE-Dropout and its intriguing “self-slimmable” property in a transfer learning scenario.
Pre-trained models from Section 4.2 are densely fine-tuned on various downstream tasks, including text classification { SST-2 } and challenging arithmetic & commonsense reasoning { CSQA , ASDiv-A , MAWPS , SVAMP } .
The performance 2 is collected in Table 2.
We find: equipped with SMoE-Dropout, Transformer-XL achieves 0 .
47% ∼ 2 .
43% accuracy improvements on SST-2 , BERT / RoBERTa obtain { 0 .
07% ∼ 9 .
42% ∼ 3 .
26% ∼ 1 .
30% , 1 . 09% ∼ 4 .
90% } and {− , 2 .
10% ∼ 5 .
07% ∼ 0 .
27% , 5 . 04% ∼ 5 .
93% } performance boosts on { CSQA , ASDiv-A , MAWPS , SVAMP } respectively, suggesting an enhanced transferability.
Similarly, we alter the model capacity during downstream fine-tuning.
Starting from one pretraining, the SMoE-based method first computes the selected times of each expert based on one feedforward pass with downstream data, then chooses the top activated experts to meet certain parameter budgets, and performs the subsequent dense fine-tuning.
As displayed in Figure 4, our SMoE-Dropout has a continually increased accuracy or problem-solving rate when involving more parameters, and clearly surpasses the rest of approaches at parameter counts beyond 0 .
8 , 8 , and 10 .
5 ( × 10 7 ) for Transformer-XL, BERT, and RoBERTa respectively.
It shows a flexible capacity adjustment, i.e. , “self-slimmable”, according to the downstream resource constraint.
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
4. E XTRA I NVESTIGATION AND A BLATION S TUDY
Q1:
When does SMoE-Dropout outperform other baselines?
A1:
Sufficient Model Capacity.
To answer Q1 and understand SMoE-Dropout’s superiority in diverse scenarios, we investigate our proposal with different model capacities by varying model depth ( e.g. , layers) & width ( e.g. , experts).
Model Depth - Different Number of Layers.
We conduct experiments on enwik8 dataset with
Transformer-XL that has 2 , 4 , 8 , 12 layers and each layer is turned into the SMoE layer through a modularization.
The comparison results of Densely Training w. Dropout and Training w. Learnable
SMoE are reported in Figure 5 (a).
We find that densely trained transformer performs the best when the network capacity is small like 2 layers, while with sufficiently large model capacity ( ≥ 4 layers),
SMoE-Dropout demonstrates a consistent advantage compared to the others.
Meantime, along with the increase of layers, the performance gap of SMoEs between the learned policy and our random policy keeps enlarging, signifying SMoE-Dropout’s better scalability.
Model Width - Different Number of Experts.
Similarly, we study the influence of model capacity by examining Transformer-XL with different widths of 2 , 4 , 8 , 16 experts.
Results are summarized in Figure 5 (b).
Consistent observations can be drawn that: ( i ) Densely Training w. Dropout outperforms SMoE-based training under small network widths such as ≤ 8 experts; ( ii ) SMoEDropout presents enhanced performance when applied to large models with 16 experts; ( iii )
Learnable routing policies are effective with a small number of experts like ≤ 8 experts, while it gets worse results than our random routing with a sufficient number of experts, e.g. , 16 experts.
Q2:
What is a better SMoE-Dropout design?
A2: Random Weight Router; Later-layer SMoE.
To answer Q2 , we focus on the main constituents of SMoE-Dropout: Modularization , Random
Routing Policies , and Gradually Increased k .
Comprehensive ablations are depicted as below.
Ablation on Diverse Random Routing Policies.
An appropriate design of random routing policies determines the achievable performance of SMoE-Dropout.
We compare our random initialized and fixed router to SMoE with fully random assignments (Zuo et al., 2022) and random hash SMoE with a pre-defined deterministic random assignment (Roller et al., 2021).
Transformer-XL results on
[Figure]
enwik8 are collected in Fig.
5 (c), where our proposed random routing obtains substantially lower
BPC of 2 .
96 ∼ 170 .
11 ( × 10 − 2 ) than the other two under different amounts of model parameters.
Ablation on w./w.o.
Gradually Increased k .
Figure 5 (d) investigates SMoE-Dropout variants with and without gradually increased k .
We see that disabling the progressive manner of enlarging the number of activated experts causes unsatisfied performance.
Also, as a result, the “selfslimmable” property completely disappears, e.g. , adopting all model parameters leads to worse BPC.
Ablation on Different Positions for Modularization.
It remains mysterious where is the best position to insert SMoE layers.
To address this question, we perform modularization to different transformer layers and record their performance in Figure 5 (e).
Specifically, given a 4 -layer
Transformer-XL, we compare four options: ( i )
Early , the first two layers are SMoE layers; ( ii )
Middle , the 2 nd and 3 rd layers are SMoE layers; ( iii ) Later , the last two layers are SMoE layers; ( iv ) Every-2 , there is one SMoE layer every two transformer layers, i.e. , the 2 nd and 4 th layers.
From the results, introducing SMoEs to later layers is in general more beneficial than modulizing earlier transformer layers.
One behind possible reason is that shallow layers might capture common features that need to be shared across input samples.
More dissections are left for future works.
Q3: Extra benefits from SMoE-Dropout? A3: Improved Distillation and Less Overfitting.
Distilling into Single Expert on Downstream.
Besides all the benefits in pre-training inference and downstream transfer, we explore additional advantages of SMoE-Dropout under the distillation scheme that is usually preferred in resource-limited applications.
As shown in Table 3, we distill all pre-trained Transformer-XLs into the same smaller variant with a single expert on the SSTdownstream task.
Our algorithm produces the most distillable models among all four methods by a clear accuracy margin of 0 .
76% ∼ 1 . 89% .
[Table]
Overfitting.
We examine SMoE-Dropout, SMoE, and densely trained transformers regarding how likely they overfit the training data distribution with the increase of model parameters.
As shown in Figure 5 (a) and (b), experiments are conducted on enwik with Transformer-XL, and three approaches are compared under the same parameter counts.
We observe both SMoE-Dropout and
Densely Training w. Dropout do not show any sign of overfitting.
That is, the performance is consistently improved as we increase the layers from 2 to 12 or experts from 2 to 16 .
In contrast, Training w. Learnable SMoE incurs BPC deterioration owing to overfitting when we expend the transformer to 12 layers, similar to the findings in Zoph et al.
We credit the reduced overfitting to the implicit regularization effect of SMoE-Dropout and Dropout.
C ONCLUSION
In this paper, we present a novel plug-and-play strategy for training over-parameterized transformer models, which leverages SMoEs to empower scaling transformers to better performance in fullcapacity settings without collapse.
Unlike existing SMoEs, we design a randomly initialized and fixed router to assign experts and gradually increase their number along with the training.
As a result, our proposal grants an appealing “self-slimmable” property to large transformers during inference and downstream fine-tuning, based on resource availability.
It implies alleviated representation collapse and delivers a once-for-all in-situ trade-off between efficiency and performance.
Extensive empirical studies across various combinations of network backbone and dataset, consistently demonstrate the significantly improved performance and training time savings from our algorithm.
Future work includes the extension of other network architectures and tasks like vision recognition. M ORE I MPLEMENTATION D ETAILS [Table] [Table]
# Apply Concrete Dropout
[Table]
=
[Table]
._concrete_dropout
[Table]
, p
[Table]
# Feed forward through the next
[Table]
= next_layer(output)
# Calculate the weight regularizer sum_of_square
=
for param in next_layer.parameters(): sum_of_square += torch.sum(torch.pow(param, 2))
weights_regularizer
=
self.weight_regularizer * sum_of_square / (1 - p)
# Calculate the dropout regularizer dropout_regularizer
= p * torch.log(p) dropout_regularizer += (1. - p) * torch.log(1. - p) input_dimensionality =
input[0].numel() dropout_regularizer *= self.dropout_regularizer
* input_dimensionality regularization =
weights_regularizer + dropout_regularizer return output, regularization def
[Table]
concrete
[Table]
[Table]
[Table]
[Table]
[Table]
= 1e
[Table]
[Table]
temp = 0.
# Calculate the dropout probability matrix unif_noise =
[Table]
[Table]
[Table]
Details of Concrete Dropout and DropBlock.
In our experiments, we use Concrete-Dropout or DropBlock to replace the original dropout layer in MLP blocks.
And for Concrete Dropout, we adopt the official implementation https://github.com/yaringal/ConcreteDropout/ blob/master/concrete-dropout-pytorch.ipynb .
For DropBlock, we follow the method in Ghiasi et al.
And the PyTorch-style pseudo codes for both methods are presented in Algorithm 1 and 2.
A M ORE E XPERIMENT R ESULTS A2. S TABILITY A NALYSIS
[Table]
To evaluate the stability of the improvement obtained by our SMoE-Dropout, we carry out further experiments of Transformer-XL on SST-2 .
The results are reported in Table A4, from which we can observe that our SMoE-Dropout achieves a statistically significant improvement of 0 .
93% ∼ 1 .
17% accuracy gains compared with other SMoE-variants and the dense network, where there is no overlap between the error bars (one standard deviation).
A2. C OMPARISON WITH L EARNABLE SM O E S W . G RADUALLY I NCREASED k
Table A5 demonstrates that both random routing policy and progressively increasing the number of activated experts are beneficial for alleviating representation collapse and providing self-slimmable property, yet not as good as combining both.
To be specific, when applying the strategy of progressively enlarging the number of activated experts, the learnable SMoEs suffer less representation collapse and achieve better performance, i.e. , 0 .
31% higher accuracy.
Meanwhile, We find that learnable SMoE with curriculum learning has the self-slimmable property only when activating experts from k = 1 to k = 8 .
However, the performance starts to degrade if using more experts like k = 16 .
As for our SMoE-Dropout with a random routing, it enjoys a better “self-slimmable” property from k = 1 to k = 16 (full model capacity), with up to 0 .
87% higher accuracy on SSTacross all scenarios, compared to its learnble varints.
[Table]
A2.
[Table]
T RANSFER S TUDY ON M ULTI -S TEP R EASONING T ASKS
We conduct a further transfer study of the pre-trained BERT networks on a multi-hop questionanswering dataset, HotpotQA Yang et al.
And We use exact match (EM) accuracy to assess networks’ performance.
Following the same metric in Press et al.
(2022), we calculate the compositionality gap, i.e., the gap of EM accuracy between multi-hop question answering and its all single-hop sub-questions Tang et al.
(2020), of each network.
As shown in Table A6, our SMoEDropout is beneficial for reducing the compositionality gap, which achieves the best performance with up to 0 .
30% higher EM score and 0 .
30% narrower compositionality gap, compared with the learnble SMoE and its dense counterpart.
[Table]
[Table]
Compositionality
[Table]
[Table]
[Table]
A