Original: Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts
Revised: Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts

Original: Anonymous Author(s)

Original: Afﬁliation email

Original: Abstract
Revised: Abstract

Original: Training large-scale mixture of experts models efﬁciently on modern hardwarerequires assigning datapoints in a batch to different experts, each with a limited capacity.
Revised: Training large-scale mixture of experts models efﬁciently on modern hardware requires assigning datapoints in a batch to different experts, each with a limited capacity.

Original: Recently proposed assignment procedures lack a probabilistic interpretationand use biased estimators for training.
Revised: Recently proposed assignment procedures lack a probabilistic interpretation and use biased estimators for training.

Original: As an alternative, we propose two unbiasedestimators based on principled stochastic assignment procedures: one that skipsdatapoints which exceed expert capacity, and one that samples perfectly balancedassignments using an extension of the Gumbel-Matching distribution [28].
Revised: 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].

Original: Bothestimators are unbiased, as they correct for the used sampling procedure.
Revised: Both estimators are unbiased, as they correct for the used sampling procedure.

Original: On atoy experiment, we ﬁnd the ‘skip’-estimator is more effective than the balancedsampling one, and both are more robust in solving the task than biased alternatives.
Revised: On a toy experiment, we ﬁnd the ‘skip’-estimator is more effective than the balanced sampling one, and both are more robust in solving the task than biased alternatives

Original:  Introduction
Revised:  Introduction

Original: Recently, large-scale sparse mixture of experts (MoE)models have been successful in scalingup model parameter count withoutsubstantially increasing the amount of computation required[37, 24, 8, 25, 34, 44].
Revised: A mixture of experts (MoE) model can be used to implement conditional computation by processing different datapoints by different expert modules. This enables increasing the MoE models’ representational capacity by adding more experts, without increasing the amount of computation per datapoint, which remains the same as for a single expert. Recently, this idea has been combined with deep neural networks, where each layer can be a separate MoE model, resulting in large-scale MoE’s that yield state-of-the-art performance in various tasks [38, 25, 8, 26, 35, 45].

Original: With MoE’s, a routing networkis trained to assign datapoints to differentexperts.
Revised: Training an MoE model involves training a routing network to assign datapoints to experts, and training individual experts to perform well on the datapoints assigned to them.

Original: In practice, for computational efﬁciency, we train on minibatches of datapoints, and eachexpert has a limited capacity,so we either have to skip datapoints exceeding the capacity of theexperts they are assigned to, or balance the assignments of datapoints to different experts [8, 25].
Revised: In practice, for computational efﬁciency, we train on minibatches of datapoints, and as each expert has a limited capacity, we either have to skip datapoints exceeding the capacity of the experts they are assigned to, or balance the assignments of datapoints to different experts [8, 26].

Original: The effect of skipping datapoints or balancing their assignments is typically not accounted for whentraining the routing network.
Revised: The effect of skipping datapoints or balancing their assignments is typically not accounted for when training the routing network.

Original: Recent work has also challenged this approach by showing that in somecases better performance can be achieved using a ﬁxed hashing-based routing strategy [35], whichsuggests that there is room for improvement in training the routing network.
Revised: Recent work has also challenged this approach by showing that in some cases better performance can be achieved using a ﬁxed hashing-based routing strategy [36], which suggests that there is room for improvement in training the routing network.

Original: As a step towards thisgoal, we present  principled methods for optimizing MoE’s under limited expert capacity.
Revised: As a step towards this goal, we present two principled methods for optimizing MoE’s under limited expert capacity.

Original: Speciﬁcally,we propose two sampling procedures and corresponding unbiased estimators in this paper:  one basedon skipping datapoints that exceed expert capacity, and  one based on balanced  assignment using anextension of the Gumbel-Matching distribution [28].
Revised: Speciﬁcally, we propose two sampling procedures and corresponding unbiased estimators in this paper: a simple one based on skipping datapoints that exceed expert capacity, and a more advanced one based on balanced datapoint assignment using an extension of the Gumbel-Matching distribution [29].

Revised: Whereas these sampling procedures ensure that each sample respects the expert capacity, we also propose to use the Sinkhorn algorithm to balance the assignment in expectation before sampling, and we connect this procedure to the Gumbel-Matching distribution, which has such Sinkhorn balancing built-in.

Revised: In this paper, we consider a single-layer MoE model, but this can be easily generalized.

Original: Formally, the problem we consider is to predict a label y for a datapoint x using an MoE modelconsisting of k individual experts p θ ( y | x, z ) indexed by z and selected using the routing networkp θ ( z | x ) .
Revised: Formally, the problem we consider is to predict a label y for a datapoint x using an MoE model consisting of k individual experts p θ ( y | x, z ) indexed by z and selected using the routing network p θ ( z | x ) .

Original: At test time, we select the most probable expert z ∗ =
Revised: At test time, we select the most probable expert z ∗ =

Original: arg max z p θ ( z | x ) , while for trainingwe optimize a smoothed objective obtained by taking the expectation over the routing decisions.
Revised: arg max z p θ ( z | x ) , while for training we optimize a smoothed objective obtained by taking the expectation over the routing decisions.

Original: Theresulting objective, ELBO, is a variational lower bound on the marginal log-likelihood: 
Revised: The resulting objective, ELBO, is a variational lower bound on the marginal log-likelihood: log p θ ( y | x, z ∗ )

Original: [Equation]
Revised: ≈ E

Revised: z ∼ p θ ( z | x ) [log p θ ( y | x, z )] (cid:124) (cid:123)(cid:122) (cid:125) ELBO ≤ log E z ∼ p θ ( z | x ) [ p θ ( y | x, z )]

Original:  We optimize this objective using minibatches of n datapoints, while respecting the expert capacityby assigning at most c = nk datapoints to each expert (for simplicity, we assume no slack capacity).
Revised: = log p θ ( y | x ) . We optimize this objective using minibatches of n datapoints, while respecting the expert capacity by assigning at most c = nk datapoints to each expert (for simplicity, we assume no slack capacity).

Original: We test our estimators on a toy experiment and ﬁnd that the  ‘skip’-estimator is more effective thanbalanced sampling, while both  are more robust than  biased alternatives using differentiable gating.
Revised: We evaluate our estimators on a toy experiment, where we ﬁnd that the simple ‘skip’-estimator is more effective than the one based on balanced sampling using the Gumbel-Matching distribution. This is a surprising result, as we designed balanced sampling (with importance weights) as a better alternative to wastefully skipping data, but it seems that the beneﬁt of using all data is outweighed by the added variance due to the importance weights. We do however ﬁnd that both estimators, which are based on REINFORCE [11, 44], are more robust than the biased alternatives using differentiable gating.

Original:  Unbiased Estimation using Balanced Assignment
Revised:  Unbiased Estimation using Balanced Assignment

Original: For simplicity, we consider the problem of optimizing a general function f ( x, z ) using gradientdescent using minibatches of datapoints x 
Revised: For simplicity, we consider the problem of optimizing a general function f ( x, z ) using gradient descent using minibatches of datapoints x = ( x 1 , ..., x n ) and

Original: = ( x 1 , ..., x n ) and expert assignments z = ( z 1 , ..., z n ) .
Revised: expert assignments z = ( z 1 , ..., z n ) .

Original: To simplify notation, we omit the dependence of f ( x, z ) on y and parameters θ (this can be addedeasily).
Revised: To simplify notation, we omit the dependence of f ( x, z ) on y and parameters θ (this can be added easily).

Original: By combining importance sampling with REINFORCE [43], we can sample from any jointproposal distribution q ( z | x ) with marginals q ( z i | x ) to estimate the gradient for a minibatch x :∇ E z ∼ p θ ( z | x ) (cid:34) [Equation]
Revised: By combining importance sampling with REINFORCE [11, 44], we can sample from any joint proposal distribution q ( z | x ) with marginals q ( z i | x ) to estimate the gradient for a minibatch x : ∇ E z ∼ p θ ( z | x ) (cid:34)n (cid:88) i f

Original: n (cid:88) [Equation]
Revised: ( x i , z i ) (cid:35)

Original: [Equation]
Revised: [Equation]

Original: Here b is a baseline which can reduce the estimator variance (see Appendix B).
Revised: Here b is a baseline which can reduce the estimator variance (see Appendix B).

Original: Taking q ( z | x ) =41 (cid:81) i p θ
Revised: Taking q ( z | x ) = (cid:81) i p θ

Original: ( z i | x i ) recovers the standard ‘on-policy’ REINFORCE 
Revised: ( z i | x i ) recovers the standard ‘on-policy’ REINFORCE estimator.

Original: estimator.2.1 Skipping Datapoints as the Simple Solution
Revised: 2.1 Skipping Datapoints as the Simple Solution

Original: Our ‘skip’-estimator respects expert capacity by sampling expert assignments independently andrandomly subsampling the datapoints assigned to experts for which capacity is exceeded.
Revised: Our ‘skip’-estimator respects expert capacity by sampling expert assignments independently and randomly subsampling the datapoints assigned to experts for which capacity is exceeded.

Original: If weassume a random order of the datapoints (or shufﬂe them ﬁrst), simply skipping the last assignments per expert is equivalent to uniform subsampling.
Revised: If we assume a random order of the datapoints (or shufﬂe them ﬁrst), simply skipping the last assignments per expert is equivalent to uniform subsampling.

Original: Let z be the vector of expert assignments,sampled independently from a proposal distribution q ( z | x ) =
Revised: Let z be the vector of expert assignments, sampled independently from a proposal distribution q ( z | x ) =

Original: (cid:81)
Revised: (cid:81)

Original: i q ( z i | x i ) .
Revised: i q ( z i | x i ) .

Original: Let n j = (cid:80) i 1 { z i = j }be the number of datapoints assigned to expert i (before subsampling) and
Revised: Let n j = (cid:80) i 1 { z i = j } be the number of datapoints assigned to expert j (before subsampling) and

Original: c = nk be the expertcapacity.
Revised: c = nk be the expert capacity.

Original: Let δ =
Revised: Let δ =

Original: ( δ 1 , ..., δ n ) with δ i ∈ { 0 , 1 } represent which datapoints are kept after per-expertsubsampling.
Revised: ( δ 1 , ..., δ n ) with δ i ∈ { 0 , 1 } represent which datapoints are kept after per-expert subsampling.

Original: Correcting for the fact that the probability of datapoint i being kept after subsamplingis min { n z i , c } /n
Revised: Correcting for the fact that the probability of datapoint i being kept after subsampling is min { n z i , c } /n

Original: z i , we obtain (see Appendix C):
Revised: z i , we obtain (see Appendix C):

Original: [Equation]
Revised: [Equation]

Original:  Here we have omitted the baseline b and used ∇ p θ ( z i | x i )
Revised: [Equation] Here we have omitted the baseline b and used ∇ p θ ( z i | x i )

Original: =
Revised: =

Original: p θ ( z i | x i

Original: [Equation]

Original:  ∇ log p θ [Equation] z i | x i ) for brevity.
Revised: p θ ( z i | x i ) ∇ log p θ ( z i | x i ) for brevity.

Original: Note that, while we skip datapoints in the gradient estimate (3), we can still propagate them tosubsequent layers, e.g. by using skip connections or no-token-left-behind routing [8].
Revised: Note that, while we skip datapoints in the gradient estimate (3), we can still propagate them to subsequent layers, e.g. by using skip connections or no-token-left-behind routing [8].

Original: In paricular, wecan also apply (3) to multilayer MoE’s, where we can skip different datapoints in different layers.
Revised: In paricular, we can also apply (3) to multilayer MoE’s, where we can skip different datapoints in different layers.

Original: 2.2 Balanced Assignment using Gumbel-Matching
Revised: 2.2 Balanced Assignment using Gumbel-Matching

Original: As an alternative to skipping, we use the n × k Gumbel-Matching distribution, a strict generalizationof the ( n × n ) Gumbel-Matching distribution [28], to sample perfectly balanced assignments.
Revised: As an alternative to skipping, we use the n × k Gumbel-Matching distribution, a strict generalization of the ( n × n ) Gumbel-Matching distribution [29], to sample perfectly balanced assignments.

Original: Wederive it by using the Gumbel-Max trick [27, 13] to view sampling of independent expert assignmentsas an optimization problem, and adding constraints to this problem to respect the expert capacity.
Revised: We derive it by using the Gumbel-Max trick [28, 14] to view sampling of independent expert assignments as an optimization problem, and adding constraints to this problem to respect the expert capacity.

Original: Letz ij = 1 { z i = j } be the one-hot representation of z i , a ij the unnormalized log-probability (logit) ofassigning datapoint i to expert j , and g ij ∼ Gumbel (0) i.i.d.
Revised: Let z ij = 1 { z i = j } be the one-hot representation of z i , a ij the unnormalized log-probability (logit) of assigning datapoint i to expert j , and g ij ∼ Gumbel (0) i.i.d.

Original: standard Gumbel variables.
Revised: standard Gumbel variables.

Original: Then z hasthe Gumbel-Matching distribution if it is the solution to the following optimization problem:
Revised: Then z has the Gumbel-Matching distribution if it is the solution to the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where τ is a temperature parameter and c
Revised: where τ is a temperature parameter and c

Original: = nk is the expert capacity.
Revised: = nk is the expert capacity.

Original: If we remove the balancing/capacity constraint (cid:80) i z ij ≤ c , the solution decomposes over i and is given by z i =arg max j ( a ij /τ + g ij ) which is equivalent to z i ∼ Categorical (exp( a ij /τ ) /
Revised: If we remove the balancing/capacity constraint (cid:80) i z ij ≤ c , the solution decomposes over i and is given by z i = arg max j ( a ij /τ + g ij ) which is equivalent to z i ∼ Categorical (exp( a ij /τ ) /

Original: (cid:80) j exp( a ij /τ ))
Revised: (cid:80) j exp( a ij /τ ))

Original: asa result of the Gumbel-Max trick.
Revised: as a result of the Gumbel-Max trick.

Original: Thus, adding the constraint can be seen as a natural way ofenforcing a balanced assignment to the otherwise independent sampling procedure.
Revised: Thus, adding the constraint can be seen as a natural way of enforcing a balanced assignment to the otherwise independent sampling procedure.

Original: Generalizingthe result from [28], the n × k Gumbel-Matching approximates the Gibbs distribution over n × kassignments z with potentials given by (cid:80) ij z ij a ij and temperature τ (see Appendix D.1):
Revised: Generalizing the result from [29], the n × k Gumbel-Matching approximates the Gibbs distribution over n × k assignments z with potentials given by (cid:80) ij z ij a ij and temperature τ (see Appendix D.1):

