Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts
Training large-scale mixture of experts models efficiently on modern hardware requires assigning datapoints in a batch to different experts, each with a limited capacity. Recently proposed assignment procedures lack a probabilistic interpretation and use biased estimators for training. As an alternative, we propose two unbiased estimators based on principled stochastic assignment procedures: one that skips datapoints which exceed expert capacity, and one that samples perfectly balanced assignments using an extension of the Gumbel-Matching distribution [29]. Both estimators are unbiased, as they correct for the used sampling procedure. On a toy experiment, we find the `skip'-estimator is more effective than the balanced sampling one, and both are more robust in solving the task than biased alternatives.
Code (0)
등록된 구현이 없습니다.
Tasks
Mixture-of-ExpertsSimilar Papers 제목 키워드 기반
A 1/R Law for Kurtosis Contrast in Balanced Mixtures
Kurtosis-based Independent Component Analysis (ICA) weakens in wide, balanced mixtures. We prove a sharp redundancy law: for a standardized projection with effective width $R_{\mathrm{eff}}$ (participation ratio), the po…
Balanced Meta-Softmax for Long-Tailed Visual Recognition
Deep classifiers have achieved great success in visual recognition. However, real-world data is long-tailed by nature, leading to the mismatch between training and testing distributions. In this paper, we show that the S…
General ClassificationInstance SegmentationLong-tail LearningSemantic SegmentationSticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference
We propose a simple and general variant of the standard reparameterized gradient estimator for the variational evidence lower bound. Specifically, we remove a part of the total derivative with respect to the variational …
Variational InferenceUnbiased Estimation of the Gradient of the Log-Likelihood for a Class of Continuous-Time State-Space Models
In this paper, we consider static parameter estimation for a class of continuous-time state-space models. Our goal is to obtain an unbiased estimate of the gradient of the log-likelihood (score function), which is an est…
parameter estimationState Space ModelsA Stochastic Gradient Method with Biased Estimation for Faster Nonconvex Optimization
A number of optimization approaches have been proposed for optimizing nonconvex objectives (e.g. deep learning models), such as batch gradient descent, stochastic gradient descent and stochastic variance reduced gradient…