H yper R outer: Towards Efficient Training and Inference of Sparse Mixture of Experts (2023)

Abstract: By routing input tokens to only a few split experts, Sparse Mixture-of-Experts has enabled efficient training of large language models. Recent findings suggest that fixing the routers can achieve competitive performance by alleviating the collapsing problem, where all experts eventually learn similar representations. However, this strategy has two key limitations: (i) the policy derived from random routers might be sub-optimal, and (ii) it requires extensive resources during training and evaluation, leading to limited efficiency gains. This work introduces HyperRouter, which dynamically generates the router’s parameters through a fixed hypernetwork and trainable embeddings to achieve a balance between training the routers and freezing them to learn an improved routing policy. Extensive experiments across a wide range of tasks demonstrate the superior performance and efficiency gains of HyperRouter compared to existing routing methods. Our implementation is publicly available at https://github.com/giangdip2410/HyperRouter .

## introduction recent years have witnessed tremendous successes of the transformer model @xcite and its variants across a wide range of tasks, ranging from natural language and speech processing @xcite , computer vision @xcite @xcite , reinforcement learning @xcite , to life sciences @xcite . since then, scaling up to larger models has become the prevailing approach for advancing the state-of-the-art in pre-training and finetuning tasks. however, training such large models comes with a high computational cost @xcite ; therefore, there is a growing need to develop efficient strategies that facilitate the training of large language models (llms) @xcite . one of the most effective strategies thus far is the sparse mixture-of-experts (smoe) @xcite , which utilizes routers to direct each input token to a small subset of network parameters (experts). smoe improves both efficiency and performance compared to dense training approaches @xcite @xcite . despite the encouraging results, smoe has been found to encounter the issue of representation collapse, where all experts converge to similar representations @xcite . this problem arises due to the learning process of the router encouraging experts to cluster around a centroid @xcite . therefore, significant efforts have been dedicated to addressing the representation collapse issue while preserving the simplicity and efficiency of smoe training. to this end, one of the most effective strategies is freezing the router, as demonstrated by @xcite , where a randomly initialized router remains fixed throughout the training process. additionally, @xcite found that progressively increasing the number of selected experts can be beneficial. however, we argue that such a naive strategy exhibits two key limitations. firstly, the random routing policy may be sub-optimal, which hinders the overall training process. secondly, we will show in section 3.3 that a fixed router will restrict the model's representation capabilities, necessitating a progressive increase in the number of chosen experts to achieve satisfactory performance. therefore, smoe-dropout inherently suffers from limited representation and does not offer efficiency gains during inference. this work introduces a novel approach called hyperrouter to address the trade-off between fixed and trainable routers in smoe training. instead of employing a random but fixed router, hyperrouter utilizes a random but fixed hypernetwork @xcite to generate the router's parameters based on a trainable router embedding. by doing so, hyperrouter can improve the routing policy throughout training while mitigating the representation collapse issues. to demonstrate its effectiveness, we conduct extensive evaluations on various nlp tasks, comparing hyperrouter to several state-of-the-art smoe routing strategies. moreover, hyperrouter achieves the same performance threshold with fewer experts (compute) during inference, thereby significantly enhancing the efficiency of deploying llms in real-world applications. fig. 1 shows that hyperrouter consistently outperforms other competitors when the same number of experts is used during inference. ## related work existing routing mechanism for smoe. one of the most important components for training smoe is the expert routing strategy, which specifies experts to process input tokens. there are two common classes of token-expert assignment algorithms for smoe training: (1) letting tokens select the top-k experts and (2) letting experts select the topk tokens. for the first approach, @xcite , @xcite , ( zuo introduced expert choice that enables selecting different experts for each token and demonstrates the potential of the second approach. ## methodology this section describes smoe training and details of the proposed hyperrouter method. ## experiments we conduct experiments to investigate the following hypotheses. first, hyperrouter improves efficiency compared to existing routing strategies, requiring less computation to achieve similar performance thresholds during evaluation. second, with the same amount of experts used during inference, hyperrouter can achieve better performances compared to other competitors. baselines use the same amount of trainable parameters, while our hyperrouter introduces a neglec ## conclusion in this work, we investigated the potentials and limitations of smoe for training llms and proposed hyperrouter to balance between two extremes of smoe: training and freezing the router. consequently, hyperrouter can learn better routing policies while alleviating the representation collapse of conventional smoe training. experiments on both pre-training and fine-tuning tasks demonstrated promising hyperrouter's capabilities in facilitating efficient, effective training and inference compared to state-of-the-art strategies.