Sigmoid MoE Router

Arcee Trinity

Easy · 🔒 Plus required

Sigmoid replaces softmax for expert routing

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

Trinity scores experts with an independent sigmoid, not a softmax over the expert axis. Section 2.3 follows Wang et al. (2024a) / DeepSeek-style aux-loss-free routing: sigmoid logits are more stable, Top-K uses score plus a decoupled expert bias, and the bias-free scores are what actually gate expert outputs after a Top-K renormalization.

The router is token-choice: each token picks K_r routed experts. On Large, N_r=256, K_r=4. There is no softmax anywhere in this path. A small constant route scale (2.448 on Large) is applied after renormalization; that scale is in Table 2 and the HF config, not in equations (16)–(18).

How it works

Let \mathbf{u}_t\in\mathbb{R}^{d} be the MoE input (post pre-MLP norm) and \mathbf{e}_i\in\mathbb{R}^{d} the i-th router vector. Equation (16):

s_{i,t}=\sigma(\mathbf{u}_{t}^{\top}\mathbf{e}_{i}),\qquad i\in\{1,\ldots,N_r\}.

AfmoeTokenChoiceRouter is the matrix form: router_logits = gate(u) with gate: d → N_r, no bias, then scores = sigmoid(logits) in fp32.

Selection uses bias, weighting does not. Equations (17)–(18):

g'_{i,t}=\begin{cases}s_{i,t} & \text{if }s_{i,t}+b_{i}\in\mathrm{TopK}(\{s_{j,t}+b_{j}\}_{j=1}^{N_r},K_r)\\ 0 & \text{otherwise,}\end{cases} \qquad g_{i,t}=\frac{g'_{i,t}}{\sum_{j=1}^{N_r}g'_{j,t}}.

Official code: topk(scores + expert_bias), then gather the raw s_{i,t} (not s+b), then divide by their sum plus 10^{-20}, then multiply route_scale. That last multiply is the Table 2 “Route scale.” Paper (18) stops at the simplex; the scale is an extra constant on the gates.

b_i is the SMEBU (Large) or sign-based (Nano/Mini) load-balance bias. It is not a trained linear bias: expert_bias is a requires_grad=False parameter, updated outside the autograd graph during training. Inference just consumes the stored vector.

Shared experts are not routed. They always run and are not columns of W_{\mathrm{gate}}. Softmax-over-N_r+1 (including a shared slot) is not this design.

The paper’s “normalizing the router scores before using them for gating” is exactly (18): normalize the selected s_{i,t}, not the full sigmoid vector and not s+b. Config route_norm: true matches that. score_func: "sigmoid" is the checkpoint’s name for (16).

Official code

AfmoeTokenChoiceRouter in modeling_afmoe.py (also modular_afmoe.py). expert_bias lives on AfmoeSparseMoeBlock. Hyperparameters: score_func, route_norm, route_scale, num_experts, num_experts_per_tok in Trinity-Large-Preview config.json. Equations (16)–(18) are in the tech report §2.3.

Watch-outs

• Softmax routing is a different model. Even “softmax then Top-K” changes both independence of s_{i,t} and the meaning of b_i.
• Add b_i only inside Top-K. Using s_{i,t}+b_i as the gate (or as the renormalization numerator) lets the load-balance bias leak into the forward mix.
• Renormalize over the K_r chosen scores, not over all N_r sigmoids. The latter dilutes gates by the mass of dropped experts.
• Route scale is applied after the unit-sum normalize. Scaling before the divide cancels; scaling logits before sigmoid is not Table 2’s 2.448.

Sources

• Paper: Arcee Trinity Large Technical Report, §2.3 equations (16)–(18), Table 2
• Code: arcee-ai/trinity-large-tech-report; transformers AfmoeTokenChoiceRouter