Original: E FFICIENT D ISCRETE M ULTI -M ARGINAL O PTIMAL T RANSPORT R EGULARIZATION
Revised: E FFICIENT D ISCRETE M ULTI -M ARGINAL O PTIMAL T RANSPORT R EGULARIZATION

Original: Anonymous authors
Revised: Ronak Mehta ∗ UW-Madison Jeffery Kline Affirm

Revised: Vishnu Suresh Lokhande

Revised: UW-Madison

Original: Paper under double-blind review
Revised: Glenn Fung Liberty Mutual Vikas Singh UW-Madison

Original: A BSTRACT
Revised: A BSTRACT

Original: Optimal transport has emerged as a powerful tool for a variety of problems in machine learning, and it is frequently used to enforce distributional constraints.
Revised: Optimal transport has emerged as a powerful tool for a variety of problems in machine learning, and it is frequently used to enforce distributional constraints.

Original: In this context, existing methods often use either a Wasserstein metric, or else they apply concurrent barycenter approaches when more than two distributions are considered.
Revised: In this context, existing methods often use either a Wasserstein metric, or else they apply concurrent barycenter approaches when more than two distributions are considered.

Original: In this paper, we leverage multi-marginal optimal transport (MMOT), where we take advantage of a procedure that computes a generalized earth mover’s distance as a sub-routine.
Revised: In this paper, we leverage multi-marginal optimal transport (MMOT), where we take advantage of a procedure that computes a generalized earth mover’s distance as a sub-routine.

Original: We show that not only is our algorithm computationally more efficient compared to other barycentric-based distance methods, but it has the additional advantage that gradients used for backpropagation can be efficiently computed during the forward pass computation itself, which leads to substantially faster model training.
Revised: We show that not only is our algorithm computationally more efficient compared to other barycentric-based distance methods, but it has the additional advantage that gradients used for backpropagation can be efficiently computed during the forward pass computation itself, which leads to substantially faster model training.

Original: We provide technical details about this new regularization term and its properties, and we present experimental demonstrations of faster runtimes when compared to standard Wasserstein-style methods.
Revised: We provide technical details about this new regularization term and its properties, and we present experimental demonstrations of faster runtimes when compared to standard Wasserstein-style methods.

Original: Finally, on a range of experiments designed to assess effectiveness at enforcing fairness, we demonstrate our method compares well with alternatives.
Revised: Finally, on a range of experiments designed to assess effectiveness at enforcing fairness, we demonstrate our method compares well with alternatives.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The use of Optimal transport (OT) is now prevalent in many problem settings including information retrieval (Balikas et al., 2018; Yurochkin et al., 2019), image processing (Bonneel et al., 2014), statistical machine learning, as well as more recently, for ethics and fairness research (KwegyirAggrey et al., 2021).
Revised: The use of Optimal transport (OT) is now prevalent in many problem settings including information retrieval (Balikas et al., 2018; Yurochkin et al., 2019), image processing (Bonneel et al., 2014), statistical machine learning, and more recently, for ethics and fairness research (Kwegyir-Aggrey et al., 2021; Lokhande et al., 2020a).

Original: OT is well-suited for tasks where dissimilarity between two or more probability distributions must be quantified; its success was made possible through dramatic improvements in algorithms (Cuturi, 2013; Solomon et al., 2015) that allow one to efficiently optimize commonly used functionals.
Revised: OT is well-suited for tasks where dissimilarity between two or more probability distributions must be quantified; its success was made possible through dramatic improvements in algorithms (Cuturi, 2013; Solomon et al., 2015) that allow one to efficiently optimize commonly used functionals.

Original: In practice, OT is often used to estimate and minimize the distance between certain (data-derived) distributions, using an appropriately defined loss functional.
Revised: In practice, OT is often used to estimate and minimize the distance between certain (data-derived) distributions, using an appropriately defined loss functional.

Original: When one seeks to operate on more than two distributions, however, newer constructions are necessary to effectively estimate distances and transports.
Revised: When one seeks to operate on more than two distributions, however, newer constructions are necessary to effectively estimate distances and transports.

Original: To this end, a well studied idea in the literature is the “barycenter,” identified by minimizing the pairwise distance between itself and all other distributions given.
Revised: To this end, a well studied idea in the literature is the “barycenter,” identified by minimizing the pairwise distance between itself and all other distributions given.

Original: The d -dimensional proxy distance is then defined as the sum of the distances to the barycenter.
Revised: The d -dimensional proxy distance is then defined as the sum of the distances to the barycenter.

Original: Computing barycenters.
Revised: Computing barycenters.

Original: Assuming that a suitably regularized form of the optimal transport loss is utilized, the pairwise distance calculation, by itself, can be efficient – in fact, in some cases, Sinkhorn iterations can be used (Cuturi, 2013).
Revised: Assuming that a suitably regularized form of the optimal transport loss is utilized, the pairwise distance calculation, by itself, can be efficient – in fact, in some cases, Sinkhorn iterations can be used (Cuturi, 2013).

Original: On the other hand, to minimize distances to the mean, most algorithms typically operate by repeatedly estimating the barycenter and those pairwise distances, and using a “coupling” strategy to push points toward the barycenter, or in other cases, summing over all pairwise distances.
Revised: On the other hand, to minimize distances to the mean, most algorithms typically operate by repeatedly estimating the barycenter and those pairwise distances, and using a “coupling” strategy to push points toward the barycenter, or in other cases, summing over all pairwise distances.

Original: As the number of distributions grows, robustness issues can exacerbate (Alvarez-Esteban et al., 2008) and the procedure is expensive (e.g., for 50 distributions, 50 bins).
Revised: As the number of distributions grows, robustness issues can exacerbate (Alvarez-Esteban et al., 2008) and the procedure is expensive (e.g., for 50 distributions, 50 bins).

Original: A potential alternative.
Revised: A potential alternative.

Original: Multi-marginal optimal transport (MMOT) is a related problem to the aforementioned task but to some extent, the literature has developed in parallel.
Revised: Multi-marginal optimal transport (MMOT) is a related problem to the aforementioned task but to some extent, the literature has developed in parallel.

Original: In particular, MMOT focuses on identifying a joint distribution such that the marginals are defined by the input distributions over which we wish to measure the dissimilarity.
Revised: In particular, MMOT focuses on identifying a joint distribution such that the marginals are defined by the input distributions over which we wish to measure the dissimilarity.

Original: The definition naturally extends the two-dimensional formulation, and recent work has explored a number of applications (Pass, 2015).
Revised: The definition naturally extends the two-dimensional formulation, and recent work has explored a number of applications (Pass, 2015).

Original: But the MMOT computation can be quite difficult, and only very recently have practical algorithms been identified(Lin et al., 2022).
Revised: But the MMOT computation can be quite difficult, and only very recently have practical algorithms been identified

Original: Additionally, even if a suitable method for computing an analogous measure of distance were available, minimizing this distance to reduce dissimilarity (push distributions closer

Original: [Figure]
Revised: [Figure]

Original:  to each other) is practically hard if standard interior point solvers are needed just to compute the distance itself.
Revised: (Lin et al., 2022). Additionally, even if a suitable method for computing an analogous measure of distance were available, minimizing this distance to reduce dissimilarity (push distributions closer to each other) is practically hard if standard interior point solvers are needed just to compute the distance itself.

Original: Why and where is dissimilarity important?
Revised: Why and where is dissimilarity important?

Original: Enforcing distributions to be similar is a generic goal whenever one wishes some outcome of interest to be agnostic about particular groups within the input data.
Revised: Enforcing distributions to be similar is a generic goal whenever one wishes some outcome of interest to be agnostic about particular groups within the input data.

Original: In applications where training deep neural network models is needed, it is often a goal to enforce distribution similarity on model outputs.
Revised: In applications where training deep neural network models is needed, it is often a goal to enforce distribution similarity on model outputs.

Original: For example, in Jiang et al.
Revised: For example, in Jiang et al.

Original: (2020), the authors define fairness measures over the probability of the prediction, given ground truth labels.
Revised: (2020), the authors define fairness measures over the probability of the prediction, given ground truth labels.

Original: However, these methods are rarely extended to continuous measures among internal neural network activations, mainly due to the strong distributional assumptions needed (product of Gaussians) and the added algorithmic complexity of estimating the barycenter.
Revised: However, these methods are rarely extended to continuous measures among internal neural network activations, mainly due to the strong distributional assumptions needed (product of Gaussians) and the added algorithmic complexity of estimating the barycenter.

Original: These issues limit application of these ideas to only the final outputs of neural network models, where the distribution is typically binomial or multinomial.
Revised: These issues limit application of these ideas to only the final outputs of neural network models, where the distribution is typically binomial or multinomial.

Original: MMOT solutions might be employed here, but suffer similar computational limitations.
Revised: MMOT solutions might be employed here, but suffer similar computational limitations.

Original: Contributions.
Revised: Contributions.

Original: (1) We identify a particular form of the discrete multi-marginal optimal transport problem which admits an extremely fast and numerically robust solution.
Revised: (1) We identify a particular form of the discrete multi-marginal optimal transport problem which admits an extremely fast and numerically robust solution.

Original: Exploiting a recent extension of the classical Earth Movers Distance (EMD) to a higher-dimensional Earth Mover’s objective, we show that such a construction is equivalent to the discrete MMOT problem with Monge costs.
Revised: Exploiting a recent extension of the classical Earth Movers Distance (EMD) to a higher-dimensional Earth Mover’s objective, we show that such a construction is equivalent to the discrete MMOT problem with Monge costs.

Original: We show that minimization of this global distributional measure leads to the harmonization of input distributions very similar in spirit to the minimization of distributions to barycenters (see Figure 1).
Revised: We show that minimization of this global distributional measure leads to the harmonization of input distributions very similar in spirit to the minimization of distributions to barycenters (see Figure 1).

Original: (3) We prove theoretical properties of our scheme, and show that the gradient can be read directly off from a primal/dual algorithm, alleviating the need for computationally intense pairwise couplings needed for barycenter approaches.
Revised: (3) We prove theoretical properties of our scheme, and show that the gradient can be read directly off from a primal/dual algorithm, alleviating the need for computationally intense pairwise couplings needed for barycenter approaches.

Original: The direct availability of the gradient enables a specific neural network instantiation, and with a particular scaffolding provided by differentiable histograms, we can operate directly on network activations (anywhere in the network) to compute/minimize the d-MMOT.
Revised: The direct availability of the gradient enables a specific neural network instantiation, and with a particular scaffolding provided by differentiable histograms, we can operate directly on network activations (anywhere in the network) to compute/minimize the d-MMOT.

Original: We establish via experiments that computing gradients used in backpropagation is fast, due to rapid access to solutions of the dual linear program.
Revised: We establish via experiments that computing gradients used in backpropagation is fast, due to rapid access to solutions of the dual linear program.

Original: We compare with barycenter-like approaches in several settings, including common fairness applications.
Revised: We compare with barycenter-like approaches in several settings, including common fairness applications.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Despite originating with (Monge, 1781), optimal transport continues to be an active area of research (Villani, 2009).
Revised: Despite originating with (Monge, 1781), optimal transport continues to be an active area of research (Villani, 2009).

Original: The literature is vast, but we list a few key developments.
Revised: The literature is vast, but we list a few key developments.

Original: Early applications.
Revised: Early applications.

