PaperlayerView exact revision on arXiv

Source-linked AI summary

DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning

Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, Ed H. Chi

arXiv:2106.03760v3cs.LGmath.OCstat.ML

TL;DR

Sparse MoE gates such as Top-k are useful but nonsmooth, creating challenges for gradient-based training. The paper introduces DSelect-k, a continuously differentiable sparse gate based on binary encoding and explicit control of selected experts. Experiments report improvements in prediction and expert selection over state-of-the-art MoE gates and multi-task learning baselines.

  • Problem

    Existing sparse MoE gates are not continuous, which can create convergence and statistical performance issues during gradient-based training.

  • Method

    DSelect-k uses smoothed binary encoding to create a continuously differentiable sparse gate that selects at most k of n experts and can be trained with first-order methods.

  • Results

    DSelect-k achieves significant improvements in prediction and expert selection compared with state-of-the-art MoE gates and multi-task learning baselines.

  • Takeaways & Limitations

    DSelect-k offers explicit sparsity control while supporting differentiable training for MoE models evaluated on synthetic and real multi-task datasets.

  • Takeaways & Limitations

    The method assumes n = 2^m experts in its basic formulation, with extensions to arbitrary n discussed separately.

Abstract

from arXiv · show

The Mixture-of-Experts (MoE) architecture is showing promising results in improving parameter sharing in multi-task learning (MTL) and in scaling high-capacity neural networks. State-of-the-art MoE models use a trainable sparse gate to select a subset of the experts for each input example. While conceptually appealing, existing sparse gates, such as Top-k, are not smooth. The lack of smoothness can lead to convergence and statistical performance issues when training with gradient-based methods. In this paper, we develop DSelect-k: a continuously differentiable and sparse gate for MoE, based on a novel binary encoding formulation. The gate can be trained using first-order methods, such as stochastic gradient descent, and offers explicit control over the number of experts to select. We demonstrate the effectiveness of DSelect-k on both synthetic and real MTL datasets with up to $128$ tasks. Our experiments indicate that DSelect-k can achieve statistically significant improvements in prediction and expert selection over popular MoE gates. Notably, on a real-world, large-scale recommender system, DSelect-k achieves over $22\%$ improvement in predictive performance compared to Top-k. We provide an open-source implementation of DSelect-k.

1 Introduction

DSelect-k addresses the limitations of nonsmooth sparse MoE gates with a continuously differentiable, explicitly sparse alternative that supports first-order training. The paper studies its use in multi-task learning, where task-specific gates control parameter sharing.

  • MoE Motivation: MoE models combine trainable expert networks through per-example weights, enabling expert specialization and potential gains in prediction and interpretability.MoE layers support efficient computation in high-capacity networks and parameter sharing in multi-task learning.
  • Problem: Existing sparse gates select only a few experts, but their lack of continuity can create convergence and statistical performance issues during gradient-based training.Top-k can be trained with SGD for end-to-end optimization, yet remains discontinuous.
  • DSelect-k: DSelect-k selects at most k of n experts using a binary encoding reformulation that implicitly enforces sparsity without retaining the original cardinality constraint.The reformulated optimization problem is unconstrained and equivalent to the cardinality-constrained formulation.
  • Efficiency: DSelect-k uses logarithmically many parameters in the number of experts, unlike existing gates such as Top-k, and computes outputs through a closed-form expression.The paper presents these properties as advantages in compactness and computational efficiency.
  • Multi-Task Learning: In multi-gate MoE models, separate task gates adaptively control parameter sharing, and the paper evaluates DSelect-k on synthetic and real multi-task datasets.The authors study both per-example and static gating because either mechanism can outperform the other in certain settings.
  • DSelect-k: DSelect-k smooths the binary encoding variables so the gate can be optimized with first-order methods such as stochastic gradient descent.The paper reports that careful initialization and regularization support optimization of the smoothed formulation.

2 Gating in the Mixture of Experts

