Mixture-of-Experts Models in Vision: Routing, Optimization, and Generalization

TL;DR

This study investigates Mixture-of-Experts (MoE) models in image classification by comparing Dense, SoftMoE, and SparseMoE heads on CIFAR-10 with a shared ResNet-18 backbone under matched capacity. It combines predictive evaluations with Hessian-based loss-landscape analyses and loss-surface perturbations to understand generalization beyond validation accuracy, finding that MoEs generalize as well as dense baselines, with SoftMoE exhibiting higher curvature. While MoE variants can yield modest validation gains, the theoretical inference savings from conditional routing do not translate into practical speedups at moderate scales due to routing overhead. The work highlights the gap between theoretical efficiency and practical performance for sparse MoEs and suggests larger benefits in heterogeneous data settings where specialization reduces interference more substantially.

Abstract

Mixture-of-Experts (MoE) architectures enable conditional computation by routing inputs to multiple expert subnetworks and are often motivated as a mechanism for scaling large language models. In this project, we instead study MoE behavior in an image classification setting, focusing on predictive performance, expert utilization, and generalization. We compare dense, SoftMoE, and SparseMoE classifier heads on the CIFAR10 dataset under comparable model capacity. Both MoE variants achieve slightly higher validation accuracy than the dense baseline while maintaining balanced expert utilization through regularization, avoiding expert collapse. To analyze generalization, we compute Hessian-based sharpness metrics at convergence, including the largest eigenvalue and trace of the loss Hessian, evaluated on both training and test data. We find that SoftMoE exhibits higher sharpness by these metrics, while Dense and SparseMoE lie in a similar curvature regime, despite all models achieving comparable generalization performance. Complementary loss surface perturbation analyses reveal qualitative differences in non-local behavior under finite parameter perturbations between dense and MoE models, which help contextualize curvature-based measurements without directly explaining validation accuracy. We further evaluate empirical inference efficiency and show that naively implemented conditional routing does not yield inference speedups on modern hardware at this scale, highlighting the gap between theoretical and realized efficiency in sparse MoE models.

Figures (7)

• Figure 1: Conceptual illustration of a Mixture-of-Experts architecture, where a gating network routes inputs to multiple expert subnetworks shazeer2017outrageously.
• Figure 2: Architectures used in this study. A shared ResNet-18 backbone produces a 512-dimensional feature vector, which is processed by either a dense classifier head, a SoftMoE head with all experts active, or a SparseMoE head using Noisy Top-$2$ routing.
• Figure 3: Illustration of flat versus sharp minima in the loss landscape, commonly used to motivate curvature-based generalization diagnostics onlargebatchtraining.
• Figure 4: Expert utilization over training.
• Figure 5: Class-conditioned expert specialization at convergence. Both variants exhibit structured specialization, with several experts strongly associated with specific classes and some acting more broadly.