Mixture of experts

The facts behind the choices in NN.Spec.Layers.MoE, in four groups.

• Routing keeps exactly min k E experts, and every kept expert beats every expert it displaces (card_kept, beats_of_kept), whenever the comparison of the scores is a strict weak order (StrictWeak). Real numbers satisfy this, and so do IEEE numbers other than NaN, on which +0 and -0 tie.
• Which experts are kept depends only on how the scores compare (kept_congr). Over ℝ the softmax preserves every comparison, so routing on probabilities, as Mixtral does, keeps the experts with the largest logits (kept_softmax), and the gates are the softmax of the kept logits (getScalar_gatesVecSpec_softmax). Rounding breaks this with floats: in binary32 the logits 0 and 2⁻³⁰ have the same probability, so routing on probabilities keeps expert 0 and routing on logits keeps expert 1.
• Over ℝ the dense layer, which runs every expert, equals Mixtral's sparse block, which runs only the kept ones (forwardSpec_eq_sparseForwardSpec).
• A token's output depends only on that token, so a batch can be split across devices anywhere. A capacity per device, as in GShard and Switch, breaks this: two tokens routed to the same expert are both kept on two devices, but on one device the second is dropped (capacity_depends_on_devices). So does a balancing loss computed on each device's tokens (balance_depends_on_devices). Computed per work unit, a fixed set of tokens, both are functions of the unit again.

Routing

theorem NN.Proofs.Models.MoE.gtBool_eq_true {α : Type} [Context α] {x y : α} : Context.gtBool x y = true ↔ x > y

theorem NN.Proofs.Models.MoE.gtBool_eq_false {α : Type} [Context α] {x y : α} : Context.gtBool x y = false ↔ ¬x > y

theorem NN.Proofs.Models.MoE.beats_eq_true {α : Type} [Context α] {E : ℕ} (scores : Fin E → α) (j i : Fin E) : Spec.MoE.beats scores j i = true ↔ scores j > scores i ∨ j < i ∧ ¬scores i > scores j

structure NN.Proofs.Models.MoE.StrictWeak {α : Type} [Context α] {E : ℕ} (scores : Fin E → α) : Prop

The comparison of scores is a strict weak order: no score is above itself, > is transitive, and whenever a > c every other score is below a or above c.

• irrefl (i : Fin E) : ¬scores i > scores i
• trans (a b c : Fin E) : scores a > scores b → scores b > scores c → scores a > scores c
• split (a b c : Fin E) : scores a > scores c → scores a > scores b ∨ scores b > scores c

theorem NN.Proofs.Models.MoE.strictWeak_real {E : ℕ} (scores : Fin E → ℝ) : StrictWeak scores

Real scores are compared by a strict weak order.

theorem NN.Proofs.Models.MoE.beats_self {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) (i : Fin E) : Spec.MoE.beats scores i i = false

theorem NN.Proofs.Models.MoE.beats_total {α : Type} [Context α] {E : ℕ} (scores : Fin E → α) {i j : Fin E} (hij : i ≠ j) : Spec.MoE.beats scores i j = true ∨ Spec.MoE.beats scores j i = true

Of two different experts, one beats the other.

theorem NN.Proofs.Models.MoE.beats_trans {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) {a b c : Fin E} (hab : Spec.MoE.beats scores a b = true) (hbc : Spec.MoE.beats scores b c = true) : Spec.MoE.beats scores a c = true

theorem NN.Proofs.Models.MoE.rank_lt {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) (i : Fin E) : Spec.MoE.rank scores i < E

theorem NN.Proofs.Models.MoE.rank_lt_rank {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) {i i' : Fin E} (hb : Spec.MoE.beats scores i i' = true) : Spec.MoE.rank scores i < Spec.MoE.rank scores i'

An expert that beats another has the smaller rank.

theorem NN.Proofs.Models.MoE.rank_injective {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) : Function.Injective (Spec.MoE.rank scores)

Different experts have different ranks, so the ranks are 0, …, E - 1 in some order.

theorem NN.Proofs.Models.MoE.card_kept {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) (k : ℕ) : {i : Fin E | Spec.MoE.kept k scores i = true}.card = min k E

Routing keeps exactly min k E experts.

theorem NN.Proofs.Models.MoE.length_keptExperts {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) (k : ℕ) : (Spec.MoE.keptExperts k scores).length = min k E

The list of kept experts has min k E entries.

