PaperlayerView exact revision on arXiv

Source-linked AI summary

From Sparse to Soft Mixtures of Experts

Joan Puigcerver, Carlos Riquelme, Basil Mustafa, Neil Houlsby

arXiv:2308.00951v2cs.LGcs.AIcs.CV

TL;DR

Sparse MoE Transformers face discrete routing challenges, including expert-balancing difficulties and unassigned tokens. Soft MoE replaces hard routing with differentiable token mixing, and it outperforms dense ViTs and popular sparse MoEs while retaining efficient scaling. Its use in autoregressive decoders remains an open research avenue.

  • Problem

    Sparse MoE Transformers require discrete token-to-expert assignments and often auxiliary losses to balance expert utilization and minimize unassigned tokens.

  • Method

    Soft MoE replaces hard routing with soft assignment by sending each expert a weighted average of all input tokens, while experts process only selected combined tokens.

  • Results

    Soft MoE outperforms ViTs and popular sparse routers across reported visual-recognition evaluations, including upstream, few-shot, finetuning, zero-shot classification, and LAION-400M training.

  • Takeaways & Limitations

    Soft MoE can increase expert capacity efficiently: Soft MoE B/16 is 5.7× faster at inference than ViT H/14 despite 5.5× more parameters.

  • Takeaways & Limitations

    Using Soft MoE in autoregressive decoders is difficult because learned token mixing must preserve causality without biasing expert training by token or slot indices.

Abstract

from arXiv · show

Sparse mixture of expert architectures (MoEs) scale model capacity without significant increases in training or inference costs. Despite their success, MoEs suffer from a number of issues: training instability, token dropping, inability to scale the number of experts, or ineffective finetuning. In this work, we propose Soft MoE, a fully-differentiable sparse Transformer that addresses these challenges, while maintaining the benefits of MoEs. Soft MoE performs an implicit soft assignment by passing different weighted combinations of all input tokens to each expert. As in other MoEs, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity (and performance) at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs (Tokens Choice and Experts Choice). Furthermore, Soft MoE scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, with only 2% increased inference time, and substantially better quality.

1 INTRODUCTION

Sparse MoEs scale model capacity by activating expert parameters selectively, but discrete routing creates optimization, balancing, and transfer challenges. Soft MoE replaces hard token-to-expert assignments with differentiable weighted token mixtures and achieves stronger visual-recognition results with lower training or inference costs.

  • Motivation: Sparse MoEs use discrete routing to match tokens with experts, often requiring auxiliary losses to balance utilization and reduce unassigned tokens.These challenges can worsen with small inference batches, novel inputs, or transfer learning.
  • Soft MoE: Soft MoE forms multiple weighted averages of all input tokens and processes each average with its corresponding expert.The weights depend on both tokens and experts, avoiding hard assignments.
  • Results: Soft MoE L/16 outperforms ViT H/14 on upstream, few-shot, and finetuning evaluations while requiring almost half the training time and 2× faster inference.Soft MoE B/16 matches ViT H/14 on few-shot and finetuning, outperforms it on upstream metrics, and is 5.7× faster at inference despite 5.5× more parameters.
  • Results: Soft routing preserves benefits for image-text alignment when representations are used in a contrastive text-tower setting.The image tower is frozen while the text encoder is trained against it.

2 SOFT MIXTURE OF EXPERTS

Soft MoE replaces discrete token routing with differentiable dispatch and combine operations over token-slot mixtures. Its cost depends primarily on the number of slots, enabling scalable expert capacity while avoiding several sparse-routing and normalization problems.

  • Algorithm: Soft MoE represents each layer’s input as m tokens and assigns n expert functions, with each expert processing p parameterized slots.The slot parameters form Φ ∈ R^d×(n·p), and experts operate on individual d-dimensional slots.
  • Algorithm: Each input slot is a convex combination of all input tokens, using dispatch weights obtained by applying a softmax over the columns of XΦ.This creates different token mixtures for the available slots.
  • Algorithm: After expert processing, each output token is a convex combination of all output slots using combine weights obtained by applying a softmax over the rows of XΦ.Slots are allocated to experts round-robin.
  • Properties: The total number of slots controls Soft MoE layer cost, while the number of experts is irrelevant when total slots remain fixed.The slot count can equal sequence length to match the FLOPs of an equivalent dense Transformer.
  • Properties: Soft MoE avoids discrete assignment operations and token-dropping or expert-unbalance issues associated with classical sparse routers.Unlike dense MoEs, each expert processes only a subset of slots, although each slot mixes all input tokens.
  • Implementation: With p = O(m/n) slots per expert, Soft MoE time complexity becomes O(m^2d + mk), allowing more experts and parameters without increasing this asymptotic cost.Throughput barely changes as experts increase from 8 to 4,096 in the cited comparison.
  • Implementation: L2 normalization stabilizes Soft MoE scaling by preventing the routing softmax from becoming nearly one-hot as model dimension increases.It enables larger model dimensions and/or higher learning rates, while having little quality impact for relatively small dimensions.

3 IMAGE CLASSIFICATION EXPERIMENTS

The experiments compare Soft MoE with dense ViTs and sparse Tokens Choice and Experts Choice models across training budgets, inference costs, expert counts, and routing ablations. Soft MoE generally delivers better quality-cost trade-offs, scales to many experts, and benefits from token mixing.

  • 3.3 TRAINING PARETO-OPTIMAL MODELS: Soft MoE strongly outperforms dense ViTs and other sparse approaches on both training cost/performance Pareto metrics.The comparison covers models on their respective Pareto frontiers, with costs measured by FLOPs and TPU-v3 training time.
  • 3.4 LONG TRAINING DURATIONS: Soft MoE outperforms ViT across JFT precision, ImageNet 10-shot accuracy, and ImageNet finetuning at fixed training budgets.Equivalent model classes have similar training costs, while Soft MoE achieves higher performance on all reported metrics.
  • 3.4 LONG TRAINING DURATIONS: 10× cheaper inference in FLOPs, Soft MoE B/16 matches or outperforms ViT H/14 at a similar training budget.The reported inference costs are 32 versus 334 GFLOPS/image; wall-clock inference is over 5× cheaper, at 1.5 versus 8.6 ms/image.
  • 3.4 LONG TRAINING DURATIONS: Soft MoE L/16 outperforms all ViT models while being almost 2× faster at inference than ViT H/14.The reported inference times are 4.8 versus 8.6 ms/image.
  • 3.5 NUMBER OF SLOTS AND EXPERTS: Increasing expert count improves Soft MoE performance at each slots-per-token setting while keeping throughput approximately constant.For Sparse MoEs, additional experts eventually encounter training difficulties and throughput falls dramatically from 1k experts.
  • 3.6 ABLATIONS: Routing ablations show that learned token mixing is important: Identity and Uniform routing substantially underperform Soft MoE, while dispatch mixing is slightly more important than combine mixing.The ablations use Soft MoE-S/14 with 256 experts trained for 300k steps.

4 CONTRASTIVE LEARNING

Soft MoE representations retain their advantages beyond image classification, improving zero-shot image recognition while showing only modest gains on open-vocabulary retrieval.

  • More than 1% and 2%: Soft MoE-L/16 outperforms ViT-L/16 on ImageNet and Cifar-100 zero-shot, respectively.The evaluation uses a ViT-g text tower trained for 18B input images, approximately five epochs.
  • Soft MoEs also surpass vanilla ViT and Experts Choice when trained from scratch on LAION-400M.With this pretraining, Soft MoEs benefit from data augmentation, unlike the compared ViT and Experts Choice models.

5 RELATED WORK

Soft MoE shares weighted-averaging ideas with token-mixing methods and multi-headed attention, but differs in its objective, expert structure, and computational organization.

  • Unlike token-fusion methods, Soft MoE does not primarily reduce sequence length and restores the original length after expert processing.
  • Multi-headed attention can be viewed as using h linear experts, whereas Soft MoE uses nonlinear experts that process d-dimensional vectors.
  • Unlike methods that average expert parameters, Soft MoE routes examples sparsely and can exploit vectorized operations more broadly.

6 CURRENT LIMITATIONS

Soft MoE has important scope and memory limitations: autoregressive decoding remains unresolved, while many experts can substantially increase model memory requirements.

  • Autoregressive decoding remains future work because learned token merging must preserve causality without biasing experts toward token indices.
  • One slot per expert tends to be optimal, but scaling to many experts can make model memory requirements large despite dense-backbone-like computational cost.The authors hypothesize that multiple slots using one expert align and provide limited informational gains.

A SOFT VS. UNIFORM VS. IDENTITY DISPATCH AND COMBINE WEIGHTS

The routing ablation compares content-independent identity and uniform dispatch/combine variants with Soft MoE’s content-dependent weighting; Soft MoE performs best.

  • Fixed-routing alternatives make expert selection or convex-combination weights independent of token content, unlike Soft MoE.
  • Identity assigns sequence tokens to experts in round-robin order, equivalent to identity dispatch and combine matrices when dimensions match.
  • Uniform averages all input tokens into every slot and all output slots into every token, with experts independently initialized.
  • Uniform / Soft keeps Soft MoE combine weights, while Soft / Uniform keeps Soft MoE dispatch weights.
  • Soft MoE is far better than all fixed-routing alternatives in the S/14 experiment with 256 experts and slots.The experiment places MoE layers in the last six backbone layers and uses one slot per expert.

B TOKEN DROPPING

