StructMoE: Augmenting MoEs with Hierarchically Routed Low Rank Experts

Authors: Zain Sarwar, Ashwinee Panda, Benjamin Thérien, Stephen Rawls, Sambit Sahu, Supriyo Chakraborty
Primary area: foundation or frontier models, including LLMs
Keywords: moe, mixture of experts, LLM, transformer

Abstract:
The traditional approach to scaling Mixture of Experts for transformer models has been to increase the total number of experts. While performance improves with more experts, the gains are diminshing whereas
memory scales linearly with the number of experts. We introduce $\textit{StructMoE}$, a scaling approach for Mixture of Experts which augments experts with additional dynamic capacity using routed structured matrices which we refer
to as $\textbf{L}$ow $\textbf{R}$ank $\textbf{E}$xprts ($\textbf{$\textit{LoRE}$}$). At a high-level, we introduce hierarchical MoEs where the first level of routing decides which expert each token
should be routed to and the second level of routing decides which $\textit{LoRE}$ should each token be routed through. The outputs of the expert and the $\textit{LoRE}$ are then entangled together to provide
the final output. This introduces more dynamism into the model which has empirically been demonstrated to improve model performance. We find this scaling approach to outperform a standard MoE baseline in terms of loss on a held out validation. Thus, we propose this to be an effective scaling technique for MoEs compared to the standard approach of adding more 
experts to the model.

Decision: reject

Reviews and discussion:

Metareview: