Drago: Primal-Dual Coupled Variance Reduction for Faster Distributionally Robust Optimization

Thank you for your concrete comments. We address them below.

I do not follow where the $bd$ term comes from in the per iteration complexity. Is it possible to provide some details?

The quantity comes from the fact that $b$ evaluations of $(\ell_i, \nabla \ell_i)$ in the computation of $v^{P}_t$ (not from $v^{D}_t$), and we assume each evaluation costs $O(d)$.

The proposed method seems to be only relevant for problems with a low dimensionality 𝑑. Otherwise, the fact that the algorithm requires to keep track of historical weights would cause a high space complexity.

Thank you for identifying this point. We have updated the manuscript to include a discussion on space complexity, which in many cases can be reduced far below $O(nd)$ for two reasons. Firstly, this is a known limitation for SAG/SAGA-type variance reduction methods for empirical risk minimization. That being said, the additional storage cost (beyond the training data) is $O(n)$ in many common cases. This is because convex losses in machine learning can typically be expressed as $\ell_i(w) = h(y_i, x_i^\top w)$ for some differentiable error function $h$ and so the gradient is given by the scalar multiplication $h’(y_i, x_i^\top w) \cdot x_i$, where $h’$ is the derivative with respect to the second argument. Thus, only these scalar derivatives need to be stored.

Secondly, a more subtle advantage of our method is that we use cyclic updates instead of randomized updates for the table of past gradients (as opposed to the randomized updates in SAGA). Conceptually, this means that for very large-scale workloads, a fraction of the gradient table can be paged out of memory until its values are updated again.

The dimensionality considered in the empirical analysis is pretty low. Can additional empirical analysis be shown where $d$ is, e.g. of a similar magnitude as $n$, if not greater than $n$?

Thank you for raising this point. We scaled up the text classification task to $d = n/2 \approx 4000$ in the attached PDF (see Figure 5). This means that the batch size is approximately $2$. Qualitatively, similar trends are observed across varying orders of magnitude for $d$, in that we show improvement over LSVRG baseline in the attached PDF. We hypothesize that the benefit over LSVRG comes from the fact that the dual variable is updated in epochs of length $n$ iterations to balance out the $O(n)$ cost of the update. We incur a higher per-iteration cost but have improved theoretical and empirical performance over the training trajectory.

Given that my questions have been satisfactorily addressed, I'd like to raise my score.

Thank you for your comments. We address the main concern below.

As the author has commented, they only solve for convex optimization formulation. The extension to non-convex setup also seems interesting and important.

Thank you for raising this point. Broadly, the study of optimization approaches for the non-convex setting (e.g., deep learning) relies on completely different technical tools and evaluations. Convexity is considered to make precise statements about global optimality, namely the linear convergence of the algorithm’s iterates to a solution in terms of problem constants. In the non-convex setting, one seeks convergence to a first-order stationary point under a gradient-domination or Polyak-Łojasiewicz condition or for large deep learning models with Gaussian random weights or other strong assumptions. Our focus here is on sound theoretical convergence results, for convex optimization algorithms, one of the subject areas of NeurIPS.

Thank you for your insightful comments. We address them below.

SAG-type methods usually have high memory requirements as $d$ increases.

Thank you for identifying this important point. We have updated the manuscript to include a discussion on space complexity, which in many cases can be reduced far below $O(nd)$ for two reasons. First, this is a known limitation for SAG/SAGA-type variance reduction methods for empirical risk minimization. That being said, the additional storage cost (beyond the training data) is $O(n)$ in many common cases. This is because convex losses in machine learning can typically be expressed as $\ell_i(w) = h(y_i, x_i^\top w)$ for some differentiable error function $h$ and so the gradient is given by the scalar multiplication $h’(y_i, x_i^\top w) \cdot x_i$, where $h’$ is the derivative with respect to the second argument. Thus, only these scalar derivatives need to be stored. Second, there is a more subtle difference between our method and SAG-type methods. We use cyclic updates instead of randomized updates for the table of past gradients (as opposed to the randomized updates in SAGA). Conceptually, this means that for very large-scale workloads, a fraction of the gradient table can be paged out of memory until its values are updated again.

The experimental evaluation is small and limited to relatively low-dimensional problems.

