core/nn/moe library

Mixture-of-Experts feed-forward block.

DeepSeek-V3-style routing: for each token the router picks the top-K of numRoutedExperts "sparse" experts and combines their outputs weighted by the gating scores. numShared always-on shared experts are added on top (unweighted). The combined output is returned in place of a dense FFN block.

The gating function is configurable via gateFunction:

  • GateFunction.softmax (default) — classical softmax router.
  • GateFunction.sigmoid — per-expert independent sigmoid gate, as used in the DeepSeek-V3 paper and the aux-loss-free paper (Wang et al. 2024, arXiv:2408.15664). The paper reports sigmoid outperforms softmax under equal load balance.

If renormalizeTopK is true the K selected weights are divided by their per-token sum so each token's expert contributions sum to 1. Mixtral does this unconditionally; DeepSeek-V3 does it for sigmoid. Default is true for sigmoid and false for softmax (matches both references).

Load balancing follows the aux-loss-free recipe: per-expert additive biases nudge the top-K comparison so under-utilized experts win more often. The biases are non-differentiable and are added ONLY to the top-K sort key, never to the weights used to combine expert outputs. Call MoEFeedForward.updateRoutingBias once per training batch (Algorithm 1 of the paper) — the counter is per-batch, not per-epoch.

biasUpdateRule selects between:

  • BiasUpdateRule.sign (default, paper's main variant): b_i += u * sign(mean_load - load_i). Better perplexity.
  • BiasUpdateRule.proportional: b_i += u * (mean_load - load_i) / mean_load. Slightly better load balance but slightly worse perplexity per §4.3 of the paper.

Design notes:

  • The router (gateW) and both routed / shared experts are ordinary trainable parameters — gradient flows through the gating function into gateW, and through the expert forward passes into their weights.
  • The discrete top-K choice is made on CPU (Dart) and applied as a [T, E] 0/1 mask multiplied into gateScores. The mask is a non-differentiable straight-through stop; only the K selected scores per row contribute to the output. When the module lives on GPU this incurs one small [T, E] device→host copy per forward — negligible for typical sequence lengths.
  • Column-broadcasting of masked_scores[:, e] to [T, embedDim] uses a precomputed one-hot selector matmul ([T, E] @ [E, D]) — this keeps everything in existing tensor ops with correct autograd, and runs on the same device as the input.
  • All experts are evaluated on all tokens (dense compute). Real sparse-MoE implementations (DeepSeek-V3, Mixtral) only run the top-K experts per token via gather/scatter; that is a much bigger rework and is not done here.
  • 2D input [T, embedDim] only. Higher-rank inputs should be reshaped by the caller. Runs on CPU or GPU (matches the device passed to the constructor).

Classes

Expert
A single expert. Either a two-layer MLP (variant = mlp, default) or a SwiGLU gated FFN (variant = swiGlu), matching DeepSeek-V3 and Mixtral.
MoEFeedForward

Enums

BiasUpdateRule
Bias update rule for MoEFeedForward.updateRoutingBias. See aux-loss-free paper (arXiv:2408.15664) §4.3.
ExpertActivation
Activation used inside an Expert. relu and silu (x * sigmoid(x)) both have fwd+bwd on CPU and GPU.
ExpertVariant
Feed-forward body inside an Expert.
GateFunction
Gating function used by MoEFeedForward to turn router logits into per-expert scores.