Sparsity-Constrained Optimal Transport
Authors: Tianlin Liu
Joan Puigcerver
Mathieu Blondel
Venue: ICLR 2023 notable top 25% (2023)
Keywords: optimal transport
sparsity

Abstract: Regularized optimal transport (OT) is now increasingly used as a loss or as a matching layer in neural networks. Entropy-regularized OT can be computed using the Sinkhorn algorithm but it leads to fully-dense transportation plans, meaning that all sources are (fractionally) matched with all targets. To address this issue, several works have investigated quadratic regularization instead. This regularization preserves sparsity and leads to unconstrained and smooth (semi) dual objectives, that can be solved with off-the-shelf gradient methods. Unfortunately, quadratic regularization does not give direct control over the cardinality (number of nonzeros) of the transportation plan. We propose in this paper a new approach for OT with explicit cardinality constraints on the transportation plan. Our work is motivated by an application to sparse mixture of experts, where OT can be used to match input tokens such as image patches with expert models such as neural networks. Cardinality constraints ensure that at most $k$ tokens are matched with an expert, which is crucial for computational performance reasons. Despite the nonconvexity of cardinality constraints, we show that the corresponding (semi) dual problems are tractable and can be solved with first-order gradient methods. Our method can be thought as a middle ground between unregularized OT (recovered in the limit case $k=1$) and quadratically-regularized OT (recovered when $k$ is large enough). The smoothness of the objectives increases as $k$ increases, giving rise to a trade-off between convergence speed and sparsity of the optimal plan.

Published as a conference paper at ICLR 2023
SPARSITY-CONSTRAINED OPTIMAL TRANSPORT
Tianlin Liu" University of Basel
Joan Puigcerver
Mathieu Blondel
Google Research, Brain team
Google Research, Brain team
ABSTRACT
Regularized optimal transport (OT) is now increasingly used as loss matching layer in neural networks. Entropy-reeullrizeed OT can be computed us- ing the Sinkhorn algorithm but it leads to fully-dense transportation plans, mean- ing that all sources are (fractionally) matched with all targets. To address this issue, several works have investigated quadratic regularization instead. This reg- ularization preserves sparsity and leads to unconstrained and smooth (semi) dual objectives, that can be solved with off-the-sheef gradient methods. Unfortunately, quadratic regularization does not give direct control over the cardinality (number nonzeros) of the transportation plan. We propose in this paper a new approach for OT with explicit cardinality constraints on the transportation plan. Our work is motivated by an application to sparse mixture of experts, where OT can be used to match input tokens such image patches with expert models such neural net- works. Cardinality constraints ensure that at most k tokens are matched with an expert, which crucial for computational performance reasons. Despite the non- convexity of cardinality constraints, we show that the corresponding (semi) dual problems are tractable and can be solved with first-order gradient methods. Our method can be thought as a middle ground between unregularized OT (recovered when & is small enough) and aeairratically-eglaariii large enough). The smoothness of the objectives increases as k increases, giving rise to a trade-off between convergence speed and sparsity of the optimal plan.
OT (recovered when k: is
INTRODUCTION
Optimal transport (OT) distances (a.k.a. Wasserstein or earth mover's distances) are a powerful computational tool to compare probability distributions and have found widespread use in machine learning (Solomon et al., 2014; Kusner et al., 2015; Arjovsky et al., 2017). While OT distances exhibit a unique ability to capture the geometry of the data, their applicability has been largely ham- pered by their high computational cost. Indeed, computing OT distances involves a linear program, which takes super-cubic time to solve using sfatc-of-the-an network-flow algorithms (Kennington & Helgason, 1980; Ahuja et al., 1988). In addition, these algorithms challenging implement and are not GPU or TPU friendly. An alternative approach consists instead solving the so-called semi-dual using (stochastic) subgradient methods (Carlier et al., 2015) or quasi-Newton methods (Mérigot, 2011; Kitagawa et al., 2019). However, the semi-dual is nonsmooth, piecewise-lined function, which can lead to slow convergence in practice.
For all these reasons, the machine learning community has now largely switched to regularized OT. Popularized by Cuturi (2013), rithm (1967) and is differentiable w.r.t. its inputs, enabling OT as a differentiable loss (Cuturi, 2013; Feydy et al., 2019) or as layer in a neural network (Genevay et al., 2019; Sarlin et al., 2020; Sander et al., 2022). A disadvantage of entropic regularization. however, is that it leads to fully-dense trans- portation plans. This is problematic in applications where undesirable to (fractionally) match all sources with all targets, e.g., for interpretability or for computational cost reasons. To address this issue, several works have investigated quadratic regularization instead (Dessein et al., 2018; Blondel et al., 2018; Lorenz et al., 2021). This regularization preserves sparsity and leads to unconstrained and smooth (semi) dual objectives, solvable with off-the-shelf algorithms. Unfortunatcly. it does not give direct control over the cardinality (number of nonzeros) of the transportation plan.
OT can be computed using the Sinkhorn algo-
Work done during an internship Google Research, Brain Team.

---

Published as a conference paper at ICLR 2023
Unregularized
Entropy regularized
Squared 2-norm regularized
Sparsity constrained (ours) k-2
k-1
Source Target
4
-2
Figure 1: OT formulation comparison (m = n 20 points), with squared Euclidean distance cost, and with uniform source and target distributions. The unregularized OT plan is maximally sparse and contains at most m-+n-1 nonzero elements. On the contrary, with entropy-regularized OT, plans are always fully dense, meaning that all points fractionally matched with another (nonzeros of transportation plan are indicated by small squares). Squared 2-norm (quadratically) regularized OT preserves sparsity but the number of nonzero elements cannot be directly controlled. Our proposed sparsity consstty-oosstrained OT allows us to set a maximum number of nonzeros per column. It recovers unregularized OT in the limit case 1 (Proposition 4) and quadratically-regularii OT when is large enough. It can be computed using solvers such as LBFGS or ADAM.
In this paper, we propose a new approach for OT with explicit cardinality constraints on the columns of the transportation plan. Our work motivated by an application to sparse mixtures of experts, in which we want each token (e.g. word or an image patch) to be matched with at most experts (e.g., multilayer perceptrons). This is critical for computational performance reasons, since the cost of processing a token is proportional to the number of experts that have been selected for it. Despite the nonconvexity of cardinality constraints, we show that the corresponding dual and semi- dual problems are tractable and can be solved with first-order gradient methods. Our method can be thought as middle ground between unregularized OT (recovered when is small enough) and OT (recovered when & is large enough). We empirically show that the dual and semi-dual are increasingly smooth as & increases, giving rise to a trade-off between convergence speed and sparsity. The rest of the paper is organized as follows.
quadratically-regulattt
We review related work in $2 and existing work on OT with convex regularization in $3.
We propose in $4 a framework for OT with nonconvex regularization, based on the dual and semi- dual formulations. We study the weak duality and the primal interpretation of these formulations.
We apply framework in to OT with cardinality constraints. We show that the dual and semi- dual formulations are tractable and that smoothness of the objective increases as A increases. We show that approach is equivalent to using squared k-support norm regularization in the primal.
We validate our framework in $6 and in Appendix A through variety of experiments.
Notation and convex analysis tools. Given matrix T E Rmxn, we denote its columns by Ł, E RT for j E [n]. We denote the non-negative orthant by RTH and the non-positive orthant by Rm We denote the probability simplex by
A {p Rm: (p. 1) 1}. We will also use 6/ m to denote the set {t E RT (t, 1) = b}. The convex conjugate of a function f: RRm + {oo} is defined by (8) == suPeedomm// () (s, f(t). It is well-known that is convex (even if is not). If the solution is unique, then its gradient is (s*(ss argmax (s. t) (t)). If the solution is not unique, then obtain subgradient. We denote the indicator function of a set C by de, i.e.. sc(t) = if E C and (t) = ao otherwise. We denote the Euclidean projection onto the set by proje (s) argmingec |ss t| The projection is unique when C is convex, while it may not be when C is nonconvex. We use [.+ to denote the non non-negative part, evaluated element-wise Given a vector E Rm we 8ig denote its i-th largest value, i.e., 8/1]
m))
2

---

Published as a conference paper at ICLR 2023
2 RELATED WORK
Sparse optimal transport. OT with arbitrary strongly convex regularization is studied by Dessein al. (2018) and Blondel et al. (2018). More specifically, quadratic regularization was studied in the discrete (Blondel et al., 2018; Roberts et al., 2017) and continuous settings (Lorenz et al., 2021). Although known that quadratic regularization leads to sparse transportation plans, does not enable explicit control of the cardinality (maximum number of nonzero elements), do. In this work, we study the nonconvex regularization case and apply it to cardinality constrained OT.
Sparse projections. In this paper, we use k-sparse projections as a core building block of our framework. Sparse projections the simplex and on the non-negative orthant were studied by Kyrillidis et al. (2013) and Bolte et al. (2014), respectively. These studies were later extended to more general sets (Beck & Hallak, 2016). On the application side, sparse projections the simplex were used for structured prediction (Pillutla et al., 2018; Blondel et al., 2020), for marginalizing over discrete variables (Correia et al., 2020) and for Wasserstein K -means (Fukunaga & Kasai, 2021).
Sparse mixture of experts (MoE). In contrast usual deep learning models where all parameters interact with all inputs, a sparse MoE model activates only a small part of the model ("experts") in an input-dependent manner, thus reducing the overall computational cost of the model. Sparse MoEs have been tremendously successful in scaling up deep learning architectures in tasks including computer vision (Riquelme et al., 2021), natural language processing (Shazeer et al., 2017; Lewis et al., 2021; Lepikhin et al., 2021; Roller et al., 2021; Fedus et al., 2022b; Clark et al., 2022), speech processing (You et al., 2022), and multimodal learning (Mustafa et al., 2022). reducing computational cost, sparse MoEs have also shown other benefits, such as an enhancement in adversarial robustness (Puigcerver et al.. 2022). See Fedus et al. (2022a) for a recent survey. Crucial to sparse MoE model is its routing mechanism that decides which experts get which inputs. Transformer-bass patches). To balance the assignments of tokens to experts, recent works cast the assignment problem as entropy-regularized OT (Kool et al., 2021; Clark et al., 2022). We go beyond entropy-regularized OT and show that sparsity-constrainee OT yields a more natural and effective router.
addition to
MoE models typically route individual tokens (embedded words or image
3 OPTIMAL TRANSPORT WITH CONVEX REGULARIZATION
In this section, we review OT with convex regularization, which also includes the unregularized case. For a comprehensive survey on computational OT, see (Peyré & Cuturi, 2019).
Primal formulation. We focus throughout this paper on OT between discrete probability distri- butions a E A and e A", Rather than performing pointwise comparison of the distributions, OT distances compute the minimal effort, according to some ground cost, for moving the probabil- ity mass of one distribution to the other. Recent applications of OT in machine learning typically add regularization on the transportation plan T. In this section, we apply convex regularization Ω: IR + IR, )oo) separately the columns t, E RT of T and consider the primal formulation
Pa(a, b, O) min TEl((a,))
(T,C) Σ s(tt),
where C E I is a cost matrix and U(a.,) b) = (T E RmXx; TI, = a, TT. 1m b) is the transportation polytope, which can be interpreted as the set of all joint probability distributions with
marginals a and b. It includes the Birkhoff polytope as a special case when m = n and a = b 1m
Dual and semi-dual formulations. Let us denote
2- (8) = (S2 + )
max (8, 1) - 22())
EER"
and
(g(s) == (Ω + 55 (())=
max (,,t) - Q(t).
(3)
tebA m
3

---

Published as a conference paper at ICLR 2023
(tt))
22 (8)
Unregularized
0
OR (s)
max 8i
Negentropy
(t, log t)
Σ i=1 exi=1
log Eis e -
Squared 2-norm
Σ 1= [s,]2
2 Ei= 11 11:20 02) Di=1 1 m/ 20 25 -22)
Sparsity-constrained (top-k)
$ ||t1l12 + uB. (t)
1 Σ i= sst1
Sparsity-constraaneed (top-1)
Jii 2 + 6z, (t)
max [m] [s:]
b maxxe [mn] 8; 32
Table I: Summary of the conjugate expressions (2) and (3) for various choices of Ω. Here, 0 and such that Ei=ls [s; 0]+ and [ -]+ sum ($5). where sijj denotes the i-th largest entry of the vector s E RR.. The top-k and top-1 expressions above assume no ties in 8.
The dual and semi-dual corresponding to (1) can then be written (Blondel et al., 2018) as
Da(a, b, C) =
max (a, a) + (3,b) - Σ j=1
24 (a + 3,1m - -c,)
(4)
a RER",
and
Sp(a,b,Cr = max (a, a) -Pala, b, = max (a, a) Rm
Σ j=1
2 (a - cj),
(5)
ask
where Pa denotes the conjugate in the first argument. When Q convex (which also includes the unregularized case Q 0), by strong duality, have that Pa(a, b, C) Da(a,b,C) Sn (a, b, C*) for all E A b E A and C' e R=x=.
With Ω 0 (without regularization), then (2) becomes the indicator function of the non-positive orthant, leading to the constraint a; + B5 S Ci.-- The dual (4) is then a constrained linear program and the most commonly used algorithm is the network flow solver. On the other hand, (3) becomes a max operator, leading to the so-called c- c-transform 35- = min.g[m] Grj a for all j E [n]. The semi-dual (5) is then unconstrained, but is a nonsmooth piecewise linear function.
Computation.
The key advantage of introducing strongly convex regularization is that it makes the corresponding (semi) dual easier to solve. Indeed, (2) and (3) become "soft" constraints and max operators.
In particular, when Ω is Shannon's 's negentropy si(t) = 7t,, log t). where controls the regulariza- tion strength, then (2) and (3) rely on the exponential and log-sum-exp operations. is well known that the primal (1) can then be solved using Sinkhorn's 's algorithm (Cuturi, 2013), which amounts to using a block coordinate ascent scheme w.r.t. a E Rm and 3 E R" in the dual (4). As pointed out in Blondel et al. (2018), the semi-dual is smooth (i.e., with Lipschitz gradients) but the dual is not.
When is the quadratic regularization ?(t) = 2 || then as shown in Blondel et al. (2018, Table 1). (2) and (3) rely on the squared relu and on the projection onto the simplex. However, is shown empirically that block coordinate ascent scheme in the dual (4) converges slowly. Instead, Blondel al. (2018) propose to use LBFGS both to solve the dual and the semi-dual. Both the dual and the semi-dual are smooth (Blondel et al., 201 8), i.c., with Lipschitz gradients. For both types of regularization, when y + O, we recover unregularized OT.
Recovering plan. f2 is strictly convex, the unique optimal solution T* of (1) can recovered from an optimal solution (a* f²*) of the dual (4) by
f (n= + 3, -c))
Vj e [n]
(6)
or from an optimal solution a' of the semi-dual (5) by
= V22 (a' c;) Vj E [n].
(7)
If Q is convex but not strictly so, recovering T more involved; Appendix B.2 for details.
4

---

Published as a conference paper at ICLR 2023
4 OPTIMAL TRANSPORT WITH NONCONVEX REGULARIZATION
In this section, we again focus on the primal formulation (1), but now study the case when the regularization ft: RT R+ U {oo} is nonconvex.
Concavity. It is well-known that the conjugate function is always convex, even when the original function is not. As a result, even if the conjugate expressions (2) and (3) involve nonconcave maxi- mization problems in the variable t, they induce convex functions in the variable We can therefore make the following elementary remark: the dual (4) and the semi-dual (5) are concave maximization problems, even if Ω is nonconvex. This means that can solve them to arbitrary precision as long as we know how to compute the conjugate expressions (2) and (3). This is generally hard but we will see in $5 setting in which these expressions can be computed exactly. We remark that the identity So(a, (a, b, C) = maXaeR" (a,, a} (a, b, C) still holds even when S is nonconvex.
The semi-dual upper-hounds the dual. Of course, if S is nonconvex, only weak duality holds, i.., the dual (4) and semi-dual (5) are lower bounds of the primal (1). The next proposition clarifies that the semi-dual is actually an upper-bound for the dual (a proof is given in Appendix B.1).
Proposition 1. Weak duality
Let Q: RT + RR U (xx) (potentially nonconvex). For all Am, b E A and C R mxx
Da(a, b, о) & So(a,b,C & Pa(a,b,C)
Therefore, if the goal is to compute approximately Pa(a, b, C), which involves intractable non- convex problem in general, may be advantageous to use Sa(a, b, C*) as a proxy, rather than Da(a, b, C). However, for the specific choice of Ω in $5, will that Da(a, b, C) and Sn (a, b, C)) actually coincide, i.e., Da(a, b, C) = Sn(a, b, C) < Pa(a, C)
Recovering a plan. Many times, the goal is not compute the quantity Pa b, C) itself, but rather the associated OT plan. If Ω is nonconvex, this is again intractable due to the nonconvex nature of the problem. As an approximation, given an optimal solution (a", , of the dual or an optimal solution a' of the semi-dual, we propose to recover transportation plan with (6) and (7), just like we would do the convex case. The following proposition clarifies that optimal transportation plan T that we get corresponds to convex relaxation of the original nonconvex problem (1). A proof is given Appendix B.3.
Proposition 2. Primal interpretation
Let Q: RTh R4 U (20) (potentially nonconvex). For all A b E A and C R
Dn
min (T,C) +
Σ
27 Pa-- (a,b,
TeRmxn T1,=0 T 1m b
1=1
Sn(a, b, = min (T,C)
Σ 1=1
= Pa-
TERmXn Xn
11,=a
In the above, f** denotes the biconjugate f. the tightest convex lower bound of f. When Ω is nonconvex, deriving expression for 2 and 0755 could be challenging in general. Fortunately, for the choice of Q in $5, we are able to do so. When function is convex and closed, its biconjugate is itself. As a result, Ω is convex and closed function, we recover Po(a, (a b C) Da(a, b, C) So (a, b, C.) for all mm, b E A and C' Rmxn.
Summary: proposed method. To approximately solve the primal OT objective (1) when Ω is non- convex, we proposed to solve the dual (4) or the semi-dual (5), which by Proposition lower-bound the primal. We do by solving the concave maximization problems in (4) and (5) by gradient- based algorithms, such as LBFGS (Liu & Nocedal, 1989) or ADAM (Kingma & Ba, 2015). When transportation plan is needed, we recover it from (6) and (7), as would do with convex Ω.
5

---

Published as a conference paper at ICLR 2023
k =
k=
k=
6
6
2
1
2
2
2
2
-2
-2
-2
-2
62
2
2
$1
Increasing k increases smoothness. Let (81, 82, 1). We visualize Q (s), defined (10) and derived in Table 1, when varying $1 and $2. It can be interpreted as relaxation of the indicator function of the non-positive orthant. The conjugate (s) (not shown) can be interpreted as a relaxed max operator, scaled by b. In both cases, the smoothness increases when &t increases.
Figure 2:
Proposition 2 clarifies what objective function this plan optimally solves. When learning with OT as a loss, necessary to differentiate through Sp(a, b, C). From Danskin's theorem, the gradients V VaSn(a,b, C* R" and Vaaaaa, b, C) E [Rmx" given by ax and 7 from (7), respectively,
5 QUADRATICALLY
OT WITH SPARSITY CONSTRAINTS
In this section, we build upon $4 to develop regularized OT formulation with sparsity constraints.
Formulation. Formally, given E Rm let us define the &o pseudo norm by
| lu = [t, * 0: j [miH].
i.e., the number of nonzero elements in t. For & (1, (1...., m). we denote the la level sets
B = {t Rm || k}.
Our goal in this section then to approximately solve the following guadratically--rgulaiizz opti- mal transport problem with cardinality constraints the columns of T
(T,C) 11T
min TEl(lab)
(8)
TEBx XB,
where y > 0 controls the regularization strength and where k is assumed large enough make the problem feasible. Problem (8) is a special case of (1) with the nonconvex regularization
We can therefore apply the methodology outlined in $4. the cardinality constraints need to be applied to the rows instead of to the columns, we simply transpose the problem.
Computation. We recall that in order to solve the dual (4) or the semi-dual (5), the main quanti- ties that we need to be able to compute are the conjugate expressions (2) and (3), as well as their gradients. While this is intractable for general nonconvex Q with the choice of Ω in (9), we obtain
2 (7(s) (s)
tE B, (s,
max
|lt
(10)
2i (8)
(s
|
where, without loss of generality, we assumed y 1. Indeed. when 7 * we can simply use the property (xf)* -/-"(/>)- From the envelope theorem of Rockafellar & Wets (2009, Theorem

---

Published as a conference paper at ICLR 2023
10.31), the gradients given by the corresponding argmax problems and obtain
1
( = argmax (s, t)
=1lell2 = projemnss (s)
BERTIBB
2
argmax (s, - = projes
ΛEbÂm'B.
Therefore, computing an optimal solution t* reduces to the k-sparse projections of onto the non- negative orthant and onto the simplex (scaled by > 0), respectively. When t' is not unique (i.e., contains ties), the argmax is set-valued. We discuss this situation in more details in Appendix B.2.
Fortunately, despite the nonconvexity the set BR, it turns out that both k-sparse projections can be computed exactly (Kyrillidis et al., 2013; Bolte et al., 2014; Beck & Hallak, 2016) by composing the unconstrained projection onto the original set with top-k operation:
projemnts ns, (s) = proje (topk $)) = [topk(s)]+
(11)
projaa mnB, (8) projea m (topk( 8)) (topk(s) T ml
(12)
for some normalization constant т E IR, such that the solution sums to b. Here, topk(s) is defined such that [topk())]= s; if s; is in the top-k elements of s and [topk(s)]; -oo otherwise. The k-sparse projection on the simplex is also known as top-k sparsemax (Pillutla et al., 2018; Blondel al., 2020; Correia al., 2020). Plugging these solutions back into (s, t) [ expressions in Table (a proof is given in Appendix B.4).
we obtain the
Computing (11) or (12) requires a top-k sort and the projection of a vector of size at most k. A top-k sort can be computed in O(m log k) time, projRz simply amounts to the non-negative part [+ and computing т, as needed for projaAm. can be computed in O(k) time (Michelot, 1986; Duchi al., 2008), by reusing the top-k sort. We have thus obtained efficient way compute the conjugates (2) and (3). The total complexity per LBFGS or ADAM iteration Omm log k).
Recovering plan. Assuming no ties in a + 1m cy or in a - Cj, the corresponding column of the transportation plan is uniquely determined by VΩ; (a* + 3 -c,) or V12i, (a* cy), respectively. From (11) and (12), this column belongs to Bg. In case of ties, ensuring that the plan belongs to (f b) requires to solve system of linear equations, detailed in Appendix B.2.
Unfortunately, the columns may fail to belong to B₁ in this situation.
Biconjugates and primal interpretation. As we discussed in $4 and Proposition 2, the bicon- jugates 22 and 22 allow to formally define what primal objective the transportation plans obtained by (6) and (7) optimally solve when f2 is nonconvex. Fortunately, for the case of Ω defined in (9), we are able to derive an actual expression. Let us define the squared k-support norm by
1 min XER
Ω (*)(t) (t)))((t))
(
Vi E [m].
(13)
2
The k-support norm is known to be the tightest convex relaxation of the to pseudo norm over the & unit ball and can be computed in O(m log m) time (Argyriou et al., 2012; McDonald et al., 2016). We then have the following proposition, proved in Appendix B.6.
Proposition 3. Biconjugate and primal interpretation
With Ω defined in (9), we have
94 (t) v(t) RT, (4'(0) otherwise.
v(t) 6 25 (t) = 2o otherwise.
Therefore, with Ω defined in (9), we have for all E A. b A and C R R
Da(a, b. C) Snla, = (v b. C) < Pa(a, C).
The last inequality is an equality if there are no ties in ax + 33 1m ej or in a* -c Cj Vj [n].

---

Published as a conference paper at ICLR 2023
Synthetic ID, 2=
Synthetic 2D, l==
Color transfer, k-2
0.02
Objective
semi-dual adamm
20
0.01
semii-dual lblgs dual adam dual lbfgs
0.0
0.00
50 Ieerations
100
50 Iterations
100
50 Iterations
100
0.02
0.1
Objective
20
0.01
0.0
0.00
100
100
100 Wallelock time (s)
Wallelock time (s)
Wallelock time (s)
Figure 3: Solver comparison for the semi-dual and dual formulation of sparsity-constrained OT with k 2 on different datasets. Columns correspond to datasets used in Figure 1, Figure 5, and Figure 7.
In other words, our dual and semi-dual approaches based on the nonconvex Ω are equivalent to using the convex relaxation V as regularization in the primal! We believe that the biconjugate expressions in Proposition 3 are of independent interest and could be useful in other works. For instance, shows that top-k sparsemax can be alternatively viewed as an argmax regularized with v.
Limit cases and smoothness. Let T be the solution of the quadratically regguaaaee OT (without cardinality constraints). If 2 ||t llo for jll E [n], then the constraint ||t; S & in (8) is vacuous, and therefore our formulation recovers the quadratically--rgguaiizz one. Since 12 is strongly convex in this case, both conjugates 2 and 97 are smooth (i.e., with Lipschitz gradients), thanks to the duality between strong convexity and smoothness (Hiriart-Urrui) & Lemaréchal, 1993). On the other extreme, when k = 1, we have the following (a proof is given in Appendix B.7).
Proposition 4. Limit cases
With Ω defined in (9) and k we have for all Rm
2 (8) = max [s and 1T(ss) = max 8 16 .E/m]
27 fe(m]
2
We then have for all Am, b E An and C Rmxx
Da(a, b, C) Sa b, C) So b, C) + 211b18 b,(a, 2/1b1
b,
When m < n, it is infeasible to satisfy both the marginal and the 1-sparsity constraints. Proposi- tion 4 shows that our (semi) dual formulations reduce to unregularized OT in this "degenerate" case. As illustrated in Figure 2, the conjugates ? and 9 become increasingly smooth as & increases. We therefore interpolate between unregularized OT (k small enough) and quadratically-regulati OT (k large enough), with the dual and semi-dual being increasingly smooth as k increases.
6 EXPERIMENTS
6.1 SOLVER AND OBJECTIVE COMPARISON
We compared two solvers, LBFGS (Liu & Nocedal, 1989) and ADAM (Kingma & Ba, 2015) for maximizing the dual and semi-dual objectives of our sparsity-constrained OT. Results are provided in Figure 3. Compared to ADAM, LBFGS is a more convenient option as it does not require the tuning of learning rate hyperparameter. In addition, LBFGS empirically converges faster than ADAM in the number of iterations (first row Figure 3). That being said, when a proper learning rate is chosen, we find that ADAM converges either as fast or faster than LBFGS in wallclock
8

---

Published as a conference paper at ICLR 2023
V-MoE B/32
V-MoE B/16
JFT prec@)
ImageNet 10 shots
JFT prec@1
ImageNet 10 shots
TopK (Riquelme et al., 2021)
43.47
65,91
48.86
72.12
S-BASE (Clark et al., 2022)
44.26
65.87
49.80
72.26
SPARSITY-CONSTRAINGN (ours)
44.30
66.52
50.29
72.76
Table 2: Performance of the V-MoE B/32 and B/16 architectures with different routers. The fewshot experiments are averaged over 5 different seeds (Appendix A.5).
time (second of Figure 3). In addition, Figure 3 shows that dual and semi-dual objectives very close to each other toward the end of the optimization process. This empirically confirms Proposition 3, which states that the dual and the semi-dual are equal at their optimum.
We have seen that a greater & leads to smoother objective landscape (Figure 2). It is known that a smoother objective theoretically allows faster convergence. We validate this empirically in Appendix A.2, where we see that greater k leads to faster convergence.
6.2 SPARSE MIXTURES OF EXPERTS
We applied sparsity- constrained OT vision sparse mixtures of experts (V-MoE) models for large- scale image recognition (Riquelme et al., 2021). A V-MoE model replaces few dense feedforward layers MLP æ E Rd + MLP(x) E Rd Vision Transformer (ViT) (Dosovitskiy et al., 2021) with the sparsely-gated mixture-of-expert. layers:
MoE(w:) = Σ (a)e ML,,(x),
(14)
where Gate, Rd + R4 is sparse gating function and feedforward layers (MLP,) are ex- perts. practice, only those experts M-P,(-) corresponding to a nonzero gate value Gate, (x) will be computed in this case, we say that the token æ is routed to the expert r. Upon a minibatch of m constrained OT match tokens with experts, so that the number of tokens routed to any expert is bounded. Following Clark et al. (2022), we backprop the gradient only through the combining weights Gate, but not through the OT algorithm (de- tails in Appendix A.5), as this strategy accelerates the backward pass of V-MoEs. Using this routing strategy, we train the B/32 and B/16 variants the V-MoE model: They refer to the "Base" variants of V-MoE with 32 x 32 patches and 16 x 16 patches, respectively. Hyperparameters of these archi- tectures are provided in Riquelme et al. (2021, Appendix B.5). We train on the JFT-300M dataset (Sun et al., 2017), which a large scale dataset that contains more than 305 million images. We then perform 10-shot transfer learning on the ImageNet dataset (Deng et al., 2009). Additional V-MoE and experimental details in Appendix A.5. Table 2 summarizes the validation accuracy on JFT- 300M and 10-shot accuracy on ImageNet. Compared to baseline methods, our sparsity-constrained approach yields the highest accuracy with both architectures on both benchmarks.
tokens {=1....... Fm). ), we apply our
CONCLUSION
We presented dual and semi-dual framework for OT with general nonconvex regularization. We applied that framework to obtain a tractable lower bound to approximately solve an OT formulation with cardinality constraints on the columns of the transportation plan. We showed that this frame- work is formally equivalent to using squared k-support norm regularization in the primal. Moreover, interpolates between unregularized OT (recovered when small enough) and quadratically- regularized OT (recovered when is large enough). The (semi) dual objectives were shown to be increasingly smooth as k increases, enabling the use of gradient-based algorithms such as LBFGS or ADAM. We illustrated our framework of mixture-of-expprts models in large-scale computer vision tasks, we showed that a direct con- trol of sparsity improves the accuracy, compared to top-k and Sinkhorn baselines. Beyond empir- ical performance, sparsity constraints may lead to more interpretable transportation plans and the integer-valued hyper-parameter may be easier to tune than the real-valued parameter 7.
variety of tasks; see $6 and Appendix A. For training

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank Carlos Riquelme, Antoine Rolet and Vlad Niculae for feedback on draft of this paper, as well as Aidan Clark and Diego de Las Casas for discussions on the Sinkhorn-Bas router. We are also grateful to Basil Mustafa, Rodolphe Jenatton, André Susano Pinto and Neil Houlsby for feedback throughout the project regarding MoE experiments. We thank Ryoma Sato for fruitful email exchange regarding strong duality and ties.
REFERENCES
Ravindra K Ahuja, Thomas L Magnanti, and James B Orlin. Network flows. Cambridge, Mass.: Alfred P. Sloan School of Management, Massachusetts, 1988. (Cited on page 1.)
Brandon Amos, Samuel Cohen, Giulia Luise, and levgen Redko. Meta optimal transport. arXiv preprint arXiv:2206.05222. 2022. (Cited on page 16.)
Dongsheng An, Na Lei. Xiaoyin Xu, and Xianfeng Gu. Efficient optimal transport algorithm by ac- celerated gradient descent. In Proceedings of the 36th AAAI Conference on Artificial Intelligence, volume 36, pp. 10119-10122 101208 2022. (Cited on page 20.)
Andreas Argyriou, Rina Foygel, and Nathan Srebro. Sparse prediction with the k-support norm. In Proceedings of the 26th Annual Conference on Neural Information Processing Systems, vol- ume 25, 2012. (Cited on pages 7 and 25.)
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Proceedings of the 34th International Conference Machine Learning (ICML), volume 70, pp. 214-223, 2017. (Cited on page 1.)
Amir Beck and Nadav Hallak. On the minimization over sparse symmetric sets: projections, opti- mality conditions, and algorithms. Mathematics of Operations Research, 41(1):196-223, 2016.
(Cited on pages 3 and 7.)
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Pro- ceedings of the s1 International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 880-889. PMLR, 2018. (Cited on pages 1, 3, 4, 15, and 24.)
Mathieu Blondel, André F. T. Martins, and Vlad Niculae. Learning with Sanchel-Yourg losses. Journal of Machine Learning Research, 21(35):1-69, 2020. (Cited on pages 3 and 7.)
Jérôme Bolte, Shoham Sabach, and Marc Teboulle. Proximal alternating lincarized minimization for nonconvex and nonsmooth problems. Mathematical Programming, 146(1):459-494, 2014.
(Cited on pages? 3 and 7.)
Guillaume Carlier, Adam Oberman, and Edouard Oudet. Numerical methods for matching for teams and Wasserstein barycenters. ESAIM: Mathematical Modelling and Numerical Analysis, 49(6): 1621-1642, 2015. (Cited on page 1.)
Felix Chern, Blake Hechtman, Andy Davis, Ruiqi Guo, David Majnemer, and Sanjiv Kumar. TPU- KNN: K nearest neighbor search at peak FLOP/s. In Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS), 2022. (Cited on page 19.)
Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In Proceedings of the 39th International Conference on Machine Learning (ICML), 2022. (Cited pages 3, 9, 18, and 19.)
Goncalo Correia, Vlad Niculae, Wilker Aziz, and André Martins. Efficient marginalization of dis- crete and structured latent variables via sparsity. In Proceedings of the 34th Annual Conference on Neural Information Processing Systems (NeurIPS), volume 33, pp. 11789-11802, 2020. (Cited on pages and 7.)
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Proceedings of the 27th Anmual Conference on Neural Information Processing Systems, pp. 2292-2300, 2013. (Cited on pages and 4.)
10

---

Published as a conference paper at ICLR 2023
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 248-255. IEEE, 2009. (Cited pages 9 and 19.)
Arnaud Dessein, Nicolas Papadakis, and Jean-Luc Rouas. Regularized optimal transport and the rot mover's distance. Journal of Machine Learning Research, 19(1))1-53, 2018. (Cited on pages and 3.)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. (Cited pages 9 and 18.)
John Duchi, Shai Shalev-Shwarde Yoram Singer, and Tushar Chandra. Efficient projections onto the ba-ball for learning in high dimensions. Proc. of ICML, 2008. (Cited on page 7.)
William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models deep learning- arXiv preprint ar arXiv:229.001667 2022a. (Cited on page
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-33 2022b. (Cited on page 3.)
Jean Feydy, Thibault Séjourné, Francois- Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and MMD using sinkhorn divergences. In Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AIS- TATS), pp. 2681-2690. PMLR, 2019. (Cited on page 1.)
Takumi Fukunaga and Hiroyuki Kasai. Wasserstein k-means with sparse simplex projection. In Proceedings of the 26th International Conference on Pattern Recognition (ICPR), pp. 1627-1634. IEEE, 2021. (Cited on page 3.)
Aude Genevay, Gabriel Dulac-Amold, and Jean-Philippe Vert. Differentiable deep clustering with cluster size constraints. arXiv preprint arXiv:1910.0900 2019. (Cited on page 1.)
Mann-Bapistt Hiriart-Urruty and Claude Lemaréchal. Convex analysis and minimization algorithms II, volume 305. Springer Science & Business Media, 1993. (Cited page 8.)
Jeff L Kennington and Richard V Helgason. Algorithms for network programming. John Wiley & Sons, Inc., 1980. (Cited on page 1.)
Diederík Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR), 2015. (Cited pages S and 8.)
Jun Kitagawa, Quentin Mérigot, and Boris Thibert. Convergence a Newton algorithm for semi- discrete optimal transport. Journal of the European Mathematical Society, 21(9):2033-2511 2019. (Cited on page 1.)
Wouter Kool, Chris J. Maddison, and Andriy Mnih. Unbiased gradient estimation with balanced assignments for mixtures of experts. NeurIPS I Can't Believe It's Not Better (ICBINB) Workshop, 2021. (Cited page 3.)
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. Proceedings of the 32nd International Conference on Machine Learning (ICML), pp. 957-966, 2015. (Cited on page 1.)
Anastasios Kyrillidis, Stephen Becker, Volkan Cevher, and Christoph Koch. Sparse projections onto the simplex. In Proceedings of the 30th International Conference on Machine Learning (ICML), pp. 235 243. PMLR, 2013. (Cited pages 3 and 7.)
Maksim Lapin, Matthias Hein, and Bernt Schiele. Top-k multiclass SVM. In Proceedings of the 29th Annual Conference on Neural Information Processing Systems, volume 28, 2015. (Cited on page 25.)
11

---

Published as a conference paper at ICLR 2023
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In Proceedings of the 9th International Conference Learning Representations (ICLR), 2021. (Cited on page 3.)
Mike Lewis. Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning (ICMI.), pp. 6265-62744 2021. (Cited page 3.)
Dong C Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programming, 45(1):503-555 528, 1989. (Cited on pages 5 and 8.)
Dirk A Lorenz, Paul Manns, and Christian Meyer. Quadratically regularized optimal transport. Applied Mathematics & Optimization, 83(3):1919-1999 2021. (Cited on pages and 3.)
Andrew M McDonald, Massimiliano Pontil, and Dimitris Stamos. New perspectives on k-support and cluster norms. The Journal of Machine Learning Research, 17(1):5376 5413, 2016. (Cited on pages 7 and 25.)
Quentin Mérigot. A multiscale approach to optimal transport. In Computer Graphics Forum, vol- ume 30, pp. 1592.. Wiley Online Library, 2011. (Cited on page 1.)
Christian Michelot. A finite algorithm for finding the projection of a point onto the canonical simplex of RR. Journal of Optimization Theory and Applications, 50(1):195-200, 1986. (Cited on page 7.)
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multi- modal contrastive learning with LIMoE: the language-image mixture of experts. In Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS), 2022. (Cited on page 3.)
Michael K Ng. A note on constrained k-means algorithms. Pattern Recognition, 3353):51--59,, 2000. (Cited page 20.)
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport: With applications Data Science, volume 11. Now Publishers, Inc., 2019, (Cited on pages 3 and 23.)
Venkata Krishna Pillutla, Vincent Roulet, Sham M Kakade, and 7aid Harchaoui. A smoother way to train structured prediction models. In Proceedings of the 32nd Annual Conference on Neural Information Processing Systems (NeurIPS), volume 31, 2018. (Cited pages 3 and 7.)
Francois Pitié, Anil C Kokaram, and Rozenn Dahyot. Automated colour grading using colour dis- tribution transfer. Computer Vision and Image Understanding, 107(1-2):123 137, 2007. (Cited on page 15.)
Joan Puigcerver, Rodolphe Jenatton, Carlos Riquelme, Pranjal Awasthi, and Srinadh Bhojanapalli. On the adversarial robustness of mixture of experts. In Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS), 2022. (Cited on page 3.)
Carlos Riquelme, Joan Puigcerver, Basil Mustafa. Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of ex- perts. In Proceedings of the 35th Annual Conference Neural Information Processing Systems (NeurIPS), 2021. (Cited on pages 3, 9, 18, and 19.)
optimal I Lucas Roberts, Leo Razoumov, Lin Su, and Yuyang Wang. Gini-regularized optimal transport with an application to spatio-temporal forecasting. arXiv preprint arXiv:1712.005112 2017. (Cited on page 3.)
Tyrrell Rockafellar and Roger J-B Wets. Variational analysis, volume 317. Springer Science & Business Media, 2009, (Cited on page 6.)
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason E Weston. Hash layers for large sparse models. In Proceedings of the 35th Anmual Conference on Neural Information Processing Systems (NeurIPS), 2021. (Cited page 3.)
12

---

Published as a conference paper at ICLR 2023
Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with doubly stochastic attention. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 3515-3530. PMLR, 2022. (Cited on page 1.)
Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. SuperGlue: Learning feature matching with graph neural networks. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 4938 4947, 2020. (Cited on page 1.)
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The spansely-gated mixture-of-experts layer. In Proceedings of the 5th International Conference Learning Representations (ICLR), 2017. (Cited on page 3.)
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matri- ces. Pacific Journal of Mathematics, 21(2):33--348 1967. (Cited on pages and 18.)
Justin Solomon, Raif Rustamov, Guibas Leonidas, and Adrian Butscher. Wasserstein propagation for semi sumi-supereiee learning. In Proceedings of the 3 International Conference on Machine Learning (ICML), pp. 306-314, 2014. (Cited on page 1.)
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable ef- fectiveness of data in deep learning In Proceedings of the IEEE international conference computer vision, pp.
2017. (Cited on page 9.)
Zhao You, Shulin Feng, Dan Su, and Dong Yu. SpeechMoE2: Mixture--of-expers model with im- proved routing. In Proceedings of the 47th IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7217-7221, 2022. (Cited on page 3.)
13

---

Published as a conference paper at ICLR 2023
Appendix
Regularization
Transportation plan
Preferred algorithm
Unregularized
None
Sparse
Network flow
regropy-regularized
Convex
Dense
Sinkhorn
Quadratie-regulaiized
Convex
Sparse
LBFGS ADAM
Sparsity-constrale
Non-Convex
Sparse & cardinality connsaaiiyy-
LBFGS ADAM
Table 3: Overview of how method compares others. Our method can be thought middle ground between unregularized OT and cendratically-eeglaarii OT.
A EXPERIMENTAL DETAILS AND ADDITIONAL EXPERIMENTS
A.1 ILLUSTRATIONS
OT between 2D points. In Figure 1, we visualize the transportation plans between 2D points. These transportation plans are obtained based on different formulations optimal transport, whose properties we recall in Table 3. The details of this experiment are as follows. We draw 20 samples from Gaussian distribution ~([81. [1 II) as source points; we draw 20 samples from a different as target points. The cost matrix R220x2 contains
Gaussian distribution [{] [-0.6 Euclidean distances between source and target points. The source and target marginals and b are both probability vectors filled with values 1/20. On the top row of Figure 1, blue lines linking source points and target points indicate nonzero values in the transportation plan obtained from each optimal transport formulation. These transportation plans are shown in the second row. Figure confirms that, by varying k in sparsity- formulation, control the columnwise sparsity of the transportation plan.
Unregularized Sparsity: 93.85%
Entropy regularized Spursity: 0%
Squared 2-norm regularized Sparsity: 92.68%
Sparsity constrained (ours) k=2, Sparsity: 94.4%
Figure 4: OT between two Gaussian distributions.
OT between two Gaussians. In Figure 4 above, we show transportation plans between two Gaus- sian distributions. The concrete set up of this experiment is as follows. We let Y and Z be categorical random variables taking values in {0, {O,...,31}, The realizations of Y and are the source and target locations on a 1D grid, respectively. Let Φ(2; m, s) = exp -2-mm) where 2, m, are all real 2
scalars with * O. The source distribution is set to be P(Y y) = d y: 10, 4)/ey with normal- izing constant cy 311 d(y: 10, 4): the target distribution
to be PZ = 2) = ((2; 16, 5)/ cz with normalizing constant cz == > d(2; 16, 5). The cost matrix C contains normalized squared Euclidean distances between source and target locations: C₁ (i - /) [0, 1]. OT formulation, we obtain transportation plan that
By setting = in our contains at most nonzeros per column (right-most panel of Figure 4).
14

---

Published as a conference paper at ICLR 2023
OT between Gaussian and bi-Gaussian. Similar to Figure 4, show transportation plans be- tween a Gaussian source marginal and a mixture of two Gaussians target marginal in Figure 5. We set the source distribution as P(Y = y) = o 16, 5) where cy = Zy= Φ(;; 16, 5) is the normalizing constant; set the target distribution P(Z 2) = (d(2; 5) + ((2; 242;24,55 2)) /cz, where ez = Σニ= (0(2;8, + (22 24, 5)) is the normalizing constant. Apart from that, we use the same settings Figure 4,
Unregularized Sparsity: 93.85%
Entropy regularized Sparsity: 0%
Squared 2-norm regularized Sparsity: 91.60%
Sparsity constrained (ours) k-2, Sparsity: 93.75%
Figure 5: OT between Gaussian and bi-Gaussian distributions.
A.2 SOLVER COMPARISON WITH AN INCREASED CARDINALITY
We have seen that an increased increases the smoothness of the optimization problem (Figure 2). This suggests that solvers may converge faster with an increased k. We show this empirically in Figure 6, where we measure the gradient norm at each iteration of the solver and compare the case k = 2 and k 4.
semi-dual, synthetic ID
semi-dual; synthetic 2D
semi-dual; color transfer
gradient norn
10
10
10 -2
10-2
200
400
0
200
400
200
400
dual; synthetic ID
dual; synthetic 2D
dual; color transfer
gradient norm
10
10
10"
10
200
400
400
200
400
Figure 6: Solvers converge faster with an increased k. We measure the gradient norm at each iteration of LBFGS applied to the semi-dual formulations (top row) and the dual formulations (bot- tom row) of different datasets. Since the gradient norm should go to zero, we see that LBFGS solver converges faster with an increased
A.3 COLOR TRANSFER
We apply our pparsity-constrainee formulation on the classical OT application of color transfer (Pitié al,, 2007). We follow exactly the same experimental setup in Blondel et al. (2018, Section 6).
15

---

Published as a conference paper at ICLR 2023
Figure shows the results obtained from sparsity-constramned approach. Similar to alternatives, our yields visually pleasing results.
Source
Reference
Sq. 12 norm reg.
Sparsity constrained
y=0.1
k=2
Plan sparsity: 99.2%
Plan sparsity: 99.3%
Plan sparsity: 99.2%
Plan sparsity: 99.3%
Figure 7: Result comparison the color transfer task. The sparsity indicated below each image shows the percentage of nonzeros in the transportation plan. For a fair comparison, we use = 2 for the sparsity-constraine formulation and the regularization weight y 0.1 for squared & formulation to produce comparably sparse transportation plan.
A.4 SUPPLY-DEMAND TRANSPORTATION ON SPHERICAL DATA
We follow Amos al. (2022) to set up a synthetic transport problem between 100 supply locations and 10,000 demand locations worldwide. Transport costs are set to be the spherical distance between the demand and supply locations. This transportation problem can solved via the entropy regu- larized optimal transport in Amos et al. (2022). We visualized this entropy- ruguopy-reullarized transport plan in panel (a) of Figure 8.
Building upon the setting in Amos al. (2022), additionally assume that each supplier has limited supplying capacity. That is, each supplier can transport goods to as many locations as possible up to certain prescribed limit. This constraint is conceivable, for instance, when suppliers operate with a limited workforce and cannot meet all requested orders. We incorporate this constraint into our formulation of pearsit--constranee optimal transport by specifying k as the capacity limit. The panel (b) of Figure 8 is the obtained transportation plan with a supplying capacity of k = 100 (each supplier can transport goods at most 100 demand locations).
Comparing panels (a) and (b) of Figure 8, we recognize that derived plans are visibly different in few ways. For instance, with the capacity constraint on suppliers, demand locations in Europe import goods from more supply locations in North America than without the capacity constraint. Similar observations go to demand locations in pacifie islands: Without the capacity constraint, demand locations in Pacific islands mostly rely suppliers in North America; with the capacity constraint, additional suppliers in South America are use.
16

---

Published as a conference paper at ICLR 2023
(a) Entropy-regularee transportation plan
(b) Sparsity-constrailn transportation plan
Figure 8: Plans obtained from the supply-demand transportation task. Blue lines show the transportation plan from the supply locations (yellow dots) to demand locations (blue dots). The top panel shows the transportation plan without a capacity limit supply: Each supplying location can transport to as many demand location as possible. This derived based on the entropy-regularized optimal transportation. The bottom panel shows the transportation plan with capacity limit on supply: Each supplying location can meet demands up to a fixed capacity. The plan in this case is derived by sparsity-constraing optimal transport with & 100.
17

---

Published as a conference paper at ICLR 2023
A.5 V-MOE EXPERIMENT
Our experiment is based on the vision MoE (V-MoB) architecture (Riquelme et al., 2021), which replaces few MLP layers of the vision Transformer (Dosovitskiy et al., 2021) by MoE layers. In this subsection, we review the background of V-MoE, with a focus on the router, which decides which experts get which input tokens.
We introduce a few notations that will be used throughout this subsection. Let (#1...., Fms C Rd be a minibatch of m tokens in R and let X E Rmxdd be corresponding matrix whose rows are tokens. Let W E Rdx= be learnable matrix of expert weights, where each column of W is a learnable feature vector of an expert. Common to different routing mechanisms is an token-expert affinity matrix II := XW E Rmxx; Its (i,j)-th entry is an inner-product similarity score between the i-th token and the j-th expert.
The TopK router. To route tokens to experts, the TopK router in Riquelme al. (2021) computes a sparse gating matrix I' that has at most K nonzeros per row, through a function top,, RR" + R that sets all but largest x values zero:
== top, (softmax (I+aae) E R with 11 XW.
(15)
Note that the integer k is not be confused with is used in the main text K here refers to the number of selected expert for each token and it can differ from the cardinality-constraiin
used in the main text in general. The vector e (0, 7) in (15) is noise injected to the token-expert affinity matrix XW with E R controlling the strength of noise. In practice, a is set to be 1/n during training and 0 in inference. To ensure that all experts are sufficiently trained, the gating matrix I' in (15) is regularized by auxiliary losses that encourage experts to taken similar amount of tokens in minibatch. A detailed description of these auxiliary losses is presented in Riquelme et al. (2021, Section A).
For an efficient hardware utilization, Riquelme et al. (2021) allocate a buffer capacity of experts, which specifies the number of tokens each expert can at most process in minibatch. With specified buffer capacity and a computed gating matrix r, the TopK router goes over the rows of T and assign each token to its top-chosen expert long the chosen expert's 's capacity is not full. This procedure is described in Algorithm of Riquelme et al. (2021, Section C.I). Finally, the outcomes of experts are linearly combined using the gating matrix I as in (14).
The S-BASE router. Clark et al. (2022) cast the token-expert matching problem as an entropy- regularized OT problem, solved using the Sinkhorn algorithm. This approach, dubbed as the Sinkhorn-BAN BASE (S-BASE) router, was originally designed for language MoEs that take text as input. In this work, we adapt it to vision MoEs. In direct parallel to the TopK gating matrix in (15), the gating matrix of entropy-regulariz OT is set to be
T ent == top, (Tem) E Rm
(16)
where
1en = argmin (T', -1]) + log T", with a = 1, and b = (m/n)1m.
(17)
())
The optimization plan Tent in (17) can be obtained using the Sinkhorn algorithm (Sinkhorn & Knopp, 1967). Note that while we formulated optimal transport problems with non-negative cost matrices in the main text, values in the cost matrix C = -11 in (17) can be both positive and nega- tive, following Clark et al. (2022). Since Iem is a dense matrix, heuristic is needed to select only a experts to form the gating matrix Teat this is achieved by using a top,, in (16). With a computed gating matrix Fents the B-BASE router assigns each token its top-chosen expert in the same way of the TopK router. This process allocates each expert an amount of tokens, up to a certain upper bound specified by the buffer capacity in the case of TopK. As in Clark et al. (2022), we lin- early combine the output of experts using softmax matrix softmax(II), In this way, the backward pass of saddienttbased training does not go through the Sinkhorn algorithm, can be faster and more numerically stable'.
Personal communications with the authors of Clark et al. (2022).
18

---

Published as a conference paper at ICLR 2023
The Sparsity-constrained router. We cast the token-expert matching problem sparsity- constrained OT problem. With a prescribed buffer capacity k. our goal is to upper- bound the number of tokens assigned to each expert by k. This amounts to adding cardinality constraint to each col- umn of the gating matrix:
I sparse argmin soltmaa) IT with a 1, and b (m/n)1n (18) Tel(aab) TEBxX..XB.
with Hottmax = (antmax(XX]) The purpose of the softmax function here is to obtain cost matrix containing values of the same sign. Otherwise, if a cost matrix contains both positive and negative values, then the obtained plan from sparsity-constrained optimal transport may contain zero at all entries corresponding to positive values in the cost matrix, so to minimize to the objective. In that case, columns this transportation may contain much fewer nonzeros than A - this is an undesirable situation as it under-uses the buffer capacity of experts. Note that, however, this was not an issue in the S-BASE router - cost matrix there can contain both positive and negative values (Clark al., 2022) because values of the transportation plan yielded by the Sinkhorn's 's algorithm strictly positive.
The sparse transportation plan I' spapee in (18), allocates each expert amount of tokens up to k. As in the S-BASE router, we linearly combine the output of experts using the matrix Iggotmaa-
To approximate I sparse: optimize its semi-dual proxy introduced in Section 4. We do by using an ADAM optimizer with a learning rate of 10-2 for 50 steps.
V-MoE architecture. We use the S-BASE router and our proposed sparsity-constrained router drop-in replacements of the TopK router in otherwise standard V-MoE architectures (Riquelme et al., 2021). l). We focus on the V-MoE B/32 and B/16 architectures, which use 32 x 32 and 16 x 16 patches, respectively. We place MoEs on every other layer, which is the Every-2 variant in Riquelme et al. (2021). We fix the total number of experts = 32 for all experiments. In the TopK and $- BASE router, we assign 2 experts to each expert, that is, λ. 2 in (15) and (16). The buffer capacity is set to be n/K / 32/ 32/2 16, that is, each expert can take 16 tokens at most. To match this setting, we use k 16 (18) for our sparsity-constraine router.
Upstream training and evaluation. We follow the same training strategy of Riquelme al. (2021) to train B/32 and B/16 models on JFT-300M, with hyperparameters reported in Riquelme al. (2021, Table 8). JFT-300M has around 305M training and 50,000 validation images. Since labels of the JFT-300M are organized in hierarchical way, an image may associate with multiple labels. We report the model performance by precision @ by checking if the predicted class with the highest probability is one of the true labels of the image.
Downstream transfer to ImageNet. For downstream evaluations, we perform 10-shot linear transfer on ImageNet (Deng et al., 2009). Specifically, with JFT-trained V-MoE model, we freeze the model up to its penultimate layer, re-initialize its last layer, and train the last layer on ImageNet. This newly initialized layer trained on 10 examples per ImageNet class (10-shot learning).
B/32
B/16
TopK (Riquelme et al., 2021)
97.11
308.14
S-BASE (Clark et al., 2022)
98.88
312.52
Sparsity-constrained (ours)
122.66
433.56
Table 4: Total Training TPUv2-core-days
Comparing the speed of routers. We note that the cparsity- router is slightly slower than baseline routers. One reason is that the topk function used for k-sparse projection steps. To further speedup the coursity-constraing
router, option is to use the approximated version of topk (Chern et al., 2022), which we did not use in this study. This approximated topk may be especially useful on large models like B/16, where the number of tokens large. Another way accelerate
19

---

Published as a conference paper at ICLR 2023
E-Step
M-Step
K-Means
ti <
Bargmin/CC()). softmax - [C())](.))
Soft K-Means
t2
1(+1) t (XT) Ø (laxm
Negentropy
Solve (1) with S2(t) = (t, log t)
Squared 2-norm
Solve (1) with Q(()) = 1111/
Sparsity-constrained Solve (4), (5) with ((t)) N =114 S8, (t)
Steps of the EM-like algorithm used to estimate the centers of the clusters with each method. In all cases, the cost matrix C is the squared distance between each data point and the current estimate of the centers, at a given step i.e. [C(*)]&s kc, ||ar, 1 [ 12 The vector et R
Table 5:
a canonical basis vector with [-th entry being 1 and all other entries being O.
the sparsity-constrained router is to explore different optimizers. Currently, we run the ADAM optimizer for 50 steps using learning rate 10-2 We suspect that with a more careful tuning of optimizer, one can reduce the number of steps without harming the performance. Variants of accelerated gradient- based methods (An et al., 2022) may also be applicable.
A.6 SOFT BALANCED CLUSTERING
OT viewpoint. Suppose we want to cluster m data points I .... Dm E Rd into n clusters with centroids 14,...., Mn Rdd. We let X Rdxm be matrix that contains data points æi, v.y 3m as columns. Similarly, we let μ E Rdx= be matrix of centroids.
The K-Means algorithm can be viewed as an OT problem with only one marginal constraint,
m2
min TER t T1,=a dERdx
th le - 15112
min (T,C),
TER mm xm
i=1
j=1
T1,=a #ER4X
where [C]s, ||ar, 12112 and a mm/m. Lloyd's algorithm corresponds alternating mini-
mization w.r.t. 7' (updating centroid memberships) and w.r.t. μ (updating centroid positions).
This viewpoint suggests two generalizations. The first one consists using two marginal constraints
min TEREX
(T,C)
min (T, C).
TEl((.b) peRdxn
T1, =a
1,, PERiXA
This is useful in certain applications impose preseribed size to each cluster (e.g., In/n) and is sometimes known as balanced or constrained K-Means (Ng, 2000).
The second generalization consists in introducing convex regularization
min (T,C) TEL((ab) pER
Σ 2(tt).
J=1
This moves the optimal plan away from the vertices of the polytope. This corresponds to a "soft" balanced K-Means, which we replace "hard" cluster memberships with "soft" ones. We can again alternate between minimization w.r.t. T (solving regularized OT problem) and minimization w.r.t. μ. In the case of the squared Euclidean distance, the closed form solution for the latter is ih x Σj= lijay for all [m].
When Ω is nonconvex, propose to solve the (semi) dual discussed in the main text.
Results on MNIST. MNIST contains grayscale images of handwritten digits, with a resolution of 28 x 28 pixels. The dataset split in 600 training and 10000 test images. As preprocessing.
20

---

Published as a conference paper at ICLR 2023
Cost
KI.
K-Means
161.43 = 2.20
0.227502 0.094380
Soft K-Means
157.10 + 0.19
0.040756 + 0.007071
Negentropy
157.45 + 0.04
0.000029 + 0.000002
Squared 2-norm
157.63 + 0.26
0.000018 + 0.000001
Sparsity-constrained
157.53 + 0.06
0.000013 + 0.000001
Table 6: Clustering results on MNIST using different algorithms. We report the average cost on the test split (average distance between each test image and its cluster), and the Kullhack-Leibler divergence between the empirical distribution of images per cluster and the expected one (a uniform distribution). We average the results over 20 runs and report confidence intervals at 95%. The algorithm proposed $5 achieves the most balanced clustering, and comparable cost to other OT- based solutions.
we simply put the pixel values in the range [-1, 1] and "flatten" the images to obtain vectors of 784 elements.
We use the training set to estimate the centers of the clusters using different algorithms. We use an EM-like algorithm to estimate the cluster centers in all cases, as described in Table (we perform 50 update steps). partieular, notice that only the E-step changes across different algorithms, as described in Table 5. Since there are 10 digits, we use 10 clusters.
We evaluate the performance on the test set. Since some of the algorithms produce "soft" clus- tering (all except K- K-Means), represented by the matrix T, for each test image i assign it to the cluster j with the largest value in T (t We measure the average cost (i.e. average squared distance between each image and its selected cluster), and the KL divergence between the empirical distribution of images per cluster and the expected one (a uniform distribution). The centers are initialized from normal distribution with mean of and standard deviation of 10- 9. Algo- rithms employing an OT-based approach perform 500 iterations to find T, using either the Sinkhorn algorithm (with the Negentropy method) or LBFGS (used by the rest of OT-based methods). We use cparsity-comstraint of k [1.15- n (recall that k is the maximum number of nonzeros per column). Notice that using & 12 and assuming that n is divisor of m, would necessary require that the number of nonzeros per row is 1. Thus, our minimization problem would be equivalent to that of the unregularized OT. Thus, we slightly soften the regularization.
a
Table shows the results of the experiment, averaged over 20 different random seeds. The best cost achieved by the Soft K-Means algorithm, but the resulting clustering is quite unbalanced, as reported by the KI. divergence metric. On the other hand, all OT-based approaches achieve similar costs, the algorithm based on $5 obtains significantly better balanced clustering.
B PROOFS
B.
WEAK DUALITY (PROPOSITION 1)
Recall that
Pala, C) min (T,C) Σ 2(t) TEl((a,b)
min
Σ (() + +(G))
TER T1.
1m
We add Lagrange multipliers for the two equality constraints but keep the constraint T E R explicitly. The Lagrangian then
L(T,
El,, +5),))()))))))))))) (a,T1, T1, -) (B, 1m b)
j=1
=
Σ
(ti, Cj. -a -331mmm 22(t))+ + (a,, a) (B,b).
1=1
21

