AI Engineer Exhaustive
MediumPlusgpt-oss

MoE Router

gpt-oss

Medium

Top-k expert routing with softmax taken after the top-k selection, not before.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Each gpt-oss block replaces the dense MLP with a token-choice Mixture-of-Experts. Section 2.2: a linear router maps the residual (after RMSNorm) to one score per expert; both released sizes pick the top (4) experts; the mixture weights are “the softmax of the router projection over only the selected experts.” That last clause is the whole point of this note. Softmax-then-top-k and top-k-then-softmax are different functions and yield different weights on the same logits.

gpt-oss-120b has 128 experts; gpt-oss-20b has 32. (k=4) in both, so a token activates (4/128) or (4/32) of the expert parameters. The card does not mention a load-balancing auxiliary loss at inference; the reference MLPBlock is inference-only and has none.

How it works

Let (\tilde{x}=\mathrm{RMSNorm}(x)\in\mathbb{R}^{T\times d}) with (d=2880). The router is an nn.Linear gate: (d\to E), so

g_t = W_r\tilde{x}_t + b_r \in \mathbb{R}^{E}.

(E=128) or (32). For each token, take the top (k=4) logits, keeping them sorted:

\mathcal{T}_t,\; z_t = \mathrm{TopK}(g_t,\,k=4).

(\mathcal{T}_t) is the index tuple, (z_t\in\mathbb{R}^{4}) the corresponding logits. Mixture weights are

w_t = \mathrm{softmax}(z_t),\qquad w_{t,i}=\frac{e^{z_{t,i}}}{\sum_{i'=1}^{k}e^{z_{t,i'}}}.

The (E-k) unselected logits never enter the softmax, so they contribute neither normalization nor a residual “not selected” class. The four weights on a token always sum to 1. If two selected logits are equal they share mass equally; if one selected logit dominates, its weight approaches 1 and the other three approach 0.

Contrast the rejected order: (\pi=\mathrm{softmax}(g_t)) over all (E), then keep (\pi_{\mathcal{T}_t}). Those four numbers generally sum to less than 1 (the discarded experts still hold probability). Using them un-normalized under-scales the MoE output; re-normalizing them is not the same as softmax-on-selected, because softmax-on-all has already mixed the discarded logits into the denominator.

The official call is torch.topk(g, k=experts_per_token, dim=-1, sorted=True) then softmax(experts.values, dim=1). For a rank-2 activation [T, E], dim=1 is the expert-of-the-k axis. Ties in topk follow PyTorch’s stable-by-value-then-index convention; the paper is silent on ties.

The router sees the normalized residual, not the raw skip stream. It does not see token id, position, or layer index. The same four-expert rule applies on every layer; there is no shared expert on the side (unlike DeepSeek-style MoE).

Official code

gpt_oss/torch/model.py: MLPBlock holds self.gate and, in forward, torch.topk then F.softmax on experts.values. ModelConfig.experts_per_token = 4, num_experts = 128 by default (120b). The educational path upcasts expert weights from MXFP4 to BF16 after load; the router itself stays BF16.

Watch-outs

Sources