PaperlayerView exact revision on arXiv

Source-linked AI summary

Soft Merging of Experts with Adaptive Routing

Mohammed Muqeeth, Haokun Liu, Colin Raffel

arXiv:2306.03745v2cs.LG

TL;DR

Discrete expert routing can require gradient estimation and learned routing may underperform heuristic routing. SMEAR averages expert parameters according to a learned router and processes inputs through one merged expert; across tested settings, it matches or outperforms other routing strategies.

  • Problem

    Discrete routing prevents the prediction loss from backpropagating through the routing decision, and learned routing has been reported to perform poorly relative to oracle routing.

    Source: §1 · p. 2

  • Method

    SMEAR weights expert parameters using the router’s distribution and sends the input through one merged expert, allowing standard gradient-based training.

    Source: §1 · p. 2

  • Results

    SMEAR matches or outperforms every other routing strategy in the reported settings, improving over tag routing by 2.7% on T5-GLUE and 0.6% on ResNet-DomainNet.

    Source: §4.2 · p. 7

  • Takeaways & Limitations

    The results support SMEAR as an effective alternative for adaptive routing among specialized expert subnetworks.

    Source: §1 · p. 2

    The cost comparison assumes one routing choice per example; token-level routing would require averaging expert parameters L times, making SMEAR’s cost similar to ensemble routing.

    Source: §3 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

Specialized parameters can support modularity, but discrete expert routing prevents ordinary gradients from updating the router. SMEAR instead uses the router’s distribution to merge expert parameters, aiming to retain adaptive routing while avoiding gradient estimation.

1 Introduction

  • Experts can specialize across diverse inputs, offering modularity that may reduce interference or let capabilities be added, removed, or modified through corresponding parameters.The paper does not focus on the separate potential benefit of decoupling computational cost from parameter count.Source: §1 · p. 1
  • Discrete routing blocks backpropagation through the routing decision, so conditional-computation models often rely on gradient-estimation techniques.Prior findings cited here include learned routing that performs poorly relative to oracle routing and random routing that can match learned routing.Source: §1 · p. 2
  • SMEAR weights expert parameters by the router’s distribution and processes inputs through one merged expert, enabling standard gradient-based training.The introduction reports better performance than discrete-routing models without significant additional computational costs.Source: §1 · p. 2, Figure 1 · p. 2
  • The paper tests whether gradient-estimation methods underperform hand-designed routing and reports that SMEAR improves the performance/cost tradeoff against learned, heuristic, and modular-model baselines.The experiments also qualitatively assess expert specialization and parameter sharing across related tasks.Source: §1 · p. 2
  • The introduction previews background, method, experiments, related work, and conclusion sections.Source: §1 · p. 2

2 Background

The background introduces expert routing and surveys gradient estimators, hand-designed routing, and methods for learning modular models. The methods provide the comparison points used to evaluate adaptive routing.

2 Background

  • The background section frames conditional computation, gradient estimators, heuristic routes, and existing modular-learning methods as context for the proposed comparisons.Source: §2 · p. 3

2.1 Routing Among Experts

  • An expert-routing block contains experts with matching input and output dimensions plus a router that selects among them for a hidden-state representation.Source: §2.1 · p. 3

2.2 Gradient Estimators

  • Gradient-estimation baselines include REINFORCE with a learned variance-reducing baseline, ST-Gumbel with straight-through gradients, top-k selection, and DSelect-k evaluated with k=1.The estimators provide approximate router gradients for discrete routing decisions.Source: §2.2 · pp. 3–4

2.3 Heuristic Routing

  • Tag routing assigns experts using example metadata, hash routing fixes a random expert choice, and single-expert baselines match either compute or parameter count.Source: §2.3 · p. 4

2.4 Methods for Learning Modular Models

  • Modular-model baselines include Adamix’s random training routes and averaged experts, task-conditioned Latent Skills, and differentiable Soft MoE input-slot routing.Source: §2.4 · p. 5

3 Soft Merging of Experts with Adaptive Routing

SMEAR targets adaptive expert routing without the gradient estimation required by discrete decisions. It averages expert parameters according to a learned router and processes each input with one merged expert, with compute costs depending on routing granularity.

3 Soft Merging of Experts with Adaptive Routing

  • Output ensembling permits exact gradients but costs more because it computes every expert’s output instead of just one.Source: §3 · p. 5
  • SMEAR forms one expert by averaging expert parameters according to router probabilities, then uses that expert’s output as the routing-block output.The experts must share an architecture so their parameters can be averaged through a one-to-one mapping.Source: §3 · p. 5
  • Because every SMEAR operation is differentiable, it supports standard end-to-end gradient training and adaptive routing without hand-specified tags.The input is processed only by the merged expert; experts start from different random parameters to break symmetry.Source: §3 · p. 6
  • For a length-L input, SMEAR’s stated cost is approximately (L × 4 + N × 2) × d × m FLOPs, versus N × L × 4 × d × m for ensemble routing.The estimate assumes a dense two-layer expert and ignores the nonlinearity’s relatively minor cost.Source: §3 · p. 6
  • SMEAR’s roughly discrete-routing-level cost assumes one routing choice per example; choosing experts separately for each token would require averaging parameters L times and make its cost similar to ensembling.The paper gives T5-GLUE’s L=128 and N=8 setting as an example where runtime difference is minimal.Source: §3 · p. 6

4 Experiments

The experiments evaluate SMEAR in modular fine-tuning of T5.1.1 on GLUE and ResNet18 on DomainNet. SMEAR matches or outperforms the other routing strategies across both settings, with qualitative routing patterns showing specialization and sharing.

4 Experiments

  • The evaluation compares learned routing with heuristic strategies in T5-GLUE and ResNet-DomainNet, two settings with metadata suitable for Tag routing.The paper adapts pretrained backbones by adding experts and leaves large-scale from-scratch experiments for future work.Source: §4 · p. 6

4.1 Settings

  • Experiments adapt frozen pretrained T5.1.1 Base across eight GLUE experts and ResNet18 across six DomainNet experts, reporting main-paper averages across datasets and domains.The study focuses on modular fine-tuned models and leaves large-scale training from scratch for future work.Source: §4 · p. 6, §4.1 · p. 7

4.2 Results

  • Gradient-estimated routing often trails heuristic routing: all considered estimators underperform Tag on ResNet-DomainNet, while REINFORCE beats Tag on T5-GLUE.Source: §4.2 · p. 7
  • SMEAR matches or outperforms every other routing strategy, improving over Tag routing by 2.7% on T5-GLUE and 0.6% on ResNet-DomainNet.Figure 2 summarizes average accuracy and inference speed across the two settings.Source: §4.2 · p. 7, Figure 2 · p. 8
  • SMEAR matches Ensemble on T5-GLUE and modestly trails it on ResNet-DomainNet, while beating Adamix by 3.2% and 4%, respectively.It also beats Soft MoE by 1.5% and DSelect-k by 2.7% on ResNet-DomainNet; Latent Skills requires oracle tags, unlike SMEAR.Source: §4.2 · p. 8
  • SMEAR's inference speed is almost identical to discrete adaptive routing, whereas Ensemble is 1.2× slower on T5-GLUE and 1.3× slower on ResNet-DomainNet.Source: §4.2 · p. 8
  • Doubling experts improves ResNet-DomainNet by 0.8%, but makes no significant difference on T5-GLUE (81.3 ± 1.1 vs. 81.6 ± 1.1).The experiments initially used one expert per task or domain.Source: §4.2 · p. 8

4.3 Qualitative Analysis

  • SMEAR routing is generally sparse; GLUE routers show varied task sharing, while DomainNet routers specialize less, with Quickdraw focused on two experts and Painting and Real correlated.Other learned routers often show degenerate or less intuitive specialization and sharing.Source: §4.3 · pp. 9–10

5 Related Work

Related work examines discrete expert-routing strategies and their training challenges, as well as parameter averaging and alternative merging methods. SMEAR connects adaptive routing with expert averaging while targeting exact gradient computation.

5 Related Work

  • Prior work trains discrete-routing models with gradual sparsity, task-specific routing, consistency regularization, or hand-designed heuristic assignments.Other work studies sparse pretrained models, expert load balancing, and converting dense models to sparse ones.Source: §5 · p. 10
  • Gradient estimators address bias and variance in discrete decisions, while SMEAR enables exact gradients through standard backpropagation.Preliminary experiments found no major gains from more sophisticated estimators, though estimators designed for discrete routing could yield stronger strategies.Source: §5 · pp. 10–11
  • Prior work reports conditional-computation problems including scaling cutoffs, representation collapse, unstable routing, and performance below simple multitask baselines.The authors present fully differentiable routing without increased computational costs as a way to avoid these issues.Source: §5 · p. 11
  • Parameter averaging can transfer or improve models, and SMEAR’s router identifies experts with similar generalization properties for each input.The cited work connects weight averaging with linear mode connectivity and task similarity.Source: §5 · p. 11
  • Unlike convolution-kernel averaging, SMEAR averages nonlinear, parameter-efficient experts and learns routing without metadata, unlike π-Tuning’s downstream-task interpolation.Other cited methods average models using validation performance, novel-domain data, or a training trajectory; SMEAR learns adaptive routing during training.Source: §5 · p. 11
  • Concurrent work applies SMEAR to LoRA- and IA3-based experts, while more sophisticated parameter-merging methods remain a future direction.The cited concurrent work reports generalization to unseen tasks that outperforms full-model fine-tuning.Source: §5 · pp. 11–12

