Wasserstein Distributionally Robust Optimization: A Three-Player Game Framework

Abstract: Wasserstein distributionally robust optimization (DRO) has recently received significant attention in machine learning due to its connection to generalization, robustness and regularization. Existing methods only consider a limited class of loss functions or apply to small values of robustness. In this paper, we present a three-player game framework for solving Wasserstein DRO problem with arbitrary level of robustness, which can handle general loss functions. Specifically, we formulate a min-max game between three players who optimize over probability measures, model parameters and Lagrange multipliers. We also propose new algorithms for finding an equilibrium of the game in convex and non-convex settings which both enjoy provable convergence guarantees. Furthermore, we prove an excess risk bound for the proposed algorithms which shows that the solution returned by the algorithms closely achieves the optimal minimax risk.  

Reviews and author responses:

This paper proposes a "three-player framework" for distributionally robust optimization (DRO), where the authors considered lifting the constraint of the ambiguity set of using the Lagrange multiplier method. The authors propose two algorithms, one for convex loss and one for non-convex, based on exponentiated gradient and gradient descent. Convergence rates are proven provided the steps in the algorithms can be carried out exactly, and some experiments are provided to demonstrate the empirical performance.

In my opinion, this paper contains a lot of strong and implicit assumptions that the authors did not address:


*****************Major issues, theoretical*****************

(i) Theorem 2 assumes that we can "update" the distribution in Algorithm 2 (for \pi_t), which is in general an impossible task. The authors merely address such an issue by saying that one can try to take samples using Metropolis-Hastings algorithm, but two problems remained unsolved. First, taking a sample of \pi_{t+1} is very different from "playing" \pi_{t+1}, which involves computing an expectation over \pi_{t+1}, and the latter is required in the theoretical analysis. Second, even if we ignore such discrepancy, the Metropolis-Hastings algorithm is known to converge extremely slow in almost all practical applications. Instead, the authors might want to try out gradient sampling methods such as Langevin dynamics https://arxiv.org/abs/1611.01838.

(ii) The authors assume the Lagrange multiplier to be bounded and did not address how this additional constraint impact the original DRO formulation. It seems to me that the boundedness of the Lagrange multiplier is purely introduced to enable regret-type analysis with bounded domains. If this is the case, the authors should make it explicit.

(iii) Both Theorem 1 and 2 require to compute, in every step of the algorithm, the best response of the Q-player exactly, which is impossible (and in fact quite daunting in practice to even get an approximation). In addition, the authors did not mention how the accuracy of approximation impact the rates.

**********************************************

By assuming (i)-(iii) above, I feel that the analysis in this work reduces to plain (exponentiated) online gradient descent and does not add much insight to the DRO problem. 


Finally, there are a lot of missing details for the experiments, where the authors did not mention how their algorithms are run (how many steps to approximate the best Q-response? how many steps for sampling? etc.). The computation aspect is totally lacking for all the benchmarks as well (how many steps are run for these methods? how much time in total did all the methods cost?)? I would suggest the authors to add these information since they are necessary to judge a new algorithm. I would also suggest the authors to use LeNet in their non-convex example as it seems to be the minimally acceptable neural networks to consider in the robust training community.

**********************************************

I thank the authors for the detailed response. After reading the rebuttal and other reviewers' comments, I feel that the consensus is the impracticality of the proposed method. I therefore retain my evaluation.

*Summary:
The papers deals with the Wasserstein DRO problem.
Their basic idea is to use a Lagrangian formulation of the problem in order to get rid of the Wasserstein distance in the constraints. Then they show how it can be solved using standard primal-dual approach while assuming an access to an oracle that can solve the transportation map.

*Detailed Comments:
-The primal dual formulation and the use of primal dual method to finding an equilibrium is quite standard. 
Translating the hard computational problem into a transportation problem is nice. Nevertheless, it preserves the hardness of the problem when $\lambda$ is small, and actually it might be the case that the optimal $\lambda$ is indeed small.

-The paper is clearly written and clearly covers prior works on this topic.

-Experiments do not establish a considerable improvement over the baselines.

*Summary of review:
The paper is clearly written and deals with an interesting problem. 
Most of the presented derivations and ideas are quite standard, there is however one nice idea which tries to handle the hardness of the problem by translating it into a  transportation map.


Post-rebuttal: I have read the authors responses and I keep my score


We thank the reviewer for the valuable comments and would like to address the reviewer’s concerns as follows.

