log

N(cid:88)

j=1

B(cid:88)

i=1

1
B

2

ex(i)

j

Lz(x) =

(5)

where B is the number of tokens, N is the number of experts, and x ∈ RB×N are the logits going
into the router. This penalizes large logits into the gating network and Section 3.4 contains a more
detailed explanation of why the z-loss before the router is useful.

7

Table 4 shows that both update clipping and the router z-loss stabilize the model in all 3 runs, but the
update clipping signiﬁcantly hurts the model quality. Therefore we use the z-loss method for ﬁxing
our model stability due to improved quality and stability4.

Method
Baseline

Update clipping (clip = 0.1)

Router Z-Loss

Fraction Stable

4/6
3/3
3/3

Quality (↑)
-1.755 ±0.02
-4.206 ±0.17
-1.741 ±0.02
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS