Top-k routing: separate a discrete choice from its differentiable weights

Let a token representation produce one finite logit per routed expert. A full softmax gives a probability for every candidate, after which a deterministic top-k rule selects distinct experts. One possible layer output is the sum of selected expert outputs multiplied by their original full-softmax probabilities. These retained weights need not sum to one. Switch uses a selected full-softmax gate for its top-one branch. This is different from selecting logits first and normalizing only the selected entries, a convention supported by other implementations.

Selected-set renormalization divides retained probabilities by their retained sum. For top one, that produces weight one, so an otherwise isolated branch has no task gradient through its gate weight inside a fixed selection region. That observation does not apply to an unnormalized selected full-softmax gate. Auxiliary losses, other differentiable paths, noise, or alternative estimators can provide additional router learning signals. A model description must specify these paths rather than infer all training behavior from the integer k alone.

Away from ties and rank crossings, the selected membership can remain constant under a small logit perturbation. Ordinary differentiation then propagates through the selected gate values and expert computations. The derivative of the hard membership indicator is not a smooth substitute for evaluating the experts that were not executed. In a full-softmax gate, an unselected logit can still receive a gate-related gradient through the denominator, even though its expert weights receive no task update from that unexecuted branch for this token.

For a scalar illustration, hold expert outputs fixed and differentiate only the router logits. With top one selecting expert a, output y=p_a F_a has derivative p_a(1−p_a)F_a with respect to z_a. It also has derivative −p_a p_j F_a with respect to another logit z_j. These are local derivatives of the stated rule. At a membership change, the branch output can jump or change derivative depending on expert values and weighting. Do not validate a local derivative with a finite-difference step that silently crosses that change.

Compare gate conventions with the same logits, expert parameters, token order, and capacity policy before drawing a learning conclusion. Inspect gradients of router weights separately from gradients of expert weights. A zero gradient in one toy configuration may result from renormalization, identical expert outputs, a saturated probability, or a detached tensor. Record selected identities and retained weight sums alongside the loss. For gradient checks, choose well-separated logits, then add explicit boundary tests whose expected result is a selection change rather than a smooth finite difference.

The network lab computes only a binary probability σ(wx−0.5). Reset, set x=0.50 and w=1, predict output 0.5, then Play or Step. Separately Reset, set x=1 and w=1, predict about 0.622, then Play or Step. It has no expert selection, capacity, or backward engine. Interpret its output as one component of logits [0,wx−0.5], and calculate the gate derivative p(1−p) by hand. This isolates a real normalization mechanism without presenting the instrument as a sparse-training simulator.