Study Plan

Top-k MoE Router with Load Statistics

Language ModelsNeural Networks

Hard

Implement top-k mixture-of-experts routing from precomputed expert outputs. For token $t$ and expert $e$ , $x_t$ is the token vector, $w_e$ is column e of w_router, and $s_{t,e}$ is the router score. Superscript T denotes transpose.

s_{t,e}=x_t^{\mathsf T}w_e.

Select k experts by descending score, breaking ties by lower expert index. Let $e_{t,j}$ be the expert at selected position j and $\alpha_{t,j}$ its weight; j and r range over the selected experts.

\alpha_{t,j}=\frac{\exp(s_{t,e_{t,j}})}{\sum_{r=1}^{k}\exp(s_{t,e_{t,r}})}.

For expert output $u_{t,e}$ and shared output $u_t^{\mathrm{shared}}$ , the routed result $y_t$ is

y_t=u_t^{\mathrm{shared}}+\sum_{j=1}^{k}\alpha_{t,j}u_{t,e_{t,j}}.

For T tokens, $f_e$ is the fraction selecting expert e and $m_e$ is its mean probability under softmax over all E experts. The indicator is one when its condition holds and zero otherwise.

f_e=\frac{1}{T}\sum_{t=1}^{T}\mathbf{1}[e\text{ is selected for }t].

m_e=\frac{1}{T}\sum_{t=1}^{T}\operatorname{softmax}(s_t)_e.

Return a dictionary of PyTorch tensors: output with shape (T, Dout), selected_experts and selected_weights with shape (T, k), and token_fraction and probability_mass with shape (E,). These contain y, the selected zero-based indices, their weights, f, and m, respectively. selected_experts is int64; all other results match x's dtype, and all results stay on x's device.

Loading visualization...

Examples

Input: x = [[1,0],[0,1]], w_router = [[1,0,1],[0,1,1]], expert_outputs = [[[10,0],[0,0],[0,10]],[[0,0],[5,5],[1,1]]], shared_output = [[1,1],[2,2]], k = 2

Output: {"output":[[6,6],[5,5]],"selected_experts":[[0,2],[1,2]],"selected_weights":[[0.5,0.5],[0.5,0.5]],"token_fraction":[0.5,0.5,1],"probability_mass":[0.288841,0.288841,0.422319]}

Explanation: Both tokens tie their own specialist with expert 2, so stable index order selects [0,2] and [1,2]. Equal selected weights average those outputs before the shared output is added.

Input: x = [[2]], w_router = [[1,0.5]], expert_outputs = [[[7],[3]]], shared_output = [[1]], k = 1

Output: {"output":[[8]],"selected_experts":[[0]],"selected_weights":[[1]],"token_fraction":[1,0],"probability_mass":[0.731059,0.268941]}

Input: x = [[0]], w_router = [[0,0]], expert_outputs = [[[2],[4]]], shared_output = [[1]], k = 2

Output: {"output":[[4]],"selected_experts":[[0,1]],"selected_weights":[[0.5,0.5]],"token_fraction":[1,1],"probability_mass":[0.5,0.5]}

Hint 1

Stable sorting preserves the original index order for tied scores.

Hint 2

torch.gather selects values using per-token indices.

Hint 3

torch.scatter can mark selected experts in a Boolean mask.

Requirements

Use PyTorch with float32 routing calculations, then cast floating-point results to the input dtype.
Keep selected experts and their weights in descending score order, with lower-index tie breaking.
Do not modify any input tensor.

Constraints

x has shape (T, D), w_router has shape (D, E), expert_outputs has shape (T, E, Dout), and shared_output has shape (T, Dout).
T is token count, D is input width, E is expert count, and Dout is output width; all are positive.
k is a Python integer from 1 through E.
All tensor inputs are PyTorch tensors sharing a device and dtype, either float32 or bfloat16.
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

