Title: Stochastic Constrained DRO with a Complexity Independent of Sample Size
Authors: Qi Qi
Jiameng Lyu
Kung-sik Chan
Er-Wei Bai
Tianbao Yang
Venue: ICLR 2023

Abstract: Distributionally Robust Optimization (DRO), as a popular method to train robust models against distribution shifts between training and test sets, has received tremendous attention in recent years. In this paper, we propose and analyze stochastic algorithms that apply to both non-convex and convex losses for solving Kullback–Leibler divergence constrained DRO problem. Compared with existing methods solving this problem, such as primal-dual methods and large mini-batch methods, our stochastic algorithms not only enjoy competitive if not better complexity independent of sample size but also just require a constant batch size at every iteration, which is more practical for broad applications. We establish a nearly optimal complexity bound for finding an $\epsilon$-stationary solution for non-convex losses and an optimal complexity for finding an $\epsilon$-optimal solution for convex losses. Empirical studies demonstrate the effectiveness of the proposed algorithms for solving non-convex and convex constrained DRO problems. 


Review discussion:
Strength:
Overall, the paper is clear and well-organized. This work proposes four efficient dual-free algorithms to solve nonconvex and convex KL divergence constrained DRO problems. Optimal convergence rates are also established in both nonconvex and convex cases. The proposed ASCDRO and BASCDRO methods are the first optimal dual-free algorithms for the KL divergence constrained DRO problem.
Weakness:
1. It seems that there are both a Kullback-Leibler (KL) divergence constraint and a KL penalty. Is there a particular reason for doing this?
2. The author claims primal-dual algorithms require maintaining and updating an O(n) dimensional vector for updating the dual variable.
But I believe some stochastic primal-dual methods can take coordinate steps on the dual side and have per-iteration complexity O(d). Please correct me if I am wrong.
3. In the convex part of Table 1, The complexity of Dual SGM is the same as the proposed method. Could the author discuss the advantage of their method compared to Dual SGM?
Clarity, Quality, and Reproducibility are good. Novelty is minor.
NO.
**Thank you for taking time to review our paper!**
> **Q1:** It seems that there are both a Kullback-Leibler (KL) divergence constraint and a KL penalty. Is there a particular reason for doing this?
**A:**  The KL penalty with a small $\lambda_0$ is introduced so that the problem is not ill imposed.
In particular, if $\lambda_0=0$, the primal objective function (after taking the maximization of $\mathbf p$) is terms of $\mathbf w$ is non-smooth, which makes optimization much more difficult. With $\lambda_0>0$, the min-max problem is non-convex in terms of $\mathbf w$ and but strongly-concave in terms of $\mathbf p$, which makes the primal objective in terms of $\mathbf w$ is smooth (cf. Lin, Jin, Jordan, 2020). This is a well-known trick for tackling DRO problems with KL constraint.
For example, the KL regularization (called known as entropic regularization) is used to smooth the classic optimal transport problem (cf. Cuturi 2013). It originates from the prestigious Nesterov's idea of non-smooth optimization by smoothing (cf. Nesterov 2005).
> **Q2:**  In addition, the author basically reformulates the minimax problem to a minimization problem by using convex conjugate, which does not require much technical effort.
**A:** Yes, the reformulating step itself actually does not require much technical effort.
However, with the compositional formulation in hand, there are  several challenges  to be addressed for achieving optimal complexities for both convex and non-convex loss functions $\ell_i(w)$.  Please refer to response to Q1 of Reviewer rwfs.  And the main contribution of our paper is  to tackle these technical challenges and achieve optimal complexities in both convex and non-convex setting.
> **Q3:**  The author claims primal-dual algorithms require maintaining and updating an $O(n)$ dimensional vector for updating the dual variable.
But I believe some stochastic primal-dual methods can take coordinate steps on the dual side and have per-iteration complexity $O(d)$. Please correct me if I am wrong.
**A:** We would like to correct you for our problem.  Coordinate ascent update can be taken.
However, the main difficulty of the min-max formulation of our considered DRO problem is that the dual variable has a constraint $D(\mathbf p, 1/n)\leq \rho$ which couples all coordinates. If you update one dual variable according to coordinate ascent method, the resulting dual variable does not necessarily satisfying the KL constraint.
As a result, we need to perform a projection on the KL constrained domain, which inevitably suffers a complexity depending on $n$.
> **Q4:**  In the convex part of Table 1, The complexity of Dual SGM is the same as the proposed method. Could the author discuss the advantage of their method compared to Dual SGM?
**A:**   It is notable that Levy et al 2020 does not provide complexity of the considered KL-contrained KL-regularized DRO problems in the non-convex setting, which is tricky due to that  the smoothness condition needs to be established.  The advantage of our method (ASCDRO) is that it has a smaller complexity of $O(1/\epsilon^3)$ than that of Dual SGM for the non-convex DRO problem.   This is due to that Dual SGM is simply the SGD algorithm applied to the dual form. We know that even with the smoothness condition (which is tricky to be shown for our problem), the complexity of SGD is only $O(1/\epsilon^4)$.
**Reference**
> Lin, Jin, Jordan. On gradient descent ascent for nonconvex-concave minimax problems. ICML 2020.
> Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. NeurIPS, 2013.
> Nesterov, Smooth minimization of non-smooth functions. Mathematical Programming, 2005.
* **Strength**
- This paper is well-written and easy to follow.
* **Weakness**
- This method is only applicable to solving problems with a particular form.
- The KL regularization on $\boldsymbol{p}$ is not a common regular term. Eq. (1) is not a generic objective function but a deliberate one.
- Computational complexity of the proposed method is not well portrayed in the experimental studies.
* **Clarity**
- This paper is well-written and easy to follow.
* **Quality**
-  It fails to meet the requirements of ICLR.
* **Novelty**
- The novelty is limited as the objective is too specific.
* **Reproducibility**
- The study reported in sufficient detail to allow for its reproducibility.
NO.
None.
**Thank you for taking time to review our paper!**
> **Q1:** This target problem, Eq. (1), is not a general formulation solved by the existing methods. The uncertainty set of Eq. (1) must be a KL-ball. The KL regularization on $\mathbf p$ is not a common regular term. Eq. (1) is not a generic objective function but a deliberate one.
Consequently, the reviewer suggests "marginally below the acceptance threshold" as the recommendation.
**A:**   We agree with you that we focus on solving a particular form of DRO.
However, that does not necessarily diminish the contributions of this paper.
- Please refer to response to Q1 of Reviewer rwfs about difficulties of solving this problem as well as our discussion in the introduction.
- The KL regularization on $\mathbf p$ is strongly convex, which is introduced to smooth the objective function in terms of $\mathbf w$. This is commonly used for solving DRO problems with KL constraint (please refer to literature of optimal transport distances, e.g., Cuturi 2013).
- Please note that we provided an experimental result in Table 2 for comparing our algorithm with other algorithms for solving different DRO formulations including DRO with CVaR-constraint, DRO with $\chi^2$ regularization,  DRO with $\chi^2$ constraint, which demonstrates the advantage of our algorithm for solving DRO with a KL constraint with a small KL regularization.
- Finally, we would like to comment that the KL divergence is a natural divergence measure between two probability vectors. The DRO with KL divergence is natural and common, which can be used to interpret  the cross-entropy loss function.   Recently, Zhu \& Yang (2021) derived a unified view of multi-class loss functions using DRO with KL divergence regularization.
In particular, let $(f_1, \ldots, f_K)$ denote the score for $K$-classes of a data $\mathbf x$. Their DRO-based loss for a data $(\mathbf x, y)$ is defined as (neglecting the margin):
$$ \psi(f, y)=\max_{\mathbf p\in\Delta} \sum_{k=1}^Kp_k(f_k   - f_y) - \lambda D_{\text{KL}}(\mathbf p, 1/K) = \lambda \log\left(\sum_{k=1}^K\exp(\frac{f_k - f_y }{\lambda})\right).$$
It is clear that when $\lambda=1$, this reduces to the cross-entropy loss. The $\lambda$ in the  above general cross-entropy loss is usually known as temperature parameter. Using the DRO-KL formulation, we can interpret it as the Lagrangian multiplier. Our framework provides a viable way to optimize the temperature parameter in the generalized cross-entropy loss.
> **Q2:**  Computational complexity of the proposed method is not well portrayed in the experimental studies.
**A:** Thanks for your suggestion, we have added running time comparison in the appendix of the revision.
**Reference**
> Zhu \& Yang (2021). A Unified DRO View of Multi-class Loss Functions with top-N Consistency.
>  Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. NeurIPS, 2013.
The algorithm for solving the constrained DRO doesn't use dual information and hence enjoys per-iteration complexity O(d). The algorithm also enjoys relatively good iteration complexiy. One drawback in the experiment is that, the per-iteration complexiy is the main improvement from this algorithm, but there is no running time comparison which would be more interesting.
Also in Assumption 1 the smoothness assumption on $l$ is pretty strong.
The paper is well-written and clear. The results are new to the reviewer though I'm not an expert in this field.
NO.
**Thank you for your positive rating! We have revised the paper according to your suggestions.**
> **Q1:** No running time comparison which would be more interesting.
**A:**  We have added running time comparison in the appendix of the revision.
Strength
I find this paper interesting and well-written. To me, the methodology behind SCDRO is highly innovative in deriving a primal-only formulation linked with existing compositional optimization works. The developed theory for SCDOR was designed for non-convex loss functions and further extended to convex loss functions, enjoying optimal complexities for finding $\epsilon$-accurate solutions for both convex ($O(1/\epsilon^2)$) and non-convex ($\tilde{O}(1/\epsilon^3)$) smooth loss functions $\ell_i(\mathbf{w})$. Both of these complexities are first in literature in the given settings.
Weaknesses
On pp. 3 ending part, the authors essentially indicate that DRO with KL divergence regularization enjoys shown superior performance for distributional shift. End of pp. 1, "They are either restricted to problems with no additional constraints on the dual variable p except for the simplex constraint (Qi et al., 2021; Jin et al., 2021), ..." On pp. 13 middle part the authors indicate that Qi et al. (2021) firstly cast the KL-regularized DRO problems into stochastic compositional optimization. They also indicate the difference with this work that they "solve KL-constrained DRO problems, which is more challenging than KL-regularized DRO problems."
As indicated in the paper, the hardness of their analyses lies on the non-smoothness of $F(\mathbf{x})$ and the non-Lipschitzness of outer function in the given unbounded regime, which breaks the assumptions of stochastic compositional optimization theory, hindering the achievement of optimal complexities for both convex and non-convex loss functions. I am wondering if the said challenge bring essential new difficulties for this change? From the perspectives of Lagrangian's multiplier these two formulations are essentially the same. Can the authors bring more details on this comparison?
Minor comments:
Page 2, middle part, the authors said $f$ being "monotone" which can confuse readers since this work does not involve variational inequalities.
Page 13, middle part, "(i) we do not need to maintain and update ... (i) we do not need to worry ... " the second (i) should be (ii).
I praise the authors for good clarity and quality of this work. The SCDRO framework is innovative but I am concerned on the technical novelty on the relationship with earlier works, including Qi et al. (2021).
NO.
**Thank you for taking time to review our paper!**
> **Q1:**  I am wondering if the said challenge bring essential new difficulties for this change?
**A:**  The first difficulty is how to tackle the unbounded smoothness constant of the objective.  To address this difficulty, we derive an upper bound of the Langrangian multiplier $\lambda$ in Lemma 1, which allows us to establish the bounded smoothness constant of the objective in Lemma 1. The second difficulty comes from the introduced domain constraint on $\lambda$ for deriving the optimal complexities. In existing analysis of compositional optimization algorithms, either (i) the problem is assumed to be unconstrained, e.g., Qi et al (2021),   or (ii) the complexity is sub-optimal, e.g.,  Ghadimi et al.
(2020), or (iii) the problem is restricted, e.g., the outer function $f$ is convex and non-decreasing as assumed in Zhang \& Lan (2020). To the best of our knowledge, this is {\bf the first result} for stochastic compositional optimization with a domain constraint that enjoys the optimal complexities for both convex and non-convex objectives.
> **Q2:**   From the perspectives of Lagrangian's multiplier these two formulations are essentially the same. Can the authors bring more details on this comparison?
**A:**  From theoretical perspective, the min-max formulation is non-convex strongly-concave, which has been studied by many previous works.
In contrast, the equivalent compositional formulation is less studied, whose optimal algorithms are not well developed.  From practical perspective, the min-max algorithms need to handle a high-dimensional vector $\mathbf p$ per-iteration by projecting an intermediate solution onto the domain $D(\mathbf p, 1/n)\leq 1/n$, which has a complexity dependent on  $n$.
In contrast, the proposed algorithms for the compositional formulation has a complexity independent of the sample size $n$.
> **Q3:**  I am concerned on the technical novelty on the relationship with earlier works, including Qi et al. (2021).
**A:**  Please refer to response of Q1. The main difference from Qi et al. (2021) is that our problem optimizes both the model parameter and the parameter $\lambda$, which incurs new difficulties that are addressed in this work.
In contrast, Qi et al (2021) only optimizes the model parameter under a fixed $\lambda$, which is an unconstrained compositional optimization problem.
> **Q4:**  confusion about $f$ being "monotone" ?
**A:**  Sorry for the confusion. In the revision, we have change "monotone" to "non-decreasing" according to Zhang \& Lan (2020).
**Reference**
> Qi Qi, Zhishuai Guo, Yi Xu, Rong Jin, and Tianbao Yang. An online method for a class of
distributionally robust optimization with non-convex objectives. Advances in Neural Information
Processing Systems, 34, 2021.
> Saeed Ghadimi, Andrzej Ruszczynski, and Mengdi Wang. A single timescale stochastic approxima-
tion method for nested stochastic optimization. SIAM Journal on Optimization, 30(1):960–979,
> Zhe Zhang and Guanghui Lan. Optimal algorithms for convex nested stochastic composite optimiza-
tion, 2020. URL https://arxiv.org/abs/2011.10076.
We thank all reviewers for their constructive comments and questions. We believe all questions have been clearly addressed or clarified. We have submitted a revision of our paper using red color to mark the revised parts.