Original: [Equation]
Revised: [Equation]

Original: For τ → 0 we obtain the deterministic assignment used in BASE layers [25].
Revised: For τ → 0 we obtain the deterministic assignment used in BASE layers [26].

Original: We propose to solvethe n × k assignment problem using a special cycle cancelling algorithm [19] (see Appendix D.2),which for k (cid:28) n is more efﬁcient than O ( n 3 ) assignment using the Hungarian Algorithm [22].
Revised: We propose to solve the n × k assignment problem using a special cycle cancelling algorithm [20] (see Appendix D.2), which for k (cid:28) n is more efﬁcient than O ( n 3 ) assignment using the Hungarian Algorithm [23].

Original: The marginals q θ ( z i | x ) of the Gumbel-Matching distribution are intractable [28], but we can computethe conditionals q θ ( z i | x , G − i ) , conditioned on the noise G − i=
Revised: The marginals q θ ( z i | x ) of the Gumbel-Matching distribution are intractable [29], but we can compute the conditionals q θ ( z i | x , G − i ) , conditioned on the noise G − i

Original:  ( g 1 , ..., g i − 1 , g i +1 , ..., g n ) of theother datapoints.
Revised: = ( g 1 , ..., g i − 1 , g i +1 , ..., g n ) of the other datapoints.

Original: These can be computed efﬁciently (see Appendix D.4) and used as a stochasticapproximations to the marginals, which still yield unbiased gradients when used in (2).
Revised: These can be computed efﬁciently (see Appendix D.4) and used as a stochastic approximations to the marginals, which still yield unbiased gradients when used in (2).

Original: Formally,let z = GM (log p ( ·| x ) , G ) be the solution for the Gumbel-Matching problem with noise G .
Revised: Formally, let z = GM (log p ( ·| x ) , G ) be the solution for the Gumbel-Matching problem with noise G .

Original: We canthen use the following estimator (see Appendix D.5):
Revised: We can then use the following estimator (see Appendix D.5):

Original: [Equation]
Revised: [Equation]

Original: 2.3 Bonus: Balancing Expectations using the Sinkhorn Algorithm
Revised: 2.3 Bonus: Balancing Expectations using the Sinkhorn Algorithm

Original: [Figure]
Revised: [Figure]

Original: As an alternative to generating balanced samples, we can balance thedistribution in expectation by using a generalization of the Sinkhornalgorithm [38, 39, 20] to non-square matrices, which iterativelynormalizes the columns and rows of the probability matrix p ij =
Revised: As an alternative to generating balanced samples, we can balance the distribution in expectation by using a generalization of the Sinkhorn algorithm [39, 40, 21] to non-square matrices, which iteratively normalizes the columns and rows of the probability matrix p ij =

Original: p θ ( z i =
Revised: p θ ( z i =

Original: j | x i ) to sum to nk and 1 respectively, until convergence.
Revised: j | x i ) to sum to nk and 1 respectively, until convergence.

Original: Wemay then use this Sinkhorn-normalized distribution as the proposalq ( z | x ) in (2) and (3), which yields more (but not perfectly) balancedsamples.
Revised: We may then use this Sinkhorn-normalized distribution as the proposal q ( z | x ) in (2) and (3), which yields more (but not perfectly) balanced samples.

Original: The Gumbel-Matching distribution is invariant to thisnormalization as it does not change the solution to (4).
Revised: The Gumbel-Matching distribution is invariant to this normalization as it does not change the solution to (4).

Original: See Figure1 for an overview.
Revised: See Figure 1 for an overview.

Original: The Sinkhorn algorithm is a direct extension of the softmax function,which solves a soft (entropy-regularized) version of the assignmentproblem [6, 29, 28].
Revised: The Sinkhorn algorithm is a direct extension of the softmax function, which solves a soft (entropy-regularized) version of the assignment problem [6, 30, 29].

Original: As a result, it can be seen as approximatingthe marginals for the Gibbs distribution (5) [10, 28] (for n = k , butthis can be generalized to k < n ).
Revised: As a result, it can be seen as approximating the marginals for the Gibbs distribution (5) [10, 29] (for n = k , but this can be generalized to k < n ).

Original: Empirically, we ﬁnd that the
Revised: Empirically, we ﬁnd that the

Original: Sinkhorn algorithm (as a ‘soft’ matching algorithm) also closelyapproximates the marginals of the Gumbel-Matching distribution(itself an approximation of (5)), at least for n ≥ 4 .
Revised: Sinkhorn algorithm (as a ‘soft’ matching algorithm) also closely approximates the marginals of the Gumbel-Matching distribution (itself an approximation of (5)), at least for n ≥ 4 .

Original: As such, wepropose to heuristically use it with (2) as an alternative to (6).
Revised: As such, we propose to heuristically use it with (2) as an alternative to (6).

Original: The Sinkhorn algorithm yields a balanced row-stochastic matrixwith probabilities/expectations ˜
Revised: The Sinkhorn algorithm yields a balanced row-stochastic matrix with probabilities/expectations ˜

Original: p ij , that can be seen as a balancedapproximation of the unbalanced probabilities p ij .
Revised: p ij , that can be seen as a balanced approximation of the unbalanced probabilities p ij .

Original: Given such a balanced matrix, there exist manydistributions over balanced assignments that have the (per datapoint) marginals equal to ˜ p ij , whichfollows from generalizing the Birkhoff theorem [4, 42].
Revised: Given such a balanced matrix, there exist many distributions over balanced assignments that have the (per datapoint) marginals equal to ˜ p ij , which follows from generalizing the Birkhoff theorem [4, 43].

Original: When using such a distribution with stochasticgradient descent, we would like to minimize the dependence between samples for different datapoints,to reduce the variance of the gradient estimates.
Revised: When using such a distribution with stochastic gradient descent, we would like to minimize the dependence between samples for different datapoints, to reduce the variance of the gradient estimates.

Original: This can be achieved by maximizing the entropy ofthe joint distribution over expert assignments with the given marginals ˜ p ij .
Revised: This can be achieved by maximizing the entropy of the joint distribution over expert assignments with the given marginals ˜ p ij .

Original: In Appendix D.6, we showthat this maximum-entropy distribution has the form (5), again motivating the Gumbel-Matchingdistribution as a principled approximation. 
Revised: In Appendix D.6, we show that this maximum-entropy distribution has the form (5), again motivating the Gumbel-Matching distribution as a principled approximation. [Figure]

Original:  Experiment
Revised:  Experiment

Original: We evaluate the proposed estimators on a toy task of modelling a discontinuous functiony = 1 x< 0 .
Revised: We evaluate the proposed estimators on a toy task of modelling a discontinuous function y = 1 x< 0 .

Original: 8 x − 0 . 2)
Revised: 8 x − 0 . 2)

Original: + 1 x ≥ 0 .
Revised: + 1 x ≥ 0 .

Original: 5 ( − 2 .
Revised: 5 ( − 2 .

Original: 0 x
Revised: 0 x

Original: + 2 . 0) .
Revised: + 2 . 0) .

Original: We sample a dataset of 100 datapointsx ∼ Uniform ( − 1 , 1) and add noise (cid:15) ∼ Normal (0 , 0 . 1 2 ) , see Figure 2a.
Revised: We sample a dataset of 100 datapoints x ∼ Uniform ( − 1 , 1) and add noise (cid:15) ∼ Normal (0 , 0 . 1 2 ) , see Figure 2a.