Original: Starting in (Peleg et al., 1989), the idea of shifting “mass” around within an image was used for comparing images to each other and applied to image retrieval (Rubner et al., 2000), where the term “Earth Mover’s Distance” (EMD) was introduced.
Revised: Starting in (Peleg et al., 1989), the idea of shifting “mass” around within an image was used for comparing images to each other and applied to image retrieval (Rubner et al., 2000), where the term “Earth Mover’s Distance” (EMD) was introduced.

Original: EMD has since been widely used in computer vision: e.g., for image warping (Zhang et al., 2011), in supervised settings (Wang & Guibas, 2012), matching point sets (Cabello et al., 2008) and in scenarios involving histogram comparisons (Ling & Okada, 2007; Wang & Guibas, 2012; Haker et al., 2004).
Revised: EMD has since been widely used in computer vision: e.g., for image warping (Zhang et al., 2011), in supervised settings (Wang & Guibas, 2012), matching point sets (Cabello et al., 2008) and in scenarios involving histogram comparisons (Ling & Okada, 2007; Wang & Guibas, 2012; Haker et al., 2004).

Original: Modern machine learning.
Revised: Modern machine learning.

Original: The continuous optimal transport problem (Monge-Kantorovich problem), was originally presented in (Kantorovich, 1942; Kantorovitch, 1958).
Revised: The continuous optimal transport problem (Monge-Kantorovich problem), was originally presented in (Kantorovich, 1942; Kantorovitch, 1958).

Original: While the continuous problem has been studied intensively (Villani, 2021), uses of optimal transport within machine learning were possible due to (Cuturi, 2013), which showed that entropic regularization enables fast algorithms for EMD (two distributions with discrete support), and contributed to the success of
Revised: While the continuous problem has been studied intensively (Villani, 2021), uses of optimal transport within machine learning were possible due to (Cuturi, 2013), which showed that entropic regularization enables fast algorithms for EMD (two distributions with discrete support), and contributed to the success of

Original: Wasserstein distances in applications.
Revised: Wasserstein distances in applications.

Original: Consequently, problems including autoencoders (Tolstikhin et al., 2018), GANs (Arjovsky et al., 2017), domain adaptation (Courty et al., 2016), word embeddings (Huang et al., 2016) and classification tasks (Frogner et al., 2015) have benefited via the use of optimal transport.
Revised: Consequently, problems including autoencoders (Tolstikhin et al., 2018), GANs (Arjovsky et al., 2017), domain adaptation (Courty et al., 2016), word embeddings (Huang et al., 2016) and classification tasks (Frogner et al., 2015) have benefited via the use of optimal transport.

Original: Multi-marginal optimal transport.
Revised: Multi-marginal optimal transport.

Original: Extending optimal transport theory to an arbitrary number of distributions has been studied on the theoretical side (Pass, 2015), and practical extensions have been proposed (Lin et al., 2022).
Revised: Extending optimal transport theory to an arbitrary number of distributions has been studied on the theoretical side (Pass, 2015), and practical extensions have been proposed (Lin et al., 2022).

Original: But implementations that integrate directly into machine learning pipelines rely on either heuristics or modifications that result in an approximation of the original
Revised: But implementations that integrate directly into machine learning pipelines rely on either heuristics or modifications that result in an approximation of the original

Original: MMOT problem (Cao et al., 2019).
Revised: MMOT problem (Cao et al., 2019).

Original: Wasserstein barycenters.
Revised: Wasserstein barycenters.

Original: One use case of our algorithms will be in formulations that involve Wasserstein barycenters.
Revised: One use case of our algorithms will be in formulations that involve Wasserstein barycenters.

Original: Given a set of probability distributions, the Wasserstein barycenter minimizes the mean of Wasserstein distances to each probability distribution in the set: a practical definition of the mean under the transportation distance (Luise et al., 2019; Cuturi & Doucet, 2014; Agueh & Carlier, 2011; Janati et al., 2020).
Revised: Given a set of probability distributions, the Wasserstein barycenter minimizes the mean of Wasserstein distances to each probability distribution in the set: a practical definition of the mean under the transportation distance (Luise et al., 2019; Cuturi & Doucet, 2014; Agueh & Carlier, 2011; Janati et al., 2020).

Original: Applications of Wasserstein barycenter include texture analysis (Rabin et al., 2011), sensor data fusion (Elvander et al., 2020), shape interpolation (Solomon et al., 2015), coupling problems (R¨uschendorf & Uckelmann, 2002) and others (Ho et al., 2017).
Revised: Applications of Wasserstein barycenter include texture analysis (Rabin et al., 2011), sensor data fusion (Elvander et al., 2020), shape interpolation (Solomon et al., 2015), coupling problems (R¨uschendorf & Uckelmann, 2002) and others (Ho et al., 2017).

Original: Very recently, polynomial time algorithms have been derived (Altschuler & Boix-Adsera, 2021).
Revised: Very recently, polynomial time algorithms have been derived (Altschuler & Boix-Adsera, 2021).

Original: Fairness.
Revised: Fairness.

Original: Proposals such as (Jiang et al., 2020;
Revised: Proposals such as (Jiang et al., 2020;

Original: Chzhen et al., 2020; Gordaliza et al., 2019) have all regularized models towards outcomes which have equal predictive power over subgroups within a population, measured using optimal transport distance.
Revised: Chzhen et al., 2020; Gordaliza et al., 2019; Lokhande et al., 2022b) have all regularized models towards outcomes which have equal predictive power over subgroups within a population, measured using optimal transport distance.

Original: Informally, the idea is to operate on distributions that are supported on disparate groups, and ask that the distributions get “pushed” towards a common central distribution.
Revised: Informally, the idea is to operate on distributions that are supported on disparate groups, and ask that the distributions get “pushed” towards a common central distribution.

Original: This requires solving an optimal transport problem.
Revised: This requires solving an optimal transport problem.

Revised: Kwegyir-Aggrey et al.

Original: Greedy algorithms and extensions.
Revised: Greedy algorithms and extensions.

Original: Hoffman et al.
Revised: Hoffman et al.

Original: (1963) observed that there exists a family of linear-time greedy algorithms that solve the classical two-dimensional transportation problem.
Revised: (1963) observed that there exists a family of linear-time greedy algorithms that solve the classical two-dimensional transportation problem.

Original: Later, Bein et al.
Revised: Later, Bein et al.

Original: (1995) extended the relevant definitions and the greedy algorithm to d -dimensional transportation problems.
Revised: (1995) extended the relevant definitions and the greedy algorithm to d -dimensional transportation problems.

Original: More recently, the results in Kline (2019), which we will use here, further extended this result to the dual program, and several theoretical properties of the generalized Earth
Revised: More recently, the results in Kline (2019), which we will use here, further extended this result to the dual program, and several theoretical properties of the generalized Earth

Original: Mover’s problem were shown.
Revised: Mover’s problem were shown.

Original: A slightly different generalized d -dimensional Earth Mover’s problem is explored in Erickson (2020), with a focus on statistical generalization properties.
Revised: A slightly different generalized d -dimensional Earth Mover’s problem is explored in Erickson (2020), with a focus on statistical generalization properties.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Denote by [ n ] := { 1 , . . . , n } , the set of positive integers no larger than n .
Revised: Denote by [ n ] := { 1 , . . . , n } , the set of positive integers no larger than n .

Original: For elements x ∈ R n , we denote the i th entry of x as x ( i ) , e.g., for any x ∈ R n , x = ( x ( i ) : i ∈ [ n ]) .
Revised: For elements x ∈ R n , we denote the i th entry of x as x ( i ) , e.g., for any x ∈ R n , x = ( x ( i ) : i ∈ [ n ]) .

Original: The positive orthant of R n is denoted R n + := { x ∈ R n : x ( i ) ≥ 0 , i ∈ [ n ] } .
Revised: The positive orthant of R n is denoted R n + := { x ∈ R n : x ( i ) ≥ 0 , i ∈ [ n ] } .

Original: We denote by e := (1 , . . .
Revised: We denote by e := (1 , . . .

Original: , 1) ∈ R n the constant vector.
Revised: , 1) ∈ R n the constant vector.

Original: For q ≥ 1 , we define the q -norm as ∥ x ∥
Revised: For q ≥ 1 , we define the q -norm as ∥ x ∥

Original: q
Revised: q

Original: =
Revised: =

Original: (cid:16)(cid:80)
Revised: (cid:16)(cid:80)

Original: i ∈
Revised: i ∈

Original: [ n ] | x ( i ) | q (cid:17) 1 /q , and if q is suppressed, then ∥ x ∥ := ∥ x ∥ 2 .
Revised: [ n ] | x ( i ) | q (cid:17) 1 /q , and if q is suppressed, then ∥ x ∥ := ∥ x ∥ 2 .

Original: A discrete probability distribution is a point p ∈ R n + with e ′ p =
Revised: A discrete probability distribution is a point p ∈ R n + with e ′ p =

Original: ∥
Revised: ∥

Original: p ∥ 1 = 1 .
Revised: p ∥ 1 = 1 .

Original: Given a pair of discrete probability distributions p 1 , p 2 ∈ R n + , we may want to quantify similarity.
Revised: Given a pair of discrete probability distributions p 1 , p 2 ∈ R n + , we may want to quantify similarity.

Original: Often, it is more convenient to quantify dissimilarity between p 1 and p 2 , and one may do this by selecting from many options including the q -norm, KL-divergence or the Earth Mover’s Distance (EMD).
Revised: Often, we do this by selecting from many measures, including the q -norm, KL-divergence or the Earth Mover’s Distance (EMD).

Original: The EMD for a pair of distributions has several equivalent interpretations.
Revised: The EMD for a pair of distributions has several equivalent interpretations.

Original: First, let p 1 be a source of mass, and p 2 be a sink for mass, and x ( i, j ) , where x ∈ R n × n , represent the flow of mass from p 1 ( i ) to p 2 ( j ) .
Revised: First, let p 1 be a source of mass, and p 2 be a sink for mass, and x ( i, j ) , where x ∈ R n × n , represent the flow of mass from p 1 ( i ) to p 2 ( j ) .

Original: Denote by c ( i, j ) the cost of moving one unit of mass from p 1 ( i ) to p 2 ( j ) .
Revised: Denote by c ( i, j ) the cost of moving one unit of mass from p 1 ( i ) to p 2 ( j ) .

Original: The EMD between p 1 and p 2 is the minimal cost to transform p 1 into p 2 , written as a linear program (LP):min x [Equation] R n ×
Revised: The EMD between p 1 and p 2 is the minimal cost to transform p 1 into p 2 , written as a linear program (LP):

Original: [Equation]
Revised: [Figure]

Original: + (cid:88)

Original: [Equation]
Revised: [Equation]

Original: The source-sink interpretation is asymmetric in p 1 and p 2 , but the LP is symmetric in p 1 and p 2 .
Revised: The source-sink interpretation is asymmetric in p 1 and p 2 , but the LP is symmetric in p 1 and p 2 .

Original: It can be shown that the objective value of this LP defines a metric (Kantorovich, 1960), and the optimal value of the objective function can be interpreted as a distance between p 1 and p 2 , and useful to quantify dissimilarity between pairs of distributions.
Revised: It can be shown that the objective value of this LP defines a metric (Kantorovich, 1960), and the optimal value of the objective function can be interpreted as a distance between p 1 and p 2 , and useful to quantify dissimilarity between pairs of distributions.

Original: In particular, p 1 = p 2 if and only if the optimal objective value of the Earth Mover’s problem vanishes.
Revised: In particular, p 1 = p 2 if and only if the optimal objective value of the Earth Mover’s problem vanishes.