theorem NN.Proofs.Models.MoE.beats_of_kept {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} (h : StrictWeak scores) {k : ℕ} {i i' : Fin E} (hi : Spec.MoE.kept k scores i = true) (hi' : Spec.MoE.kept k scores i' = false) : Spec.MoE.beats scores i i' = true

Every kept expert beats every expert that is not kept.

theorem NN.Proofs.Models.MoE.kept_dominates {E : ℕ} {scores : Fin E → ℝ} {k : ℕ} {i i' : Fin E} (hi : Spec.MoE.kept k scores i = true) (hi' : Spec.MoE.kept k scores i' = false) : scores i' < scores i ∨ scores i' = scores i ∧ i < i'

Over ℝ: a kept expert's score is at least that of every expert it displaces, and larger unless the displaced expert has the larger index.

theorem NN.Proofs.Models.MoE.kept_congr {α : Type} [Context α] {E : ℕ} {scores : Fin E → α} {β : Type} [Context β] {t : Fin E → β} (h : ∀ (i j : Fin E), t i > t j ↔ scores i > scores j) (k : ℕ) : Spec.MoE.kept k t = Spec.MoE.kept k scores

Two score vectors that compare alike keep the same experts.

Gates over the reals

theorem NN.Proofs.Models.MoE.kept_softmax {E : ℕ} (k : ℕ) (logits : TorchLean.Tensor ℝ [E]) : (Spec.MoE.kept k fun (e : Fin E) => (Activation.softmaxVecSpec logits).getScalar e) = Spec.MoE.kept k fun (e : Fin E) => logits.getScalar e

Over ℝ the softmax preserves every comparison, so routing on the probabilities keeps the experts with the largest logits.

@[reducible, inline]

noncomputable abbrev NN.Proofs.Models.MoE.keptSet {E : ℕ} (k : ℕ) (scores : Fin E → ℝ) : Finset (Fin E)

The kept experts as a finite set.

theorem NN.Proofs.Models.MoE.getScalar_gatesVecSpec {E : ℕ} (k : ℕ) (logits : TorchLean.Tensor ℝ [E]) (e : Fin E) : have p := fun (j : Fin E) => (Activation.softmaxVecSpec logits).getScalar j; (Spec.MoE.gatesVecSpec k logits).getScalar e = if Spec.MoE.kept k p e = true then p e / ∑ j ∈ keptSet k p, p j else 0

Over ℝ, the gate of expert e is its probability divided by the total probability of the kept experts, or zero when e is not kept.

theorem NN.Proofs.Models.MoE.getScalar_gatesVecSpec_softmax {E : ℕ} (k : ℕ) (logits : TorchLean.Tensor ℝ [E]) (e : Fin E) : have l := fun (j : Fin E) => logits.getScalar j; (Spec.MoE.gatesVecSpec k logits).getScalar e = if Spec.MoE.kept k l e = true then Real.exp (l e) / ∑ j ∈ keptSet k l, Real.exp (l j) else 0

Over ℝ, the gates are the softmax of the logits of the kept experts, the experts with the largest logits.

theorem NN.Proofs.Models.MoE.gatesVecSpec_simplex {E : ℕ} (k : ℕ) (hk : 0 < k) (hE : 0 < E) (logits : TorchLean.Tensor ℝ [E]) : (∀ (e : Fin E), 0 ≤ (Spec.MoE.gatesVecSpec k logits).getScalar e) ∧ ∑ e : Fin E, (Spec.MoE.gatesVecSpec k logits).getScalar e = 1

Over ℝ, with at least one expert kept, the gates are nonnegative and sum to one.

The dense layer is the sparse block

theorem NN.Proofs.Models.MoE.combineSpec_apply {E : ℕ} {s : Spec.Shape} (gates : TorchLean.Tensor ℝ [E]) (outputs : Fin E → TorchLean.Tensor ℝ s) (c : s.Coord) : TorchLean.Tensor.Internal.Rep.get (Spec.MoE.combineSpec gates outputs) c = ∑ e : Fin E, gates.getScalar e * TorchLean.Tensor.Internal.Rep.get (outputs e) c

One coordinate of the combined output, over ℝ.

theorem NN.Proofs.Models.MoE.forwardSpec_eq_sparseForwardSpec {E : ℕ} {s σ : Spec.Shape} (k : ℕ) (router : TorchLean.Tensor ℝ σ → TorchLean.Tensor ℝ [E]) (experts : Fin E → TorchLean.Tensor ℝ σ → TorchLean.Tensor ℝ s) (x : TorchLean.Tensor ℝ σ) : Spec.MoE.forwardSpec k router experts x = Spec.MoE.sparseForwardSpec k router experts x

Over ℝ, the dense layer computes exactly Mixtral's sparse block.

Capacity and balancing

def NN.Proofs.Models.MoE.capacityKeeps {E : ℕ} (C : ℕ) (group : List (Fin E)) : Array Bool

Which tokens of group, each given by the expert it is routed to, a capacity of C tokens per expert keeps: a token is kept when fewer than C earlier tokens of the group went to its expert.

def NN.Proofs.Models.MoE.switchCapacity (E n : ℕ) : ℕ

The capacity Switch gives each expert for a group of n tokens and E experts with capacity factor one, ⌈n / E⌉.

theorem NN.Proofs.Models.MoE.capacity_depends_on_devices : capacityKeeps (switchCapacity 2 2) [0, 0] = #[true, false] ∧ capacityKeeps (switchCapacity 2 1) [0] ++ capacityKeeps (switchCapacity 2 1) [0] = #[true, true]

Two tokens routed to the first of two experts. On one device the capacity is one, so the second token is dropped; split over two devices, each device has capacity one and keeps its token.

def NN.Proofs.Models.MoE.switchBalance {E : ℕ} (group : List (Fin E × (Fin E → ℚ))) : ℚ

Switch's balancing loss on a group of tokens, each given by its expert and the router's probabilities: E times the sum over experts of the fraction of tokens sent to the expert times its mean probability.

theorem NN.Proofs.Models.MoE.balance_depends_on_devices : have a := (0, ![3 / 4, 1 / 4]); have b := (1, ![1 / 4, 3 / 4]); switchBalance [a, b] = 1 ∧ switchBalance [a] = 3 / 2 ∧ switchBalance [b] = 3 / 2

Two tokens, each sent to the expert it prefers with probability 3/4. On one device the balancing loss is 1; split over two devices, each device's loss is 3/2, and so is their mean.