How to Loop MoE: Flatten the Experts, Untie the Attention

Shouren Wang; Chuang Ma; Mohsen Hariri; Debargha Ganguly; Wang Yang; Xiaoqing Tong; Qianying Liu; Xiaotian Han

cs.LGSep 28, 2026

How to Loop MoE: Flatten the Experts, Untie the Attention

Authors: Shouren Wang, Chuang Ma, Mohsen Hariri, Debargha Ganguly, Wang Yang, Xiaoqing Tong, Qianying Liu, Xiaotian Han, +1 more

Organizations: Case Western Reserve University · Kyoto University · NII LLMC

Abstract

Looped Transformers reuse one block of layers several times: by spending extra computation they push a model of fixed size further, and so use its parameters more fully; while sparse mixture-of-experts (MoE) models activate only a few of many experts for each token. Looped MoE bridges these two design philosophies and gives MoE models new potential for better expert usage, but it raises a question: how to loop a MoE? We answer it with Foil. With the expert parameters and the expert compute per token held fixed, Foil (1) flattens the experts, halving the expert layers, doubling the experts per layer and doubling the passes, so that every routing decision chooses from a larger pool, and (2) unties the attention, giving each pass its own attention parameters while the experts and routers stay shared. Experiments show that Foil clearly outperforms the unflattened looped baseline: at 20B tokens every Foil model has lower pretraining loss than the baseline; at 100B tokens the loss improves monotonically with the degree of flattening, the most flattened Foil ending 0.012 nat below the baseline at equal parameters and compute, with downstream accuracy on par or better; untying the attention also yields more balanced and more confident routing at equal shape. Our ablations analyse why Foil works and turn the findings into design guidance for looped MoE: the returns of looping and of widening the expert layers amplify each other, routing confidence tracks healthy expert use better than load balance, and a sparse looped MoE should therefore use more experts per layer and more passes. Code and configurations are available at https://github.com/SR-A-W/how-to-loop-moe.

Figures & tables

Figure 1: Motivation of designing Foil. Left: why loop a sparse MoE . (1) At equal parameters, looping lowers the loss, and matches non-looped models with much less parameters; (2) Tokens reaches more distinct experts the more passes it makes. Right: Foil’s design . (1) Flattens the experts, (2) Unties the attention. (3)Loops more times, (4)Keep parameters and compute unchanged.

Figure 2: Vanilla looped MoE versus Foil . Both models share the same skeleton: a prelude (embedding and one regular MoE layer), a recurrent core applied for several passes, and a coda (one regular MoE layer and the LM head). The vanilla core ties both attention and experts across passes. Foil flattens the core—it keeps the total number of experts fixed while using fewer layers, more experts per layer and more passes—and unties the attention: the experts are shared across passes, but every pass has its own attention set, selected by the pass index. Both models call the same number of experts per token. When the core has several layers, every layer carries one attention set per pass; the ellipsis marks the remaining sets. Red: attention; blue: experts.

Figure 3: Flattening with untied attention: loss and downstream accuracy of the Foils and Base. Top row: 20B tokens; bottom row: after continued training to 100B tokens. In each row, solid bars on the left show the final loss (lower is better) and tinted bars on the right the zero-shot accuracy on the three representative tasks (higher is better), each panel with its own vertical axis; bars from left to right: Foil-1, Foil-2, Foil-3, Base. At 20B all three Foils are below Base ; at 100B the loss falls monotonically with flattening and Foil-1 ends

0.012

nat below Base . Downstream accuracy is on par or slightly better; further tasks are in Appendix E .

Figure 4: Tied versus untied attention at equal shape: loss (a), downstream accuracy (b) and routing (c). Horizontal axis: flattening steps from

(8,8,2)

(64,1,16)

; each step to the right doubles

E

and

L

and halves

D

. Red squares: untied attention (Base and Foil-3, -2, -1); grey-blue circles: tied attention (Base-tied and proto-Foil-3, -2, -1). (a1, a2) Final loss at 20B and 100B tokens; error bars are twice the standard error of the paired difference from the

(8,8,2)

model of the same line. (b1, b2) Mean zero-shot accuracy over the three representative tasks,

\pm 1

standard error. (c1, c2) Load balance

B_{2}

and routing confidence MMR of the recurrent core at 20B tokens, from a single-seed probe without error bars; since both depend on

E

, only the two points at the same shape are compared. With tied attention the loss rises again after the first flattening step, while with untied attention it levels off at 20B and keeps falling at 100B; the gap is largest for the most flattened pair (

0.049

nat at 100B). At every shape the untied model is both more balanced and more confident.

						$\Delta$ loss vs. model with
$D$	$E$	$L$	$E_{\mathrm{real}}$	$k/E$	Final loss (nat)	half the passes	half the experts per layer
$4$	$8$	$2$	32	$1/4$	$2.7874$	–	–
		$4$	32	$1/4$	$2.7196$	$0.068$	–
		$8$	32	$1/4$	$2.6875$	$0.032$	–
	$16$	$2$	64	$1/8$	$2.7490$	–	$0.038$
		$4$	64	$1/8$	$2.6793$	$0.070$	$0.040$

Table 1: Final loss (nat) of the tied-attention models at 20B tokens and its decrease from the model with half the passes or half the experts per layer (standard errors

0.0006

–

0.0012

); –: not applicable.

Figure 5: Routing confidence (MMR) per loop pass at 20B tokens (geometric mean over the core layers). (a) At fixed

(E,D)=(8,8)

, more passes lower the confidence of every pass, and beyond

L=2

the model-level MMR falls from

4.15

3.36

;

(8,8,8)

peaks at pass 7. (b, c) Along the flattening sequence, confidence generally rises over the passes; proto-Foil-2, proto-Foil-1 and Foil-1 reach a peak (black triangles) and then fall, whereas Foil-2, the Foil with the lowest loss at 20B, has no peak. After a peak, further flattening or more passes very likely gain least ( Section 4.4 ).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Metric	Definition	Range, direction	Source
normalised load entropy $B_{H}$	$-\sum_{i}q_{i}\ln q_{i}/\ln E$	$[\ln k/\ln E,1]$ , higher more even	( Nguyen et al., 2026 )
MaxVio	$E\max_{i}q_{i}-1$	$[0,E/k-1]$ , lower less overload	( Wang et al., 2024 )
normalised Gini $G^{\ast}$	$\sum_{i,j}\|q_{i}-q_{j}\|/[2(E-1)]$	$[0,(E-k)/(E-1)]$ , lower more even	( Chen et al., 2026a )
dropped-assignment rate	dropped / requested assignments	$[0,1]$ ; not applicable (no drops)	( Fedus et al., 2022 ; Lepikhin et al., 2021 )
selected-expert probability	$p_{i}$ , $i\in S$	$[0,1]$ , higher more support	( Fedus et al., 2022 ; Riquelme et al., 2021 )
selected-set mass	$\sum_{i\in S}p_{i}$	$[k/E,1]$ , higher more concentrated	( Shahout et al., 2025 )

Appendix

Table 2: Related metrics used in prior work: load balance (top four rows) and routing confidence (bottom five rows). Ranges are for top-

k

routing with

k<E

. Cited works show prior use of each quantity; the exact formulas and normalisations are as stated here.

Model	Run id	$(E,D,L)$	Attention	$E_{\text{real}}$	$E_{\text{comp}}$	$E_{\text{eq}}$	Total	Loop experts	Loop attn.	Head/ tail	Other
Base width ( $d_{\text{model}}=1024$ ): Foil family
Base-tied	S1	(8,8,2)	shared	64	32	128	520.2	302.1	33.6	83.9	100.7
proto-Foil-3	S2	(16,4,4)	shared	64	32	256	503.4	302.1	16.8	83.9	100.7
proto-Foil-2	S3	(32,2,8)	shared	64	32	512	495.0	302.1	8.4	83.9	100.7
proto-Foil-1	S4	(64,1,16)	shared	64	32	1024	490.8	302.1	4.2	83.9	100.7
Base	U1	(8,8,2)	untied	64	32	128	553.7	302.1	67.1	83.9	100.7

Appendix

Table 3: Per-configuration hyperparameters and parameter counts.

(E,D,L)

: experts per MoE layer, layers in the loop block, loop passes.

E_{\text{real}}=E\cdot D

(distinct experts),

E_{\text{comp}}=D\cdot L\cdot k

(expert calls per token),

E_{\text{eq}}=E\cdot D\cdot L

. Attention: shared = one set of attention weights reused by every pass; untied = one set per pass. Parameter columns in millions, rounded independently: total; experts (including their routers) and attention inside the loop block; the non-looped head and tail layers; “other” = token embeddings and output projection. Run ids are the identifiers used in our logs. Base width:

d_{\text{model}}=1024

, 16 attention heads, expert hidden width 1536, top-

k=2