Original: The LP in (1) has an equivalent dual LP,
Revised: The LP in (1) has an equivalent dual LP,

Original: [Equation]
Revised: [Equation]

Original: By strong duality, the optimal value of the primal program (1) equals the optimal value of the dual program (2).
Revised: By strong duality, the optimal value of the primal program (1) equals the optimal value of the dual program (2).

Original: Many practical relaxations have been proposed for (1), including entropic regularization (Cuturi, 2013).
Revised: Many practical relaxations have been proposed for (1), including entropic regularization (Cuturi, 2013).

Original: The EMD is available in  Python Optimal Transport (POT)  (Flamary et al., 2021).
Revised: Computation of the EMD is readily available, as in the Python Optimal Transport (POT) library (Flamary et al., 2021).

Original: 3. D ISCRETE M ULTI -M ARGINAL O PTIMAL T RANSPORT
Revised: 3. D ISCRETE M ULTI -M ARGINAL O PTIMAL T RANSPORT

Original: The foregoing approach applies only to d = 2 distributions, namely p 1 and p 2 .
Revised: The foregoing approach applies only to d = 2 distributions, namely p 1 and p 2 .

Original: We briefly review the extension to d > 2 distributions; the literature calls this multi-marginal optimal transport (MMOT) .
Revised: We briefly review the extension to d > 2 distributions; the literature calls this multi-marginal optimal transport (MMOT) .

Original: Definition 3.1 (Discrete Multi-Marginal Optimal Transport (d-MMOT)) .
Revised: Definition 3.1 (Discrete Multi-Marginal Optimal Transport (d-MMOT)) .

Original: Let p 1 , . .
Revised: Let p 1 , . .

Original: p d ∈ R n + be discrete probability distributions.
Revised: p d ∈ R n + be discrete probability distributions.

Original: Let C d : R n d → R + .
Revised: Let C d : R n d → R + .

Original: The discrete multi-marginal optimal transport problem (d-MMOT) can be written as
Revised: The discrete multi-marginal optimal transport problem (d-MMOT) can be written as

Original: [Equation]
Revised: [Equation]

Original:  ×···×
Revised: n ×···×

Original: [Equation]
Revised: [Equation]

Original: where X i ∈ R n is the i -th marginal of X ∈ R n ×···× n = R n d .
Revised: where X i ∈ R n is the i -th marginal of X ∈ R n ×···× n = R n d .

Original: Following the original formulation (Kantorovich, 1942), we will restrict the cost function C d ( · ) to the linear map, C d ( X ) :
Revised: Following the original formulation (Kantorovich, 1942), we will restrict the cost function C d ( · ) to the linear map, C d ( X ) :

Original: = ⟨ c, X ⟩ ⊗ , where c ∈ R n ×···× n + is nonnegative.
Revised: = ⟨ c, X ⟩ ⊗ , where c ∈ R n ×···× n + is nonnegative.

Original: Here, the d-MMOT is the LP, min x ∈ R nd
Revised: Here, the d-MMOT is the LP, min x ∈ R nd

Original: [Equation]
Revised: [Equation]

Original: This linear program (LP) is central to the regularization schemes that are discussed below.
Revised: This linear program (LP) is central to the regularization schemes that are discussed below.

Original:  E FFICIENT D -MMOT C OMPUTATION
Revised:  E FFICIENT D -MMOT C OMPUTATION

Original: The linear d-MMOT problem (3) suffers from the curse of dimensionality: the LP has n d variables, and even modest choices of n and d can result in a LP with billions of variables, making standard
Revised: The linear d-MMOT problem (3) suffers from the curse of dimensionality: the LP has n d variables, and even modest choices of n and d can result in a LP with billions of variables, making standard

Original: LP solvers inapplicable.
Revised: LP solvers inapplicable.

Original: Alternatively, specific algorithms have been proposed (Benamou et al., 2015), and relaxations via entropic regularization have become more widespread, with very recent extensions to the d-MMOT setting (Tupitsa et al., 2020; Lin et al., 2022).
Revised: Alternatively, specific algorithms have been proposed (Benamou et al., 2015), and relaxations via entropic regularization have become more widespread, with very recent extensions to the d-MMOT setting (Tupitsa et al., 2020; Lin et al., 2022).

Original: In practice, the cost c in (3) typically takes one of two forms.
Revised: In practice, the cost c in (3) typically takes one of two forms.

Original: In the case where the distributions p 1 , . . .
Revised: In the case where the distributions p 1 , . . .

Original: , p d are over categorical variables, the cost is typically defined as c ( i 1 , . .
Revised: , p d are over categorical variables, the cost is typically defined as c ( i 1 , . .

Original: , i d ) = 0 when i 1 = · · · = i d and 1 otherwise.
Revised: , i d ) = 0 when i 1 = · · · = i d and 1 otherwise.

Original: However, and importantly, if the distributions p i are histograms over some ordinal or discretized space, the cost typically has a structure closer to that of a “tensorized” distance, characterized by the Monge property.
Revised: However, and importantly, if the distributions p i are histograms over some ordinal or discretized space, the cost typically has a structure closer to that of a “tensorized” distance, characterized by the Monge property.

Original: Definition 4.1 (Monge Property) .
Revised: Definition 4.1 (Monge Property) .

Original: A tensor c is Monge if for all valid i 1 , . . .
Revised: A tensor c is Monge if for all valid i 1 , . . .

Original: i d and j 1 , . .
Revised: i d and j 1 , . .

Original: , j d ,
Revised: , j d ,

Original: [Equation]
Revised: [Equation]

Original: 4) where s k = min( i k , j k ) and t k = max( i k , j k ) .
Revised: where s k = min( i k , j k ) and t k = max( i k , j k ) .

Original: Our focus is on a specific cost, which is known to be Monge: c ( i 1 , i 2 , . .
Revised: Our focus is on a specific cost, which is known to be Monge: c ( i 1 , i 2 , . .

Original: , i d ) :
Revised: , i d ) :

Original: = max { i k : k ∈ [ d ] }− min { i k : k ∈ [ d ] } .
Revised: = max { i k : k ∈ [ d ] }− min { i k : k ∈ [ d ] } .

Original: When d = 2 , this cost reduces to c ( i 1 , i 2 ) =
Revised: When d = 2 , this cost reduces to c ( i 1 , i 2 ) =

Original: | i 1 − i 2 | , which agrees with the classical EMD cost.
Revised: | i 1 − i 2 | , which agrees with the classical EMD cost.

Original: This choice of c is called the generalized EMD cost .
Revised: This choice of c is called the generalized EMD cost .

Original: Remark 4.2 .
Revised: Remark 4.2 .

Original: Limiting our attention to this cost is not as restrictive as it may appear.
Revised: Limiting our attention to this cost is not as restrictive as it may appear.

Original: Indeed, Bein et al.
Revised: Indeed, Bein et al.

Original: (1995) shows that the optimal solution to the LP (3) is independent of the cost, as long as it is
Revised: (1995) shows that the optimal solution to the LP (3) is independent of the cost, as long as it is

Original: Monge.
Revised: Monge.

Original: Additionally, when c is the generalized EMD cost, Kline (2019) describes a greedy algorithm that solves both (3) and its dual (5) in linear time.
Revised: Additionally, when c is the generalized EMD cost, Kline (2019) describes a greedy algorithm that solves both (3) and its dual (5) in linear time.

Original: It is also shown that the optimal objective value is a continuous nonnegative function of each probability distribution p j (i.e., small changes in one distribution cause small changes in the objective value).
Revised: It is also shown that the optimal objective value is a continuous nonnegative function of each probability distribution p j (i.e., small changes in one distribution cause small changes in the objective value).

Original: Continuity is critical for numerical stability.
Revised: Continuity is critical for numerical stability.

Original: Next, when c is the generalized Earth Mover’s array, the optimal objective value vanishes if and only if p i = p j for all i, j ∈ [ n ] .
Revised: Next, when c is the generalized Earth Mover’s array, the optimal objective value vanishes if and only if p i = p j for all i, j ∈ [ n ] .

Original: This separability property is useful in applications where we wish to iteratively “step towards” the barycenter of a set of distributions, see Fig.
Revised: This separability property is useful in applications where we wish to iteratively “step towards” the barycenter of a set of distributions, see Fig.

Original: In order to step towards (red arrows) a barycenter, we require a descent direction.
Revised: In order to step towards (red arrows) a barycenter, we require a descent direction.

Original: We observe that the following result gives us this functionality.
Revised: We observe that the following result gives us this functionality.

Original: Theorem 4.3.
Revised: Theorem 4.3.

Original: The dual linear program of the d-MMOT problem (3) is
Revised: The dual linear program of the d-MMOT problem (3) is

Original: [Equation]
Revised: [Equation]

Original: where the indices in the constraints include all i j ∈ [ n ] , j ∈
Revised: where the indices in the constraints include all i j ∈ [ n ] , j ∈

Original: [ d ] .
Revised: [ d ] .

Original: Denote by ϕ ( p 1 , . .
Revised: Denote by ϕ ( p 1 , . .

Original: . , p d ) , the optimal objective value of the LP in (3).
Revised: . , p d ) , the optimal objective value of the LP in (3).

Original: Let z ∗ be an optimal solution to the dual program (5).
Revised: Let z ∗ be an optimal solution to the dual program (5).

Original: Then, ∇ ϕ ( p 1 , . . .
Revised: Then, ∇ ϕ ( p 1 , . . .

Original: , p d ) =
Revised: , p d ) =

Original: z ∗ , and for any t ∈ R , ϕ ( p 1 , p 2 , . .
Revised: z ∗ , and for any t ∈ R , ϕ ( p 1 , p 2 , . .

Original: . , p d ) = (cid:88)
Revised: . , p d ) = (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where η :=
Revised: where η :=

Original: ( z ∗ 1 ( n ) e, z ∗ 1 ( n ) e, · · · , z ∗ d ( n ) e ) .
Revised: ( z ∗ 1 ( n ) e, z ∗ 1 ( n ) e, · · · , z ∗ d ( n ) e ) .

Original: Proof.
Revised: Proof.

Original: The main observation invokes perturbation analysis (Mangasarian & Meyer, 1979;
Revised: The main observation invokes perturbation analysis (Mangasarian & Meyer, 1979;

Original: Ferris &
Revised: Ferris &

Original: Mangasarian, 1991) of LPs to assert that, under mild uniqueness conditions, small changes to a LP’s input data does not change its optimal solution.
Revised: Mangasarian, 1991) of LPs to assert that, under mild uniqueness conditions, small changes to a LP’s input data does not change its optimal solution.

Original: The full proof is in Appendix 7.1.
Revised: The full proof is in Appendix 7.1.

Original: Remark 4.4 .
Revised: Remark 4.4 .

Original: The first part of this result provides what we require: a direction of descent.
Revised: The first part of this result provides what we require: a direction of descent.

Original: Thus, if we can solve the d-MMOT problem and also find the optimal solution to its dual, then we can step (or move) our distributions in the opposite direction of the dual variables to push them together, see
Revised: Thus, if we can solve the d-MMOT problem and also find the optimal solution to its dual, then we can step (or move) our distributions in the opposite direction of the dual variables to push them together, see

Original: Fig.
Revised: Fig.

Original: The second claim is somewhat technical, and reconciles particular affine shifts that result in equivalent objective values.
Revised: The second claim is somewhat technical, and reconciles particular affine shifts that result in equivalent objective values.

Original: 4. O PTIMIZATION OF D -MMOT
Revised: 4. O PTIMIZATION OF D -MMOT

Original: Setup.
Revised: Setup.

Original: We can now instantiate d-MMOT as an add-on term in a standard machine learning formulation.
Revised: We can now instantiate d-MMOT as an add-on term in a standard machine learning formulation.

Original: Concretely, it can be positioned alongside typical learning losses L ( f ( x ) , y ; θ ) to encourage minimizing distances among d different distributions g i ∈ G, i ∈
Revised: Concretely, it can be positioned alongside typical learning losses L ( f ( x ) , y ; θ ) to encourage minimizing distances among d different distributions g i ∈ G, i ∈

Original: [ d ] , d-MMOT ( f ( x ) , g ; θ ) , i.e., min θ L ( f ( x ) , y ; θ ) + d-MMOT ( f ( x ) , g ; θ ) .
Revised: [ d ] , d-MMOT ( f ( x ) , g ; θ ) , i.e., min θ L ( f ( x ) , y ; θ ) + d-MMOT ( f ( x ) , g ; θ ) .

Original: Within a deep neural network (DNN) architecture, as in some of our experiments, several properties of the d-MMOT module are useful: our results above naturally provide clean operations for computing both the required objective in the “forward” pass and gradients in the “backward” pass.
Revised: Within a deep neural network (DNN) architecture, as in some of our experiments, several properties of the d-MMOT module are useful: our results above naturally provide clean operations for computing both the required objective in the “forward” pass and gradients in the “backward” pass.

Original: Using the primal and dual variables.
Revised: Using the primal and dual variables.

Original: If the optimal objective value of (3) serves in regularizing a deep neural network, then we can train the network as follows.
Revised: If the optimal objective value of (3) serves in regularizing a deep neural network, then we can train the network as follows.

Original: During the forward pass, i.e., computing the d-MMOT objective, we can employ a version of the aforementioned primal/dual algorithm that solely computes the function ϕ and stores the dual variables z .
Revised: During the forward pass, i.e., computing the d-MMOT objective, we can employ a version of the aforementioned primal/dual algorithm that solely computes the function ϕ and stores the dual variables z .

Original: As backpropagation proceeds, when the EMD module encounters an incoming gradient, it is simply multiplied by the stored dual variables (see Theorem 4.3 and Algorithm 1).
Revised: As backpropagation proceeds, when the EMD module encounters an incoming gradient, it is simply multiplied by the stored dual variables (see Theorem 4.3 and Algorithm 1).

Original: We call our procedure the d − dimensional
Revised: We call our procedure the d − dimensional

Original: Earth Mover’s Distance, or in short, the DEMD algorithm .
Revised: Earth Mover’s Distance, or in short, the DEMD algorithm .

Original: Complexity.
Revised: Complexity.

Original: Computing the DEMD distance in the forward pass is exactly O ( nd ) : linear in the number of distributions and number of bins.
Revised: Computing the DEMD distance in the forward pass is exactly O ( nd ) : linear in the number of distributions and number of bins.

Original: This property follows directly from the algorithm, needing only a single pass through all of the data.
Revised: This property follows directly from the algorithm, needing only a single pass through all of the data.

Original: Bein et al.
Revised: Bein et al.

Original: (1995) provides this result for greedy algorithms that solve OT programs as in Def. 3.1.
Revised: (1995) provides this result for greedy algorithms that solve OT programs as in Def. 3.1.

Original: In contrast to methods that derive gradients
Revised: In contrast to methods that derive gradients

Original: Algorithm 1 Our proposed method: d  Distance (DEMD) function F ORWARD ( p 1 , . .
Revised: Algorithm 1 Our proposed method: d − Dimensional Earch Mover’s Distance (DEMD) function F ORWARD ( p 1 , . .

Original: . , p j )
Revised: . , p j )

Original: Compute (cid:80) s k t k and ( z 1 , . .
Revised: Compute (cid:80) s k t k and ( z 1 , . .

Original: . , z d ) via the primal/dual algorithm in Kline (2019).
Revised: . , z d ) via the primal/dual algorithm in Kline (2019).

Original: save ( z 1 , . . .
Revised: save ( z 1 , . . .

Original: , z d ) .
Revised: , z d ) .

Original: return (cid:80) k
Revised: return (cid:80) k

Original: s k t k function B ACKWARD (gradOutput) load ( z 1 , . .
Revised: s k t k function B ACKWARD (gradOutput) load ( z 1 , . .

Original: . , z d ) return ( z1 , . .
Revised: . , z d ) return ( z

Original: . , zd ) · gradOutput via entropic regularization schemes, i.e., relaxations of the optimal transport problem (Luise et al., 2018;
Revised: , z via entropic regularization schemes, i.e., relaxations of the optimal transport problem (Luise et al., 2018;

Original: Xie et al., 2020; Cuturi et al., 2020), this approach solves the distance computation exactly in linear time.
Revised: Xie et al., 2020; Cuturi et al., 2020), this approach solves the distance computation exactly in linear time.

Original: This linear time analysis is not only provided by the theory in prior work, but is also explicit in the number of iterations defining our algorithm (see Appendix 7.4 for more details).
Revised: This linear time analysis is not only provided by the theory in prior work, but is also explicit in the number of iterations defining our algorithm (see Appendix 7.4 for more details).

Original: For minimizing the DEMD (computing the updates, i.e., red arrows in Fig. 2), a convergence analysis would follow from the properties of the optimization scheme chosen.
Revised: For minimizing the DEMD (computing the updates, i.e., red arrows in Fig. 2), a convergence analysis would follow from the properties of the optimization scheme chosen.

Original: Our tool can be dropped in exactly as any other module in modern learning applications (using the observation that gradients are easily computed, i.e., O (1) time to read stored dual variables).
Revised: Our tool can be dropped in exactly as any other module in modern learning applications (using the observation that gradients are easily computed, i.e., O (1) time to read stored dual variables).

Original: A few practical adjustments.
Revised: A few practical adjustments.

Original: It often happens during training that the optimal solutions may, through updates by stepping in the direction of the gradient, acquire entries that are negative.
Revised: It often happens during training that the optimal solutions may, through updates by stepping in the direction of the gradient, acquire entries that are negative.

Original: This violates an assumption that entries must be nonnegative.
Revised: This violates an assumption that entries must be nonnegative.

Original: However, Thm. 2.2 in Kline (2019) shows that the optimal objective value, ϕ , possesses a type of translation invariance.
Revised: However, Thm. 2.2 in Kline (2019) shows that the optimal objective value, ϕ , possesses a type of translation invariance.

Original: We can leverage this result to ensure that, in the event that a point escapes the nonnegative orthant of R n , an appropriately constructed constant vector may be added to the current iterate so that it again lies in the nonnegative orthant, without changing the objective value .
Revised: Leveraging this alongside the second part of Thm 4.3, we ensure that, in the event that a point escapes the nonnegative orthant of R n , an appropriately constructed constant vector may be added to the current iterate so that it again lies in the nonnegative orthant, without changing the objective value .

Original: Further, with a relatively small learning rate, by the continuity of ϕ , normalization is a small perturbation of the original point and can be applied as necessary to enforce that updates result in valid distributions.
Revised: Further, with a relatively small learning rate, by the continuity of ϕ , normalization is a small perturbation of the original point and can be applied as necessary to enforce that updates result in valid distributions.

Original: Remark 4.5 .
Revised: Remark 4.5 .

Original: Another useful property is that if the convex hull of ( p 1 , . . .
Revised: Another useful property is that if the convex hull of ( p 1 , . . .

Original: , p d ) is contained in the convex hull of (ˆ p 1 , . . .
Revised: , p d ) is contained in the convex hull of (ˆ p 1 , . . .

Original: , ˆ p ˆ d ) , then ϕ ( p 1 , . . .
Revised: , ˆ p ˆ d ) , then ϕ ( p 1 , . . .

Original: , p d ) ≤ ϕ (ˆ p 1 , . . .
Revised: , p d ) ≤ ϕ (ˆ p 1 , . . .

Original: , ˆ p ˆ d ) , with strict inequality if containment is strict.
Revised: , ˆ p ˆ d ) , with strict inequality if containment is strict.

Original: A direct consequence of this property is that points within the interior of the convex hull of the data fed to the optimization model have vanishing gradients.
Revised: A direct consequence of this property is that points within the interior of the convex hull of the data fed to the optimization model have vanishing gradients.

Original: Practically, this leads to sparse gradients w.r.t.
Revised: Practically, this leads to sparse gradients w.r.t.

Original: the distributions, and nonzero gradients correspond to points on the hull, i.e., distributions which are maximally different from the rest.
Revised: the distributions, and nonzero gradients correspond to points on the hull, i.e., distributions which are maximally different from the rest.

Original: Minimization proceeds by iteratively moving mass such that these maximally different distributions are pushed toward each other.
Revised: Minimization proceeds by iteratively moving mass such that these maximally different distributions are pushed toward each other.

Original: Fig.
Revised: Fig.

Original: 2 shows the convex hull during minimization of the generalized EMD objective via our DEMD algorithm.
Revised: 2 shows the convex hull during minimization of the generalized EMD objective via our DEMD algorithm.

Original: Need for Histogramming.
Revised: Need for Histogramming.

Original: Note that outputs f ( x ) and intermediate activations are continuous values (layer shape by batch size).
Revised: Note that outputs f ( x ) and intermediate activations are continuous values (layer shape by batch size).

Original: So, we must transform activations into normalized histograms (i.e., discrete distributions).
Revised: So, we must transform activations into normalized histograms (i.e., discrete distributions).

Original: While libraries such as PyTorch and Tensorflow typically provide histogram functions, they are not differentiable.
Revised: While libraries such as PyTorch and Tensorflow typically provide histogram functions, they are not differentiable.

Original: The discontinuous operation of bin assignment does not allow for an end-to-end pipeline where we can push upstream parameters in the direction of minimizing the
Revised: The discontinuous operation of bin assignment does not allow for an end-to-end pipeline where we can push upstream parameters in the direction of minimizing the

Original: EMD objective over histograms.
Revised: EMD objective over histograms.

Original: To address this, we use a simple relaxed/differentiable histogramming operation.
Revised: To address this, we use a simple relaxed/differentiable histogramming operation.

Original: First, all outputs are mapped to [0 , 1] using a Sigmoid activation.
Revised: First, all outputs are mapped to [0 , 1] using a Sigmoid activation.

Original: Then, for each bin location, the “count” in that bin is determined by a ReLU function applied to the difference between the gap between the activation and the distance to the bin boundary.
Revised: Then, for each bin location, the “count” in that bin is determined by a ReLU function applied to the difference between the gap between the activation and the distance to the bin boundary.

Original: In other words, if an activation falls in a bin, the count of that bin increases based on the distance to the bin boundary, otherwise the count remains the same.
Revised: In other words, if an activation falls in a bin, the count of that bin increases based on the distance to the bin boundary, otherwise the count remains the same.

Original: The full procedure is detailed in Appendix 7.2.
Revised: The full procedure is detailed in Appendix 7.2.

Original: The ReLU activations defining bin boundaries allow for gradients to move samples towards neighboring bins as needed.
Revised: The ReLU activations defining bin boundaries allow for gradients to move samples towards neighboring bins as needed.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate our construction in a number of settings.
Revised: We evaluate our construction in a number of settings.

Original: First, we demonstrate the computational speedup associated with evaluating d-MMOT using our algorithm, along with speedups associated with directly computing the gradient.
Revised: First, we demonstrate the computational speedup associated with evaluating d-MMOT using our algorithm, along with speedups associated with directly computing the gradient.

Original: Next, we compute the construction in a series of neural network tasks associated with ensuring distributional similarity: fairness, invariant representations, and multi-domain matching.
Revised: Next, we compute the construction in a series of neural network tasks associated with ensuring distributional similarity: fairness, invariant representations, and multi-domain matching.

Original: We provide a complete PyTorch Network Module that packages the above
Revised: We provide a complete PyTorch Network Module that packages the above

Original: [Figure]
Revised: [Figure]

Original: 5. P ERFORMANCE B ENCHMARKS
Revised: 5. P ERFORMANCE B ENCHMARKS

Original: Figure 4 presents NumPy and Torch instantiations of both forward and backward passes using automatic differentiation and the gradients computed using the dual as in §4.1.
Revised: Figure 4 presents NumPy and Torch instantiations of both forward and backward passes using automatic differentiation and the gradients computed using the dual as in §4.1.

Original: As expected, the forward (distance computation) times are comparable[Figure] but the backwards computation scales poorly with the number of distributions to be updated using automatic differentiation.
Revised: As expected, the forward (distance computation) times are comparable but the backwards computation scales poorly with the number of distributions to be updated using automatic differentiation.

Original: Our dual setting allows the gradients to simply be read off (based on the forward pass), leading to no additional computation overhead during backpropagation
Revised: Our dual setting allows the gradients to simply be read off (based on the forward pass), leading to no additional computation overhead during backpropagation

Original: [Figure]
Revised: [Figure]

Original:  times, where automatic differentiation scales poorly with the number of distributions.
Revised: +AutoGrad). Direct reading of the gradient via the dual leads to significant gains in backward pass times, where automatic differentiation scales poorly with the number of distributions.

Original: We compare our DEMD computation to what one may use given existing Optimal Transport methods in Figure 3 (right).
Revised: We compare our DEMD computation to what one may use given existing Optimal Transport methods in Figure 3 (right).

Original: Using standard off-the-shelf methods as a baseline, when the cost is Monge our algorithm provides significant speedups in computation time, on the scale of orders of magnitude!
Revised: Using standard off-the-shelf methods as a baseline, when the cost is Monge our algorithm provides significant speedups in computation time, on the scale of orders of magnitude!

Original: Further, if the number of distributions and bins increases, the time-cost for existing methods using more generic LP solvers can become significant, and may become infeasible with generic solutions via
Revised: Further, if the number of distributions and bins increases, the time-cost for existing methods using more generic LP solvers can become significant, and may become infeasible with generic solutions via

Original: CVXPY (Diamond & Boyd, 2016).
Revised: CVXPY (Diamond & Boyd, 2016).

Original: We compare our approach to two off-the-shelf implementations of barycenters via the Python Optimal Transport (POT) Library (Flamary et al., 2021), along with directly solving the Earth Mover’s problem via CVXPY.
Revised: We compare our approach to two off-the-shelf implementations of barycenters via the Python Optimal Transport (POT) Library (Flamary et al., 2021), along with directly solving the Earth Mover’s problem via CVXPY.

Original: Even for 10 distributions over 10 bins,
Revised: Even for 10 distributions over 10 bins,

Original: CVXPY is unable to allocate the necessary memory using a direct instantiation.
Revised: CVXPY is unable to allocate the necessary memory using a direct instantiation.

Original: 5. G ENERALIZED EM F AIRNESS ON F AIRNESS D ATASETS
Revised: 5. G ENERALIZED EM F AIRNESS ON F AIRNESS D ATASETS

Original: With a viable tool in hand, we move to practical applications in machine learning fairness, which naturally requires enforcing closeness in model outputs.
Revised: With a viable tool in hand, we move to practical applications in machine learning fairness, which naturally requires enforcing closeness in model outputs.

Original: Here, we construct networks with our
Revised: Here, we construct networks with our

Original: DEMD regularizer, where we discretize the final activation output prior to classification, and push the distributions of this activation to be similar among sensitive attributes.
Revised: DEMD regularizer, where we discretize the final activation output prior to classification, and push the distributions of this activation to be similar among sensitive attributes.

Original: Data.
Revised: Data.

Original: We identify 4 common fairness datasets often used to benchmark fair machine learning algorithms: (1) the German Credit Dataset (Hofmann, 1994), (2) the Adult Income Dataset (Dua &
Revised: We identify 4 common fairness datasets often used to benchmark fair machine learning algorithms: (1) the German Credit Dataset (Hofmann, 1994), (2) the Adult Income Dataset (Dua &

Original: Graff, 2017), (3) the Communities and Crime Dataset (Redmond, 2009), and (4) The ACS Income dataset, recently made available as a large, population-level demographic dataset (Ding et al., 2021)
Revised: Graff, 2017), (3) the Communities and Crime Dataset (Redmond, 2009), and (4) The ACS Income dataset, recently made available as a large, population-level demographic dataset (Ding et al., 2021)

Original: [Table]
Revised: [Table]

Original: containing as many as nine sensitive attributes.
Revised: containing as many as nine sensitive attributes.

Original: We set up a simple three-layer neural network for classification tasks with the addition of a fairness-type regularizer.
Revised: We set up a simple three-layer neural network for classification tasks with the addition of a fairness-type regularizer.

Original: We compare our construction with 4 off-the-shelf plug-in regularizers: (1) No regularization, (2) Demographic Parity (DP), (3)
Revised: We compare our construction with 4 off-the-shelf plug-in regularizers: (1) No regularization, (2) Demographic Parity (DP), (3)

Original: Equalized Odds (EO), and (4) a histogrammed barycenter construction.
Revised: Equalized Odds (EO), and (4) a histogrammed barycenter construction.

Original: DP and EO regularizers were computed using a PyTorch version of FairLearn (Bird et al., 2020), and the barycenter version was implemented using POT library (with GPU backend).
Revised: DP and EO regularizers were computed using a PyTorch version of FairLearn (Bird et al., 2020), and the barycenter version was implemented using POT library (with GPU backend).

Original: Because the scale of the regularization term is not directly comparable, we sweep regularization weights and select the best over all measures for a each dataset/regularizer pair.
Revised: Because the scale of the regularization term is not directly comparable, we sweep regularization weights and select the best over all measures for a each dataset/regularizer pair.

Original: We use 10 bins and replicate all experiments over three random seeds.
Revised: We use 10 bins and replicate all experiments over three random seeds.

Original: Models.
Revised: Models.

Original: We set up two model settings with a standard logistic regressor and a 2-layer neural network.
Revised: We set up two model settings with a standard logistic regressor and a 2-layer neural network.

Original: We compare three types of plug-in regularizers: (1) Demographic Parity (DP), (2) Equalized Odds (EO), and (3) the Generalized EMD.
Revised: We compare three types of plug-in regularizers: (1) Demographic Parity (DP), (2) Equalized Odds (EO), and (3) the Generalized EMD.

Original: DP and EO regularizers were computed using a PyTorch implementation of FairLearn (Bird et al., 2020).
Revised: DP and EO regularizers were computed using a PyTorch implementation of FairLearn (Bird et al., 2020).

Original: Results.
Revised: Results.

Original: In the summary in Table 1, models were selected with the largest regularization weight before accuracy dropped significantly.
Revised: In the summary in Table 1, models were selected with the largest regularization weight before accuracy dropped significantly.

Original: When accuracies are comparable, we see good performance (when minimizing DEMD regularizer) against baseline methods.
Revised: When accuracies are comparable, we see good performance (when minimizing DEMD regularizer) against baseline methods.

Original: Notably, when accuracies are similar across methods, minimizing DEMD tends to give better (lower) fairness measures across all datasets.
Revised: Notably, when accuracies are similar across methods, minimizing DEMD tends to give better (lower) fairness measures across all datasets.

Original: Summary: DEMD on the final network layer helps control multiple fairness measures.
Revised: Summary: DEMD on the final network layer helps control multiple fairness measures.

Original: 5. H ARMONIZATION FOR I NVARIANT R EPRESENTATIONS
Revised: 5. H ARMONIZATION FOR I NVARIANT R EPRESENTATIONS

Original: Having evaluated the use of the DEMD layer to constrain the neural network for fairness measures, we will now move to a more general problem of deriving invariant representations from the datasets.
Revised: Having evaluated the use of the DEMD layer to constrain the neural network for fairness measures, we will now move to a more general problem of deriving invariant representations from the datasets.

Original: Here, invariance is sought in regard to the sensitive attributes.
Revised: Here, invariance is sought w.r.t. the sensitive attributes.

Original: Recent works (Lokhande et al., 2022) on invariant representation learning propose leveraging an encoder-decoder architectures to map the dataset features to latent space representations.
Revised: Recent works (Lokhande et al., 2022a; Akash et al., 2021) on invariant representation learning propose leveraging an encoder-decoder architectures to map the dataset features to latent space representations.

Original: The latent space representations are penalized to match or harmonize the distributions across several groups in the dataset.
Revised: The latent space representations are penalized to match or harmonize the distributions across several groups in the dataset.

Original: In contrast to the previous section, the goal here is to identify a good mapping in the latent space that is devoid of any group-related information in the dataset.
Revised: In contrast to the previous section, the goal here is to identify a good mapping in the latent space that is devoid of any group-related information in the dataset.

Original: Consequently, our evaluation metrics test if the latent representations have a lower value of (i) ADV, adversarial measure and (ii) MMD, maximum mean discrepancy measure.
Revised: Consequently, our evaluation metrics test if the latent representations have a lower value of (i) ADV, adversarial measure and (ii) MMD, maximum mean discrepancy measure.

Original: The measure ADV tests to what extent can a separate neural network predict the group information from the latent features.
Revised: The measure ADV tests to what extent can a separate neural network predict the group information from the latent features.

Original: Alternatively, the MMD scores measure the distance between the probability distributions across groups.
Revised: Alternatively, the MMD scores measure the distance between the probability distributions across groups.

Original: Prior works such as Li et al.
Revised: Prior works such as Li et al.

Original: (2014) and Xie et al.
Revised: (2014) and Xie et al.

Original: (2017) optimize each of these measures separately.
Revised: (2017) optimize each of these measures separately.

Original: Our experiments (Table 2) show the DEMD layer performs competitively with the baselines when applied on the latent space.
Revised: Our experiments (Table 2) show the DEMD layer performs competitively with the baselines when applied on the latent space.

Original: Interestingly, these performance gains come despite not directly optimizing the harmonization measures, in contrast to baselines, which require several practical adjustments (batch variants and secondary neural networks).
Revised: Interestingly, these performance gains come despite not directly optimizing the harmonization measures, in contrast to baselines, which require several practical adjustments (batch variants and secondary neural networks).

Original: Summary: DEMD as an intermediate layer can be used to derive invariant representations efficiently.
Revised: Summary: DEMD as an intermediate layer can be used to derive invariant representations efficiently.

Original: 5. M ULTI - DOMAIN I MAGE T RANSLATION
Revised: 5. M ULTI - DOMAIN I MAGE T RANSLATION

Original: We apply our construction to a recent multi-marginal GAN framing of multi-domain image translation.
Revised: We apply our construction to a recent multi-marginal GAN framing of multi-domain image translation.

Original: With the goal of learning a mapping for a source image to multiple target domains, a focal point of recent literature has been to reduce the computational requirements of training individual models for each target, and learn a concurrent matching problem over a number of shared parameters or[Table] networks.
Revised: With the goal of learning a mapping for a source image to multiple target domains, a focal point of recent literature has been to reduce the computational requirements of training individual models for each target, and learn a concurrent matching problem over a number of shared parameters or networks.

Original: MWGAN (Cao et al., 2019) set up a Multiple Marginal Matching problem in this context.
Revised: MWGAN (Cao et al., 2019) set up a Multiple Marginal Matching problem in this context.

Original: Extending upon their key observation that inter-domain constraints can be measured via the gradients  for each domain, we instantiate our DEMD layer over the gradient norms computed for each sample in a batch per group.
Revised: Extending upon their key observation that inter-domain constraints can be measured via the gradients [Table] for each domain, we instantiate our DEMD layer over the gradient norms computed for each sample in a batch per group.

Original: Specifically, DEMD minimizes the differences between the distributions over the gradient norms across all target domains.
Revised: Specifically, DEMD minimizes the differences between the distributions over the gradient norms across all target domains.

Original: Using a weight of 100 for both regularizers, we observe similar performance when compared to the original MWGAN construction.
Revised: Using a weight of 100 for both regularizers, we observe similar performance when compared to the original MWGAN construction.

Original: In Figure 5 we show a few samples generated over an image translation task on the CelebA dataset.
Revised: In Figure 5 we show a few samples generated over an image translation task on the CelebA dataset.

Original: Here, we aim to translate original dataset images to ones with specific target attributes.
Revised: Here, we aim to translate original dataset images to ones with specific target attributes.

Original: Summary: Multi-domain image translation can be conducted by adding a DEMD layer in the native GAN construction.
Revised: Summary: Multi-domain image translation can be conducted by adding a DEMD layer in the native GAN construction.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We presented an efficient solution for solving common practical multi-marginal optimal transport problems.
Revised: We presented an efficient solution for solving common practical multi-marginal optimal transport problems.

Original: Our construction is significantly cheaper to compute compared to similar methods, and allows for large numbers of distributions to be matched in common DNN pipelines.
Revised: Significantly cheaper to compute compared to similar methods, it allows for large numbers of distributions to be matched in common DNN pipelines.

Original: Our implementation allows imposing fairness constraints for a variety of applications,including those with many groups, without the need for pairwise measures.
Revised: Our implementation allows imposing fairness constraints for a variety of applications, without the need for pairwise measures.

Original: As such, subgroup fairness (Kearns et al., 2018) is an interesting problem setting that we believe can benefit.
Revised: As such, subgroup fairness (Kearns et al., 2018) is an interesting problem setting that we believe can benefit.

Original: Other properties such as Minkowski additivity that have not been explicitly leveraged in our experiments may also be a worthwhile direction to explore.
Revised: Other properties such as Minkowski additivity that have not been explicitly leveraged in our experiments may also be a worthwhile direction to explore.

Original: One limitation of DEMD is its inability to be directly applied to distributions over multi-dimensional discrete spaces, such as  latent spacescommon in generative models.
Revised: In its current form, DEMD cannot be directly applied to distributions over multi-dimensional discrete spaces, such as images or latent spaces in generative models.

Original: Slicing is a heuristic that has been shown to work well.
Revised: Slicing is a heuristic that has been shown to work well.

Original: To evaluate feasibility, we embed distributions over multidimensional continuous spaces, take random projections over 1-D spaces, and recompute our DEMD measure.
Revised: To evaluate feasibility, we embed distributions over multidimensional continuous spaces, take random projections over 1-D spaces, and recompute our DEMD measure.

Original: Our gains in the many-distribution setting extend here as well: over a 64-dimensional latent space embedding of CelebA, we can efficiently compute our DEMD measure over all 40 attribute subgroups, and observe convergent behavior w.r.t.
Revised: Over a 64-dimensional latent space embedding of CelebA, we can efficiently compute our DEMD measure over all 40 attribute subgroups, and observe convergent behavior w.r.t.

Original: the number of projections.
Revised: the number of projections.

Original:  A PPENDIX 7. P ROOF OF T HEOREM 4. 
Revised:  A PPENDIX 7. P ROOF OF T HEOREM 4.

Original: Here, we present the proof of Theorem 4.2.
Revised: Here, we present the proof of Theorem 4.2.

Original: The main observation in the proof invokes perturbation analysis Mangasarian & Meyer (1979); Ferris & Mangasarian (1991) of linear programs to assert that, under mild uniqueness conditions, small changes to a linear program’s data does not change the linear program’s optimal solution.
Revised: The main observation in the proof invokes perturbation analysis Mangasarian & Meyer (1979); Ferris & Mangasarian (1991) of linear programs to assert that, under mild uniqueness conditions, small changes to a linear program’s data does not change the linear program’s optimal solution.

Original: The informal reason why such a result is possible can be explained using a short, geometric argument as follows.
Revised: The informal reason why such a result is possible can be explained using a short, geometric argument as follows.

Original: The feasible set of a nontrivial linear program is a polytope, and, as a rule, an optimal solution to a linear program lives at the point where a hyperplane defined by the objective functional intersects a vertex of the polytope.
Revised: The feasible set of a nontrivial linear program is a polytope, and, as a rule, an optimal solution to a linear program lives at the point where a hyperplane defined by the objective functional intersects a vertex of the polytope.

Original: A small perturbation of the hyperplane does not alter the intersecting vertex.
Revised: A small perturbation of the hyperplane does not alter the intersecting vertex.

Original: The derivative of a linear program’s optimal objective value, treated as a function of the data fed to the LP, has been described several times in prior work, and under a variety of conditions De Wolf &
Revised: The derivative of a linear program’s optimal objective value, treated as a function of the data fed to the LP, has been described several times in prior work, and under a variety of conditions De Wolf &

Original: Smeers (2019); Freund (1985); Agueh & Carlier (2011); Mills (1957).
Revised: Smeers (2019); Freund (1985); Agueh & Carlier (2011); Mills (1957).

Original: The second claim of Theorem 4.2 is useful because dual solutions to the generalized EMD linear program are not unique.
Revised: The second claim of Theorem 4.2 is useful because dual solutions to the generalized EMD linear program are not unique.

Original: The claim explains how one can modify a given solution (found by a direct gradient computation or else from an interior point solver, for example) so that it agrees with the solution yielded by the primal/dual greedy algorithm.
Revised: The claim explains how one can modify a given solution (found by a direct gradient computation or else from an interior point solver, for example) so that it agrees with the solution yielded by the primal/dual greedy algorithm.

Original: Proof of Theorem 4.2.
Revised: Proof of Theorem 4.2.

Original: Let z ∗ j , for j ∈ [ d ] , denote an optimal solution to the dual linear program of equation  5 in the main paper.
Revised: Let z ∗ j , for j ∈ [ d ] , denote an optimal solution to the dual linear program of equation equation 5 in the main paper.

Original: Standard sensitivity analysis of linear programs implies that z ∗ j ∈ R n ( j ∈ [ d ] ) is also optimal for the perturbed linear program, 
Revised: Standard sensitivity analysis of linear programs implies that z ∗ j ∈ R n ( j ∈ [ d ] ) is also optimal for the perturbed linear program, maximize

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

Original: where the indices in the constraints include all i j ∈ [ n ] , j ∈ [ d ] , ε > 0 is sufficiently small and h j ∈ R n are held fixed.
Revised: where the indices in the constraints include all i j ∈ [ n ] , j ∈ [ d ] , ε > 0 is sufficiently small and h j ∈ R n are held fixed.

Original: If ϕ ( x 1 + εh 1 , . . .
Revised: If ϕ ( x 1 + εh 1 , . . .

Original: , x d + εh d ) represents the optimal objective value of this program, then by linearity,
Revised: , x d + εh d ) represents the optimal objective value of this program, then by linearity,

Original: [Equation]
Revised: [Equation]

Original: Thus, we can form the directional derivative of ϕ as
Revised: Thus, we can form the directional derivative of ϕ as

Original: [Equation]
Revised: [Equation]

Original: where ∥ h ∥ 2 := (cid:80) j ∥ h j ∥ 2 .
Revised: where ∥ h ∥ 2 := (cid:80) j ∥ h j ∥ 2 .

Original: From this, it follows that
Revised: From this, it follows that

Original: [Equation]
Revised: [Equation]

Original: This shows the first claim of the theorem.
Revised: This shows the first claim of the theorem.

Original: To see the second claim, we have from Theorem 3.1 item 2 of Kline (2019)
Revised: To see the second claim, we have from Theorem 3.1 item 2 of Kline (2019)

Original: that
Revised: that

Original: [Equation]
Revised: [Equation]

Original: for all i ∈ [ n ] .
Revised: for all i ∈ [ n ] .

Original: Consequently, if one defines
Revised: Consequently, if one defines

Original: [Equation]
Revised: [Equation]

Original: then since we assume that e ′ x j = 1 for all j ∈ [ d ] ,
Revised: then since we assume that e ′ x j = 1 for all j ∈ [ d ] ,

Original: [Equation]
Revised: [Equation]

Original: where the last equality holds by equation (6).
Revised: where the last equality holds by equation (6).

Original: This shows the second claim.
Revised: This shows the second claim.

Original: 7. D IFFERENTIABLE H ISTOGRAMMING
Revised: 7. D IFFERENTIABLE H ISTOGRAMMING

Original: 1: function I NIT ( n ) // bins, (discretization) 2: r := 1 /n //
Revised: 1: function I NIT ( n ) // bins, (discretization) 2: r := 1 /n //

Original: bin size 3: locs :
Revised: bin size 3: locs :

Original: = arange (0 , 1 , r ) //
Revised: = arange (0 , 1 , r ) //

Original: bin boundaries 4: function F ORWARD ( acts ) 5: cdfs = σ ( acts ) // compute CDFs 6: counts = [] 7: for loc in locs do 8: dist = | cdfs − loc | // dist.
Revised: bin boundaries 4: function F ORWARD ( acts ) 5: cdfs = σ ( acts ) // compute CDFs 6: counts = [] 7: for loc in locs do 8: dist = | cdfs − loc | // dist.

Original: to boundary 9: ct = (cid:80) i ∈ [ nbins ] ReLU( r − dist [ i ]) //
Revised: to boundary 9: ct = (cid:80) i ∈ [ nbins ] ReLU( r − dist [ i ]) //

Original: soft bucket count 10: counts.append ( ct ) 11: out = stack ( counts ) 12: out = out/sum ( out ) 13: return out
Revised: soft bucket count 10: counts.append ( ct ) 11: out = stack ( counts ) 12: out = out/sum ( out ) 13: return out

Original: While gradients are now readily available, typical ML pipelines do not have distributions or histograms predefined at outputs which can be fed directly into our EMD loss.
Revised: While gradients are now readily available, typical ML pipelines do not have distributions or histograms predefined at outputs which can be fed directly into our EMD loss.

Original: Applying existing binning procedures over the batch to estimate histograms will break the ability to autodifferentiate: soft thresholds are necessary at bin boundaries such that samples within a bin may move smoothly as needed.
Revised: Applying existing binning procedures over the batch to estimate histograms will break the ability to autodifferentiate: soft thresholds are necessary at bin boundaries such that samples within a bin may move smoothly as needed.

Original: Algorithm 2 provides a differentiable histogram implementation.
Revised: Algorithm 2 provides a differentiable histogram implementation.

Original: Using a rectified linear relaxation allows for samples to have a continuous gradient towards neighboring bins.
Revised: Using a rectified linear relaxation allows for samples to have a continuous gradient towards neighboring bins.

Original: 7. E XPERIMENTAL D ETAILS
Revised: 7. E XPERIMENTAL D ETAILS

Original: Results reported in tables in the main paper are of the form M ( SD ) , where M is the mean and SD is the standard deviation calculated over replications.
Revised: Results reported in tables in the main paper are of the form M ( SD ) , where M is the mean and SD is the standard deviation calculated over replications.

Original: Setup details.
Revised: Setup details.

Original: Experiments were conducted using NumPy and PyTorch on a Intel(R) Xeon(R)CPU
Revised: Experiments were conducted using NumPy and PyTorch on a Intel(R) Xeon(R)

Original:  E5-2620 v3 @ 2.40GHz with an Nvidia Titan Xp GPU.
Revised: CPU E5-2620 v3 @ 2.40GHz with an Nvidia Titan Xp GPU.

Original: Particular parameter settings and experimental runs can be found below or in the scripts included with the code provided in the supplement.
Revised: Particular parameter settings and experimental runs can be found below or in the scripts included with the code provided at https: //github.com/ronakrm/demd .

Original: 7.3. D ATA AND L ICENSES
Revised: 7.3. D ATA AND L ICENSES

Original: The datasets Adult, Communities and Crime, and German datasets are all available under Creative Commons Attribution 4.0 International (CC BY 4.0) licenses via the UCI Machine Learning Dataset Repository https://archive.ics.uci.edu/ml/index.php .
Revised: The datasets Adult, Communities and Crime, and German datasets are all available under Creative Commons Attribution 4.0 International (CC BY 4.0) licenses via the UCI Machine Learning Dataset Repository https://archive.ics.uci.edu/ml/index.php .

Original: The CelebA dataset http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html is available for noncommercial purposes.
Revised: The CelebA dataset http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html is available for noncommercial purposes.

Original: See the website for more details.
Revised: See the website for more details.

Original: ACS Data.
Revised: ACS Data.

Original: The American Census Survey (ACS) has recently made available a large set of demographic data.
Revised: The American Census Survey (ACS) has recently made available a large set of demographic data.

Original: The original UCI Adult dataset Dua & Graff (2017) was curated from this data, however recent work by Ding et al.
Revised: The original UCI Adult dataset Dua & Graff (2017) was curated from this data, however recent work by Ding et al.

Original: (2021) has identified temporal shifts in demographic data, and recommends using a more recent collection as a baseline when evaluating biases and adjusting for fairness.
Revised: (2021) has identified temporal shifts in demographic data, and recommends using a more recent collection as a baseline when evaluating biases and adjusting for fairness.

Original: Part of their contribution includes APIs to directly interface with the data provided by the ACS, and the ability to identify and construct similar problems associated with the original UCI-provided dataset, albeit with updated data.
Revised: Part of their contribution includes APIs to directly interface with the data provided by the ACS, and the ability to identify and construct similar problems associated with the original UCI-provided dataset, albeit with updated data.

Original: Data for the income prediction task was downloaded from 2018, localized to Louisiana.
Revised: Data for the income prediction task was downloaded from 2018, localized to Louisiana.

Original: Race is the provided group label, which we wish to be agnostic towards, over some measure of our output.
Revised: Race is the provided group label, which we wish to be agnostic towards, over some measure of our output.

Original: Data was accessed using the folktables codebase https://github.com/zykls/folktables with MIT License.
Revised: Data was accessed using the folktables codebase https://github.com/zykls/folktables with MIT License.

Original: The US Census data accessed is available for use so long as it is not used in combination with other data “to identify any particular respondent to a Census Bureau survey.”
Revised: The US Census data accessed is available for use so long as it is not used in combination with other data “to identify any particular respondent to a Census Bureau survey.”

Original: See https: //www.census.gov/data/developers/about/terms-of-service.html for more details.
Revised: See https: //www.census.gov/data/developers/about/terms-of-service.html for more details.

Original: German Data.
Revised: German Data.

Original: The German dataset classifies people as good or bad credit risks.
Revised: The German dataset classifies people as good or bad credit risks.

Original: There are aboutfeatures ( 7 numerical and 13 categorical).
Revised: There are aboutfeatures ( 7 numerical and 13 categorical).

Original: These features represent the economic status of the person, such as, credit history, savings account, year of present employment, property and others.
Revised: These features represent the economic status of the person, such as, credit history, savings account, year of present employment, property and others.

Original: Adult Data.
Revised: Adult Data.

Original: The Adult dataset is comprised of demographic characteristics from the UCI repository
Revised: The Adult dataset is comprised of demographic characteristics from the UCI repository

Original: Dua & 
Revised: Dua & Graff (2017).

Original: The protected attribute here is gender.
Revised: The protected attribute here is gender.

Original: It contains 44 , 842 samples.
Revised: It contains 44 , 842 samples.

Original: The features that were used in the experiments include “age”, “workclass”, “fnlwgt”, “education”, “educationnum”, “marital-status”, “occupation”, “relationship”, “race”, “sex”, “capital-gain”, “capital-loss”, “hours-per week”, “native-country”, “income”.
Revised: The features that were used in the experiments include “age”, “workclass”, “fnlwgt”, “education”, “educationnum”, “marital-status”, “occupation”, “relationship”, “race”, “sex”, “capital-gain”, “capital-loss”, “hours-per-week”, “native-country”, “income”.

Original: A positve target label in this dataset is indicated by the attribute “income-bracket” being above $
Revised: A positve target label in this dataset is indicated by the attribute “income-bracket” being above $50K.

Original: Communities and Crime Data.
Revised: Communities and Crime Data.

Original: The Communities and Crime dataset consists of summary statistics of per-capita measures from a wide variety of communities across the United States, measured from a number of US census and surveys from the 1990’s.
Revised: The Communities and Crime dataset consists of summary statistics of per-capita measures from a wide variety of communities across the United States, measured from a number of US census and surveys from the 1990’s.

Original: The original goal of the dataset was to predict crime rates in communities as a function of various demographic, socioeconomic, and other features.
Revised: The original goal of the dataset was to predict crime rates in communities as a function of various demographic, socioeconomic, and other features.

Original: The dataset has widely become known as the prototypical example in which using racial population distributions can be extremely harmful in perpetuating stereotypes and lead to models that continue to exacerbate inequity that may exist within the data.
Revised: The dataset has widely become known as the prototypical example in which using racial population distributions can be extremely harmful in perpetuating stereotypes and lead to models that continue to exacerbate inequity that may exist within the data.

Original: As such, the dataset has become a de facto tool in evaluating fairness metrics and methods that attempt to account for these inequities and biases.
Revised: As such, the dataset has become a de facto tool in evaluating fairness metrics and methods that attempt to account for these inequities and biases.

Original: The preprocessed data contains 1994 samples, and we attempt to predict the violent crime rate (binarized at 0.3 after normalization between 0 and 1), and the sensitive attribute is a similarly binarized version of the percentage black population variable.
Revised: The preprocessed data contains 1994 samples, and we attempt to predict the violent crime rate (binarized at 0.3 after normalization between 0 and 1), and the sensitive attribute is a similarly binarized version of the percentage black population variable.

Original: CelebA Data.
Revised: CelebA Data.

Original: CelebA
Revised: CelebA

Original: (Liu et al., 2015) consists of 200 K celebrity face images from the internet annotated by a group of paid adult participants.
Revised: (Liu et al., 2015) consists of 200 K celebrity face images from the internet annotated by a group of paid adult participants.

Original: There are up to 40 labels available in the dataset, each of which is binary-valued.
Revised: There are up to 40 labels available in the dataset, each of which is binary-valued.

Original: 7.3. F AIRNESS E XPERIMENT D ETAILS
Revised: 7.3. F AIRNESS E XPERIMENT D ETAILS

Original: All experiments related to fairness use a three-layer fully-connected neural network classifier with a hidden layer size of 100.
Revised: All experiments related to fairness use a three-layer fully-connected neural network classifier with a hidden layer size of 100.

Original: Numbers reported in Table 1 in the main paper are means and standard deviations over three replicate runs with different seeds.
Revised: Numbers reported in Table 1 in the main paper are means and standard deviations over three replicate runs with different seeds.

Original: Hyperparameters were selected as described in the main paper, by taking the best result for each dataset over the parameter range λ ∈ [1 . 0 , 0 . 1 , 10 , 0 . 01 , 100 , 0 . 001] .
Revised: Hyperparameters were selected as described in the main paper, by taking the best result for each dataset over the parameter range λ ∈ [1 . 0 , 0 . 1 , 10 , 0 . 01 , 100 , 0 . 001] .

Original: The full results of this sweep are presented in Table 3 and
Revised: The full results of this sweep are presented in Table 3 and

Original: Table 4.
Revised: Table 4.

Original: The | DP | and | EO | measures are the gap between the corresponding conditional probability distributions when there is a single binary group attribute, and the largest gap ( max − min ) when there are more than 2 groups.
Revised: The | DP | and | EO | measures are the gap between the corresponding conditional probability distributions when there is a single binary group attribute, and the largest gap ( max − min ) when there are more than 2 groups.

Original: The DEMD loss is computed as described in the main paper with a discretization/bin level of 10.
Revised: The DEMD loss is computed as described in the main paper with a discretization/bin level of 10.

Original: 7.3. H ARMONIZATION E XPERIMENT D ETAILS
Revised: 7.3. H ARMONIZATION E XPERIMENT D ETAILS

Original: For all our experiments related to harmonization, we use an encoder-decoder framework comprising of fully-connected layers.
Revised: For all our experiments related to harmonization, we use an encoder-decoder framework comprising of fully-connected layers.

Original: The hidden layers comprised of 64 nodes and the latent space was of dimension 30 .
Revised: The hidden layers comprised of 64 nodes and the latent space was of dimension 30 .

Original: We report the mean and standard deviation on unseen test datasets for three random runs.
Revised: We report the mean and standard deviation on unseen test datasets for three random runs.

Original: The hyper-parameter selection has been done on a validation split obtained from the training dataset.
Revised: The hyper-parameter selection has been done on a validation split obtained from the training dataset.

Original: The model that achieves the best ADV (the adversarial evaluation measure), with the test set accuracy remains within 5% of the vanilla (titled “None” in the paper) model, is chosen.
Revised: The model that achieves the best ADV (the adversarial evaluation measure), with the test set accuracy remains within 5% of the vanilla (titled “None” in the paper) model, is chosen.

Original: Recall that the harmonization experiments aimed at minimizing the distributional differences of the latent features across the groups.
Revised: Recall that the harmonization experiments aimed at minimizing the distributional differences of the latent features across the groups.

Original: Below, we will provide details on the evaluation metrics, ADV and
Revised: Below, we will provide details on the evaluation metrics, ADV and

Original: MMD measures, that aptly assess distributional differences of the latent features.
Revised: MMD measures, that aptly assess distributional differences of the latent features.

Original: Evaluation Metric: Adversarial Measure (ADV)
Revised: Evaluation Metric: Adversarial Measure (ADV)

Original: The ADV measure corresponds to the accuracy obtained by training a separate neural network to predict groups from the latent features.
Revised: The ADV measure corresponds to the accuracy obtained by training a separate neural network to predict groups from the latent features.

Original: This step is conducted post harmonization.
Revised: This step is conducted post harmonization.

Original: A lower value of the ADV accuracy denotes that the latent features are free of any group related information.
Revised: A lower value of the ADV accuracy denotes that the latent features are free of any group related information.

Original: This implies successful harmonization and suggests minimal distributional differences of the latent features across groups.
Revised: This implies successful harmonization and suggests minimal distributional differences of the latent features across groups.

Original: We follow (Xie et al., 2017) for training the adversary used for reporting ADV measure.
Revised: We follow (Xie et al., 2017) for training the adversary used for reporting ADV measure.

Original: We use a three-layered fully-connected network with batch normalization and train it with Adam optimizer for 200 epochs.
Revised: We use a three-layered fully-connected network with batch normalization and train it with Adam optimizer for 200 epochs.

Original: The learning rate for the adversary is decreased multiplicatively by a factor of 0 .
Revised: The learning rate for the adversary is decreased multiplicatively by a factor of 0 .

Original: 65 every 10 epochs for convergence.
Revised: 65 every 10 epochs for convergence.

Original: Evaluation Metric: Maximum Mean Discrepancy (MMD)
Revised: Evaluation Metric: Maximum Mean Discrepancy (MMD)

Original: We simply use the following MMD criterion as described in (Gretton et al., 2006) and evaluate the metric on the latent features obtained from the test set.
Revised: We simply use the following MMD criterion as described in (Gretton et al., 2006) and evaluate the metric on the latent features obtained from the test set.

Original: Each group is considered as a different distribution and a lower value of this metricsuggests minimal distributional differences across the groups. For latent feature vector ℓ and groups
Revised: Each group is considered as a different distribution and a lower value of this metric

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  i/j , we have
Revised: suggests minimal distributional differences across the groups. For latent feature vector ℓ and groups i/j , we have

Original: [Equation]
Revised: [Equation]

Original: The criterion is defined using a Reproducing Kernel Hilbert Space with norm ∥ · ∥ H and kernel K .
Revised: The criterion is defined using a Reproducing Kernel Hilbert Space with norm ∥ · ∥ H and kernel K .

Original: 7.3. MWGAN D ETAILS AND A DDITIONAL R ESULTS
Revised: 7.3. MWGAN D ETAILS AND A DDITIONAL R ESULTS

Original: The MWGAN code https://github.com/deepmo24/MWGAN was used under the MIT license extended from the original StarGAN codebase https://github.com/yunjey/ stargan .
Revised: The MWGAN code https://github.com/deepmo24/MWGAN was used under the MIT license extended from the original StarGAN codebase https://github.com/yunjey/ stargan .

Original: The original paper constructs an inter-domain penalty added to the multi-domain GAN discriminator loss as follows:
Revised: The original paper constructs an inter-domain penalty added to the multi-domain GAN discriminator loss as follows:

Original: [Equation]
Revised: [Equation]

Original: The sampling for expectation in practice is done by interpolating randomly between real data and generated data for each domain.
Revised: The sampling for expectation in practice is done by interpolating randomly between real data and generated data for each domain.

Original: In contrast to this, we directly push the gradient norms computed over samples from each domain to be close using our DEMD regularizer:[Equation]
Revised: In contrast to this, we directly push the gradient norms computed over samples from each domain to be close using our DEMD regularizer:

Original: [Equation]
Revised: [Equation]

Original: ) where [ i ] indicates the domain associated with the fake samples generated in the same manner for each domain.
Revised: (9) where [ i ] indicates the domain associated with the fake samples generated in the same manner for each domain.

Original: Convergence Comparison.
Revised: Convergence comparison.

Original: Interestingly we observe that our construction tends to lead to more stable training.
Revised: Interestingly we observe that our construction tends to lead to more stable training.

Original: In Figure 7, we can see that the variance of the total discriminator and generator loss fluctuate much less.
Revised: In Figure 7, we can see that the variance of the total discriminator and generator loss fluctuate much less.

Original: We suspect that this can be directly attributed to the boundedness of the gradients of the DEMD regularizer, specific to the method where we compute the gradients via the dual of the
Revised: We suspect that this can be directly attributed to the boundedness of the gradients of the DEMD regularizer, specific to the method where we compute the gradients via the dual of the

Original: LP: the dual variables are bounded by definition and in this particular LP, are bounded exactly by the number of bins used for discretization.
Revised: LP: the dual variables are bounded by definition and in this particular LP, are bounded exactly by the number of bins used for discretization.

Original: While additional investigation is needed, smoothness has been identified as a desirable property of GAN training (Arora et al., 2017; Chu et al., 2020).
Revised: While additional investigation is needed, smoothness has been identified as a desirable property of GAN training (Arora et al., 2017; Chu et al., 2020; Lokhande et al., 2020b; Ravi et al., 2019).

Original: [Table]
Revised: [Table]

Original: Quantitative Results.
Revised: Quantitative Results.

Original: Using the Frechet Inception Distance we quantitatively measure the resulting generative samples.
Revised: Using the Frechet Inception Distance we quantitatively measure the resulting generative samples.

Original: Following the description in Cao et al.
Revised: Following the description in Cao et al.

Original: (2019), we present the FID scores of our model and theirs in Table 5.
Revised: (2019), we present the FID scores of our model and theirs in Table 5.

Original: Our metric as a drop-in replacement performs comparably.
Revised: Our metric as a drop-in replacement performs comparably.

Original: 7. D -D IMENSIONAL E ARTH M OVER ’ S D ISTANCE B ACKGROUND AND A LGORITHM
Revised: 7. D -D IMENSIONAL E ARTH M OVER ’ S D ISTANCE B ACKGROUND AND A LGORITHM

Original: Algorithm 3 describes the greedy algorithm that solves both primal and dual generalized Earth mover’s programs, also see (Kline, 2019).
Revised: Algorithm 3 describes the greedy algorithm that solves both primal and dual generalized Earth mover’s programs, also see (Kline, 2019).

Original: The algorithm accepts d distributions (i.e., histograms) p 1 , . .
Revised: The algorithm accepts d distributions (i.e., histograms) p 1 , . .