Original: We model this datasetusing a mixture of two linear experts and a Bernoulli router distribution with probability given by thesigmoid of a linear function.
Revised: We model this dataset using a mixture of two linear experts and a Bernoulli router distribution with probability given by the sigmoid of a linear function.

Original: This model is able to solve the task perfectly, but the training methodsneed to take into account the fact that the solution involves an imbalanced partitioning of the datasetbetween the experts.
Revised: This model is able to solve the task perfectly, but the training methods need to take into account the fact that the solution involves an imbalanced partitioning of the dataset between the experts.

Original: We train the model to minimize the mean squared error (MSE) under samplingof the experts, which corresponds to the objective (1) where p θ ( y | x, z ) is Gaussian with a ﬁxedvariance.
Revised: We train the model to minimize the mean squared error (MSE) under sampling of the experts, which corresponds to the objective (1) where p θ ( y | x, z ) is Gaussian with a ﬁxed variance.

Original: We consider the task solved successfully if the MSE < 0 .
Revised: We consider the task solved successfully if the MSE < 0 .

Original: 02 after 10K steps of trainingusing Adam [17] with the learning rate of 0 . 1 (found to work best overall using a grid search).
Revised: 02 after 10K steps of training using Adam [18] with the learning rate of 0 . 1 (found to work best overall using a grid search).

Original: We compare the following sampling strategies combined with REINFORCE:
Revised: We compare the following sampling strategies combined with REINFORCE:

Original: Sample Skip IW is theskipping estimator with the importance-weighting correction (Equation (3)), whereas Sample Skip [Figure]
Revised: Sample Skip IW is the skipping estimator with the importance-weighting correction (Equation (3)), whereas Sample Skip is the biased alternative that simply averages the gradient over the remaining datapoints.

Original: REINFORCE [Figure]
Revised: GumbelMatching IW uses the conditional distributions for the importance weights (Equation (6)),

Original: is the biased alternative that simply averages the gradient over the remaining datapoints.
Revised: whereas

Original: GumbelMatching IW uses the stochastic conditional distributions for the importance weights (Equation(6)), whereas Gumbel-Matching SH is the (biased) version that uses the ‘Sinkhorn marginals’ (see Section 2.3) with Equation (2), and Gumbel-Matching is the biased estimator that does not correctfor balancing using importance weights.
Revised: Gumbel-Matching SH is the (biased) version that uses the ‘Sinkhorn marginals’ (see Section 2.3) with Equation (2), and Gumbel-Matching is the biased estimator that does not correct for balancing using importance weights.

Original: Lastly, to quantify the reduction in performance due tothe expert capacity constraints, we also include the results for Sample , the ideal baseline that doesnot respect expert capacity.
Revised: Lastly, to quantify the reduction in performance due to the expert capacity constraints, we also include the results for Sample , the ideal baseline that does not respect expert capacity.

Original: We run each experiment  for a range of sampling temperatures τ , takinginto account their effect on the proposal distribution in the importance weights for all estimators.
Revised: We run each experiment with 10 seeds, for a range of sampling temperatures τ , taking into account their effect on the proposal distribution in the importance weights for all estimators.

Original: Toreduce variance, we include an exponential moving average baseline [40] with decay 0.99.
Revised: To reduce variance, we include an exponential moving average baseline [41] with decay 0.99.

Original: Figure 2b plots the ﬁnal training MSE forfor differentsampling temperatures τ , where we observethat Sample Skip IW is the only estimator that matches the imbalanced  Sample estimator, bothsolving the task for τ ≥ 1 .
Revised: Figure 2b plots the ﬁnal training MSE for different temperatures τ of the sampling (proposal) distribution, where we observe that Sample Skip IW is the only estimator that matches the imbalanced (unconstrained) Sample estimator, both solving the task for τ ≥ 1 .

Revised: The skipping estimator thus provides a simple and effective way to deal with limited expert capacity, but it is important to upweight the remaining samples for experts which have skipped datapoints, as Sample Skip , which does not use this weighting, does not achieve the same performance.

Revised: Confounding our expectation, the

Original: The other estimators solve the task only for  largeτ values, which avoidlarge importance weights by smoothing the proposal distribution.
Revised: Gumbel-Matching based estimators turned out to be less effective, because of the increased variance due to the importance weights. Investigating the issue, we found that a datapoint x can have high probability p θ ( z | x ) for an expert z according to the router , but a low probability under the proposal q ( z | x ) of actually being assigned to the expert z , due to the balancing constraint. While the latter probability is small, occasionally the datapoint will get assigned to z , resulting in a large importance weight p θ ( z | x ) q ( z | x ) . This effect can be mitigated by increasing the temperature of the proposal distribution, making it more uniform and avoiding large importance weights, which explains the good results for large τ values for all estimators except Gumbel-Matching .

Original: We also experiment with Sinkhornbalancing (Section 2.3) before sampling, which only works for high temperatures (see Appendix E).
Revised: We also experiment with all estimators with Sinkhorn balancing (Section 2.3) before sampling, which only works for high temperatures (see Appendix E).

Original: 3.1 Biased training using  gating
Revised: 3.1 Biased training using differentiable gating

Original: Large scale MoE’s used in practice [37, 24, 8, 25, 34, 44] do not use REINFORCE, but insteadmultiply the output of an expert by the router probability p θ ( z | x ) , also known as  gating .
Revised: Large scale MoE’s used in practice [38, 25, 8, 26, 35, 45] do not use REINFORCE, but instead multiply the output of an expert by the router probability p θ ( z | x ) , which we refer to as differentiable gating .

Original: Thisway, the router becomes more coupled with the experts and gets a gradient signal directly fromthe objective.
Revised: This way, the router becomes more coupled with the experts and gets a gradient signal directly from the objective.

Original: Different strategies for injecting noise to encourage exploration have been proposed,e.g. perturbing router inputs or (log-)probabilities with multiplicative or additive (Gaussian) noise[37, 8].
Revised: Different strategies for injecting noise to encourage exploration have been proposed, e.g. perturbing router inputs or (log-)probabilities with multiplicative or additive (Gaussian) noise [38, 8].

Original: Empirically, we ﬁnd that we get similar results by perturbing log-probabilities with (scaled) Gumbel noise, which, since arg max j ( a ij + τ · g ij ) =
Revised: Empirically, we ﬁnd that we get similar results by perturbing log-probabilities with (scaled) Gumbel noise, which, since arg max j ( a ij + τ · g ij ) =

Original: arg max j ( a ij /τ + g ij ) , has the advantage ofbeing interpretable as sampling from a categorical distribution with the temperature τ
Revised: arg max j ( a ij /τ + g ij ) , has the advantage of being interpretable as sampling from a categorical distribution with the temperature τ

Original: (see Section2.2).
Revised: (see Section 2.2).

Original: We ﬁnd training using  gating succeeds only if we additionally include a load balancing loss[8] with a weight of 0 . 01 or 0 . 03 , and use balanced sampling with importance weights in a lowtemperature regime, as can be seen in Figure 2c.
Revised: We ﬁnd training using differentiable gating succeeds only if we additionally include a load balancing loss [8] with a weight of 0 . 01 or 0 . 03 , and use balanced sampling with importance weights in a low temperature regime, as can be seen in Figure 2c.

