Expert shape breaks free from routing granularity: near-square experts plus layerwise capacity allocation top nine downstream tasks at 8B

Latent mixture-of-experts moves expert computation into a lower-dimensional latent space, so the expert input dimension no longer equals the residual stream width and expert shape becomes a design choice that can be set on its own. Earlier fine-grained mixture-of-experts increased expert count under a fixed feed-forward budget, which could only divide that budget into smaller per-expert intermediate dimensions, leaving the aspect ratio a by-product of routing granularity rather than an independent choice. Controlled shape sweeps at 3.4B and 8B scales under two routing granularities show a U-shaped validation-loss curve over aspect ratio, with both highly unbalanced directions degrading and the best region near square.

A single expert's local Jacobian rank is bounded jointly by its input dimension and its intermediate width, and at fixed per-expert parameters balancing the two maximizes that rank ceiling. The original latent mixture-of-experts formulation fixed a four-fold compression empirically and gave no general principle for jointly choosing the latent dimension and the expert intermediate dimension. Across 27 sweep checkpoints, expert weights and routed-token outputs are essentially full rank, and trained local Jacobians attain 0.909 to 1.000, showing the ceiling is approached in trained models rather than vacuous.

Allocating per-expert capacity by the layerwise intrinsic-dimension ordering keeps per-layer parameter count, activated floating-point operations, and communication volume nearly constant while giving higher-intrinsic-dimension layers more capacity. Prior heterogeneous designs changed expert count or feed-forward width alone, creating per-layer compute and communication imbalance that complicates pipeline partitioning and efficient expert-parallel kernels. At 3.4B and 8B, the intrinsic-dimension-guided allocation achieves the highest mean score across three runs, beating both uniform latent mixture-of-experts baselines and the uncompressed base model, while exchanged-regime controls worsen validation loss at both scales.

The report covers pretraining at 3.4B-A250M and 8B-A560M scales for 20B and 40B tokens, plus the nine-task OLMES core_9mcqa downstream suite. The layerwise allocation uses a two-segment implementation that splits shallow and deep blocks by the ordinal intrinsic-dimension profile rather than calibrating a configuration per layer. The near-square rule allows up to about 30% aspect-ratio deviation to align per-layer compute and communication and to satisfy grouped-GEMM and expert-parallel divisibility requirements. Readers interested in larger scales or communication-bound settings can watch the extension the source proposes: a finer-grained mapping from representation geometry to per-layer capacity.

The rank ceiling is only a single-expert capacity proxy, and intrinsic dimension is only an ordinal allocation signal; neither is a complete performance model. Transposed shapes that share the same rank ceiling are not always equivalent, so shape orientation, optimization, and interactions with the shared path still matter. Uniform-shape models with a fixed ceiling still show depth-wise variation in realized rank aligned with the intrinsic-dimension profile, showing that the available ceiling and its realized use are distinct. The two-segment allocation is an implementation choice for kernel and configuration simplicity, not a restriction of the allocation rule.