Mixture of Experts Softens the Curse of Dimensionality in Operator Learning

2024-04-13

AI Summary

To address the computational and memory bottlenecks imposed by the curse of dimensionality in high-dimensional operator learning, this paper proposes the Mixture-of-Experts Neural Operator (MoNO): a framework that decomposes a global nonlinear operator into multiple lightweight expert sub-operators, routed input-adaptively via a learnable decision-tree mechanism. Theoretically, we establish the first distributed universal approximation theorem, proving that MoNO uniformly approximates any Lipschitz-continuous nonlinear operator in Sobolev spaces; each expert’s depth, width, and rank scale as O(ε⁻¹), ensuring controllable memory footprint compatible with standard hardware. We further derive the first quantitative approximation rate for classical neural operators. Experiments and theory jointly demonstrate that MoNO achieves ε-accuracy with significantly reduced complexity, overcoming both expressive and deployability limitations inherent to monolithic neural operators.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Economics, Online Markets and Human Computation: Social networks and social learning
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We study the approximation-theoretic implications of mixture-of-experts architectures for operator learning, where the complexity of a single large neural operator is distributed across many small neural operators (NOs), and each input is routed to exactly one NO via a decision tree. We analyze how this tree-based routing and expert decomposition affect approximation power, sample complexity, and stability. Our main result is a distributed universal approximation theorem for mixture of neural operators (MoNOs): any Lipschitz nonlinear operator between $L^2([0,1]^d)$ spaces can be uniformly approximated over the Sobolev unit ball to arbitrary accuracy $varepsilon>0$ by an MoNO, where each expert NO has a depth, width, and rank scaling as $mathcal{O}(varepsilon^{-1})$. Although the number of experts may grow with accuracy, each NO remains small, enough to fit within active memory of standard hardware for reasonable accuracy levels. Our analysis also yields new quantitative approximation rates for classical NOs approximating uniformly continuous nonlinear operators uniformly on compact subsets of $L^2([0,1]^d)$.

Problem

• Analyzes mixture-of-experts for operator learning to reduce model complexity.
• Studies how tree-based routing affects approximation power and sample complexity.
• Proves distributed universal approximation for Lipschitz operators using small neural operators.

Innovation

• Mixture of Experts distributes complexity across small neural operators
• Tree-based routing directs each input to a single expert
• Each expert scales with accuracy, fitting hardware memory constraints