Original: With  gating, we may wonder if we need importanceweights at all, since existing methods do not use importance weights to correct for the samplingtemperature (noise scale) τ .
Revised: With differentiable gating, we may wonder if we need importance weights at all, since existing methods do not use importance weights to correct for the sampling temperature (noise scale) τ .

Original: In Appendix E we show however, that with  gating, we fail to train themodel when not using importance weights.
Revised: In Appendix E we show however, that with differentiable gating, we fail to train the model when not using importance weights.

Original:  Discussion
Revised:  Discussion

Original: In this paper we proposed several new estimators for training MoE models with a limited computational capacity per expert.
Revised: In this paper we proposed several new estimators for training MoE models with a limited computational capacity per expert.

Original: We expected balanced sampling with importance weighting to correct forassignment of datapoints to low-probability experts to perform at least as well as skipping, whicheffectively uses a weight of 0 for the skipped datapoints.
Revised: We expected balanced sampling with importance weighting to correct for assignment of datapoints to low-probability experts to perform at least as well as skipping, which effectively uses a weight of 0 for the skipped datapoints.

Original: We found this not to be the case in practice,with the added variance from importance weights eliminating the beneﬁt from the increased expertcapacity utilization due to balanced sampling.
Revised: We found this not to be the case in practice, with the added variance from importance weights eliminating the beneﬁt from the increased expert capacity utilization due to balanced sampling.

Original: Fortunately, the skipping estimator turned out to bea simple and effective alternative.
Revised: Fortunately, the skipping estimator turned out to be a simple and effective alternative.

Original: We hope this work will be useful for training MoE models inpractice.
Revised: We hope this work will be useful for training MoE models in practice.

Revised: Acknowledgments and Disclosure of Funding

Revised: We would like to thank Michalis Titsias, Jörg Bornschein, Matthias Bauer and Yee Whye Teh for helpful comments, discussions and support.

Original: A Related work
Revised: A Related work

Original: Mixtures of experts [7, 37] have a long history as a method for conditional computation [3, 2, 11]where different experts are used for different datapoints.
Revised: Mixtures of experts [7, 38] have a long history as a method for conditional computation [3, 2, 12] where different experts are used for different datapoints.

Original: Typically, the router or gating network thatassigns datapoints to experts is learned jointly with the experts themselves.
Revised: Typically, the router or gating network that assigns datapoints to experts is learned jointly with the experts themselves.

Original: Mixtures of experts andthe more general routing or modular networks [36, 18, 33], sometimes treat expert assignments aslatent variables and use EM or variational methods [18] for training them.
Revised: Mixtures of experts and the more general routing or modular networks [37, 19, 34], sometimes treat expert assignments as latent variables and use EM or variational methods [19] for training them.

Original: The beneﬁt of optimizinga single-sample bound (ELBO) rather than marginal likelihood, is that in additional to being moretractable, it can help to avoid stochasticity [32] in datapoint assignments to experts.
Revised: The beneﬁt of optimizing a single-sample bound (ELBO) rather than marginal likelihood, is that in additional to being more tractable, it can help to avoid stochasticity [33] in datapoint assignments to experts.

Original: Whereas the conditional distribution over experts given a datapoint should ideally have low entropy,the marginal distribution over experts should be balanced for efﬁcient training and use of modelcapacity.
Revised: Whereas the conditional distribution over experts given a datapoint should ideally have low entropy, the marginal distribution over experts should be balanced for efﬁcient training and use of model capacity.

Original: In some cases this is explicitly encouraged by using a load balancing loss [2, 8].
Revised: In some cases this is explicitly encouraged by using a load balancing loss [2, 8].

Original: Othermethods algorithmically balance the assignment [25] or use a ﬁxed distribution [35].
Revised: Other methods algorithmically balance the assignment [26] or use a ﬁxed distribution [36].

Original: Many recentlarge scale MoE models use heuristics to train the router module (see Section 3.1) but have yieldedstate-of-the-art performance in different domains [37, 24, 8, 25, 34, 44].
Revised: Many recent large scale MoE models use heuristics to train the router module (see Section 3.1) but have yielded state-of-the-art performance in different domains [38, 25, 8, 26, 35, 45].

Original: Sampling and optimization of balanced assignments is closely related to the sampling and optimizationof permutations ( n × n assignments)
Revised: Sampling and optimization of balanced assignments is closely related to the sampling and optimization of permutations ( n × n assignments)

Original: [23, 1, 26, 29, 28, 12, 31], which relate to estimation of thematrix permanent [15, 16, 21].
Revised: [24, 1, 27, 30, 29, 13, 32], which relate to estimation of the matrix permanent [16, 17, 22].

Original: In a different setting, the problem can also be seen as random fairassignment [5], with the difference being that in random fair assignment one typically is not concernedabout dependence between the individual assignments.
Revised: In a different setting, the problem can also be seen as random fair assignment [5], with the difference being that in random fair assignment one typically is not concerned about dependence between the individual assignments.

Original: B Derivation of (off-policy) REINFORCE
Revised: B Derivation of (off-policy) REINFORCE

Original: First, we derive (off-policy) REINFORCE for a single datapoint x :
Revised: First, we derive (off-policy) REINFORCE for a single datapoint x :

Original: [Equation]
Revised: [Equation]

Original: Now consider a minibatch x = ( x 1 , ..., x n ) and let p θ ( z | x ) = (cid:81)
Revised: Now consider a minibatch x = ( x 1 , ..., x n ) and let p θ ( z | x ) = (cid:81)

Original: i p θ
Revised: i p θ

Original: ( z i | x i ) be the distribution thatsamples expert assignments independently.
Revised: ( z i | x i ) be the distribution that samples expert assignments independently.

Original: Let q ( z i | x ) = (cid:80)
Revised: Let q ( z i | x ) = (cid:80)

Original: z − i q ( z | x ) be the marginal of any jointproposal distribution q ( z | x ) , which we can use to estimate the minibatch gradient:
Revised: z − i q ( z | x ) be the marginal of any joint proposal distribution q ( z | x ) , which we can use to estimate the minibatch gradient:

Original: [Equation]
Revised: [Equation]

Original: Since ∇ E z ∼ p θ ( z | x ) [ b ] = 0 , we can subtract any constant baseline b from f ( x, y ) , resulting in (2).
Revised: Since ∇ E z ∼ p θ ( z | x ) [ b ] = 0 , we can subtract any constant baseline b from f ( x, y ) , resulting in (2).

Original: C Unbiasedness of the skipping estimator
Revised: C Unbiasedness of the skipping estimator

Original: First, consider a function h ( x i , z i ) .
Revised: First, consider a function h ( x i , z i ) .

Original: Let z ij = 1 { z i = j } be the one-hot representation of z i and leth ij = h
Revised: Let z ij = 1 { z i = j } be the one-hot representation of z i and let h ij = h

Original: ( x i , z i ) | z i = j .
Revised: ( x i , z i ) | z i = j .

Original: Let n j = (cid:80)
Revised: Let n j = (cid:80)

Original: i z ij be the number of datapoints assigned to expert j (beforesubsampling).
Revised: i z ij be the number of datapoints assigned to expert j (before subsampling).

Original: Now let δ i ∈ { 0 , 1 } represent which datapoints are kept after we, for each expert j ,uniformly subsample min { n j , c } datapoints, where c
Revised: Now let δ i ∈ { 0 , 1 } represent which datapoints are kept after we, for each expert j , uniformly subsample min { n j , c } datapoints, where c

Original: = nk is the expert capacity.
Revised: = nk is the expert capacity.

Original: If z ij = 1 (beforesubsampling), then the probability that datapoint i remains after subsampling is min { n j ,c } n j , so wehave E δ [Equation] | z [ δ i ] [Equation]
Revised: If z ij = 1 (before subsampling), then the probability that datapoint i remains after subsampling is min { n j ,c } n j , so we have E δ ∗ | z [ δ i ] n j min { n j ,c } z ij = z ij .

Original: [Equation]
Revised: [Equation]

Original: Now substituting h ( x i , z i )
Revised: Now substituting h ( x i , z i )

Original: = ∇ p θ ( z i | x i ) q
Revised: = ∇ p θ ( z i | x i ) q

Original: ( z i | x ) f ( x i , z i ) and combining with Equation (2) results in (3).
Revised: ( z i | x ) f ( x i , z i ) and combining with Equation (2) results in (3).

Original: D The Gumbel-Matching distribution 
Revised: D The Gumbel-Matching distribution D.1 Approximation

Original: D.1 Approximation to the Gibbs distribution
Revised: distribution

Original: Here we essentially reproduce the argument from [28] for the n × k Gumbel-Matching distribution.
Revised: Here we essentially reproduce the argument from [29] for the n × k Gumbel-Matching distribution.

Original: By sampling i.i.d.
Revised: By sampling i.i.d.

Original: Gumbel noise g z for every assignment z , we can sample from (5) by maximizing
Revised: Gumbel noise g z for every assignment z , we can sample from (5) by maximizing

Original: [Equation]
Revised: [Equation]

Original: subject to the constraints given by (4).
Revised: subject to the constraints given by (4).

Original: Comparing this to the objective for the Gumbel-Matchingproblem (4):
Revised: Comparing this to the objective for the Gumbel-Matching problem (4):

Original: [Equation]
Revised: [Equation]

Original: we observe how the Gumbel-Matching distribution approximates (5) through the use of rank-oneperturbations [30, 14, 41] (cid:80) ij z ij g ij instead of g z .
Revised: we observe how the Gumbel-Matching distribution approximates (5) through the use of rank-one perturbations [31, 15, 42] (cid:80) ij z ij g ij instead of g z .

Original: D.2 Solving n × k matching using cycle cancelling with Floyd-Warshall
Revised: D.2 Solving n × k matching using cycle cancelling with Floyd-Warshall

Original: The n × k assignment problem can be modelled as a minimum cost ﬂow problem which can besolved using cycle cancelling [19] as follows:
Revised: The n × k assignment problem can be modelled as a minimum cost ﬂow problem which can be solved using cycle cancelling [20] as follows:

Original: • Find an initial (heuristic) feasible assignment z .
Revised: • Find an initial (heuristic) feasible assignment z .

Original: We use the auction algorithm used in [25] with (cid:15) = 1 . 

Revised: We use the auction algorithm used in [26] with (cid:15) = 1 .

Original: such that it ﬁnds a good (but suboptimal) solution quickly.
Revised: such that it ﬁnds a good (but suboptimal) solution quickly.

Original: • For every combination of experts j, j (cid:48) , ﬁnd the lowest cost d jj (cid:48) to move a datapoint from expert j to j (cid:48) .
Revised: • For every combination of experts j, j (cid:48) , ﬁnd the lowest cost d jj (cid:48) to move a datapoint from expert j to j (cid:48) .

Original: Let s ij = a ij /τ +
Revised: Let s ij = a ij /τ +

Original: g
Revised: g

Original: ij be the score for assigning datapoint i to j , such thatmoving datapoint i from j to j (cid:48)
Revised: ij be the score for assigning datapoint i to j , such that moving datapoint i from j to j (cid:48)

Original: we lose s ij but gain s ij (cid:48) , incurring a net ‘cost’ of s ij − s ij (cid:48) .
Revised: we lose s ij but gain s ij (cid:48) , incurring a net ‘cost’ of s ij − s ij (cid:48) .

Original: Therefore, the minimum cost to move any of the currently assigned datapoints from j to j (cid:48)is d jj (cid:48) =
Revised: Therefore, the minimum cost to move any of the currently assigned datapoints from j to j (cid:48) is d jj (cid:48) =

Original: min i :
Revised: [Equation]

Original:  z ij =1 s ij − s ij (cid:48) .
Revised: i : z ij =1 s ij − s ij (cid:48) .

Original: • Use the Floyd-Warshall algorithm 1 [9] to ﬁnd all indirect shortest paths in the fully connectedgraph with k nodes (one per expert) and distance from j to j (cid:48) given by d jj (cid:48) .
Revised: • Use the Floyd-Warshall algorithm 2 [9] to ﬁnd all indirect shortest paths in the fully connected graph with k nodes (one per expert) and distance from j to j (cid:48) given by d jj (cid:48) .

Original: Stop as soon asa negative cycle is found (distance from j to j smaller than 0).
Revised: Stop as soon as a negative cycle is found (distance from j to j smaller than 0).

Original: • If no negative cycle exists, the assignment is optimal, stop.•
Revised: • If no negative cycle exists, the assignment is optimal, stop.

Original:  For each edge ( j, j (cid:48) ) in the negative cycle, 2 move the datapoint i that minimizes s ij − s ij (cid:48)from j to j (cid:48) .
Revised: • For each edge ( j, j (cid:48) ) in the negative cycle, 3 move the datapoint i that minimizes s ij − s ij (cid:48) from j to j (cid:48) .

Original: This will improve the assignment by incurring a negative total cost.•
Revised: This will improve the assignment by incurring a negative total cost.

Original:  Repeat until no negative cycle exists.
Revised: • Repeat until no negative cycle exists.

Original: Depending on the initial assignment, only a small number of O ( k 3 ) improvements is needed and weﬁnd in practice for k (cid:28) n this algorithm is much faster than the O ( n 3 ) Hungarian [22] algorithm.
Revised: Depending on the initial assignment, only a small number of O ( k 3 ) improvements is needed and we ﬁnd in practice for k (cid:28) n this algorithm is much faster than the O ( n 3 ) Hungarian [23] algorithm.

Original: D.3 Computing conditionally optimal assignments
Revised: D.3 Computing conditionally optimal assignments

Original: If the assignment is optimal, we denote the entries of the all-pairs shortest path matrix resultingfrom the Floyd-Warshall algorithm by d ∗ jj (cid:48)
Revised: If the assignment is optimal, we denote the entries of the all-pairs shortest path matrix resulting from the Floyd-Warshall algorithm by d ∗ jj (cid:48)

Original: (see Appendix D.2).
Revised: (see Appendix D.2).

Original: We can use this to efﬁciently obtainconditionally optimal assignments , conditioning on z ij = 1 for all i, j , as follows.
Revised: We can use this to efﬁciently obtain conditionally optimal assignments , conditioning on z ij = 1 for all i, j , as follows.

Original: If we conditionon z ij = 1 , we may move datapoint i from the globally optimal assignment j ∗ to a suboptimalassignment j , incurring a cost s ij ∗ − s ij , and move another datapoint from j to j ∗ for a cost ofd ∗ jj ∗ , indirectly via the path found by the Floyd-Warshall algorithm.
Revised: If we condition on z ij = 1 , we may move datapoint i from the globally optimal assignment j ∗ to a suboptimal assignment j , incurring a cost s ij ∗ − s ij , and move another datapoint from j to j ∗ for a cost of d ∗ jj ∗ , indirectly via the path found by the Floyd-Warshall algorithm.

