Mixture of experts

nn.moe k router experts applies Spec.MoE.forwardSpec to every token of a [tokens, width] input. The router scores the experts, the k experts with the largest softmax probabilities are kept, ties going to the lower index, and the outputs of the kept experts are added in increasing expert index, each weighted by its probability divided by the total probability of the kept experts. Over ℝ this is Mixtral's sparse block (NN.Proofs.Models.MoE).

Every expert runs on every token and the experts that are not kept get gate zero, so all shapes are static. When the router and the experts act on each token separately, as linear layers and MLPs over the token axis do, a token's output depends only on that token, whatever else is in its batch; nn.moe accepts any [tokens, width] networks, and one that mixes tokens, such as attention, gives up that independence.

nn.moeMlp is the same layer for the common case where every expert is a two-layer MLP. It keeps the experts' weights in single tensors and runs them all with one batched product per layer, so a step issues a few dozen operations instead of several per expert, and is the one to train with. nn.moe takes any experts.

def TorchLean.nn.stackBranches {σ τ : Shape} (E : ℕ) : (Fin (E + 1) → Sequential σ τ) → Sequential σ (τ.prependDim (E + 1))

The outputs of E + 1 branches that read the same input, stacked along a new first axis in branch order.

def TorchLean.nn.Internal.moeMix {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {tokens width E : ℕ} (k : ℕ) (logits : Runtime.Autograd.Torch.Ref [tokens, E + 1]) (outputs : Runtime.Autograd.Torch.Ref [E + 1, tokens, width]) : m (Runtime.Autograd.Torch.Ref [tokens, width])

The gating and mixing of nn.moe: from the router's logits [tokens, E + 1] and the experts' outputs stacked as [E + 1, tokens, width], compute Spec.MoE.combineSpec of Spec.MoE.gatesVecSpec k for every token. The total kept probability and the output are both summed in increasing expert index, starting from zero.

def TorchLean.nn.Layer.moe {tokens width outWidth E : ℕ} (k : ℕ) (router : Sequential [tokens, width] [tokens, E + 1]) (experts : Fin (E + 1) → Sequential [tokens, width] [tokens, outWidth]) : Layer [tokens, width] [tokens, outWidth]

The nn.moe layer. Its parameters are the router's followed by each expert's, in order.

def TorchLean.nn.moe {tokens width outWidth E : ℕ} (k : ℕ) (router : Sequential [tokens, width] [tokens, E + 1]) (experts : Fin (E + 1) → Sequential [tokens, width] [tokens, outWidth]) : Sequential [tokens, width] [tokens, outWidth]

A mixture-of-experts layer over the tokens of a [tokens, width] input, as in Mixtral: router scores the experts, the k with the largest softmax probabilities are kept (ties to the lower index), and their outputs are added in increasing expert index, each weighted by its probability divided by the kept total. Spec.MoE.forwardSpec is the layer on one token.

Example:

-- Eight experts, two kept per token, on sequences of 16 tokens of width 32.
def model : nn.Builder (nn.Sequential [16, 32] [16, 32]) := do
  let router ← nn.linear 32 8 (batchShape := [16])
  let experts ← (List.range 8).mapM fun _ =>
    nn.Sequential![nn.linear 32 64 (batchShape := [16]), nn.relu,
      nn.linear 64 32 (batchShape := [16])]
  pure (nn.moe 2 router fun e => experts.getD e.val (nn.Sequential.identity _))

def TorchLean.nn.Internal.moeGates {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {tokens E : ℕ} (k : ℕ) (logits : Runtime.Autograd.Torch.Ref [tokens, E + 1]) : m (Runtime.Autograd.Torch.Ref [tokens, E + 1])

The gates of nn.moe for every token of logits: Spec.MoE.gatesVecSpec k on each row. The kept probabilities of a row are summed by one product with a column of ones, a left fold from zero as sumSpec is.

def TorchLean.nn.Layer.moeMlp {tokens width E : ℕ} (k hidden : ℕ) (activation : Activation.Kind) (routerSeed firstSeed secondSeed : ℕ) : Layer [tokens, width] [tokens, width]

The nn.moeMlp layer with E + 1 experts. Its parameters are the router's weight [E + 1, width] and bias [E + 1], the experts' first weights [E + 1, width, hidden] and biases [E + 1, hidden], and their second weights [E + 1, hidden, width] and biases [E + 1, width].

def TorchLean.nn.moeMlp {tokens width : ℕ} (experts k hidden : ℕ) (activation : Activation.Kind := Activation.Kind.relu) : Builder (Sequential [tokens, width] [tokens, width])

A mixture of experts two-layer MLP experts over the tokens of a [tokens, width] input, with a linear router; k experts are kept per token. The layer computes Spec.MoE.forwardSpec on every token, as nn.moe with the same router and experts does, with all experts batched. experts must be positive.

Example:

-- Eight experts of 32-64-32, two kept per token, on sequences of 16 tokens.
def model : nn.Builder (nn.Sequential [16, 32] [16, 32]) :=
  nn.moeMlp (tokens := 16) (width := 32) (experts := 8) 2 64