PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.14644v1cs.AIcs.LG

TL;DR

Fine-grained MoEs often make expert shape a consequence of routing granularity, leaving its independent role in expert capacity underexplored. This paper analyzes expert shape and uses layerwise intrinsic dimension to guide capacity allocation across depth; controlled sweeps favor near-square experts, and ID-guided allocation improves validation loss and downstream performance over layer-homogeneous baselines.

  • Problem

    In standard MoEs, expert input dimension is fixed and increasing expert count narrows intermediate width, leaving expert shape tied to routing granularity rather than an independent design choice.

    Source: §1 · p. 2

  • Method

    The paper derives a local-rank ceiling to guide near-square expert shapes, then uses layerwise intrinsic-dimension ordering to allocate expert capacity across depth under matched resource constraints.

    Source: §1 · pp. 2–3, §3.1 · p. 4

  • Results

    Near-square shapes consistently reduce validation loss across model scales and routing granularities, while ID-guided allocation achieves the highest mean downstream score across three runs at both scales.

    Source: §5.3 · p. 7, §5.2 · p. 7

  • Takeaways & Limitations

    Expert shape and coarse layerwise capacity allocation provide practical design guidance for latent MoEs while preserving aggregate compute and communication.

    Source: §7 · p. 9

    The local-rank ceiling is only a single-expert proxy and intrinsic dimension an ordinal allocation signal, so a fine-grained mapping from representation geometry to per-layer capacity remains unresolved.

    Source: §7 · p. 9, §5.3 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

LatentMoE makes expert input dimension and intermediate width explicit design choices, revealing that expert shape and layerwise capacity allocation can improve MoE performance. The paper motivates near-square experts through a Jacobian-rank bound and uses layerwise intrinsic dimension to allocate capacity across depth under matched system budgets.

1 INTRODUCTION

  • In standard fine-grained MoEs, matching activated expert width to the dense FFN budget narrows each expert as expert count rises, potentially limiting its representational capacity despite unchanged aggregate activated width.Standard experts inherit the model hidden dimension, so increasing expert count makes the input-to-intermediate aspect ratio increasingly large.Source: §1 · p. 2
  • Controlled shape sweeps show a U-shaped validation-loss curve over expert aspect ratio, with the best performance near square, where the fixed-budget Jacobian-rank ceiling is largest.The ceiling is bounded by both expert input dimension and intermediate width, and balancing them maximizes it.Source: §1 · pp. 2–3
  • LatentMoE decouples expert input dimension from model hidden size by performing expert computation in a lower-dimensional latent space, exposing expert shape as a design choice.This work revisits its original fixed four-fold compression by jointly selecting expert input and intermediate dimensions under a fixed per-expert parameter budget.Source: §1 · p. 2
  • ID-guided allocation co-varies expert number and shape across layers, improving validation loss and downstream performance while keeping per-layer parameter count, activated FLOPs, and communication volume nearly constant.Higher-ID layers are hypothesized to benefit from greater per-expert capacity; the allocation follows ID ordering under matched-budget and hardware constraints.Source: §1 · pp. 2–3

2 PRELIMINARIES

LatentMoE decouples expert input dimension from residual width, making expert shape configurable at fixed per-expert parameter count. The section introduces the routed-path accounting and motivates choosing expert shape and allocating it layerwise according to input intrinsic dimension.

2 PRELIMINARIES

  • A standard MoE’s routed experts use parameters proportional to Ndm, activated computation proportional to Kdm, and communication proportional to Kd.The residual width d determines both expert input dimension and communicated representation width.Source: §2 · p. 3
  • With latent dimension ℓ, routed-path parameters, activated computation, and communication scale as Nℓm, Kℓm, and Kℓ, respectively.The routed path, including dispatch and aggregation, operates in the ℓ-dimensional space.Source: §2 · p. 3
  • Setting N′ = αN and K′ = αK at fixed m, where α = d/ℓ, preserves routed-path parameters, computation, and communication while providing α times as many candidate and activated experts.The original study selected α = 4 empirically (Elango et al., 2026), but gave no general principle for jointly choosing ℓ and m.Source: §2 · p. 3
  • LatentMoE makes expert shape ℓ×m configurable independently of residual width d; at fixed ℓm, its aspect ratio can vary without changing per-expert parameters.The paper asks which shape is most effective when shared across layers and when allocated layerwise to match input intrinsic dimension.Source: §2 · p. 3