Hard routing in Experts Choice and Tokens Choice increasingly drops tokens as expert counts grow, while buffer slack or Batch Priority Routing can reduce dropping at added cost. Soft MoE avoids token dropping and gains little from extra slots compared with its cost.

  • Increasing experts causes progressively more token dropping in both Experts Choice and Tokens Choice.
  • With C = 1, dropping quickly increases with experts; Experts Choice drops more but remains more performant than Tokens Choice.
  • At 100–1000 experts, Experts Choice can drop over 40–50% of tokens in some layers, while Tokens Choice can completely drop up to ∼25%.
  • C = 1.125 reduces dropping by around ∼5% and slightly improves performance at large expert counts, but increases step time.
  • Batch Priority Routing uniformly reduces Tokens Choice dropping and significantly improves performance.
  • Increasing Soft MoE slots yields only modest performance gains while quickly increasing cost.
  • Soft MoE experiments place experts across layers; configurations with expert counts near the 196 input tokens, concentrated in later layers, perform best.

E.2 EMPIRICAL ANALYSIS

The analysis finds that softmax routing can collapse toward one-hot assignments as model dimension or learning rate increases without suitable normalization. The proposed normalization produces better accuracy and greater stability.

  • Standard initialization does not fix the softmax scaling issue as model dimension increases.
  • As model dimension grows, unnormalized Soft MoE dispatch and combine weights approach one-hot behavior during training.
  • Higher learning rates similarly drive unnormalized dispatch and combine maxima toward 1.0, severely hurting accuracy.
  • The proposed normalization gives better accuracy and makes training less sensitive to learning-rate configuration.

F ADDITIONAL RESULTS

Additional experiments show that Soft MoE transfers well to contrastive learning and public-data pretraining, outperforms dense and sparse baselines across reported metrics, and maintains advantages as expert counts increase. Its training and inference efficiency also compares favorably with dense Transformers.

  • F.1 CONTRASTIVE LEARNING ON LAION-400M: On LAION-400M pretraining, Soft MoE significantly outperforms vanilla ViT and Experts Choice across downstream metrics.
  • F.1 CONTRASTIVE LEARNING ON LAION-400M: Soft MoE benefits from data augmentation, unlike the vanilla ViT and Experts Choice models in this experiment.
  • F.2 ADDITIONAL TABLES AND PLOTS: BPR is especially useful for Tokens Choice with K = 1.
  • F.3 ADDITIONAL TABLES AND PLOTS: The LAION-400M experiment used a downloadable subset of 275M image-text pairs rather than the full dataset.
  • F.3 ADDITIONAL TABLES AND PLOTS: Soft MoE outperforms ViT on all metrics in long runs with equivalent model classes and similar training costs.
  • F.3 ADDITIONAL TABLES AND PLOTS: With total slots fixed, Soft MoE improves consistently as experts increase while cost remains roughly constant; too many Experts Choice or Tokens Choice experts hurt performance and raise cost.

G MODEL INSPECTION

Model inspection shows that Soft MoE uses uneven token contributions and slot-specific dispatch patterns, while later layers become nearly uniform in dispatch. Combine weights remain concentrated on selected slots, indicating unequal expert importance.

  • Tokens contributions to slots: Around 2–5% of tokens provide summed dispatch weight above 2, while 15–20% contribute at most 0.25.
  • Cumulative sum of dispatch weights: The cumulative-dispatch analysis measures how many sorted input tokens are needed to cover a given fraction of each slot’s weighted-average mass.
  • Cumulative sum of dispatch weights: In later layers, each slot is close to a uniform average of all tokens, whereas in layer 28 a few tens of tokens cover 80% of dispatch mass.
  • Cumulative sum of dispatch weights: Dispatch weights vary substantially across slots and images, except possibly in the last two layers.
  • Cumulative sum of combine weights: Combine weights remain non-uniform in the last two layers, so some slots and their experts are more important for output tokens than others.

I SLOT CORRELATION

The section examines how Soft MoE slot parameters relate to one another and finds that adding multiple slots per expert leads consecutive slots to become highly aligned. This alignment means the extra slots compute similar linear combinations and add limited useful information.

  • I SLOT CORRELATION: Soft MoE evaluates pairwise inner products between normalized slot parameter vectors across six layers and different slots-per-expert settings.Figures 28–30 compare models with 1, 4, and 16 slots per expert.
  • I SLOT CORRELATION: With one slot per expert, slots from different experts show no clear relationship.This behavior is shown for Soft MoE S/16 with one slot per expert.
  • I SLOT CORRELATION: With four or sixteen slots per expert, consecutive slots belonging to the same expert are extremely aligned.The alignment indicates that these slots learn similar parameter directions.
  • I SLOT CORRELATION: Adding more slots per expert provides limited useful information because aligned slots compute somewhat similar linear combinations.The authors conclude that increasing slots per expert does not work very well under this learned alignment pattern.
Loading 2308.00951v2…