Router Z-Loss

Problem

Compute the router z-loss that keeps routing logits small: for each token compute the stable log-sum-exp of the router logits and average its square over tokens.

The signature is router_z_loss(router_logits). Return a scalar.

Hint

Penalizing the z-loss improves numerical stability and prevents logit blow-up.

Test cases

case 1

in: [[[0,0]]]

expected: 0.4804530139

case 2

in: [[[0,1,2]]]

expected: 5.79656648

case 3

in: [[[-1,0,1]]]

expected: 1.9813545511

case 4

in: [[[0,0],[1,2]]]

expected: 2.9158163244

Starter codePython

import math
def router_z_loss(router_logits):
    # Your code here
    pass

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