\> Theorem 2 assumes that we can "update" the distribution in Algorithm 2 (for \pi_t), …, but two problems remained unsolved.

Thanks. Algorithm 2 only requires sampling from $\pi_{t+1}$ rather than calculating $\pi_{t+1}$ exactly. See Remark 1 in Page 5 of the manuscript. Therefore, we do not assume to update $\pi_{t+1}$ and compute the integral in the denominator when implementing the algorithm and instead use MCMC sampling method which is a technique to sample a distribution without calculating the integral in the denominator, and in Theorem 2 we only use $\pi_{t+1}$ as an intermediate step, and the final convergence results actually hold for sampled $\omega_t$. In addition, Algorithm 2 can be implemented with any sampling method which allows us to sample from $\pi_{t+1}$. In the manuscript, we just take Metropolis-Hastings algorithm as an example to show that Algorithm 2 works. We thank the reviewer for the suggestion and agree that more advanced MCMC methods like Langevin dynamics-based MCMC as the reviewer suggests might be more practical.

\> The authors assume the Lagrange multiplier to be bounded and did not address how this additional constraint impact the original DRO formulation.

Thanks. We first clarify that we do not assume that the Lagrange multiplier is bounded in the manuscript. To derive the three-player game formulation, we introduce an additional constraint $B$ on the Lagrange multiplier, which can be seen as a source of optimization error. Then in Section 3.3, we show that this optimization error caused by the introduction of $B$ can be driven arbitrary small via a careful selection of $B$. In particular, in Theorem 3, we prove that by choosing $B=\frac{(2\mu+M)^2}{\mu\tilde{\epsilon}}+\frac{2\mu+M}{\tilde{\epsilon}}$ the solution returned by our algorithms can approximately solve the original DRO problem (3) (or (4)) up to the error $3\mu$.

\> Both Theorem 1 and 2 require computing the best response of the Q-player exactly, which is impossible. In addition, the authors did not mention how the accuracy of approximation impact the rates.

We apologize for the confusion caused. In fact, in Theorem 1 and 2, we only require computing $\mu_1$-approximate best response of the $Q$-player, which is proven in Section 3.2.3 to be equivalent to solving a transportation map approximately. See the notion $BEST_{Q, \mu_1}$ in Algorithm 1\&2 and its definition in the 8th line of Section 3.2.1. Moreover, our theoretical results have shown that after a certain number of iterations our algorithms return a $(\mu:=\mu_0+\mu_1)$-approximate equilibrium, where $\mu_1$ is the accuracy parameter of $Q$-player.

\> I feel that the analysis in this work reduces to plain (exponentiated) online gradient descent and does not add much insight to the DRO problem.

We highlight our novelty as follows. First, our three-player formulation and the proposed algorithms are totally new to the DRO problem. Second, our algorithms enjoy provable convergence guarantees in both convex and non-convex cases. Finally, prior work either focuses on a limited class of losses or considers a relaxation of the DRO problem. Our framework, on the other hand, applies to general DRO problem and can also be used to develop new algorithms by designing different strategies for each player.

\> There are a lot of missing details for the experiments. And I would also suggest the authors to use LeNet in their non-convex example.

Thanks for the suggestion. We have updated the experimental section to include more implementation details and added a new Appendix E which includes the running time comparison and additional experimental results using LeNet.

We appreciate the insightful and valuable comments by all reviewers. We have revised the manuscript as suggested by the reviewers, and summarize the major changes as follows:

\> Add error bars to Figure 1 and Table 1.

\> Add a new Appendix E which includes running time comparison and additional experimental results using LeNet.

For particular comments, please see answers to individual reviewer below. We thank the reviewers again for their effort in reviewing the manuscript.

We thank the reviewer for the valuable comments and would like to clarify our contributions as follows.

\> Most of the presented derivations and ideas are quite standard.

Despite that the idea of finding an equilibrium has been investigated a lot in the literature, it has never been explored by prior studies on distributionally robust optimization. To the best of our knowledge, we are the first to formulate the Wasserstein DRO problem as a game between three-player. Our three-player game formulation and the proposed algorithms are totally new to the DRO problem. Besides, we provide theoretical guarantees for our algorithms in both convex and non-convex cases. Finally, we highlight that one of our main contribution is that we propose a framework which applies to general Wasserstein DRO problem, whereas prior work considers either restricted classes of loss functions or relaxed versions of the DRO problem.