3 EXPERT SHAPING

An expert’s local Jacobian rank is bounded by min(ℓ,m), with the ceiling maximized by square shapes at fixed per-expert parameters. Controlled sweeps favor near-square experts across the tested model scales and routing granularities.

3.1 A PER-EXPERT LOCAL-RANK CEILING

  • For gated FFN experts, local Jacobian rank is bounded by q = min(ℓ, m), which is maximized at ℓ = m when the per-expert parameter product ℓm is fixed.The nonlinearity changes the local linearization but cannot bypass the m-dimensional intermediate representation.Source: §3.1 · p. 4
  • Across shape sweeps, trained local Jacobians nearly attain q, supporting the local-rank ceiling as an empirically relevant capacity proxy.Expert weights and routed-token outputs are also essentially full rank.Source: §3.1 · p. 4

3.2 CONTROLLED SHAPE SWEEPS

  • At both model scales and routing granularities, validation loss forms a U-shaped curve across aspect ratios, with highly unbalanced shapes worse and the best region near square.Shapes with |ℓ/m −1| ≤0.4 lie in or near the best-performing plateau; sweeps held ℓm fixed and adjusted the shared-expert dimension to keep total per-layer parameters nearly constant.Source: Figure 3 · p. 4, §3.2 · pp. 4–5

4 LAYERWISE ID-GUIDED CAPACITY ALLOCATION

LatentMoE uses layerwise intrinsic dimension (ID) as a directional signal for allocating per-expert capacity across depth. The allocation varies expert capacity while keeping per-layer parameter count, activated FLOPs, and communication volume nearly constant.

4 LAYERWISE ID-GUIDED CAPACITY ALLOCATION

  • The method uses qt = min(ℓt, mt) as a local-capacity proxy and tests whether higher-ID regions benefit from larger qt while lower-ID regions tolerate smaller qt.Because ID and qt measure different quantities, their ordinal relation serves as a directional signal.Source: §4 · p. 5

4.1 ESTIMATING LAYERWISE INTRINSIC DIMENSION

  • TwoNN profiles show ID rising through shallow layers, peaking in the middle-to-deep model, then declining near the output; this ordering appears across scales and from an early checkpoint.The estimate fits a slope through the origin after discarding the largest 10% of neighbor ratios, allowing allocation from a short proxy run.Source: §4.1 · pp. 5–6, Figure 4 · p. 5

4.2 CONSTANT-COST CAPACITY ALLOCATION

  • With routed compute, expert parameters, and dispatch volume fixed, increasing ℓt is compensated by decreasing Kt and Nt while mt stays approximately constant and Nt/Kt is preserved.The construction uses near-square experts with up to 30% aspect-ratio deviation to align per-layer compute and communication while meeting hardware constraints.Source: §4.2 · p. 6
  • For implementation, the model partitions blocks at a separation in the measured ID profile, assigning lower capacity to the shallow prefix and higher capacity to the remaining blocks.This two-regime split avoids a scale-specific numerical threshold and is an implementation choice, not a restriction of the allocation rule.Source: §4.2 · p. 6

5 EXPERIMENTS

The experiments separate complete-model downstream evaluation from controlled validation-loss analyses of expert shape, layerwise allocation, and regime boundaries. This distinction reflects that pretraining-architecture differences may be small relative to individual downstream-task resolution.

5.1 EXPERIMENTAL SETUP

  • The evaluation matches total and activated parameter budgets within each scale and uses OLMES-v0.1 5-shot mc-or-rc scoring across nine core tasks.Layerwise experiments compare ID-guided placement with uniform and exchanged controls while keeping data, optimizer, and training schedule unchanged.Source: §5.1 · pp. 6–7

