Structure Aware Neural Architecture Search for Mixture of Experts

Petr Babkin; Oleg Bakhteev

Source-linked AI summary

Structure Aware Neural Architecture Search for Mixture of Experts

Petr Babkin, Oleg Bakhteev

arXiv:2608.29817v1Aug 30, 2026cs.LG

TL;DR

NAS for MoE has not explicitly controlled alignment between experts and data structure. This paper jointly searches cluster assignments and expert architectures through a surrogate-assisted generalized EM procedure, recovering 95% of image-mixture clusters and outperforming label-free baselines across image and time-series tasks.

Problem
NAS remains relatively underexplored for MoE, whose expert alignment with distinct data modalities is usually left to emerge rather than explicitly controlled.
Method
The method formulates joint routing and architecture selection as cluster-aware likelihood maximisation and solves it with generalized EM using an adaptively refined surrogate.
Results
95% of clusters recover the ground-truth source partition on the heterogeneous image mixture, while SGEM is the best label-free method under the cluster gate on both image classification and four-domain forecasting.
Takeaways & Limitations
Explicitly searching cluster-to-expert assignments alongside expert architectures improves label-free MoE performance across the evaluated image and time-series mixtures.
Takeaways & Limitations
The framework is currently limited to search spaces with K ≤4, and matching the validation objective does not guarantee matching test-time generalisation on the low-contrast time-series mixture.

Abstract

from arXiv · show

Neural Architecture Search (NAS) has so far rarely been applied to Mixture-of-Experts (MoE) models, and existing MoE designs leave the alignment between experts and the structure of the data to emerge on its own. We propose an architecture search framework that makes this alignment an explicit search variable: the assignment of data clusters to experts is optimised jointly with the per-expert architectures. We cast the joint problem as a cluster-aware likelihood maximisation, show that it coincides with the incomplete-data maximum likelihood of a latent-variable mixture, and solve it by a generalised Expectation-Maximisation procedure whose otherwise intractable expert-quality term is supplied by an adaptively refined surrogate. We prove that the iterates converge whenever the surrogate errors are summable, and that at every limit point no candidate the search produces improves the true objective. On a heterogeneous image-classification mixture the method recovers the underlying domain partition on 95% of clusters without ever observing domain labels, and on that benchmark and a four-domain time-series forecasting one alike it outperforms the MoE and NAS baselines that likewise use no label information.

1. Introduction

NAS for MoE remains underexplored, while existing methods leave expert–data alignment to emerge rather than controlling it explicitly. The paper formulates this alignment as a joint architecture-and-routing search problem and evaluates it against label-free baselines.

NAS has matured separately from MoE, but applying architecture search specifically to MoE remains relatively underexplored.
MoE architectures suit strongly clustered data, yet existing methods do not explicitly control how experts align with distinct modalities.
The paper asks how expert architectures can adapt to specific data clusters and modalities.
The proposed formulation jointly searches MoE architectures through cluster-aware likelihood maximisation, equivalent to incomplete-data maximum likelihood for a latent-variable mixture.
The algorithm uses Expectation–Maximisation with an adaptively refined surrogate, accompanied by a convergence theorem.
The experiments compare mixture-of-experts approaches against several baselines on datasets designed to expose their advantages.

2. Problem Statement

The search space combines discrete per-expert architectures with stochastic cluster-level routing. The resulting cluster-aware likelihood objective jointly selects architectures and assignments, while surrogate evaluation replaces repeated candidate training.

2.1. Search Space of MoE: The search space contains per-expert architectures and cluster-level routing, which jointly specify an MoE model.
2.1. Search Space of MoE: Each expert uses a cell-based architecture, where a finite architecture vector selects operations and predecessor nodes within a directed acyclic graph.
2.1. Search Space of MoE: Routing is stochastic at the cluster level: a routing matrix assigns each cluster a distribution over experts, with hard assignments read out only after optimisation.
2.1. Search Space of MoE: Unlike standard MoE training, the search jointly determines routing and per-expert architectures, coupling architecture quality to each expert’s assigned cluster subset.
2.2. MoE Objective: The objective maximises mixture likelihood by marginalising latent cluster-to-expert assignments and factoring each term into routing probability and conditional likelihood.
2.2. MoE Objective: Exact conditional likelihood evaluation would require training every candidate on its routed subset, so the method replaces it with a vector-valued learned surrogate.
2.2. MoE Objective: The cluster-aware objective weights routing and expert likelihood by cluster size and solves the result with generalised EM and an adaptively refined surrogate.