Thank you for the valuable input. Below, we provide point-to-point responses.

\> In terms of the solution and the theory, all the methods seem to be generic methods for solving generic optimization problems.

Although finding an equilibrium is a generic idea in machine learning and has been investigated a lot in the literature, it has never been explored by prior studies on distributionally robust optimization. To the best of our knowledge, we are the first to formulate the Wasserstein DRO problem as a game between three-player. Our three-player game formulation and the proposed algorithms are totally new. In particular, it involves solving a minimization problem of a functional over the infinite dimensional space of probability measures, which can be proven to be equivalent to solving a transportation map. Besides, our algorithms enjoy provable convergence guarantees in both convex and non-convex cases. We also prove an excess risk bound which shows that the solution returned by our algorithms can approximately solve the original DRO. Finally, we highlight that one of our main contribution is that we propose a framework which applies to general Wasserstein DRO problem, whereas prior work considers either restricted classes of loss functions or relaxed versions of the DRO problem.

\> The proposed MCMC method for the non-convex problem setup is not scalable. Can you please report running times for your algorithm and compare to the other baselines?

Thanks. We have added a new Table 2 in Appendix E which reports the running time. Moreover, our algorithm 2 can be implemented with any MCMC sampling method which allows us to sample from $\pi_{t+1}$. In the manuscript, for simplicity, we just take the naïve (non-scalable) Metropolis-Hastings algorithm as an example to show that Algorithm 2 works. To improve the scalability, we may use more advanced MCMC methods instead like scalable Metropolis-Hastings [1].

[1]. Bardenet, Rémi, Arnaud Doucet, and Chris Holmes. "Towards scaling up Markov chain Monte Carlo: an adaptive subsampling approach." International Conference on Machine Learning. 2014.

\> Can you please add error bars to all tables and figures and explain whether there is any statistically significant advantage associated with the proposed algorithm?

Thanks. We have added error bars to Table 1 and Figure 1 in the experimental section. Despite that in some dataset the performance of our method is just comparable but not dominate other baselines significantly, we highlight that our three-player game formulation and the proposed algorithms are totally new to the DRO problem. Besides, our algorithms also enjoy theoretical guarantees in both convex and non-convex cases.

\> Can you please share your code with reviewers and ACs?

Thanks. We will surely reorganize and release our code once the work is published. For the pseudo-code, please see Algorithm 4 in Appendix D.

We thank the reviewer for the valuable comments and would like to answer the reviewer’s questions as follows.

\> In Lemma 1, the authors state that the strong duality holds. Is there any requirement on the l function?

Thanks. In Lemma 1, the strong duality holds for general loss functions, whether convex or nonconvex. The reason for achieving this generality is that we enlarge the solution space from deterministic models $\omega\in \Omega$ to stochastic models $\pi\in \Delta$. And by considering this larger space, we can get that $-R_Q(\pi)$ is linear in $\pi$ and thus the minimax theorem can be applied.

\> For Theorem 1 and 2, can the authors change alpha and beta/eta to a range that the results still hold?

Thanks. $\alpha$ and $\beta/\eta$ can be set to a range of values and the results still hold. For example, in Theorem 1, we can let $\alpha=constant*\frac{\mu_0}{\rho^2B}$ and $\beta=(1/2-constant)*\frac{\mu_0}{K^2}$, where $constant$ can be any number in $(0, 1/2)$. See the last three lines of the proof of Theorem 1 in Section B.1.

\> For theorem 3, is there any distributional assumption to get 1/\sqrt{n} rate? For example, if the problem is a high-dimensional estimation problem, such rate doesn't hold.

Thanks. There is no distributional assumption to get $1/\sqrt{n}$ for Theorem 3. But the constant hidden in the notion $\tilde{\mathcal{O}}$ might depend on the dimension of the data distribution.


We appreciate the insightful comments by the reviewer. Below, we provide detailed responses to the concerns.

\> Lemma 1 requires further assumptions.

