[QST] MoE auxiliary loss

[osayamenja]

Hello!

My understanding is that the gate layer implements Algorithm 1 of GShard; however, our auxiliary loss computation seems to deviate from the algorithm; please help me understand.

That is, we compute l_aux here by sum(me * ce) * num_experts, while line 13 of the algorithm specifies mean(me * ce) or, equivalently, sum(me * ce) / num_experts.

Why do we deviate there?