Top-k gate renormalisation (Kamath Ch 2)

Keeps the top-k gates and renormalises them – NOT a masked softmax. selected_experts is 0-BASED and sorted.

Usage

morie_kamath_moe_top_k_gating(gates, k)

Arguments

gates

Non-negative gate scores.

k

Experts kept.

Value

List with weights, selected_experts, kept_mass.

Examples

morie_kamath_moe_top_k_gating(gates = c(1, 2, 3, 4, 5, 6, 7, 8), k = 5L)
#> $weights
#> [1] 0.0000000 0.0000000 0.0000000 0.1333333 0.1666667 0.2000000 0.2333333
#> [8] 0.2666667
#> 
#> $selected_experts
#> [1] 3 4 5 6 7
#> 
#> $kept_mass
#> [1] 0.8333333
#> 
#> $n_active
#> [1] 5
#> 
#> $estimate
#> [1] 0.2666667
#> 
#> $k
#> [1] 5
#> 
#> $n
#> [1] 8
#> 
#> $method
#> [1] "Top-k gate renormalisation (not a masked softmax)"
#>