Problems
Loading...
Step 7 of 11
Stable LatentMoE
Hard

Stable LatentMoE keeps common transformations at model width while running specialized routed experts in a smaller latent width. For one token x∈Rdx\in\mathbb{R}^{d}, the down-projection maps model width dd to latent width ℓ\ell, the up-projection maps ℓ\ell back to dd, and the router produces one score for each of the EE routed experts.

Let WrW_r be the router projection, bb the current selection bias, and kk the number of selected routed experts. Compute raw router scores

r=sigmoid⁡(Wrx).r=\operatorname{sigmoid}(W_rx).

Select expert indices with the biased scores

Tk(x)=TopK⁡(r+b,k).\mathcal{T}_k(x)=\operatorname{TopK}(r+b,k).

The bias changes selection only. Normalize the selected raw scores to obtain mixture weights

pi=ri∑j∈Tk(x)rj,i∈Tk(x).p_i=\frac{r_i}{\sum_{j\in\mathcal{T}_k(x)}r_j},\qquad i\in\mathcal{T}_k(x).

For latent representation z=W↓xz=W_{\downarrow}x and routed expert Eirouted:Rℓ→RℓE_i^{\mathrm{routed}}:\mathbb{R}^{\ell}\rightarrow\mathbb{R}^{\ell}, aggregate the selected expert outputs as

u=∑i∈Tk(x)piEirouted(z).u=\sum_{i\in\mathcal{T}_k(x)}p_iE_i^{\mathrm{routed}}(z).

Kimi K3 uses exactly two shared experts at model width. With shared experts E1sharedE_1^{\mathrm{shared}} and E2sharedE_2^{\mathrm{shared}}, the final output is

y=E1shared(x)+E2shared(x)+W↑RMSNorm⁡(u).y=E_1^{\mathrm{shared}}(x)+E_2^{\mathrm{shared}}(x)+W_{\uparrow}\operatorname{RMSNorm}(u).

Every shared and routed expert uses the SiTU-GLU gate, up, and down projections. Apply the computation independently to each token row. Return a dictionary with output, selected_experts, mixture_weights, and latent_routed_aggregate. The floating tensors preserve the token dtype and device, while selected_experts uses integer dtype.

Loading visualization...

Examples

Input: tokens = [[1]], latent_down_projection = [[1]], latent_up_projection = [[1]], router_projection = [[1],[-1]], current_bias = [0,0], routed_gate_weights = [[[1]],[[0.5]]], routed_up_weights = [[[1]],[[1.5]]], routed_down_weights = [[[1]],[[0.75]]], shared_gate_weights = [[[0.5]],[[-0.5]]], shared_up_weights = [[[1]],[[1.5]]], shared_down_weights = [[[1]],[[0.5]]], selected_count = 1, eps = 1e-06, gate_cap = 4, up_cap = 25

Output: {"output": tensor([[1.1687765696141315]]), "selected_experts": tensor([[0]]), "mixture_weights": tensor([[1.0]]), "latent_routed_aggregate": tensor([[0.7158178280748405]])}

Explanation: The router selects expert 0 with weight 1. The returned values expose the final output and the routed latent value used before normalization.

Input: tokens shape (2, 2), latent width = 1, routed experts = 3, shared experts = 2, selected count = 2

Output: {"output": tensor of shape (2, 2), "selected_experts": tensor of shape (2, 2), "mixture_weights": tensor of shape (2, 2), "latent_routed_aggregate": tensor of shape (2, 1)}

Input: tokens shape (1, 2), latent width = 1, routed experts = 3, shared experts = 2, selected count = 1

Output: {"output": tensor of shape (1, 2), "selected_experts": tensor of shape (1, 1), "mixture_weights": tensor of shape (1, 1), "latent_routed_aggregate": tensor of shape (1, 1)}

Hint 1

Choose routed experts with biased scores, but normalize their mixture weights from raw sigmoid scores.

Hint 2

Run routed experts on the latent representation and shared experts directly on the full-width tokens.

Hint 3

Normalize the routed aggregate before its up-projection, then add both shared-expert outputs.

Requirements

  • Return a dictionary with exactly output, selected_experts, mixture_weights, and latent_routed_aggregate.
  • Final output has the same shape, dtype, and device as the token tensor.
  • Selected mixture weights sum to one per token.
  • Both shared experts contribute directly at model width.

Constraints

  • Exactly two shared experts are supplied.
  • Selected count is positive and no larger than the routed expert count.
  • Routed experts operate entirely in latent width.
  • All projection and expert weights have compatible dimensions.
Prev7 / 11Next