Gated DeltaNet State UpdateHard Parallel and Recurrent Linear AttentionMedium Mamba 2 Gated State ScanMedium RMSNorm Forward PassEasy Decode Attention with a KV CacheHard Recover an IsoFLOP Compute-Optimal FrontierHard Fit a Data Power Law with a Noise FloorHard Fit Elo Ratings from Pairwise VotesHard Vision Transformer Patch EmbeddingHard LSTM Cell from ScratchHard Early StoppingHard Backpropagation Through a Computation GraphHard

Top-k MoE Router with Load Statistics

Language ModelsNeural Networks

Hard

s_{t,e}=x_t^{\mathsf T}w_e.

Select k experts by descending score, breaking ties by lower expert index. Let $e_{t,j}$ be the expert at selected position j and $\alpha_{t,j}$ its weight; j and r range over the selected experts.

\alpha_{t,j}=\frac{\exp(s_{t,e_{t,j}})}{\sum_{r=1}^{k}\exp(s_{t,e_{t,r}})}.

For expert output $u_{t,e}$ and shared output $u_t^{\mathrm{shared}}$ , the routed result $y_t$ is

y_t=u_t^{\mathrm{shared}}+\sum_{j=1}^{k}\alpha_{t,j}u_{t,e_{t,j}}.

For T tokens, $f_e$ is the fraction selecting expert e and $m_e$ is its mean probability under softmax over all E experts. The indicator is one when its condition holds and zero otherwise.

f_e=\frac{1}{T}\sum_{t=1}^{T}\mathbf{1}[e\text{ is selected for }t].

m_e=\frac{1}{T}\sum_{t=1}^{T}\operatorname{softmax}(s_t)_e.

Loading visualization...

Examples

Input: x = [[1,0],[0,1]], w_router = [[1,0,1],[0,1,1]], expert_outputs = [[[10,0],[0,0],[0,10]],[[0,0],[5,5],[1,1]]], shared_output = [[1,1],[2,2]], k = 2

Output: {"output":[[6,6],[5,5]],"selected_experts":[[0,2],[1,2]],"selected_weights":[[0.5,0.5],[0.5,0.5]],"token_fraction":[0.5,0.5,1],"probability_mass":[0.288841,0.288841,0.422319]}

Explanation: Both tokens tie their own specialist with expert 2, so stable index order selects [0,2] and [1,2]. Equal selected weights average those outputs before the shared output is added.

Input: x = [[2]], w_router = [[1,0.5]], expert_outputs = [[[7],[3]]], shared_output = [[1]], k = 1

Output: {"output":[[8]],"selected_experts":[[0]],"selected_weights":[[1]],"token_fraction":[1,0],"probability_mass":[0.731059,0.268941]}

Input: x = [[0]], w_router = [[0,0]], expert_outputs = [[[2],[4]]], shared_output = [[1]], k = 2

Output: {"output":[[4]],"selected_experts":[[0,1]],"selected_weights":[[0.5,0.5]],"token_fraction":[1,1],"probability_mass":[0.5,0.5]}

Hint 1

Stable sorting preserves the original index order for tied scores.

Hint 2

torch.gather selects values using per-token indices.

Hint 3

torch.scatter can mark selected experts in a Boolean mask.

Requirements

Use PyTorch with float32 routing calculations, then cast floating-point results to the input dtype.
Keep selected experts and their weights in descending score order, with lower-index tie breaking.
Do not modify any input tensor.

Constraints

x has shape (T, D), w_router has shape (D, E), expert_outputs has shape (T, E, Dout), and shared_output has shape (T, Dout).
T is token count, D is input width, E is expert count, and Dout is output width; all are positive.
k is a Python integer from 1 through E.
All tensor inputs are PyTorch tensors sharing a device and dtype, either float32 or bfloat16.
Inputs contain finite values and may be noncontiguous.

Try Similar Problems

x =

Accepts: array

w_router =

Accepts: array

expert_outputs =

Accepts: array

shared_output =

Accepts: array

k =

Accepts: number