Tricks for Training Sparse Translation Models (2022)

Abstract: Multi-task learning with an unbalanced data distribution skews model learning towards high resource tasks, especially when model capacity is fixed and fully shared across all tasks. Sparse scaling architectures, such as BASELayers, provide flexible mechanisms for different tasks to have a variable number of parameters, which can be useful to counterbalance skewed data distributions. We find that that sparse architectures 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 @xcite @xcite , with recent progress driven by training transformer models on a wide range of tasks @xcite @xcite . 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 overfitting on low-resource tasks whilst simultaneously underfitting on tasks with abundant training data. in this work, we study multilingual machine translation as a multi-task learning problem @xcite , where a single model is trained to translate between many language pairs @xcite . multilingual learning has the potential of crosslingual transfer, allowing lowresource languages to benefit from high-resource data @xcite . however, in practice, this positive transfer is often mitigated by interference between languages @xcite @xcite . this is because all languages, irrespective of the amount of data, are trained with a fixed model capacity @xcite , leading to insufficient specialized capacity. recent efforts have focused on sparse architectures @xcite to train high capacity models, but these overfit to low-resource languages and have worse performance than dense architectures @xcite . 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 @xcite for multitask learning: first, we slowly ramp the number of instances from low-resource tasks over epochs rather than having a fixed sampling ratio @xcite . this promotes cross-lingual transfer and reduces over-fitting 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 first with a dense model and then gradually sparsify and specialize the experts to different tasks. overall with these two modifications, we observe improvement in low-resource performance by 0.6 bleu on wmt-15 benchmark and 1.1 bleu on ml-50 benchmark -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. ## experiments and results we experiment with @xmath0, wmt-15 1 and ml-50 @xcite ) -the first includes 15 languages and the second 50 languages. we use a transformer @xcite sequence-tosequence model with 6 encoder and decoder layers. we replace the final feed-forward layer of every alternate transformer block with a baselayer. for ml50, we increase model capacity to 12 transformer layers following @xcite . we implement our methods in fairseq @xcite and evaluate performance with bleu. ## 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 significantly increase convergence rate of mixture-of-expert models and improve their performance relative to dense baselines on two benchmarks.