Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference

• Token-level routing can require model-parallel serving, inter-device communication during autoregressive decoding, and large batches to avoid device under-utilization.Source: §1 · p. 1
• For multilingual translation, the goal is to preserve cross-language transfer while reducing task interference and extracting per-language-pair subnetworks for inference.Source: §1 · p. 2
• Prior work cited in the introduction found that distillation preserved only a small fraction of large sparse models' quality gains.Source: §1 · p. 2
• The proposed routing sends all tokens for a task collectively to the same experts, then decodes each task using only its associated expert subset.Source: §1 · p. 2
• +3.6 BLEU over the multilingual model trained from scratch and +2.1 BLEU over the distilled dense student at comparable inference cost are reported for task-level routing.The comparison is reported for the multilingual WMT task.Source: §1 · p. 2
• Comparable quality gains to token-MoE come with 1.9x peak throughput and 6.3% of the decoder size.Source: §1 · p. 2
• On the in-house dataset, task-level routing reports +3.6 BLEU over the dense baseline and comparable quality gain to token-level routing with 2.6x higher peak throughput.The task-level model uses 1.6% of the token-level model's decoder size.Source: §1 · p. 2
• The paper also analyzes MoE routing decisions to motivate its method.Source: §1 · p. 2

• A Transformer sequence-to-sequence model contains an encoder and decoder, each with multiple Transformer layers.Source: §2 · p. 2
• Each MoE layer contains E feed-forward networks, with expert outputs combined according to the gating network's selected weights.Source: §2 · p. 2
• The gating network routes each input token to selected experts, and this work chooses the top two weights for each example.The gating function must balance expert utilization and remain efficient to implement at scale.Source: §2 · p. 2

3 Methods

• Task-level routing is motivated by balancing cross-language transfer with interference from extensive parameter sharing in multilingual translation.MoE variants partition parameters across languages or tasks; the authors study routing to mitigate interference while enabling transfer and efficient inference.Source: §3 · p. 3

3.1 Routing Strategies

• Sentence-level routing sends every token in a sentence to experts selected from the sentence's average token representation.Token-level routing instead makes routing decisions independently for each token.Source: §3.1 · p. 3
• Task-level routing selects experts at task boundaries, defined in multilingual translation by target language or language pair.For example, French-to-English and German-to-English can count as the same task under target-language routing or separate tasks under language-pair routing.Source: §3.1 · p. 3

3.2 Inference Implications of Routing Strategies

• For large MoEs, host-to-device expert loading incurs communication costs, while model parallelism introduces inter-device communication overhead.Task-specific sub-networks are presented as an alternative to complex serving strategies or compression.Source: §3.2 · p. 4
• Token- and sentence-level routing may require loading all E experts when inputs choose disjoint experts, whereas task-level top-K routing preloads only K experts.The task-level assignment is static for tokens in an input sentence; token- and sentence-level decisions are dynamic.Source: §3.2 · p. 4

4.1 Experimental Setup

• WMT experiments cover 15 languages to and from English, with training data ranging from more than 60 million sentence pairs (en-cs) to roughly 150k (en-gu).The models use the 142M Transformer Base architecture or enhanced versions with MoE layers.Source: §4.1 · p. 4

4.2 Comparison of different Routing Strategies on WMT

• 1 BLEU over Token MoE, decoder-only Task MoE (Token/Target) works best, with gains consistent across directions, resource levels, and two zero-shot pairs.The routing strategies were compared at identical network size: 32 experts and 533M parameters.Source: §4.2 · p. 5

4.3 Comparison of Throughput of Sparse Models

• 1.87 times higher peak throughput for Task-MoE than Token-MoE, with 3.75 times less decoder parameters (142M vs 533M) and 0.0% versus 26.9% of step time in communication overhead.The throughput comparison decodes the WMT14 English-German test set on 32 Cloud TPU V3 cores; decoder-only task routing has roughly the same inference cost as routing on both encoder and decoder.Source: Figure 2 · p. 5, §4.3 · pp. 5–6

4.4 Comparison of Extracting Task MoE Models to Distillation

• 29.0 average BLEU for best-performing Task-MoE versus 26.9 for distilled Token-MoE (+2.1 BLEU); distillation preserves only 32.25% of gains over dense multilingual baselines.The comparison uses models with similar inference cost and covers several language pairs.Source: Table 2 · p. 6, §4.4 · p. 6

5.1 Experimental Setup

• The experiments scale to 204 direct language pairs and 25 billion sentence pairs, using a 473M Transformer Big architecture and comparing token and task routing configurations.The web-crawled corpus is diverse, noisy, and heavily imbalanced across language pairs; scaled models use encoder token routing and decoder routing by target language or language pair.Source: §5 · p. 6, §5.1 · p. 6, §5.2 · p. 7

5.2 Results

