Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization

The gist The proposed framework establishes a reusable, cost-efficient experimental paradigm for systematically studying MoE architectural properties at scale by employing a mathematical decoupling strategy to reduce the search space from O(n 16) to two sequential phases of O(n 3) and O(n 2)

How it works

The paper introduces a general-purpose framework for holistic MoE architectural optimization that is itself a central contribution of this work by being a reusable, cost-efficient experimental paradigm for systematically studying MoE architectural properties at scale The analysis first reveals that exclusive reliance on FLOPs per token (M) is an inadequate scaling metric for MoE models because heterogeneous computational densities between Attention and FFN layers can yield misleading performance gains driven by parameter inflation rather than effective compute utilization To tame the combinatorial explosion of the resulting high-dimensional search space, a mathematical decoupling strategy is employed that leverages inherent structural constraints and a rank-preserving property of the hidden dimension, factorizing the global optimization problem into a highly efficient, proxy-based search paradigm

Constraint Triad and Dimension Reduction

To ensure fair and robust structural optimization within the MoE design space, the architectural search space must be rigorously constrained by a joint triad: FLOPs per token (M), active parameters (Na), and total parameters (N) This constraint is necessary because anchoring solely on M introduces an inherent bias because MoE models can inflate their total parameter count (N) and active parameters (Na) while keeping M constant The framework first adopts the established methodology of Hoffmann et al. to characterize the classical (C, M, D) scaling laws by identifying six compute-optimal (C, M, D) combinations This approach effectively eliminates two dimensions from the overarching optimization problem

Two-Phase Search Strategy

The dimension decomposition methodology transforms an intractable O(n 16) search space into a manageable O(n 3) + O(n 2) two-phase search through mathematical constraints, rank-preserving properties, and coordinate transformation The framework first locks the macroscopic parameters (M, Na, N), reducing the complexity to four fundamental degrees of freedom: FLOPs per token (M), active parameters (Na), total parameters (N), and the hidden dimension d Phase 1 involves evaluating optimal scaling behaviors of Nopt a(M) and Nopt(M) using a d-proxy, incurring a manageable cost of O(n 3) Phase 2 focuses on determining the optimal scaling law for the hidden dimension d given the fixed (M, Na, N), requiring only O(n 2) additional runs

Key Empirical Findings and Engineering Insights

A key empirical finding is that the near-optimal configuration band widens as compute scale increases, providing a quantitative basis for practical trade-offs between scaling law recommendations and infrastructure engineering constraints This widening provides a quantitative basis for practitioners to navigate trade-offs between scaling law recommendations and real-world infrastructure engineering constraints Furthermore, the near-optimal band widens with increasing compute scale, indicating that large-scale models are increasingly tolerant of architectural deviations from the strict optimum The near-optimal band provides a principled, quantitative basis for evaluating such deviations by quantifying the performance cost of such deviations

Generality and Future Directions

The dimension decomposition methodology is designed to be reusable, as after locking the macroscopic parameters (M, Na, N), the remaining architectural degrees of freedom collapse to exactly one This single degree of freedom can be filled by any structural ratio of interest—such as the MoE FFN intermediate width relative to the hidden size—enabling controlled, fair comparisons across structural variants The framework is designed to be extended, and the dimension reduction and two-phase search strategy remain valid regardless of which structural property fills the final degree of freedom This research provides a deterministic pipeline from compute budget to complete MoE architecture by deriving the complete optimal allocation (M, D, Na, N, d) from a given compute budget C

Conclusion

The derived scaling laws provide a deterministic pipeline from compute budget to complete MoE architecture by determining optimal Mopt and Dopt via Eq.

1. Holistic architectural optimization framework: The proposed framework allows for determining strictly optimal MoE architectures under arbitrary compute budgets by deriving a complete optimal allocation (M, D, Na, N, d) from a given compute budget C, providing actionable blueprints for designing MoE models that minimize training loss.

2. Dimension decomposition methodology: The strategy reduces the search complexity by transforming an intractable O(n 16) search space into a manageable optimization via a two-phase search of O(n 3) and O(n 2), enabling systematic exploration of MoE architectural properties efficiently.

3. Robust performance trade-off assessment: The empirical finding that the near-optimal configuration band widens as compute scale increases allows practitioners to make informed decisions, providing a quantitative basis for practical trade-offs between scaling law recommendations and infrastructure engineering constraints.