塑造潜在混合专家模型：从专家几何结构到逐层容量分配

# Shaping Latent Mixture-of-Experts: From Expert Geometry to Layerwise Capacity Allocation
Source: [https://arxiv.org/html/2610.14644](https://arxiv.org/html/2610.14644)
Yaoxiang Wang, Kai Zhang, Ruobing Xie, Lanrui Wang,Guoliang Zhao, Junhui Li, Jinsong SuXiamen 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\.

## 1Introduction

Sparse mixture\-of\-experts \(MoE\) models have become a primary route to scaling model capacity without paying the computational cost of activating every parameter for every token\([Shazeer et al\., 2017](https://arxiv.org/html/2610.14644#bib.bib1);[Lepikhin et al\., 2021](https://arxiv.org/html/2610.14644#bib.bib2);[Fedus et al\., 2022](https://arxiv.org/html/2610.14644#bib.bib3);[Jiang et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib4)\)\. In a dense Transformer, the feed\-forward network is intentionally wider than the residual stream, reflecting a block\-level allocation of parameters between attention and feed\-forward computation\([Vaswani et al\., 2017](https://arxiv.org/html/2610.14644#bib.bib28)\)\. Recent MoE LLMs commonly replace the dense FFN with fine\-grained experts while matching the sum of the activated experts’ widths to the dense FFN width\([Dai et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib5);[DeepSeek\-AI, 2024a](https://arxiv.org/html/2610.14644#bib.bib6);[DeepSeek\-AI, 2024b](https://arxiv.org/html/2610.14644#bib.bib7)\), as illustrated in Figure[1](https://arxiv.org/html/2610.14644#S1.F1)\.

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\. LetMMdenote the dense FFN intermediate width andKKthe number of experts selected per token\. Activated\-width matching gives each selected expert intermediate widthm=M/Km=M/K, making each expert much narrower\.Once the block\-level budget of attention and the FFN is chosen, the attention configuration fixes the model hidden sizedd\(the width of the residual stream\), and every expert in a standard MoE inherits this input dimension\. Increasing the number of experts under a matched FFN budget must therefore divide the budget into smaller per\-expert intermediate dimensionsmm, making the aspect ratiod/md/ma consequence of expert granularity rather than an independent design choice\. As Figure[2](https://arxiv.org/html/2610.14644#S1.F2)shows, representative architectures with more experts consequently exhibit increasingly larged/md/mratios, a trend that may create a hidden bottleneck during training\. MatchingK​mKmto the dense intermediate width preserves the activated parameter and compute budget, but does not preserve how intermediate features are used\. A dense FFN applies the same fixed set of channels to every token, whereas an MoE router selects different combinations of experts\. Each expert must therefore remain useful across many routing combinations and retain independent representational capacity\. Reducingmmmay constrain what each expert can represent even when the aggregate activated widthK​mKmremains unchanged\.

Figure 2:Fine\-graining trends across representative MoE architectures\.Left:the increasing per\-expert aspect ratiosd/md/m\.Right:the increasing numbers of routed expertsNNon a log scale\. Values are taken from the corresponding model specifications\([Jiang et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib4);[Databricks AI Research Team, 2024](https://arxiv.org/html/2610.14644#bib.bib22);[Dai et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib5);[DeepSeek\-AI, 2024a](https://arxiv.org/html/2610.14644#bib.bib6);[DeepSeek\-AI, 2024b](https://arxiv.org/html/2610.14644#bib.bib7);[Team, 2025](https://arxiv.org/html/2610.14644#bib.bib8)\)\.LatentMoE\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\)reduces routed computation and communication by applying experts in a lower\-dimensional latent space\. This design was introduced for efficiency, but it decouples the model hidden dimension from the expert input dimension and thus exposes expert shape as a design choice\. Since its original formulation fixes a four\-fold compression, we ask how to jointly choose the expert input and intermediate dimensions under a fixed per\-expert parameter budget\. Our expert\-shape rationale starts from a simple local property: an expert’s Jacobian rank is bounded by both its input dimension and its intermediate width\. Under a fixed parameter budget, this ceiling is largest when the two dimensions are balanced\. We test this implication in controlled shape sweeps and observe a U\-shaped validation\-loss curve over aspect ratio, with the best performance in a near\-square region\.

Within the near\-square region, we further explore the effect of varying per\-expert capacity across depth\. Intrinsic dimension \(ID\), which estimates the number of locally independent degrees of freedom occupied by a layer’s input representations, varies substantially with Transformer depth\([Valeriani et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib15);[Joshi et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib16)\)\. We therefore hypothesize that higher\-ID layers benefit from greater per\-expert capacity, whereas repeating one expert shape at every layer may allocate capacity inefficiently\. We use the ID ordering to determine the direction of a coarse allocation, with the expert dimensions selected under matched\-budget and hardware constraints\. Prior heterogeneous designs vary expert counts or feed\-forward widths across layers\([Jawahar et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib21)\), but doing so naively creates pronounced per\-layer compute and communication imbalance\. This imbalance poses severe systems\-optimization challenges for pipeline partitioning and efficient MoE kernels, making such designs difficult to deploy efficiently at scale\. We instead co\-vary the number and shape of experts across depth to match per\-expert capacity to layerwise ID while keeping per\-layer parameter count, activated FLOPs, and communication volume nearly constant\. The resulting allocation is heterogeneous in capacity but nearly homogeneous from the system’s perspective\. Our experiments show that this allocation improves validation loss and downstream performance; ablations on uniform endpoints, allocation direction, and regime boundaries further support its effectiveness and robustness\.

Our contributions are summarized as follows:

- •We establish expert shape as an overlooked but important MoE design variable and empirically validate the advantage of near\-square experts, providing a low\-cost rule for choosing the LatentMoE compression dimension\.
- •Under the near\-square expert principle and guided by layerwise intrinsic dimension, we extend expert shaping from within\-layer design to across\-layer capacity allocation\. This perspective challenges the usual assumption that every MoE layer should use the same expert configuration\.
- •We introduce an infrastructure\-friendly layerwise allocation that keeps per\-layer computation and communication nearly unchanged while achieving lower validation loss and stronger downstream\-task performance\.

## 2Preliminaries

Let𝒙∈ℝd\\bm\{x\}\\in\\mathbb\{R\}^\{d\}be the input to an MoE layer\. A router selectsKKexperts from a pool ofNN, denoted by𝒯K​\(𝒙\)\\mathcal\{T\}\_\{K\}\(\\bm\{x\}\), and assigns weightspip\_\{i\}\. The routed path of a standard MoE is

𝒚std=∑i∈𝒯K​\(𝒙\)pi​Ei​\(𝒙\),Ei:ℝd→ℝd,\\bm\{y\}\_\{\\mathrm\{std\}\}=\\sum\_\{i\\in\\mathcal\{T\}\_\{K\}\(\\bm\{x\}\)\}p\_\{i\}E\_\{i\}\(\\bm\{x\}\),\\qquad E\_\{i\}:\\mathbb\{R\}^\{d\}\\rightarrow\\mathbb\{R\}^\{d\},where each expert expands through an intermediate dimensionmm\. Ignoring architecture\-dependent constants, the routed experts contain parameters proportional toN​d​mNdm\. For each token, activated expert computation is proportional toK​d​mKdm, while the all\-to\-all communication volume for dispatch and aggregation is proportional toK​dKd\. The residual widthddtherefore determines both the expert input dimension and the width of the communicated representations\.

LatentMoE introduces learned projectionsW↓∈ℝℓ×dW\_\{\\downarrow\}\\in\\mathbb\{R\}^\{\\ell\\times d\}andW↑∈ℝd×ℓW\_\{\\uparrow\}\\in\\mathbb\{R\}^\{d\\times\\ell\}\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\):

𝒛=W↓​𝒙,𝒚latent=W↑​∑i∈𝒯K​\(𝒙\)pi​Ei​\(𝒛\),Ei:ℝℓ→ℝℓ\.\\bm\{z\}=W\_\{\\downarrow\}\\bm\{x\},\\qquad\\bm\{y\}\_\{\\mathrm\{latent\}\}=W\_\{\\uparrow\}\\sum\_\{i\\in\\mathcal\{T\}\_\{K\}\(\\bm\{x\}\)\}p\_\{i\}E\_\{i\}\(\\bm\{z\}\),\\qquad E\_\{i\}:\\mathbb\{R\}^\{\\ell\}\\rightarrow\\mathbb\{R\}^\{\\ell\}\.The routed expert path, including all\-to\-all dispatch and aggregation, now operates in theℓ\\ell\-dimensional space\. Under the same accounting, routed expert parameters, per\-token activated computation, and per\-token communication volume are proportional toN​ℓ​mN\\ell m,K​ℓ​mK\\ell m, andK​ℓK\\ell, respectively\. Appendix[B\.2](https://arxiv.org/html/2610.14644#A2.SS2)provides the full accounting\.

Letα=d/ℓ\\alpha=d/\\elldenote the compression factor\. The accuracy\-oriented LatentMoE variant keepsmmfixed and reinvests the compression savings by settingN′=α​NN^\{\\prime\}=\\alpha NandK′=α​KK^\{\\prime\}=\\alpha K\. It then satisfies

N′​ℓ​m=N​d​m,K′​ℓ​m=K​d​m,K′​ℓ=K​d\.N^\{\\prime\}\\ell m=Ndm,\\qquad K^\{\\prime\}\\ell m=Kdm,\\qquad K^\{\\prime\}\\ell=Kd\.Thus, under this routed\-path accounting, it providesα\\alphatimes as many candidate and activated experts without changing routed\-expert parameters, activated expert computation, or all\-to\-all communication volume\. The original study selectedα=4\\alpha=4empirically\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\), but did not provide a general principle for jointly choosing the latent dimensionℓ\\elland expert intermediate dimensionmm\. LatentMoE variants have been adopted by frontier models including MAI\-Thinking\-1 and Kimi K3\([The Microsoft AI Team, 2026](https://arxiv.org/html/2610.14644#bib.bib10);[Kimi Team, 2026](https://arxiv.org/html/2610.14644#bib.bib11)\), demonstrating their practical relevance at large scale\.

The key structural consequence is that LatentMoE decouples the expert input dimensionℓ\\ellfrom the residual widthdd, making expert shape a configurable design variable\. We define the*expert shape*asℓ×m\\ell\\times m\. Since the per\-expert parameter count is proportional toℓ​m\\ell m, fixing this product allows the aspect ratioℓ/m\\ell/mto vary without changing the parameter budget\. We ask which shape is most effective under this constraint, first when shared across layers and then when allocated layerwise to match the intrinsic dimension of the inputs\.

## 3Expert Shaping

### 3\.1A Per\-Expert Local\-Rank Ceiling

Consider a bias\-free gated FFN expert

Ei​\(𝒛\)=Wo,i​\[ϕ⁡\(Wg,i​𝒛\)⊙Wu,i​𝒛\],E\_\{i\}\(\\bm\{z\}\)=W\_\{o,i\}\\\!\\left\[\\phi\(W\_\{g,i\}\\bm\{z\}\)\\odot W\_\{u,i\}\\bm\{z\}\\right\],whereWg,i,Wu,i∈ℝm×ℓW\_\{g,i\},W\_\{u,i\}\\in\\mathbb\{R\}^\{m\\times\\ell\}andWo,i∈ℝℓ×mW\_\{o,i\}\\in\\mathbb\{R\}^\{\\ell\\times m\}\. This form includes the SwiGLU experts used in our models\([Shazeer, 2020](https://arxiv.org/html/2610.14644#bib.bib30)\)\. Whereverϕ\\phiis differentiable, its Jacobian can be written as

JEi​\(𝒛\)=Wo,i​\[Du,i​\(𝒛\)​Wg,i\+Dg,i​\(𝒛\)​Wu,i\],J\_\{E\_\{i\}\}\(\\bm\{z\}\)=W\_\{o,i\}\\\!\\left\[D\_\{u,i\}\(\\bm\{z\}\)W\_\{g,i\}\{\}\+\{\}D\_\{g,i\}\(\\bm\{z\}\)W\_\{u,i\}\\right\],whereDu,iD\_\{u,i\}andDg,iD\_\{g,i\}are diagonalm×mm\\times mmatrices determined by the activations\. The bracketed term has shapem×ℓm\\times\\ell, and therefore

rank⁡JEi​\(𝒛\)≤min⁡\(ℓ,m\)\.\\operatorname\{rank\}J\_\{E\_\{i\}\}\(\\bm\{z\}\)\\leq\\min\(\\ell,m\)\.The nonlinearity changes the entries of the local linearization but cannot bypass themm\-dimensional intermediate representation\. We denote this ceiling byq=min⁡\(ℓ,m\)q=\\min\(\\ell,m\)\. At a fixed per\-expert parameter budget,ℓ​m\\ell mis fixed, soqqis maximized whenℓ=m\\ell=m\. The bound motivates the hypothesis that strongly unbalanced shapes will underperform near\-square ones when per\-expert capacity is limiting\.

Across our shape sweeps, expert weights and routed\-token outputs are essentially full rank, and trained local Jacobians nearly attainqq\. This supportsqqas an empirically relevant capacity proxy; Appendix[E\.1](https://arxiv.org/html/2610.14644#A5.SS1)reports the estimators, controls, and full results\.

### 3\.2Controlled Shape Sweeps

We test this hypothesis with sparse LatentMoEs at 3\.4B and 8B scales\. At each scale, we sweepℓ/m\\ell/munder both\(N,K\)=\(256,8\)\(N,K\)=\(256,8\)and\(512,16\)\(512,16\); each token therefore activates only 3\.1% of the routed experts, reflecting the very sparse, fine\-grained regime used by frontier MoEs\. Within each sweep,NNandKKare fixed,ℓ​m\\ell mis held constant, and all other training choices are shared\. Since changingℓ\\ellalso changes the projection parameters, we adjust the full\-width shared\-expert dimensionmsm\_\{s\}to keep the total per\-layer parameter count nearly constant\. Appendix[B](https://arxiv.org/html/2610.14644#A2)gives the training details\. Appendix[C\.2](https://arxiv.org/html/2610.14644#A3.SS2)reports a complementary fixed\-msm\_\{s\}control that preserves the same qualitative result without this parameter\-matching adjustment\.

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 ceilingq=min⁡\(ℓ,m\)q=\\min\(\\ell,m\), with loss averaged over shapes sharing the sameqq\. Lower is better; Appendix[C\.1](https://arxiv.org/html/2610.14644#A3.SS1)gives all configurations and losses\.Figure[3](https://arxiv.org/html/2610.14644#S3.F3)shows the same U\-shaped dependence on aspect ratio at both scales and routing granularities: highly unbalanced experts deteriorate in either direction, while the best region lies near a square\. Operationally, the tested shapes with\|ℓ/m−1\|≤0\.4\\lvert\\ell/m\-1\\rvert\\leq 0\.4lie in or near the best\-performing plateau across these sweeps, providing a practical tolerance around the idealized square shape\. Across the parameter\-matched sweeps, the resulting validation\-loss gaps are substantial relative to empirical run\-to\-run variation; Appendix Section[B](https://arxiv.org/html/2610.14644#A2)gives the supporting replication details\. Under fixed per\-expert parameters, validation loss broadly improves as the local\-rank ceilingqqgrows and then plateaus\. They yield a simple design rule: begin with a hardware\-feasible shape nearℓ=m\\ell=m, then tune within the resulting plateau if needed\.

## 4Layerwise ID\-Guided Capacity Allocation

Section[3](https://arxiv.org/html/2610.14644#S3)holds the expert budget fixed and optimizes its shape\. We next ask whether the per\-expert budget itself should be uniform across depth\. LatentMoE exposes a trade\-off: compressing token representations and the per\-expert parameter budget can hurt performance by constraining the capacity of each individual expert, while reinvesting the savings into more experts can improve performance\. A uniform compression implicitly assumes that this trade\-off is the same at every layer\. However, prior work shows that the geometry and intrinsic dimension \(ID\) of token representations vary systematically with depth\([Valeriani et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib15);[Joshi et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib16)\), suggesting that different layers may tolerate different degrees of representation and per\-expert compression\. We therefore useqt=min⁡\(ℓt,mt\)q\_\{t\}=\\min\(\\ell\_\{t\},m\_\{t\}\)as a local\-capacity proxy for layerttand test the directional hypothesis that higher\-ID regions may benefit from a largerqtq\_\{t\}, whereas lower\-ID regions may tolerate a smaller one\. Because ID andqtq\_\{t\}measure different quantities, the allocation uses their ordinal relation as a directional signal\.

### 4\.1Estimating Layerwise Intrinsic Dimension

Intrinsic dimension estimates the number of local degrees of freedom occupied by representations, which may be much smaller than their ambient width\([Ansuini et al\., 2019](https://arxiv.org/html/2610.14644#bib.bib29)\)\. Prior work has found that this quantity varies substantially with Transformer depth\([Valeriani et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib15);[Joshi et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib16)\)\. We measure it using TwoNN\([Facco et al\., 2017](https://arxiv.org/html/2610.14644#bib.bib12)\)\.

Given representation samples\{𝒉j\(t\)\}\\\{\\bm\{h\}^\{\(t\)\}\_\{j\}\\\}at a block boundary, letrj,1r\_\{j,1\}andrj,2r\_\{j,2\}be the distances from samplejjto its first and second nearest neighbors, and setμj=rj,2/rj,1\\mu\_\{j\}=r\_\{j,2\}/r\_\{j,1\}\. Under locally constant sampling density on anrtr\_\{t\}\-dimensional manifold,

F⁡\(μ\)=1−μ−rt,−log⁡\(1−F⁡\(μ\)\)=rt​log⁡μ\.F\(\\mu\)=1\-\\mu^\{\-r\_\{t\}\},\\qquad\-\\log\(1\-F\(\\mu\)\)=r\_\{t\}\\log\\mu\.We estimatertr\_\{t\}by fitting this slope through the origin, discarding the largest 10% of neighbor ratios for robustness\.

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\.Figure[4](https://arxiv.org/html/2610.14644#S4.F4)shows a consistent qualitative pattern: ID rises through the shallow layers, peaks in the middle\-to\-deep portion of the model, and declines gradually near the output\. The same depth\-wise ordering is visible across model scales and from an early training checkpoint onward\. Consequently, the allocation can be chosen from a short proxy run rather than requiring a fully trained base model\. Appendix[E\.2](https://arxiv.org/html/2610.14644#A5.SS2)gives the estimator, sampling, and checkpoint details\.

### 4\.2Constant\-Cost Capacity Allocation

Given a target ordering of per\-expert capacities, a naive heterogeneous model changes FFN width or expert count alone and thereby creates layerwise compute and communication imbalance\. LatentMoE provides enough degrees of freedom to vary per\-expert capacity while avoiding this imbalance\. Up to constants shared by all layers, define the routed computation, routed\-expert parameters, and one\-way dispatch volume as

Ct=Kt​ℓt​mt,Pt=Nt​ℓt​mt,Vt=Kt​ℓt\.C\_\{t\}=K\_\{t\}\\ell\_\{t\}m\_\{t\},\\qquad P\_\{t\}=N\_\{t\}\\ell\_\{t\}m\_\{t\},\\qquad V\_\{t\}=K\_\{t\}\\ell\_\{t\}\.RequiringCt=C⋆C\_\{t\}=C^\{\\star\},Pt=P⋆P\_\{t\}=P^\{\\star\}, andVt=V⋆V\_\{t\}=V^\{\\star\}implies

mt=C⋆V⋆,NtKt=P⋆C⋆,Kt=V⋆ℓt\.m\_\{t\}=\\frac\{C^\{\\star\}\}\{V^\{\\star\}\},\\qquad\\frac\{N\_\{t\}\}\{K\_\{t\}\}=\\frac\{P^\{\\star\}\}\{C^\{\\star\}\},\\qquad K\_\{t\}=\\frac\{V^\{\\star\}\}\{\\ell\_\{t\}\}\.Thusmtm\_\{t\}remains approximately constant, the sparsity ratioNt/KtN\_\{t\}/K\_\{t\}is preserved, and increasingℓt\\ell\_\{t\}is compensated by decreasing bothKtK\_\{t\}andNtN\_\{t\}\. This can raise the per\-expert ceilingqtq\_\{t\}without changing the aggregate routed budget\.

We additionally keep

Qt=2​ℓt\+Nt\+3​ms,tQ\_\{t\}=2\\ell\_\{t\}\+N\_\{t\}\+3m\_\{s,t\}nearly constant\. Multiplication byddaccounts for the two latent projections, router, and full\-width shared expert, so matchingQtQ\_\{t\}controls the remaining dominant MoE parameters and computation\. Appendix[B\.2](https://arxiv.org/html/2610.14644#A2.SS2)expands this accounting, with particular attention to expert\-parallel communication\. Finally, we use near\-square rather than exactly square experts, allowing up to a 30% aspect\-ratio deviation to align per\-layer compute and communication while satisfying grouped\-GEMM and expert\-parallel divisibility requirements\. The equalities are approximate where integer hardware constraints require rounding\. The construction is therefore heterogeneous in per\-expert capacity but nearly homogeneous in per\-layer parameter count, activated FLOPs, and communication volume\.

#### Practical two\-segment allocation\.

Although a future calibrated method could assign a distinct configuration to every layer, our experiments deliberately use only two regimes\. For a simple implementation with minimal kernel and configuration complexity, we partition the MoE blocks at a clear separation in the measured ID profile: the largest ID in the shallow prefix remains below the smallest ID in the remaining blocks\. In our profiles, the latter minimum typically occurs at the final block given the late\-depth decline\. We assign the lower\-capacity configuration to the shallow prefix and the higher\-capacity configuration to the remaining blocks\. This data\-dependent split avoids a scale\-specific numerical threshold and is an implementation choice rather than a restriction of the allocation rule\. Table[1](https://arxiv.org/html/2610.14644#S5.T1)lists the resulting configurations\.

## 5Experiments

We separate end\-task evaluation from controlled diagnosis\. We first compare complete models on downstream benchmarks in Section[5\.2](https://arxiv.org/html/2610.14644#S5.SS2)\. We then use validation loss in Section[5\.3](https://arxiv.org/html/2610.14644#S5.SS3)to isolate the effects of expert shape, allocation direction, and the regime boundary\. This distinction is important because the expected differences between pretraining architectures are small relative to the resolution of individual downstream tasks\.

### 5\.1Experimental Setup

#### Pretraining\.

The controlled shape sweeps and layerwise experiments use the models and training protocols detailed in Appendix[B](https://arxiv.org/html/2610.14644#A2)\. For the layerwise experiments, we use the configurations in Table[1](https://arxiv.org/html/2610.14644#S5.T1)and compare their ID\-guided placement with uniform and exchanged controls\. Within each comparison, the data, optimizer, and training schedule are unchanged\.

#### Compared models\.

Table[1](https://arxiv.org/html/2610.14644#S5.T1)lists the practical configurations resulting from the two\-segment allocation in Section[4](https://arxiv.org/html/2610.14644#S4)\. At 3\.4B, routed computation and routed parameters match exactly between regimes; dispatch differs by 1\.8%\. At 8B, routed computation, routed parameters, and dispatch each differ by at most 0\.32%\. We evaluate four settings at each scale, matching their total and activated parameter budgets within that scale\. The base model’s experts receive the uncompressed residual\-stream input, with width 1,280 at 8B and 1,024 at 3\.4B\. The two uniform LatentMoE models use compression factorsα=2\\alpha=2andα=4\\alpha=4, scaling\(N,K\)\(N,K\)to\(512,16\)\(512,16\)and\(1024,32\)\(1024,32\), respectively, following the original LatentMoE scaling rule\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\)\. Note that theα=2\\alpha=2baseline is already near\-square\. The layerwise model uses the two\-regime allocation in Table[1](https://arxiv.org/html/2610.14644#S5.T1)\. Appendix Table[4](https://arxiv.org/html/2610.14644#A2.T4)gives the full configurations\.

Table 1:Two\-segment ID\-guided configurations used in the experiments\. Block 0 is dense; the listed indices are MoE blocks\. Heremsm\_\{s\}is the full\-width shared\-expert intermediate dimension\.
#### Downstream protocol\.

We use the full OLMEScore\_9mcqasuite\([Gu et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib13)\): ARC\-Easy, ARC\-Challenge, BoolQ, CommonsenseQA, HellaSwag, OpenBookQA, PIQA, SocialIQA, and WinoGrande\. We follow the standard OLMES\-v0\.1 5\-shot protocol and report themc\_or\_rcscore, which selects the better of the multiple\-choice and ranked\-classification formats for each model–task pair\.

### 5\.2Main Downstream Results

Table 2:Task\-level downstream results using themc\_or\_rcscore \(%\)\. Task columns and Avg\. report means over three training runs; Avg\. additionally reports the standard deviation across runs\.SettingARC\-EARC\-CBoolQCSQAHella\.OBQAPIQASIQAWino\.Avg\.±\\pms\.d\.3\.4BBase60\.3433\.7955\.6650\.2145\.1340\.1369\.0845\.2152\.2750\.20±0\.1550\.20\\pm 0\.15LatentMoEα=2\\alpha=261\.6434\.8159\.7851\.4945\.9041\.4769\.8845\.5851\.9351\.39±0\.1651\.39\\pm 0\.16LatentMoEα=4\\alpha=461\.9333\.9361\.1151\.0045\.6239\.8069\.3745\.4751\.3851\.07±0\.2451\.07\\pm 0\.24Ours62\.1836\.0160\.5752\.1446\.8642\.4770\.6645\.1752\.1652\.02±0\.30\\mathbf\{52\.02\\pm 0\.30\}8BBase65\.6740\.1057\.1956\.7954\.6440\.7371\.5450\.0256\.9654\.85±0\.4354\.85\\pm 0\.43LatentMoEα=2\\alpha=266\.5039\.3363\.5656\.8455\.1344\.5372\.9149\.3855\.4355\.96±0\.2155\.96\\pm 0\.21LatentMoEα=4\\alpha=466\.3641\.8461\.3457\.8254\.4344\.6771\.7348\.1154\.4155\.63±0\.2055\.63\\pm 0\.20Ours68\.8441\.8160\.4960\.0656\.3747\.0072\.3449\.7855\.1256\.87±0\.37\\mathbf\{56\.87\\pm 0\.37\}

At both scales, the ID\-guided layerwise allocation achieves the highest mean score across three runs, outperforming both uniform LatentMoE baselines and the base model\. This provides a descriptive downstream complement to the matched validation\-loss evidence, with the reported standard deviations summarizing run\-level variation in the aggregate score\.

### 5\.3Analysis

Downstream accuracy is too coarse for the controlled fine\-grained ablations\. We therefore use validation loss exclusively in this subsection\.

#### Shape must be optimized before granularity\.

Across the four sweeps in Figure[3](https://arxiv.org/html/2610.14644#S3.F3), moving from an extreme aspect ratio toward the near\-square region consistently reduces validation loss\. The trend appears at both routing granularities and model scales, and survives the fixed\-msm\_\{s\}control in Appendix[C\.2](https://arxiv.org/html/2610.14644#A3.SS2)\. These comparisons isolate expert shape from the number of available and activated experts: routing granularity alone leaves the shape\-dependent per\-expert constraint in place\.

#### ID\-guided placement beats both matched endpoints\.

Table[3](https://arxiv.org/html/2610.14644#S5.T3)compares the layerwise allocation with the two uniform models obtained by repeating either of its regimes at every MoE block\. The uniform narrow model applies a lower per\-expert rank ceiling even to the higher\-ID later blocks, where it can become restrictive\. The uniform wide model raises this ceiling throughout but also assigns the same capacity to the lower\-ID shallow blocks, where the additional per\-expert capacity is less useful\. The ID\-guided model instead places the two regimes according to the observed ordinal profile and improves over both uniform endpoints at both scales\. This establishes a benefit beyond interpolation between a strong and a weak uniform model\. Since both uniform endpoints are themselves near\-square \(\|ℓ/m−1\|≤0\.4\\lvert\\ell/m\-1\\rvert\\leq 0\.4\), this comparison holds expert shape fixed and isolates the across\-depth allocation as the source of the gain\. Broader comparisons among depth\-correlated allocation signals can further refine the choice of guidance signal\.

Table 3:Configurations and validation loss for matched layerwise controls\. Narrow and wide denote the lower\- and higher\-capacity regimes in Table[1](https://arxiv.org/html/2610.14644#S5.T1)\. 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\.
#### Direction, not heterogeneity alone, produces the gain\.

The two exchanged controls isolate complementary alternatives\. The fixed\-split version preserves the ID\-guided boundary but reverses which regime appears on each side\. The matched\-count version preserves the exact multiset of narrow and wide blocks and changes only their order\. Both reversals worsen validation loss at both scales and perform below both uniform endpoints\. In particular, the matched\-count result rules out the possibility that the gain comes merely from mixing the two regimes or from using more wide blocks; within the tested controls, their shallow\-to\-deep ordering is important\. This result supports the allocation direction suggested by the ID profile\.

#### The regime boundary is robust\.

Moving the shallow–deep boundary by one block in either direction changes validation loss only slightly, and every tested boundary remains better than the stronger uniform endpoint\. Assigning the final block back to the shallow configuration also degrades performance, suggesting that the late ID decline does not warrant an additional regime switch\. A coarse two\-part allocation therefore captures most of the benefit without fitting an exact boundary to validation loss\. Appendix[D](https://arxiv.org/html/2610.14644#A4)reports the complete configurations and results\.

#### Scope of the rank proxy\.

When the four shape sweeps are grouped byq=min⁡\(ℓ,m\)q=\\min\(\\ell,m\), mean validation loss decreases asqqgrows and then plateaus, consistent with reducing a per\-expert capacity constraint\. However, transposed shapes with the sameqqare not always equivalent\. The fixed\-msm\_\{s\}experiment confirms that the broad near\-square advantage remains after removing shared\-width reallocation, but it also shows thatqqis not a complete performance model\. Figure[5](https://arxiv.org/html/2610.14644#S5.F5)provides two complementary diagnostics\. First, local Jacobians closely approachqqacross the full shape sweep, showing that the single\-expert ceiling is attained rather than vacuous\. At the mixture level, the selected expert Jacobians can collectively span the latent space, makingqqspecifically a per\-expert capacity proxy\. Second, uniform\-shape models with fixedqqexhibit depth\-wise variation in realized rank aligned with the base\-model ID profile, showing that the available ceiling and its realized use are distinct\. Moreover, transposed shapes realize nearly the same rank without reaching the same loss\. Shape orientation \(ℓ<m\\ell<mversusℓ\>m\\ell\>m\), optimization, and interactions with the shared path therefore remain relevant once the rank ceiling is no longer limiting\. Appendix[E\.1](https://arxiv.org/html/2610.14644#A5.SS1)reports the estimator, controls, and complete measurements\.

Figure 5:Realized local\-Jacobian rank\.Left:depth\-averaged rank normalized byqq; markers show means, bars show ranges, and the dashed line is the ceiling\.Right:blockwise normalized rank for two fixed\-qq3\.4B models; the dashed curve is base\-model TwoNN ID\.

## 6Related Work

Transformer capacity can be made nonuniform across layers, tokens, or experts\. AutoMoE searches over expert placement, per\-layer expert count, and FFN width under global FLOP and latency constraints\([Jawahar et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib21)\)\. Fixed\-budget methods taper MLP width or derive layerwise FFN widths from representation geometry\([Bayat et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib18);[Mudarisov et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib20)\); MAPLE reallocates routed experts by layer sensitivity\([Li et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib19)\), and Prism schedules attention heads across depth\([Aggarwal, 2026](https://arxiv.org/html/2610.14644#bib.bib17)\)\. Dynamic alternatives route selected tokens through blocks or vary token\-level expert activation\([Raposo et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib27);[Gülmez, 2026](https://arxiv.org/html/2610.14644#bib.bib26)\)\. Together, these studies establish nonuniform allocation as useful, but generally change per\-layer or per\-token workloads, or target a component other than latent expert shape\.

Within a layer, heterogeneous\-expert methods route tokens among different\-width experts and require mixed\-size execution or explicit load balancing\([Wang et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib23);[Sun et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib24);[Ma et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib25)\)\. Our construction instead uses one static expert shape and a fixed Top\-KKrule per layer, co\-varying\(ℓ,m,N,K\)\(\\ell,m,N,K\)to keep routed computation, routed parameters, and dispatch volume nearly constant\. Its schedule follows the ordinal intrinsic\-dimension profile\([Valeriani et al\., 2023](https://arxiv.org/html/2610.14644#bib.bib15);[Joshi et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib16)\), using input ID to set the local\-rank ceiling while retaining homogeneous grouped GEMMs and predictable per\-layer workloads\.

## 7Conclusion

We identify expert shape as a first\-class design variable in latent MoEs\. The per\-expert local\-rank ceiling favors a near\-square factorization, and controlled sweeps consistently support this design rule\. We further propose an intrinsic\-dimension\-guided allocation that varies per\-expert capacity across depth while preserving aggregate compute and communication\. Its practical instantiation improves matched uniform controls, whereas reversing the capacity placement removes the gain\. The rank ceiling remains a single\-expert proxy and intrinsic dimension an ordinal allocation signal\. Together, they provide practical guidance for expert shaping and coarse layerwise allocation, while a fine\-grained mapping from representation geometry to per\-layer capacity remains a promising extension\.

### AI use statement

In this work, we used generative AI tools for language polishing, LaTeX formatting, and drafting and refining code\. The authors manually reviewed and edited all AI\-assisted text, inspected and tested the generated code, and checked the resulting figures\. We did not use generative AI tools for any of the following purposes: generating synthetic datasets; developing theoretical models or conceptual frameworks; formulating mathematical claims; providing critical ingredients for proving mathematical claims; assisting in the writing of proofs; proposing or refining hypotheses; designing or providing feedback on research methodology or experiments; implementing methods; assisting with translation; cleaning or reformatting datasets; supporting qualitative or thematic data analysis; or interpreting results\. We take responsibility for the final content of this work, including all text, claims, code, and artifacts produced with the aid of generative AI\.

### Ethics statement

This work uses existing language\-model corpora and benchmarks and involves no human\-subject experiments or new personal\-data collection\. It inherits potential bias, privacy, and misuse risks from the data and general\-purpose language models\. Training also incurs computational and energy costs\.

### Reproducibility statement

Sections[2](https://arxiv.org/html/2610.14644#S2)–[5](https://arxiv.org/html/2610.14644#S5)define the architecture and budget quantities, the shaping and allocation methods, and the evaluation protocol\. Appendix[B](https://arxiv.org/html/2610.14644#A2)reports the datasets, optimization settings, and model configurations\. The remaining appendices provide complete shape sweeps, matched controls and ablations, parameter, compute, and communication accounting, and details of the realized\-rank and intrinsic\-dimension estimators\.

## References

- S\. AggarwalPrism Transformer: progressive head schedules for hierarchical attention processing\.arXiv preprint arXiv:2606\.27449\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Ansuiniet al\.\(2019\)A\. Ansuini, A\. Laio, J\. H\. Macke, and D\. ZoccolanIntrinsic dimension of data representations in deep neural networks\.InAdvances in Neural Information Processing Systems,Cited by:[§4\.1](https://arxiv.org/html/2610.14644#S4.SS1.p1.1)\.
- Bayatet al\.\(2026\)R\. Bayat, A\. Behrouz, and A\. CourvilleTapered language models\.arXiv preprint arXiv:2606\.23670\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Daiet al\.\(2024\)D\. Dai, C\. Deng, C\. Zhao, R\. X\. Xu, H\. Gao, D\. Chen, J\. Li, W\. Zeng, X\. Yu, Y\. Wu, Z\. Xie, Y\. K\. Li, P\. Huang, F\. Luo, C\. Ruan, Z\. Sui, and W\. LiangDeepSeekMoE: towards ultimate expert specialization in mixture\-of\-experts language models\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[Figure 2](https://arxiv.org/html/2610.14644#S1.F2),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Databricks AI Research Team \(2024\)Databricks AI Research TeamIntroducing DBRX: a new state\-of\-the\-art open LLM\.Note:Databricks BlogExternal Links:[Link](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm)Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[Figure 2](https://arxiv.org/html/2610.14644#S1.F2)\.
- DeepSeek\-AI \(2024a\)DeepSeek\-AIDeepSeek\-V2: a strong, economical, and efficient mixture\-of\-experts language model\.arXiv preprint arXiv:2405\.04434\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[Figure 2](https://arxiv.org/html/2610.14644#S1.F2),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- DeepSeek\-AI \(2024b\)DeepSeek\-AIDeepSeek\-V3 technical report\.arXiv preprint arXiv:2412\.19437\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§B\.1](https://arxiv.org/html/2610.14644#A2.SS1.p1.1),[Figure 2](https://arxiv.org/html/2610.14644#S1.F2),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Elangoet al\.\(2026\)V\. Elango, N\. Bhatia, R\. Waleffe, R\. Shafipour, T\. Asida, A\. Khattar, N\. Assaf, M\. Golub, J\. Guman, T\. Mitra, R\. Zhao, R\. Borkar, R\. Zilberstein, M\. Patwary, M\. Shoeybi, and B\. RouhaniLatentMoE: toward optimal accuracy per FLOP and parameter in mixture of experts\.arXiv preprint arXiv:2601\.18089\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§B\.3](https://arxiv.org/html/2610.14644#A2.SS3.p1.1),[§1](https://arxiv.org/html/2610.14644#S1.p3.1),[§2](https://arxiv.org/html/2610.14644#S2.p2.1),[§2](https://arxiv.org/html/2610.14644#S2.p3.2),[§5\.1](https://arxiv.org/html/2610.14644#S5.SS1.SSS0.Px2.p1.1)\.
- Faccoet al\.\(2017\)E\. Facco, M\. d’Errico, A\. Rodriguez, and A\. LaioEstimating the intrinsic dimension of datasets by a minimal neighborhood information\.Scientific Reports\.Cited by:[§4\.1](https://arxiv.org/html/2610.14644#S4.SS1.p1.1)\.
- Feduset al\.\(2022\)W\. Fedus, B\. Zoph, and N\. ShazeerSwitch transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Guet al\.\(2025\)Y\. Gu, O\. Tafjord, B\. Kuehl, D\. Haddad, J\. Dodge, and H\. HajishirziOLMES: a standard for language model evaluations\.InFindings of the Association for Computational Linguistics: NAACL 2025,Cited by:[§5\.1](https://arxiv.org/html/2610.14644#S5.SS1.SSS0.Px3.p1.1)\.
- Gülmez \(2026\)G\. GülmezDynaMoE: dynamic token\-level expert activation with layer\-wise adaptive capacity for mixture\-of\-experts neural networks\.arXiv preprint arXiv:2603\.01697\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Jawaharet al\.\(2023\)G\. Jawahar, S\. Mukherjee, X\. Liu, Y\. J\. Kim, M\. Abdul\-Mageed, L\. V\. S\. Lakshmanan, A\. H\. Awadallah, S\. Bubeck, and J\. GaoAutoMoE: heterogeneous mixture\-of\-experts with adaptive computation for efficient neural machine translation\.InFindings of the Association for Computational Linguistics: ACL 2023,Cited by:[§1](https://arxiv.org/html/2610.14644#S1.p4.1),[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Jianget al\.\(2024\)A\. Q\. Jiang, A\. Sablayrolles, A\. Roux, A\. Mensch, B\. Savary, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, E\. B\. Hanna, F\. Bressand, G\. Lengyel, G\. Bour, G\. Lample, L\. R\. Lavaud, L\. Saulnier, M\. Lachaux, P\. Stock, S\. Subramanian, S\. Yang, S\. Antoniak, T\. Le Scao, T\. Gervet, T\. Lavril, T\. Wang, T\. Lacroix, and W\. El SayedMixtral of experts\.arXiv preprint arXiv:2401\.04088\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[Figure 2](https://arxiv.org/html/2610.14644#S1.F2),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Joshiet al\.\(2025\)A\. Joshi, D\. Bhatt, and A\. ModiGeometry of decision making in language models\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2610.14644#S1.p4.1),[§4\.1](https://arxiv.org/html/2610.14644#S4.SS1.p1.1),[§4](https://arxiv.org/html/2610.14644#S4.p1.1),[§6](https://arxiv.org/html/2610.14644#S6.p2.1)\.
- Kimi Team \(2026\)Kimi TeamKimi K3: open frontier intelligence\.arXiv preprint arXiv:2607\.24653\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§2](https://arxiv.org/html/2610.14644#S2.p3.2)\.
- Lepikhinet al\.\(2021\)D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. ChenGShard: scaling giant models with conditional computation and automatic sharding\.InInternational Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Liet al\.\(2026\)L\. Li, W\. Li, J\. Shen, and G\. HuMAPLE: MoE adaptive plug\-and\-play layer\-wise expert allocation\.arXiv preprint arXiv:2608\.15299\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Liuet al\.\(2025\)Z\. Liu, H\. Wu, R\. She, X\. Fu, X\. Han, T\. Zhong, and M\. YuanMoLAE: mixture of latent experts for parameter\-efficient language models\.arXiv preprint arXiv:2503\.23100\.Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1)\.
- Maet al\.\(2026\)Z\. Ma, X\. Liu, Z\. Liu, N\. Wang, Y\. Shen, K\. Wang, S\. Shi, and S\. LianMixture of heterogeneous grouped experts for language modeling\.InProceedings of the 64th Annual Meeting of the Association for Computational Linguistics \(Volume 6: Industry Track\),Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p2.1)\.
- Mudarisovet al\.\(2026\)T\. Mudarisov, M\. Burtsev, and R\. StateGeometry\-guided layerwise FFN width allocation in transformers\.arXiv preprint arXiv:2608\.02064\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Raposoet al\.\(2024\)D\. Raposo, S\. Ritter, B\. Richards, T\. Lillicrap, P\. C\. Humphreys, and A\. SantoroMixture\-of\-depths: dynamically allocating compute in transformer\-based language models\.arXiv preprint arXiv:2404\.02258\.Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p1.1)\.
- Shazeeret al\.\(2017\)N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. DeanOutrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.InInternational Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Shazeer \(2020\)N\. ShazeerGLU variants improve transformer\.arXiv preprint arXiv:2002\.05202\.Cited by:[§3\.1](https://arxiv.org/html/2610.14644#S3.SS1.p1.2)\.
- Suet al\.\(2025\)D\. Su, K\. Kong, Y\. Lin, J\. Jennings, B\. Norick, M\. Kliegl, M\. Patwary, M\. Shoeybi, and B\. CatanzaroNemotron\-CC: transforming common crawl into a refined long\-horizon pretraining dataset\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),Cited by:[§B\.1](https://arxiv.org/html/2610.14644#A2.SS1.p1.1)\.
- Sunet al\.\(2024\)M\. Sun, W\. Liu, J\. Luan, P\. Gao, and B\. WangMixture of diverse size experts\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track,Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p2.1)\.
- Team \(2025\)Q\. TeamQwen3 technical report\.External Links:2505\.09388,[Link](https://arxiv.org/abs/2505.09388)Cited by:[Figure 2](https://arxiv.org/html/2610.14644#S1.F2)\.
- The Microsoft AI Team \(2026\)The Microsoft AI TeamMAI\-Thinking\-1: building a hill\-climbing machine\.Technical reportMicrosoft AI\.External Links:[Link](https://microsoft.ai/pdf/mai-thinking-1.pdf)Cited by:[Appendix A](https://arxiv.org/html/2610.14644#A1.p1.1),[§2](https://arxiv.org/html/2610.14644#S2.p3.2)\.
- Valerianiet al\.\(2023\)L\. Valeriani, D\. Doimo, F\. Cuturello, A\. Laio, A\. Ansuini, and A\. CazzanigaThe geometry of hidden representations of large transformer models\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2610.14644#S1.p4.1),[§4\.1](https://arxiv.org/html/2610.14644#S4.SS1.p1.1),[§4](https://arxiv.org/html/2610.14644#S4.p1.1),[§6](https://arxiv.org/html/2610.14644#S6.p2.1)\.
- Vaswaniet al\.\(2017\)A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. PolosukhinAttention is all you need\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2610.14644#S1.p1.1)\.
- Wanget al\.\(2025\)A\. Wang, X\. Sun, R\. Xie, S\. Li, J\. Zhu, Z\. Yang, P\. Zhao, W\. Han, Z\. Kang, D\. Wang, N\. Okazaki, and C\. XuHMoE: heterogeneous mixture of experts for language modeling\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,Cited by:[§6](https://arxiv.org/html/2610.14644#S6.p2.1)\.

## Appendix AAdditional Related Work

Sparse MoEs scale parameter capacity by routing each token to a small subset of experts\([Shazeer et al\., 2017](https://arxiv.org/html/2610.14644#bib.bib1)\)\. GShard and Switch Transformer established scalable expert parallelism and load\-balancing strategies for Transformer models\([Lepikhin et al\., 2021](https://arxiv.org/html/2610.14644#bib.bib2);[Fedus et al\., 2022](https://arxiv.org/html/2610.14644#bib.bib3)\), and later systems such as Mixtral and DBRX demonstrated the effectiveness of sparse MoEs at large language\-model scale\([Jiang et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib4);[Databricks AI Research Team, 2024](https://arxiv.org/html/2610.14644#bib.bib22)\)\. DeepSeekMoE and its successors pursue finer expert segmentation by increasingNNandKKwhile reducing individual expert width\([Dai et al\., 2024](https://arxiv.org/html/2610.14644#bib.bib5);[DeepSeek\-AI, 2024a](https://arxiv.org/html/2610.14644#bib.bib6);[DeepSeek\-AI, 2024b](https://arxiv.org/html/2610.14644#bib.bib7)\)\. This accounting couples routing granularity to increasingly narrow experts\. LatentMoE instead uses a projected space to decoupleNNandKKfrom the residual width\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\), and appears in systems including MAI\-Thinking\-1 and Kimi K3\([The Microsoft AI Team, 2026](https://arxiv.org/html/2610.14644#bib.bib10);[Kimi Team, 2026](https://arxiv.org/html/2610.14644#bib.bib11)\)\. Related latent\-factorization work compresses pretrained MoE experts into shared projections and expert\-specific low\-dimensional transformations\([Liu et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib31)\); it targets post\-training parameter efficiency rather than expert shape and routed communication during pretraining\.

These developments primarily optimize routing, expert count, compression, or aggregate budget\. Our work is orthogonal: it isolates the two\-dimensional factorizationℓ×m\\ell\\times m, relates it to a per\-expert local\-rank ceiling, and allocates that capacity across depth\. Rank matching can therefore be combined with existing routing and load\-balancing methods\.

## Appendix BTraining, Evaluation, and Systems Details

### B\.1Pretraining Details

The shape sweeps use a 3\.4B\-A250M model with hidden size 1,024 and 12 Transformer blocks, and an 8B\-A560M model with hidden size 1,280 and 16 blocks\. The first block is dense and all remaining blocks use sparse LatentMoE\. The backbone follows the DeepSeek\-V3 architecture\([DeepSeek\-AI, 2024b](https://arxiv.org/html/2610.14644#bib.bib7)\), with its routed experts replaced by the LatentMoE structure described in Section[2](https://arxiv.org/html/2610.14644#S2)\. Both scales use a sequence length of 4,096 and the Adam optimizer with\(β1,β2\)=\(0\.9,0\.95\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.9,0\.95\)\. The 3\.4B\-A250M model uses a learning rate of4×10−44\\times 10^\{\-4\}, a batch size of 2M tokens, and 1,500 warmup steps; the 8B\-A560M model uses a learning rate of3×10−43\\times 10^\{\-4\}, a batch size of 2M tokens, and 2,000 warmup steps\. The models train for 20B and 40B tokens, respectively\. We train on NVIDIA Nemotron\-CC\([Su et al\., 2025](https://arxiv.org/html/2610.14644#bib.bib14)\)\. Following DeepSeek\-V3, we combine an auxiliary load\-balancing loss with coefficient10−410^\{\-4\}and expert\-bias updates with rate10−310^\{\-3\}\. Within a sweep, the data mixture, initialization protocol, optimizer configuration, batch construction, router, attention blocks, and expert activation are held fixed\. The layerwise experiments use the same scale\-specific training setup; within each set of controls, only the placement of the MoE configurations in Table[1](https://arxiv.org/html/2610.14644#S5.T1)changes\.

#### Run\-to\-run variation\.

We estimate the validation\-loss noise scale from five independent training runs of the base configuration under the same protocol\. Their final validation losses have a standard deviation of1\.2×10−31\.2\\times 10^\{\-3\}at 3\.4B\-A250M and1\.5×10−31\.5\\times 10^\{\-3\}at 8B\-A560M\. Because repeating every point in the two\-scale shape sweep would be computationally prohibitive, each shape is trained once\. The largest within\-sweep gaps,1\.04×10−21\.04\\times 10^\{\-2\}–3\.99×10−23\.99\\times 10^\{\-2\}, are roughly 7–27 times this empirical noise scale, making the observed separation between extreme and near\-square shapes substantially larger than ordinary run\-to\-run variation\.

#### Downstream model configurations\.

Table[4](https://arxiv.org/html/2610.14644#A2.T4)lists the four downstream settings at each scale\. Total and activated parameter budgets are matched within each scale; internal checkpoint identifiers are omitted\.

Table 4:Models used in the downstream evaluation\. For uniform models,α=d/ℓ\\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](https://arxiv.org/html/2610.14644#S5.T1)\.

### B\.2Parameter, Compute, and Communication Accounting

Section[4](https://arxiv.org/html/2610.14644#S4)defines the matched quantities up to constants shared across layers\. Here we expand those constants and distinguish model parameters, activated computation, and expert\-parallel communication\. We consider one LatentMoE layer; unchanged attention and backbone components are outside this comparison\. A SwiGLU expert contains three weight matrices: two input projections and one output projection\.

Figure 6:Expert\-parallel communication paths\.Top:a standard MoE dispatches and returnsKKvectors of widthddper 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 widthddtoℓ\\ell\. Gray dashed arrows show the local routing computation and Top\-KKassignments; colored arrows show the activation path\.#### Parameter count\.

TheNtN\_\{t\}routed experts contain3​Nt​ℓt​mt3N\_\{t\}\\ell\_\{t\}m\_\{t\}parameters\. The latent down\- and up\-projections contain2​d​ℓt2d\\ell\_\{t\}parameters, the router containsd​NtdN\_\{t\}, and the full\-width shared expert contains3​d​ms,t3dm\_\{s,t\}\. The layer\-dependent unique trainable parameters are therefore

Paramst≈3​Nt​ℓt​mt\+d⁡\(2​ℓt\+Nt\+3​ms,t\)=3​Pt\+d​Qt\.\\operatorname\{Params\}\_\{t\}\\approx 3N\_\{t\}\\ell\_\{t\}m\_\{t\}\+d\\\!\\left\(2\\ell\_\{t\}\+N\_\{t\}\+3m\_\{s,t\}\\right\)=3P\_\{t\}\+dQ\_\{t\}\.Thus,PtP\_\{t\}in the main text controls the routed\-expert parameter pool, whileQtQ\_\{t\}controls the remaining dominant parameters\.

#### Activated computation\.

For one token, onlyKtK\_\{t\}routed experts are evaluated\. Counting multiply–accumulate operations \(MACs\), the corresponding layer cost is

MACst/token≈3​Kt​ℓt​mt\+d⁡\(2​ℓt\+Nt\+3​ms,t\)=3​Ct\+d​Qt\.\\operatorname\{MACs\}\_\{t\}/\\text\{token\}\\approx 3K\_\{t\}\\ell\_\{t\}m\_\{t\}\+d\\\!\\left\(2\\ell\_\{t\}\+N\_\{t\}\+3m\_\{s,t\}\\right\)=3C\_\{t\}\+dQ\_\{t\}\.Under the convention that one MAC is two floating\-point operations, the FLOP count is twice this quantity\. MatchingCtC\_\{t\}andQtQ\_\{t\}therefore matches the dominant activated computation\. Activation functions, routing softmax, weighted aggregation, and other lower\-order elementwise operations are not included in this approximation\.

#### Expert\-parallel communication\.

Figure[6](https://arxiv.org/html/2610.14644#A2.F6)contrasts the communication paths\. In a standard MoE, routing one token toKKexperts dispatchesKKcopies of add\-dimensional representation to the workers hosting those experts\. The one\-way logical payload is thereforeK​dKdscalars per token\. The selected expert outputs have the same width and produce anotherK​dKdscalars on the return all\-to\-all\.

LatentMoE applies the down\-projection on the source worker before dispatch\. Each expert assignment consequently sends anℓt\\ell\_\{t\}\-dimensional vector rather than add\-dimensional one\. Expert outputs return in the same latent space and are aggregated before the up\-projection\. The shared projections therefore remain outside the all\-to\-all, giving

Vtone\-way=Kt​ℓtandVtround\-trip=2​Kt​ℓtscalars/token\.V\_\{t\}^\{\\text\{one\-way\}\}=K\_\{t\}\\ell\_\{t\}\\quad\\text\{and\}\\quad V\_\{t\}^\{\\text\{round\-trip\}\}=2K\_\{t\}\\ell\_\{t\}\\qquad\\text\{scalars/token\}\.ForTTrouted tokens represented withbbbytes per scalar, these becomeb​T​Kt​ℓtbTK\_\{t\}\\ell\_\{t\}bytes one way and2​b​T​Kt​ℓt2bTK\_\{t\}\\ell\_\{t\}bytes for dispatch plus return\. Routing requiresO⁡\(T​Kt\)O\(TK\_\{t\}\)indices and weights locally\. Implementations that transmit this metadata add anO⁡\(T​Kt\)O\(TK\_\{t\}\)term, which is small relative to theO⁡\(T​Kt​ℓt\)O\(TK\_\{t\}\\ell\_\{t\}\)activation payload at the widths studied here\.

### B\.3System Throughput

As shown in Table[5](https://arxiv.org/html/2610.14644#A2.T5), we additionally compare end\-to\-end training throughput\. Our heterogeneous allocation achieves throughput within measurement variation of the uniform LatentMoE baseline, indicating that layerwise capacity heterogeneity introduces negligible additional systems overhead\. All LatentMoE variants exhibit a modest throughput gap relative to the base model under our current implementation, which does not use dedicated kernel optimizations\. This gap is consistent with implementation\-level overhead from finer\-grained routing, smaller grouped GEMMs, and the latent projections\. Prior LatentMoE analysis shows that throughput\-oriented large scale distributed MoEs can become strongly communication\-bound and kernel\-related overhead constitutes a smaller fraction of end\-to\-end execution time and may therefore be better amortized\([Elango et al\., 2026](https://arxiv.org/html/2610.14644#bib.bib9)\)\. Thus, the heterogeneous allocation has a near\-homogeneous execution cost relative to uniform LatentMoE, although LatentMoE itself retains a small wall\-clock overhead relative to the base model\.

Table 5:End\-to\-end training throughput\.

## Appendix CAdditional Expert\-Shaping Results

### C\.1Full Shape\-Sweep Results

Table 6:Complete configurations and validation losses for Figure[3](https://arxiv.org/html/2610.14644#S3.F3)\.K/NK/Ndenotes the number of activated experts over the total number of experts\.
### C\.2Fixed Shared\-Expert\-Width Control

To remove shared\-path reallocation as a confound, we repeat the 3\.4B\-A250M\(N,K\)=\(256,8\)\(N,K\)=\(256,8\)sweep withms=768m\_\{s\}=768and the per\-expert productℓ​m≈393,216\\ell m\\approx 393\{,\}216fixed\. All other training choices match the corresponding shape sweep\. Because the down\- and up\-projection parameter count grows linearly withℓ\\ell, larger latent widths receive a modest parameter advantage in this control\. Table[7](https://arxiv.org/html/2610.14644#A3.T7)reports this advantage relative toℓ=256\\ell=256as a fraction of the base model’s total trainable parameters, rather than as a fraction of activated parameters\. The largest advantage is only 0\.47%\.

Table 7:Validation loss in the fixed\-msm\_\{s\}control\. The parameter advantage is the additional down/up\-projection parameters relative toℓ=256\\ell=256, normalized by the total parameter count of the 3\.4B\-A250M base model\.The validation\-loss span is 0\.0151, compared with 0\.0166 in the parameter\-matched sweep, and the broad near\-square advantage is preserved\. Thus the shape effect remains even though the fixed\-msm\_\{s\}design increasingly favors settings with larger latent widths\.

## Appendix DAdditional Layerwise\-Allocation Ablations

Table[8](https://arxiv.org/html/2610.14644#A4.T8)reports the 3\.4B sensitivity experiments under the same training protocol as the main layerwise model\. Allocations are listed from shallow to deep, and each coefficient gives the number of MoE blocks using that configuration\.

Table 8:Boundary and shape sensitivity at 3\.4B\. The first four rows use the narrow and wide regimes from Table[1](https://arxiv.org/html/2610.14644#S5.T1); the final row is a matched\-budget shape stress test\. Lower loss is better\.Moving the boundary by one block changes validation loss by at most1\.9×10−31\.9\\times 10^\{\-3\}, whereas returning the final block to the narrow regime incurs a larger penalty\. The strongly unbalanced matched\-budget control performs substantially worse, consistent with the near\-square design rule\.

## Appendix ECapacity Diagnostics

### E\.1Realized Rank

#### Estimator\.

For a SwiGLU expert, let𝒂=Wg​𝒛\\bm\{a\}=W\_\{g\}\\bm\{z\}and𝒃=Wu​𝒛\\bm\{b\}=W\_\{u\}\\bm\{z\}\. We evaluate the exact local Jacobian

JE​\(𝒛\)=Wo​\[diag⁡\(𝒃⊙ϕ′​\(𝒂\)\)​Wg\+diag⁡\(ϕ⁡\(𝒂\)\)​Wu\]J\_\{E\}\(\\bm\{z\}\)=W\_\{o\}\\\!\\left\[\\operatorname\{diag\}\\\!\\big\(\\bm\{b\}\\odot\\phi^\{\\prime\}\(\\bm\{a\}\)\\big\)W\_\{g\}\+\\operatorname\{diag\}\\\!\\big\(\\phi\(\\bm\{a\}\)\\big\)W\_\{u\}\\right\]at tokens that were actually routed to the expert\. The derivative includes the implementation’s activation\-clamp masks, although no sampled activation reaches the clamp\. For singular valuesσ1≥⋯≥σs\\sigma\_\{1\}\\geq\\cdots\\geq\\sigma\_\{s\}, we estimate rank with the relative threshold

rankτ⁡\(A\)=\#⁡\{j:σj​\(A\)\>τ​σ1​\(A\)\},τ=10−3\.\\operatorname\{rank\}\_\{\\tau\}\(A\)=\\\#\\\{j:\\sigma\_\{j\}\(A\)\>\\tau\\sigma\_\{1\}\(A\)\\\},\\qquad\\tau=10^\{\-3\}\.We apply the same estimator to each expert weight, to centered matrices of expert inputs and outputs over routed tokens, and to the combined expert\-map Jacobian∑i∈TKpi​JEi​\(𝒛\)\\sum\_\{i\\in T\_\{K\}\}p\_\{i\}J\_\{E\_\{i\}\}\(\\bm\{z\}\)\. The last quantity holds the selected experts and routing weights locally fixed; it tests whether combining the selected expert maps spans the latent space, rather than differentiating through the router\.

#### Sampling and controls\.

We collect 120,219 non\-padding tokens from the same held\-out data mixture used for the ID study\. In each MoE block we sample 16 experts with at least 512 routed tokens, evaluate 16 token\-level Jacobians per expert, and retain up to 4,096 routed tokens per expert for the input and output clouds\. The combined expert\-map statistic uses 24 tokens per block\. Checkpoint weights are loaded in their native bfloat16 representation and converted to float32 for the SVD\.

Two controls test the estimator\. First, we construct matrices with known ranks0\.25​q0\.25q,0\.5​q0\.5q, and0\.9​q0\.9q, geometrically decaying nonzero spectra, and then round them to bfloat16\. The estimator recovers the true rank exactly in 33/33 cases, ruling out a spurious full\-rank floor from bfloat16 quantization\. Second, shape\-matched iid Gaussian experts attain Jacobian ranks between 0\.953qqand 1\.000qq\. Their small relative deficit near square shapes shows that the thresholded estimate there is affected by the expected poorer conditioning, rather than implying unused trained channels\.

#### Full shape\-sweep measurements\.

Tables[9](https://arxiv.org/html/2610.14644#A5.T9)and[10](https://arxiv.org/html/2610.14644#A5.T10)use exactly the configurations and row order of Table[6](https://arxiv.org/html/2610.14644#A3.T6); validation loss is not repeated here\.YYdenotes the centered cloud of expert outputs,ZZthe centered latent inputs to a block, andJmixJ\_\{\\mathrm\{mix\}\}the combined expert\-map Jacobian defined above\. Every entry is a mean over MoE depth\.

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\.Table 10:Realized ranks for the complete 8B\-A560M shape sweep, using the same definitions as Table[9](https://arxiv.org/html/2610.14644#A5.T9)\.
#### The per\-expert ceiling is attained\.

Across the 27 sweep checkpoints, 1001 of 1059 block–matrix means equalqqexactly, and the minimum is 0\.992qq\. Local Jacobians attain 0\.909–1\.000qq, while expert output clouds attain 0\.988–1\.000qq\. The latent inputs themselves span 0\.995–1\.000 of their ambient width\. The full\-rank weights and non\-degenerate inputs makeq=min⁡\(ℓ,m\)q=\\min\(\\ell,m\)an empirically relevant upper bound in trained models\. Near\-square experts can show a slightly larger relative thresholding deficit, as they do in the Gaussian control, but their largerqqstill yields the largest absolute realized ranks\.

Table 11:Ceiling–realized\-rank relation within each controlled shape sweep\.JJis the depth\-averaged realized local\-Jacobian rank\. The final column counts all within\-sweep pairs with distinct ceilings for which a largerqqalso gives a largerJJ\.Table[11](https://arxiv.org/html/2610.14644#A5.T11)separates the four sweeps so that model scale, routing configuration, training protocol, and approximate per\-expert parameter budget are fixed within each row\. Across all 72 within\-sweep pairs with different ceilings, the configuration with the largerqqrealizes the larger absoluteJJ\. At the same time, the variation inJ/qJ/qconfirms that the ceiling and realized rank remain distinct\. The high correlations show that trained experts make use of an increased ceiling, while optimization and validation loss retain additional dependencies\.

#### Per\-expert capacity versus mixture\-level span\.

The combined expert\-map statistic is already 0\.931–0\.994 ofℓ\\ellin every configuration\. This is expected and does not contradict the per\-expert capacity hypothesis\. Holding the selected experts and routing weights locally fixed,

rank⁡\(Jmix\)≤min⁡\(ℓ,∑i∈TKrank⁡\(JEi\)\)≤min⁡\(ℓ,K​q\)\.\\operatorname\{rank\}\(J\_\{\\mathrm\{mix\}\}\)\\leq\\min\\left\(\\ell,\\sum\_\{i\\in T\_\{K\}\}\\operatorname\{rank\}\(J\_\{E\_\{i\}\}\)\\right\)\\leq\\min\(\\ell,Kq\)\.Every shape\-sweep configuration satisfiesK​q≥ℓKq\\geq\\ell, so this mixture\-level bound is non\-restrictive: the selected experts can collectively span the latent space even though each constituent expert is individually limited byqq\. Full mixture rank therefore certifies aggregate local span; it does not imply that reducing the per\-expert ceiling is lossless or that the constituent conditional maps have equal capacity\. TheJmixJ\_\{\\mathrm\{mix\}\}statistic consequently rules out a mixture\-level span bottleneck, but it is not a test of the per\-expert capacity hypothesis\.

Across 9 within\-sweep transposed pairs\(ℓ,m\)\(\\ell,m\)and\(m,ℓ\)\(m,\\ell\), which share the sameqq, mean Jacobian rank differs by at most 3\.4% ofqq, while Table[6](https://arxiv.org/html/2610.14644#A3.T6)shows different validation losses\. The rank proxy is therefore consistent with one aspect of local expert capacity\. Shape orientation, optimization, interactions with the shared path, and mixture composition account for the remaining variation\.

#### Interpretation for layerwise allocation\.

The allocation hypothesis concerns tolerance to per\-expert compression rather than numerical rank matching\. We do not equate ID withqtq\_\{t\}, nor do we require a low\-ID layer’s mixture Jacobian to be rank deficient\. Instead, the directional hypothesis is that reducing the per\-expert ceiling incurs a smaller loss penalty in lower\-ID regions than in higher\-ID regions\. The matched\-count exchange preserves the same multiset of narrow and wide block configurations while changing their placement across depth, providing a controlled test of this directional claim rather than a calibrated mapping from ID values to rank ceilings\.

#### Depth\-wise evidence and the final\-block exception\.

Uniform\-shape models keepqqfixed across depth, so any layerwise variation in realized rank is not induced by a changing rank budget\. Correlating each checkpoint’s block\-averaged local\-Jacobian rank with the corresponding base\-model TwoNN profile gives a positive Spearman correlation in 27/27 checkpoints, with medianρ=\+0\.76\\rho=\+0\.76\. Numerical rank is already near saturation, so the variation is small; this is directional rather than magnitude evidence\. The final MoE block frequently departs from the preceding trend\. Reporting the prespecified sensitivity analysis rather than silently dropping it, the correlation is positive in 27/27 checkpoints with medianρ=\+0\.85\\rho=\+0\.85when that block is excluded\. The boundary ablation in Section[5\.3](https://arxiv.org/html/2610.14644#S5.SS3)nevertheless favors the higher\-capacity configuration for the final block, consistent with its ID declining only mildly and remaining well above the low\-ID shallow prefix \(Table[13](https://arxiv.org/html/2610.14644#A5.T13)\)\.

Finally, the two\-regime models produce the intended absolute\-rank step while retaining the matched routed budgets\. Table[12](https://arxiv.org/html/2610.14644#A5.T12)shows that the deep regime realizes a larger local Jacobian rank at both scales\.

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

The segment averages above verify that the allocated rank changes across the regime boundary\. To separate this designed change inqtq\_\{t\}from a possible effect of heterogeneous training, we additionally compare the 3\.4B two\-segment checkpoint at every block with the uniform checkpoint that repeats the same\(ℓ,m,K,N\)\(\\ell,m,K,N\)configuration at all blocks\. Thus, the matched counterfactual is the uniform narrow model in blocks 1–4 and the uniform wide model in blocks 5–11\. Figure[7](https://arxiv.org/html/2610.14644#A5.F7)shows that the ID\-guided checkpoint closely follows this piecewise uniform reference: the mean signed residual is−0\.10%\-0\.10\\%ofqq, the mean absolute residual is0\.38%0\.38\\%ofqq, and the maximum absolute residual is1\.54%1\.54\\%ofqq\. This is a mechanism check that the heterogeneous model realizes its allocated ceiling similarly to the corresponding uniform models\.

Figure 7:Matched\-endpoint realized\-rank check at 3\.4B\.Left:the ID\-guided checkpoint’s blockwise local\-Jacobian numerical rank, its allocated ceilingqtq\_\{t\}, and the realized ranks of the two uniform endpoint checkpoints\. The uniform narrow and wide checkpoints repeat their respective\(ℓ,m,K,N\)\(\\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 byqtq\_\{t\}\.

### E\.2Intrinsic\-Dimension Estimation

At each block boundary, we collect 20,000 non\-padding token representations from the same held\-out corpus and randomly retain 5,000 for TwoNN estimation\. Both scales use a maximum sequence length of 2,048, seed 0, and mean reduction over the model’s parallel chain dimension\. We exclude the embedding output from regime selection because repeated discrete token vectors do not satisfy TwoNN’s local continuous\-manifold assumption\.

The cross\-scale comparison uses the final 3\.4B and 8B base checkpoints\. Table[13](https://arxiv.org/html/2610.14644#A5.T13)gives the complete layerwise values, including the embedding outputL​0L0for reference\.

Table 13:Layerwise TwoNN intrinsic\-dimension values underlying the left panel of Figure[4](https://arxiv.org/html/2610.14644#S4.F4)\. Relative depth ist/Lt/L, whereLLis the number of Transformer blocks;L​0L0denotes the embedding output\.To measure evolution during training, we use a separate 8B checkpoint series at 2k, 4k, 10k, 16k, and 20k steps under the same estimation protocol\. Table[14](https://arxiv.org/html/2610.14644#A5.T14)reports every layer and checkpoint\.

Table 14:Layerwise TwoNN intrinsic dimension over training for the separate 8B checkpoint series underlying the right panel of Figure[4](https://arxiv.org/html/2610.14644#S4.F4)\.Across the two model sizes and all five 8B checkpoints, the depth\-wise trends are close: ID rises through the early blocks, peaks in the middle\-to\-deep region, and declines toward the output\.