Train CF

Step Time (s) (↓)

Model
ST-MoE-L
ST-MoE-L
ST-MoE-32B
ST-MoE-32B

1.25
2.0
1.25
2.0

2.397
2.447 (+7%)
4.244
4.819 (+14%)

Table 9: Proﬁling sparse models on TPUs. Increasing the train capacity factor from 1.25 to 2.0
increases the step-time by +7% for the large (1B) model but by +14% for our 32B model. As the
model size increases, we ﬁnd the small quality gains of higher train capacity factors from Table 8
are more than offset by the signiﬁcant 14% slow-down. Note: the step time between ST-MoE-L and
ST-MoE-32B are not comparable because they used a different number of cores.

Our results in this section focus on top-n routing, but we also experimented with a variety of other
routing techniques in Appendix J. We found most performed similarity or worse compared to top-n
routing. However we found Batch Prioritized Routing (BPR), introduced in Riquelme et al. (2021),
signiﬁcantly helps performance for capacity factors less than one (Appendix D). We recommend

15
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS