Soft Merging of Experts with Adaptive Routing

Mohammed Muqeeth, Haokun Liu, Colin Raffel

Primary Area: transfer learning, meta learning, and lifelong learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Modular models, Routing, Weight Averaging

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We propose SMEAR, a performant and efficient method for training modular models that allows standard gradient-based training.

Abstract:

Neural networks that learn to route their inputs through different "expert" subnetworks provide a form of modularity that standard dense models lack.
Despite their possible benefits, modular models with learned routing often underperform their parameter-matched dense 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 modular models that use non-differentiable discrete routing decisions.
To address this issue, we introduce $\textbf{S}$oft $\textbf{M}$erging of $\textbf{E}$xperts with $\textbf{A}$daptive $\textbf{R}$outing (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 routing through gradient estimation.
Furthermore, we provide qualitative analysis demonstrating that the experts learned via SMEAR exhibit a significant amount of specialization.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6566