Geometric Metrics for MoE Specialization: From Fisher Information to Early Failure Detection

TL;DR

This paper introduces an information-geometric framework for analyzing expert specialization in MoE models, providing theoretically grounded metrics that outperform existing heuristics in predicting training failure and correlating with performance.

Contribution

It offers the first rigorous geometric analysis of MoE specialization dynamics, proposing new metrics with strong theoretical justification and practical effectiveness.

Findings

Fisher Specialization Index correlates with downstream performance (r=0.91).

Fisher Heterogeneity Score predicts training failure with AUC=0.89.

Proposed metrics outperform validation-loss-based early stopping by 23%."

Abstract

Expert specialization is fundamental to Mixture-of-Experts (MoE) model success, yet existing metrics (cosine similarity, routing entropy) lack theoretical grounding and yield inconsistent conclusions under reparameterization. We present an information-geometric framework providing the first rigorous characterization of MoE specialization dynamics. Our key insight is that expert routing distributions evolve on the probability simplex equipped with the Fisher information metric, enabling formal analysis via Riemannian geometry. We prove that standard heuristic metrics violate parameterization invariance (Theorem 1), establish that specialization corresponds to geodesic flow with quantified approximation bounds (Theorem 2), and derive a failure predictor with theoretical threshold justification (Theorem 3). The framework introduces two principled metrics: Fisher Specialization Index (FSI) achieving r=0.91+/-0.02 correlation with downstream performance, and Fisher Heterogeneity Score (FHS) predicting training failure at 10% completion with AUC=0.89+/-0.03 -- outperforming validation-loss-based early stopping by 23% while requiring 40x fewer compute cycles. We validate intervention protocols achieving 87% recovery rate when FHS>1 is detected. Comprehensive experiments across language modeling (WikiText-103, C4), vision MoE (ImageNet), and scaling studies (8-64 experts, 125M-2.7B parameters) validate our theoretical predictions.