Expert Routing for Communication-Efficient MoE via Finite Expert Banks

Technical Summary: Expert Routing for Communication-Efficient MoE via Finite Expert Banks

Problem Statement
Sparse Mixture-of-Experts (MoE) architectures are increasingly used for resource-efficient machine learning, where a gating mechanism selects a subset of experts for each input. A critical theoretical challenge is interpreting the gate under communication, privacy, or compression constraints. While recent work frames the gate as a stochastic channel $X \to T$ with an operational information rate $R_g = I(X; T)$, applying this framework to standard deep learning benchmarks is difficult. In high-dimensional continuous parameter spaces (e.g., CNNs with $\sim 10^5$ parameters), the mutual information between the training sample and the model parameters, $I(S; \Theta)$, and the associated generalization bounds are typically intractable and lack closed-form estimators. Consequently, there is a need for a benchmark-oriented extension that makes these information-theoretic quantities measurable to analyze resource-aware MoE inference.

Methodology
The authors propose a finite-bank construction on the MNIST dataset to render the information-theoretic analysis tractable. The methodology involves three core components:

1. Finite Expert Bank: Instead of training a continuous parameter space, the authors utilize a bank of $R=25$ pre-trained, independent CNN classifiers. Each candidate is trained on a distinct subset of MNIST data. The "learned parameter" $\Theta$ is restricted to the discrete index $W$ of the selected expert from this finite set.
2. Data-Dependent Selection Rule: A randomized learning rule is defined via an $\alpha$-mixture posterior. Given a training sample $S$, the empirical risk minimizer $r^*(S)$ is identified. The selection probability for expert $r$ is given by $q_\alpha(r|S) = \frac{1-\alpha}{R} + \alpha \cdot \mathbb{1}\{r = r^*(S)\}$. The parameter $\alpha \in [0, 1]$ controls the dependence of the selected expert on the training data: $\alpha=0$ yields a uniform, data-independent selection, while $\alpha=1$ corresponds to deterministic Empirical Risk Minimization (ERM).
3. Information Estimation:
   • Algorithmic Mutual Information ($I(S; W)$): Because $W$ is discrete, $I(S; W) = H(W) - H(W|S)$ is estimated using a plug-in estimator over Monte Carlo samples. The conditional entropy $H(W|S)$ is analytically constant for a fixed $\alpha$, meaning variation in the estimate arises solely from the marginal distribution of selected experts.
   • Routing Information ($I(X; T)$): The paper extends the framework to input-dependent routing using a Blahut-Arimoto solver. This solves a constrained optimization problem to trace an empirical rate-distortion curve $\hat{D}(\rho)$, balancing the routing rate $\rho$ against the expected loss.

Key Contributions

1. Finite-Bank Protocol: Introduction of a tractable MNIST-based protocol that recasts deep model selection as an information-theoretic experiment, allowing for exact entropy-based evaluation of $I(S; W)$.
2. Estimator Derivation: Development of a Monte Carlo estimator for $I(S; W)$ induced by the $\alpha$-mixture rule, accompanied by a Miller-Madow bias correction and bootstrap confidence intervals.
3. Baseline Comparison: A rigorous comparison against a uniform union-bound baseline. The authors demonstrate that while the union bound is tighter for deterministic ERM ($\alpha=1$), the information-theoretic bound provides unique insight into the continuous dependence of the generalization gap on gate randomization ($\alpha < 1$).
4. Empirical Rate-Distortion: Implementation of a Blahut-Arimoto solver on the fixed expert bank to generate an empirical accuracy-rate curve, tracing the trade-off between routing information and distortion.

Results

• Monotonic Tracking: As the data-dependence parameter $\alpha$ increases from 0 to 1, the estimated mutual information $\hat{I}(S; W)$ increases monotonically (from 0 to $\approx 2.53$ nats). Crucially, the generalization gap also increases monotonically, confirming that the information-theoretic bound qualitatively tracks the behavior of the learning algorithm.
• Bound Looseness: The Xu-Raginsky bound proxy $\sqrt{2\hat{I}(S; W)/m}$ is found to be loose by a factor of approximately 14–20$\times$ compared to the observed generalization gap. The authors attribute this to the known constant-factor looseness of the Xu-Raginsky bound rather than a failure of the estimator.
• Union Bound Comparison: For deterministic selection ($\alpha \to 1$), the information-theoretic bound converges to $\sqrt{2 \log R / m}$, which is exactly twice the standard union bound. However, for randomized gates ($\alpha < 1$), the MI bound captures the reduction in generalization error due to randomization, a regime where the union bound offers no such granularity.
• Rate-Distortion Curve: The Blahut-Arimoto solver successfully traces a strictly decreasing convex curve $\hat{D}(\rho)$, demonstrating that reducing the routing rate (by collapsing the gate to a single expert) increases distortion, while maximizing the rate minimizes distortion.

Significance and Claims
The paper positions this work as an engineering validation of the information-generalization mechanism in MoE systems, rather than a claim of solving the full high-dimensional rate-distortion problem for continuous neural networks.

• Tractability: The primary significance is making $I(S; W)$ observable and computable in a deep-learning context, bridging the gap between synthetic information-theoretic examples and standard benchmarks.
• Design Proxy: The framework provides a practical tool for interpreting $I(X; T)$ and the rate-distortion term $D(R_g)$ as design proxies for resource-aware inference. This is particularly relevant for communication-constrained domains (e.g., aeronautics, aerospace, edge computing) where routing decisions involve trade-offs between bandwidth, latency, and accuracy.
• Theoretical Insight: The work clarifies the distinct roles of two information quantities: $I(S; W)$, which measures the algorithm's dependence on the training data (generalization), and $I(X; T)$, which measures the information communicated to the expert bank at inference time (efficiency).

The authors explicitly note limitations, including the use of the relatively simple MNIST dataset, the restriction to a discrete expert bank (which precludes direct claims about continuous parameter spaces), and the inherent looseness of the Xu-Raginsky bound. Future work is suggested to involve end-to-end training of input-dependent gates and testing on more complex datasets and architectures.