Skip to content
Medium Research

Top-K Gating

An expert gate chooses which experts should process each token and how much weight each selected expert receives. Implement top_k_gating(logits, k) to turn scores into sparse routing weights.

Inputs and output

logits is a floating-point tensor of shape (batch, n_experts), with finite values. Each row represents one token and each column identifies an expert. k is an integer with 1 <= k <= n_experts.

Return a tensor of the same shape, dtype and device. Keep the original expert-column order and leave logits unchanged.

Select and normalize

For each row independently:

  1. Select the k largest logits. Break equal-score ties by choosing lower expert indices first.
  2. Apply softmax over the selected scores only.
  3. Put those probabilities in their original columns. Every unselected expert receives exactly zero.

For logits [[5, 3, 1]] and k=2, select experts 0 and 1. Their normalized weights are approximately [[0.880797, 0.119203, 0]]. For [[2, 2, 2]] and k=2, the tie rule gives [[0.5, 0.5, 0]].

Each row sums to one. At k=1, the chosen expert gets weight one; at k=n_experts, the result is ordinary softmax over the full row.

One approach is to replace unselected scores with negative infinity and then apply a numerically stable softmax. Subtract the row maximum before exponentiating. Do not use a finite mask value such as -1e9, which can outrank valid, more negative scores. Do not apply full-row softmax and simply zero entries afterward without renormalizing.