6 Conclusion

SMEAR replaces discrete expert routing with a learned weighted average of expert parameters. Across the reported settings, it outperforms discrete and heuristic routing approaches without increasing computational costs, while its experts specialize to different inputs.

6 Conclusion

  • SMEAR outperformed discrete-routing and performant heuristic-routing models across settings, with no increase in computational costs.The paper also reports matching or outperforming state-of-the-art expert-averaging methods without oracle task labels.Source: §6 · p. 12
  • SMEAR computes a weighted average of expert parameters using a learned router, enabling exact-gradient routing instead of discrete decisions.The work targets shortcomings that can lead discrete-routing models to underperform heuristic non-learned routing.Source: §6 · p. 12
  • Qualitative analysis found that SMEAR’s experts specialize across input types and that its router exploits commonalities between examples.Source: §6 · p. 12
  • The authors identify alternative expert architectures and merging methods as future directions.They also propose testing large-scale settings if more compute becomes available, to assess whether SMEAR helps with poor scaling properties of discrete routing.Source: §6 · p. 12

A Compute resources used

Experiments used A6000 GPUs, except for the Ensemble method, and required substantially different training times across the two settings.

A Compute resources used

B Experiment Details

The appendix reports implementation and memory details, then specifies the T5-GLUE datasets, prompting protocol, and training configuration.

B Experiment Details

  • The appendix describes Adamix and Latent Skills implementation choices and hyperparameter selection.For Latent Skills, the skill-matrix learning-rate ratio was selected by sweeping 1, 10, and 100 in both settings.Source: B Experiment Details · p. 19
  • SMEAR and other methods use around 30GB of memory for T5-GLUE, with no significant memory differences observed for ResNet-DomainNet.Experts are prepared per example and applied in parallel using torch.bmm, without an extra weight-sized tensor.Source: B Experiment Details · p. 19

B.1 T5-GLUE

  • T5-GLUE uses eight GLUE datasets, prompted with randomly selected templates during training and all templates during evaluation.WNLI is excluded; training concatenates the eight datasets and runs T5 for 600k steps with a 3e−4 learning rate, 2k warmup steps, and batch size 128.Source: §B.1 · p. 19

B.2 ResNet-DomainNet

The ResNet-DomainNet appendix specifies training settings and estimator-specific hyperparameters.

B.2 ResNet-DomainNet

  • ResNet models were trained for 100k steps with batch size 128, learning rate 1e-3, no warm-up, and Adam.Source: §B.2 · p. 19
  • The ST-Gumbel estimator used τ = 10 and an anneal rate of 1e-4, selected by sweeping the stated candidate values.Source: §B.2 · p. 19
  • Dselect-k used entropy-regularization weight 0.1, and the adapters used an intervening swish non-linearity.Source: §B.2 · p. 19

C Expert dropout

This section reports expert-dropout effects, dataset and model licenses, and the scope of the ethics statement.

C Expert dropout

  • On a single seed, dropout raises SMEAR by 2.9% on T5-GLUE and Top-k by 1.2% on T5-GLUE and 0.2% on ResNet-DomainNet, but hurts ST-Gumbel and REINFORCE.Source: C Expert dropout · pp. 19–20

D License

  • T5 is Apache 2.0 licensed and the ResNet model is BSD 3-Clause; seven named datasets allow non-commercial research use.Source: D License · p. 20

E Ethics Statement

  • The authors report no known negative ethical implications or human subjects, and state that their work focuses on smaller-scale models despite the high costs of extremely large conditional-computation networks.Source: E Ethics Statement · p. 20

F Full results on T5-GLUE and ResNet-DomainNet

The full results are presented for both evaluation settings.

F Full results on T5-GLUE and ResNet-DomainNet

G Routing distribution in all routing blocks

The section presents routing distributions across SMEAR and other routing methods, alongside full-results tables for both evaluation settings.

G Routing distribution in all routing blocks

Loading 2306.03745v2…