---

Published as a conference paper at ICLR 2023
Using the inequality min, max. f(u, v) max, miny f(u, v) twice, we have
Pn(a, b, C) = min max max L(T, a,
TER aER" BER"
max min max L(T, a,
acR" TERMXn BER"
TER+
max max min L(T,a,
RcRT BER" TER
For the first inequality, have
nax min
max a, = max (oe, 0) + min max (3,b)
Σ
(4,, ej -- - + +s(())
OER" TERT Kn BER
aERm
TERT BER" j=1
xn
max (oe, a) + aERm
Σ
min (ti Cj a) + ((t))
Eb, A
==1
max (a, a) agRm
M (a -c²)
j=l
Sa b, C).
For the second inequality, have
max max min L(T, KT,,,,)) B) = max max (a, a) + (3,b) + R BER
min (tj, Cj. a - -Bylm) + 12(t,) t, ERR
acR" BER" TERT
j=1
max max (ar, a) + (3,b) - Σ aER BER
o (ax + +,11m -c))
j=1
= Do(a, b, C). To summarize, we showed Pa(a, C) Sn b, C) Do(a, b, C).
B.2 DUAL-PRIMAL LINK
When the solution of the maximum below is unique, tj can be uniquely determined for j e [n] from
tj VΩ (a* + +, c2) = argmax (a + 1.1m -g. t,) - Qt(t)
E
Vn2 (a* -e)) argmax
(x* Cj. t2) - ((t;).
(19)
tyE
See Table for examples. When the maximum is not unique, t is jointly determined by
tj E 022 (a" + 1m - c2) = argmax (a* + 1, c,, t,) :2(tt) EEERY
E Enn (a cz) = argmaax (a cj b1) - R(t;), t
where Ô indicates the subdifferential, and by the primal feasability T E Ua, b), or more explicitly
T T Rmxx
T*1, = a and (1 = b.
This also implies (T" 1 111=1 1. Unregularized case.
When Ω = 0, for the dual, we have
000 (8;) = argmax (8))4))
t
22