Original: As such, the total cost ofenforcing z ij = 1 is s ij ∗ − s ij + d ∗ jj ∗ .
Revised: As such, the total cost of enforcing z ij = 1 is s ij ∗ − s ij + d ∗ jj ∗ .

Original: We denote the value of the globally optimal assignment by v ∗ .
Revised: We denote the value of the globally optimal assignment by v ∗ .

Original: Subtracting the cost for enforcing z ij = 1 , we ﬁnd that the value v ∗| z ij =1 of the conditionally optimalassignment is given by
Revised: Subtracting the cost for enforcing z ij = 1 , we ﬁnd that the value v ∗| z ij =1 of the conditionally optimal assignment is given by

Original: [Equation]
Revised: [Equation]

Original: D.4 Computation of the conditionals
Revised: D.4 Computation of the conditionals

Original: The dependence of the Gumbel-Matching distribution on x is only through the logits
Revised: The dependence of the Gumbel-Matching distribution on x is only through the logits

Original: A =
Revised: A =

Original: ( a ij ) ,which allows us to slightly simplify notation in the rest of this section.
Revised: ( a ij ) , which allows us to slightly simplify notation in the rest of this section.

Original: Let v ∗ ( A , G ) be the valueof the optimal assignment for the Gumbel-Matching problem with logits
Revised: Let v ∗ ( A , G ) be the value of the optimal assignment for the Gumbel-Matching problem with logits

Original: A =
Revised: A =

Original: ( a ij ) and noise
Revised: ( a ij ) and noise

Original: G = ( g ij ) , and let v ∗| z ij =1 ( A , G ) be the value of the conditionally optimal assignment with theadditional constraint that z ij = 1 (see Appendix D.3).
Revised: G = ( g ij ) , and let v ∗| z ij =1 ( A , G ) be the value of the conditionally optimal assignment with the additional constraint that z ij = 1 (see Appendix D.3).

Original: Assuming a capacity c
Revised: Assuming a capacity c

Original: = nk , then giventhat z ij = 1 , the problem reduces to assigning the remaining n − 1 datapoints to the remaining( k − 1) · nk + ( nk − 1) = n − 1 ‘slots’ ( nk for experts j (cid:48) (cid:54) =
Revised: = nk , then given that z ij = 1 , the problem reduces to assigning the remaining n − 1 datapoints to the remaining ( k − 1) · nk + ( nk − 1) = n − 1 ‘slots’ ( nk for experts j (cid:48) (cid:54) =

Original: j and nk − 1 for expert j ).
Revised: j and nk − 1 for expert j ).

Original: As this reducedproblem does not depend on datapoint i , we denote with A − i and G −
Revised: As this reduced problem does not depend on datapoint i , we denote with A − i and G −

Original: i the logits and Gumbels withrow i removed and we let v ∗− ij ( A − i , G − i ) be the value of the optimal assignment of the reducedproblem.
Revised: i the logits and Gumbels with row i removed and we let v ∗− ij ( A − i , G − i ) be the value of the optimal assignment of the reduced problem.

Original: From the principle of optimality it follows that
Revised: From the principle of optimality it follows that

Original: [Equation]
Revised: [Equation]

Original: We can use this to compute the desired conditionals.
Revised: We can use this to compute the desired conditionals.