Meta-review:
role: PC
data: metareview:_summary,_strengths_and_weaknesses: The paper considers stochastic distributionally robust optimization problems under KL constraint and penalty and with both convex and nonconvex loss functions. The proposed algorithms improve upon primal-dual algorithms in terms of their dependence on the sample size (complexity of primal-dual algorithms would scale with the number of samples, whereas there is no such dependence in the proposed algorithms) and overall either match or improve upon the state of the art. The paper argues that the proposed algorithms are (near-)optimal in terms of the target error epsilon for approximately solving convex/nonconvex problems (finding an eps-stationary point for nonconvex problems or eps-optimal solutions for convex problems).

The overall evaluation of this work was that technically there were new results that appeared nontrivial. However, even following a discussion, it was not completely clear what was the motivation for studying the specific setting, which contained both a KL penalty and a KL constraint. The authors clarified in the discussion that the KL penalty was necessary for making the problem well posed, however, the significance of using the KL constraint remained unclear. It appeared as if the KL constraint was there only to make alternative algorithms non-competitive. However, from a modeling perspective and under the right parameter choice, there should be no difference between imposing the constraint and imposing the penalty. The paper itself did not provide a good explanation for using the KL constraint. Hence, the problem setting gave the impression of being contrived making it hard to accept the paper as is. Additionally, technical novelty compared to the cited paper Qi et al. (2021) was not completely clear and the authors are advised to more carefully compare to this work in a revision.
justification_for_why_not_higher_score: Even after all the discussion on Open Review and on Zoom, and after looking at the paper myself, I still do not understand why the paper considers this specific problem. I do think the paper has some novelty, but in its current state it is hard to judge the significance of the results. 
justification_for_why_not_lower_score: N/A
summary_of_AC-reviewer_meeting: 3 out of 4 reviewers (Z1u3, E9G7, rwfs) showed up for an online zoom meeting. Reviewer p2MP never followed up to the OpenReview message or the email I sent. Reviewer rwfs had audio problems during the call but followed up over email.

Reviewer Z1u3, who gave the highest score (8), explained their reasoning but noted low confidence and insufficient familiarity with relevant literature. None of the reviewers could explain the motivation for the setting considered in the paper and were under the impression that the setting was contrived, chosen in such a way to make other algorithms inapplicable. Reviewer rwfs agreed with this evaluation in a follow-up email and additionally noted that they were not sure about the paper's novelty compared to Qi et al. (2021).  

Decision: Reject