2026-09-04
SUT adds sparse experts and stick-breaking halt to a Universal Transformer. 66M params hit 29.2 BLEU on WMT'14 En-De, matching a 210M Transformer's 29.3.
Vanilla Transformers give each layer its own weights. Universal Transformers reuse one block at every depth. Sharing has two measured upsides: fewer parameters, and better compositional generalization on formal-language tests. Csordás et al. showed that an operation learned at one depth can fire at another, so unseen orderings at test time are less of a trap. Finite-depth vanilla Transformers also run into expressivity limits in the theory literature; a UT with unbounded depth is Turing-complete.
The bill is quadratic compute. An L-layer VT with P parameters per layer costs about LP per forward pass. A UT matched on total parameters packs LP weights into one block and runs it L times, so the cost is about L²P. Takase and Kiyono measured roughly 2x training time and much higher GPU memory for UT versus VT on WMT English-German. Kaplan et al. already split the same curve: shared-parameter models scale well in parameters and poorly in FLOPs. This paper keeps the sharing and cuts the quadratic blow-up.
Sparse Universal Transformer still reuses one block per layer. Inside the block, two pieces become token-conditional:
Attention has to be sparse for the same reason the FFN does. When a UT grows capacity, both pieces inflate. Sparsing only the FFN leaves the attention bill intact. MoMHA turns heads into routable experts, so a fixed k keeps multiply-accumulates flat as expert count grows.
A Mutual Information Maximization auxiliary loss balances the router: raise the entropy of the marginal expert distribution so every expert gets work, and lower the conditional entropy so each routing decision is sharp. Same job as Switch Transformer's load-balancing loss, written as mutual information.
Halting is rewritten as a stick-breaking process. The original UT mixes states as a convex combination, which multiplies (1-α) across depth and starves gradients. SUT treats halt decisions as a probability distribution and computes an expected halted state: leftover halt mass times the current hidden state, plus the weighted sum of states that already halted. Queries read the unhalted state; keys and values read the expected halted state. An Adaptive Computation Time loss penalizes expected depth. At train time, once cumulative halt mass exceeds 0.999, that position is routed to a no-op expert and skipped. After training, the same threshold can be lowered to trade accuracy for fewer steps.
On WMT'14 English-German:
| Model | Params | BLEU | MACs |
| Transformer base | 65M | 27.3 | 604M |
| UT | 65M | 28.9 | not reported |
| UT base + stick-breaking | 64M | 29.3 | 1998M |
| SUT base | 66M | 29.2 | 787M |
| Transformer big | 210M | 29.3 | 2090M |
| SUT big | 110M | 29.4 | 787M |
| UT big + stick-breaking | 105M | 29.6 | 3707M |
SUT base matches Transformer big's 29.3 BLEU with 66M parameters and 787M MACs, against 210M and 2090M. Against a dense UT of similar size it loses 0.1-0.2 BLEU and spends about one-fifth to two-fifths the MACs. Because k is held fixed, scaling SUT from 66M to 110M leaves MACs unchanged at 787M. Admin 60L-12L still leads at 30.1 BLEU with 256M parameters; SUT does not cross that line.
Ablations from SUT base (29.2 BLEU): drop MIM to 28.9, replace MoMHA with plain MoA to 28.7, drop the ACT loss to 29.0, drop halting to 29.1. On translation, mixture attention and MIM move the needle more than halting. Expert-token co-occurrence shows coarse buckets: some experts fire on determiners, others on pronouns, suffixes, or nouns. That is a hint of modularity, not evidence that an expert owns an algorithmic step.
On CFQ, natural language maps to SPARQL and compound divergence measures how different train and test are in token combinations. The hyperparameter search collapsed to E=1, a dense UT. Without pretraining, UT with halting averages 58.4 across MCD splits, against 21.3 for Bergen et al.'s T5-based UT and 21.4 for Keysers' Transformer. Their own T5-style UT reimplementation already reaches 52.3; switching to T5 relative bias and pre-LN lifts it further. Pretrained Dangle still leads at 66.1, at 51033M MACs, because it reruns the encoder for every decoded token.
Logical inference, train on 0-6 operators, test on 7-12: SUT scores 98/97/94/90/88/81, LSTM 88/84/80/78/71/69, Transformer stuck near 51. Compositional splits A/B/C, easy to hard: SUT 97/94/52, LSTM 80/60/59, Transformer 53/51/51. Recurrence plus sharing helps a lot on A and B. On C, SUT's 52 sits below LSTM's 59. Mean halt depth rises with operator count.
After training, lowering the halt threshold saves about 50% of steps on the formal-language tasks with almost no accuracy loss, about 33% of steps on CFQ at threshold 0.8, and only about 9% on WMT while holding 29.1 BLEU. Translation models halt late.
If layer sharing was dropped because of the L²P tax, this is a workable patch: sparse experts decouple parameter count from MACs, and a fixed k means extra experts do not add multiply-accumulates. A 110M SUT tying a 200M-class dense model on WMT is a parameter-efficiency result, not a new translation SOTA.
The recurrence-plus-sharing bias is real on the synthetic side. Length extrapolation and the easier compositional splits move a lot. Split C does not, so shared weights plus early exit do not cover the whole compositionality hole. The post-hoc threshold is a deployment knob that does not require retraining; on translation that knob barely turns.
The authors say some compositional generalization is still unsolved, and that larger-scale tests are missing.
A few sharper gaps. The CFQ headline numbers come from a dense UT; sparsity did not help on that small benchmark. SUT is consistently a bit worse than a same-size dense UT on WMT; the win is MACs. Against Transformer base, SUT base is actually heavier (787M vs 604M MACs). Early exit on translation saves 9%. The largest model here is 110M, so whether the same story holds at a billion parameters is untested. Expert-token co-occurrence looks like POS buckets, which is weak evidence of algorithmic modularity.