3. Solving the Cluster-aware Objective by Surrogate-assisted EM

The method solves the cluster-aware MoE architecture-search problem with generalized EM, using an adaptively refined surrogate for otherwise expensive expert-quality evaluations. Its analysis establishes convergence despite surrogate refinement and limits improvement claims to candidates the stochastic search can produce.

3.1. Learning the Surrogate: The surrogate is learned from architecture/cluster-subset pairs collected by genuine training, then used to score candidate architectures before selectively training them.A single run supplies per-cluster loss targets, while the surrogate is refined through active learning.
3.1. Learning the Surrogate: Monte-Carlo dropout produces predictive means and uncertainties, and the Lower Confidence Bound selects candidates with low predicted loss and high epistemic uncertainty for genuine training.The resulting observations are appended to the surrogate dataset for subsequent refinement.
3.2. Surrogate-assisted Generalized EM (SGEM): Routing is treated as a latent variable, so the cluster-aware objective becomes an incomplete-data likelihood optimized by alternating E- and M-steps.Responsibilities are proportional to routing weights multiplied by the expert’s per-cluster likelihood, and all samples in a cluster share one responsibility.
3.2. Surrogate-assisted Generalized EM (SGEM): Each expert’s architecture is selected by minimizing its responsibility-weighted surrogate loss over a freshly sampled candidate pool.The stochastic pool search gives an improvement M-step rather than an exact maximization, while selected candidates can also refine the surrogate.
3.3. Convergence under a Drifting Surrogate: The procedure converges when surrogate errors are summable, and every limit point admits no improvement from candidates the algorithm keeps offering under the true objective.The guarantee applies despite a moving surrogate objective and is relative to the reachable M-step candidate set.
3.3. Convergence under a Drifting Surrogate: The analysis requires bounded expert-loss outputs across architectures, routing masks, clusters, and iterations.The paper notes that upper-quantile clipping of surrogate targets enforces this boundedness in practice.

4. Experiments

Experiments test SGEM on toy, heterogeneous image, and multi-domain forecasting mixtures, showing strong source recovery and accuracy gains in sharp-contrast settings but weaker domain recovery in low-contrast forecasting.

Toy experiment: SGEM recovers the ideal 7-ring/13-linear cluster split in 20 EM iterations on the toy mixture.The recovered assignment matches the hand-defined oracle partition.
CIFAR-100 and SVHN mixture: 0.6402 ± 0.002 test accuracy lets SGEM outperform DARTS×K, random-MoE, and the DARTS oracle MoE under cluster gating.The gains over DARTS×K and random-MoE are +4.9 and +8.5 percentage points, respectively.
CIFAR-100 and SVHN mixture: 95% ± 2% of clusters match the oracle source split across five seeds, versus approximately 54% for balanced random routing.A typical run produces a pure CIFAR expert while the other expert absorbs all SVHN clusters and two CIFAR clusters.
CIFAR-100 and SVHN mixture: Source recovery rises from 53–67% initially to approximately 97% within two EM iterations, then remains on a 0.93–0.97 plateau.This empirical behavior is consistent with the stated convergence guarantee.
Multi-domain time-series forecasting: On the four-domain forecasting mixture, SGEM improves cluster-gated NLL over BestArch×K by 0.042, NAS-MoE by 0.025, and geometry routing by 0.175.NLL is treated as the primary criterion because it matches the likelihood-based objective; CRPS can rank methods differently.
Multi-domain time-series forecasting: 45% mean Hungarian agreement is only marginally above the 40% random-split base, while SGEM remains competitive in end metrics in the low-contrast forecasting mixture.The likelihood landscape contains near-optimal partitions that need not coincide with domain boundaries, unlike the sharper CIFAR/SVHN setting.

5. Conclusion

The framework jointly searches expert architectures and data-cluster assignments, with convergence guarantees despite surrogate refinement. It recovers heterogeneous image domains and outperforms label-free baselines, while low-contrast time-series domains expose limits of domain recovery and test generalisation.

