We gratefully acknowledge support from
the Simons Foundation and member institutions.
Full-text links:

Download:

Current browse context:

cs.LG

Change to browse by:

References & Citations

DBLP - CS Bibliography

Bookmark

(what is this?)
CiteULike logo BibSonomy logo Mendeley logo del.icio.us logo Digg logo Reddit logo ScienceWISE logo

Computer Science > Machine Learning

Title: Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts

Abstract: 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.
Comments: I (Still) Can't Believe It's Not Better Workshop at NeurIPS 2021
Subjects: Machine Learning (cs.LG); Machine Learning (stat.ML)
Cite as: arXiv:2109.11817 [cs.LG]
  (or arXiv:2109.11817v2 [cs.LG] for this version)

Submission history

From: Wouter Kool [view email]
[v1] Fri, 24 Sep 2021 09:02:12 GMT (245kb,D)
[v2] Wed, 8 Dec 2021 13:40:28 GMT (157kb,D)

Link back to: arXiv, form interface, contact.