Thanks. We may verify that $R_Q(\pi)$ is continuous in $\pi$ and $Q$ as follows. First given $Q$, by Holder inequality, we have $|R_Q(\pi)-R_Q(\pi’)|=|\int_{\omega}R_Q(\omega)(\pi(\omega)-\pi’(\omega))|d\omega\leq ||R_Q(\cdot)||\_{L^\infty}||\pi-\pi’||\_{L^1}\leq M||\pi-\pi’||\_{L^1}$ where $l(\omega, z)\leq M$. Then, if $\pi’\rightarrow \pi$, $R_Q(\pi’)\rightarrow R_Q(\pi)$, which shows that $R_Q(\pi)$ is continuous in $\pi$. Similarly, given $\pi$, denote $l_{\pi}(z)=\int_{\omega} l(\omega, z)\pi(d \omega)$. By Holder inequality again, we get $||R_Q(\pi)-R_{Q’}(\pi)||\leq ||l_{\pi}||\_{L^\infty}||Q-Q’||\_{L^1}\leq M||Q-Q’||\_{L^1}$. And if $Q’\rightarrow Q$, $R_{Q’}(\pi)$ goes to $R_{Q}(\pi)$, which again prove that $R_Q(\pi)$ is continuous in $Q$.

\> Algorithm 2 contains some obscure points.

Thanks. First, we clarify that both $\pi_1$ and $\pi_{t+1}$ are continuous distributions on the whole set of model parameters; Second, Algorithm 2 only requires sampling a $\omega_{t}$ from $\pi_{t}$ in each step rather than calculating $\pi_{t}$ exactly. Therefore, we do not compute the integral in the denominator and instead use MCMC sampling method which is a technique to sample a distribution without calculating the integral in the denominator. Since in each step we only need a sampled $\omega_{t}$, the final output is a discrete measure over these sampled $\omega_{t}$s. Finally, despite that the Metropolis-Hastings sampling method in subroutine 3 needs to store the whole distributions that the $Q$-player has produced, in practice we can truncate past history and use only a fixed number of timesteps. We admit that when $\lambda$ is small it is hard to solve the transport problem (9) exactly as it is a non-concave problem, and the gradient ascent can only guarantee to find a local maximum. But for large $\lambda$, computing the transportation map is a concave optimization problem, which can be efficiently solved with gradient methods.

\> The paper thus relies on a number of heuristics such as bounding lambda by B.

We clarify that introducing an additional constraint $B$ on the Lagrange multiplier is not a heuristic, but rather has theoretical guarantees. We will introduce a source of optimization error by bounding $\lambda$ by $B$. In Theorem 3, we prove that this optimization error can be driven arbitrary small via a careful selection of $B$, and by choosing $B=\frac{(2\mu+M)^2}{\mu\tilde{\epsilon}}+\frac{2\mu+M}{\tilde{\epsilon}}$ the solution returned by our algorithms can approximately solve the original DRO problem (3) (or (4)) up to the error $3\mu$.


\> Please justify the claim that Problem (3) and (2) are equivalent for convex loss (after equation (3)).

Thanks. Denote the solution to (3) by $\pi*$ and $\omega*=E_{\pi*}(\omega)$. Since the loss function (also $R_Q(\omega)$) is convex, we have $R_Q(\omega*)\leq R_Q(\pi*)$ for $Q$ and thus $\min_{\omega\in\Omega}\sup_{Q: W(Q, P)\leq \tilde{\epsilon}} R_Q(\omega)\leq\sup_{Q: W(Q, P)\leq \tilde{\epsilon}} R_Q(\omega*)\leq \sup_{Q: W(Q, P)\leq \tilde{\epsilon}} R_Q(\pi*)$. On the other hand, by definition of $\Delta$, we have $(3)\leq (2)$. Therefore, we can conclude that $(3)=(2)$.

\> The formulation (7) is thus a bit misleading because (7) can be regarded as a two-player game. Is it possible to pose a min-max-min formulation first and arrive at problem (7) as an equivalence?

We apologize for the confusion caused. Our game formulation can also be regarded as two-player game where one player is in charge of the parameters that need to be minimized, and the other player controls the parameters that need to be maximized. However, we still would like to call it “the three-player game” in the text as we find it to be a useful way to think about the problem algorithmically in that the three sets of parameters can be optimized with different algorithms; If we directly work with the original DRO problem (4) by lifting the Wasserstein constraint , a min-max-min game can be obtained. But it is not equivalent to (7) as here we cannot freely change the order between min and max in the min-max-min game.

This paper proposes two algorithms to solve Wasserstein distributionally robust optimization (DRO), posed as a min-max-max problem, with three players taking actions at each round. The authors derive some theoretical bounds and do some experiments to showcase the effectiveness of the proposed algorithms. Unfortunately, the algorithms in this paper are quite standard and not scalable. In terms of performance, the results are within statistical error of the existing baselines. The theoretical results are also standard adaptations of online learning literature. Hence, I conclude that the paper falls short of the threshold for acceptance. Here are some specific comments for the authors to improve the quality of the paper.