5. Conclusion: SGEM jointly searches data-cluster assignments to experts and per-expert architectures through a latent-variable likelihood and surrogate-refined EM procedure.The true likelihood is never evaluated directly by the algorithm, yet Theorem 1 establishes convergence.
5. Conclusion: 95% of clusters recover the ground-truth source partition on a heterogeneous CIFAR-100 / SVHN mixture without source labels, with 0.640 cluster-gate test accuracy.This exceeds the oracle MoE baseline’s 0.632 accuracy and all label-free baselines.
5. Conclusion: SGEM is the best label-free method on the four-domain time-series mixture under the cluster gate, achieving test NLL 0.883.The method also recovers the ideal split on a controlled toy problem.
5. Conclusion: On the low-contrast time-series mixture, SGEM matches the oracle partition in real objective value but agrees with it barely above chance.The remaining gap to the oracle ceiling of 0.789 NLL leaves test-time generalisation as an open question.

CRediT authorship contribution statement

The authors divide contributions across conceptualisation, methodology, software, investigation, supervision, and writing.

CRediT authorship contribution statement: Petr Babkin contributed conceptualisation, methodology, software, investigation, and the original draft.
CRediT authorship contribution statement: Oleg Bakhteev contributed conceptualisation, methodology, supervision, and review and editing.

A. Proof of Theorem 1

The proof represents the likelihood through a variational free energy and bounds the surrogate’s deviation from the true objective and responsibilities. These bounds support convergence analysis of the surrogate-based procedure.

A. Proof of Theorem 1: The proof defines Lu(θ) as the maximised variational free energy, representing the size-weighted log-likelihood induced by loss u.The maximum is taken over the variational distribution q.
A. Proof of Theorem 1: The ELBO formulation attains its maximum at the responsibility distribution induced by the candidate architecture and loss.The KL-divergence argument shows equality precisely at the corresponding optimal responsibilities.
A. Proof of Theorem 1: Uniform ELBO deviation is bounded by the per-iteration surrogate error εt, yielding |Ltrue(θ) − L(t)(θ)| ≤ Nεt.The bound follows because shared likelihood terms cancel between the true and surrogate ELBOs.
A. Proof of Theorem 1: Softmax stability transfers surrogate loss-error bounds to responsibility-error bounds when εt ≤ 1.The normalised-vector deviation is bounded by (e^2ε − 1)pk and by 3εpk for ε ≤ 1.
A. Proof of Theorem 1: The surrogate Q-function differs from the true Q-function by at most CQεt under Assumption 1.The constant depends on N, K, and Umax.

A.2. Main Proof

The main proof combines approximate EM ascent with summable surrogate errors to establish convergence of the true likelihood and characterize limit points. At each limit point, no candidate in the attainable M-step set improves the true objective.

A.2. Main Proof: Approximate EM ascent follows by chaining variational maximisation, surrogate deviation bounds, the GEM step, and the E-step identity.The resulting inequality controls true-likelihood changes using two Nεt error terms.
A.2. Main Proof: Because the parameter space is compact and the surrogate errors are summable, the true likelihood sequence converges.Robbins–Siegmund is applied to the nonnegative gap from the finite likelihood supremum.
A.2. Main Proof: Continuity of the true Q-function transfers any hypothetical improvement at a limit point to sufficiently late iterates.Continuity follows from continuous losses, softmax, and the extension of r log r at zero.
A.2. Main Proof: The M-step maximisation and Q-function deviation bound imply a fixed-point contradiction for any candidate that would improve the true objective.The proof uses the surrogate M-step choice and εtj → 0 to preserve a positive improvement gap.
A.2. Main Proof: No candidate in the attainable limit M-step set improves the true objective at any limit point.Otherwise, the accumulated positive increments contradict summability of the error-controlled increments.

A.3. Convergence Rate

The appendix describes exploration mechanisms that refine the surrogate and expand routing search, then refits each expert architecture after the EM loop. The convergence analysis gives a rate when surrogate errors decay sufficiently quickly.