, one head and one tail layer with 8 experts each, vocabulary 49,152. a

E=4

with

k=2

: half of the experts in a layer are active for every token. b Two head and two tail layers. c Not looped (

L=1

Model	$(E,D,L)$	attention	params (M)	$E_{\mathrm{real}}$	$E_{\mathrm{comp}}$	$E_{\mathrm{eq}}$	loss 20B	ppl 20B	loss 100B
Flattening, tied attention
Base-tied	$(8,8,2)$	tied	$520.2$	64	32	128	$2.688$	$14.71$	$2.522$
proto-Foil-3	$(16,4,4)$	tied	$503.4$	64	32	256	$2.679$	$14.58$	$2.520$
proto-Foil-2	$(32,2,8)$	tied	$495.0$	64	32	512	$2.684$	$14.64$	$2.531$
proto-Foil-1	$(64,1,16)$	tied	$490.8$	64	32	1024	$2.700$	$14.89$	$2.548$
Flattening, untied attention

Appendix

Table 4: Main results of all small-tier models (

d_{\mathrm{model}}=1024

) at 20B tokens, and at 100B tokens where the model was continued. Loss: mean task loss over the final window (nat); perplexity is its exponential. Parameters in millions. § Repeated from an earlier group for comparison. Single seed except the seed-43 repeat of Base-tied.

20B tokens
Task	Base-tied	proto-Foil-3	proto-Foil-2	proto-Foil-1	Base	Foil-3	Foil-2	Foil-1
PROST	$31.27$ $\pm 0.34$	$31.00$ $\pm 0.34$	$29.50$ $\pm 0.33$	$30.73$ $\pm 0.34$	$30.26$ $\pm 0.34$	$31.71$ $\pm 0.34$	$31.10$ $\pm 0.34$	$29.80$ $\pm 0.33$
LAMBADA (standard)	$23.17$ $\pm 0.59$	$24.22$ $\pm 0.60$	$25.48$ $\pm 0.61$	$23.73$ $\pm 0.59$	$25.89$ $\pm 0.61$	$26.14$ $\pm 0.61$	$26.66$ $\pm 0.62$	$26.14$ $\pm 0.61$
LAMBADA (OpenAI)	$32.08$ $\pm 0.65$	$32.33$ $\pm 0.65$	$32.23$ $\pm 0.65$	$31.09$ $\pm 0.64$	$35.30$ $\pm 0.67$	$34.54$ $\pm 0.66$	$33.86$ $\pm 0.66$	$34.29$ $\pm 0.66$
SWAG	$55.34$ $\pm 0.35$	$55.47$ $\pm 0.35$	$55.36$ $\pm 0.35$	$54.69$ $\pm 0.35$	$55.74$ $\pm 0.35$	$55.90$ $\pm 0.35$	$56.73$ $\pm 0.35$	$56.41$ $\pm 0.35$
HellaSwag	$39.55$ $\pm 0.49$	$40.41$ $\pm 0.49$	$40.72$ $\pm 0.49$	$39.86$ $\pm 0.49$	$40.82$ $\pm 0.49$	$41.30$ $\pm 0.49$	$41.07$ $\pm 0.49$	$41.46$ $\pm 0.49$

Appendix

Table 5: Zero-shot downstream accuracy (%,

\pm 1

standard error) of the eight models of Table 3 . SWAG, HellaSwag and PROST are scored by length-normalised accuracy, the other tasks by accuracy; the mean over tasks treats tasks as independent (standard error of a difference between two models about

0.31

20B tokens
Task	Base-tied	proto-Foil-3	proto-Foil-2	proto-Foil-1	Base	Foil-3	Foil-2	Foil-1
PROST †	$33.24$ $\pm 0.34$	$31.50$ $\pm 0.34$	$26.96$ $\pm 0.32$	$27.88$ $\pm 0.33$	$30.35$ $\pm 0.34$	$29.96$ $\pm 0.33$	$29.12$ $\pm 0.33$	$29.50$ $\pm 0.33$
LAMBADA (standard) †	$22.71$ $\pm 0.58$	$23.17$ $\pm 0.59$	$24.84$ $\pm 0.60$	$22.34$ $\pm 0.58$	$24.86$ $\pm 0.60$	$24.41$ $\pm 0.60$	$24.80$ $\pm 0.60$	$24.24$ $\pm 0.60$
LAMBADA (OpenAI)	–	–	–	–	–	–	–	–
SWAG †	$54.41$ $\pm 0.35$	$55.22$ $\pm 0.35$	$54.85$ $\pm 0.35$	$54.33$ $\pm 0.35$	$54.96$ $\pm 0.35$	$55.30$ $\pm 0.35$	$55.59$ $\pm 0.35$	$55.55$ $\pm 0.35$
HellaSwag	$40.13$ $\pm 0.49$	$41.12$ $\pm 0.49$	$41.01$ $\pm 0.49$	$40.17$ $\pm 0.49$	$40.79$ $\pm 0.49$	$41.15$ $\pm 0.49$	$41.50$ $\pm 0.49$	$41.38$ $\pm 0.49$

Appendix

Table 6: 5-shot downstream accuracy (%,

\pm 1

standard error). HellaSwag is the mean over three evaluation seeds; † seed 0 only; – marks tasks without a 5-shot variant. The mean over five tasks excludes these two and treats tasks as independent (standard error of a difference between two models about

0.38

). Metrics as in Table 5 .

20B tokens
Model	$(E,D,L)$	Training stream	FineWeb-Edu	Wikipedia	C4	arXiv
Base-tied	(8,8,2)	2.6884	2.6307	2.6840	3.1288	2.9151
proto-Foil-3	(16,4,4)	2.6793	2.6198	2.6782	3.1190	2.9147
proto-Foil-2	(32,2,8)	2.6836	2.6240	2.6933	3.1251	2.9395
proto-Foil-1	(64,1,16)	2.7005	2.6414	2.7052	3.1410	2.9733
Base	(8,8,2)	2.6659	2.6077	2.6703	3.1043	2.8997

Appendix

Table 7: Held-out loss (nat) of the Foil family on a FineWeb-Edu validation slice and three domain slices, final checkpoints. The FineWeb-Edu slice reproduces the training-stream ordering: Foil-2 lowest at 20B, Foil-1 lowest at 100B, and the untied model below the tied one at every shape. At 100B, Wikipedia and C4 agree in direction: every Foil is below Base and every untied model is below its tied counterpart. arXiv is the exception: at 20B Foil-3 and the

(16,4,4)

untied model are above their comparators, and Foil-3 remains above Base at 100B. Slice sizes 3–8M tokens; the evaluation reports means only, so no standard errors. Code and book slices are absent from the training corpus and are not reported. The slices are small and no standard errors are available, so differences of about 0.001 nat or less may be within noise.

Widening	$L=2$	$L=4$	$L=8$
$(8,4,L)\to(16,4,L)$	$0.0384\pm 0.0011$	$0.0402\pm 0.0010$	$0.0491\pm 0.0012$
$(4,8,L)\to(8,8,L)$	$0.0393\pm 0.0010$	$0.0460\pm 0.0009$	–
$(8,8,L)\to(16,8,L)$	$0.0470\pm 0.0009$	–	–

Appendix

Table 8: Widening and looping, 20B tokens, tied attention. Gains are decreases of the final-window loss (nat),

\pm 1

standard error of the paired difference. Top: widening doubles

E

at fixed

D

and

L

; there are no

(4,8,8)

(16,8,4)

(16,8,8)

models. Bottom: for each

2\times 2

block, looping doubles

L

and “both” does the two steps together; interaction = both

-

widening

-

looping, with a standard error of about

0.0014

combined from the independent paired differences.

$L$	$1$	$2$	$3$	$4$	$8$
MMR	$4.06$	$4.15$	$4.02$	$3.73$	$3.36$
Gain	–	$0.0640\pm 0.0013$	$0.1022\pm 0.0013$	$0.1237\pm 0.0014$	$0.1564\pm 0.0014$
Per pass	–	$0.0640$	$0.0511$	$0.0412$	$0.0223$

Appendix

Table 9: The

(8,8,L)

series, 20B tokens, tied attention. MMR: model-level routing confidence of the recurrent core (geometric mean). Gain: decrease of the final-window loss relative to the non-looped

(8,8,1)

(nat,

\pm 1

standard error); per pass: gain divided by the

L-1

added passes.

Flattening, tied attention
Model	$(E,D,L)$	attention	distinct experts	upper bound	ratio	$B_{2}$	MMR
Base-tied	$(8,8,2)$	tied	$23.47$	32	$0.73$	$0.921$	$4.15$
proto-Foil-3	$(16,4,4)$	tied	$15.67$	32	$0.49$	$0.849$	$5.74$
proto-Foil-2	$(32,2,8)$	tied	$11.61$	32	$0.36$	$0.807$	$7.46$
proto-Foil-1	$(64,1,16)$	tied	$9.42$	32	$0.29$	$0.731$	$8.39$
Flattening, untied attention

Appendix

Table 10: Routing metrics of the recurrent core for all small-tier models at 20B tokens (step 50,000). Class 0: distinct real experts used per token, its upper bound

D\min(E,kL)

and the ratio of the two; class 1: model-level

B_{2}

; class 2: MMR. § Repeated from an earlier group. Values depend on the pool size

E

; compare absolute values only at equal shape ( Section 2.3 ).

Foil family
			Load			Confidence
Model	Run id	$(E,D,L)$	$B_{H}$	MaxVio	$G^{*}$	Rank-1 prob.	Selected mass $M_{k}$	Boundary margin	$C_{\text{full}}$	$C_{\text{sel}}$
Base-tied	S1	(8,8,2)	0.979	0.523	0.162	0.382	0.556	0.417	0.164	0.130
proto-Foil-3	S2	(16,4,4)	0.971	1.099	0.221	0.269	0.403	0.374	0.143	0.103
proto-Foil-2	S3	(32,2,8)	0.972	1.812	0.236	0.173	0.274	0.295	0.130	0.069
proto-Foil-1	S4	(64,1,16)	0.968	3.449	0.270	0.095	0.159	0.198	0.110	0.040

Appendix

Table 11: Values of related-work routing metrics for every base-width (

d_{\text{model}}=1024

) model at the end of 20B-token training (step 50,000). Loop block only. Load metrics:

B_{H}

= normalised load entropy

H(q)/\log E

; MaxVio

=E\max_{i}q_{i}-1

;

G^{*}

= Gini coefficient normalised by

E-1

; each is computed per physical layer from selection counts pooled over all passes, then averaged arithmetically over physical layers. Confidence metrics (per token, then averaged over tokens): rank-1 mean probability; selected-set mass

M_{k}

(sum of the

k

selected full-pool probabilities); boundary margin (logit of the

k

-th minus the

(k{+}1)

-th expert);

C_{\text{full}}=1-H(p)/\log E

;

C_{\text{sel}}

= one minus the normalised entropy of the renormalised weights of the selected experts; each is averaged with equal weight over all (physical layer, pass) cells of the loop block (arithmetic mean). Selections are recomputed as the top-

k

of the stored float16 router logits. Token drop rate is 0 for every model and not applicable (no capacity limit, no dropping), so it is not listed. Probe set:

255{,}500

tokens from

500

documents,

50

from each of ten domains of the Pile. Model names and footnotes as in Table 3 .

Model	$(E,D,L)$	experts masked	$\Delta L_{T}$	random mean	random range	above random
20B tokens
Base-tied	(8,8,2)	11	$0.160$	$0.104$	$[0.068,\,0.165]$	19/20
proto-Foil-3	(16,4,4)	11	$0.173$	$0.137$	$[0.047,\,0.229]$	15/20
proto-Foil-2	(32,2,8)	11	$0.221$	$0.231$	$[0.129,\,0.386]$	9/20
proto-Foil-1	(64,1,16)	12	$0.309$	$0.271$	$[0.140,\,0.464]$	15/20
Base	(8,8,2)	9	$0.141$	$0.087$	$[0.044,\,0.143]$	19/20

Appendix

Table 12: Masking the least-used experts at the 10% traffic tier.

\Delta L_{T}

: loss increase (nat) when the least-used group is masked; random: loss increase for 20 traffic-matched random groups (mean and range); last column: number of the 20 random groups whose loss increase is below

\Delta L_{T}

Model	$(E,D,L)$	attention	ratio	random groups below
Base-tied	$(8,8,2)$	tied	$1.54$	19/20
proto-Foil-3	$(16,4,4)$	tied	$1.26$	15/20
proto-Foil-2	$(32,2,8)$	tied	$0.96$	9/20
proto-Foil-1	$(64,1,16)$	tied	$1.14$	15/20
Base	$(8,8,2)$	untied	$1.61$	19/20
Foil-3	$(16,4,4)$	untied	$1.69$	20/20

Appendix

Table 13: Least-used group versus random groups at the 10% traffic tier, 20B tokens. Ratio: loss increase from masking the least-used group divided by the mean loss increase of the 20 traffic-matched random groups; last column: number of random groups whose loss increase is below that of the least-used group. Seven of the eight ratios exceed one; none of the least-used groups falls below the range of the random groups.

Figure 6: Masking the least-used experts at the 10% traffic tier, 20B tokens. For each shape, the left bar is the tied-attention model and the right bar the untied-attention model; bar height is the loss increase from masking the least-used group and the short dark dash beside it the mean loss increase of the 20 traffic-matched random groups. The ranges of the random groups are given in Table 12 , the ratios to the random mean in Table 13 .

Explore similar work

Oct 1, 2026cs.LG

Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts

Looped Transformers introduce recurrent depth as a new scaling axis for LLMs: by repeatedly applying shared Transformer blocks, they increase effective depth without increasing parameter count. However, the benefits of looping remain unclear for large MoE LLMs under FLOPs-matched comparisons. The main reason is that the gains from additional iterations diminish quickly and can even turn into degradation, so the extra FLOPs spent on looping yield little substantial improvement. Consequently, prior work typically settles on two loops. We identify two main obstacles to scaling looped MoE. First, looping inherits and amplifies the curse of depth: hidden-state variance grows with each iteration as residual updates accumulate, which destabilizes deep recurrence and causes representations to drift. Second, looped MoE suffers from expert selection collapse: routers repeatedly select the same experts across loops, so extra iterations add computation without adding computational diversity. Guided by this diagnosis, we propose LOOM, built on a single principle: each loop should contribute new computation while keeping the recurrent state stable. LOOM stabilizes recurrence by scaling residual updates to bound variance growth and re-injecting the input embedding at every loop, and diversifies it through per-loop routers that engage different experts and a Looping Residual that carries earlier outputs forward. Experiments across 100M-1.7B models show stable scaling to 9-12 loops. Under near-iso-FLOP, the 700M model performs best at 5 loops, reducing perplexity from 18.36 to 16.54 and improving average zero-shot accuracy from 38.84% to 39.53% over the non-looped baseline. Without FLOP matching, the 1.7B model trained on 60B tokens peaks at 9 loops, reducing perplexity from 9.62 to 7.77 and improving average zero-shot accuracy from 42.4% to 47.7%. Code is available https://github.com/hed-ucas/LOOM.

Di He, Pengxiang Li, Da Chang +3

Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences · Peng Cheng Laboratory · University of Chinese Academy of Sciences +5

Mixture-Of-Experts Large Language Models Mixture-Of-Experts Transformer Architectures Hidden States Sequential Scaling Recipe

Ask PDF

May 9, 2026cs.LG

Sparse Layers are Critical to Scaling Looped Language Models

Looped language models repeat a set of transformer layers through depth, reducing memory costs and providing natural early-exit points at loop boundaries. However, looped models do not scale as favorably as standard transformers with unique layers. We compare standard and Mixture-of-Experts (MoE) transformers, with and without looping, and find two main results. First, we find Looped-MoE models scale better than the standard baseline while dense looped models do not. We trace this to routing divergence between loops: in Looped-MoE models, different experts are activated on each pass through the same shared layers, recovering expressivity without additional parameters. Our second finding is that looped models have better compute-quality trade-offs with early exits than standard models. Because each loop ends with the same layers that produce the final output, loop boundaries are superior exit points, as confirmed by earlier output convergence at these points. In sum, we provide a clear direction for scaling looped models: a Looped-MoE model with early exits can not only beat standard transformers at scale, but also enable significant memory and inference savings with minimal degradation in quality.

Ryan Lee, Jacob Biloki, Edward J. Hu +1

1USC Information Sciences Institute · 2Netflix · 3Independent Researcher

Transformer Architectures Layer-Wise Localized Lora-Moe Mixture-Of-Experts Peak Memory

Ask PDF

Date pendingcs.LG

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size, conflating architectural advantage with extra FLOPs. We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache. Through a series of ablations, we arrive at a recipe we call SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets. We scale SMELT across four sizes up to 54B non-embedding parameters and fit a separate Chinchilla-style scaling law for each architecture. SMELT's loss drops faster with compute, saving 6.8--18.0% of training FLOPs on the compute-optimal frontier. The advantage transfers to downstream benchmarks beyond what validation loss predicts, is largest on Code, and grows with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens, an inductive bias that may underlie the observed performance gains. These results show that looping can improve Transformers even under budget matching, offering a practical recipe that turns depth reuse into measurable gains.

Shaowen Wang, Ge Zhang, Kairong Luo +6

Tsinghua University · ByteDance Seed · M-A-P

Transformer Architectures Scaling Laws

Ask PDF