Softmax top-k MoE router (Kamath Ch 2)

Delegates to morie_kamath_moe_topk_gating (km040) and morie_kamath_moe_output (km039); unselected experts are never run.

Usage

morie_kamath_moe_router_softmax(x, Wr, experts, k)

Arguments

x

Token.

Wr

Router projection.

experts

Expert list.

k

Experts kept.

Value

List with output, gate_weights, selected_experts (0-based).

Examples

x <- c(1, 2, 3, 4)
morie_kamath_moe_router_softmax(c(1, 0), matrix(c(1, 0, 0, 1,
    2, 0), 2), list(function(x) x * 1, function(x) x * 2, function(x) x *
    3), 2)
#> $output
#> [1] 2.462117 0.000000
#> 
#> $gate_weights
#> [1] 0.2689414 0.0000000 0.7310586
#> 
#> $selected_experts
#> [1] 0 2
#> 
#> $n_active
#> [1] 2
#> 
#> $experts_evaluated
#> [1] 2
#> 
#> $estimate
#> [1] 2.462117
#> 
#> $k
#> [1] 2
#> 
#> $n
#> [1] 3
#> 
#> $method
#> [1] "Softmax top-k MoE router (delegates to km040 + km039)"
#>