Original: In a slight abuse of notation 3 we writeq ( z ij | A , G − i)
Revised: In a slight abuse of notation 4 we write q ( z ij | A [Equation] G − i

Original: [Equation]
Revised: [Equation]

Original: where we have used the Gumbel-max trick.
Revised: where we have used the Gumbel-max trick.

Original: Although g ij appears in (9), q ( z ij | A , G − i ) does notdepend on g i as this value cancels against g ij in (8).
Revised: Although g ij appears in (9), q ( z ij | A , G − i ) does not depend on g i as this value cancels against g ij in (8).

Original: The values v ∗ | z ij =1 ( A , G ) can be computed efﬁciently using Equation (7) with the method described in Appendix D.3.
Revised: The values v ∗ | z ij =1 ( A , G ) can be computed efﬁciently using Equation (7) with the method described in Appendix D.3.

Original: D.5 Unbiased of the Gumbel-Matching estimator
Revised: D.5 Unbiased of the Gumbel-Matching estimator

Original: Let z = GM (log p ( ·| x ) , G ) be the solution for the Gumbel-Matching problem with noise G .
Revised: Let z = GM (log p ( ·| x ) , G ) be the solution for the Gumbel-Matching problem with noise G .

Original: Theidea behind the Gumbel-Matching estimator is that we can derive an unbiased estimate of the gradientfor each datapoint as follows.
Revised: The idea behind the Gumbel-Matching estimator is that we can derive an unbiased estimate of the gradient for each datapoint as follows.

Original: First sample the Gumbel noise G − i for all datapoints except i , andcompute the conditional distribution q θ ( z i | x , G − i )
Revised: First sample the Gumbel noise G − i for all datapoints except i [Equation]

Original:  (see Appendix D.4 where A = log p ( ·| x ) is thematrix with log-probabilities and z ij = 1 ⇔ z i = j ).
Revised: and compute the conditional distribution q θ ( z i | x , G − i ) (see Appendix D.4 where A = log p ( ·| x ) is the matrix with log-probabilities and z ij = 1 ⇔ z i = j ).

Original: We can then use this conditional distributionover
Revised: We can then use this conditional distribution over z

Original: z i as a proposal distribution in (2), which we can then reparameterize in terms of the Gumbelnoise g i for the i -th datapoint.
Revised: i as a proposal distribution in (2), which we can then reparameterize in terms of the Gumbel noise g i for the i -th datapoint.

Original: Finally, we can use this estimator for all datapoints i , where we mayreuse the same Gumbel noise and use their average as the estimate:
Revised: Finally, we can use this estimator for all datapoints i [Equation] where we may reuse the same Gumbel noise and use their average as the estimate:

Original: [Equation]
Revised: [Equation]

Original: D.6 Maximum-entropy distribution over balanced assignments
Revised: D.6 Maximum-entropy distribution over balanced assignments

Original: For simplicity, we assume n = k , but this can be easily generalized.
Revised: For simplicity, we assume n = k , but this can be easily generalized.

Original: Assume that we have a balanced(square) matrix (see Section 2.3) P =
Revised: Assume that we have a balanced (square) matrix (see Section 2.3) P =

Original: ( p ij ) with probabilities p ij ≥ 0 such that (cid:80) i p ij = (cid:80) j p ij =1 , i.e. the matrix P is doubly stochastic .
Revised: ( p ij ) with probabilities p ij ≥ 0 such that (cid:80) i p ij = (cid:80) j p ij = 1 , i.e. the matrix P is doubly stochastic .

Original: The Birkhoff decomposition [4, 42] decomposes such amatrix as a convex combination over permutation matrices
Revised: The Birkhoff decomposition [4, 43] decomposes such a matrix as a convex combination over permutation matrices

Original: Z = ( z ij ) , for which z ij ∈ { 0 , 1 } and364 (cid:80) i z ij = (cid:80)
Revised: Z = ( z ij ) , for which z ij ∈ { 0 , 1 } and (cid:80) i z ij = (cid:80)

Original: j
Revised: j

Original: z ij = 1 .
Revised: z ij = 1 .

Original: Such permutation matrices represent n × n matchings z in one-hot encoding(i.e. z ij = 1 ⇔ z i = j ), which is why we will use z to denote them.
Revised: Such permutation matrices represent n × n matchings z in one-hot encoding (i.e. z ij = 1 ⇔ z i = j ), which is why we will use z to denote them.

Original: A Birkhoff decompositionα z > 0 , (cid:80) z α z = 1 thus represents a joint probability distribution over matchings z (represented aspermutation matrices) with marginal distributions (per datapoint) given by P :
Revised: A Birkhoff decomposition α z > 0 , (cid:80) z α z = 1 thus represents a joint probability distribution over matchings z (represented as permutation matrices) with marginal distributions (per datapoint) given by P :

Original: [Equation]
Revised: [Equation]

Original: Many such decompositions/distributions exist and can be found using the Birkhoff algorithm[4], butthese will yield sparse α z and thus have low entropy and high dependence between the marginaldistributions for different i .
Revised: Many such decompositions/distributions exist and can be found using the Birkhoff algorithm[4], but these will yield sparse α z and thus have low entropy and high dependence between the marginal distributions for different i .

Original: We aim to minimize this dependence between the marginal distributions,which is achieved by maximizing the entropy − (cid:80)
Revised: We aim to minimize this dependence between the marginal distributions, which is achieved by maximizing the entropy − (cid:80)

Original: z
Revised: z

Original: α z log α z subject to the constraint (10) and372 (cid:80) z
Revised: α z log α z subject to the constraint (10) and (cid:80) z

Original: α z = 1 , which has the Lagrangian:
Revised: α z = 1 , which has the Lagrangian:

Original: [Equation]
Revised: [Equation]

Original: This has ﬁrst order conditions
Revised: This has ﬁrst order conditions

Original: [Equation]
Revised: [Equation]

Original: If we let u ij =
Revised: If we let u ij =

Original: exp
Revised: exp

Original: ( λ ij ) and w = exp( η − 1) , and convert from the ‘one-hot’ representation z ij tothe ‘indexing’ representation
Revised: ( λ ij ) and w = exp( η − 1) , and convert from the ‘one-hot’ representation z ij to the ‘indexing’ representation

Original: z i (i.e. (cid:80) ij λ ij z ij = (cid:80) i λ i,z i ), we ﬁnd the solutionα z =
Revised: z i (i.e. (cid:80) ij λ ij z ij = (cid:80) i λ i,z i ), we ﬁnd the solution α z =

Original: [Equation]
Revised: [Equation]

Original: Here Perm ( U ) is the permanent of the matrix U =
Revised: Here Perm ( U ) is the permanent of the matrix U =

Original: ( u ij ) , and U − ij is the matrix U with rows i and jremoved.
Revised: ( u ij ) , and U − ij is the matrix U with rows i and j removed.

Original: Since w = 1 / Perm ( U ) we ﬁnd
Revised: Since w = 1 / Perm ( U ) we ﬁnd

Original: [Equation]
Revised: [Equation]

Original: which can, in theory, be solved using a (very expensive) ﬁxed point iteration scheme.
Revised: which can, in theory, be solved using a (very expensive) ﬁxed point iteration scheme.

Original: Empiricallywe found that the solution takes the form u ij ≈ p 1 /τij for some τ , such that the probability of anassignment z is given by
Revised: Empirically we found that the solution takes the form u ij ≈ p 1 /τij for some τ , such that the probability of an assignment z is given by

Original: [Equation]
Revised: [Equation]

Original: Here P (1 /τ ) is the Hadamard power , which raises the entries p ij to the power 1 /τ element-wise.
Revised: Here P (1 /τ ) is the Hadamard power , which raises the entries p ij to the power 1 /τ element-wise.

Original: Bygeneralizing permanents (sums over permutations) to non-square matrices as sums over balancedassignment matrices, the same result can be derived for n (cid:54) =
Revised: By generalizing permanents (sums over permutations) to non-square matrices as sums over balanced assignment matrices, the same result can be derived for n (cid:54) =

Original: k .
Revised: k .

Original: Using a ij = log p ij and converting to‘one-hot’ representation z ij , we ﬁnd that (11) is equal to (5).
Revised: Using a ij = log p ij and converting to ‘one-hot’ representation z ij , we ﬁnd that (11) is equal to (5).

Original: If we do not use the approximation u ij ≈ p 1 /τij , then it still holds that the maximum entropydistribution has the form (5), but we should let τ = 1 and a ij = log u ij , so in this case a ij (cid:54)
Revised: If we do not use the approximation u ij ≈ p 1 /τij , then it still holds that the maximum entropy distribution has the form (5), but we should let τ = 1 and a ij = log u ij , so in this case a ij (cid:54)

Original: = log p ijare not the marginal log-probabilities we aim to sample from. 
Revised: = log p ij are not the marginal log-probabilities we aim to sample from. [Figure]

Original: E Experiment
Revised: E Experiment

Original: Figure 3 presents results for both REINFORCE (top) and  gating (bottom) both with (middle) andwithout (left) the Sinkhorn normalization before sampling.
Revised: Figure 3 presents results for both REINFORCE (top) and differentiable gating (bottom) both with (middle) and without (left) the Sinkhorn normalization before sampling.

Original: When using Sinkhorn normalizationbefore sampling, we take it into account when computing the importance weights.
Revised: When using Sinkhorn normalization before sampling, we take it into account when computing the importance weights.

Original: With REINFORCE,not using Sinkhorn normalization works better, as using Sinkhorn normalization requires a highsampling temperature to work well, i.e. close to uniform proposal samples.
Revised: With REINFORCE, not using Sinkhorn normalization works better, as using Sinkhorn normalization requires a high sampling temperature to work well, i [Figure] e. close to uniform proposal samples.

Original: With  gating (and load balancing loss with weight 0 . 01 ), we observe the opposite: the results are betterwith Sinkhorn normalization than without Sinkhorn normalization (Section 3.1), but, contrasting REINFORCE, require a very low sampling temperature to succeed, so close to deterministic training.
Revised: With differentiable gating (and load balancing loss with weight 0 . 01 ), we observe the opposite: the results are better with Sinkhorn normalization than without Sinkhorn normalization (Section 3.1), but, contrasting REINFORCE, require a very low sampling temperature to succeed, so close to deterministic training.

Original: Existing methods often do not use importance weights to take into account the noise scale ortemperature [37, 8], so we experiment with this as well by completely dropping importance weightsfor all estimators (right column in Figure 3).
Revised: Existing methods often do not use importance weights to take into account the noise scale or temperature [38, 8], so we experiment with this as well by completely dropping importance weights for all estimators (right column in Figure 3).

Original: We ﬁnd this only works for REINFORCE, whensampling with a low temperature but not when we use the Gumbel-Matching estimators.
Revised: We ﬁnd this only works for REINFORCE, when sampling with a low temperature but not when we use the Gumbel-Matching estimators.

Original: In all cases, we found REINFORCE succeeds both with and without a load balancing loss, whereasgating requires a load balancing loss with a weight of 0.01 to succeed at all.
Revised: In all cases, we found REINFORCE succeeds both with and without a load balancing loss, whereas differentiable gating requires a load balancing loss with a weight of 0.01 to succeed at all.

Original: If the load balance lossis too high (0.1) all methods fail to model the unbalanced data.
Revised: If the load balance loss is too high (0.1) all methods fail to model the unbalanced data.

Original: [Figure]