This section contrasts dense softmax and discontinuous Top-k gates with DSelect-k's sparse, differentiable selection. The comparison emphasizes training stability, expert recovery, inference efficiency, and the trade-off in conditional training.

  • MoE Architecture: A Mixture of Experts contains neural-network experts and a gate that assigns them nonnegative weights summing to one, producing a weighted expert combination.The gate maps inputs to probability vectors, and the MoE output is formed from the corresponding expert outputs.
  • Popular Gates: Softmax gates assign nonzero probabilities to all experts, whereas Top-k gates retain exactly k nonzero entries and can support conditional training.Top-k can compute only the selected experts' loss gradients during backpropagation, enabling computational savings with careful implementation.
  • Top-k Limitations: Top-k's discontinuity means gradients do not exist at certain inputs and can produce oscillatory expert selection during SGD training.Small input changes can cause jumps in the gate output.
  • DSelect-k Comparison: DSelect-k is continuously differentiable, can select k of n experts, and provides more stable expert selection during training than Top-k.The synthetic-data figure reports that DSelect-k recovers the true experts used by the data-generating model, whereas Top-k does not.
  • DSelect-k Comparison: DSelect-k evaluates only a subset of experts during inference, but conditional training is unavailable initially and becomes possible only after convergence to a small subset.The paper characterizes this partial support for conditional training as a trade-off against Top-k's full support.

3 Differentiable and Sparse Gating

DSelect-k formulates sparse expert selection with an explicit at-most-k constraint, then replaces the constrained problem with a binary-encoded and smoothed formulation suitable for first-order optimization. The construction uses k single-expert selectors and supports both static and per-example gating, with regularization or annealing used to enforce binary selections and cardinality constraints.

  • Gating settings: DSelect-k develops static and per-example gates that select at most k of n experts, with k specified by the user.The section first presents static gating and then generalizes the construction to per-example gating.
  • Cardinality-constrained formulation: The static-gating objective represents the gate as a probability vector w with at most k nonzero expert weights.The L0 constraint counts nonzero entries, while simplex constraints require nonnegative weights summing to one.
  • Binary encoding: Binary variables encode expert indices: each selector entry equals one exactly when its binary vector represents the corresponding expert index, yielding a one-hot vector.The general construction uses m = log2 n binary variables for n experts in the power-of-two case.
  • Binary encoding: The constrained optimization problem is transformed into an equivalent unconstrained problem using k binary-encoded single-expert selectors combined through a softmax-weighted convex combination.Each selector returns a one-hot expert choice, so their convex combination remains a probability vector with at most k nonzero entries.
  • Smooth gating: The binary formulation is smoothed with a continuously differentiable smooth-step function, producing a continuous relaxation whose objective can be optimized with first-order methods.The smooth-step function can reach 0 and 1 exactly for sufficiently large inputs, while the relaxed objective is differentiable when experts are differentiable.
  • Constraint enforcement: Entropy regularization encourages selectors to converge to binary encodings, while annealing the smooth-step scaling parameter toward zero can enforce the cardinality constraint for every example.Without strict per-example enforcement, the method tunes regularization so the average number of selected experts is at most k.

4 Experiments

The experiments evaluate DSelect-k against sparse and dense MoE gates across synthetic, benchmark, and large-scale recommendation settings. DSelect-k generally improves prediction and expert selection, with performance depending on the gating mechanism and dataset.

  • Experimental design: DSelect-k and Top-k are evaluated with static and per-example gates, alongside Gumbel-softmax, softmax-MoE, and shared-bottom baselines.The study covers MovieLens, Multi-MNIST, Multi-Fashion, a large-scale recommender system, and synthetic data with up to 128 tasks.
  • MovieLens: MovieLens results show that a DSelect-k gate outperforms competing methods in both test loss and average number of selected experts for every loss weighting.DSelect-k outperforms both Top-k and Gumbel Softmax for two of three α values in static gating, while sparse gates uniformly outperform softmax MoE.
  • Multi-MNIST and Multi-Fashion: On Multi-MNIST, static DSelect-k improves Task 2 accuracy by over 1% versus static Top-k and approaches softmax-MoE accuracy using 1.7 rather than 8 experts.Static DSelect-k also performs best on Multi-Fashion, while its per-example variant does not improve over static gating on Multi-MNIST.
  • Large-scale recommender system: The recommender system contains hundreds of millions of items and billions of users, with six classification and two regression ranking tasks.Its multi-gate MoE uses separate static gates for each task and eight dense-layer experts.
  • Large-scale recommender system: DSelect-k outperforms Top-k on all eight recommender-system tasks, with the largest improvements on satisfaction tasks.The DSelect-k gates share at least one expert across all engagement tasks, and two satisfaction tasks share an expert.

5 Conclusion

The paper concludes that DSelect-k is a continuously differentiable sparse MoE gate that supports first-order training and direct sparsity control. Experiments across synthetic and real MTL datasets indicate improvements in prediction and expert selection over state-of-the-art gates and baselines.

  • Conclusion: DSelect-k is a continuously differentiable sparse MoE gate that can be trained using first-order methods.The gate selects at most k of n experts for a user-specified k.
  • Conclusion: A binary encoding scheme enables selecting k experts without imposing explicit optimization constraints.The paper states that the resulting unconstrained reformulation is equivalent to the cardinality-constrained problem.
  • Conclusion: Experiments on synthetic and real datasets indicate significant improvements in prediction and expert selection over state-of-the-art MoE gates and MTL baselines.The synthetic experiments include settings with up to 128 tasks.
  • Societal impact: The paper states that DSelect-k can improve MoE interpretability and efficiency in underlying applications.It reports no direct negative societal impacts from the proposal.

B.1.1 Proof Proposition 1

The proof establishes equivalence between the cardinality-constrained formulation and the binary-encoded unconstrained formulation. It constructs binary codes and continuous parameters that reproduce any feasible sparse expert-weight vector.

  • Equivalence proof: The proof verifies both directions of equivalence while preserving feasibility and objective value.One direction maps an optimal reformulated solution through q(α, Z), while the other constructs α and Z from an optimal sparse weight vector.
  • Binary construction: For each nonzero expert weight, the proof assigns z(i) the binary representation of the expert index minus one.If fewer than k experts are nonzero, remaining codes duplicate the last selected expert’s code.
  • Parameter construction: The proof treats separately the cases where the optimal vector has exactly k nonzero entries or fewer than k.In each case, α is chosen so q(α, Z) exactly reproduces the optimal sparse vector.

B.1.2 Proof of Proposition 2

The proof establishes that the recursive selector r(S(v)) lies in the probability simplex, and explains how DSelect-k handles expert counts that are not powers of two.

  • Proof of Proposition 2: Induction shows that r(S(z)) has nonnegative entries whose sum equals 1, so it belongs to the probability simplex.The base case uses one smooth binary selector, and the inductive step preserves nonnegativity and total mass.
  • Proof of Proposition 2: The recursive construction expands a valid selector over 2^t experts into one over 2^(t+1) experts while preserving simplex membership.The new entries multiply existing selector values by S(v_t+1) or 1 − S(v_t+1).
  • Non-power-of-two expert counts: When n is not a power of two, DSelect-k embeds the problem in the next larger power-of-two size and penalizes probability assigned to nonexistent experts.The penalty strength is controlled by the nonnegative parameter ξ, and sufficiently large ξ makes the probability on valid experts converge to 1 in the experiments.

C.1 Prediction and Expert Selection Performance on Synthetic Data

The synthetic experiment evaluates prediction and expert selection as task count grows, comparing DSelect-k, Top-k, and softmax-based ablations. DSelect-k increasingly outperforms Top-k as the number of tasks rises.

  • Experiment design: The experiment compares DSelect-k, Top-k, and two softmax-based ablation gates across 16, 32, 64, and 128-task prediction problems.Each task-specific gate is connected to all experts and selects a convex combination of 4 experts.
  • Synthetic data: The synthetic data contain 128 regression tasks arranged into four mutually exclusive groups, with stronger relatedness within groups than across groups.Each group has 16 tasks generated from a group-specific MoE, and task-weight correlations within groups are set to 0.8.
  • Evaluation: The evaluation reports test MSE and average Jaccard indices for related and unrelated task pairs over 100 random initializations.Higher Jaccard is preferred for related tasks, while lower Jaccard is preferred for unrelated tasks.
  • Results: As task count increases, Jaccard indices decrease because using more experts gives task gates greater flexibility to choose mutually exclusive subsets.This trend is reported for both related and unrelated task comparisons.
  • Results: At 128 tasks, DSelect-k achieves over 40% improvement in MSE and 76% improvement in related-task Jaccard index compared to Top-k.The reported differences become more pronounced as the number of tasks increases, and DSelect-k significantly outperforms Top-k across the considered measures.
  • Ablation results: Replacing binary encoding with softmax-based selectors reduces selection effectiveness, while the entropy ablation often selects approximately 2 experts despite k = 4.DSelect-k and Top-k select 4 experts in the reported repetitions.