**Q1.** In terms of the solution and the theory, all the methods seem to be generic methods for solving generic optimization problems. Can you please explain which part of the algorithms particularize to Wasserstein DRO and would not be applicable to other types of loss functions? Could you also explain what the novelty is in the theoretical results with respect to existing results in the literature?

**Q2.** The proposed MCMC method for the non-convex problem setup is not scalable. Can you please report running times for your algorithm and compare to the other baselines?

**Q3.** The reported numbers seem to be within statistical error with baselines. Can you please add error bars to all tables and figures and explain whether there is any statistically significant advantage associated with the proposed algorithm?

**Q4.** Can you please share your code with reviewers and ACs?
Typo: In Algorithm 2 $Q_t \leftarrow \text{BEST}_{Q,\mu_1} (w_t, \lambda_t)$ needs to be updated. 


This paper consider the distributionally robust optimization (DRO) problem with Wasserstein ambiguity set. Instead of formulating the problem as a two-player game (min-max form), this paper lifts the problem to a three-player game: the first player (the decision maker) chooses the decision, the second player (the nature) chooses a probability measure, and the third player (Lagrange) chooses the multiplier. The introduction of the third player aims to remove the Wasserstein distance constraint. The paper also explore the case where the loss function is non-convex. The paper proposes iterative algorithms to solve the three-player game, along with providing guarantees (both numerical complexity and statistical).

Positive points:
- The problem is pertinent to the field (robustness, non convex loss) and any rigorous results along this line can really help pushing the field forward.


Major concerns:
- Lemma 1 requires further assumptions. The results from Sion (1958) require that R_Q(\pi) is a lower semicontinuous function in \pi, and is a upper semicontinuous function in Q. As both \pi and Q now live in functional space, these semicontinuity is not trivial (a linear function is not necessarily continuous).
- Algorithm 2 contains some obscure points:
i) What is the nature of the prior distribution pi_1? Is pi_1 a discrete, continuous or mixed distribution?
ii) What is the nature of the distribution pi_{t+1} in each iteration?
iii) If pi_1 is a uniform distribution on Omega, should pi_{t+1} be continuous as well? In that case, why is the output \hat{\pi_t} a discrete measure?
iv) There is some hidden complexity in computing pi_{t+1} for the denominator. First, evaluating the integral is difficult. Second, even subroutine 3 requires storing Q_{1}, ..., Q_{t} which is highly undesirable
v) The operation BEST_{Q, mu} relies on finding the optimal transport plan (9), and problem (9) is notoriously difficult. The authors did not resolve this difficult but "leave this problem to future work".

The paper thus relies on a number of heuristics (bounding lambda by B, using heuristic ascent to solve (9), etc.), and which renders the framework significantly less attractive.


Minor concerns:
- Please justify the claim that Problem (3) and (2) are equivalent for convex loss (after equation (3))
- The title of the paper says a three-player game, so I was expecting a min-max-min formulation. The formulation (7) is thus a bit misleading because (7) can be regarded as a two-player game. Is it possible to pose a min-max-min formulation first and arrive at problem (7) as an equivalence?


This work presents a three-player game framework for solving Wasserstein DRO problems, and it can handle general loss functions with new algorithms. I find this work interesting and potentially applicable to many problems. Below, I list my questions.

1. In Lemma 1, the authors state that the strong duality holds. Is there any requirement on the l function?

2. For Theorem 1 and 2, can the authors change alpha and beta/eta to a range that the results still hold?

3. For theorem 3, is there any distributional assumption to get 1/\sqrt{n} rate? For example, if the problem is a high-dimensional estimation problem, such rate doesn't hold.

Meta-review: The paper provides a reformulation of the distributionally robust optimization problem into a (difficult in general) transportation map problem. In the new reformulation, the authors provide strong convergence results albeit requiring strong conditions, such as solvability of the auxiliary problems in reasonable time or complexity per iteration. At this point, it is very difficult to accept the paper   given the issues that reviewers pointed. In particular, authors' method which requires an expensive MCMC implementation is not scalable. Based on the run-times reported in Table 2, the proposed method is 10x slower than their baseline (Volpi et al 2018) achieving performance within their statistical error.

Decision: withdrawn-rejected-submissions