• Hybrid task-level routing performs slightly better on En-Xx pairs, while pure token-level routing performs slightly better on Xx-En pairs.Both scaled sparse models outperform their dense baselines; the authors hypothesize that leaving decoder experts undivided by task improves transfer for Xx-En.Source: §5.2 · p. 7
• At inference, token-level decoder routing uses 200M parameters, although both routing forms have 13B parameters at training time.The paper also points to practical serving considerations discussed in Section 3.1.Source: §5.2 · p. 7

5.3 Comparison of Throughput on Massive Models

• 2.6 times higher peak throughput is achieved by Task-MoE, using 201M decoder parameters versus Token-MoE's 6.5B.On 128 Cloud TPU V3 cores, communication overhead was 0.2% of step time for Task-MoE versus 36% for Token-MoE.Source: §5.3 · p. 7, Figure 4 · p. 8

5.4 A Closer Look at the Routing Decisions

• Encoder tokens across Xx-En tasks slightly prefer the same few experts, while decoder tasks slightly prefer different experts.Related languages appear to have similar expert distributions; the authors interpret decoder task preferences as evidence that task-level decisions emerge naturally under token routing.Source: §5.4 · p. 8

6 Related Work

• Conditional computation routes examples through only a sub-network, with routing based on input, computational budget, or simplified expert allocation.Prior work applied conditional computation in NLP and computer vision.Source: §6 · p. 8
• Multi-task learning activates sub-networks according to task and can share parameters; prior applications include classification, recommender systems, and NLP.Some approaches were designed with inference benefits in mind.Source: §6 · p. 8
• Multilingual NMT work explores language-specific components and language-group parameter sharing, while this approach targets scaling without harming inference efficiency.A related inference-focused method learns language-specific layer combinations; this work instead focuses on inference efficiency in mixture-of-experts models.Source: §6 · p. 8

7 Conclusions

• Task-based routing performs as well as or better than token-based routing on WMT's 30 language pairs and an internal dataset covering 200 pairs, measured with BLEU.The conclusion compares translation quality across the two datasets.Source: §7 · p. 9
• Task-based routing is superior to token-based MoE serving in peak throughput and to distillation in BLEU.The authors compare throughput across routing approaches and distilled models.Source: §7 · p. 9
• The authors identify more granular hybrid or hierarchical routing as future research directions for large, sparsely gated, multi-task networks.They characterize inference-friendly routing that retains MoE quality gains as a promising direction.Source: §7 · p. 9

A.1 WMT Model and Training Details

• The multilingual Transformer Base model has 142M parameters, with parameters shared across language pairs, including embeddings and the softmax layer.The model uses a 64k-token vocabulary.Source: §A.1 · p. 12
• A shared 64k-token Sentence Piece vocabulary is learned using temperature-based sampling with T = 5.Source: §A.1 · p. 12
• Training uses Adafactor with momentum factorization, norm clipping at 1.0, and 40K warm-up steps; distillation uses a reduced learning rate of 0.2.BLEU scores are calculated with SacreBLEU on WMT test sets.Source: §A.1 · p. 12

A.2 WMT Dataset Details

• The section presents training-set details for the WMT 4 setup used in the experiments.Source: §A.2 · p. 12
• It reports data sizes for the training setup.Source: §A.2 · p. 12
• It identifies the WMT years used for the Train, Dev, and Test sets.Source: §A.2 · p. 12

A.3 Individual WMT BLEU Scores

• Tables 5 and 6 provide individual BLEU scores for the models reported in Table 1.Source: §A.3 · p. 12, Table 5 · p. 21, Table 6 · p. 21

A.7 Gating Decisions for task-level and token-level MoEs

• Position-wise MoE figures show broadly distributed token routing, with later layers preferring a few experts slightly; decoder figures also show similar distributions for some related language pairs.Source: Figure 6 · p. 14, Figure 7 · p. 15, Figure 8 · p. 16
• In the target-language-wise decoder, some language pairs share experts, while English-Spanish and English-Catalan have different experts.Source: Figure 11 · p. 19

A.4 Detailed Breakdown of Parameter Counts on WMT

• Table 7 breaks down the parameter counts by component for the models compared in Section 4.2.Source: §A.4 · p. 12, Table 7 · p. 22

A.5 Detailed Breakdown of Parameter Counts

• Table 8 breaks down parameter counts by component for the models discussed in Section 5.2.Source: §A.5 · p. 13, Table 8 · p. 22

A.6 Results on Large MoE Model

• Table 9 summarizes aggregate BLEU results for the scaled-up 128-expert MoE models, grouped into High-25, Low-25, and Mid-52 language averages.High-25 and Low-25 are the 25 highest- and lowest-resource languages; Mid-52 covers the remaining 52.Source: §A.6 · p. 13, Table 9 · p. 23