A.3. Convergence Rate: If surrogate errors satisfy ε_t = O(t^-p) with p > 1, the true objective converges at rate O(T^(1−p)).This rate matches that observed in stochastic-EM analyses.
A.3. Convergence Rate: The search uses cluster affinities, exploration masks, and partition jumps to improve surrogate coverage and escape routing-local optima.Affinity groups generate near-pure training masks and candidate partitions, while partition replacements require real training outcomes.
A.3. Convergence Rate: After EM terminates, each expert architecture is re-searched from scratch within its fixed assigned region.Once routing is fixed, the remaining problem is ordinary per-region NAS.

C. Additional Experimental Results

The appendix reports ablations and a convergence trace on the CIFAR-100/SVHN mixture under the same evaluation protocol as the headline table.

C. Additional Experimental Results: The ablations use a scalar surrogate predicting one region-average value per candidate instead of the method’s per-cluster vector surrogate.The scalar reduction changes the surrogate used for the reported ablation analyses.

C.1. Decomposition of the cluster-gate accuracy

The decomposition identifies cluster-to-expert routing as the main contributor to accuracy on the heterogeneous image mixture, while architectural heterogeneity adds little once routing is correct.

C.1. Decomposition of the cluster-gate accuracy: 3.0 percentage points: oracle routing raises accuracy over a random cluster split when the architecture is fixed.This isolates the routing contribution under a single DARTS architecture searched on the full mixture.
C.1. Decomposition of the cluster-gate accuracy: Architectural variants remain within 1 p.p. of one another when oracle routing is fixed.The compared variants are a full-mixture architecture and architectures searched separately for CIFAR-100 and SVHN.
C.1. Decomposition of the cluster-gate accuracy: 92% source recovery and 0.620 accuracy: full SGEM learns routing without source labels, versus 59% and 0.598 when routing is frozen randomly.The learned result is essentially the oracle-routing accuracy reported in the ablation.
C.1. Decomposition of the cluster-gate accuracy: “src. match” measures agreement between hard routing and the ground-truth source partition, reaching 100% by construction for oracle routing.The metric is reported in the ablation table’s caption.

C.2. The scalar surrogate and load balancing

The scalar surrogate can collapse routing because it removes cluster-specific performance signals, while load balancing counteracts collapse but can conflict with an imbalanced domain partition. The full per-cluster surrogate avoids this conflict and performs better in the reported ablation.

C.2. The scalar surrogate and load balancing: A scalar surrogate makes the routing objective linear in expert assignment counts, favoring one expert for every cluster when predicted losses differ.When predicted expert losses are equal, the objective is flat over the routing simplex instead.
C.2. The scalar surrogate and load balancing: The per-cluster surrogate preserves cluster-dependent signals, so its routing maximizer can vary across clusters without requiring a balancing penalty.The method therefore uses the vector-valued surrogate throughout.
C.2. The scalar surrogate and load balancing: A load-balancing penalty is minimal for uniform routing and saturates at K under full collapse.The penalty is added to counteract collapse in the scalar case.
C.2. The scalar surrogate and load balancing: The oracle split wins the penalised objective only for λLB < 3.77 because balancing penalises its imbalanced cluster allocation.The unpenalised advantage is ΔL = 7874, while the balancing penalty difference is ΔLBN = 2091 per unit of λLB.
C.2. The scalar surrogate and load balancing: The full method uses no balancing term and shows no collapse, whereas the scalar-surrogate variant reaches 0.6199 ± 0.009 accuracy and 92% ± 2% source recovery.The full method reaches 0.6402 ± 0.002 accuracy and 95% ± 2% source recovery, a 2.0 p.p. accuracy gap.
C.2. The scalar surrogate and load balancing: The balancing conflict reflects cluster imbalance rather than sample imbalance: the oracle split is 19/11 clusters but nearly sample-balanced.Using sample fractions instead of cluster fractions would leave the oracle split essentially unpenalised.

C.3. Convergence of a single run

During a scalar-surrogate run, the objective increased up to optimisation noise while source recovery improved from 67% to 100%. Their peaks differed by one iteration, so selecting the best objective produced a near-ideal but not always perfect split.

The source-recovery rate rose from 67% to 100% as the size-weighted objective increased up to optimisation noise.This matches the behaviour predicted by Theorem 1.
The objective peaked at t = 8, while source recovery peaked at t = 9.The one-iteration mismatch reflects the objective/quality mismatch near the optimum discussed in Section 4.2.
Selecting the iterate with the best objective returned a near-ideal but not always perfect split.