---

Published as a conference paper at ICLR 2023
We note that the problem is coordinate wise separable with
if si.j > 0
argmax fig
Si.j
R++ (0)
8i.j = 0
f.g ER+
if 81.j < 0
With 81, 0 + 88 Ci we therefore obtain
1. 20 0 if of + 8 = Ci. [ 0 if of + + < ci.j
since Si,j > 0 is dual infeasible. We can therefore use a' and 3 to identify the support of T* The size of that support constraints T* a and (T*) mm b, we can therefore form a system of linear equations of size + n to recover 2**
most + (Peyré & Cuturi, 2019, Proposition 3.4). Using the marginal
Likewise, for the semi-dual, with 8, = a Cj. we have
tj E a0, (s,)
argmax (sj, tz) EYE
conv(( uI....... U|s,
conv(Ss),
where Sj = conv(fer: E argmaaxggmm 81.))) Let us gather U1,...,00 U|sll as a matrix Vj E Rmx|ss], There exists wj ا such that t = Vjwj. Using the primal feasability, solve with respect to wy for [n]. This leads to (potentially undertermined) system of linear equations with Σ [Sy] unknowns and m + equations.
Squared k-support norm. We now discuss Ω = T as defined in (13). When the maximum is unique (no ties), 1j is uniquely determined by (19). We now discuss the case of ties.
For the dual, with 8j ax + 131, Cj. we have
t (sj)
argmax (sj, t2) st(t)) E
conv((v)........................ s,,))
= conv([[u]]++ uj E S;}),
where S, := topk(s,) ) a set containing all possible top-k vectors (the set is a singleton there are ties in 8j, meaning that there only one possible top-k vector).
For the semi-dual, with 8, a Cj. we have
tj E a0g (81)
argmmax (8,, tt) - 12(t,) t,
conv(((),...., conv((lu, -T]++ új S,}),
where is such that E Li [8 - 7yl+ = by. Again, we can combine these conditions with the primal feasability T E 2d(a, b) to obtain a system of linear equations. Unfortunately, in case of ties, ensuring that T* L( b) by solving this system may cause t; В₁. Another situation causing t & B, is if is set to a smaller value than the maximum number of nonzero elements in the columns of the primal I. solution.
23

---

Published as a conference paper at ICLR 2023
B.3 PRIMAL INTERPRETATION (PROPOSITION 2)
For the semi-dual, we have
So (a, b, C) = max (ax, a)
Ri (a ccj cz)
amR
j=l
(a,a)
22 (Mz)
acR
By M4 a
j=1
min acE" TERmxn Hy RyERm
max
(a, a)
Σ )=l
S (11) + "
(b. - a + cj)
j=c
min TER"X"
(T,C) +
max (a, T1n - a) + Σ
max MYER j=1
(t,, pby) - (06,(H1s)
acRm
min (T,C) Σ 257(4)
TER =
j=1
where we used that strong duality holds, since the conjugate is always convex, even if Ω is not. Likewise, for the dual, we have
Dn(a, b, C)
max
(e.
Σ
Mi (ax + 3m cÿ)
aER" QER",BER"
j=l
max aER" BER" 1y (4y a 1m mg
(B,b)
S4 (11)
j=1
mNax min (a.a) (B,b) aER" .ER" TeRuxn HRERR
Σ
ot (mz) + >
(li Mj --- -,, + +c²)
j=1
min (T,C) + max (a, T1, a) + max (B,T" T 1m b) + T
Σ
max (tz M5) (H,))
asRtm
BER
My ERm
j=1
min (T,C) + Σ (t)) TERmXn T1,=a 1m.=b
j=1
B.4 CLOSED-FORM EXPRESSIONS (TABLE 1)
The expressions for the unregularized, negentropy and quadratic cases provided in (Blondel et al., 2018, Table 1). We therefore focus on the top-k case.
Plugging (11) back into (s,t) ||t|| we obtain
& 2ss1
1 &
1 [s11]2
27 (8) = Σ [sp]] + [[]]
2 [sn1
>
[s
2
2
i=1
24

---

Published as a conference paper at ICLR 2023
Plugging (12) back into (s, t)
2 we obtain
k
k
1
26 (s)
[[8] - +)+
[1
7|2
2
i=1
k
k
1
11 21 (8/3) - -5)8
1l 1m [8 r)²
2
2m1
2m
Lay 2 (6/g) 72).
B.5 USEFUL LEMMAS
Lemma Conjugate of the squared k-support norm
Let us define the squared k-support norm for all t E IR by
1
m
(t) == min AER
s.t. (2,1)= Vi E [m].
2
Its conjugate for all s E Rm is the squared k-support dual norm:
(s) =
8
Proof. This result was proved in previous works (Argyriou et al., 2012; McDonald al., 2016). We include here an alternative proof for completeness.
Using Lapin al. (2015, Lemma 1), have for all
Rm
m [a u]+.
c
min kv eER
We therefore obtain the variational formulation
m [1sal
(>)
s
min ku VER
w/+
min kv + vER
Σ 1s? v|+
We rewrite the problem in constrained form
min EERT
ku + 1) s.t.
& 2
Vi [m].
We introduce Lagrange multipliers À E RT. for the inequality constraints but keep the non non-negetive constraints explicitly
(>8)
max kv + (E,1) + 8A, E)- VER. AERT
min
Using strong duality, we have
m
(
max
ssss + min ulk (A,
min (E. A)
XER R
vER
EER
2 s.t. (A,1) k. N 5 [m].
max
1=1
25

---

Published as a conference paper at ICLR 2023
We therefore obtain
M
(t)
max (8, t)
max AER
(A.1) k A 1IWE [m]
SER
min max (8, AERI sER
118 s.L. =, X - 3E m
1=1
min 4 XER"
s.t.
(, -, [m].
Using (t) (2t) gives the desired result.
Lemma 2. For all E 1Rth and b >> 0
max (s, t) tEb4 m
= min
0 + 0b 2
[s
VER
with optimality condition Σ1 [s, 0|+ b.
Proof.
max (8, t) - -0111 max min(s, t) - [1011 - O(t 1) - 6)
EER DER
max min(s 01,t) 2 || 0b
E BER
1 ||t1
min Ob + max (s BER
81,
tER
m Σis; min PER
sss 0 + Ab.
From which we obtain the optimality condition
m
8; 0 b.
B.6 BICONIUGATES (PROPOSITION 3)
We use Ω defined in (9).
Derivation of ?... Recall that
22 (s)
8 [ 12
2
im1
26

---

Published as a conference paper at ICLR 2023
Using Lemma 1, then have
(t)
max (s, t) -22((s
SER
max (s, t)
_ [1
HERm
i=1
max *ER .ERmm
(8,t) V (E) s.t. _
max min (s, t) (E) - (1s, $ E
sER EERT MER
min max (s, - 11) + max (M, E) - v/ (E) EERM
MERT aEk
(t)
i RD otherwise
x,,
Note that & is not constrained be non-negative because is squared in the objective.
Derivation of 92 Recall that
(5(ss
max (s.t) ; 111122 *E
Using Lemma 1, Lemma 2 and (12), we obtain
(g)ss
min rER 2
H - - + Tb
1=1
min EER
Lr V'(E) + Tb
s.t. 8-8-11
min YER.EERT
max & (() rb + (1, r1 E)
MER
We then have
(t)
max (8, t) - 12% (8) aERm
min max (( (1,, 1) - b) + max +(8, 1 11) + max (E, 11) - (E) PERT TER
SERm
CER
(t) ift E bA
D0, otherwise
Proof of proposition. We recall that T is defined in (13). From Proposition 2, we have Do(a, b, C) = Sn(a, C) = Pv C). From Proposition (weak duality), we have Po b, C* Pa(a, C).
Assuming no ties in a + 33 1m ej or in a cy for all j E [m]. we know that E E Bg for all j E [n.. Furthermore, from (13), we have for all t, E Ba that l(tt) vett)) 1 | 11. Therefore, without any ties, we have Py (a. b, C) = Pa b, C).
B.7 LIMIT CASES (PROPOSITION 4)
In the limit case 1 with Ω defined in (9), we have
(7((s) EEb
max
(8, t) ] [|tt1[
7
(s, 7 ||t1|
max tEbfei....mm!
(s, b2
= max
tE(e)....em
2
Y 622
b max iE[m]
2
27

---

Published as a conference paper at ICLR 2023
We therefore get
(a b, = max (a. a)
Σ
2 (a - c₁)
agRm
j=1
max (a, a)
Σ
bj max 0r Ctj + 2 ==1
63
agkm
mm
==1
b,
2 f=1
Likewise, have
2 |/
27
max (s, t)
tER nB;
2
2
max max s,t rE[m] KER+ 2
= max [s:]2.
2y iE
We therefore get
b)
max (a, a) + (3,
Σ 7i +, - Cj)
aER" BERn
n
max (a, a) + (B,b) aER" ,ER"
Σ
max (0s + 3g c.. rE(mi
2% j=1
From Proposition 3, we have Dn S2 So + 1b1b Po + 2 18
28


Reviews:
[{"review_id": "wfgvZsJTGP", "reviewer": "Reviewer_hZUP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an optimization algorithm for sparse optimal transport. The optimal transport is a linear program with additional regularization on the columns on the transportation plan matrix. In this work, the authors add further constraints to ensure each column has no more than k non-zeroes. They then propose using gradient based optimizers to directly solve the non-convex problem. ", "summary_of_the_review": "After the rebuttal and reading other reviews, I am updating my score. I seem to have misunderstood the paper's contributions earlier.\n\nSolving non-convex problems using convex duals is not new, the paper explains the setup for sparse constrained optimization in context of optimal transport. I have concerns for novelty, limited empirical evaluation and the relevance to the iclr community. ", "strength_and_weaknesses": "Strengths:\n\n-- Largely well written and the relevant works are well-cited. \n-- The duality connection is well-explained, albeit is standard for solving such problems. Similarly, the section explaining the development of the dual specific to the sparsity constrained problem is also well-explained.\n\nWeaknesses:\n\n-- I have concerns for novelty. Given enormous progress already made with iterative hard thresholding and its variants, I am not sure if they optimal transport setup adds enough novelty above that. I suggest the authors look at works of Anastasios Kyrillidis (one of whose papers the authors have cited) to look at iterative hard thresholding methods that are directly employed on the primal. Given those works, I am unfortunately not excited about this paper.\n\n-- Similarly the section on double conjugate mentioning what problem is solved when solving the dual and the relationship with k-support norms is known. \n\n-- There are limited empirical evaluations, especially considering the relevance to this conference. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XYeQ50rUaO", "reviewer": "Reviewer_yCJL", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Entropy-regularized optimal transport is very popular today, but it leads to fully-dense transport plans. To address this problem, this work proposes a new approach for OT with explicit cardinality constraints on the transportation plan. Despite the non-convexity of the cardinality constraints, the work shows that the corresponding dual or semi-dual problems are tractable and can be solved with first-order gradient methods. This method can be treated as a middle ground between unregularized OT (case k=1) and quadratically-regularized OT( k is big enough). The smoothness of the objectives increases as k increases, giving rise to a trade-off between convergence speed and sparsity of the optimal plan.\n\nThe contributions are\n1. presents a dual and semi-dual framework for OT with general nonconvex regularization. The framework is applied to obtain a tractable lower bound to solve an OT formulation with cardinality constraints on the transportation plan. \n2. shows this framework is formally equivalent to using squared k-support norm regularization in the primal, this model is the middle point between the unregularized OT and the quadratically-regularized OT.\n3. The (semi) dual objectives were shown to be increasingly smooth as k increases, enabling the use of gradient-based algorithms such as LBFGS or ADAM.\n4. the framework is demonstrated on a variety of tasks, which shows a direct control of sparsity improves the accuracy, lead to more interpretable transportation plans and the integer-valued hyperparameter k may be easier to tune.\n", "summary_of_the_review": "This work presents a dual and semi-dual framework for OT with general nonconvex regularization. The framework is applied to obtain a tractable lower bound to solve an OT formulation with cardinality constraints on the transportation plan. The work also shows this framework is formally equivalent to using squared k-support norm regularization in the primal, this model is the middle point between the unregularized OT and the quadratically-regularized OT. The (semi) dual objectives were shown to be increasingly smooth as k increases, enabling the use of gradient-based algorithms such as LBFGS or ADAM. The framework is demonstrated on a variety of tasks, which shows a direct control of sparsity improves the accuracy, lead to more interpretable transportation plans and the integer-valued hyperparameter k may be easier to tune.\n\nThe sparsity problem is fundamental in deep learning, this work gives a nice solution to the problem. The work has solid theoretic foundation, and high practical value. The mathematical formulation and proofs are rigorous and elegant. The description and analysis are clear and easy to follow. The experimental results are encouraging. Especially, this work unifies the existing methods, and present a coherent theoretic framework, this global view is very valuable.", "strength_and_weaknesses": "The work has many strengths:\n1. The problem to tackle is fundamental: control the sparsity of optimal transportation plans, which is crucial for deep learning.\n2. The idea of adding direct control to the cardinality of the images of transportation plan is novel and effective.\n3. The high level perspective, which unifies the proposed method and the existing ones, is inspiring and promising.\n4. The work has solid theoretic foundation, all the main claims are formulated as theorems and proved rigorously.\n5. The experimental results are thorough and convincing.\n\nThe weakness of the work is that it requires advanced background knowledge,  currently, detailed discussions in the supplementary materials are very helpful, it will be more helpful to give more explanations to the key concepts and the main theorems.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "NzbWWCMHlc", "reviewer": "Reviewer_aL1u", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies optimal transport with prescribed sparsity. Specifically, based on the well known quadratic regularized optimal transport, the paper further constrains the cardinality of non-zeros entries of each column of the optimal plan. The problem of optimal transport with constrained sparsity is well motivated from expert assignment problem, and based on weak duality of the proposed optimization problem, efficient algorithm is proposed, with experiments provided.", "summary_of_the_review": "The paper is overall well written and organized. Notable results are proven, with novel regularization that leads to a desired level of sparsity. Some weaknesses are listed for possible improvement. In summary The paper is a solid contribution to the sparse optimal transport problem.", "strength_and_weaknesses": "The paper has the following strengths:\n1. The paper answers, in some sense, the question of how to make optimal transport plan more sparse, especially with proposition 3 where the paper introduces the k-support norm, which is of general interest and is a more principled way to enforce sparsity.\n2. The most notable result in the reviewer's opinion, is that it reveals the relations between the k-support norm regularization, and the sparsity of the OT plan. Though the strong duality does not hold, the sparsity constraint is still nicely carried from the original primal problem into the dual solution.\n3. The paper provides pretty strong application motivation of the sparsity-constrained optimal transport, e.g. the sparse mixture of experts.\n\nSome weaknesses:\n1. Though the paper claim it proposed a 'constrained' problem, it does not solve the primal constrained problem exactly. In fact, the paper finds feasible solution(s) to the original problem through minimax principle, but no approximate optimality is provided. It would be more natural if some tightness result between primal and dual formulations can be provided, as for the current paper the 'constraint' point of view is more of a heuristic formulation rather than a problem solved in the paper. (Maybe considering the tightness between $\\ell_2$ and $\\Psi$ can lead to it.)\n2. The paper claimed that the proposed formulation has increased smoothness when the constraint is looser, but the reviewer didn't find a concrete characterization of the increase of smoothness except for a figure illustrating this phenomenon. This point would be more valid and clear if some concrete measurement of smoothness can be added, and even a formal proof of such increase of smoothness.\n\nMinor comment: the paper claimed in [section 3][computation] that the entropic formulation is not smooth, which is unclear to the reviewer, since it is well known that the dual of entropic OT is sum of exponentiation of dual potentials which is very smooth, and the whole quantity is even smooth in marginals, see e.g. https://arxiv.org/pdf/2207.08683.pdf , https://arxiv.org/pdf/2207.07427.pdf .", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2qKblr_fdi", "reviewer": "Reviewer_1nGA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Optimal transport (OT) distances (a.k.a. Wasserstein's distances) are a powerful computational tool to compare probability distributions and they play an increasingly preponderant role in machine learning, statistics, and computer vision. Since computing OT distances involves a linear program, which takes super-cubic time in the data size, its large-scale applications have mainly been hampered by its high computational cost. A classical and popular approach to tackling this problem consists in adding an entropic regularization that can be solved using Sinkhorn iterations. Comparing to the (unregularized) OT that produces a sparse transport plan, the entropic penalty blurs the OT solution, namely, it entails a dense transport plan being undesirable for interpretability. Towards realizing a sparse solution, Blondel et al. 2018 (ICML) proposed to substitute the entropy penalty with the squared-$2$norm.\n\nIn this spirit of obtaining a sparse OT, this paper proposes a novel approach by working on a nonconvex penalty with control of the sparsity. In other words, the authors define the $k$-cardinality constraint penalty that controls the cardinality of non-zeros elements in each column of the transport plan. The idea is quite \"nice\" and seems interesting for many applications involving OT metrics. ", "summary_of_the_review": "Overall, I like the approach and I think it deserves to be published at ICLR conference since it gives advances to the OT community. \n", "strength_and_weaknesses": "### Strength ###\n- To the best of my knowledge, working with nonconvex regularization on OT problems is not well-known. This paper investigates the $k$-cardinality constraint, that forces each column to get at most $k$ non-zeros elements, avoids the spreading of masses in the OT solution. I really appreciate the finite analytic results about the relations between the primal, dual, and semi-dual forms of the objective functions.\n- The approach is strongly supported by theoretical findings. \n- As the OT community is a big fan of acronyms, so I can suggest the following acronym for this paper **Top-$k$-OT**.\n- Extensive numerical experiments are conducted on sparse mixture experts with good performances and on supply-demand transportation spherical data.\n\nI have a single question: what is the computational complexity of **Top-$k$ OT**?  One can complete Table 2 by adding time for the proposed algorithm. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "08d8ZPa6AZ", "reviewer": "Reviewer_xd11", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors propose to add a non-convex sparsity constraint to the $l^2$ regularized discrete OT problem. In such a way, they can control the sparsity of the OT plan. Though the sparsity constraint makes the original problem nonconvex, it is still tractable through the dual and semi-dual problem. With the k-sparse projection, the proposed problem can be solved by L-BFGS or ADAM. Furthermore, to obtain the approximate OT plan, they turn to the convex bi-conjugate problem. Experiments show that the proposed method successfully induce a controlled sparse solution.", "summary_of_the_review": "The paper propose a novel method to solve the discrete OT problem with controlled sparsity. Experiments show that the proposed method work well. To make the results more convincible", "strength_and_weaknesses": "Strength\n- It is novel to add the k-sparsity constraint to the regularized OT problem and use the k-sparse projection to make the solution tractable.\n- To reconstruct the OT plan, the use of bi-conjugate tech is interesting.\n- The extensive experiments help show the performance of the proposed method.\n\nWeakness\n- Is there any guarantee for the sparsity of the reconstructed OT plan through bi-conjugate?\n- Since the authors are trying to solve an OT problem, it is necessary to investigate the influence of $\\gamma$ in equation (8). I am extremely curious about the convergence/performance with a small $\\gamma$. Ablation study is needed here.\n- Besides sparsity, the accuracy of the propose method should also be reported. With small scale problems like $784\\times 784$, the authors can use linear programming to obtain the groundtruth and make the comparison.\n- It is important to report the scale of the problems in the experiments. For small scale problems, linear programming may be more efficient.\n- Lack of important reference paper \"An et al. Efficient Optimal Transport Algorithm by Accelerated Gradient descent, AAAI 2022\". The authors may also try the accelerated gradient decent method to see if it helps improve the performance.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%