Soft Merging of Experts with Adaptive Routing

TL;DR

This work tackles the inefficiencies of gradient-estimated routing in sparsely activated models by introducing SMEAR, a differentiable approach that builds a single merged expert from a router's distribution over multiple experts. By routing through a merged parameter set, SMEAR preserves end-to-end trainability while maintaining comparable compute to single-expert routing and outperforming heuristic and gradient-estimator baselines. Empirical results on T5-GLUE and DomainNet show SMEAR achieving competitive or superior accuracy, with clear signs of expert specialization and robust performance across tasks. The approach offers a practical route to modular models with adaptive routing that scales without the cost penalties typical of ensembles or complex gradient estimators.

Abstract

Sparsely activated neural networks with conditional computation learn to route their inputs through different "expert" subnetworks, providing a form of modularity that densely activated models lack. Despite their possible benefits, models with learned routing often underperform their parameter-matched densely activated counterparts as well as models that use non-learned heuristic routing strategies. In this paper, we hypothesize that these shortcomings stem from the gradient estimation techniques used to train sparsely activated models that use non-differentiable discrete routing decisions. To address this issue, we introduce Soft Merging of Experts with Adaptive Routing (SMEAR), which avoids discrete routing by using a single "merged" expert constructed via a weighted average of all of the experts' parameters. By routing activations through a single merged expert, SMEAR does not incur a significant increase in computational costs and enables standard gradient-based training. We empirically validate that models using SMEAR outperform models that route based on metadata or learn sparse routing through gradient estimation. Furthermore, we provide qualitative analysis demonstrating that the experts learned via SMEAR exhibit a significant amount of specialization. All of the code used in our experiments is publicly available.

Figures (28)

• Figure 1: The discrete routing decisions commonly used in models that route activations among experts require the use of gradient estimation (left). We propose SMEAR (right), which uses a given router's distribution to average the parameters of the corresponding experts and then routes the input through a single merged expert. SMEAR achieves better performance than models with discrete routing, can be trained with standard backpropagation, and does not incur significant additional computational costs.
• Figure 2: Average accuracy and inference speed (in examples processed per second) for models using different routing approaches on our T5-GLUE and ResNet-DomainNet settings. Routing approaches are grouped by color; groups are (in order of the legend) our method (SMEAR), methods that do not use adaptive routing ($1\times$ compute, $1\times$ parameters, Adamix, and Hash), methods that make use of metadata (Tag and Latent Skills), methods that learn adaptive routing (Top-$k$, ST-Gumbel, REINFORCE, DSelect-$k$, and Soft MoE), and methods that ensemble expert outputs (Ensemble). We omit Hash routing from the plots because its poor performance (66.9% on T5-GLUE and 52.4% on ResNet-DomainNet) hampers readability. Exact numerical results for all methods and standard deviation across five runs are provided in \ref{['sec:full_results']}.
• Figure 3: Average routing distributions produced by SMEAR for two routers from the T5-GLUE model and two from the ResNet-DomainNet model. For a given router, we average all routing distributions across all examples from a given dataset.
• Figure 4: Routing distribution learnt by SMEAR in the routing blocks of ResNet-DomainNet
• Figure 5: Routing distribution learnt by SMEAR in the encoder routing blocks (1-24) of T5-GLUE