MoE Aux Load-Balance Loss

Problem

Compute the Switch-Transformer auxiliary loss: num_experts * sum over experts of (fraction of tokens routed to the expert) * (mean routing probability of the expert), using the top-1 expert per token.

The signature is moe_aux_loss(router_probs, expert_indices). Each row of router_probs is a distribution over experts.

Hint

The auxiliary loss is minimized at 1.0 when routing is perfectly uniform.

Test cases

case 1

in: [[[0.5,0.5],[0.5,0.5]],[0,1]]

expected: 1

case 2

in: [[[0.9,0.1],[0.8,0.2]],[0,0]]

expected: 1.7

case 3

in: [[[0.25,0.25,0.25,0.25],[0.25,0.25,0.25,0.25],[0.1,0.2,0.3,0.4]],[1,2,3]]

expected: 1.0666666667

case 4

in: [[[0.6,0.4]],[1]]

expected: 0.8

Starter codePython

def moe_aux_loss(router_probs, expert_indices):
    # Your code here
    pass

4 test cases · run your code against them in the editor