Skip to content

← Flax step 56 of 100

Hard Primitives

Mixture-of-Experts FFN

Jump to the task

Let each token choose expert outputs

A mixture-of-experts feed-forward layer contains several independent networks and a router that scores them for each token. The output combines the selected experts' results using weights derived from those scores.

This exercise evaluates every expert on every token, then selects the outputs to combine. It teaches routing and weighted aggregation. It does not implement the token dispatch needed to avoid computing unused experts.

Compute the experts and router

Let x have shape (T, D), E = num_experts, F = d_ff and K = top_k. Each expert is a separate Dense(F) -> ReLU -> Dense(D) network with its own weights and biases. Stacking their outputs on axis 1 gives expert_outs of shape (T, E, D).

A separate Dense(E) router produces (T, E) logits. Apply softmax over the expert axis to get probabilities, then call jax.lax.top_k(probs, K). It returns the selected probabilities and their expert indices, each with shape (T, K), in descending score order. Equal scores favor lower expert indices.

Each token makes its own selection. Keep the token axis when gathering; one token's expert indices must not select another token's outputs.

Weight the selected outputs

Use the same normalization as the reference:

weights = selected_probs / (sum(selected_probs, axis=-1, keepdims=True) + 1e-9)
out[t] = sum(weights[t, k] * expert_outs[t, selected_indices[t, k]])

The sum in the denominator is over the selected experts for each token. The weights are nonnegative and sum approximately to one. The 1e-9 term and floating-point rounding mean exact unit sums are not promised. Do not divide raw router logits by their sum; logits can be negative and are not probabilities.

K=1 selects one expert per token, with weight approximately one. K=E combines every expert. In both cases, this implementation still evaluates all E experts.

Work through two tokens

Here the expert outputs and routing probabilities are supplied directly so you can check the gather and weighted sum without random initialization:

import jax
import jax.numpy as jnp

probs = jnp.array([[.6, .3, .1], [.2, .2, .6]])
expert_outs = jnp.array([
    [[1., 0.], [0., 3.], [10., 10.]],
    [[4., 0.], [0., 8.], [0., 2.]],
])
values, indices = jax.lax.top_k(probs, 2)
weights = values / (values.sum(axis=-1, keepdims=True) + 1e-9)
chosen = jnp.take_along_axis(expert_outs, indices[:, :, None], axis=1)
out = (chosen * weights[:, :, None]).sum(axis=1)

assert jnp.array_equal(indices, jnp.array([[0, 1], [2, 0]]))
assert chosen.shape == (2, 2, 2)
assert jnp.allclose(out, jnp.array([[2/3, 1.], [1., 1.5]]), atol=1e-6)

For the second token, experts 0 and 1 tie for the remaining slot; expert 0 wins. The unselected expert output does not enter the final sum.

Selection does not imply sparse computation

Each expert has 2*D*F + F + D parameters. The router adds D*E + E, for a total of E*(2*D*F + F + D) + D*E + E. Changing K does not change that parameter count or the number of expert forward passes here. The intermediate expert-output array contains T*E*D values.

A sparse implementation can send tokens only to selected experts. Its actual speed and memory use also depend on routing, dispatch and hardware; this dense exercise establishes no inference speedup. It has no capacity limit, token dropping, load-balancing loss or training step.

Problem

Implement moe_ffn_forward(seed, x, num_experts, d_ff, top_k):

  1. Define Expert(nn.Module) with a d_ff field. Inside its compact call, apply unnamed nn.Dense(F), nn.relu, then nn.Dense(D).
  2. Define MoEFFN(nn.Module) with integer num_experts, d_ff and top_k fields. Create E separate unnamed Expert instances, in order, and stack their outputs on axis 1.
  3. Apply an unnamed nn.Dense(E) router to x, followed by softmax over the final axis. Select K probabilities with jax.lax.top_k and normalize using the denominator above, including 1e-9.
  4. Gather on expert axis 1 with indices of shape (T, K, 1). Multiply by weights of shape (T, K, 1) and sum over axis 1.
  5. Initialize the module directly with jax.random.PRNGKey(int(seed)) and x, apply the returned variables to x, and flatten the output.

Keep all Dense defaults, including biases. Preserve scopes Expert_0 through Expert_{E-1}, each with Dense_0 and Dense_1, and the router's root Dense_0. These names affect the seeded initialization.

x is a float32 array of shape (T, D) with positive dimensions. E and F are positive, and 1 <= K <= E. Return a one-dimensional array of length T * D, with features varying fastest. Keep seed, num_experts, d_ff and top_k static if JIT-compiling the whole function.