5.2 MAIN DOWNSTREAM RESULTS

  • At both scales, ID-guided layerwise allocation achieves the highest mean downstream score across three runs, outperforming both uniform LatentMoE baselines and the base model.Reported standard deviations summarize run-level variation in aggregate score.Source: Table 2 · p. 7, §5.2 · p. 7

5.3 ANALYSIS

  • Across four sweeps, moving from extreme aspect ratios toward near-square experts consistently reduces validation loss at both routing granularities and model scales.The trend survives a fixed-ms control; routing granularity alone does not remove the shape-dependent per-expert constraint.Source: §5.3 · p. 7
  • ID-guided placement improves over both uniform endpoints at both scales, with near-square endpoint shapes holding expert shape fixed and isolating across-depth allocation.The uniform narrow regime can constrain higher-ID later blocks, while the uniform wide regime assigns less-useful capacity to lower-ID shallow blocks.Source: §5.3 · p. 8, Table 3 · p. 8
  • Both reversed allocations worsen validation loss at both scales and fall below both uniform endpoints, showing that shallow-to-deep ordering matters beyond mixing regimes or using more wide blocks.The controls preserve, respectively, the ID-guided boundary or the exact count of narrow and wide blocks.Source: §5.3 · p. 8, Table 3 · p. 8
  • Moving the shallow–deep boundary one block either way changes validation loss only slightly, and every tested boundary remains better than the stronger uniform endpoint.Returning the final block to the shallow configuration degrades performance, so a coarse two-part allocation captures most of the benefit without exact boundary fitting.Source: §5.3 · p. 8
  • Local Jacobians approach q = min(ℓ, m), while selected experts can collectively span latent space, making q a per-expert capacity proxy rather than a complete performance model.Fixed-q transposed shapes can realize nearly the same rank without the same loss; shape orientation, optimization, and shared-path interactions remain relevant.Source: §5.3 · pp. 8–9, Figure 5 · p. 9

6 RELATED WORK

Prior work reallocates Transformer capacity across layers, tokens, or heterogeneous experts, often changing workloads or requiring mixed-size execution. This construction instead uses a static expert shape and fixed Top-K rule per layer, guided by intrinsic dimension while keeping routed computation, parameters, and dispatch volume nearly constant.

6 RELATED WORK

  • Prior allocation methods vary layer, token, or component capacity, but generally change per-layer or per-token workloads or target components other than latent expert shape.Examples include AutoMoE’s search over expert placement, per-layer expert count, and FFN width under global FLOP and latency constraints, and layerwise sensitivity-based expert reallocation (Li et al., 2026).Source: §6 · p. 9
  • Within-layer heterogeneous-expert methods route tokens among different-width experts and require mixed-size execution or explicit load balancing.The cited approaches include Wang et al. (2025), Sun et al. (2024), and Ma et al. (2026).Source: §6 · p. 9
  • The construction uses one static expert shape and fixed Top-K per layer, covarying (ℓ, m, N, K) to keep routed computation, parameters, and dispatch volume nearly constant.Its schedule follows the ordinal intrinsic-dimension profile, using input ID to set the local-rank ceiling while retaining homogeneous grouped GEMMs and predictable per-layer workloads.Source: §6 · p. 9

7 CONCLUSION … B.1 PRETRAINING DETAILS

The paper treats expert shape and layerwise capacity allocation as practical MoE design choices: near-square experts maximize the local-rank ceiling, and ID-guided allocation improves matched uniform controls. These signals have limits: rank is a single-expert proxy, while intrinsic dimension gives only an ordinal allocation signal.

7 CONCLUSION

  • ID-guided allocation improves matched uniform controls, but reversing capacity placement removes the gain; controlled sweeps favor near-square factorization.The local-rank ceiling favors a near-square factorization, while the rank ceiling remains a single-expert proxy and intrinsic dimension an ordinal allocation signal.Source: §7 · p. 9
  • A fine-grained mapping from representation geometry to per-layer capacity remains an extension beyond the paper’s coarse allocation guidance.The work identifies expert shape as a design variable and varies per-expert capacity across depth while preserving aggregate compute and communication.Source: §7 · p. 9
  • Generative AI tools assisted with language polishing, LaTeX formatting, and code drafting and refinement; authors reviewed the text, inspected and tested code, and checked figures.The authors state that AI tools were not used for research design, method implementation, hypothesis refinement, or interpretation of results.Source: AI USE STATEMENT · p. 10
  • The work uses existing language-model corpora and benchmarks without human-subject experiments or new personal-data collection, while inheriting data and model risks and incurring training costs.The stated risks include bias, privacy, misuse, and computational and energy costs.Source: ETHICS STATEMENT · p. 10

