Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization
<|reference_start|>Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization: The Mixture of Experts (MoE) paradigm provides a powerful way to decompose dense layers into smaller, modular computations often more amenable to human interpretation, debugging, and editability. However, a major challenge lies in the computational cost of scaling the number of experts high enough to achieve fine-grained specialization. In this paper, we propose the Multilinear Mixture of Experts ($\mu$MoE) layer to address this, focusing on vision models. $\mu$MoE layers enable scalable expert specialization by performing an implicit computation on prohibitively large weight tensors entirely in factorized form. Consequently, $\mu$MoEs (1) avoid the restrictively high inference-time costs of dense MoEs, yet (2) do not inherit the training issues of the popular sparse MoEs' discrete (non-differentiable) expert routing. We present both qualitative and quantitative evidence that scaling $\mu$MoE layers when fine-tuning foundation models for vision tasks leads to more specialized experts at the class-level, further enabling manual bias correction in CelebA attribute classification. Finally, we show qualitative results demonstrating the expert specialism achieved when pre-training large GPT2 and MLP-Mixer models with parameter-matched $\mu$MoE blocks at every layer, maintaining comparable accuracy. Our code is available at: https://github.com/james-oldfield/muMoE.<|reference_end|>
@article{oldfield2024multilinear,
    title={Multilinear Mixture of Experts: Scalable Expert Specialization through
  Factorization},
    author={James Oldfield, Markos Georgopoulos, Grigorios G. Chrysos, Christos
  Tzelepis, Yannis Panagakis, Mihalis A. Nicolaou, Jiankang Deng, Ioannis
  Patras},
    journal={arXiv preprint arXiv:2402.12550},
    year={2024},
    archivePrefix={arXiv},
    eprint={2402.12550},
    primaryClass={cs.CV cs.LG}
}