AI Engineer Exhaustive
EasyPlusArcee Trinity

SMEBU Load Balancing

Arcee Trinity

Easy · 🔒 Plus required

Soft-clamped momentum expert bias updates

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

SMEBU — Soft-clamped Momentum Expert Bias Updates — is Trinity Large’s replacement for the sign-step aux-loss-free bias of Wang et al. (2024a). Nano and Mini still use that sign update plus a re-centering of \mathbf{b}. Large does not: §2.3 argues that a constant-size \pm\lambda step cannot settle on a fixed ideal bias, so the bias oscillates, and the oscillation gets worse as N_r grows to 256.

SMEBU keeps the same interface as aux-loss-free balancing: a vector \mathbf{b}\in\mathbb{R}^{N_r} added only inside Top-K (equation (17)). The change is how \mathbf{b} is updated each optimizer step. The released Hugging Face model stores expert_bias but does not implement the update; SMEBU is specified only in the report.

How it works

Let n_i be tokens routed to expert i this step and \bar{n}=N_r^{-1}\sum_j n_j the mean load. Nano/Mini (equations (20)–(22)) take

\Delta b_i=\gamma\,\mathrm{sign}(\bar{n}-n_i),\qquad b_i\leftarrow b_i+\Delta b_i-\frac{1}{N_r}\sum_j b_j.

Large instead forms a normalized violation and tanh-clamps it (equations (23)–(24)):

v_i=\frac{\bar{n}-n_i}{\bar{n}},\qquad \tilde{v}_i=\tanh(\kappa\,v_i).

Division by \bar{n} makes v_i dimensionless: a 2\times overload is v_i=-1 whether the batch is 10^3 or 10^7 tokens. \tanh(\kappa v_i) is a bounded, odd, smooth stand-in for \mathrm{sign}(v_i). Near balance (v_i\approx 0) the update shrinks toward 0 instead of remaining \pm\lambda. They set \kappa=2 so moderate violations still saturate slowly.

Then (25)–(28):

\Delta b_i=\lambda\tilde{v}_i,\qquad \Delta b_i\leftarrow\Delta b_i-\frac{1}{N_r}\sum_j\Delta b_j,

m_i\leftarrow\beta m_i+(1-\beta)\Delta b_i,\qquad b_i\leftarrow b_i+m_i.

\mathbf{m} is a persistent momentum buffer (not an Adam moment on the router weights). Re-centering is applied to \Delta\mathbf{b} before momentum, not to \mathbf{b} after the add — that differs from the Nano/Mini recenter of \mathbf{b} itself. Trinity Large used \lambda=5\times 10^{-4}, \beta=0.5, \kappa=2.

Under-used experts (n_i<\bar{n}) get v_i>0, hence a positive \Delta b_i, which raises s_{i,t}+b_i and makes Top-K more likely to pick them next step. Over-used experts get the opposite. Gates g_{i,t} still use s_{i,t} only, so SMEBU never rescales expert outputs directly.

They also keep a small sequence-wise balance loss (29)–(32) with coefficient \alpha (load_balance_coeff: 5e-5 on the Large checkpoint). SMEBU is the primary balancer; that loss is a within-sequence nudge.

Linear unclamped \Delta b_i\propto v_i reduced MaxVio quickly in their tests and later went unstable. Hard \mathrm{sign} cannot converge. Tanh plus momentum is the compromise they actually trained.

Official code

SMEBU is §2.3 of arcee-ai/trinity-large-tech-report / arXiv:2602.17004. There is no smebu function in transformers. Inference only reads AfmoeSparseMoeBlock.expert_bias in modeling_afmoe.py. Any recreation of (23)–(28) is from the paper, not from a public trainer in that repo.

Watch-outs

Sources