To illustrate performance, we scaled up the text classification task to $d = n/2 \approx 4000$ and show the improvement over LSVRG in the attached PDF (see Figure 5). This means that the batch size is approximately $2$. We hypothesize that the benefit over LSVRG comes from the fact that the dual variable is updated in epochs of length $n$ iterations to balance out the $O(n)$ cost of the update. We incur a higher per-iteration cost but have improved theoretical and empirical performance over the training trajectory.

The paper organization makes it quite difficult to understand the different algorithmic components used in Drago.

Thank you for the helpful suggestions on paper organization. In the revised version, we have introduced the algorithm early on and described it using the reference points you have mentioned. We have dedicated more of the main text toward the comparisons noted in Appendix B.

The assumptions necessary for convergence are somewhat unrealistic, although these appear to be standard in the literature. I think that assuming each $\ell_i$ is both Lipschitz continuous and Lipschitz smooth is quite unrealistic. Typically only one or the other of these conditions holds.

Because the losses exist within the statistical learning context, assuming both Lipschitz continuity and smoothness is valid in common scenarios. A canonical loss in robust statistics is the Huber loss, which is both Lipschitz and smooth. Similarly, any smooth loss will be Lipschitz continuous if $\mathcal{W}$ is compact. Compactness is a common assumption, as statistical guarantees (uniform convergence, minimax rate optimality, etc.) rely on it.

Isn't $nq_{\max{}} \leq n$ trivially satisfied since 𝑞 is constrained to the probability simplex?

