30

A TOKEN LOAD BALANCE DESCRIPTION

The auxiliary load balancing loss from Shazeer et al. (2017) is also used to here to balance tokens
across experts. Assume we have N experts indexed by i = 1 to N and a batch B with T tokens. The
auxiliary loss is computed as the scaled dot-product between vectors f and P,

i=1
where fi is the fraction of tokens dispatched to expert i,

and Pi is the fraction of the router probability allocated for expert i, 2

loss = α · N · N(cid:88)
(cid:88)

fi · Pi

fi =

1
T

x∈B

1{argmax p(x), i}

(cid:88)

x∈B

1
T

Pi =

pi(x)

(7)

(8)

(9)

constant as the number of experts varies since under uniform routing(cid:80)N
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS