Shaping Latent Mixture-of-Experts: From Expert Geometry to Layerwise Capacity Allocation

Yaoxiang Wang; Kai Zhang; Ruobing Xie; Lanrui Wang; Guoliang Zhao; Junhui Li; Jinsong Su

cs.AIOct 10, 2026

Shaping Latent Mixture-of-Experts: From Expert Geometry to Layerwise Capacity Allocation

Authors: Yaoxiang Wang, Kai Zhang, Ruobing Xie, Lanrui Wang, Guoliang Zhao, Junhui Li, Jinsong Su

Organizations: Xiamen University · Tencent Hunyuan · Soochow University

Abstract

Fine-grained mixture-of-experts (MoE) models commonly increase expert count by narrowing each expert, treating expert shape as a by-product of routing granularity. LatentMoE moves expert computation into a lower-dimensional space, exposing the input dimension and intermediate width of each expert as explicit design choices. We show that an expert's local Jacobian rank is jointly bounded by these two dimensions, and that balancing them maximizes the rank ceiling at fixed per-expert parameters. Controlled shape sweeps across model scales and routing granularities consistently favor near-square experts. We further use layerwise intrinsic dimension (ID) to guide capacity allocation across depth, co-varying the number and shape of experts while keeping per-layer parameter count, activated FLOPs, and communication volume constant. Compared with layer-homogeneous baselines, our ID-guided allocation achieves lower validation loss and stronger downstream performance. These results establish expert shape and layerwise allocation as practical design axes for MoEs.

Figures & tables

Figure 1: From a wide dense FFN to fine-grained routed experts. Blue denotes the dense FFN and the experts activated for one token; gray denotes inactive experts. Let

M

denote the dense FFN intermediate width and

K

the number of experts selected per token. Activated-width matching gives each selected expert intermediate width

m=M/K

, making each expert much narrower.

Figure 2: Fine-graining trends across representative MoE architectures. Left: the increasing per-expert aspect ratios

d/m

. Right: the increasing numbers of routed experts

N

on a log scale. Values are taken from the corresponding model specifications ( Jiang et al., 2024 ; Databricks AI Research Team, 2024 ; Dai et al., 2024 ; DeepSeek-AI, 2024a ; DeepSeek-AI, 2024b ; Team, 2025 ) .

Figure 3: Controlled expert-shape sweeps at two model scales and two routing granularities. Left: validation loss versus aspect ratio; the dashed line marks a square expert. Right: the same runs plotted against the per-expert local-rank ceiling

q=\min(\ell,m)

, with loss averaged over shapes sharing the same

q

. Lower is better; Appendix C.1 gives all configurations and losses.

Figure 4: Layerwise TwoNN profiles. Left: final profiles at two model scales exhibit closely aligned variation over relative depth; the color-matched horizontal dashed lines mark the thresholds used to separate the shallow and deeper regimes. Right: the pointwise range across a separate 8B checkpoint series shows that the depth-wise profile is already visible early and persists through training.

Scale	Regime (blocks)	$\ell$	$m$	$N$	$K$	$m_{s}$
3.4B	shallow (1–4)	352	448	640	20	490
	deeper (5–11)	512	440	448	14	448
8B	shallow (1–5)	464	512	704	22	1077
	deeper (6–15)	640	512	512	16	1024

Table 1: Two-segment ID-guided configurations used in the experiments. Block 0 is dense; the listed indices are MoE blocks. Here

m_{s}

is the full-width shared-expert intermediate dimension.

Setting	ARC-E	ARC-C	BoolQ	CSQA	Hella.	OBQA	PIQA	SIQA	Wino.	Avg. $\pm$ s.d.
3.4B
Base	60.34	33.79	55.66	50.21	45.13	40.13	69.08	45.21	52.27	$50.20\pm 0.15$
LatentMoE $\alpha=2$	61.64	34.81	59.78	51.49	45.90	41.47	69.88	45.58	51.93	$51.39\pm 0.16$
LatentMoE $\alpha=4$	61.93	33.93	61.11	51.00	45.62	39.80	69.37	45.47	51.38	$51.07\pm 0.24$
Ours	62.18	36.01	60.57	52.14	46.86	42.47	70.66	45.17	52.16	$\mathbf{52.02\pm 0.30}$
8B

Table 2: Task-level downstream results using the mc_or_rc score (%). Task columns and Avg. report means over three training runs; Avg. additionally reports the standard deviation across runs.

3.4B (11 MoE blocks)
Setting	$(\ell,m)$	$(K,N)$	Val. loss
Uniform narrow	$11\times(352,448)$	$11\times(20,640)$	2.2391
Uniform wide	$11\times(512,440)$	$11\times(14,448)$	2.2305
ID-guided	$4\times(352,448)+7\times(512,440)$	$4\times(20,640)+7\times(14,448)$	$\mathbf{2.2234}$
Exchange (fixed split)	$4\times(512,440)+7\times(352,448)$	$4\times(14,448)+7\times(20,640)$	2.2422
Exchange (matched count)	$7\times(512,440)+4\times(352,448)$	$7\times(14,448)+4\times(20,640)$	2.2433

Table 3: Configurations and validation loss for matched layerwise controls. Narrow and wide denote the lower- and higher-capacity regimes in Table 1 . The fixed-split exchange preserves the ID-guided boundary while swapping regimes; the matched-count exchange reverses their order while preserving how many blocks use each regime. Allocations are ordered from shallow to deep. Lower loss is better.

Figure 5: Realized local-Jacobian rank. Left: depth-averaged rank normalized by

q

; markers show means, bars show ranges, and the dashed line is the ceiling. Right: blockwise normalized rank for two fixed-

q

3.4B models; the dashed curve is base-model TwoNN ID.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

	Ours: ID-guided layerwise	two regimes in Table 1
Scale	Setting	Expert input	$m$	$(N,K)$	$m_{s}$
8B	Base	uncompressed (1,280)	512	$(256,8)$	768
	LatentMoE, $\alpha=4$	latent ( $\ell=320$ )	512	$(1024,32)$	1067
	LatentMoE, $\alpha=2$	latent ( $\ell=640$ )	512	$(512,16)$	1109
3.4B	Base	uncompressed (1,024)	384	$(256,8)$	384
	LatentMoE, $\alpha=4$	latent ( $\ell=256$ )	384	$(1024,32)$	597

Appendix

Table 4: Models used in the downstream evaluation. For uniform models,

\alpha=d/\ell

. Expert input denotes the representation presented to each expert: the base receives the uncompressed residual stream, whereas LatentMoE operates at latent width

\ell

. The layerwise entries use the two configurations in Table 1 .

Figure 6: Expert-parallel communication paths. Top: a standard MoE dispatches and returns

K

vectors of width

d

per token. Bottom: LatentMoE retains routing decisions in the model hidden space but places the shared down- and up-projections outside the two all-to-all collectives, reducing each routed vector from width

d

\ell

. Gray dashed arrows show the local routing computation and Top-

K

assignments; colored arrows show the activation path.

Setting	Throughput (tokens/GPU/s)
3.4B
Base	17991.2
LatentMoE $\alpha=2$	17436.0
LatentMoE $\alpha=4$	17319.2
Ours	17144.1
8B

Appendix

Table 5: End-to-end training throughput.

8B-A560M				3.4B-A250M
$K/N$	$\ell$	$m$	Val. loss	$K/N$	$\ell$	$m$	Val. loss
$8/256$	320	2048	2.1151	$8/256$	256	1536	2.2539
$8/256$	512	1280	2.1064	$8/256$	512	768	2.2385
$8/256$	640	1024	2.1047	$8/256$	640	616	2.2373
$8/256$	816	800	2.1065	$8/256$	768	512	2.2407
$8/256$	1024	640	2.1074	$8/256$	1024	384	2.2487

Appendix

Table 6: Complete configurations and validation losses for Figure 3 .

K/N

denotes the number of activated experts over the total number of experts.

Latent width $\ell$	Expert width $m$	Val. loss	Parameter advantage (% total)
256	1536	2.2517	0.00
384	1024	2.2418	0.08
512	768	2.2374	0.16
640	616	2.2371	0.23
768	512	2.2366	0.31
1024	384	2.2456	0.47

Appendix

Table 7: Validation loss in the fixed-

m_{s}

control. The parameter advantage is the additional down/up-projection parameters relative to

\ell=256

, normalized by the total parameter count of the 3.4B-A250M base model.

Setting	$(\ell,m)$	$(K,N)$	Val. loss
Three-block shallow	$3\times(352,448)+8\times(512,440)$	$3\times(20,640)+8\times(14,448)$	2.2253
ID-guided	$4\times(352,448)+7\times(512,440)$	$4\times(20,640)+7\times(14,448)$	$\mathbf{2.2234}$
Five-block shallow	$5\times(352,448)+6\times(512,440)$	$5\times(20,640)+6\times(14,448)$	2.2242
Final block narrow	$5\times(352,448)+5\times(512,440)$ $+\,1\times(352,448)$	$5\times(20,640)+5\times(14,448)$ $+\,1\times(20,640)$	2.2267
Strongly unbalanced shapes	$4\times(160,704)+7\times(880,224)$	$4\times(28,896)+7\times(16,512)$	2.2498

Appendix

Table 8: Boundary and shape sensitivity at 3.4B. The first four rows use the narrow and wide regimes from Table 1 ; the final row is a matched-budget shape stress test. Lower loss is better.

$K/N$	$\ell$	$m$	$q$	$W_{g}$	$W_{u}$	$W_{o}$	$J$	$J/q$	$Y/q$	$Z/\ell$	$J_{\mathrm{mix}}/\ell$
$8/256$	256	1536	256	256	256	256	253.6	0.991	1.000	1.000	0.994
$8/256$	512	768	512	512	512	512	494.9	0.967	1.000	1.000	0.985
$8/256$	640	616	616	616	616	616	566.8	0.920	0.999	1.000	0.980
$8/256$	768	512	512	512	512	512	509.8	0.996	1.000	1.000	0.974
$8/256$	1024	384	384	384	384	384	383.8	0.999	1.000	0.995	0.935
$16/512$	192	1024	192	192	192	192	188.9	0.984	1.000	1.000	0.993

Appendix

Table 9: Realized ranks for the complete 3.4B-A250M shape sweep. Weight and Jacobian columns are absolute ranks; the remaining columns are normalized by their stated ceilings.

$K/N$	$\ell$	$m$	$q$	$W_{g}$	$W_{u}$	$W_{o}$	$J$	$J/q$	$Y/q$	$Z/\ell$	$J_{\mathrm{mix}}/\ell$
$8/256$	320	2048	320	320	320	320	316.7	0.990	1.000	1.000	0.993
$8/256$	512	1280	512	512	512	512	502.7	0.982	1.000	1.000	0.990
$8/256$	640	1024	640	640	640	640	619.0	0.967	0.999	1.000	0.987
$8/256$	816	800	800	799.64	799.91	800	727.3	0.909	0.995	1.000	0.980
$8/256$	1024	640	640	640	640	640	637.1	0.995	1.000	1.000	0.967
$8/256$	1280	512	512	512	512	512	511.6	0.999	1.000	0.995	0.934

Appendix

Table 10: Realized ranks for the complete 8B-A560M shape sweep, using the same definitions as Table 9 .

Scale	$K/N$	Checkpoints	$r(q,J)$	$\rho(q,J)$	$J/q$ range	Ordered pairs
3.4B	$8/256$	5	0.994	0.975	0.920–0.999	9/9
3.4B	$16/512$	8	0.996	0.982	0.921–1.000	25/25
8B	$8/256$	6	0.992	0.971	0.909–0.999	13/13
8B	$16/512$	8	0.997	0.982	0.909–1.000	25/25

Appendix

Table 11: Ceiling–realized-rank relation within each controlled shape sweep.

J

is the depth-averaged realized local-Jacobian rank. The final column counts all within-sweep pairs with distinct ceilings for which a larger

q

also gives a larger

J

Scale	Blocks	$\ell\times m$	$q$	Mean ID	$J$	$J/q$
3.4B	1–4	$352\!\times\!448$	352	4.78	334.3	0.950
3.4B	5–11	$512\!\times\!440$	440	12.97	431.2	0.980
8B	1–5	$464\!\times\!512$	464	4.87	425.0	0.916
8B	6–15	$640\!\times\!512$	512	13.07	504.9	0.986

Appendix

Table 12: Realized rank in the shallow and deep regimes of the two-segment models. ID and rank are averaged within each segment.

Figure 7: Matched-endpoint realized-rank check at 3.4B. Left: the ID-guided checkpoint’s blockwise local-Jacobian numerical rank, its allocated ceiling

q_{t}

, and the realized ranks of the two uniform endpoint checkpoints. The uniform narrow and wide checkpoints repeat their respective

(\ell,m,K,N)

configurations at every MoE block. Base-model TwoNN ID is shown only as the ordinal signal used to choose the allocation; its numerical scale is not comparable to rank. Right: the ID-guided rank minus the layer-matched uniform endpoint, normalized by

q_{t}

3.4B-A250M ( $L=12$ )			8B-A560M ( $L=16$ )
Layer	$t/L$	ID	Layer	$t/L$	ID
L0	0.000	0.194	L0	0.000	0.183
L1	0.083	2.092	L1	0.062	2.081
L2	0.167	3.835	L2	0.125	3.631
L3	0.250	5.587	L3	0.188	4.839
L4	0.333	7.607	L4	0.250	5.970

Appendix

Table 13: Layerwise TwoNN intrinsic-dimension values underlying the left panel of Figure 4 . Relative depth is

t/L

, where

L

is the number of Transformer blocks;

L0

denotes the embedding output.

Layer	2k	4k	10k	16k	20k
L0	0.161	0.183	0.189	0.192	0.192
L1	1.761	1.753	1.921	2.094	2.163
L2	2.757	3.085	3.413	3.637	3.705
L3	4.137	4.265	4.755	4.993	5.144
L4	5.061	4.884	5.246	5.571	5.614
L5	6.235	6.059	6.505	6.873	6.928