A ADDITIONAL RELATED WORK

  • The work is orthogonal to routing, load-balancing, compression, and aggregate-budget approaches, isolating ℓ× m factorization, local-rank ceilings, and capacity allocation across depth.Rank matching can be combined with existing routing and load-balancing methods.Source: A ADDITIONAL RELATED WORK · p. 13

B.1 PRETRAINING DETAILS

  • Shape sweeps used 3.4B-A250M and 8B-A560M models trained for 20B and 40B tokens, respectively, while holding key training and routing factors fixed within each sweep.Layerwise controls used the same scale-specific setup and changed only the placement of MoE configurations.Source: §B.1 · p. 13
  • The largest within-sweep validation-loss gaps, 1.04 × 10−2–3.99 × 10−2, were roughly 7–27 times the estimated run-to-run noise scale.Noise scales from five base-configuration runs were 1.2 × 10−3 at 3.4B-A250M and 1.5 × 10−3 at 8B-A560M; each shape was trained once.Source: §B.1 · p. 13
  • Downstream evaluation lists four model settings at each scale, with total and activated parameter budgets matched within each scale.The table caption states that uniform models use α = d/ℓ and layerwise entries use the two configurations in Table 1.Source: §B.1 · p. 13, Table 4 · p. 14

B.2 PARAMETER, COMPUTE, AND COMMUNICATION ACCOUNTING · B.3 SYSTEM THROUGHPUT · C ADDITIONAL EXPERT-SHAPING RESULTS

LatentMoE separates routed-expert parameters, activated computation, and communication through explicit expert dimensions, while heterogeneous layer allocation maintains throughput close to uniform LatentMoE. Its current implementation nevertheless has a modest throughput gap relative to the base model.

B.2 PARAMETER, COMPUTE, AND COMMUNICATION ACCOUNTING

  • Each layer’s routed experts contain 3N_tℓ_tm_t parameters; P_t controls the routed-expert parameter pool, while Q_t controls the remaining dominant parameters.The layer also has 2dℓ_t latent-projection parameters, dN_t router parameters, and 3dm_s,t full-width shared-expert parameters.Source: §B.2 · p. 14
  • Activated MACs per token are approximately 3C_t + dQ_t, so matching C_t and Q_t matches dominant activated computation.The FLOP count is twice the MAC count; the approximation excludes lower-order elementwise operations such as activation functions and routing softmax.Source: §B.2 · p. 15
  • Standard MoE dispatches and returns K copies of a d-dimensional representation, totaling 2Kd scalars per token across both directions.The dispatch and return each carry Kd scalars per token.Source: Figure 6 · p. 14, §B.2 · p. 15
  • LatentMoE sends and returns latent vectors, reducing the two-way payload to 2K_tℓ_t scalars per token by placing shared projections outside the all-to-all.For T tokens and b bytes per scalar, dispatch plus return uses 2bTK_tℓ_t bytes; transmitted routing metadata adds an O(TK_t) term.Source: §B.2 · p. 15

B.3 SYSTEM THROUGHPUT

  • Heterogeneous allocation achieves throughput within measurement variation of uniform LatentMoE, while LatentMoE variants retain a modest throughput gap relative to the base model.The reported gap is consistent with implementation-level overhead from finer-grained routing, smaller grouped GEMMs, and latent projections.Source: Table 5 · p. 15, §B.3 · p. 15

C.1 FULL SHAPE-SWEEP RESULTS

This section presents the complete shape-sweep configurations and validation losses for Figure 3; K/N denotes activated experts over total experts.

