Tricks for Training Sparse Translation Models
Abstract
Multi-task learning with an unbalanced data distribution skews learning towards high resource tasks, especially when model capacity is ﬁxed and fully shared across all tasks.
Sparse scaling architectures, such as BASELayers, provide ﬂexible mechanisms for tasks to have a variable number of parameters, which can be useful to counterbalance skewed data distributions.
However, we ﬁnd that that BASELayers sparse model for multilingual machine translation can perform poorly out of the box, and propose two straightforward techniques to mitigate this — a temperature heating mechanism and dense pre-training.
Overall, these methods improve performance on two multilingual translation benchmarks compared to standard BASELayers and dense scaling baselines, and in combination, more than 2x model convergence speed.
Introduction
Training a universal model capable of handling many different tasks is a longstanding ambition in natural language processing (Collobert and Weston, 2008; Ruder, 2017; McCann et al., 2018), with recent progress driven by training transformer models on a wide range of tasks (Xue et al., 2021;
Khashabi et al., 2020; Lu et al., 2020).
A central challenge in multi-task learning is accounting for the dramatically varying amounts of training data available for different tasks, which can lead to overﬁtting on low-resource tasks whilst simultaneously underﬁtting on tasks with abundant training data.
In this work we study multilingual machine translation as a multi-task learning problem (Dong et al., 2015; Firat et al., 2016), where a single model is trained to translate between many language pairs (Fan et al., 2021).
Multilingual learning has the potential of crosslingual transfer, allowing lowresource languages to beneﬁt from high-resource data (Conneau et al., 2020).
However, in practice, this positive transfer is often mitigated by interference between languages (Arivazhagan et al., 2019; Tan et al., 2019; Zhang et al., 2020).
This is because all languages, irrespective of the amount of data, are trained with a ﬁxed model capacity (Lepikhin et al., 2020), leading to insufﬁcient specialized capacity.
Recent efforts have focused on sparse architectures (Lewis et al., 2021) to train high capacity models, but these overﬁt to low-resource languages and have worse performance than dense architectures (Fan et al., 2021; Tran et al., 2021).
We analyze the learning patterns of experts throughout training and identify a fundamental problem: experts specialize early on and rarely change specialization.
We propose two straightforward techniques to improve BASELayers-based sparse architectures (Lewis et al., 2021) for multitask learning: ﬁrst, we slowly ramp the number of instances from low-resource tasks over epochs rather than having a ﬁxed sampling ratio (Arivazhagan et al., 2019).
This promotes cross-lingual transfer and reduces over-ﬁtting as the model witnesses low-resource task instances in the later epochs.
Second, we train a dense architecture before switching to sparse training.
Intuitively, we learn a generalized representation that can transfer across all tasks ﬁrst with a dense model and then gradually sparsify and specialize the experts to different tasks.
Overall with these two modiﬁcations, we observe improvement in low-resource performance by 0.6 BLEU on WMT-15 benchmark and 1.1 BLEU on MLbenchmark — whilst halving the training time.
Methods
We motivate the need for preventing early expert specialization and describe our proposal to circumvent it and more than double convergence speed.
Expert Utilization Rarely Changes
Sparse scaling schemes, such as BASELayers or
Mixture-of-Experts, enable sparse computation by distributing model capacity across sets of experts.
In each forward pass, only a small subset of experts are utilized, leading to incredibly computeefﬁcient scaling.
The challenge, however, is the routing function — or how experts can be balanced so they actually specialize to learn different things (Kudugunta et al., 2020).
When the routing mechanism is unbalanced, all the tasks degenerate to using only a single speciﬁc expert for all tasks (Lepikhin et al., 2020) — essentially wasting parameters.
BASELayers (Lewis et al., 2021) employ a simple mechanism that learns a balanced routing without the need for additional auxiliary losses.
We focus on BASELayers as it has straightforward and simple training and has previously been shown to have strong performance.
[Figure]
Even though BASELayers leads to effective utilization of all parameters, it limits parameter sharing across tasks, which is crucial when the data distribution is unbalanced — if the number of tasks and experts are the same, all tasks end up using a different set of experts.
As a result, when applied to multilingual machine translation, the performance is worse than a corresponding dense architecture.
Figure 1 demonstrates that the main reason for limited parameter sharing is that expert assignment is ﬁxed incredibly early on in training and rarely changes.
Instead of learning how to better utilize capacity across high and low-resource languages over the training process, expert capacity is essentially frozen.
We describe two strategies for more effective utilization of expert capacity, which can be easily applied to improve both low and highresource translation performance.
2. Balancing Low-Resource Tasks Temperature Sampling:
To ensure that lowresource tasks are well represented during model training, temperature sampling (Arivazhagan et al., 2019) is used to upsample low-resource tasks.
If the data distribution across different tasks is p , then temperature sampling re-scales this distribution:
[Equation]
As we increase temperature from 1 to ∞ , the sampling distribution changes from the original data distribution (e.g. highly skewed) to a uniform distribution (e.g. tasks are equally represented).
Temperature Heating:
Instead of keeping the temperature ﬁxed while sampling data for each task, we deﬁne temperature as a function of current epoch, t
= f ( e )
We deﬁne a minimum starting temperature t s , which is gradually increased at each epoch e , with a square root factor deﬁned over maximum number of epochs C .
The conduction coefﬁcient k inﬂuences the rate at which the temperature is increased over epochs.
[Equation]
In particular, we adopt square root scaling of temperature with each epoch, instead of linear to allow for gradual changes to the sampling distribution.
During the initial steps of training, this trains with lower temperatures, meaning high-resource tasks are better represented than low-resource tasks.
As a result, the experts are more uniformly assigned across high-resource tasks.
Upon slowly introducing low-resource tasks by increasing temperature during the learning process, the gating mechanism learns to route low-resource tasks through experts which were initially trained with high-resource tasks.
This promotes positive cross-lingual transfer from high-resource languages to linguistically similar low-resource languages.
2. Dense Pre-training
Architecturally, the sparsity in the output feedforward layer of the transformer block can be viewed as a version of the same transformer on multiple GPUs with two main differences: the sparse
[Table]
[Table]
[Table]
k=3 k= k=1.5 k= feed-forward layers do not share parameters (have different initialization and gradients) and an additional gating mechanism decides which token should be routed to which expert. The alternative dense architecture would fully share parameters, so all parameters are utilized for each training example rather than routing to sparse parameters.
We propose ﬁrst training a dense model for a ﬁxed number of updates.
Afterwards, we add a randomly initialized gating module and continue training the (output) feed-forward layers with sparsity, e.g. we do not average their gradients across compute nodes before back-propagating but update the weights individually in each node.
As the sparse weights slowly diverge, they become more specialized towards speciﬁc tasks.
Thus, models ﬁrst learn a generalized representation when all parameters are fully shared, and then gradually specialize to handle different tasks.
Training in this fashion not only improves the learning of specialized experts, but also increases convergence.
Experiments and Results
We experiment with English →
Many multitasking on two benchmarks, WMT-15 1 and ML50 (Tang et al., 2020) — the ﬁrst includes 15 languages and the second 50 languages.
We use a
Transformer (Vaswani et al., 2017) sequence-tosequence model with 6 encoder and decoder layers.
We replace the ﬁnal feed-forward layer of every alternate transformer block with a BASELayer.
For
ML50, we increase model capacity to 12 Transformer layers following Tang et al.
We implement our methods in fairseq (Ott et al., 2019) and evaluate performance with BLEU.
3. Effectiveness of
Temperature Heating
On WMT-15, training with BASELayers as a baseline has worse low-resource performance compared to a similarly sized dense model, losing 0.6 BLEU.
However, as we increase temperature, we recover the loss in low-resource task performance and also see improvements in the high-resource languages.
The heating technique improves the overall BASELayers model performance by +0.7 BLEU (at t s = 0.8) (see Table 1).
We observe similar trends in ML-50, where adding heating improves lowresource performance by +1.4 BLEU.
Furthermore, temperature heating improves convergence speed.
Given ﬁxed t s , the higher the k , the faster the model converges.
As shown in Figure 2, the model converges to same validation perplexity with k =3 at 50k updates as 100k updates with k =1.
Expert distribution for English to Romanian
Epoch=1 Epoch=
Epoch=
[Figure]
3. Dense Pre-training with Heating
For the WMT-15 benchmark, Table 2 demonstrates that with dense pre-training, the best performing model improves by +0.75 BLEU over baseline
BASELayers model but at the cost of 12% more compute time.
To resolve this, we reduce compu- [Table] tation time by increasing the temperature, keeping the +0.7 BLEU improvement but reducing the computation time by ∼ 60%.
Table 2 conﬁrms a similar trends on ML-50.
By combining Dense Pretraining with heating, we improve over baseline
BASELayers model by +0.5 BLEU and 2.5x in convergence speed.
However, heating can also be applied to the baselines.
In those cases, on both benchmarks, we ﬁnd that utilizing Dense Pre-training in combination with heating still has slightly better performance with signiﬁcantly faster convergence.
3. Effect on Expert Distribution
In standard BASELayer training, the learned expert distributions rarely change over training (see
Figure 1).
This prevents effective expert capacity utilization resulting in low-resource overﬁtting.
In contrast, with our proposed techniques, the expert distribution changes and learns over training.
Figure 3 compares expert distribution between ﬁxed temperature sampling and temperature heating over epochs for a low-resource language, demonstrating that temperature heating leads experts to change and learn over time.
Figure 4 shows that by utilizing dense pre-training, we observe a high entropy in the expert distribution and increased expert sharing, indicating positive cross-lingual transfer from similar high to low-resource languages.
Related Work Data Sampling
Low-resource tasks are upsampled to balance their representation when pooled with high-resource tasks.
Temperature sampling (Arivazhagan et al., 2019) upsamples the data distribution based on ﬁxed temperature, but can result in overﬁtting.
Dynamic sampler (Gottumukkala et al., 2020) selects instances based on current performance of task on dev set, which is [Figure] useful in case of catastrophic forgetting.
Learned data samplers (Wang et al., 2020) choose better are sample efﬁcient but computationally expensive.
Sparse
Scaling
Sparsely-gated
MoE models (Shazeer et al., 2017; Du et al., 2021) use a routing mechanism that decides which expert a task should be routed to.
This is the key element that governs effective (better representation) and efﬁcient (balanced assignment) resource utilization.
To promote a balanced assignment, routing techniques (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021) add a number of auxiliary task to encourage the use of diverse set of experts.
BASELayers (Lewis et al., 2021) circumvents this by treating the routing mechanism as a linear expert-to-task assignment problem, without the need of auxiliary loss.
Routing networks (Rosenbaum et al., 2018) learn better task representations by clustering and disentangling parameters conditioned on input.
Conclusion
We analyze the problem of balancing shared and specialized capacity in multitask learning, focusing on multilingual machine translation.
We present two straightforward tricks to signiﬁcantly increase convergence rate of mixture-of-expert models and improve their performance relative to dense baselines on two benchmarks