Original: p d ∈ R n + with e ′ p
Revised: p d ∈ R n + with e ′ p

Original: j
Revised: j

Original: = 1 for all j ∈ [ d ] .
Revised: = 1 for all j ∈ [ d ] .

Original: Although the algorithm states that all histograms have the same number of bins, the algorithm can be easily adapted to accept as inputs p i ∈ R n i + with n i ̸ = n j .
Revised: Although the algorithm states that all histograms have the same number of bins, the algorithm can be easily adapted to accept as inputs p i ∈ R n i + with n i ̸ = n j .

Original: Algorithm 3 EMD
Revised: Algorithm 3 EMD

Original: Primal/
Revised: Primal/

Original: Dual Algorithm input p j ∈ R n + with e ′ p
Revised: Dual Algorithm input p j ∈ R n + with e ′ p

Original: j = 1 ,
Revised: j = 1 ,

Original: ( ∀ j ∈ [ d ]) //iteration index, active indices, primal variables, dual variables k := 

Revised: ( ∀ j ∈ [ d ]) //iteration index, active indices, primal variables, dual variables k :=

Original: ∈ Z , I := 

Revised: ∈ Z , I :=

Original: ∈ Z d , x := 0 ∈ R n d , z j := 0 ∈ R n , ( ∀ j ∈ [ d ]) while I ( j ) ≤ n , ( ∀ j ∈ [ d ]) do s k :=
Revised: ∈ Z d , x := 0 ∈ R n d , z j := 0 ∈ R n , ( ∀ j ∈ [ d ]) while I ( j ) ≤ n , ( ∀ j ∈ [ d ]) do s k :=

