Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Liao, Fangshuo, Kyrillidis, Anastasios
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866909831514816512
author Liao, Fangshuo
Kyrillidis, Anastasios
author_facet Liao, Fangshuo
Kyrillidis, Anastasios
contents Mixture-of-Experts (MoE) architectures have emerged as a cornerstone of modern AI systems. In particular, MoEs route inputs dynamically to specialized experts whose outputs are aggregated through weighted summation. Despite their widespread application, theoretical understanding of MoE training dynamics remains limited to either separate expert-router optimization or only top-1 routing scenarios with carefully constructed datasets. This paper advances MoE theory by providing convergence guarantees for joint training of soft-routed MoE models with non-linear routers and experts in a student-teacher framework. We prove that, with moderate over-parameterization, the student network undergoes a feature learning phase, where the router's learning process is ``guided'' by the experts, that recovers the teacher's parameters. Moreover, we show that a post-training pruning can effectively eliminate redundant neurons, followed by a provably convergent fine-tuning process that reaches global optimality. To our knowledge, our analysis is the first to bring novel insights in understanding the optimization landscape of the MoE architecture.
format Preprint
id arxiv_https___arxiv_org_abs_2510_07205
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts
Liao, Fangshuo
Kyrillidis, Anastasios
Machine Learning
Optimization and Control
Mixture-of-Experts (MoE) architectures have emerged as a cornerstone of modern AI systems. In particular, MoEs route inputs dynamically to specialized experts whose outputs are aggregated through weighted summation. Despite their widespread application, theoretical understanding of MoE training dynamics remains limited to either separate expert-router optimization or only top-1 routing scenarios with carefully constructed datasets. This paper advances MoE theory by providing convergence guarantees for joint training of soft-routed MoE models with non-linear routers and experts in a student-teacher framework. We prove that, with moderate over-parameterization, the student network undergoes a feature learning phase, where the router's learning process is ``guided'' by the experts, that recovers the teacher's parameters. Moreover, we show that a post-training pruning can effectively eliminate redundant neurons, followed by a provably convergent fine-tuning process that reaches global optimality. To our knowledge, our analysis is the first to bring novel insights in understanding the optimization landscape of the MoE architecture.
title Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2510.07205