C.1 FULL SHAPE-SWEEP RESULTS

  • Table 6 lists the complete configurations and validation losses for Figure 3, with K/N denoting activated experts over total experts.Source: Table 6 · p. 16

C.2 FIXED SHARED-EXPERT-WIDTH CONTROL

With shared-path reallocation controlled, the broad near-square advantage remains, despite the fixed-ms design increasingly favoring larger latent widths.

C.2 FIXED SHARED-EXPERT-WIDTH CONTROL

  • The control repeats the 3.4B-A250M (N, K) = (256, 8) sweep with ms = 768 and ℓm ≈ 393,216 fixed; the largest parameter advantage is only 0.47%.Larger latent widths receive a modest parameter advantage because down- and up-projection parameter count grows linearly with ℓ.Source: §C.2 · p. 16
  • 0.0151 validation-loss span in the fixed-ms control, versus 0.0166 in the parameter-matched sweep, with the broad near-square advantage preserved.The fixed-ms design increasingly favors settings with larger latent widths.Source: Table 7 · p. 16, §C.2 · p. 16

D ADDITIONAL LAYERWISE-ALLOCATION ABLATIONS · E CAPACITY DIAGNOSTICS

At 3.4B, layerwise-allocation ablations show that small boundary shifts have limited impact, while restoring a narrow final block or using a strongly unbalanced shape is more costly.

D ADDITIONAL LAYERWISE-ALLOCATION ABLATIONS

E.1 REALIZED RANK · E.2 INTRINSIC-DIMENSION ESTIMATION

Realized local-Jacobian rank generally tracks the per-expert ceiling q, while mixture-level span does not establish that per-expert capacity is dispensable. Across model sizes and training checkpoints, intrinsic dimension rises early, peaks in the middle-to-deep layers, and declines toward the output.

E.1 REALIZED RANK

  • Across 27 sweep checkpoints, 1001 of 1059 block–matrix means equaled q and the minimum was 0.992q; local Jacobians reached 0.909–1.000q.Expert output clouds reached 0.988–1.000q and latent inputs spanned 0.995–1.000 of their ambient width, making q = min(ℓ, m) an empirically relevant ceiling.Source: §E.1 · p. 18, Table 9 · p. 18, Table 10 · p. 18
  • Across all 72 within-sweep pairs with distinct ceilings, the configuration with larger q realized larger absolute local-Jacobian rank J.Across nine transposed (ℓ, m) and (m, ℓ) pairs sharing q, mean J differed by at most 3.4% of q even though validation losses differed.Source: §E.1 · p. 19, Table 11 · p. 19
  • Combined expert-map Jacobians spanned 0.931–0.994 of ℓ in every configuration, ruling out a mixture-level span bottleneck but not per-expert capacity limits.Every configuration had Kq ≥ ℓ, so aggregate local span does not show that reducing individual experts’ ceilings is lossless.Source: §E.1 · p. 19
  • With uniform q across depth, block-averaged local-Jacobian rank correlated positively with the base-model TwoNN profile in all 27 checkpoints, with median ρ = +0.76.Excluding the frequently divergent final MoE block, the correlation remained positive in 27/27 checkpoints with median ρ = +0.85; the allocation hypothesis concerns greater tolerance to compression in lower-ID regions, not ID equaling qt.Source: §E.1 · p. 19
  • In two-regime models, the deep regime realized a larger local-Jacobian rank than the shallow regime at both scales.At 3.4B, the ID-guided checkpoint closely followed its layer-matched uniform endpoint reference: mean signed residual −0.10% of q, mean absolute residual 0.38% of q, and maximum absolute residual 1.54% of q.Source: §E.1 · pp. 19–20, Table 12 · p. 20, Figure 7 · p. 20

E.2 INTRINSIC-DIMENSION ESTIMATION

  • Across both model sizes and all five 8B checkpoints, ID rose through early blocks, peaked in the middle-to-deep region, and declined toward the output.The cross-scale estimates used 20,000 non-padding token representations per block boundary, randomly retaining 5,000 for TwoNN estimation.Source: §E.2 · pp. 20–21

Related papers

Loading 2610.14644v1…