Distributionally Robust Optimization with Bias and Variance Reduction
Authors: Ronak Mehta
Vincent Roulet
Krishna Pillutla
Zaid Harchaoui
Venue: ICLR 2024 spotlight (2024)
Keywords: stochastic optimization
convex optimization
distributionally robust learning
spectral risk measures
incremental optimization
Primary area: optimization
TL;DR: We propose a linearly convergent stochastic (a.k.a. incremental) algorithm that optimizes spectral risk measures, which are distributionally robust objectives that include the superquantile.

Abstract: We consider the distributionally robust optimization (DRO) problem, wherein a learner optimizes the worst-case empirical risk achievable by reweighing the observed training examples. We present Prospect, a stochastic gradient-based algorithm that only requires tuning a single learning rate hyperparameter, and prove that it enjoys linear convergence for smooth regularized losses. This contrasts with previous algorithms that either require tuning multiple hyperparameters or potentially fail to converge due to biased gradient estimates or inadequate regularization. Empirically, we show that Prospect can converge 2-3x faster than baselines such as SGD and stochastic saddle-point methods on distribution shift and fairness benchmarks spanning tabular, vision, and language domains.

Published as a conference paper at ICLR 2024
DISTRIBUTIONALLY ROBUST OPTIMIZATION WITH BIAS AND VARIANCE REDUCTION
Ronak Mehta
Vincent Roulet
Krishna Pillutla
Zaid Harchaoui'
ABSTRACT
We consider the distributionally robust optimization (DRO) problem. wherein a learner optimizes the worst-case empirical risk achievable by reweighing the ob- served training examples. We present Prospect, a stochastic gradient-based algo- rithm that only requires tuning a single learning rate hyperparameter, and prove that it enjoys linear convergence for smooth regularized losses. This contrasts with previous algorithms that either require tuning multiple hyperparameters potentially fail to converge due to biased gradient estimates or inadequate reg- ularization. Empirically, we show that Prospect converge 2-3x faster than baselines such as SGD and stochastic saddle- saddle-point methods on distribution shift and fairness benchmarks spanning tabular, vision, and language domains.
INTRODUCTION
The ingredients of empirical risk minimization (HRM) are generally considered be: model with parameters w E TRd (e.g. a neural network), a training set 21,..., En E 2 of independent and identically distributed realizations of random variable 7 P, loss function & Rd 2 R, and an optimization algorithm that solves
min Ez~PA [8(w, Z)]. WER
(1)
where Pn is the empirical distribution of {-:)7-1. The fourth ingredient-often taken for granted-is the choice of risk functional, which aggregates the distribution of E(w, 2) into univariate summary to be minimized. The objective (1) (the expected loss under P,) is an unbiased estimate of the expected loss under an underlying gettereeneraiing distribution P: however, deployed model often observes data from distributions other than P. Motivated by this practical phenomenon, we consider instead an objective that explicitly captures sensitivity to such distribution shifts:
min max Ez-q (f(w, Z) - (.())Pp), wERd QEQ
in which Q is an uncertainty set of probability measures, V/ 2 is hyperparameter, and P(Q represents the divergence of Q from the original training distribution P. (e.g. the x2 or Kullback Leibler divergence). The objective (2) emulates game in which nature pays a price of per unit D(Q to replace the expected loss under P., with the expected loss Ez-q [E(w, Z)] associated with the shifted distribution Q. Since penalizes these shifts, we shall refer to it as the shift cost.
Objectives of the form (2), known as distributionally robust optimization (DRO) problems, have seen wave of recent interest in machine learning theory and practice (Chen & Paschalidis, 2020). Historically used in quantitative finance, a popular such objective is the conditional value-at-risk (CVaR, a.k.a. superquantile/xxpeeted shortfallaaverge top-k loss). In terms of methods, the CVaR has been used as a canonical DRO objective (Fan et al., 2017: Kawaguchi & Lu, 2020; Rahimian & Mehrotra, 2022), well as in unsupervised learning (Maurer et al., 2021), reinforcement learning (Singh et al., 2020), and federated learning (Pillutla et al., 2023). In applications, it has also been employed for robust language modeling (Liu et al., 2021) and robotics (Sharma al., 2020). The superquantile/CCaa falls into the broader category of spectral risk measures (SRMs), a class of DRO objectives that includes the extremile and exponential spectral risk measure (ESRM) (Acerbi & Tasche, 2002; Cotter & Dowd, 2006; Daouia et al., 2019).
"University of Washington, 'Google DeepMind, 'Google Research.

---

Published as a conference paper at ICLR 2024
Motivated by 1) the success of the superquantile in applications and 2) the importance of stochastic optimization in machine learning. the principal goal of this paper is to develop stochastic' algo- rithms for spectral risk minimization.
Contributions. We propose Prospect, a stochastic algorithm for optimizing spectral risk measures with only one tunable hyperparameter: a constant learning rate. Theoretically, Prospect converges linearly for any positive shift cost on regularized convex losses. This contrasts with previous stochas- tic methods that fail to converge due bias (Levy et al., 2020; Kawaguchi & Lu, 2020), may not converge for small shift costs (Mehta et al., 2023), or have multiple hyperparameters (Palaniappan & Bach, 2016). Experimentally, Prospect demonstrates equal or faster convergence than competi- tors on the training objective on nearly all objectives and datasets considered, and exhibits higher stability with respect to external metrics on faimess and distribution shift benchmarks.
Related Work. Examples of DRO formulations range throughout diverse contexts such as rein- forcement (Liu al., 2022b; Kallus et al., 2022: Liu et al., 2022c; Xu et al., 2023; Wang et al., 2023; Lotidis et al., 2023; Kallus et al., 2022; Ren & Majumdar, 2022; Clement & Kroer, 2021), continual ( (Wang et al., 2022), interactive (Yang et al., 2023; Mu et al., 2022: Inatsu et al., 2021: Sinha et al., 2020), Bayesian (Tay et al., 2022; Inatsu et al., 2022), and federated (Deng et al., 2020; Pillutla et al., 2023) learning. along with dimension reduction (Vu et al., 2022), computer vision (Samuel& & Chechik, 2021: Sapkota et al., 2021), and structured prediction (Li et al., 2022; Fathony et al., 2018). Various forms of these objectives are parameterized by the uncertainty set Q. includ- ing those based on f-divergences (Levy et al., 2020; Ben-Tal et al., 2013), the Wasserstein metric (Blanchet et al., 2019b; Kuhn et al., 2019), maximum mean discrepancy (Kirschner et al., 2020; Staib & Jegelka, 2019; Nemmour al., 2021), more general classes of metrics (Husain, 2020; Shapiro, 2017). We focus on SRM objectives, as motivated in detail in Sec. 2 and Appx. B.
These objectives also yield connections other areas in modern machine learning. They are spe- cial case of subpopulation shift, wherein the data data-generating distribution is modeled as mixture of subpopulations, and the distribution shift stems from changes the mixture. In our case, the sub- populations are point masses at the observed data points. In the context of algorithmic fairness, the subpopulations may represent data conditioned on some protected attribute (e.g. race, gender, age range), and common notations of fairness such as demoggappic/statiiiice parity (Agarwal et al., 2018; 2019) impose (informally) that model performance with respect to each subpopulation should be roughly equal. As such, robustness reweighting and algorithmic faimess often aligned no- tions ( (Williamson & Menon, 2019). with recent research arguing that distributionally robust models more fair (Hashimoto et al., 2018: 8; Vu et al., 2022) and that fair models are more distribution- ally robust (Mukherjee et al., 2022). In supervised learning, the data distribution is modeled as P = Px,y for feature-label pair (X, Y) and related settings of covariate shift (changes in Px and not Pylx) (Sugiyama et al., 2007) as well as label shift (changes in Py and not Pxly) (Lipton et al., 2018) may also modeled with distributional robustness (Zhang et al., 2021) as illustrated in Fig. 1. In these settings, distributional robustness may be described as a property of learned representations that are transferable to multiple tasks (Stowik & Bottou, 2022).
In comparisons, we include stochastic algorithms that either are pogle-byperparameter "out-of-the- box" methods such as stochastic gradient descent and stochastic regularized dual averaging (Xiao, 2009), or multi-hyperparamen jectives, such as LSVRG (Mchta et al., 2023) and stochastic saddle-point SAGA (Palaniappan & Bach, 2016). Note that LSVRG may not converge for small shift costs. Other methods may achieve sublinear convergence rates, even with multiple hyperparameters (Yu et al., 2022: Ghosh et al., 2021; Carmon & Hausler, 2022; Li et al., 2019; Shen et al., 2022; Yazdandoost Hamedani & Jalilzadeh, 2023; Namkoong & Duchi, 2016). Non-convex settings have also been studied (Jin et al., 2021; Jiao et al., 2022; Sagawa et al., 2020; Luo et al., 2020), as well statistical aspects of resulting minimizes of DRO objectives (Liu et al., 2022a; Blanchet et al., 2019a; Zeng & Lam, 2022; Maurer et al., 2021; Lee et al., 2020: Khim al., 2020; Zhou & Liu, 2023; Zhou et al., 2021: Cranko al., 2021; Slowik & Bottou, 2022). Our goal is to achieve unconditional linear convergence for smooth, strongly convex (regularized) losses with a single hyperparameter.
methods that converge linearly on strongly convex SRM-hased ob-
We use stochastic interchangeably with incremental, meaning algorithms that make O11) calls per iteration a fixed set of oracles (f4, ( and not streaming algorithms that sample fresh data at each iteration.
2

---

Published as a conference paper at ICLR 2024
Original Distribution
Subpopulation Shift
Label Shift
Adversarial Reweighting
Uniform weight examples.
Weight Group
Weight shifts toward positive class.
Weight shifts arbitraniy..
Posibive Class (Group
Positive Class (Group
Negative Class (Group
Negative Class (Group
Notions of Distribution Shift. Illustration of various forms of distribution shift that are characterized by maintaining the same training data but changing the weight of each example.
Figure 1:
MINIMIZING SPECTRAL RISK WITH BIAS AND VARIANCE REDUCTION
We first describe the key technical challenges in constructing a stochastic optimizer for spectral risk measures and how Prospect tackles them. Because we consider a fixed training set of size n and use f-divergences for D (reviewed in Appx. B), we may parameterize probability measures Q by a set of weights P C {p R" p; 2 0, Σx= =1} indicating the probability mass allotted to training examples (21, (21,..., zn). Similarly, the empirical distribution
=
Pa is represented by the 1/n). We can then rewrite (2) in the convenient form
uniform weights 1, = (1/n, ...,, ...,
min Rx (f(w)) for R() == max { EL gEP
q:4e DD(allan/ni.
wER'
where f(w) ( (f) (w)) E R" and (((w) = f(w, ) is the loss on training example i. In order to build a convergent stochastic algorithm, we will construct an estimate v for the gradient of (3) based on a single data index i, such that e; + (eR)((()) as the iteration counter approaches infinity. Precisely, we require that for i ~ Unif(n].
[E| V.RPp((()) - wy 1 = VV (l(w)) E (r: 112 + EE [[ - Us
(4)
bias
variance
decreases to zero asymptotically. In the remainder of this section, we first identify concretely our target estimand (i.e. (((()( for the spectral risk uncertainty set), construct an estimate, and then describe individual procedures to ensure that the bias and variance terms in (4) vanish.
Spectral Risk Measures and their Gradients. The conventional p-superquantil/CCaR (Rock- is defined as Σ (f where (w) < <... < (t)) are the ordered losses and a; = 1/k for the largest k = np val- ues of and a, = 0 otherwise (see Equation (21)). Other spectral risk measures (SRMs) are generated by constructing any vector called the spectrum. Traditionally used in economies and
afellar & Royset, 2013) of a loss vector (& (w)
(aI...., Gn) of non-negative weights satisfying
σi S S On and > o, finance, SRMs strike a user-defined balance between measuring average and tail loss. Recogniz- ing that Σm_ ( is the maximum of Σ (w)) for q permutation of (01...... Gn), or equivalently, convex combination of such permutations, ConvexHull perrutttiion of o}} to achieve the regularized formulation in (3). See Appx. B for fur- ther details on SRMs including the extremile (Daouia al., 2019) in Equation (22) and ESRM (Cot- ter & Dowd, 2006) in Equation (23). SRMs are motivated not only by their historical use but also the efficiency with which we can compute their exact gradient, as we now describe.
define the uncertainty set P P())
Define R2 : Rp).). When > 0 and the map q + D(g) 1,/n) is strongly convex over P(a), we have that (Lem. 5, Appx. B) R2 fact differentiable with gradient given by
(R,(1) = arg max {q T 1 1D-DD(gg)) 1m)))) E R
(5)
gEP(o)
This means the full-batch gradient w + (& E Rd can be computed by solving the inner maximization retrieve 1 + VR, (I) E RD, calling the oracles to retrieve + V((w) E Raxd
3

---

Published as a conference paper at ICLR 2024
Biased Update Direction
Stochastic Grad. Estimate
Control Variate Correction
Asymptotic
Unbiasedness
Expected Trajectory
Trajectory w/ Var. Reduction Trajectory w/o Reduction
Gradient Descent
Figure 2: Prediction Error Reduction. Optimization trajectories on the DR objective. Darker shades indicate lower objective value. Left (Bias): Average trajectory of Prospect over 20 seeds compared to full-batch gradient descent. Right (Variance): Single trajectory of Prospect with/without adding control variate.
and multiplying them by the chain rule. To solve for the maximizer, we prove by standard convex du- ality arguments (Prop. 3, Appx. B) that when = Dj an f-divergence, the maximum over q can be expressed minimization problem that reduces to isotonic regression problem involving f", the convex conjugate of f. Isotonic regression can solved exactly by the Pool Adjacent Violators algorithm (Best al., 2000), which runs in Onn) time when the losses are sorted; see Appx. C.
Bias Reduction Loss Estimation. With a formula for the gradient in hand, we proceed to es- timation. Denote by g' == VR. (1) from (5), and observe that by the chain rule, VR, (f(w)) In words, we compute the "most adversarial" reweighting o ((w)) E (()) for
Em= V((w). v(())) given set of losses f(w), and then take convex combination of the gradients Vei Vln(w)
weighted by the probability mass function gf(w), While the gradient is computable, however, cessing ((w) and V((w) requires n calls to the function/ function/gradient oracles {,, V&, ) which can be prohibitive. While using a plug-in estimate with mini- batch of size < is natural choice in ERM (making the first term in (4) zero), this will be biased for our objective due to the maximization over q. For example, for m which are unbiased estimates of the ERM objective and gradient, respectively (not the SRM objec- tive). However, note that if the optimal weights g((w) were known, then for sampled uniformly (f(w)) While computing g ((w) again requires computing f(wr), the key ingredient of bias reduction in Prospect is maintaining a table E IR" of losses such that zs {(w) for the current iterate w E Rd and using g' as running estimate of (ww). This is justified when q + (gg 1,, is strongly convex, we have by Lem. 6 the map 1 + q is Lipschitz continuous in with respect to 11-112. Thus,
we have that Ry (&ww)) ((w)) and R. (()(w) w (w),
on [n]. that ng,
((w)) ( (w) is an unbiased estimate for Σ 9:
((") V(ww)
=s (ww) = g¹ 2 ((w)
Eg-Unit(a]]
[gq] V4, (we) R VR=
We prove in Sec. 3 that ((w) + 0 as the iterate counter goes to infinity for our particular choice of L, yielding asymptotically unbiased gradient estimate as illustrated in Fig. 2 (left).
Variance Reduction via Control Variates. The final ingredient our stochastic gradient estimate variance reduction scheme. Given any estimator à of Rd a control variate is another estimator b over the same probability space with known expectation EE( b Rdd such that E[(à a) (6- b)] > 0. We can exploit this positive correlation to construct an estimator with strictly smaller variance than à. Indeed, for 0, we have that
E à (6 Ella
[ā
b]] + o(7) < 1 for small 7,
2
(6)
all
E.
demonstrating the improvement of à 76 b) over à. Note that à need not be unbiased. In
our case, we have = ng. Vl((w, where 1 is the table of losses approximating {(w). We
also keep approximations g E Roxd of V((w) and p of g(ww), and define b np:gs and b Unif [npaas] Σー=1 Pz9j. In the unrealistic case in which i à, the optimal multi-
4

---

Published as a conference paper at ICLR 2024
Algorithm Prospect
Initial wo, spectrum number of iterations T, regularization >> 0, shift cost v > 0. Hyperparameter:: Stepsize 0. 1: Initialize / +- ((wo) and di - (w,(w)) + wwo for = 1, .....
Inputs:
2: Set f- arg (()))
I (4 |n/n) and p k q.
3: Set < Pig: Pig: E IR
4: Set w f- wo.
5: for T iterations do
6: Sample i. Unif[n] independently. x (q.(VL(() + pw) np.9i, + g. w + w -qu.
7:
8:
Iterate Update
9:
Ij Ez(w).
10:
q arg max [ 2 -D
Bias Reducing Update
11:
g - Pag + qi (V,( pmc)
12:
91 + ((((w) + www
P 6 q.
Variance Reducing Update
Output: Final point w.
plier is 7 1, trivially achieving zero variance. Similar we prove in Sec. that g VVV((w))))) and p f(w) + 0, so we have in the notation (6) that b- Thus, by using 7 1, our final stochastic gradient estimate is
76 b) =gg np.gt + Σ}= Pıgs.
with asymptotic variance reduction factor (with respect to i [n]]:
Var ()) - np:9, t (Beeouuainn|((og/VV44 - V((w) 9) (mpin: - - g]
Var ng (w))
Ep || ngq: [(( Vi(w) al2
This results asymptotically vanishing variance without decreasing the learning rate, as illustrated in Fig. 2 (right). This scheme generalizes (and is inspired by) the employed in the SAGA optimizer (Defazio et al., 2014) for ERM, in which = 1,/n. Finally, while ignored in this section for ease of presentation, each will actually store the gradients of the regularized losses & + 1-112
3 THE PROSPECT ALGORITHM
By combining the bias reduction and variance reduction schemes from the previous section, we build an algorithm that achieves overall prediction error reduction. Thus, we now present the Prediction Optimizer for Spectral Risk Measures (Prospect) algorithm to solve
Fa (w) == RR (f(w)) /1 [|]
min WERd
where μ > 0 is regularization constant. The full algorithm is given in Algorithm 1.
Instantiating Bias and Variance Reduction. Consider a current iterate w E Rd. As mentioned in Sec. 2, bias and variance reduction relies on the three approximations: the losses [ for E(w) R". each gradient gi for ((,(w) + paw E Rd, and the weights p for gf(w) P. Given initial point We E Rdd we initialize 1 ((wo), Vl(we) + w., and p = ((())) (including := p).
At each iterate, we sample indices j ~ Unif[n] independently. The index i is used to compute the stochastic gradient estimate (7), yielding the update direction in line 7 at the cost of call to a (4,, Vl,) oracle. Then, is updated by replacing I, with (w)) costing another call to (Lz, Vly), and we reset q (the variable that stores q¹). Next, we use i again to make the replacements of 9: with (e,(w) + pa and p, with 9: d. In summary, each approximation is updated every iteration by changing one component based on the current iterate w. The indices i,j are "decoupled" for theoretical convenience, but in practice using only works similarly, which in Sec. 4.
5

---

Published as a conference paper at ICLR 2024
Computational Aspects. The weight update in Line 10 is solved exactly by (i) sorting the vector of losses in On logn). n). (ii) plugging the sorted loss table into the Pool Adjacent Violators (PAV) algorithm rumning in Onn) time, mentioned in Sec. 2. Because only one element of changes every iterate, we may simply bubble sort / starting from the index that was changed. While in the worst case, this cost is (n)) it is exactly ((s) where is the number of swaps needed to resort I. We find in experiments that the sorted order of stabilizes quickly. The storage of the gradient table g requires ((nd) space in general, but it can be reduced to ( for generalized linear models and nonlinear additive models. For losses of the form (,()) = h(xz, w, Ui), for a differentiable loss h and scalar output yi. we have Vl((w)= = x; h h1(x,] w. ys). We only need to store the scalar h (xx] w, yi), so Prospect requires On + d) memory. In terms of computational complexity, Lines 8 and 13 require ((d) operations and Line 10 requires at most ((n) operations, so that in total the iteration complexity is Onn + d). In comparison, a full batch gradient descent requires (nd) operations so Prospect decouples efficiently the cost of computing the losses, gradients, and weights.
Convergence Analysis. We assume throughout that each & convex, G-Lipschitz, and L-smooth. We also assume that the D = D is an f-divergence with the generator f being a, a,-stnnggly convex on the interval [0, r2] (e.g. Ct2 = 2n for the x² -ivveregeece and a, 1 for the KL-divergence).
The convergence guarantees depend on the condition numbers k = 1 + L/μ of the individual regularized losses, as well as a measure Ka = no, of the skewness of the spectrum. Note that both and Ko are necessarily larger than or equal to one. Define w == arg min.. Fo(w), which exists and is unique due to the strong convexity of F.. The proof is given in Appx. D.6.
Theorem Prospect with a small enough step size is guaranteed to converge linearly for all > O. If. in addition, the shift cost is v 2 22(2/uaa then the sequence of iterates )(2) generated by Prospect and learning rate r = (12p(1 + K))R)) converges linearly at a rate
2 max (n. 2 + 1)}, i.e.
(l) w (1+,, + t
w²11.
The number of iterations t required by Prospect to achieve E| [)( w0 S E (provided that is large enough) is f = O(n + kKa) In(1/s)). This exactly matches the rate of the LSVRG (Mehta et al., 2023), the only primal stochastic optimizer that converges linearly for spectral risk measures. However, unlike LSVRG, Prospect is guaranteed to converge linearly for any shift cost and has single hyperparameter, the stepsize 77. Similarly, compared to primal-dual stochastic saddle-point methods, our algorithm requires only one learning rate, streamlining its implementation.
Prospect Variants for Non-Smooth Objectives. We may wonder about the convergence behavior 0, or the underlying losses l, non-smooth. While
of Prospect when either the shift cost the smoothness of the objective is then lost, Prospect can still converge to the minimizer w6 as we prove below. The first setting relevant as historically, SRMs such as the superquantile have been employed as coherent risk measures for loss distributions (Acerbi & Tasche, 2002) in the form of an L-estimator Σi= 034(s) (as seen in Sec. 2). If these losses are separated at the optimum, however, we may achieve linear convergence with Prospect even with v/ = O, due to "hidden smoothness" of the objective, i.e. differentiability at points for which ((w) has distinct components. Assume that each E; is convex and that μ O.
Proposition 2. Let wt be the unique minimizer of (8) with shift cost 2 O. Assume that the values (ww)),.... en (wā) are all distinct. Then, there exists consiant vo > 0 such that w w exactly for all v < vo- Thus, running Prospect with v E (0, vo] converges to the minimizer wy-
In particular, up is chosen so that 2nw (a++1 -o;) < ((41)) (wg)) ((y for each i, or as the multiplicative factor that relates gaps in the spectrum to gaps the optimal losses (see Appx. B). When & itself may be non-smooth, we generalize Prospect by applying it to the Moreau envelope of each loss &, and its gradient (Bauschke & Combettes, 201 I; Rockafellar, 1976), allowing for losses such as those containing an f penalty. Specifically, we consider oracles returning V (E))(&)(w) where env (f))(w) := infeeRd Ea(v) + || - vIß3: the update steps can be expressed in terms of the proximal operators of the losses (Bauschke & Combettes, 2011). These oracles can easily be accessed either in closed form or by efficient subroutines in common machine learning settings (De- fazio, 2016; Frerix et al., 2018; Roulet& & Harchaoui. 2022). The resulting algorithm enjoys linear convergence guarantee similar to Thm. 1 with a more liberal condition on the shift cost v while providing competitive performance in practice (see Appx. E).

---

Published as a conference paper at ICLR 2024
yacht
energy
concrete
kin8nm
power
10"
10"1
10
2072
CVaR
1073
10-3
10"4
10-1
10
10-1
10'4
20-7
20-7
10"
10"1
10
2013
Extremile
1033
10-2
10-4
10-1
10m4
20
10-7
20-7
10
10
10
201
ESRM
10-4
1074
10
10
10-7
10-7
10-7
30-7
SGD
SRDA
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 3: Regression benchmarks. The y-axis measures suboptimality as given by (9), while the a-axis measures the number of calls to the function value/gradient oracle divided by (i.e. passes through the training set). Rows indicate different SRM objectives while columns indicate datasets.
EXPERIMENTS
We compare Prospect against baselines variety learning tasks. While we focus attention on its performance optimizer of its training objective, also highlight metrics of interest the test set in fairness and distribution shift benchmarks. The algorithm implementation and data preparation code is made publicly available online: https//ggthuh..om/ooaknnm
Setting, Baselines, Evaluation. We consider supervised learning tasks with input-label example (I;, W1). Losses are of the form f((w) :- (yv, w c(x;)), with fixed feature embedding ø, and measuring prediction loss. Uncertainty sets considered are the CVaR, extremile, and ESRM. We compare against four baselines: minibatch stochastic gradient descent (SGD), stochastic regularized dual averaging (SRDA) (Xiao, 2009), Saddle-SAGE SAGA (Palaniappan & Bach, 2016), and LSVRG (Mehta et al., 2023). For SGD and SRDA, we use batch size of 64, and for L.VRRG we use an epoch length of n. For Saddle-SAGA, we find that allowing different learning rates for the primal and dual variables improves thcoretically and experimentally (Appx. F) and compare against an improved heuristic (setting the dual stepsize as 10n times smaller than the primal stepsize). We plot
Suboptimality(u)) (P& (w)
(w")) (F& (w( (o))))
where w* is approximated by running LBFGS (Nocedal & Wright, 1999) the objective until convergence. The r-axis displays the number calls to any first first-order oracle w + (e,(w), (w)) divided by n, i.e. the number of passes through the training set. We fix the shift cost 1 and regularization parameter μ = 1/n. Further details of the setup such as hyperparameter tuning, and additional results are given in Appxs and respectively.
4.1 TABULAR LEAST-SQUARES REGRESSION
We consider five tabular regression benchmarks under square loss. The datasets used are yacht (n 244) (Tsanas & Xifara, 2012), energy (n 614) (Baressi Segota et al., 2020), concrete (n = 824) (Yeh, 2006), kin8nm nm (n = 6553) (Akujuobi & Zhang, 2017), and power (n = 7654) (Tüfekei, 2014). The training curves are shown in Fig. 3.
Results. Across datasets and objectives, we find that Prospect exhibits linear convergence at a rate worse than SaddleSAGA and LSVRR but that often much better. For example, Prospect

---

Published as a conference paper at ICLR 2024
diabeted CVoR
diabetes FSRM
acsincoea
RCCAACCCO FSRM
CVaR
110
107
13
10-1
Suboptimality
Suboptimality
10m3
10-1
10-1
111
10
Unfair
Unfoir
0.06
Statistical Parity
(Classiicattoo) fication
Statisical Parity
0.78
[Regresion)
0.76
0.04
G.05
0.70
0.02
ww
wW
Fair
0.74
Fair
LSVRG
SaddleSAGA
Prospect (Ours)
Training curves for optimizers the CVaR and extrem- ile for diabetes (left) and CVaR and extremile for acsincome (right). Bottom: Statistical parity scores for the classification objectives diabetes (left) and regression objectives acsincome. Smaller values indicate better performance for all metrics.
Figure 4: Fairness benchmarks. Top:
converges to precision 10-8 for the CVaRR concrete and the extremile on power within half the number of passes that LSVRG takes for the same suboptimality. Similarly, for the ESRM on yacht, SaddleSAGA requires 64 epochs to reach the same precision Prospect at 40 epochs. The direct stochastic methods, SGD and SRDA, are biased and fail to converge for any learning rate.
4.2 FAIR CLASSIFICATION AND REGRESSION
Inspired by Williamson & Menon (2019), explore the relationship between distributional robust- ness and group fairness on 2 common tabular benchmarks. Diabetes 130-Hospitals (di adiabetes es) classification task of predicting readmission for diabetes patients based clinical data from US hospitals (Rizvi et al., 2014). Adult Census (acs income) is regression task of predicting income of US adults from data compiled by the American Community Survey (Ding al., 2021).
Evaluation. We evaluate fairness with the statistical parity score, which compares predictive dis- tributions of model given different values of particular protected attribute Agarwal et al. (2018; 2019). Letting = (X,Y, Y, A) denote a random (input, label, metadata attribute) triplet, a model said satisfy statistical parity (SP) if the conditional distribution of gXX) over predictions given A = a is equal for any value Intuitively, SP scores measure the maximum deviation between these distributions for any over a, so values close zero indicate SP-fairness. In diabetes. we use gen- der as the protected attribute A, whereas in acs inmo we use race the protected attribute. Note that the protected attributes are not supplied to the models. Results are given in Fig. 4.
Results. Firstly, we note that Prospect converges rapidly on both datasets while LS fails to con- verge on di abetes and SaddleSAGA fails to converge on acs incoe.. Secondly. LSVRG does not stabilize with respect to classification SP, showing mean/std SP score of 1.38 + 0.25% within the final ten passes on the abetes CVaR, whereas Prospect gives 0.82 0.00%. i.e., a 40% rel- ative improvement with greater stability. While SaddleSAGA does stabilize in SP on diabetes, fails to qualitatively decrease at all on the acs icsincome. Interestingly, while suboptimality and fa-rainesss are correlated for Prospect, SGD (reaching only 10-1 suboptimality with respect the CVaR objectives on acs i ncome) achieves a lower fairness score. Again, across both suboptimality and fairnsss, Prospect either the best close the best.
4.3 IMAGE AND TEXT CLASSIHCATION UNDER DISTRIBUTION SHIFT
We consider two tasks from the WILDS distribution shift benchmark (Koh et al., 2021). The Ama- Reviews (amazon) task (Ni et al., 2019) consists of classifying text reviews of products
8

---

Published as a conference paper at ICLR 2024
Superquantile
Superquantile
Extremile
Extremile
Worst Group Error
1.0
10a
Worst Group Error
1.0
10-1
Suboptimaling lity
Suboptimalin
0.9
10
0.9
10-3
0.8
0.8
yaala MA
10-1
X
20
M
0
20
Passes CVaR
Passes CVaR
Posses Extremile
Passes Extremile
Median Group Error
Median Group Error
Suboptimality lity
0.77
Suboptimaligg
6:70
0:70
6.00
8.00
0.90
10
(0
50
60
Passes
Passes
SGD
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 5: Distribution shift benchmarks. Top row: Training curves and worst group misclassi-
fication еггог on amazon test. Bottom row: Training curves and median group misclassification emror on the iwi 1deam test set. Smaller values indicate better performance for all metrics.
rating of 1-5, with disjoint train and test reviewers. The iWildCam (iwi 1dcam) image classifica- tion challenge (Beery et al., 2020) contains labeled images of animals, flora, and backgrounds from cameras placed in wilderness sites. Shifts are due to changes in camera angles, locations, lighting... We use 10000 and n = 20000 examples respectively. For both datasets, we train a linear probe classifier, i.c., a linear model over a frozen deep representation. For amazon, we use a pretrained BERT model (Devlin et al., 2019) fine-tuned on held-out subset of the Amazon Reviews training set for 2 epochs. For iwi dcam, we use a ResNet50 pretrained on ImageNet (without fine-tuning).
Evaluation.
Apart from the training suboptimality, evaluate the spectral risk objectives on their robustness to subpopulation shifts. We define each subpopulation group based on the true label. For the worst group misclasiification error the iwi wwilddmm, we use the median group error owing to its larger number of classes.
ama zon,
(Sagawa al., 2020). For
Results. For both amazon and twildcam, Prospect and SaddleSAGA (with our heuristic) outper- form LSVRG in training suboptimality. We hypothesize that this phenomenon due to checkpoints LSVRG getting stale over the n-length epochs for these datasets with large n. (leading slow duction of bias). In contrast, Prospect and SaddleS avoid this issue by dynamically updating the running estimates of the importance weights. For the worst group error for amazo zon.on. Prospect and SaddIeSAGA outperform LSVRG. Prospect has mean/std worst group eгтоr of 77.38:00.00% over the last ten passes on the extremile, whereas SaddleSAGA has slightly worse 77.53+1.57% In- terestingly, on 1wiam LSVRG and Prospect give stronger generalization performance, nearly 1pp better, than SaddleSAGA in terms of median group misclassification rate. summary, across tasks and objectives, Prospect demonstrates best or close best performance.
5 DISCUSSION
We introduced Prospect, a distributionally robust optimization algorithm for minimizing spectral risk measures that has a linear convergence guarantee. Prospect demonstrates rapid lincar convergence on benchmark examples and has the practical benefits of converging for any shift cost while only having a single hyperparameter. Promising avenues for future work include extensions to the non- convex setting by considering the regular subdifferential, variations using other uncertainty sets, and further exploring connections to algorithmic fairness.

---

Published as a conference paper at ICLR 2024
Acknowledgements. This work was supported by NSF DMS- 2023166, CCF-2019844, DMS- 2134012, NIH, and ODNI's IARPA program via 2022-22072200003, Part this work was done while Z. Harchaoui visiting the Simons Institute for the Theory of Computing.
REFERENCES
Carlo Acerbi and Dirk Tasche. On the Coherence of Expected Shortfall. Journal of Banking & Finance, 26, 2002.
Alekh Agarwal, Alina Beygelzimer, Miroslav Dudik, John Langford, and Hanma Wallach. Re- ductions Approach to Fair Classification. In ICML, 2018.
Alekh Agarwal, Miroslav Dudik, and Zhiwei Steven Wu. Fair Regression: Quantitative Definitions and Reduction- Algorithms. In ICML, 2019.
Uchenna Akujuobi and Xiangliang Zhang. Delve: A Sataset-Driven Scholarly Search and Analysis System. SIGKDD Explor. Newsl., 19, 2017.
Sandi Baressi Segota, Nikola Andelic, Jan Kudlacek, and Robert Cep. Artificial Neural Network for Predicting Values of Residuary Resistance per Unit Weight of Displacement. Journal of Maritime & Transportation Science, 57, 2020.
Heinz H Bauschke and Patrick L Combettes. Convex Analysis and Monotone Operator Theory in Hilbert Spaces. Springer, 2011.
Sara Beery, Elijah Cole, and Arvi Gjoka. The iWildCam 2020 Competition Dataset. arXiv preprint arXiv:2004. /0340, 2020.
Aharon Ben-Tal. Tal, Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust Solutions of Optimization Problems Affected by Uncertain Probabilities. Management Science, 59, 2013.
Dimitri P Bertsekas. Nonlinear Programming. Journal of the Operational Research Society, 1997.
Michael Best, Nilotpal Chakravarti, and Vasant A Ubhaya Minimizing Separable Convex Func- tions Subject to Simple Chain Constraints. SIAM Journal on Optimization,, 10, 2000.
Jose Blanchet, Peter W Glynn. Jun Yan, and Zhengqing Zhou. Multivariate Distributionally Robust Convex Regression under Absolute Error Loss. In NeurlPS, 2019a.
Jose Blanchet, Yang Kang, and Karthyek Murthy. Robust Wasserstein Profile Inference and Appli- cations to Machine Learning. Journal of Applied Probability, 56, 2019b.
Yair Carmon and Danielle Hausler. Distributionally Robust Optimization via Ball Oracle Accelera- tion. In NeurlPS, 2022.
Ruidi Chen and loannis Ch. Paschalidis. Distributionally robust learning. Foundations and Trends in Optimization, 2020.
Julien Grand Clement and Christian Kroer. First-Order Methods for Wasserstein Distributionally Robust MDP. In ICML, 2021.
John Cotter and Kevin Dowd. Extreme Spectral Risk Measures: Application Futures Clear- inghouse Margin Requirements. Journal of Banking & Finance, 30, 2006.
Zac Cranko, Zhan Shi, Xinhua Zhang, Richard Nock, and Simon Korbblith. Generalised Lipschitz Regularisation Equals Distributional Robustness. In ICMI., 2021.
Abdelaati Daouia, Irêne Gijbels, and Gilles Stupfler. Extremiles: A New Perspective Asymmetric Least Squares. Journal of the American Statistical Association, 1 11,, 2019.
Aaron Defazio. Simple Practical Accelerated Method for Finite Sums. In NeurlPS, 2016.
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: Fast Incremental Gradient
Method With Support for Non-Strongly Convex Composite Objectives. NeurlPS, 27, 2014.
10

---

Published as a conference paper at ICLR 2024
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A Large-scale Hierarchical Image Database. In CVPR, 2009.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionallyy Robust Feder- ated Averaging. In NeurIPS, 2020.
Jacob Devlin, Ming-Was Mei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACI, 2019.
Frances Ding, Moritz Hardt, John Miller, and Ludwig Schmidt. Retiring Adult: New Datasets for Fair Machine Learning. NeurIPS, volume 34. Curran Associates, Inc., 2021.
Yanbo Fan, Siwei Lyu, Yiming Ying, and Baogang Hu. Learning with Average Top-k Loss, In NeurlPS, volume 30, 2017.
Rizal Fathony, Ashkan Rezaei, Mohammad Ali Bashiri, Xinhua Zhang, and Brian Ziebart. Distri- butionally Robust Graphical Models. In NeurlPS, 2018.
Thomas Frerix, Thomas Möllenhoff, Michael Möller, and Daniel Cremers. Proximal Backpropaga- tion. In ICLR, 2018.
Soumyadip Ghosh, Mark Squillante, and Ebisa Wollega. Efficient Generalization with Distribution- ally Robust Learning. In NewrlPS, 2021.
Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without Demographies in Repeated Loss Minimization. In ICML, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In CVPR, 2016.
Jean-Baptiste Hiriart-Urruty and Claude Lemaréchal. Fundamentals of Convex Analysis. Springer Science & Business Media, 2004.
Hisham Husain. Distributional Robustness with IPMs and links to Regularization and GANs. In NeurlPS, 2020.
Yu Inatsu, Shogo Iwazaki, and Ichiro Takeuchi. Active Learning for Distributionally Robust Level- Set Estimation. In ICML, 2021.
Yu Inatsu, Shion Takeno, Masayuki Karasuyama, and Ichiro Takeuchi. Bayesian Optimization for Distributionally Robust Chante-constrained Problem. In JCML, 2022.
Yang Jiao, Kai Yang, and Dongjin Song. Distributed Distributionally Robust Optimization with Non-Convex Objectives. In NeurlPS, 2022.
Jikai Jin, Bohang Zhang, Haiyang Wang, and Liwei Wang. Non-convex Distributionally Robust Optimization: Non-asymptotic Analysis. In NeurlPS, 2021.
Nathan Kallus, Xiaojie Mao, Kaiwen Wang, and Zhengyuan Zhou. Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning. In ICML, 2022.
Kenji Kawaguchi and Haihao Lu. Ordered SGD: New Stochastic Optimization Framework for Empirical Risk Minimization. In AISTATS, volume 108, 2020.
Justin Khim, Liu Leqi, Adarsh Prasad, and Pradeep Ravikumar. Uniform Convergence of Rank- weighted Learning. In ICML, 2020.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In ICLR, 2015.
Johannes Kirschner. Ilija Bogunovic, Stefanie Jegelka, and Andreas Krause. Distributionally Robust Bayesian Optimization. In AISTATS, 2020.

---

Published as a conference paper at ICLR 2024
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsub- ramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A Benchmark of 1i-the-Wild Distribution Shifts, 2021.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadch-Abadeh Wasserstein Distributionally Robust optimization: Theory and Applications in Machine Learning. Operations research & management science in the age of amalytics, 2019.
Yassine Laguel, Krishna Pillutla, Jérôme Malick, and Zaid Harchaoui. Superquantiles Work: Ma- chine Learning Applications and Efficient Subgradient Computation. Set- Valued and Variational Analysis, 2021.
Jacho Lee, Sejun Park, and Jinwoo Shin. Learning Bounds for Risk-sensitive Learning. In NeurIPS, volume 33, 2020.
Daniel Levy, Yair Carmon, John Duchi, and Aaron Sidford. Large-Scale Methods for Distribution- ally Robust Optimization. In NeurlPS, 2020.
Jiajin Li, Sen Huang, and Anthony Man-Cho So. A First-Order Algorithmic Framework for Distri- butionally Robust Logistic Regression. In NeutIPS, 2019.
Yeshu Li, Danyal Saeed, Xinhua Zhang, Brian Ziebart, and Kevin Gimpel. Moment Distributionally
Robust Tree Structured Prediction. In NeurIPS, volume 35. Curran Associates, Inc., 2022.
Zachary C. Lipton, Yu-Xiang Wang, and Alexander J. Smola. Detecting and Correcting for Label Shift with Black Box Predictors. In ICML, 2018.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just Train Twice: Improving Group Robustness without Training Group Information. ICML, 2021.
Jiashuo Liu, Zheyan Shen, Peng Cui, Linjun Zhou, Kun Kuang, and Bo Li. Distributionally Robust Learning With Stable Adversarial Training. IEEE TKDE, 2022a.
Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally Robust Optimization with Data Geometry. In NeurIPS, volume 35. Curran Associates, Inc., 2022b.
Zijian Liu, Qinxun Bai, Jose Blanchet, Perry Dong, Wei Xu, Zhengqing Zhou, and Zhengyuan Zhou. Distributionally Robust Q-Learning. In ICML, 2022c
Kyriakos Lotidis, Nicholas Bambos, Jose Blanchet, and Jiajin Li. Wasserstein Distributionally Ro- bust Linear- Quadra-Q Estimation under Martingale Constraints. In AISTATS, 2023.
Luo Luo, Haishan Ye, Zhichao Huang, and Tong Zhang. Stochastic recursive gradient descent ascent minimax problems. In NewrlPS, 2020.
for stochastic maaconvex-sttronlly-om
Andreas Maurer, Daniela Angela Parletta, Andrea Paudice, and Massimiliano Pontil. Robust Unsu- pervised Learning via L-statistic Minimization. In /CML, 2021.
Ronak Mehta, Vincent Roulet, Krishna Pillutla, Lang Liu, and Zaid Harchaoui. Stochastic Opti- mization for Spectral Risk Measures. In AISTATS, 2023.
Tong Mu, Yash Chandak, Tatsunori B Hashimoto, and Emma Brunskill. Factored DRO: Factored Distributionally Robust Policies for Contextual Bandits. In NeurIPS, 2022.
Debarghya Mukherjee, Felix Petersen, Mikhail Yurochkin, and Yuckai Sun. Domain Adaptation meets Individual Fairness. And they get along. In NeurIPS, 2022.
Hongseok Namkoong and John C Duchi. Stochastic Gradient Methods for Distributionally Robust Optimization with f-divergences. NeurlPS, 2016.
12

---

Published as a conference paper at ICLR 2024
Yassine Nemmour, Bernhard Schölkopf, and Jia-Jie Zhu. Approximate Distributionally Robust Nonlinear Optimization with Application to Model Predictive Control: A Functional Approach. In L4DC, 2021.
Yurii Nesterow. Smooth Minimization of Non-Smooth Functions. Mathematical Programming, 2005.
Yurii Nesterov. Lectures on Convex Optimization. Springer Publishing Company, Incorporated, 2nd edition, 2018.
Jianmo Ni, Jiacheng Li, and Julian McAuley. Justifying Recommendations using Distantly-Labeled Reviews and Fine-grained Aspects. EMNLP, 2019.
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
Balamurugan Palaniappan and Francis Bach. Stochastic Variance Reduction Methods for Saddle- Point Problems. NeurlPS, 29,201, 2016.
Krishna Pillutla, Yassine Laguel, Jérôme Malick, and Zaid Harchaoui, Federated Learning with Superquantile Aggregation for Heterogeneous Data. Mach. Learn., 2023.
B.T. Polyak. Some methods of speeding up the convergence iteration methods. USSR Computa- lional Mathematics and Mathematical Physics, 1964.
Hamed Rahimian and Sanjay Mehrotra. Frameworks and Results in Distributionally Robust Opti- mization. Open Journal of Mathematical Optimization, 3, 2022.
Allen Z. Ren and Anirudha Majumdar. Distributionally Robust Policy Learning via Adversarial Environment Generation. JEEE RAL, 2022.
Ali Rizvi, Beata Strack, Jonathan P. DeShazo, Chris Gennings, Juan L. Olmo, Sebastian Ventura, Krzysztof J. Cios, and John N. Clore. Impact of HbAle Measurement on Hospital Readmission Rates: Analysis of 70,000 Clinical Database Patient Records. BioMed Research International, 2014, 2014.
Tyrrell Rockafellar. Monotone Operators and the Proximal Point Algorithm. SIAM Journal on Control and Optimization, 14, 1976.
Tyrrell Rockafellar and Johannes O Royset. Superquantiles and Their Applications Risk, Ran- dom Variables, and Regression. In Theory Driven by Influential Applications. Informs, 2013.
Vincent Roulet and Zaid Harchaoui. Differentiable Programming la Moreau. In ICASSP. IEEE, 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally Robust Neural Networks. In ICLR, 2020.
Dvir Samuel and Gal Chechik. Distributional Robustness Loss for Long-Tail Learning. In ICCV, 2021.
Hitesh Sapkota, Yiming Ying, Feng Chen, and Qi Yu. Distributionally Robust Optimization for Deep Kernel Multiple Instance Learning. In AISTATS, 2021.
Alexander Shapiro. Distributionally Robust Stochastic Programming. SIAM Journal on Optimiza- tion, 2017.
Vishnu D. Sharma, Maymoonah Toubeh, Lifeng Zhou, and Pratap Tokekar. Risk-Aware Planning and Assignment for Ground Vehicles using Uncertain Perception from Aerial Vehicles. In 2020 JEEE/RS/ IROS, 2020.
Xuli Shen, Xiaomei Wang, Qing Xu, Weifeng Ge, and Xiangyang Xue. Towards Scalable and Fast Distributionally Robust Optimization for Data-Driven Deep Learning. In IEEE ICDM, 2022.
Galen Shorack. Probability for Statisticians. Springer Texts in Statisties, 2017.
13

---

Published as a conference paper at ICLR 2024
Rahul Singh, Qinsheng Zhang, and Yongxin Chen. Improving Robustness via Risk Averse Distri- butional Reinforcement Learning. In L4DC, 2020.
Aman Sinha, Matthew O' O'Kelly, Hongrui Zheng, Rahul Mangharam, John Duchi, and Russ Tedrake. FormulaZero: Distributionally Robust Online Adaptation via Offline Population Synthesis. In ICML, 2020.
Agnieszka Stowik and Léon Bottou. On Distributionally Robust Optimization and Data Rebalancing In AISTATS, 2022.
Matthew Staib and Stefanie Jegelka. Distributionally Robust Optimization and Generalization in Kerel Methods. In NeurIPS, 2019.
Masashi Sugiyama, Shinichi Nakajima, Hisashi Kashima, Paul Buenau, and Motoaki Kawanabe. Direct Importance Estimation with Model Selection and Its Application to Covariate Shift Adap- tation. In NeurIPS, 2007.
Sebastian Shenghong Tay, Chuan Sheng Foo, Urano Daisuke, Richalynn Leong, and Bryan Kian Hsiang Low. Efficient Distributionally Robust Bayesian Optimization with Worst-case Sen- sitivity. In ICML, 2022.
Athanasios Tsanas and Angeliki Xifara. Accurate Quantitative Estimation of Energy Performance of Residential Buildings Using Statistical Machine Learning Tools. Energy and Buildings, 49, 2012.
Pinar Tüfekei. Prediction of Full Load Electrical Power Output of Base Load Operated Combined Cycle Power Plant using Machine Learning Methods. International Journal of Electrical Power & Energy Systems, 60, 2014.
Hieu Vu, Toan Tran, Man-Chung Yue, and Viet Anh Nguyen. Distributionally Robust Fair Principal Components via Geodesic Descents. In ICLR, 2022.
Shengbo Wang, Nian Si, Jose Blanchet, and Zhengyuan Zhou. A Finite Sample Complexity Bound for Distributionally Robust Q-learning In AISTATS, 2023.
Zhenyi Wang, Li Shen, Le Fang, Qiuling Suo, Tiehang Duan, and Mingchen Gao. Improving Task- free Continual Learning by Distributionally Robust Memory Evolution. In ICML, 2022.
Robert Williamson and Aditya Menon. Faireess Risk Measures. In ICML, 2019.
Lin Xiao. Dual Averaging Method for Regularized Stochastic Learning and Online Optimization. In NeurlPS, volume 22, 2009.
Mengdi Xu, Peide Huang, Yaru NIu, Visak Kumar, Jielin Qiu, Chao Fang, Kuan-Hui Lee, Xuewei Qi, Henry Lam, Bo Li, and Ding Zhao. Group Distributionally Robust Reinforcement Learning with Hierarchical Latent Variables. In AISTATS, 2023.
Zhouhao Yang, Yihong Guo, Pan Xu, Anqi Liu, and Animashree Anandkumar. Distributionally Robust Policy Gradient for Offline Contextual Bandits. In AISTATS, 2023.
Erfan Yazdandoost Hamedani and Afrooz Jalilzadeh. A Stochastic Variance-Reduce Accelerated Primal-Dual Method for Finite-Sum Saddle-Point Problems. Comput. Optim. Appl., 2023.
Ivy Yeh. Analysis of Strength of Concrete Using Design of Experiments and Neural Networks. Journal of Materials in Civil Engineering, 18, 2006.
Yaodong Yu, Tianyi Lin, Eric V Mazumdar, and Michael Jordan. Fast Distributionally Robust Learning with Variance-Reduced Min-Max Optimization. In AISTATS, 2022.
Yibo Zeng and Henry Lam. Generalization Bounds with Minimal Dependency on Hypothesis Class via Distributionally Robust Optimization. In NeurlPS, 2022.
Jingzhao Zhang, Aditya Menon, Andreas Veit, Srinadh Bhojanapalli, Sanjiv Kumar, and Suvrit Sra. Coping with Label Shift via Distributionally Robust Optimisation. In ICLR, 2021.
14

---

Published as a conference paper at ICLR 2024
Zhengqing Zhou, Zhengyuan Zhou, Qinxun Bai, Linhai Qiu, Jose Blanchet, and Peter Glynn. Finite- Sample Regret Bound for Distributionally Robust Offline Tabular Reinforcement Learning. In AISTATS, 2021.
Zhengyu Zhou and Weiwei Liu. Sample Complexity for Distributionally Robust Learning under chi-square divergence. JMLR, 2023.
15

---

Published as a conference paper at ICLR 2024
Appendix
The Appendix sections are organized as follows. We summarize notation Appx. and provide intuition and results regarding the primal/dual objective function in Appx. B. We describe in detail efficient implementations of the proposed algorithm in Appx. C. In Appx. D, we describe the convergence analyses of the main algorithm. In Appx. E and Appx. F, we describe an Moreau envelope-hasee based variant of our method and an improved version of an existing saddle-point method, respectively. Appx. contains technical results shared to multiple proofs. We then describe the experimental setup in detail in Appx. H and give additional results in Appx.
Table of Contents
A Summary of Notation
17
B Properties of the Primal and Dual Objectives
17
Efficient Implementation of Prospect
25
Convergence Analysis Prospect
30
D.I Overview
30
D.2 Key Technical Lemma: Bias of the Prospect Gradient Estimate
31
D.3 Bounding the Distance-to-Optinne
33
D.4 Bounding the Evolution of the Lyapunov Function Terms
35
D.5 Tuning Constants and Final Rate
39
D.6 Proof Main Result
46
E Improving Prospect with Moreau Envelopes
46
E.1 Overview
46
E.2 Convergence Analysis
48
F Improving the Direct Saddle-Point Baseline
51
Fl Overview
51
Convergence Analysis
52
G Technical Results from Convex Analysis
57
H Experimental Details
59
H.] Tasks & Objectives
59
H.2 Datasets
60
H.3 Hyperparameter Selection
61
Compute Environment
61
I Additional Experiments
61
16

---

Published as a conference paper at ICLR 2024
A
SUMMARY OF NOTATION
We collect the notation used throughout the paper in Tab. 1.
Symbol
Description
120 0
Standard regularization constant.
Shift cost.
On
Strong convexity constant for any f generating an f-divergence.
Shorthand 7/ (used in the convergence proofs).
((w)..... ...., (a(w)
Loss functions & [R R.
(ww)
Vector losses f(w) = (fi(w)..... fn (wr) for w Rd.
ri(w)
Regularized loss r:(w) (ww) -
r(w)
Vector regularized losses r(w) (ri (w). Tu (w)).
V((w)
Jacobian matrix of f: IRd + R at w. (shape x d).
a
The vector o (d)..... .n) E [0, 1 where each 01 .... S on and they sum to The set (TTo IT [0, 10x ITI, 1,, IT 1, 1,,, known the permutahedron.
P(a)
f
Convex function f: [0, oc) + R U { generating an fddivergence.
f
Convex conjugate (v) == suPacR (ay fx))]-
Shift penalty function 22 P(a) + [0, zo).
Gy or S
We consider f-divergence penalties (g)(g) = (g
Fo
Main objective (w(w) maxgep(o) {2 [T(((w) DD 1. 1
Most unfavorable reweighting for a given vector of losses, i...
(1) of
( arg max pcP()) q T <g In//).
of used only in main text for readability.
w
Optimal weights arg min maX ( (DDgllll (4/2) l
Most unfavorable reweighting f(w"), ), i.c., 9" = (e(W"))
G
Lipschitz constant of each w.r.t. 11-112-
T.
Lipschitz constant of each VE, w.r.t 11:112- M L + μ. the Lipschitz constant of each Vr; w.r.t. 11:112-
M
E, [+]
Shorthand for [ii i.e., expectation conditioned w!(t)
Table 1: Notation used throughout the paper.
B PROPERTIES OF THE PRIMAL AND DUAL OBJECTIVES
Recall that we are interested in the optimization problem
min WERd
(ww) max 4" ((w) -Dy(alll./n) gEP(o)
+ 2 |
(10)
where (y(q|]1n./ denotes an f-divergence between the distribution associated to the reweighting q and the discrete uniform weights 1,/n = .L/n...,,1I) and P(o) is the spectral risk measure uncertainty set.
The first goal for this section will be to derive properties of the function F, (w), or the primal objective, as well as the inner maximization problem, which refer to as the dual objective. Both will be useful in motivating and analyzing Prospect (used for the primal minimization) and various subroutines used to compute the maximally unfavorable reweighting (i.e., the maximizer q in
17

---

Published as a conference paper at ICLR 2024
the inner maximization). These properties do not depend on the structure of the P(r) itself, only that it is a closed, convex set. The second goal the section is to provide additional background on the choice of P(a)) from statistical modeling perspective. The uncertainty set P(o) is then described further from a computational perspective in Appx. C.
Review of f-Divergences. Let q and p be any probability mass functions defined on atoms {1...., n.. Consider a convex function f: [0, o0) + U {+00} such that f(1) = 0, f(z) is finite for x > O, and limz-+ot f(x) 0. The f-divergence from to p generated by this function f is
Dy(allp) = /
Pj Pi:
where we define 0f (0/0) := 0 in the formula above. If there is an i such that p. 0 but qi 0, say (q(lpp) = oc. The x² -d-diverge is generated by fx²(x) ([(=)== x² and the KL divergence is generated by (x) = x In 4 (z) where 1+ denotes the convex indicator that is zero for and too otherwise, and we define x In x = O for all x < 0.
The Dual Problem. We describe the inner maximization first, that is
gEP(o) -- (11-vD((((11.)
max
(11)
Its properties will inform the algorithmic implementation for the minimization over w in (10), In the case of f-divergence between q and the uniform weights 1n/n, we have
Dy(al11./m) Σ f (nqi).
(12)
We now derive the dual problem to Equation (11). This will lead to an algorithm to solve the optimization problem efficiently. Throughout, we denote (") = suPrer (ry - [xx) as the convex conjugate of f-
We consider the following functions whose conjugates are strictly convex. Recall that if f is smooth, i.e., with Lispehtiz continuous gradients, then its conjugate is strongly convex, hence strictly convex. More generally f is strictly convex if f convex and essentially smooth, that is, with gradient norm tending to +oo at its boundaries, see e.g. (Rockafellar, 1976) for detailed presentation. For simple cases such as the x² or KL divergence presented, strict convexity of the convex conjugate is naturally satisfied:
fx((x)= =x= - and (V) -/2/4
(x2 (di-ivergeecee)
fk((z) Inx + + (x) and fir(2) exp (y 1).
(KL-divergence)
Proposition 3. Leil E TR" be vector and π be a permutation that sorts its entries in non-decreasing order, i.e. (x(1) .... I L(()). Consider a function f strictly convex with strictly convex conjugate defining a divergence Dy. Then, the maximization over the permutahedron subject to the shift penalty can be expressed as
max gEP(a)
(q-1- HD-DDD(gll11n/ = min CERn
gi(cst))
(13)
fe (166)-
where we define gi gi(e;; 020 n
The optima of both problems, denoted
(P)( arg min "E Sc....cn
gi(egl), Tpth Σ
-arg max - (eP(a)
are related as gmth VuDD (-111/11))) (n/ln))))((( -1- ((), that is.
1 | (2 - (1))) n
a. ("'
(14)
18

---

Published as a conference paper at ICLR 2024
Proof. Let LP(a) denote the indicator function of the permutahedron P(o),, which is 0 inside P(o) and t00 outside of P(o). Its convex conjugate is the support function of the permutahedron, i.e.,
(P))
max q 1. (EP(m)
For two closed convex functions hı and h2 that are bounded from below, the convex conjugate of their sum the infimal convolution of their conjugate (Hiriart-Uruuy & Lemaréchal, 2004, Propo- sition 6.3.1):
(h1 + (2) inf (1((2) h5(x (
vER4
Provided that fı + h2 is strictly convex, we have that the maximizer defining the conjugate is unique and equal to the gradient, that is,
Vhh + hx) = arg max {2T - (h1 + h2)(2)}
=ERd
If, in addition, 41 + h is strictly convex and h is differentiable, we have, by Danskin's theo- rem (Bertsekas, 1997).
Vhh h22) Vhū(a v (x:) for (x) arg min (h7(v) h5(1 y)) yER*
Consider then hi (4) = (g) and h2(2) S2f(g) == DDy 1n/n), Provided that f is strictly convex with f strictly convex, D, also strictly convex with Dj strictly convex since D just decomposes
of independent variables. We have then
sup {q 1 ()(1)) sup {q (LP(G)(0) +g((4)))
(EP(a)
vER"
(UP(0) + (1)
inf yER" { { ((p() (v) + 27 y)
inf yERe
qTy y 97 3) -3) (9E(())
max
OW + 97 v)
inf yE
(15)
where V(1) S ... S V(n) are the ordered values of y E RR. Moreover we have that
arg max (1-1- v1 (1) for y (1) arg min (EP(a)
Σ owvo + 51 v)
DER"
Since for any R", P2 is decomposable into sum of identical functions evaluated at the coor- dinates (#1...., p. that is, (z) Σn= wI;)), its convex conjugate is (17 = Σn w (v:). In our case, wxx)) = / (na;) from Equation (12), so sow" sw)) = (M/n)S"
Next, by convexity of w", have that if for scalars 14, Vi. yy such that 4 S Ij and yi 2 yjr then using Lem. 33, we have that
w' (4 - y:) + w" (1, - -13) 2 (4 - 9;) + (1, v:).
Hence for to minimize (1 - y) = Ei= w (1: - y1). the coordinates of y must be ordered as 1. That is. if π is an argsort for I, s.t. (+(1) S ... < [n(m)) then yr(1) Σ ... < Y=(n)- Since (v) = Ei (il(() does not depend on the ordering of y, the solution of (15) must also be
19

---

Published as a conference paper at ICLR 2024
ordered 1 such that the dual problem (15) can be written
inf yER Mx(1) SVr(n)
Σ O.Un(()) + v fr (Lx()
Ma(())
inf
Σ oic
+
f
(Lx(1)
-
-i
V
CE
1
uS..Sc 1m1
min
Σ 9:(c;:))
1=1
where we used a change of variables such that the solutions of the left and right hand sides are related y" (1) (1).
While strict convexity of the function / ensures naturally the existence the maximizer n defined above, this does not help quantify the continuity of the maximizer with respect the given vector of losses, For that, we need to consider strong convexity of the f-divergence on the maximization The following proposition simply links the strong convexity of f to the strong convexity of the associated weighted divergence in & norm.
Proposition 4. Assume that f : R is an --rSllgy convex n). Then, q + vDy(gll In/n) is -strongy convex with respect to |1-11g
Proof. Due the at Strong convexity of f, for any q,p [0, 1]" and any 0 (0, 1) and any e [n],
(ōnq. + (1 0)mp.) (f(q) + (1 (1=)))((mp))) f)p(p)))
On 0(1 () np.)
2
We average this inequality over i, yielding
I (n(0g: + - ()):)) = 0 n
(mqgs) + (1 0)) s(pp)) On 0(1 (0]1-0][laa 21001-0)l1ng -ppil??
1=1
Defining l2(()) = (y(g 1, nn), the statement above can be succinctly written as (2 + (1 (1) - (g)(() + - ())
0(1 0)((a - -pall².
2
is Therefore, 22 is nstonn-y convex with respect to 11-112 [0, 11" so q vD, (q|l11// (n) -stonn-strongly convex.
The Pool Adjacent Violators (PAV) algorithm is designed exactly for the minimization (13). The algorithm described for the divviergee and KL--iivrgence with implementation steps in Appx. C. Both the argsort and the inverse argort -1 mappings from [n] {l,..- .n) onto itself, but the interpretation
these indices are different for the input and output spaces [n]. The argsort can be thought of
index finder, in the sense that for a vector / R" because (a(1) S ... < (a(n). (i) can be interpreted as the index of an element of which achieves the rank i in the sorted vector. On the other hand, x -i) can be thought of as a rank finder, in that (i) rank(i) the position that 1 takes in the sorted form of To summarize:
i
[n]]
while π
[n]
[ manks of losses
ranks of losses
indices of training examples
indices of training examples
We may equivalently write (14) as
1 1 IP (204--00 opt (ank()) (1))
(1) (1)
(16)
n
Finally, as seen in Appx. C, will be helpful to compute पht in sorted order. Because the f- divergence is agnostic the ordering of the q vector (as is being compared to the uniform weights).
20

---

Published as a conference paper at ICLR 2024
P3>>333> 14
aRs
4/
n
n
Uncertainty Set
Figure 6: Visualization of dual solution for varying shift cost. The hexagon represents a three- dimensional uncertainty set (see Fig. 7) whereas the white dots signify solutions of (11) for fixed loss vector and varying shift cost v. As y increases, the solution tends closer to the uniform vector 1n/n, whereas decreases, the shifted distribution is allowed to drift closer the boundary of the uncertainty set.
q can also sorted by w. Thus, we may also write
(5 (1) 1 1B (c))
(17)
The effect of varying the shift cost the solution of (11) is visualized in Fig. 6.
The Primal Function. When the divergence generator
strongly convex and the loss function show next.
é: R2d + RR is convex and differentiable, we have that (10) is differentiable,
Lemma 5. When the map + vD(q)l,./n) is strongly corvex over P(a), then Ro is continuously differentiable with gradient given by
VRo = arg max {q _ - ( e R". gEP(O)
Proof. Because P(o) as defined in (20) is closed and convex, the strongly concave function g H> q -- 1,/n) has unique maximizer. Because P(o)) is closed subset of the compact set A is also compact. By Danskin's theorem (Bertsekas, 1997, Proposition B,25), we have that Fg is continuously differentiable with the given gradient formula.
Lemma Let Fdd Fa be differentiable with Jacobian V((w) Raxd Let each l; IRd + IR. be convex. Assume > O. Let be on--slrnglyy convex the interval [0, n). Then, the function Fa from Equation (10) is differentiable with gradient equal
FF (w) = (VE(w))T (e(W)(()) + pw.
Furthermore + (1) is -I-pschhi-ipppiice continuous w.E.t. ||-112-
Proof. First, by Prop. 4, we have that q (D 1,)) is (pan,n)- convex with respect to Il-ll2 on [0, 1]". Next, due to the convexity of each l, and the non-negativity of any E P(a), we have that
max {gT((w)) v(ly(1)) PEP(O)
21

---

Published as a conference paper at ICLR 2024
is convex, is its pointwise maximum (over q) of family of convex functions q ((w). We have by Lem. that Fg is continuously differentiable with
(Fw) = v((w) T ((r + pw.
Moreover, by Nesterov (2005, Theorem 1), we have that | + (1) is Lipschitz continuous with Lipschitz constant equal to the inverse of the strong convexity constant of vSly, which is va,n. M
Returinng to our canonical examples, we have that for the x², (x)) = x² 2-strongly convex on R and that (xr) = xr In æ is (1/n)-stroggl convex on [0, m). Thus, the function 1 + (T) will have Lipschitz constant 2nv and v, respectively.
Smoothness Properties. By applying Lem. 6 to Lipschitz continuous losses, may achieve the following guarantee regarding the changes in 90sl with respect to w.
Lemma 7. Let f be stronngyy convex on the interval [0, n]. For any w1,..., Wn. wi...., w', E Rdd construct E(w1,.... w.) = (fi(w;)) E IR", as well as E(wi...., w'n) where each l is G- Lipschitz w.r.t. I1-112- Then, we have
G²
((("(( (f(wi,..., wn)) (e('',, whii) n²a v2
Σ 1=1
Ilu; - w/l1
Proof. By the Lipschitz property of TM (Lem. 6), have,
([ ((ftw,...,.,,. wn)) wh))1
wr) - E(wo "or w/
n n²azv2
n
(ls(wr) L,(w())2
n 22
Imll -w(ll
n²a2
As special case of Lem. 7, we may consider wh Wh w Rdd
and w/ = =... w/y
w/ E Rdd in which case the result reads
G² 22
(((w)) - 9TTh (((w'))
| w 1
na
Properties under No Shift Penalty. Next, we use the smoothness properties above to prove Prop. by virtue of the following proposition, which states the equivalence of the minimizers of "no-cost" and "low-cost" objectives.
Proposition Let wt, the unique minimizer of (8) with shift cost v 2 O and -ivvvergeence penalty. Define (()(w5) <..., L(n) (w)) to be the order statistics of & (w Ln (w6), which are assumed to be distinct. Consider vo such that
1wg (0 - o.)) (wz) (w) i 1.....n.
(18)
We have that w6- wt, for all vo-
22

---

Published as a conference paper at ICLR 2024
Proof. For vector I E R" and 0, consider
(1)
max q vn 11 1n/n112 gEP(a)
q' (1 + 2v1nn) vn 11al12 - (v/n) 1|1n112
max gEP(a)
max q_6 vn ||q112 + + (EP(a)
(v) +V,
where we used that q 1 = 1 for all q E P(o)). For v > by Danskin's 's theorem (Bertsekas, 1997, Proposition B.25),
Vh (1) = Vg (1) arg maxq vn 14112 arg maxq (1/2mv) |191|2 (EP(O)
gEP(o)
Without loss of generality, assume that is sorted, so that by applying the duality given by Prop. and that (t) = /4/4+ 1, we have that
copl arg min 024; +
1
Ci
2
4n
c Sc...5n
9.
ml (1)
2nv (4 - C. (I))
By differentiating with respect c, have that if unconstrained, ()) - + 2nwa, the primal solution and the dual solution is given by 9,th o, which is equal the gradient of ( + go (I) when / has distinct elements. Thus, we derive a condition under which c: (1) = 4 + 2nva; is monotonically non-decreasing (i.e. has the same sorted order as !), which will be true if is small enough. Specifically, we have that
2nvo (a++1 - o.) < (161))(6)) E(s) (wx) for i I,...,n,
(19)
for some v/ > 0, then o (1) 1, 2nvo; is monotonically non-decreasing. Consequently, for any V S P/.
Vge(((w))) Vgo(L(w))).
Denote our objective
(w hr (E(w)) + [|
where we explicitly show the dependence on the shift cost 2 O. For O, since the losses are differentiable and Elwz)) is composed of distinct coordinates. La.0 is differentiable at wi with gradient V((w:)) Vho (f(w))) + +ww (Mehta al., 2023, Proposition 2), where V((wg) E Raax denotes the Jacobian of l at wy- Using the chain rule, we successively deduce
(w)
v((wg) Vha(f(w()) + pmg 0
V((wō) Vgo(((w:))) + ymm
V((wE) (g(()()))) + mmm = O
v((wE) (h, + pwo = 0
(.r = 0.
Applying the first-order optimality conditions of Co,0 and Co,u, well the uniqueness of w completes the proof.
Prop. 2 of the main paper then follows by combining Prop. 8 above with the convergence guarantee
Thm. 1 of Prospect. Indeed, Thm. shows that Prospect is able to converge linearly for arbitrarily
small > 0 and long up- Under Prop. 8, the minimizer will be equal to w/-
23

---

Published as a conference paper at ICLR 2024
CVaR P(a)
Extremile P(a)
ESRM P(a)
0
02
03
0
02
03
01
02
03
Figure 7: Visualization of uncertainty sets. Top: Illustration of the permutahedra P(o) within three-dimensional probability simplex for the CVaR (21) (top left) = 0.25 (light purple) and p 0.3 (light purple), the extremile (22) (top center) with 2.5 (light yellow) and b 1.5 (dark yellow). and the ESRM (23) (top right) with 7 = (light pink) and y = 1 (dark pink). The size of Each point within the colored shapes represents vector of probabilities over n = 3 samples. The vertices of the polytopes are various permutations of the vectors, whereas the interior of the polytopes
P(a)) increases for more non- con-uniform spectra
elements of the convex hull of these permutations. Bottom: The a associated to each colored shape is given by the same-colored bars the bottom rows, i.e. the CVaR (bottom left) p 0.25 (light purple) and p = 0.5 (light purple) from Equation (21), the extremile (bottom center) with b = 2.5 (light yellow) and b 1.5 (dark yellow) from Equation (22), and the ESRM (bottom right) with y = (light pink) and ? 1 (dark pink) from Equation (23). The parameters (p, b, 7) are risk-sensitivity parameters that increase or decrease the size of the uncertainty sets.
We interpret this phenomenon the "hidden smoothness" of Fo. in that the non-differentiable points of the map w - maXgeP(o) q {(w) are precisely the points at which ((w)) (w(w) for some i * j, the subdifferential may contain multiple elements (Mehta al., 2023, Proposition 2). Thus, if the losses are well-separated enough (in comparison to the spectrum σ)) at the minimizer ww. the objective for the non-smooth setting minimizer.
and regularized setting v > 0 result in the same
Spectral Risk Measures. Recall that for spectral risk measure with spectrum (a), -01,-..,00 Rn), the uncertainty set (ar) is given by:
P(o) = ConvexHull ( ((=())........ .a(m)) is permutation on [n]]
(20)
As for particular examples: for p E [0, 1], the P-CRaR (a.k.a. superquantile) (Rockafellar & Royset, 2013; Kawaguchi & Lu, 2020; Laguel et al., 2021) requires that = up elements of a non-zero with equal probability and that the remaining and expxoneetiaal spectral risk measure (Cotter & Dowd, 2006) define their spectra the equations below.
zero. The b-extremile (Daouia et al., 2019)
1
if E [[n(1 .......,n) [n(1 - p)] <i< [n(1 - p)1
0,
ap
up
if
p-CVaR, p (0,
(21)
0,
(1)
1-1
b-extremile, b b2
(22)
n
(E ()-ua)
-ESRM, >>0
(23)
1-
24

---

Published as a conference paper at ICLR 2024
The multiple cases in the CVaR definition account for the instance in which np is not integer. In light of Lem. 7, when v > O, we have that the objective based on the uncertainty set (20) is differentiable with Lipschitz continuous gradient (i.e. smooth). To be more specific, may call this the regularized or smoothed spectral risk measure. On the other hand, as descrihed in Sec. 2, SRMs were historically computed for any loss vector 1 R as
Rd
max (EP(2)
q
Σ O,10) 1=1
(24)
where ((1) .... ((m) are the order elements of I. For this reason, SRMs may also be called L-risks (Maurer et al., 2021), based classical L-estimators (linear combinations of order statistics) from the statistics literature (Shorack, 2017). In fact, they were originally introduced as functionals L, on arbitrary real-valued probability measures (Acerbi & Tasche, 2002) given by a weighted integral of the quantile function
T. [.,[F]] J.'' (D)))))
where (p) in F(x) p} is the quantile function of cumulative distribution
function (CDF) F and is a non-negative, non-decreasing function satisfying Jo s(p) dp = 1. To recover (24) view I, a random sample and define Fn the empirical distribution induced by the sample. Then, the quantile function (p) is given by the order statistic (()) when
p ( and the discontinuity points of F, are defined make it left-continuous. Applying Lg to the empirical CDF yields
s/n
Lg Fn
F (p)s Σ
FØ (p)s(p))
(1-1)/n
i/n
Σ
(0)
s(p) dp Σ 0,4(3)
(i-1)/n
1=1
for o; = - discretization of a continuous spectrum over [0, 1). Notably, the permutation-bast description of spectral risk measures in (24) is unique feature of the their discrete versions. For a visualization of the feasible set P(rr) for the CVaR, extremile, and ESRM, see Fig. 7.
s(p) dp. Thus, the formulae for d; for each SRM are defined by taking n-bin
C EFFICIENT IMPLEMENTATION OF PROSPECT
In this section, we describe Prospect including computational details, in a way that is amenable to implementation. For convenience, the conceptual description the algorithm from Sec. 3 restated in Algorithm 1.
Efficient Implementation. We exactly solve the maximization problem
g = q qpt (l) arg max
q T (v/n) Σ f(ng)
(25)
gEP(O)
by sequence of three steps:
Sorting: Find π such that (x11) S... S ([()).
Isotonic regression: Apply Pool Adjacent Violators (PAV) (Subroutine 1) to solve the isotonic regression minimization problem (13), yielding solution cPP(().
Conversion: Use (14) to convert e back (I)
The sorting step runs in On In) elementary operations whereas the isotonic regression and con- version steps run in O(n) operations. Crucially, retrieving q from the output cPM(() in the third
25

---

Published as a conference paper at ICLR 2024
Algorithm Prospect (restated from Sec. 3)
Inputs:
Initial wo, spectrum number of iterations T, regularization >> 0, shift cost v > 0. Hyperparameter:: Stepsize
0.
1: Initialize / + ((wo) and dJ < (w,(w)) + wwo for = 1, .....
2: Set f- arg (())
I (4 |n/n) and p k q.
3: Set < Pig: Pig: E IR"
4: Set w f- wo.
5: for T iterations do
6: Sample i. Unif[n] independently. x (q.(VL(() + pw) np.9i, + g. w + w -qp.
7:
Iterate Update
8:
9:
Ij Ez(w).
Bias Reducing Update
10:
q arg max 7T6 -D@
11:
g - Pag: + qi (V,(w) pmc)
Variance Reducing Update
12:
91 + ((((w) + www
A q.
Output: Final point w.
step can be done by single ((n)-time pass by setting
(m()) = - -
-c:)
for i I,...,,t n, opposed to computing the inverse - and (14) directly, which in fact requires another sorting operation and can be avoided. Because only one element of changes on every iteration, may sort it by simply bubbling the value of the index that changed into its correct position to generate the newly sorted version. The full algorithm is given Algorithm 2. We give brief explanation on the PAV algorithm for general f-divergences below.
Pool Adjacent Violators (PAV) Algorithm. First, recall the optimization problem we wish to solve:
n Σ
" (7(1) -G 9i(cal), where go(c;;)) == o; + n ((
min CER
(26)
Son
The objective can be thought of fitting real-valued monotonic function the points (1, (1,4=()))...,, (n,) (())). which would require specifying its values (cI,..., Cn) on (I....,n) and defining the function as any I E [ej, c1+1] on (j.j + 1). Because ((1) S... I (x(n), if we evalu- ated our function (cI....,, Cn) on a loss such as En (ba(2) - c;)², , we may easily solve the problem by returning C (x(1)--.. On Lx(n)- However, by specifying functions -1.-..: In. allow our loss function to change in different regions of the inputs space {l,... n). In such cases, the monotonicity constraint Cj I ... S C₂ is often violated because individually minimizing ge(c,) for each c; has no guarantee of yielding function that is monotonic.
The idea behind the PAV algorithm is to attempt pass at minimizing each 9s individually, and correcting violations as they appear. To provide intuition, define ot arg min, eR g;(c;), and consider i < j such that ci > G.. If f is strictly convex, then gi(x) > g:(c²) for any x < o, and
similarly gj (x) > gj g;(c)) for any > Gy. Thus, to correct the violation, we decrease c to ē and
increase oj to εj until & Cj. We determine this midpoint precisely by
c = 2 = arg min (z) + (z)
zER
as these are exactly the contributions made by these terms in the overall objective. The computation above is called pooling the indices i and We may generalize this viewpoint to violating chains, that is collections of contiguous indices (i, + 1,..., + m) such that of < e for all j < and oj cf4 for all j > i + m, but of c/4 One approach is use dynamic programming identify
26

---

Published as a conference paper at ICLR 2024
such chains and then compute the pooled quantities
C: = arg min
Σ g++x(z). FER 1=1
This requires two passes through the vector: one for identifying violators and the other for pooling. The Pool Adjacent Violators algorithm, the other hand, is able perform both operations in one pass by greedily pooling violators as they appear. This can be viewed as a meta-algorithm, as hinges the notion that the solution of "larger" pooling problems can be easily computed from solutions of "smaller" pooling problems. Precisely, for indices S C
[ra] = (I,...,m nj define
Sol(S) = arg min Eg((r). =ER
$ES
We rely on the existence of an operation Pool, such that for any S, T,7 S [ni] such that SnT = 0, we have that
Sol(S Pool (Sol(S), SolT)) m())),
(27)
where mS) denotes "metadata" associated to S, and that the number of elementary operations in the Pool function is ((1) with respect to |S| + (T). We review our running examples.
For the diverrggee we have that (xx) x² - 1 and (T) = y²² /4 4 + 1,
Sol(S) = arg min .ER
x
Σ oi iES
+ |S| + 1
Σ (1a()) - xx²
4nv
iES
1
(i))
2nv
Σ
S
LiES
'ES
1
US
(x(())
2nw
Σ ai
[ + [T]
LESUT
VESUT
|S| Sol(S) [T] Sol(T)
[S] + [T]
Thus, the metadata m(S) = |S| used in the pooling step eq. (27) the size of each subset.
For the KI. divergence, fKL. (z) = a: In and (v) = e- exp (y), so
Sol(S) = arg min sER
Σo, NES
exp (1( exp (-2/ x/v)
iES
In
exp (Lx(s//2)
In Σ iES
o. Inn
iES
Sol(SUT)
In
Σ exp (la10//)= In Σ oi
Inn-
dESUT
$ESUT
v In
Σ VES
exp (1/0//)) + Σ exp ((200)/o)
-In Σo;+
Σo;
In
iET
iES
iET
Here, we carry the metadata mS)) (In Eies exp (lz()/v) In Eies o;), which can easily be combined and plugged into the function
(m1, m2), (m/, m/)) + v [In (expm + exp - In (exp + exp - Inn 1]. (28)
for two instances of metadata (m), and (m/, ml). We carry the "logsumexp" instead just the sum of exponential quantities for numerical stability, and Equation (28) applies this operation as well. It might be that Eieso, o. = 0, e.g. for the superquantile. In this case, we may interpret
27

---

Published as a conference paper at ICLR 2024
Algorithm 2 Prospect (with exact implementation details)
Inputs: Initial points Wo. spectrum stepsize > 0, number of iterations T, regularization parameter μ > O, shift cost O, loss/gradient oracles .1,...,6, bn and Ve,,..,, Vln.
1: ((wo)RR"
2: 9 - (VE,(w)) + pwo) R=xd
3: argsort(!).
4: c t PAV(L, T.
> Subroutine or Subroutine Subroutine
5: 9 I, M, On)
6: P + 9.
7: 9 < Di=1 Pig: E Rd.
8: for T' iterations do
9: Sample i.j [ni
10:
+ ng. (Ve((w) + paw) np:9i -
Iterate Update
11:
w + w -ne(t).
I 4 (w().
b Bias Reducing Update b Subroutine
13:
+ Bubble(T, l).
14:
C + PAV(I, T).
15:
- Convert(e,,, q).
16:
9 + 9- P.g: + (V((V + nw)
Variance Reducing Update
17:
9: + V(((w) + pww
18:
P < ςi-
Output: Final point w.
Sol(S) = oo and evaluate cxp (-30) = in the conversion formula (26). Two examples of the PAV algorithm are given in Subroutine 1 and Subroutine 2, respectively. These operate by selecting the unique values of the optimizer and partitions of indices that achieve that value.
Hardware Acceleration. Finally, note that all of the subroutines in Algorithm 2 (Subroutine 1/Sub- routine 2, Subroutine 3, and Subroutine 4) all require primitive operations such as control flow and linear scans through vectors. Because these steps multiplications, they benefit from just-in-time time compilation on the CPU. We accelerate these subrou- tines using the Numba package in Python and are able to achieve an approximate 50%-60% decrease in runtime across benchmarks.
outside of the purview of oracle calls or matrix
Subroutine Pool Adjacent Violators (PAV) Algorithm for x2 divergence
Inputs: Losses (4)) argsort æ, and spectrum (d:)
Initialize partition endpoints (bo, b1) (0, 1), partition value ii (m(1) 2nvo,, number of
parts k = 1.
2: for 2,
do
3: Add part h: = k + 1.
4: Compute U22- (x(() 2nvo;-
5: while & > 2 and Uk-1 Uk do
6: Uk- (bg hy 1)14 (i-be))
b- bk-1
7: Set
8:
by
Output: Vector e containing & uk for b <i g bc.
28

---

Published as a conference paper at ICLR 2024
Subroutine Pool Adjacent Violators (PAV) Algorithm for KL. divergence
Inputs: Losses (4)) argsort π, and spectrum (di))
1: Initialize partition endpoints (bn, b1) = (0, 1), number of parts & = 1.
2: Initialize partition value U1 = // (4/1)// 1n -In - 1).
3: Initialize metadata m = (x(1)/ and 6 In
4: for i = ...... do
5: Add part = k ++1 1.
6: Compute U/
(lx() /v - Ino; Inn 1). Compute m, = La(t)/ and 4 = In
7:
8:
while & 2 and lik 2 Uk do
9:
m (ogsumexppmm--1.nee m2) and 44- logsumexp(te-1,
10:
Uk-1 = V (m;-1 4--1 Inn
11:
Set k 1.
12:
by i.
Output: Vector e containing Q uk for b <is S by-
Subroutine Convert
Require: Sorted vector e E R. vector E R", argsort o, shift cost v 2 0, vector E IR. 1: for irri= In....,n..... do
2: Set Qm(s) = (1/n)lS*1 ((<+()) ci)/v).
3: return q.
Subroutine Bubble
Require: Index Jimis sorting permutation loss table 1.
Jmit-
(a(aa) too small, bubble left.
2: while > and (x() < (m1j-1) do
3: Swap T()) and ( - 1).
4: Jinit
(m(jma) too large, bubble right.
5: while j n and (Ú) > (xj+1) do
6: Swap T() and ( + 1).
7: return
29

---

Published as a conference paper at ICLR 2024
Algorithm 3 Prospect (with iteration counters specified to accompany convergence analysis)
Inputs: Initial points w((), stepsize > 0, number of iterations T.
(0) - w(o) for all E [n].
1: Set 2;
2: (0) arg max) ((w =l(()), pl = =(()
3: Set (10) (4(C) (0) ) E ER", .R" g) (Vri(z) () 1)=21 E Rx .g =0==Ei=>> (0) 9;
E R..
4: for = O...., do
5:
it [n]]([[]]), Unif([n]).
6:
7:
(1) Vr;, (wlt)- (mp, Vri, (2(!)) g(o)). w(t) nqi.
Iterate Update
8:
((+1) = w(() - mu(t).
9:
10:
(1+1) = w and (/(+1) = - for j + jr-
b Bias Reducing Update
11:
((+1) (((((t1))
g (++1) arg maPoE 9 (((++)) (8(()).
13:
14:
(++1) w and 2; [(+1) zi
(1)
for i + in.
Variance Reducing Update
15:
g (+ = (Vri(2(++1))) ));-
16:
Pi (I+1) gi,
(1) and Pi ((+1)
Pi
(1) for i + it.
g(++l) E Pi (1-1) (++1)
17:
9i
Output: Final point (T)
D CONVERGENCE ANALYSIS OF PROSPECT
This section provides the main convergence analysis for Prospect. For readability, we reference the version of algorithm presented Alg. 3, which written to match quantities appearing in proof. We begin with high-level overview, whereas the remaining subsections contain technical lemmas of interest along with key steps in the proof.
D.1 OVERVIEW
Notation used throughout the proof is collected in Tab. 1, and is also introduced as it appears. In the following, we denote A L + μ the smoothness constant of the regularized losses r,. We denote E, the expectation w.r.t to the randomness induced by picking it. given w!!), i.e. the conditional expectation given w((() The optimum of (8) denoted w" and satisfies
V(a" r(w")) = 0, for q arg maxq ((ww) - gEP(o)
(29)
Denote in addition 1' = f(w*). For simplicity, we use the shorthand 2(g)
1 (D((((n//))
nan
for any f-divergence Dy. where on is the strong convexity constant of the generator f over the in- terval n). By Prop. 4, this gives that Ω 1-strongly convex function over the probability simplex. All forthcoming statements will reference the setting of Algorithm 3. Note that when implementing
the algorithm, storing the iterates (2 (4) Ji=1 and {c, Ji=1 is not necessary.
Proof Outline.
We argue convergence by way of defining a Lyapunov function (tt,, which will upper bound the quantity (t) - w which will be called the "main term". Specifically, define
vi = [[] wn lla 5(0) + (0) + ct c.aR R(R())
30

---

Published as a conference paper at ICLR 2024
where C1 C22 Czs and C₁ constants to be determined later, and
1
5(0
Imp!
Vr:(z,))
ng Vri(w")]
T Σ
11 w 12
m
U
|
1/. R = -r) 17).
Though not included in the Lyapunov function, we will also introduce
Inq. " Vre(w( nq. (r
When the shift cost is large, we will be able to simplify the analysis by excluding the terms (7(1) and R(()), The colors are used for the convenience the reader so that the quantities above are easy to track from result to result. Each term in the Lyapunov function is motivated by terms that appear when bounding the main term [| w"1 which appears in Step 2 below. The outline of the proof is follows.
1. We introduce a lemma of general interest which is the key technical step in the analysis: bounding the bias of the gradient estimate w(() given in line 7 of Algorithm
2. We expand the main term and identify "descent" and "noise" terms, as standard anal- ysis of stochastic gradient methods. The descent term will be treated using the technical lemma from the previous step, whereas the noise term will be upper bounded using stan- dard techniques. In either case, additional non-negative terms are introduced that will be incorporated into the Lyapumov function.
3. We the bound the evolution of the Lyapunov terms that are not the main term. For the large shift cost setting, only 5()) and ( are needed, while (()) and R(O) can be ignored.
4. In the final step, we tune any free constants that are encountered in previous steps. At this point, we split the proof into two subsections, one for the large shift cost and one for any shift cost.
D.2 KEY TECHNICAL LEMMA: BIAS OF THE PROSPECT GRADIENT ESTIMATE
In this subsection, we present the key technical step that allows for the unconditional linear conver- gence of Prospect. When analyzing stochastic gradient methods in the smooth and strongly convex setting, we typically expand
E, (t+ w01 [| w011 2n [E ( w²) + 17 Eel E
(30)
dosccnt term
noise Term
First, note that the expectation of the primal gradient estimate oft) is Vr(w(t)) (0) where (()) arg maXgeP(o) g (()) #Q(g), where ((0) R= denotes the estimate of the full loss vector. Applying standard convex inequalities to the descent term yields
- -(Vr(w(0) (,)))
(()) - w²) S J.M
|| - w.ll
M
1
Σ
(0)
Vr.(1o(t)) Vri(w=))
9
# I
+
(vr(w") (0 )
The first two terms on the right-hand side are negative, which provide decrease in the expected distance-to-optimum value on every iterate. In the empirical risk minimization setting, the final term on the right-hand side would be zero due to the first-order optimality conditions on w*, as (()) 1n/n, implying the decrease of Ex]lu[(++1) w0113 1|3 for small enough y) However, because
31

---

Published as a conference paper at ICLR 2024
(()) is potentially non-uniform vector estimated using the table of losses ((0) (as opposed to the loss vector f(w*) at optimum), the term - (Vr(w*)T
w(() - w²) is non-zero. Additionally, this term multiplied only by the learning rate 7, instead of the noise terms which are multiplied by n². Thus, this bias term must be bounded carefully in order to achieve the convergence guarantec under this regime. This is the subject of Lem. 9, the main result of this subsection.
Lemma 9 (Bias Bound), Consider any E IR 1 E R", and @ P(a). Define
q = q1) arg max p #81(p). pEP(o)
For any any B1 E [0, 11.
-Vr((w) T q- vr(w") 9) T) (w -w*)
5 T) (E(1e) ((ww)) | w' w²11
31
Σ
Inq, Vr.(w) - nn Vr.("" 112 +
7(M + 1)) n1 (1 (1-1))
G3
4(M + 10)15
1=1
Proof. First, for any 9: > 0, we have that w + (ww is (q:M)-smooth and (q.uu)-ttroglly convex, so by applying standard convex inequalities (Lem. 30) we have that
grr 2 (w) + 4 (r: -w)
+ 24 (M + p)) 1/: (w) Vr.ww"))) 21ft w| w= 1
4
ar)(w)) 4: (w) (w- w)
1 (w) vri(w>)11 +
9ift I|
w01
20, (M 14)
4
as q: S Gn. The second inequality holds for q; = 0 as well, so by summing the incquality over and using that E. qi 1, we have that
q T (T(w") 2 q r(w) + q T Vr(w)(w* w)
+ n0 4)) Vr:(w 4i (r,(ww)) )12 + # w1 w0ll
Applying the same argument replacing & by and swapping and w* yields
q r(w)= q r(ww) q Vr(w"") (w w²)
+ 1 (M 11)
Σ Vrsw)) - g. Vri(w")) 112 + 1|w0 w²
Summing the two inequalities yields
- (q- - ②) (r(w) r(w*))
2 - 2 9- Vr(w*) TT)) ww) + 2 w| w=1
+ (M (M
1
Σ l1 (w:(w) v. Vr, (w 012
| (r: Vr: (w
p1)
32

---

Published as a conference paper at ICLR 2024
Dropping the Σx= l1 (rr q: Vrrr(w")) 21l term and applying weight of B1 [0.1] to E l1 Vri(w) - 91 Vr: (w=" (12 still satisfies the inequality, which can equivalently be written as
(Vr(w) - vr(ww) 4-V(w)) (w - w I -q ②) (r(w) - r((")) - / || w01
2
Vrrr) - a Vri(w )| (31)
20 + jL)
Next, because
l1 Vr.(w) o Vr, (w" )12 S 2 11, Vr:(w)) 2 Vr.(w=) 2(q. -3 (r:
we have that (by summing over i) that
n
1 Vrr(ww - q. Vr. 2 Σ [l..vr(()) - 9 Vrr(w>)1 + 4G2 114 - 21, (32)
where we used that each [|rr("")]: )| 2G. To see this, use that Vr(w") q² = 0 and Vr(w")= V((w") + pww", so
VVr Ill = [[ + pew 1l2 = v[( - EE a (w=) || s 2G.
Because the map nht is the gradient of a convex and (1/7)-smooth smooth map, we also have that
11g q" 11 = ( (e(())))) S { - 10-9))(( {(w*)),
(33)
so we apply the above to (32) achieve
1/ (w) ai vriw"11
4G² q (w) -: vrr((")))
D (a (1
(34)
We also use (33) claim non-negativity of (q (q-"")"Teppp TT (1-((w")) )). Finally, because Eg Es % 1, we have that
(q - (r(w)) - r(ww))r(ww)) = (4 - -T (f(w) + p
|u l 1n - -((w") / = (q - T) (e(w) ((w") )) + (q 7 In
2 || 1 1n
[| 1| 112
= (q - T (f(w) E(w")).
(35)
Combine (31), (34), and (35) along with Kg = no,, to achieve the claim.
In the next step, we apply Lem. 9, as well as noise bound to give the initial per-iterate progress bound on the distance-to-optimum quantity Ex, ()(])) w" [.
D.3 BOUNDING THE COTAANCETTO-PTTINN
To outline the remainder of the proof, observe the expansion (30). With the bias bound for the descent term in hand, we now upper bound the noise term.
Lemma 10 (Noise Bound). Consider the notations of Alg. 3, we have for any B >
(0) 1 < (1 + B)E, Inq (c) Vr., (w(()) ng;; Vr,, (w²)
+ (1 + 9 D, E, Imptr Vra (2 ng. Vri (w+) 112-
33

---

Published as a conference paper at ICLR 2024
Proof. In the following, we use the identity E|X 1| = E| - ||[ in equations
denoted with (+). We denote by 8 an arbitrary positive number stemming from using Young's s in equation (o). Noting that Σ 9, Vr;(ww)
inequality ||a+bll = (1+ (1+ ||6||2 0, we get,
Ee [10000 V(g" (( r)
Ee IInggl (t) Vri, (w) - -ngi (w=)
+ ng Vrs, Vri, (2*) - !p! Vri, (2;, (1)) (V(g" (V(o"T)((w)) ()(")) - -
(2)
(+) ((((t)) (2)( V(g'' T r) )|2
+ Fr |ind ( Vri, (w - nn, Vri, (w - (V(10) ()(
- V(a") T)((10) r)())) r)a²))
+ +ng, Vri, (w=)(w") - np: Vri, (10) (V(a") (T(aTTr)(us)) r)(ww) - )
(b)
[V( ()) - -(('' T (m==))1
+ (1 + B) |! (1) Vri, (1) mut. Vr; (w*) - (V(('') (w)()) V(0¹ r)
+ +1+ 8 LE, |l Vr,, (w*) - -mp," " Vrr. (2 (V(x' (rww)) 002]133
0 (9 T ()(())))) - V (w) (r 11
(T)
+ +(1 + T ||nun ( Vn
nng Vra (w) 11
+ +1+ + 8 31,, [11 bB Vra (w*) - -mp (0) Vri (z >|
(1 + 8 || V (ww(t) - 1/.
We then combine the analyses of the first and second-order terms to yield the main result of this subsection.
Lemma 11 (Analysis of distance-to-opiimun term). For any constants B1 E [0, 1] and B2 > 0, and any q E P(a), we have that
E, (t w01l2 < (1 - m711 (wo()) 11
(w((w - - (V(w=*j
31
17(1 + B2) Q(() + 1 + 13, 1800
2(M + M)tio
2/3 G2 R(() (0)
+
- ()) (e(w)) - f(w"))
DM + pi
Proof. Recall the expansion given in (30), which is:
(e]lw((+1)_ w01l (1oo(() w01l 2n/ ( [E.. (u w0)) -w0) + +, 11
(36)
Observe that
E, [u(0))]]
(1)
Vr( = (r(w((0)TT())
34

---

Published as a conference paper at ICLR 2024
By Lem. 9 with / = 1(6), 9 q(), w== wlt), and multiplying by 21,, we have that
-2n(w w²) (v((w()) g(()) S -2n(w w²) Vr(w") ) 2n(4 a) (e(w((()) {(w"))
un [w(()) - w/ 112
n Q
2(M 10
23, G² R8
7(M + 11)6a
The noise term is bounded by applying Lem. 10, so that for some B2 > 0,
1f , 2 11 (2) + (1+ (5()
Combine the two displays above to get the desired result.
Now, we see the appearance of S(0,, T(O, UO, and R(O in Lem. 11. We incorporate them into the Lyapunov function and describe their evolution in the next section.
D.4 BOUNDING THE EVOLUTION OF THE LYAPUNOV FUNCTION TERMS
We now describe the evolution of the terms ( T(0 Ut)), R(()) from iterate / to iterate + 1.
The first two terms simply measure the closeness of the iterates {2 H within the table to the optimum w", measured either in closeness in weighted gradients (S = Σ Impe
and {c, ) )
(0 (0Vr,, (=1() (() ngt, Vr., (l3) or directly (T(1) = E |1c/ w())
Recall that Q ()) = Σ Ing. (1) (r,(((!)) - ng, Vr.(w=)133
Lemma 12. The following hold.
Eg s((+1)
1 -00 + -1-11 5(1))
11
E2 [(++1)
= 11w)( w²12 + 1 1 T()
Proof.
Write
E, 5(++1)
n Eg !np] ((p(( (r:((((-1)) bbn Vr.(w=))
|ng (1) (w)(w(()) oF Vr:((")) + (1- !!mp!" (ri((()) nn. Vr:
Q0)00 1 (1- n
5(O)
Similarly,
(1c(ee11) 20 212
E,
E2
i=1
20 + (1-1)
1
||
w" 113
n1
w=1 112 + (1-1) )) ((t)
(t)
35

---

Published as a conference paper at ICLR 2024
Next, we handle Ut)) = 1 شー= [1x(0) - which can be ignored if we assume a particular lower bound on V.
Lemma 13. For any value of B2 > 0, we have that
Es (
( + 5 (1+ + (1++82))5 (5))5(0)
ny (1-1 T(o + (1- n2 20pm
G2
+11-2) U70.
R +
un 72
n2
Proof. First, note that a separate index je (independent of i,) is used to update (c])) ) so we may first take the expected value with respect to je conditioned on iz:
[((+1]]
1
- 12
Eg
Et
n
1
Ex
ro(e+11]
- 1 2
+
11
1
E2
1
[w(tt)) _ -,,t22)13 1/ je F
j=1
1
E2
(((+))
t1 |22 + 1 -
1
E₁
(e+(1 - 5555 112
j=1
172
(1))
27
+
(1-
1
Er
(((+++) - -,,-0
==1
Next, we expand the second term.
E
[1o(++1)
5.50 11
n
1
E,
[wo(t+11) - 10
2 Ex
n
-000))))o((0)Tooaooo) (wo - (o(())((0())) ((000))()))) -c()))))))))))))))))))) -5/0)
- Ee
1
12
1c
(t)
n. (0
n
n
n
/=1
= =
Es
[10]
2n
V ()) (()))))))))))))))))))))))))))( - ((t)
115
w0) 12.
The first term is simply the noise term that appears in Lem. 10, whereas the last term is (t). Next, we have
-2V(2 (w) = 2( T (1000) - v(g ))((()))) (wol0) - (
T) (w!!t) (1)) 2(V(g(()) T r)(w*) V(a" TT (r)(ww) )) (w)()) - (1))
2(V ()((/0) -VGit )))
where the last term is introduced because v(a* Tr) = 0. We bound each of the three terms. First,
2(V (w))))) V ()(()))))(() (/ S -2 2m/|mm (1) 2
because g(o) is p-strongly convex (Nesterov, 2018, Theorem 2.1.9). Second,
-2 r)(()0) v(g(0) ()())) (wn(t) - -0 < Ball (t r) - v(g (w0 + + w0
< BuM221 (0) w11 + , - (0) 1
36

---

Published as a conference paper at ICLR 2024
by Young's inequality with parameter 84 and the M-Lipschitz continuity of v(y(0) r). Third, -V((() r)("") ()("""))((""))) - V(a") T r))("))) (w)l)) - (30)) = (0)((((())'')(((())) )) (w - =²)
(0) (3/))
S {
-00((((0))0'0000)) () (w )13 + + 11, w 101
s & - 2112 + 1,-551155 ||c w 10
by Young's inequality with parameter Bg and the G-L.ipschitz.coniinee continuity of each &. Combining with the above, we have
2 v(g T ()( - (,9) S (2M + (3T + 5/5 - 22 + 3ssG2nlla(0) - g'11
S ي s M2 + 2 G nllgg - -011
when we set 81 = 33= = =x Hence, we get
[(++t))
Eg
2 E2
27
+
1
E,
n [w
- 12
n
72
n 3=1
S B
1000
1
2
[(g(t) T
r)(c -c,0) + (1-1) (1))
_
12
11
71
S E
1
1 11-- [ 1--M2T10
+ 1 G nn - x 11] + (1- n
((0))
= 2
1 11_1 1
nM T(O)
+
(1
-
1
G2
2nnll
- --- +
1
1
IT(1)
um
n1
2np
n
17 E
n22 T(O
1
1
G² R(1)
(1-1) 1
rtt)
n
um
n
2mp
n
S (1 + 22) (1)
+ 1/ (7 + B3 151)
+ 2yM²
1 (1-
2
-1-1)
um
n
T + (1- n
20pm
R
+
n
U(0)
where the two last steps follow from Lem, 10 and Lem. 34 to claim )*).
-0010 s (g (0))(()-
Finally, we consider R(O) = ()) - q²) T - 1t). This can be viewed as a measurement of orthogonality between the vectors (()) - q² and ((t) which in turn can be viewed as the directions to the optimal gradient and optimal solution of constrained optimization problem. Indeed, we may define
[" arg min h(1)= = max 4 -10((())))) lEC
qEP(O)
and
C= (= 12 l(w) for some E Rd).
where the inequality is taken element-wise. The set a convexification of the set ((R:)) which shares minimizer and has the same minimum value. Indeed, letting | be any minimizer of h. select i such that / = f(w). Define q = Vh(l) = arg max max)E q i- =fl()), and due to the non-negativity of y. we have that
min h() = h0) = q i - - 29 (T((T))))) - =l(()) lec
37

---

Published as a conference paper at ICLR 2024
Taking the maximum over shows that mingg h(I) = h(((@)). For convexity, for any 4 E L
satisfying / 2 {(w1) and lz 2 f(w2), and any 8 E (0, 1), apply the following inequalities element-
wise:
01 + (1 - 0) 2 0f(w1) + (1 - 0)(((z)) 2 f(fur + (1 - 0)w2).
with owz + -1- E Rd. By convexity, (g(0) - 9¹) T (1(0) 1t) 2 0. Finally, this term is of particular importance because the term -(g 9²) (f(w)) {(w*)) that appears in Lem. 9 can be used for cancellation in this case. The next result describes its evolution.
Lemma 14. For any B3 > 0, it holds that
Es
[ S 2m(a - 9¹) (((x -1-) + 1 1
R(0
n
' 17(((
27 P3 + 21IG2n (1 + (2)
Proof. First, decompose
(g --) ((((((tt)) -1) = (g))) (((t))))))1177
- -1² + (1 - ((())(((1)))))))))))))1))) (9:1))100
(37)
+ (q(+))) - 10 (()00)1111))
(38)
Because g(0) = (10)) for all t, and (-) is the gradient of a convex and (1/0) function, we have for the second term of (38) that
(q ( - ((())((()))))) S - [ 11
Next, using Young's 's inequality, that is, T S 83 1|al11 + 1013 for any 83 > 0, we have for the third term term of (38) that
-0 (110 (²) - 33
[1 - 10 + 1
1-1 11 - 11
-
2
2
Ag
100 + 3/
|||(+]])
(1(00) - 13.
2
Note that we have
Es [0+1]] = ((w(())) + (1- (1-1) 1) 1) 100.
17
Hence, we get,
1 Ex 21m1
R(++1)
1 ()))) - (T) - 1²) + (1 (1- (g)) 1 -100000 12)
n
n2
+ Be (g (+11) - 100 (((+t))) - ]
+ F. (11) - (0) T(10) - rr)]
()) - ((((000)) --- -1+)+ + (1
1
(0)) - --) 1100 110--))
1
Ba
Es ||(ct+)) - [
8-1
+ b
+ 20
+
11(00 - 11
2
(1))
g (e(w -1r) + (1-1)
(q -11-2)
1
- TT)
T (100) - 1T)
Ba
n
1+
Σ (e)((019)-1 - =2(()))
nw
27
1=1
35 1
(Cg(C)) - =2
2
1=1
38

---

Published as a conference paper at ICLR 2024
Then, apply the G-Lipschitz continuity of each & to achieve
1
R
(g
--"([((! - -1)+ (1- (0 -²) T(()) 1160-1 1))
Ee
25mn
G²
1
83
|1
1/2112 112
ni
27
j=1
G
3g
15)
w
112.
2
/=1
Replacing 83 by 2003 gives the claim.
D.5 TUNING CONSTANTS AND FINAL RATE
Recall that our Lyapunov function is given by
V ||w w + ci + e2T(1) +cgU(0 +
Recall in addition the definitions
500
Imp, (2) Vri(z(") - ng. Vri(w")1E
T( = Σ
11 - -w- =='12
1=1
U(
(1w(())
112 R -0)) 1*).
n
We will derive a value > 0 such that for all / 2 0,
E2 [v(++]] (1 (1-7-1)
In order to describe our rates, we define the condition number == M/U and recall that Kg = 10
D.5.1 STEP 3A: ANALYZE LARGE SHIFT COST SETTING.
The following gives the convergence rate for large shift cost.
Theorem 15. Suppose the shift cost satisfies
D > SnG2 /n-
Then, the sequence of iterates produced by Algorithm 3 with 1j = 1/(11(/ + +M)ra)) achieves
(l) w011 s (1 + o. o
+ t (xp((tt))) -w:-ww) w='13
with
T 2 max (n, 24kig (n + 1)].
Proof. First, invoke Lem. 11 with o' = (0()) and B1 = 1 to obtain
E, (t w²1 w < (1 m w'
(39)
20
-2n(ww( w Vr(w")a 7(M + p) R((0)
(40)
1
1711 + Ba)) Q(0)) + n 83 13 15(0)
(41)
2(M p)no
39

---

Published as a conference paper at ICLR 2024
We will first bound (40), by using that Vr(w")q² 0 and Young's inequality with parameter 0 to write
[(w(() w²) - (r(w")((1)
= (vo(1) - w²) Vr(w") (r(0) - 9r){
- 2 Vr(w") )T (g - - 2 + 2a
12
1
(1)
w0
TG T 2a
w w²
202
where we used in the second inequality that:
(r(w") - 8311 -
V((w")) ((
-9)1 S 2 02.10000 -
9-)
1,2
22 D2
[00)
1- |
00 2² 2 Tt)) x
We also have by Cauchy-Schwand and Lipschitz continuity that
R = 2 --) -1) -
27m
[(0] - 1"
2
<
2mm
T()
1
Combining the above displays yields
2G
21(w(( "="" "o"))r(("")a
12(e)
J(M 11)
G
4nG²
7 wo(l) w0 w011
52 ay + (M + p)
T(()) +
We take 32 = 2, C3 c4 0, and apply Lem. 12 to achieve
E, [v(tt)]
(1 - ()))
< [r -1
1
+ 02) |10 - w='13
qu na
37/²
1
+
C1 S(1)
201
n1
G22
4nG²
02211)
07
12
(M + u)tia
+
1
37/²
C1
Q(1))
2(M + 11)
n
where O is to-be-specified rate constant. We now need to set the various free parameters a, C1. 422 and 17 to make each of the squared bracketed terms be non non-positive We enforce 7 2 2n throughout. By setting
and C1 = 4(1r + M M)Ra²
12(μ + M)s,
have that the bracketed constants before (())) and Qit) vanish. Then, setting
u and C2 48(K + 1)k,
2
make the bracketed constant before (t)(t) I12. assuming that we enforce
2 48(k + 1)ti.
We turn to the final constant after substituting the values of a, 02, and v). We need that
G/2
8nG²
8G2
2nG2
< 1
an + (M + 11)roa)
12 + (N + 1)Ra
J22
22
2n
40

---

Published as a conference paper at ICLR 2024
which occurs when
16nG² 82
2nG²
2 +
(16 + 1)6
Because 1 s v 2nG2, this is achieved when
8nG
1
completing the proof of the claim
E,
(1 (1-T (1-
To complete the proof, we bound the initial terms. Because Cg = C1 O, we need only to bound ((0) and T(O).
S(O)
mp, Vr:((, no; Vra(w*)1B
||q
Vri(w (0)) n4 Vr;(w")) )13
||nq:
(0)
()) Vr:(w"))
2 2
Σ
n
|n(0; (0)
Vr( )12
n i=1
2n Σ q 2 MM (10 w0112 +
d²1
8n²
wn 1n
2n ||o 1/ _ + D2
| - 112
2n | 1/ 21 + 2²² 42188 (10(0) - w01 I 3nM ( ()) 11.
-w-
This means ultimately that
n
(
|w w 112.
16(1 + x 122
Next, we have
(2))
11
(0) w0
48( + 1)10
Thus, we can write
n²
(00) 1 1+
(0 w²
16(1 + x 12
48(K + 1)xr
(1 + on -1 + x
(0) == 2
completing the proof.
D.5.2 STEP 3B: ANALYZE SMALL SHIFT COST SETTING.
To describe the rate, define == nGG2/(D) (ND). The quantity ó captures the effect of the primal regu- larizer 4 and dual regularizer 7 as compared to the inherent continuity properties of the underlying losses.
41

---

Published as a conference paper at ICLR 2024
Theorem 16. Assume that 2 and that the shift cost D 8nG2/P /p. The sequence of iterates produced by Algorithm 3 with
1
"
16mp
min
482 max (2nx² 83 6|88 + (x + 1)Hrrl
we have
E, p((+t)) S ((1-
R2 (0 wn 1| 1 5 166
6x²
exp (-1/T) |00000 w0
on
for
T = 32m max {6|85 + (k + 1) 46 max f2nx² ,
Proof. First, apply Lem, 11 with g' 9", well as Lem. 14, Lem. 12, and Lem. 13, set C4 and consolidate all constants to write
1,
E, [( (1 - - (0
$ (T-1
17p1 + +22
(1) 1
(42)
+
+ 7(M 1)mio
23
1
1
G2 C3
R(O)
(43)
n
2ūju
+
+ C3 ( + 1
(10(())
(44)
G n
C3
1
1
+
3
1
(1)
(45)
2028
62/mn
71
7
2n n
1
(1
+ B3)
CgU())
(46)
CzD
m
n
2
c3)(1 + (2)
C1
Q00.
(47)
2(M + p)ma
We first set C1
40M+ and (1 1/n) < 1. Then, notice in (43) that to achieve
n
nu/2 to clean up (42) and (47). We also drop the terms
G2
1
S
v(M + u)na
4n
we need that B1 S ((M + (SnG) /v). Combined with the requirement that 3 E [0. we
set 81 ((M + (1)),))/(nn(C) + (M p)).)). We set B2 2, and can rewrite the expression
42

---

Published as a conference paper at ICLR 2024
above.
Es [(ttt]] - (1 -11---1vvo
2)
/w(t) -00
1
1
3
G2cz
R(0
4n
2
6(1 cM
+ 11)
1
n
n,
72
Gnn
33' + c2M
Cg (()))))
μD
12
Gn
+
r
(1 + Ba) n2
cz((1)
CzU
m3
+
+ +377 + cz)
Q
4(M + 1))no
Next, set the learning rate to be
B1
" S
(48)
12(1 + c2) + 11)
to cancel out Q( and achieve
Eg y(41 (1 - ----1 vVM)
T711
|w(()
w 2
2
GCe3
R((0)
4n
20p
S1
T
2n
GZn
c3
T
3/5
+
T())
j2n
n
Gg n
1
(1 + 33)
c2U
C3U
Requiring now that 2 2n, we may also cancel the ((t) term. We substitute 8 nG2 n22//pp)) to achieve
E, [yitt]] - (1 -1---111 (1-
7/11
(1(
w"
2
+
Cz
R(t)
4n
2n
1
8
(1)
+
2m
pºn
2pmo
(1
+ Ba)
2n
G3
43

---

Published as a conference paper at ICLR 2024
It remains select e3 and By. As such, set Bg 4nd and use that 1 + 4nd Σ 8n8 when and 2 1/8 as assumed, and so
Eg
[(itt]
(1 (1- (1-T -
-
A/
w
w²
2
2
C3d
R(0)
4n
2n
C3k
C2 (1)
+
4n
16
Cg((1)
+
2n
43
We require now that
n
C3 2
min
2xx2 8
which cancels ('()) and R(*), leaving
Eg p - (1 - - 1)((( s (3 7H
(o())
w(1
2
22 max (2nx², 03 cgtiles.
+
+
2n
From the above, we retrieve the requirement that
(49)
64nj62 max (2nx², ,8)
It now remains to set 17. By substituting in the values for 31 and c3 into (48), have that
B1
S
12(1 + ca)(M 1)no
12(1 + 1818 + (M + 1))ool
(12 6n/6)(Suè + (M + 1)
(12 18p + (M m)nol
96n (80 + (M + n)mo]]
The combination of (49) and the above display yields
1
1
min
96n [8p
(M + m)mo]' 64n488 max {2mx? 8)
1
min
16ny
6(85 (N: 1)ma) 482 max {2mx? 88
We need finally that 2 2/(m)) resulting in the requirement
r ス 32n max {6[85 + (w+ + 1)M=), 2 max {2nk3 {8}} }
This is achieved by setting
T = 32m max [6 + (N + 1)Hso), 452 max {2nxx .0] 1/16).
completing the proof of the claim
E, [v(tt]] (1 -1- <11-7-1)v'
44

---

Published as a conference paper at ICLR 2024
Next, bound the initial terms achieve the final rate. First, we bound 17 which is used in all of the terms. Because 2 1/8,
1
7 - 16np 48 max (2nk², 83
(50)
64np0'
ny
Then,
((O)
npp
(0)
(r: gon Vrr(w*)1BB
71
||rqq
(0) Vr, (w)() nhn Vr, (r,(w")1 1
||nq;
(0)
())((((0)) Vri(w vr:(w"))11+ 2
|n(2 (0)
Vr)vr,("""11
i=1
$ 2n Σ qg )))) 2 1l/00 - w011 + -
Sn²
3 2n ||a M² + D²
|w - w01
w 2n | .2 + +22 1²/88 (1 - -wwwl S 33 - -w=11 12
Continuing with 31 S and (50),
mnB1
S())
((O)
4(M + p)
2
3n (1:
w= 112
MMM + 1))
6nR
||
(1
K)
6n²
(w((0) 11.
on
Next, we have ( n 11|1:(0) w 1|2 and by (50),
(2) !!14
2 (0)
w²
2
< 4 (0) w²
Because (T(0) =0, 0, is bounded trivially. For R((0), with C4 = we have
R(0) = 2nn(g
"T (e(w) u (e(w*)) T)) (E(w)(0)) - {(w*))
2007
(()
[(ww)
1
2n Gn²
(0)
w"
16n
(0)
2
168
||
(0)
w
45

---

Published as a conference paper at ICLR 2024
Combining each of these terms together, we have that
V(O)
5 166
6r²2)
1| |
-==
on
completing the proof.
D.6 PROOF OF MAIN RESULT
The objective is once again
F (w) max q [(w( DD vD2(gll,,/n) 1 qEP(o)
|1
[[(( nan Dy(gllin/n) (EP(a)
max
+ 1
1
n0nn
2
q ([((w) na, (())) + |
max (EP(a)
2
max 9 [(((() Tg + (EP(a)
where (gg) = Dy 1n//n/nna, is the penalty scaled to be 1-strongly convex and we simply notate n..,v V. The previous subsections give the convergence analysis the cases of large and small values of D. They are combined below.
Theorem 1. Prospect with small enough step size is guaranteed to converge linearly for all > If. in addition, the shift cost is v Z (2 uaa), then the sequence of iterates (w(t)) J(21 generated by Prospect and learning rate n (12p(1 + k)ka)=1 converges linearly at a rate T max (n, 24 ( + 1)], ie..
(lI wn S + t )xpp eexp(-t/)\ll w0112
Proof. Combine Thm. 15 (the analysis for E 2 SnG²/p) and Thm. 16 (the analysis for & < SnG2 (1) to achieve convergence for any value of D. Substitute no, v so that the condition > 8nG2 Iu reads G2 (woun)
E IMPROVING PROSPECT WITH MOREAU ENVELOPES
As mentioned in Sec. 3, we may want generalize Prospect to non-smooth settings which arise either when the shift cost v or when the underlying losses (4,) are non-smooth. The former case is already addressed by Prop. in Appx. B. The latter case can be handled by considering variant of Prospect applied to the Moreau envelope this extend the algorithm to non-smooth losses, it also allows even in the smooth setting for less stringent lower bound on v/ required for the O((n + x)) (n(1/e)) rate. The rest of this section contains necessary background, implementation details, and the adjustments to the analysis.
the losses, as defined below. Not only does
E. OVERVIEW
We first describe the method and practical details of the implementation, followed by the conver- gence analysis in the next section.
Notation. The Moreau envelope and the proximal (prox) operator of convex function ʃ Rd
R are respectively defined for a constant 77 > 0 as
Mn[f[(w)- = min =ERd
f(z) + 29
1100 -0122
x112
(51)
(w) = arg min
f(x) + 2m
||w 2112
(52)
=ER"
46

---

Published as a conference paper at ICLR 2024
A fundamental property is that the gradient of the Moreau envelope is related to the proximal oper- ator:
V.M,[](w)) -w - p (w))
(53)
For simplicity, we denote y 2nv.
Algorithm Description. The algorithm nearly equivalent Algorithm but makes the follow- ing changes. We sample it ~ (((()) non-uniformly in the sense that TP [& = = q We do not store an additional vector of weights ,( and use ( in all associated steps. This does not change the expectation of the update direction the control variate, but creates minor changes in the analysis
of the variance term. In the iterate update, we replace the gradient descent-like update with
(0) (1 + nlyn -)
w (++1) prox,,, (1!
The vector (1) adds the control variate to w((t) before passing it to the proximal operator. The second change is that the elements of the gradient table are updated using je and the gradients the Moreau envelope. That is,
((+1) 9j. = Ir, + nn -0))),
9j (++1) = o3 for j jc-
Plugging these changes into Algorithm produces the Prospect- Moreau variant.
Implementation Details. The proximal operators can be computed in closed form or algorithmi- cally for common losses. We list here the implementations for some losses of interest. The proximal operators for the binary or multiclass logistic losses cannot be obtained in closed form, we approxi- mate them by one Newton step.
Squared loss. For the squared loss, defined as {(w) {w -- for x E IRd E R, then
(w)
+ 17 |0:|2 (x -u-
17
Binary logistic loss. For the binary logistic loss defined for x E TRd, E (0,1), w E Rd as f(w) = -y (= w)) (1 y) In(1 o(x ()) yr w + In(1 c" ""), approximate the proximal operator by one Newton step, whose formulation reduces to
19
prox Rs w
+ 77 x
Multinomial logistic loss. For the multinomial logistic loss of linear model defined by W sample (r,y) as ((W) = -y Wx + In In(exp(Wxx) T 1). for E Rd, y {0, ; y 1 1, W e k we consider approximating the proximal operator by one Newton-step. whose formulation reduces to
prox (rxx,e(W)
(W) 25 W - =
z² = 21 A²22:
y 23 22, 22 ((Wx) z3, 23 = (1 1 = =T1
Regularized losses. For a convex Idd + +R, define r(w) = f(w) + (n/2) || Then, we have,
w 1 1p44
prox by (w) = prox
IT
47

---

Published as a conference paper at ICLR 2024
E.2 CONVERGENCE ANALYSIS
Prospect-Moreau satisfies the following convergence bound. Recall that 7+
/V((w))12
Theorem 17. Suppose the smoothing parameter i is set large enough as
{
7+
2nk
min
(V
Am7 1 2x
M
and define constant
T= ==2+ max - 1), 1e(Aki 1)}
for Ki = On/o1. Then, the sequence of iterates (w(t)) generated by Prospect-Moreau with learning
rate 1j M min {1/(4ng
1), k)((n 1)} satisfies
|| E w w² 2- S (n 3/2) (t/pp | w0
We now prove Thm. 17, using similar techniques as in Appx. D.
Additional Notation. Further, we define w = w* + vvrr:"")). By analyzing the first-order
conditions of the prox, it is easy to see that
prOx (wi) = w.
(54)
We will use the Lyapunov function
V = |100) w=
w²
C2
Vr:(w=>)
(55)
M
The first step is to analyze the effect of the update on w( as the first term of the Lyapunov function.
Proposition 18. The iterates of Prespect-Moreau satisfy
(1 + En)
w²
1|
(
w-1 + 2n²on, > 10000 - Vra(w*))
022
(1) w²
+
D2
7?
1+ Mn
g(0)) - - - vri(w=))?
Proof. We use the co-coercivity of the prox operator (Thm. 31) to get
(1+ + pn) E, ((+1)) w² || = (1 + un) E, prox (z) prox mi (w /
< Ee (t)(t) w., prox nr, (n()) prox spr (w,,)
(t) wi, ((++1) w²)
(u((t) w ((()) w²) + Er (1(
wi, w ((+1) (1))
=:
(56)
where we added and subtracted (t) in the last step.
For the first term, we observe that Ex[n[((] = (1) and Er [w:]]] w² E, Vri, (w*) so that
T (E,ju(() wi,] (0 w0) (t w/ll n E, Vri (w (1) - -w-) w²) (57)
48

---

Published as a conference paper at ICLR 2024
For T2, note that
w (t+1) wl0 = -n
(VMnnrr,lu())) (
- - + x
We manipulate T2 to set ourselves up to apply co-coercivity of prox- by adding and sub- tracting V,M, (we,) ) as follows:
T2 (ul) w, V.M,, (rie (u())) Ri + g(o)
(1)
- = Be - wi I (VM,pri(((tt))
77 Be((() wio V.MpIrs](EE,)) - 0/ + g0)
=
Now, co-coercivity of the prox-gradient (Thm. 32) of the M-smooth function ri, gives
T/d S nn² (1+
1 Mn
B. \.Mplri](u(()) Ir,
(58)
Next, we use u1() = 10 + +n )))) and w. = n= Vr, and (.,,(rslw:;)) = Vr((w") to get
T2² -E (w) w² VVr (w=) - gi. +0))), Vri, (rr,(w)) Di.
(1)
(8) g()))
=1 (w w", E. Vr. (w }}} + 1/² Es
9/.
(1) gl Vri, (w")
2
where we used that Eg glt), Next, we use || + y/ s 2 || + 2 |1yl1 2 for any vectors y,
and EIIX- - E[X]]| X|| 2 for any random vector X get
Th S -n (w w E. Vrs (w") }}} + 217 E, gh.
(0)
Vrig(w>)/
+ 227 | (rr
(59)
Plugging (59), (58), and (59) into (56) gives us
(1 + pm) Br ((++1) w/ w²12 s (0(() w0 + 2n7² Ee |0 vra,(e")/''
+ 227 1 Vrr((")))
1 Mn) Es
n22 11
Vria(())) -vr
(60)
Next, we note that P(o) C [0 onl" to get,
Es 199, - Vr., (w=))= Σ
(1) 19:11-v4 vrl(w)111 s Son, Σ 119 - vri(w*)1 and
1=1
2=1
E, VM,prri((u!)) - -Vra(w") }1²
= Σ
1:
( T.M,[ril(utt 10:00-ooon :(o? - (0)) Vr:(1o')
i=1
201 n
V.M,lrr((u 1796"" - 00) Vri(ww)) 1|²
49

---

Published as a conference paper at ICLR 2024
Moreover, also have that
|| (1=)
Ve(w>)] (1 - gTth
(e(w 1)1]
=
- (10() (0((((((" 1))
7. G2
w" ||² (0)
2
D2
Plugging these back into (60) completes the proof.
Next, we analyze the other two terms of the Lyapunov function. The proof is trivial, so we omit it. Proposition 19. We have,
I
((+1)
w²
(1 n
w/
+ -1 -
w/1
Yi Ex,
(r:(w=)) [
g
(++1)_
=(1 - n -1)
Σ
9i vri(w=))
==1
8
V.M (ro((o(()) -79/ -
vrs(w>)
We are now ready to prove Thm. 17.
Proof of Thm. 17. Let denote the effect of the smoothing. Combining Props. 18 and 19, we can write
be constant be determined later and let r 220 (MP2)
E. - (1 - -1-- 1I =-||u|0- - w'Il
0
(1+ pm
c1 ----
n
172(1 + (Mm)-1)
V.M,(rel(uut - -g())) Vri(w>)|
C2
01
1 um
no n
(1) w-1l (a((-11--1)
27727222
(1 u2)))2
(1) Vri((ot))((
Tiz (1 - (
2nj² on 1 1 um)
(
(61)
Let = 6/M. Our goal is to set the constants b, C2, > 0 so that the right side above is non- positive and as small as possible. We will require 2n so that -1 -1-111 (2n) -1 Thus, we can have the right side nonpositive with
61 -
20 0
(62a)
b b+ K
02
1 b K
b(b
+ 1) >
(62b)
no1
T
C1
2 0
(62c)
2n
1+ b/K
202.n
02
0.
(62d)
2n
+
b/16
50

---

Iterating (61) over Published as a conference paper at ICLR 2024
Let us set C1 = =T- By setting 02 Akno,, b² /(b (6+ k). ensure that (62d) is satisfied. Next, satisfy (62a) with
b 2T
2x b
O+ K
T -
Now, (62b) is an incquality only т. is satisfied with r 2 Tx : + 2x (4 - 1). This lets us fix max 2n, T throughout, which leads to the value of r) as claimed in the theorem statement. Finally, (62c) requires
4m
VTNYA
2 <(4xg - 1)
2
D
min
M
vn
Thus, under these conditions, the right-hand side of (61) is non-negative. updates, get
E[ -1 -1- exp((t/rr\(())
To complete the proof, we note that - 1/(2n) and
4
Khii b n:( 1) Krt
1
C2 b+K
This lets us use the fact that Vr; is M- M-Lipschitz to bound
( |1000 w²1
+ c1 Σ
(0) - w0 11
C2
Vr: (uo -Vr((w)) 12
+
M
(n+ 3/2))
(000) w
F IMPROVING THE DIRECT PODNTT BASELINE
In Sec. compared against the existing method of Palaniappan Bach (2016), which views our objective (8) in its min-max form directly and applies variance reduction techniques to both the primal and dual sequences. this method both theoretically and empirically by utilizing different learning rates for the primal and dual sequences. In this section, we provide a new convergence analysis for this improved two- hyperparameter variant, which dub SaddleSAGA, under the x² -2-vieergence penalty.
order to make the comparison more convincing, we also improve
F.I OVERVIEW
As in Appx. E, we begin with the additional notation and description of the algorithm, followed by the convergence analysis.
Notation. For simplicity, denote 2nn,, and consider directly the min-max problem
min max WEB' (EP())
(W q) ((w( + 2 u 114
1n//11
(63)
Note that the function V is strongly convex in its first argument and strongly concave in its second argument. pair (w*, q²") is called a saddle point the convex-concat function
max (w"", ;0)) S VWw", (g" gEP(a)
min (w, WERd
In our setting, we can verify that the pair = arg min F, and g" = (PT'(E(W))) is the unique saddle point of V.
51

---

Published as a conference paper at ICLR 2024
Algorithm Description. The algorithm makes of proximal operators (as described in Appx. in addition), which is defined for a convex function f : R - R, and æ E Rd as
1 | -v12.
proxy (z) arg min f(y) + yERa
2
The method is nearly equivalent to Algorithm 3, but applies the update
(t) n4., (1)
Vl, (w/ (np, 9/))) g(
w (++1) proxnull (w)( my(())
to the primal iterates and
π (1) nli, (w - -(n1 (0 Ci2 - -())
q (c+1) po +o8ll
8x
to the dual iterates, where 5 O is the dual learning rate. The vector x (()) plays the role of an update direction, and the proximal update on (((=1) can be solved with the PAV algorithm, in Appx. C. Overall, the time and space complexity identical to that Prospect.
Rate of Convergence. We prove the following rate of convergence for SaddleSAGA
Theorem 20, The iterates (wo(), (0) of SaddleSAGA with learning rates
= min 1 6(Lsa + 2C2n 8 min
Sn
converge linearly to the saddle point of (63). In particular, for non-trivial regularization μi S Gn² and (for some constant e) is at most
n/v,, the number of iterations to get |w(() -w" +lllgl))(()))
-113))='133+ -q"
11 <
o ((n+ KKg +
n2 n²2
1
nD
The proof of this statement is given as Cor. 26 later this section. compare to the original variant, the rate obtained by Palaniappan & Bach (2016) in terms of our problem's 's constants is
o ((~
nG² n
1
nu²
In
Compared to this, the rate we prove for SaddleSAGA improves x² KKg while suffering an ad- ditional factor of n in the n n² G² (ND) term. Empirical comparisons between SaddleSAGA and the original algorithm are given Appx. I. As compared to Prospect, SaddleSAGA matches the rate of Thm. I only when the shift cost / is large enough.
F.2 CONVERGENCE ANALYSIS
In the following, we denote by E, [+] the expecation of a quantity according to the randomness of ie conditioned on nn(t)g(() Throughout the proof, we consider that the losses are L-smooth and G-Lipschitz continuous.
Evolution of Destances-to-Optime operator to bound the evolution of the distances to the saddle point (w*,g").
We start by using the contraction properties of the proximal
52

---

Published as a conference paper at ICLR 2024
Lemma 21. We have,
E [w( w011 (1+ + nu)² (Iko()) - -1
2m (7(((( T ()(w()) V O)ww))
( (w ((9')))(²*)) wn)
+ B [10000 V(g e)(
E, [1(((11) -0'11] - (1
4¹1 6σ)² (19 (1+8
+ 25 ()( - [((w)) (9 (9)
+ 5 [| )(w
Proof. By considering the first-order optimality conditions of the problem one verifies that w", :9" satisfy for any 17, 6,
wn = prox n |1 11 (w" - 1/ ( ())
9"' = prox +-0)
13/2(9 + (e((w)))
Recall that the proximal operator of a c-strongly convex function h contractive such that (2) - prox, (2))x I +tallz 2' 112. In our case, it means that
[| w01 112 -
1 || mp((1) - (w- V/V e)(w") 122, )ll,. ()
+ 1111
[g 8-112 S
+ 5ū l1g + 8,((1) (9* + (l(w2))
By taking the squared norm, the expectation, expanding the squared norms and using that E, [0(0] = (()()) (w((t)) E, [x( = E(wl!),, we get the result.
Evolution of Variance Term. We consider the evolution of the additional variance term added to the dual variables.
Lemma 22. We have for any 92
E, [| - ((w(www))) >M3 ((w0))1
S (n + (n (2) w0 112
+ (n 1) ((11+331)1145 (3'))1(((=)) 10
Proof. As in the proof of Lem. 10, we have for some B2 > O, (1)()) (w)
E2 [[mm(0) - E(w >1
Re, [[nea
nni,(w") (w)")
+ (nli (w*) -nly (2 )es, (((w*) - ((0))113
S --B2lle(w - =((w))12
+ (1+ B. [[(nli - nei, (w (w))))))) )) 12
+ +1+ 1/ [1(ne, [1(ne,(w")) (w*) nl, (2 1)
- -(1+ (1+ - -200
(n (n (n-1))3))(((w - f(w")
+ (n (n-11((1=)))) + 35 331))/1/11011 ))(((())) - 100
< (n + (n - 10.-1)33)) (1 w" 112
+ (n 1)(1 + 35 )|[((=+) 119
53

---

Published as a conference paper at ICLR 2024
Incorporating Smoothness and Convexity of the Losses. The improved algorithm devel- oped here, for the purpose of a fair comparison to our own algorithm, differs from the original one from Palaniappan & Bach (2016) by Cor. 24 stemming from Lem. 23. We exploit the smoothness and convexity of the losses to get negative term -B, [llng, Vl,, (w(+) reg Vl,, (w)) used
to temper the variance of the primal updates at the price of an additional positive term l|q(0) q" 12.
The sum of both being positive we can dampen the effecc of the additional positive term [] q-
the price of getting less negative term Es [lnq, Vls, (w(+) nq, Vl, (w*) [12]
Lemma 23. For any q1, 92 P(o), W1, W2 RR we have,
(q) (((wu)))))) - [(w2)) - (V(g) (w)) V (w2)) [ (wi - w2)
1 (Bi-Unif(m) [1|m41,, (w,) - n4.,, V4, (w2)) 11 + || Ve(w2) - n92. v((w)))1)
2Lnomax max
G
+ Lo Lomax
112 - -92112
Proof. For any 9 E P(r) and any E Rd, we have by smoothness and convexity of q, for
9: >
qli(v)) 2 ((((ww)))))))))))))))))) + 9. VL(()) (v-
2Lg; 11. V(((w) - 9. VV,('))
(64)
1
( 2: (lw (w-w) w) 2Ln 2L m
ngg (w)(w) na (
(65)
Note that the second inequality is then true even qi = O. since in that case all terms are 0. Therefore, for any Qı, Q2 E P(r), and any w1, w2, we have
4. ((w2) 2 [T((w))> + v (q ()w) (T)))((n (w2 - w1) + 2Lo max
E Unifin] [[] Ve,(w1) no1, (w2(w))
Th ((w1) 2 q {(w)) + V (q ()2) (wi w2) + 2Ln Uniiff]
[[ngg,.VE(u)) 1292, Ve(wg) |1
Combining these inequalities, we get
(q1 42) (e(w)) E(w2)) + (V (w)) V(1z ()22) (wi - -22)
1
2Lnomax (EnwUnif(m)
For any 4 vectors a, b, c, d.
||a - b122 + |c d| -|o d|ž 2(a (b--
Applying this for a = 41. (w)),, b= n. (4)())) c= c= (12)), d = 42. (l((w, we get
(q) (2)T(((())) ((w2)) + (V(gT()(()) v(g30) T (w) -22)
2 2Tnommax Er-Unitii)) -
- V4 + ||q V((w2) ngz. V
Reorganizing the terms and bounding (4,(w1) (w,)w)) (V((w))T(((w) by G2 we get the result.
54

---

Published as a conference paper at ICLR 2024
Corollary 24. We have for any [0. 1]
(1 + mp)2
12 (1 + 1z)
[w (+++))===11
110 (et)) -/
"
$ 1) 'I =2'12 8-1 20G²
119 9" 11
Lamax
+ riR [1 V(q" " ()(w") )|6
+ ER [1 E(w )13
Et [Ing], (w(()) -nq;, Vl,, (w)))
Lno
Proof. Follows from Lem. 23
Lyapunov Function and Overall Convergence. Thm. 25 shows that appropriately defined Lyapunov function incorporating the distances to the optima, decrease exponentially.
Theorem 25. Define the Lyapunov function
v(t) (1 + ?
(1+ 8D)
lw w=
11 9 1=112
Inp (1) (+((2(()
-ng ng Vli 2 +
G |1 f(ww')())
1=1
with C1 = 2(Eme n/D and C2 = with No no max By taking
1
min
1 6(Lris + 222
8 = min
v 8n
we have
Eg [vatt]] $11---)
for some 1. In particular, for small regularizations, i.e., u S Sn² and μ 6(Liio + 2G²n/D), DV, we have
24LKa
48G² n2
16G
7 == max
2n, 4+
2
1
Dp
Proof. Let denote
T
1mmp Ve((2(0) (10 - n4: (e s(t) = /1(00 - E(w=)113
we have,
E, (
1
1nq (1) (l((1()) - nn. Vel >11 + (1- TTt).
<
n²
E, [((+1)
< (2) - 20=13 113 + ( (1- n
1
5(D.
55

---

Published as a conference paper at ICLR 2024
By combining Cor. 24, Lem. 10, Lem. 22 we have, denoting Ri manax
E, [
< (n 8(n + (n -
+c2) (100()) -w"1
2anG²²
8-1 +
119
-0012
Like
(11 (1) +
c1
Ep-Unit(v)
[Ing; Ves, (w)(t) -ngi, Vli, (10=))
n
Indmmax
+ n(1 + 8T + t
(1- -
1
!mpp (2) [l)(2()) mqq (e,(
1=1
6(n () + 1 +
C2
1 (1- (1-)) [1((w0)) 10.0
G
Therefore for some T > 1, we have
E, [vitt]]] - (1 - ()(()) S Ki /|/l(()) - w011 + K2lll - 1-1
-
[|nng, Vll,, (v)(0)) - nq.; Vl,, [1)
+ K3
K₁ 1 n
Imp V - noz V(( 012 kolle(w") - 11
71
with,
K₃ (1 + nu)² m1+
1 17 ((n + (n (n-11) (133))3) 02) - (1 -(1--)))
(1+nu)²
K2- (1+ +0)z
1 1+ 24
(1 - -1-
8
(1060)2 C a
K3 (1 + (1)) +
Lto
1 - - (1 -11--1)
K4 = CI
(1 + 8T
+
1- n1
Ci
02
G² + (1
(1_ 1
- (1 (1---1)
K5 =
8(n - 1)(1 + B5
G2
n
Fix 31 = 2, = 1. Denote also π
I+ma E (0, 1) and 8
50
E (0, with e.g. 17 "((+0))
1+50
We have then for c = /// (2
and C2 = / 1/2,
2 G28
Ki S m1²1 112
n) + r -1
nGnn n Ltia
K2 S 6728
1 8- 2
8+7-1
a
K3 3n
2L.io
Ltig
K4 CI 37
+ 7-1
na
n
22 K5 < T
8
4nG²
+
7-1
1
n
56

---

Published as a conference paper at ICLR 2024
We can further take 3n - a2(2)))) and ò I (/(8nn By imposing the constraint 7 2 2n, can simplify
3
K1 S mu
1) -
7 +--
4
8 8- 2 (1
aG²n
K2 S 5628
LRgNÜ
K3 - 0, K1 $ 0, K; I 0.
Recall that a must be chosen in [0, 1]. Taking then
LKo L 2G"n/D
LKoN
2G'
we get
K1 S np²i
172
7 1"+ 1 K2 S 50 (82 8+ T
3
By taking < 1/1, ô S 1/D, we get π s 1/2, S 1/2 and so - - - - and 0² - 5 - Therefore taking
1
8 min we get K2 < for all as long as t 2 max 4/7, 2/0). In our case,
n min
j 6(Lii + 2G2n/DD
Bn²
+ 6Lmn + 120*n)
if u 6(Lka + 2G2n/D), otherwise,
8
2
1 +
8G D2
if uD S Gn otherwise.
The result follows.
Corollary 26. Under the setting of Thm. 25, after iterations of SaddleSAGA, we have
(1 + n1)a
(1 + 2
E
[|w!()) -=11
1
5
(1 + nuup
(1 + 50)2
< exp(-t/7)
w w0112
11
92'12
7
5
+ n
||q (0)
(l:(w( (0)) q. VVi("" 112
2 (le(w(0)) [((w") =(w==)
G TECHNICAL RESULTS FROM CONVEX ANALYSIS
Herein, we collect several results, mostly from Nesterov (2018), used throughout the paper. In the following, let |-||| denote an arbitrary norm on Rd and let Il- denote its associated dual norm.
The first concerns L-smooth function, or those with L-Lipschitz continuous gradient.
Theorem 27. (Nesterov, 2018, Theorem 2.1.5) The conditions below considered for any x, Rd and a E [0, 1). The following are equivalent for a differentiable function f : Rd. - R.
1. f is convex and L-smooth with respect to ||-I|-
2. 0 fly) f(x) (Vf()), x) < & || y/1
3. f(x) (Vf(x), =)) + + [ v/(2)1 f(v).
57

---

Published as a conference paper at ICLR 2024
4. 1 VVf()) V (Vf(²) v
5. 0 < (Vf(z) Vf(w) x y) S y1
Next, we detail the properties of strongly convex functions.
Theorem 28. (Nesterov, 2018, Theorem 2. 10) If f IR tiable, then for any I, Rd.
is p-sttrongly convex and differen-
((u) S f(x) (f(z). y- + [V - V
(V/(z))
S [V/(z) (V/(c)-V/(O)) v
#1| - [Vf(=) Vf(i)l..
Finally, functions that are both smooth and strongly convex enjoy number of relevant primal-dual properties.
Theorem 29. (Nesterov, 2018, Theorem 2.1.12 If f is both L-smooth and p-strongly convex, then for any I, E R4
(Vi(z),
u+E || yl1 a+L - ((()))-f((i1-- (V)())====)==-00) y). (66)
Lemma 30. Let f Rdd + R be p-strongly convex and M-smooth. Then, we have for any w. v E Rd
f(v) 2 f(w) + V[(w)) (v ww) + 2(M + [V[(w) V f e112.
Proof. The function g = pll 112/2 is convex and M / smooth. Hence, we have by line 3 of Thm. 27 for any w, v E IRd,
g(o) 2 g(w) + Vg(w) w)+ 1 2(M 1)
V Vg(w))33
Expanding g and Vg. we get
f(v) 2 f(w) + (.(w) (w-w) 2(M - 11) [Vf(w) f
2(M PM
1)) ||0 (ww-vl v12
M (V](w) vf(v)) (w v).
p
Using Young's inequality, that is, a b < ||||| + a- |bbl2. we have
f(v) f(w) ()( T w)
ap
[Vf(w) -
2(M p)
M(M 21 1| 0112. 2(M p)
Taking a #+M gives the claim.
We state a few properties of the prox operator.
Theorem 31 (Co-coercivity of the prox). If Rd E. is u-strongly convex, then we have for any constant y) >> 0 that
(x - y. v.proxx,y(a (x) prox,, (v)) (1 m71) [/rrox,)(x) prox, (w)]1
The same result applied to the convex conjugate f"r of f and noting that v.M[f((x)) = prox, (n((//))))))))))///)))/))()))( gives the following result:
58

---

Published as a conference paper at ICLR 2024
Dataset
d
Thrain
Thest
Task
Source
yacht
6
244
62
Regression
UCI
energy
614
154
Regression
UCI
concrete
8
824
206
Regression
UCI
kin8nm
8
6,553
1,639
Regression
OpenML
power
7,654
1,914
Regression
UCI
diabetes
33
4,000
1,000
Binary Classification
Fairleam
acsincome
202
4,000
1,000
Regression
Fairlear
amazon 535
10,000
10,000
Multiclass Classification
WILDS
iwildcam
9420
20,000
5,000
Multiclass Classification
WILDS
Table 2: Dataset attributes and dimensionality d, train sample size Tumins and test sample size Thest-
Theorem 32 (Co-coercivity of the prox). If Rdd + IR is L-smooth, then we have for any constant
?) > that
(x y, V.M,f(a)) V.M,f(ly)) 2n 1+ 1n) |V.M,1f((() -MM,))
Lemma 33. For convex function f: R ( (+ if x1 2 x2 and 1/2 2 yr, then
f1 x11++ + (yy x2) 2 (y2 x1) f x2).
Proof. First, observe that
1/2 12 x2
y1
I and
// 22 2 /1 21
Thus, 3/2 1 and both lie between - and /1 Iı and can be expressed as a convex combination of the two endpoints, that is
1/2 I1 ayy x2) + (1 a) =1) V1 x2 = B(yz - x2) + +(1 - 8) - 21)
for some a, B E (0, 1). By solving for we get a 1 3. Apply the definition of convexity to get
[y2 x1) af x2) (1 a) f(y 21) f(v1 - x2) - (1
x2) + of xz
Sum both inequalities to achieve the desired result.
Lemma 34. Define for E g".
1
h(I) = max IT q- 11n/n1/
gEP(a)
The function h is 1/U-smooth and convex such that for any 1, E ER"
v|(v((() - V (V^()) - VA(()) -1) - 1
-|1 -1'11. -1
H EXPERIMENTAL DETAILS
H.1 TASKS & OBJECTIVES
In all settings, we consider supervised learning tasks specified by losses of the form
((w)) h(vi, w (xx:)))
59

---

Published as a conference paper at ICLR 2024
where we consider input z; X. feature map 4 X RR, and label y, y. The function h: y x R - measures the eπror between the true label and another value which is the prediction in regression and the logit probabilities of the associated classes in classification. In the regression tasks, = R2 and we used the squared loss
1 -9: ((x)))² ((w)) 2
For binary classification, we have {-11 1}. denoting negative and positive class. We used the binary logistic loss
((w) - yo _ In(1 + e ")
For multiclass classification, multinomial logistic loss:
{I,...,Cy C} where C' is the number of classes. We used the
(.. (W.T.r.)
l((w) In Pay (3045 w), where Py, (Ir4; w)
exp
W E Baxe
Ey= exp
w (.
The design matrix unit variance, and the estimated mean and variance from the training set used to standardize the test sets as well. Our final objectives are of the form
E Ig is standardized to have columns with zero mean and
(w = max gli(w) - on 119 1n/11/2 + 1 || (E)(()
2
for shift cost 2 and regularization constant 2 O.
H.2 DATASETS
We detail the datasets used in the experiments. If not specified below, the input space X = IR and sp the identity map. The sample sizes, dimensions, and source of the datasets are summarized in Tab. 2, where refers to the dimension of each sor;))
(a) yacht: prediction of the residuary resistance of a sailing yacht based on its physical attributes Tsanas & Xifara (2012).
(b) energy: prediction of the cooling load of a building based on its physical attributes Ba- ressi Segota et al. (2020).
(c) concrete: prediction of the compressive strength of a concrete type based on its physical and chemical attributes Yeh (2006).
(d) kin8nm: prediction of the distance of an 8 link all-revolute robot arm to a spatial endpoint (Akujuobi & Zhang, 2017).
(e) power: prediction of net hourly electrical energy output of a power plant given environmental
factors (Tüfekei, 2014).
(f) di abetes: prediction of readmission for diabetes patients based on 10 years worth of clinical care data at 130 US hospitals (Rizvi et al., 2014).
acsincome: prediction of income of US adults given features compiled from the American Community Survey (ACS) Public Use Microdata Sample (PUMS) (Ding et al., 2021).
(h) amazon: prediction of the review score of a sentence taken from Amazon products. Each input x E is a sentence in natural language and the feature map ((x) E IR generated by the following steps:
A BERT neural network Devlin al. (2019) (fine-tuned 10, 000 held-out examples) is applied to the text I;, resulting in vector z/--
The 2/......y o', are normalized to have unit norm.
Principle Components Analysis (PCA) applied, resulting 105 components that explain 99% of the variance, resulting in vectors x." 1205 The d in Tab. 2 refers to the total dimension the parameter vectors for all classes.
ii) 1wildcam: prediction of an animal or flora in an image from wilderness camera traps, with heterogeneity in illumination, camera angle, background, vegetation, color, and relative animal frequencies Beery et al. (2020). Each input x E X is an image the feature map (px)) E Rd is generated by the following steps:
60

---

Published as a conference paper at ICLR 2024
A ResNet50 neural network He et al. (2016) that is pretrained on ImageNet Deng et al. (2009) applied to the image 24, resulting in vector z/-
The x!,...,,,, I', are normalized to have unit norm.
Principle Components Analysis (PCA) is applied, resulting in d 157 components that explain 99% of the variance. The Tab, refers to the total dimension of the parameter vectors for all 60 classes.
H.3 HYPERPARAMETRR SELECTION
We fix a minibatch size of 64 SGD and SRDA and an epoch length of N for LSVRG. For SaddleSAGA we consider three schemes for selecting the primal and dual learning rates that reduce to searching for single parameter 7 > O, described in Appx. I. In practice, the regularization parameter l and shift cost v are tuned by a statistical metric, i.e. generalization error as measured on validation set. We study the optimization performance of the methods for multiple values of each in Appx. I.
the following method. Let k (1,..., K} be seed that determines algorithmic randomness. This corresponds to sampling a minibatch without replacement for SGD and SRDA and a single sampled index for SaddleSAGA, LSVRG, and Prospect. Letting (1) denote the average value of the training loss of the last ten passes using learning rate 1/ and seed k, the quantity [(n)
For the tuned hyperparameters,
t Σ (x) was minimized to select 1). The learning rate 1) is chosen in the set {1 x 10 10 x 10-1
10-3,3 x 10-3, 10-2 3 x 10-2, 1 10-1, x 10 x 100, 3 x 10 with two orders of magnitude lower numbers used in acsincome due to its sparsity. We discard any learning rates that cause the optimizer diverge for any seed.
H.4 COMPUTE ENVIRONMENT
No GPUs used in the study; Experiments
CPU workstation with Intel i9 processor, a clock speed of 2.80GHz, 32 virtual cores, and 126G of memory. The code used in this project was written in Python 3 using the PyTorch and Numba packages for automatic differentiation and just- post-in-tiin in-time compilation, respectively.
I ADDITIONAL EXPERIMENTS
Varying Risk Parameters. We study the effect of varying the risk parameters, that is (p, b, for the p-CVaR (Equation (21)), b-extremile (Equation (22)), 7-ESRM (Equation (23)), choosing the spectrum to increase the condition number Ka no, compared to the experiments in the main text. We use p 0.25, b 2.5, and 7 1/c-2 to generate "hard" version of the superquantile, extremile, and ESRM. Fig. 8 plots the corresponding training curves for four datasets of varying sample sizes: yacht, energy, concrete, and iwildcess methods is the same the original methods, that is that Prospect performs the best close best in terms of optimization trajectories. Except on concrete, SaddleS AGA generally matches the performance of Prospect. The trajectory of LSVRG is noticeably noisier than on the original settings; we hypothesize that the bias accrued by this epoch-based algorithm is exacerbated by the skewness the spectrum, as mentioned in Mehta et (2023, Proposition
We see that the comparison of
Lowering or Removing Shift Cost. A relevant setting is the low or shift cost regime, as this allows the adversary make arbitrary distribution shifts (while still constrained to P(a)). These settings correspond to 10-3 and 0, respectively. The low-cost experiment is displayed in Fig. 9 while Fig. 10 displays these curves for the no-cost experiment. When v the optimization problem can equivalently be written as
min WERd
max [((w)) + 1 || 2 gEP(G)
Σ
Co (w) +
1 /wll2
In this case, always have that (I) (a (o--1(1);- (n)) where sorts l. Here, w is chosen to optimize a linear combination order statistics of the losses. the low shift cost settings, performance trends are qualitatively similar to those seen from 1. Interestingly, for the no-cost.
61

---

Published as a conference paper at ICLR 2024
yacht
energy
concrete
ivildcam
10P
CVaR (Hard)
10-1
19'1
1074
10"4
10"44
10-7
10-7
Extremile (Hard)
10
10:1
100
10f
10-5
1074
10--1
10
10-1
10-7
10m7
10f
ESRM (Hard)
10-1
10
10-4
10'14
10-11
10-1
10-4
10111
10
32
SGD
SRDA
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 8: Harder risk parameter settings. Each row represents a different "hard" variant of the superquantile, extremile, and ESRM spectra. Columns represent different datasets. Suboptimality (9) is measured on the y-axis while the æ-axis measures the total number of gradient evaluations made divided by n, i.e. the number of passes through the training
yacht
energy
concrete
krn8Sn
power
1vi1dcam
10
CVaR
10-0
38
1074
3274
1072
Extrenile
10-
30-1
10
33.7
10-7
13.
00
ESRM
10
Tiams
Ph00
SGD
SRDA
LSVRG
Sadeli-AAA
Prospect (Ours)
Figure 9: Low shift cost settings. Each row represents different spectral risk objective with y 10-3 (instead of === 1) while each column represents a different datasets. Suboptimality (9) measured the y-axis while the z-axis measures the total number of gradient evaluations made divided by n, i.c. the number of passes through the training set.
setting, LSVRG, SaddleAAGA, and Prospect seem to converge linearly empirically even without smoothness of the objective.
Lowering Regularization. Next, we decrease the l-rregularization from μ = 1/n p 1/(10m) and μ 1/(100m). These settings are plotted in Fig. 11 and Fig. 12, respectively. Performance rankings among methods reflect those the original parameters. For five the six datasets, that is yacht, energy, concrete, kin8nm, and power, the regression tasks involve optimizing the squared error. This function is already strongly convex, with a constant depending on the smallest eigenvalue of the empirical second-moment matrix. When assuming that the input data vectors
62

---

Published as a conference paper at ICLR 2024
yacht
amargy
Srus
1wildcam
1077
1071
CVRR
32-
1077
30"1
1011
60
2274
10-1
100
Extremile
1011
2277
10--1
30"1
1077
2275
1074
10-1
ESRM
10
10"
10
10-7
22
10
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 10: No shift cost settings. Each row represents a different spectral risk objective with v/ = (instead of 1) while each column represents different datasets. Suboptimality (9) is measured on the y-axis while the z-axis measures the total number of gradient evaluations made divided by n, Le. the number of passes through the training set.
yacht
energy
kin88m
power
iwildcam
12
10
10
10'7
CVaR
101
12
10"
10
3274
107
10"1
Extemile
00
10
10*
3277
ESRM
16
22
20
10
10
24
15
LSVRG
SaddkSAGGA
Prospect (Ours)
Figure 11: Reduced l--reulaarization settings (p = 1/(10n). Each row represents different spectral risk objective with μ 1/(10n) (instead of μ 1/n) while each column represents different dataset. Suboptimality (9) is measured on the y-axis while the z-axis measures the total number of gradient evaluations made divided by n, i.e. the number of passes through the training set.
bounded, this function is also G-Lipschitz. Thus, if the problem is already well-conditioned, may observe similar behavior even at negligible regularization (/4 = 5 10-7 for iwi 1dcam, for example).
Comparison of Saddle-Point and Moreau Variants. Finally, observe in Fig. 13 the comparison of SaddleSAGA variants (Appx. F), well the Moreau version of Prospect using Moreau envelope- based oracles (Appx. F). There are variants shown.
Primal LR Dual LR: The original variant of Palaniappan Bach (2016), which the primal and dual learning rates are set to be equal and searched as a single hyperparameter.
Search Dual LR: Here, the primal learning rate is fixed as the optimal one for Prospect, and the dual leamnnng rate is searched as single hyperparameter.
63

---

Published as a conference paper at ICLR 2024
yacht
smargy
In
1wildcam
1374
1077
1071
1074
CVoR
10
1274
16"1
32'5
10-7
10
107
Extremile
10
1017
33"4
2277
20177
ESRM
10-1
10-7
15
22
24 10
SRDA
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 12: Low fz-reuularization settings (u = 1/(00)). Each row represents a different spectral risk objective with 1100m) (instead of μ 1/n) while each column represents different dataset. Suboptimality (9) is measured on the y-axis while the z-axis measures the total number of gradient evaluations made divided by n, i.e. the number of passes through the training set.
Primal-Dual Heuristic: In this version, used as the "SaddleSAGa baseline in the main text, the dual learning rate set to be 10n times smaller than the primal learning rate.
Prospect-Moreau: The Moreau-envelope version of Prospect using proximal oracles.
We find that all methods besides the original variant (primal LR = dual LR) perform comparably on yacht, energy, concrete, kin8nm, and power. Notably, the ProXSAGA method performs similarly to Prospect and the saddle point-based baselines. While using the Moreau envelope results in accelerated rates in the ERM setting Defazio (2016), we find that the convergence rate is the same empirically. This phenomenon is in agreement with Thm. 17, which states that ProXSAGA will achieve the same linear convergence rate as Prospect, but will require a much less stringent condition on the shift cost than in the case of Prospect.
Mcasuring Wall Time. In Fig. 14, we measure the suboptimality for each method as a function of the wall time taken for training on the compute environment described in Appx. H. Across tasks, SGD is approximately 2 orders of magnitude faster than Prospect (and SaddleSAGA) but fails converge due bias/variance. In lowtt--moerrate sample size settings, LSVRG achieves a similar suboptimality to Prospect in a 2-3x smaller wall time. However, in the large sample size setting, LSVRG achieves essentially the same suboptimality as SGD compared for Prospect, indicating that while fast, LSVRG does not demonstrate adequate accuracy. Also, at the large sample size setting, Prospect performs about 3x faster than Sad but achieves the same suboptimality. Finally, wall time cost that is not represented by the attached Figure 2 is the hyperparameter tuning cost. Both LSVRG and SaddleSAGA involve a grid search for tuning two hyperparameters, whereas Prospect has single hyperparameter that can be tuned with binary search.
Comparison to Seneral-Purposs Learning Algorithms. In Fig. 15, we compare the Prospect and baselines against optimization algorithms with widespread usage in generic learning problems, namely stochastic gradient descent (SGD) with momentum (Polyak, 1964; Nesterov, 2018) and Adam (Kingma & Ba, 2015). We use the default hyperparameters such momentum 0.9 for SGD and (31, 32) (0.9, 0.999) for Adam. For both methods, we use a batch size of 64 and tune the learning rate. Because of the bias and variance issues identified in Sec. 2, Adam and SGD with momentum perform comparably to standard SGD and are unable to converge on the objectives considered. Across objectives/datasets, final suboptimality often caps 10 for both methods.
64

---

Published as a conference paper at ICLR 2024
yacht
energy
kin8nm
10-
101
16-7
CV&R
1174
117
-077
10
10"
1077
Extremile
10"
30-4
23m
1071
117
30m4
ESRM
10m11
30-1
13
Panss
Pans
Search Dual LR
Primal LR
Dual
Primal-Dual Houristic
Prospect-Moreau
Prospect (Ours)
Figure 13: SaddleSAGA and Prospect-Moreau method comparisons. Each row represents a different spectral risk objective while each column represents different dataset. Suboptimality (9) is measured on the y-axis while the æ-axis measures the total number of gradient evaluations made divided by n, i.c. the number of passes through the training set.
yacht
energy
power
iwildcar
10
1023
10
CVaR
10-44
10a2
10-7
16/
20--
10:
Extremile
10-
10"4
10-44
30m4
10-77
10-22
107
1177
107
10P
10m3
10"1
ESRM
20-4
10
10
10"
20
10"
20
10
11f
102
10 VHll Time
WSl Time (6)
Well mi
Well Time (4)
Vell Time (4)
SGD
LSVRG
SaddleSAGA
Prospect (Ours)
Figure 14: Wall time comparisons. Each row represents a different spectral risk objective while each column represents a different dataset. Suboptimality (9) is measured on the y-axis while the r-axis measures the total wall time in seconds on the environment described in Appx. H.
65

---

Published as a conference paper at ICLR 2024
yacht
energy
concrete
kin8nm
power
CVR
10:
10"4
10"9
19"9
10
10-7
13-9
10'
10-1
10-1
19
10
Extremile
20-2
13a4
10-1
10--1
10-1
10.0
1374
10"
ESRM
13:4
13-4
10m7
13-7
15
45
64
22
-6
45.
64
16
Passm
Pasms
SGD
SGD
Momentum
Adam
LSVRG
Sadd
Prospect (Ours)
Figure 15: Comparison to SGD w/ Momentum and Adam. Each row represents a different spec- tral risk objective while each column represents different dataset. Suboptimality (9) is measured on the y-axis while the z-axis measures the total number of gradient evaluations made divided by n, i.c the number of passes through the training set.
66


Reviews:
[{"review_id": "SlgNZIECHC", "reviewer": "Reviewer_12Ff", "summary": "The authors propose Prospect, a distributionally robust algorithm that requires the tuning of only one hyperparameter. The algorithm's formulation includes a reweighed empirical risk and an f-divergence term to account for the cost requires to shift from a uniform representation of the training data. The authors propose both bias and variance reduction procedures to ensure that both quantities vanish, hence guaranteeing convergence.  Prospect is proved to enjoy a linear convergence for any positive shift cost on regularized convex losses and is empirically competitive with the considered baselines.", "questions": "Please refer to the weaknesses section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper has multiple strengths in terms of originality, quality and significance. The clarity component is lacking in the ways that I will explain in the weaknesses section. \n\n- The paper has solid contributions compared to prior work. \n- The theoretical analysis is sound and well supported. \n- The authors also discuss the case where the hypotheses of this work are violated and argue that their algorithm still converges in that case. \n- The empirical evaluation considers 3 important problems and covers both classification and regression. \n- The proposed algorithm and all baselines are trained to convergence as shown in all figures. The hyperparameter selection seems to also have been done in a fair manner.", "weaknesses": "I enumerate below the weaknesses of this work, which to me are important to address but do not undermine the overall quality of this work. I hope the authors will be able to address them during the rebuttal.\n\n- Presentation and clarity: Although the authors clearly attempt to make the paper as clear as possible, some key notions are never introduced. For instance, CVaR was never formally introduced. It is also unclear to me what 0.5-CVaR, 2-extremile, and 1-ESRM really mean mathematically speaking. The entire notion of spectral risk-based needs to be explained much more clearly. I also don't understand Figure 6, and overall the captions of the figures could be more descriptive. The abstract is heavily jargony, and so is a lot of the introduction text. I would appreciate a clearer presentation of this work, its motivation, and key notions at least in the introduction. \n\n- Baselines: I find the baselines considered by this work to be very restrictive. Some of the problems considered in the experiments section violate the hypotheses that lead to Prospect's theoretical guarantees, and therefore it would only be fair to compare it to other algorithms that do not enjoy the same guarantees and do not have only one hyperparameter to tune. In fact, the number of hyperparameters being only one is not as big of an advantage as portrayed in my opinion, since more practical algorithms would have default values for their hyperparameters that would work well in practice. I recommend comparing additionally to at least Adam and SGD with momentum.", "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": "", "extra_scores": {}}, {"review_id": "7X37vH5PRa", "reviewer": "Reviewer_ucUf", "summary": "The authors studied DRO with spectral risk measure and modeled the ambiguity set as $f$-divergence. Strong duality result, together with a stochastic optimization algorithm, was proposed to optimize the formulation. The idea of their proposed algorithm is to reduce the bias and variance of gradient estimators while maintaining small computational costs when performing iteration updates. Convergence analysis revealed that their algorithm converges linearly for \"smooth regularized losses\". A comprehensive numerical study is performed to show the superior performance of their algorithm.", "questions": "- In contrast to the standard DRO that models $\\mathcal{P}$ as a probability simplex on $n$ atoms, the authors consider the spectral risk measure such that $\\mathcal{P}=\\text{CovexHull}(\\text{Permutaions of }(\\sigma_1,\\ldots,\\sigma_n))$. I was confused about the motivation for using such a spectral measure. In which scenarios will people focus on this type of measure? The authors should give more justification on this part.\n- When deriving the strong duality result in Proposition 3, the authors assume that the conjugate function of $f$, denoted as $f^*$ satisfies $|f^*(y)|<\\infty, \\forall y\\in\\mathbb{R}$. However, when studying the standard $f$-divergence DRO with $\\mathcal{P}$ being the probability simplex on $n$ atoms (see Section 3.2 in (Shapiro Alex, 2017)), one does not require such an Assumption. Could the authors justify why we need this restrictive assumption? Also, I think the authors missed this important citation.\n\nRef: Shapiro A. Distributionally robust stochastic programming[J]. SIAM Journal on Optimization, 2017, 27(4): 2258-2275.\n- In Proposition 4, the assumption that $f$ is $\\alpha_n$-strongly convex may be restrictive. Is it possible to relax this assumption into strictly convex?\n- It is difficult to tell why Proposition 5 holds. The authors should add complete proof regarding this proposition. Besides, in Eq. (15) I find the authors implicitly assume the conjugate function $f^*$ is differentiable. I am wondering if the authors assume the differentiability of the function $f$, or the condition $f^*$ is differentiable can be derived based on some conditions?\n- In page 28-32, page 35-37, and page 39-40, some equations are highlighted in color. What is the meaning of those highlighted colors?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- It seems the strong duality of $f$-divergence DRO when considering spectral risk measure in Proposition 3 is new in the literature. Appendix B shows a range of nice properties regarding the formulation (2).\n-  The designed Prospect Algorithm is novel and operates by reducing the bias and variance of gradient estimators while maintaining small computational costs. Nice convergence guarantees are established in Theorem 1 provided that the regularization for $f$-divergence is lower bounded and the regularization value of the model parameter $w$ is positive.\n- A comprehensive numerical study is performed to show the superior performance of their algorithm.", "weaknesses": "See Questions part.", "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": "", "extra_scores": {}}, {"review_id": "tOQ10JJ9xH", "reviewer": "Reviewer_ss6k", "summary": "In this paper the authors have developed an iterative algorithm to solve optimization problems with spectral risk measures. \nThis is equivalent to optimization a distributionally robust optimization problem where the ambiguity set is the space of all probability distributions with a term to penalize the distance from the uniform distribution.\nTheir algorithm consists of 3 key components\n1. Gradients steps of the spectral risk measure\n2. Bias reduction for the gradient estimates\n3. Variance reduction for the gradient estimates\n\nEach of these steps can be done efficiently. \nThe authors also prove that the algorithm converges to the correct solution linearly\nFinally, the authors illustrate the performance of their algorithm through numerical experiments on a variety of tasks.", "questions": "1. What is the scale of the problems you can solve. How much time does it take to solve the problem. \n2. Will the algorithm work if used along with a projection step to solve constrained optimization problems?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "**originality**: I think this work is original. This paper develops a new iterative algorithm for DRO problems with ambiguity sets connected to spectral risk measures. \n\n**quality**: The work is sound and presents justifies the algorithm with both theoretical and experimental results\n\n**clarity**: the work is well presented and the numerical results are clearly explained.\n\n**significance**: I believe the work is significant since it expands the type of DRO problems that can be solved with iterative algorithms and hence the size of DRO problems that can be solved which has always been a key limitation of it.", "weaknesses": "I feel the discussion on the connection to DRO is quite limited. The ambiguity sets for equation (2) is the entire space of distributions which is quite large and not very useful. I believe the presence of the penalty term shows that this problem can be equivalent to a tighter ambiguity set (maybe restricted by the divergence metric used in the objective) and it would be good if the authors can discuss this. \n\nf-divergences are quite a broad type of divergence as discussed in the appendix. The 3 types of divergences considered in the experimental section are quite limited. It would be good if authors can also discuss some of the other popular divergences such as KL-divergence etc. and compare the developed algorithm to existing approaches for DRO problems with these divergences.", "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": "", "extra_scores": {}}, {"review_id": "x0lC1jsG0I", "reviewer": "Reviewer_Ed47", "summary": "The paper provides a practical algorithm to solve DRO problems. The algorithm enjoys a linear convergence rate and it does not require tuning multiple hyperparameters. The author further demonstrated the proposed algorithm can achieve 2-3x speedup compared to other methods vias read-data sets.", "questions": "The papers (including the appendix) are long. I am not sure whether it is suitable for a conference review.\n\nI am curious that whether there are some unbiased estimators could be used, e.g., Wang et al (2023).\n\nWang, Shengbo, et al. \"A finite sample complexity bound for distributionally robust Q-learning.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2023.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and the contribution is significant. The algorithm is easy to understand and easy to implement. The authors also provide intuition behind the algorithm via texts and figures.", "weaknesses": "1. The paper needs to solve an optimization problem with n variables in each iteration (line 10 in Algorithm 1). It is very computationally expensive if n is large. I wonder whether the algorithm can reduce to another one-step stochastic gradient descent in this step. For example, the two-time scale method in distributionally robust RL [1].\n\n[1] Yang, Wenhao, et al. \"Avoiding model estimation in robust Markov decision processes with a generative model.\" arXiv preprint arXiv:2302.01248 (2023).\n\n2. The paper only provides theoretical results for the case $\\mu>0$. I am not sure if it is necessary or if it is only for theoretical convenience. It is known that the DRO objective $\\mathcal{R}_\\sigma$ could induce smoothness. Some discussions are needed and numerical examples with $\\mu=0$ will also help.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)