Original: min j ∈ [ d ] p j ( I ( j )) //
Revised: min j ∈ [ d ] p j ( I ( j )) //

Original: the mass to move x
Revised: the mass to move x

Original: ( I ) ← s k // update the EMD solution p j
Revised: ( I ) ← s k // update the EMD solution p j

Original: ( I ( j ))
Revised: ( I ( j ))

Original: ← p j
Revised: ← p j

Original: ( I ( j ))
Revised: ( I ( j ))

Original: − s k , ( ∀ j ∈ d ) // shrink the data j ∗ ← arg min j ∈
Revised: − s k , ( ∀ j ∈ d ) // shrink the data j ∗ ← arg min j ∈

Original: [ d ] p j ( I ( j ))
Revised: [ d ] p j ( I ( j ))

Original: I  z d  , and the 
Revised: I ( of x , ( z 1 , . . . , z d ) , and the objective value (cid:80) k s k t k .

Original: The algorithm has explicit terminal conditions for the main while loop.
Revised: The algorithm has explicit terminal conditions for the main while loop.

Original: In the worst case the number of iterations can be nd .
Revised: In the worst case the number of iterations can be nd .

Original: 7. E XTENDED E THICS D ISCUSSION
Revised: 7. E XTENDED E THICS D ISCUSSION

Original: As discussed in the main paper discussion, the primary application of our proposed construction is to reduce invariance over a particular set of features.
Revised: As discussed in the main paper discussion, the primary application of our proposed construction is to reduce invariance over a particular set of features.

Original: In practice, with respect to typical machine learning models and pipelines, this corresponds to minimizing performance difference as measured across subgroups within the data corresponding to a minority or protected subsets of samples or individuals.
Revised: In practice, with respect to typical machine learning models and pipelines, this corresponds to minimizing performance difference as measured across subgroups within the data corresponding to a minority or protected subsets of samples or individuals.

Original: While the construction can be applied to any existing ML pipeline, we do not claim to provide a catch-all solution for group disparity that may be inherent to the data or exacerbated by the choice of the ML model that is being used.
Revised: While the construction can be applied to any existing ML pipeline, we do not claim to provide a catch-all solution for group disparity that may be inherent to the data or exacerbated by the choice of the ML model that is being used.

Original: As always, care needs to be taken when working with sensitive data or models which may have disparate impacts on different groups.
Revised: As always, care needs to be taken when working with sensitive data or models which may have disparate impacts on different groups.

Original: We point interested readers to the following extensive surveys and references therein regarding various methods and procedures for addressing and dealing with bias and unfairness in ML problems, and the potential danger associated with using models without care: (Mehrabi et al., 2021; Leavy, 2018; O’neil, 2016; d’Alessandro et al., 2017; Rakova et al., 2021).
Revised: We point interested readers to the following extensive surveys and references therein regarding various methods and procedures for addressing and dealing with bias and unfairness in ML problems, and the potential danger associated with using models without care: (Mehrabi et al., 2021; Leavy, 2018; O’neil, 2016; d’Alessandro et al., 2017; Rakova et al., 2021).

Original: We also note that the final multi-marginal GAN image translation application could be used to generate so-called “deepfakes.”
Revised: We also note that the final multi-marginal GAN image translation application could be used to generate so-called “deepfakes.”

Original: While the results of our algorithm are comparable to existing works, we believe that existing methods of identifying deepfakes would work well, and that the methods provided here and in the original paper Cao et al.
Revised: While the results of our algorithm are comparable to existing works, we believe that existing methods of identifying deepfakes would work well, and that the methods provided here and in the original paper Cao et al.

Original: (2019) would require significant effort to be made practical for much larger scale images.
Revised: (2019) would require significant effort to be made practical for much larger scale images.

Revised:  A CKNOWLEDGMENTS

Revised: This work was supported in part by NIH grants RF1AG059312, RF1AG062336, RF1AG059869, and

Revised: NSF award CCF 1918211.

Revised: The work presented here was partially developed while Ronak Mehta,

Revised: Jeffery Kline, and Glenn Fung were at the Machine Learning Group at American Family Insurance.

Revised: The authors thank Sathya N. Ravi (UIC) for time reading the paper and providing suggestions.