Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts
This work addresses training efficiency for large-scale mixture of experts models, but it is incremental as it builds on existing assignment methods with a focus on unbiased estimation.
The paper tackles the problem of biased gradient estimation in mixture of experts models by proposing two unbiased estimators based on stochastic assignment procedures, finding that the 'skip'-estimator is more effective and both are more robust than biased alternatives in a toy experiment.
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.