MoE auxiliary load-balancing loss (Kamath Ch 2)

alpha * N * sum f_i P_i; equals alpha under perfect balance.

Usage

morie_kamath_moe_load_balance_loss(
  fractions,
  gate_means,
  N,
  alpha,
  tol = 1e-06
)

Arguments

fractions

Dispatch fractions.

gate_means

Mean router probs.

N

Expert count.

alpha

Weight.

tol

Sum tolerance.

Value

List with estimate, per_expert, imbalance_ratio.

Examples

morie_kamath_moe_load_balance_loss(c(0.7, 0.3), c(0.6, 0.4),
    2, 0.01)
#> $estimate
#> [1] 0.0108
#> 
#> $loss
#> [1] 0.0108
#> 
#> $per_expert
#> [1] 0.42 0.12
#> 
#> $balanced_floor
#> [1] 0.01
#> 
#> $imbalance_ratio
#> [1] 1.08
#> 
#> $alpha
#> [1] 0.01
#> 
#> $n
#> [1] 2
#> 
#> $method
#> [1] "MoE auxiliary load-balancing loss alpha*N*sum f_i P_i"
#>