Appendix

Table 14: Layerwise TwoNN intrinsic dimension over training for the separate 8B checkpoint series underlying the right panel of Figure 4 .

Explore similar work

Apr 9, 2026

Alloc-MoE: Budget-Aware Expert Activation Allocation for Efficient Mixture-of-Experts Inference

Mixture-of-Experts (MoE) has become a dominant architecture for scaling large language models due to their sparse activation mechanism. However, the substantial number of expert activations creates a critical latency bottleneck during inference, especially in resource-constrained deployment scenarios. Existing approaches that reduce expert activations potentially lead to severe model performance degradation. In this work, we introduce the concept of activation budget as a constraint on the number of expert activations and propose Alloc-MoE, a unified framework that optimizes budget allocation coordinately at both the layer and token levels to minimize performance degradation. At the layer level, we introduce Alloc-L, which leverages sensitivity profiling and dynamic programming to determine the optimal allocation of expert activations across layers. At the token level, we propose Alloc-T, which dynamically redistributes activations based on routing scores, optimizing budget allocation without increasing latency. Extensive experiments across multiple MoE models demonstrate that Alloc-MoE maintains model performance under a constrained activation budget. Especially, Alloc-MoE achieves

1.15\times

prefill and

1.34\times

decode speedups on DeepSeek-V2-Lite at half of the original budget.

Baihui Liu, Kaiyuan Tian, Wei Wang +3

National Key Laboratory of Parallel and Distributed Computing, College of Computer Science and Technology, National University of Defense Technology

Mixture-of-Experts Inference Mixture of Experts Efficient Language Model Inference

Ask PDF

Mar 6, 2026

MoE Lens -- An Expert Is All You Need

Mixture of Experts (MoE) models enable parameter-efficient scaling through sparse expert activations, yet optimizing their inference and memory costs remains challenging due to limited understanding of their specialization behavior. We present a systematic analysis of expert specialization in MoEs through two complementary approaches: domain-specific routing patterns and an early decoding framework that tracks expert contributions to output representations. Our analysis of the DeepSeekMoE model reveals that despite having 64 routed experts with 6 active for each layer's computation, the model predominantly relies on a few specialized experts, with the top-weighted expert's output closely approximating the full ensemble prediction. We quantitatively validate these findings through a systematic analysis of the token routing distribution, demonstrating that very few experts handle over 50% of routing decisions across different specialized domains. Hidden state similarity between single and ensemble experts for every layer is extremely high, with some layers having cosine similarity as high as 0.95 and perplexity increasing by only 5% when using a single expert across all three domains. Our results indicate that Mixture of Experts models exhibit concentrated expertise highlighting potential opportunities for inference optimization through targeted expert pruning while maintaining model performance and opening avenues towards studying localization of learned knowledge in these models.

Marmik Chaudhari, Idhant Gulati, Nishkal Hundia +2

The Pennsylvania State University, University Park, PA 16802 USA · University of Maryland, College Park, MD 20742 USA · Harvard University, Cambridge, MA 02138 USA

Mixture-of-Experts Models Mixture-of-Experts Language Models Mixture of Experts LLM Routing

Ask PDF

Sep 28, 2026

MoRE: Scaling mixture of experts with hardware-aware low-rank routing

Mixture-of-Experts (MoE) layers are central to frontier language models, and recent architectures push toward more and smaller experts. In this regime, the standard linear router becomes a bottleneck: with

M

experts and hidden dimension

h

, its per-token cost

Θ(Mh)

dominates the MoE layer once

M

is large. We introduce MoRE (Mixture of Rank-reduced-routed Experts), which factorizes the router weight matrix at rank

r

and reduces the routing cost to

O((h + M)r)

. We prove that rank logarithmic in

M

suffices for routing expressivity when the number of active experts is fixed, and is necessary up to precision factors. We also prove that logarithmic rank preserves load balance in a Gaussian memorization model, and training on a synthetic phonebook task shows that low rank does not hurt memorization. At matched active FLOPs, the factorization allows a factor of

Θ(h/r)

more experts. To realize this gain in wall-clock time, we design a fused Triton kernel at inference that avoids expensive memory operations on HBM. Empirically, MoRE improves memorization on the phonebook task and performance on knowledge-intensive Q&A benchmarks after pretraining, while matching reasoning ability. Code available at https://github.com/Matheart/MoRE_code.

Honam Wong, Surbhi Goel, Enric Boix-Adserà

University of Pennsylvania · The Wharton School, University of Pennsylvania

Parameter-Free Mixture-of-Experts Routing Low-Rank Matrix Decomposition Mixture-of-Experts Inference LLM Inference Acceleration Mixture-of-Experts Models

Ask PDF