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

Alloc-MoE: Budget-Aware Expert Activation Allocation for Efficient Mixture-of-Experts Inference
Apr 9, 2026Baihui Liu, Kaiyuan Tian, Wei Wang +3Mixture-of-Experts Inference Mixture of Experts
National Key Laboratory of Parallel and Distributed Computing, College of Computer Science and Technology, National University of Defense Technology
MoE Lens -- An Expert Is All You Need
Mar 6, 2026Marmik Chaudhari, Idhant Gulati, Nishkal Hundia +2Mixture-of-Experts Models Mixture-of-Experts Language Models
The Pennsylvania State University, University Park, PA 16802 USA · University of Maryland, College Park, MD 20742 USA · Harvard University, Cambridge, MA 02138 USA
MoRE: Scaling mixture of experts with hardware-aware low-rank routing
Sep 28, 2026Honam Wong, Surbhi Goel, Enric Boix-AdseràParameter-Free Mixture-of-Experts Routing Low-Rank Matrix Decomposition
University of Pennsylvania · The Wharton School, University of Pennsylvania

Apr 9, 2026J/K move · Enter open · S save

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

Baihui Liu, Kaiyuan Tian, Wei Wang, Zhaoning Zhang+2

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

MoE Lens -- An Expert Is All You Need

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