C.2.2 Synthetic Data

The synthetic gate-convergence analysis measures how quickly DSelect-k reaches binary selectors and compares its training computation with Top-k. Convergence speed varies substantially by dataset, while FLOPS can be lower on some datasets.

  • Synthetic-data setup: The synthetic-data setup trains a 16-expert MoE with four copied data-generating experts and freezes all experts while learning only gate parameters.The expected outcome is recovery of the four experts used to generate the data.
  • Metrics: Table C.4 reports the percentage of training steps needed for DSelect-k’s S(Z) to converge to a binary matrix and its training FLOPS relative to Top-k.The table also includes the MovieLens loss-weight parameter α.
  • Convergence results: Convergence time varies by dataset, ranging from 0.04% of training steps on MovieLens to 80% on Multi-Fashion MNIST.These values are reported for convergence of S(Z) to a binary matrix.
  • FLOPS results: On MovieLens with α = 0.9 and Multi-MNIST, DSelect-k requires fewer training FLOPS than Top-k.The passage identifies these datasets as cases where DSelect-k is effective at conditional training.
  • MovieLens results: Table C.5 reports MovieLens test loss for task 2 and accuracy for task 1, with standard errors shown beside each mean.The test loss is multiplied by 10^4, and α determines the weight of task 1’s loss.

D Experimental Details

The experiments compare MoE gates under specified architectures and tuning procedures, including a Gumbel-softmax gate whose penalty does not directly control the number of selected experts.

  • Gumbel-softmax gate: The Gumbel-softmax gate uses binary variables to determine which experts are selected and a learnable vector α for expert weights.Its selection distribution is learned with the Gumbel-softmax trick.
  • Gumbel-softmax gate: The Gumbel-softmax sparsity penalty uses λ to control the number of nonzeros but cannot directly impose a specified number as DSelect-k or Top-k.The penalty is based on the Bernoulli parameters ψ_i of the selection variables.
  • Architecture: The experiments use a multi-gate MoE architecture with separate gates for tasks and eight experts.The expert and task-subnetwork architectures vary across experimental settings.
  • Hyperparameters and tuning: Models are tuned with random grid search averaging five trials per grid point, using Adagrad or Adam with batch sizes of 128 or 256.The searched hyperparameters include learning rate, epochs, k, smooth-step γ, architecture, and gate-specific λ values.
  • Hyperparameters and tuning: In one setting, DSelect-k and Top-k fix k to 2, tune learning rate and architecture, and stop when validation loss shows no significant improvement.This setting uses eight experts composed of ReLU-activated dense layers with 256 units each.

D.4 Synthetic Data

The synthetic-data experiments tune models with random grid search and Adam, while varying learning rates, training duration, smooth-step parameters, and regularization settings.

  • Hyperparameters and tuning: Synthetic-data models are tuned with random grid search averaging five trials per grid point.The optimizer is Adam with batch size 256.
  • Hyperparameters and tuning: The synthetic-data search varies learning rate over 0.001, 0.01, and 0.1, and epochs over 25, 50, 75, and 100.These ranges define the principal optimization and training-duration settings reported in the passage.
  • Hyperparameters and tuning: The search varies smooth-step γ over 5, 10, and 15 and DSelect-k λ over 0.001, 0.005, 0.01, and 0.1.An entropy-ablation λ is also searched over values from 10^-6 through 100.
Loading 2106.03760v3…