Mixture-of-experts top-k routing with auxiliary load-balance loss

Router and combination per Lepikhin et al. (2021) GShard, ICLR 2021, arXiv:2006.16668, Section 2.1/Algorithm 1: G(x) = softmax(W_g x), each token dispatched to its top-k experts, output the renormalised gate-weighted sum. Auxiliary loss per Fedus, Zoph and Shazeer (2022), Switch Transformers, JMLR 23(120), arXiv:2101.03961, Eqs 4-6: aux = alpha N sum_i f_i P_i; uniform routing gives exactly alpha (the anchor).

Usage

Moetop(x, W_g, experts, k = 2L, alpha = 0.01)

Arguments

x: Token batch (T x d_in).
W_g: Router weights (d_in x N).
experts: List of N linear expert maps (d_in x d_out each).
k: Experts per token (default 2, GShard).
alpha: Aux-loss coefficient (default 0.01, Switch).

Value

List with output, gates, topk_indices (1-based), topk_gates, aux_loss, f, P, n_experts, k, estimate, n, method.

References

Lepikhin, D. et al. (2021), arXiv:2006.16668, Sec 2.1; Fedus, W., Zoph, B. and Shazeer, N. (2022), JMLR 23(120), arXiv:2101.03961, Eqs 4-6.

Examples

set.seed(6)
x <- matrix(rnorm(8 * 4), 8, 4)
experts <- lapply(1:3, function(i) matrix(rnorm(4 * 2), 4, 2))
r <- Moetop(x, W_g = matrix(rnorm(4 * 3), 4, 3), experts = experts, k = 2L)
r$aux_loss
#> [1] 0.01128073