We have added to Appendix B a precise claim regarding the fact that $n q_{\max{}}$ is constant. In summary, we mean that $nq_{\max{}}$ is bounded by a constant independent of $n$. The two canonical examples of $\mathcal{Q}$ are divergence-ball based and spectral risk measure-based feasible sets. It is true that $\mathcal{Q}$ is always a subset of the probability simplex so $q_{\max{}} \leq 1$. However, this bound is only tight when $\mathcal{Q}$ is the entire probability simplex, regardless of the choice of $n$. Because this feasible set is often defined using statistical properties of the weights $q$, each element can be upper bounded as $q_i \leq c/n$ for some $c \geq 1$. For spectral risk measures, this can be seen by the argument in Section 2 of Mehta (2023) - every weight will be the integral of a bounded function over an interval of length $1/n$. For the $\theta$-CVAR (see our Appendix B.2$, this bound is$1/(\theta n)$. For divergence balls, in [Namkoong and Duchi (2017)](https://papers.nips.cc/paper_files/paper/2017/hash/5a142a55461d5fef016acfb927fee0bd-Abstract.html) Eq (4), the divergence ball ambiguity set is defined using radius$\rho/n$, which enforces the condition$\frac{n}{2} \Vert q - \mathbf{1}/n \Vert _2^2 \leq \frac{\rho}{n}$on every element of$\mathcal{Q}$in the case of$\chi^2$-divergence.

Why is $b=1$ only show for three datasets and only in terms of oracle complexity?

As mentioned toward the end of section 4.1, the $b=1$ is only competitive with baselines for small datasets. As the dataset size grows, the batch size needs to be increased to maintain performance. While the theoretically optimal batch size $n/d$ is a good default value, $b$ can be tuned experimentally to achieve the best performance (with respect to $b$) for a particular empirical setting.

Display after Line 124: It looks you are using the equality ... but I don't think this is true because the terms in the product are correlated.

This equality holds because the expression in the expectation is the inner product between two vectors $\ell(w\_{t+1}) - \hat{\ell}\_t$ and a “one-hot” vector $(\ell\_{t, j\_t} - \hat{\ell}\_{t-1, j\_t})e\_{j\_t}$ (see Eq (8)). Thus, the only terms that are multiplied are in the $j_t$ coordinate.

Minor Comments: ...

Thank you for the additional “Minor Comments”. We have incorporated them into the manuscript. To address the questions:

• Line 102: This is the full Jacobian - thank you for the suggestion.
• Algorithm Description: These constants are defined by simplifying the algorithm description used in the analysis. We have made the algorithm description consistent with the analysis so that these constants can be avoided.

Thank you for your detailed suggestions—we have incorporated the comments about comparisons, formal statements of complexity results, and minor typos into the manuscript.

The cases with $\mu=0$ or $\nu=0$ are handled in Appendix C.4, but not in a very formal way.

Thank you for raising this point. We have changed the analysis to account for the settings of $\mu = 0$, $\nu = 0$, and $\mu = \nu = 0$ formally. Because achieving an unconditional linear convergence rate that improves upon gradient descent was the main motivation of the work, we focused on the strongly convex-strongly concave setting. In the final version, we will provide a unified analysis for all settings.

...it is always best to choose $b=1$ to get the best first-order oracle complexity, but your runtime may be better with a different choice of $b$?

That is correct. If we assume that the cost of querying the oracle is $O(d)$ and $n \geq d$, then taking into account the $O(n)$ cost of updating the dual variable in each iteration, a batch size of $b = n/d$ balances the costs of the two updates. This affords us a $O(n^{3/2})$ dependence on $n$ in the overall runtime.

Line 81 says that $n q_{\max{}}$ is upper bounded by "a constant in 𝑛," but isn't it just bounded by $n$ because $q_{\max{}} \leq 1$?

We have added to Appendix B a precise claim regarding the fact that $n q_{\max{}}$ is constant. In summary, we mean that $nq_{\max{}}$ is bounded by a constant independent of $n$. The two canonical examples of $\mathcal{Q}$ are divergence-ball based and spectral risk measure-based feasible sets. It is true that $\mathcal{Q}$ is always a subset of the probability simplex so $q_{\max{}} \leq 1$. However, this bound is only tight when $\mathcal{Q}$ is the entire probability simplex, regardless of the choice of $n$. Because this feasible set is often defined using statistical properties of the weights $q$, each element can be upper bounded as $q_i \leq c/n$ for some $c \geq 1$. For spectral risk measures, this can be seen by the argument in Section 2 of Mehta (2023)—every weight will be the integral of a bounded function over an interval of length $1/n$. For the $\theta$-CVAR (see our Appendix B.2), this bound is $1/(\theta n)$. For f-divergence balls, in Namkoong and Duchi (2017) Eq (4), the divergence ball ambiguity set is defined using radius $\rho/n$, which enforces the condition $\frac{n}{2} \Vert q - \mathbf{1}/n \Vert _2^2 \leq \frac{\rho}{n}$ on every element of $\mathcal{Q}$ in the case of $\chi^2$-divergence.

Is this the first work you are aware of which applies these cyclic coordinate-style updates to a non-separable dual feasible set?

We are not aware of any previous work that does so in a cyclic fashion. The key element at play is that $\hat{q}_t$ (see Algorithm 1) does not need to be in the feasible set to be used in the gradient estimate in the primal update. Thus, we may apply cyclic updates to this vector even if we cannot do the same for $q_t$.

...a classic trick to recover a non-strongly convex rate from the strongly-convex case is to just set something like, e.g., $\mu = \frac{\epsilon}{2D}$ where $D$ is a bound on the diameter of $\mathcal{W}$.

Thank you for identifying this trick. This would be one way to approach this, but it relies on knowing $D$, the diameter or initial distance-to-optimum in an unbounded setting. Furthermore, it leads to an extra log factor on the complexity. We chose to provide analysis for a direct method, which does not require knowledge of $D$ and leads to the improved convergence rates.

Thank you for these clarifications. The one answer I still don't fully follow is the one related to $n q_{max}$. I now realize Line 81 is connected to Line 468 in Appendix B. Line 468 reads: "$n q_{max}$ is upper bounded by a constant independent of $n$ in common cases" and then goes on to cite $\theta$-CVaR and $\chi^2$-DRO with radius $\rho$ as examples. I don't agree with this though in general. For $\theta$-CVaR for example, the constraint is $q_{max} \le \frac{1}{\theta n}$ as you wrote in your response. Then it is true that $n q_{max}$ is a constant if you restrict to $\theta = \Omega(1)$. But to my knowledge, the $\theta$-CVaR DRO problem encompasses all $\theta \in [1 / n, 1]$, in which case it is not correct to claim that $n q_{max}$ is always a constant for $\theta$-CVaR. (Indeed, the entire range $\theta \in [1 / n, 1]$ is considered in Levy et al. 2020. And it is even pointed out later in your Appendix B - Line 489 - that previous methods in the literature don't beat gradient descent when $\theta = 1 / n$.)

It doesn't seem to me that the strength of your contribution relies on $n q_{max}$ being a constant as opposed to being as large as $n$ per the rates in Table 1. However, I would avoid claiming $n q_{max}$ is a constant in general for $\theta$-CVaR, $\chi^2$-DRO, etc. It seems to me like the point you want to make (?) is that if $\theta = \Omega(1)$, then your rate improves, but then you must make the $\theta = \Omega(1)$ restriction clear (and similarly for $\chi^2$-DRO, etc.).

Please let me know if I'm still misunderstanding something however.

(Also just to add to my original question - I would suggest avoiding the phrase "a constant in $n$" and just say "a constant" or "a constant independent of $n$." I understand it now, but when I first read "a constant in $n$," I thought it meant $cn$ for some constant $c$.)

Thank you for your thorough comments. We address them below.

Is it possible to extend the analysis to non-strongly convex settings, or the case with convex individual functions but strongly convex summation?

Thank you for raising this important point. We discuss in Appendix C.4 modifications to the analysis which would handle if either or both of the strongly convex (strongly concave) parts of the objective were changed to convex (concave). In summary, we would not achieve linear convergence of course, but $O(1/t^2)$ if either $\mu = 0$ or $\nu = 0$ or $O(1/t)$ if $\mu = \nu = 0$. Because achieving an unconditional linear convergence rate that improves upon gradient descent was the main motivation of the work, we focused on the strongly convex-strongly concave setting. In the final version, we will provide a unified analysis for all settings.

The “strongly convex summation” could be handled with a similar overall structure, but different inequalities to produce the upper and lower bounds on the primal-dual objective function. For the lower bound, rather than applying the smoothness of $\ell_i$ and then the definition of the proximal operator, we may apply an interpolation inequality that combines smoothness and strong convexity. One challenge is that we require $\ell_i$ to be Lipschitz continuous, so we would need to assume a compact domain for it to also be strongly convex.

...if we only track the dependence on $n$, it's worse than the vanilla subgradient methods.

The vanilla subgradient method has an $n^2$ dependence due to the second term, whereas our method achieves an $n^{3/2}$ dependence.

In many cases, this costs the same complexity as storing all the training data.

Thank you for identifying this point. We have updated the manuscript to include a discussion on space complexity, which in many cases can be reduced far below $O(nd)$ for two reasons. Firstly, this is a known limitation for SAG/SAGA-type variance reduction methods for empirical risk minimization. That being said, the additional storage cost (beyond the training data) is $O(n)$ in many common cases. This is because convex losses in machine learning can typically be expressed as $\ell_i(w) = h(y_i, x_i^\top w)$ for some differentiable error function $h$ and so the gradient is given by the scalar multiplication $h’(y_i, x_i^\top w) \cdot x_i$, where $h’$ is the derivative with respect to the second argument. Thus, only these scalar derivatives need to be stored.

Secondly, a more subtle advantage of our method is that we use cyclic updates instead of randomized updates for the table of past gradients (as opposed to the randomized updates in SAGA). Conceptually, this means that for very large-scale workloads, a fraction of the gradient table can be paged out of memory until its values are updated again.

In order to minimize the total computational cost... what is the optimal choice of the minibatch size?

The optimal batch size is $b = n/d$, which is used in the red lines in the experiments. As mentioned in the “Global Complexity” paragraph on page 7, because the dual variables must normalize to 1, we have an $O(n)$ computational cost per iteration regardless. Thus, under the assumption that $n \geq d$, we can compute $n/d$ gradients at $O(d)$ cost without changing the per-iteration complexity. This allows for the $n^{3/2}$ dependence in the sample size in Table 1, as opposed to $n^2$ for gradient descent.

If we choose the parameters optimally, how does the total computational cost compare to other algorithms? In particular, as the per-iteration cost depends on $n$ as well, it is important to compare with the full-batch gradient-based algorithms.

Thank you for the suggestion; we have added a table for this quantity as well for clarity (see the supplemental PDF Table 6). Among the linearly convergent methods of the manuscript, this can be measured by multiplying the complexity of gradient descent, LSVRG, and Drago by $d$ and multiplying the complexity of LSVRG/Prospect by $n+d$. Thus, under $b = n/d$, Drago achieves the best complexity without conditions on $\nu$. A comparison to general-purpose primal-dual min-max algorithms is also given in Table 4 of the manuscript, with runtime.