Min-Max Multi-objective Bilevel Optimization with Applications in Robust Machine Learning
Authors: Alex Gu
Songtao Lu
Parikshit Ram
Tsui-Wei Weng
Venue: ICLR 2023 poster (2023)
Keywords: robust optimization
bilevel optimization
multi-objective optimization

Abstract: We consider a generic min-max multi-objective bilevel optimization problem with applications in robust machine learning such as representation learning and hyperparameter optimization. We design MORBiT, a novel single-loop gradient descent-ascent bilevel optimization algorithm, to solve the generic problem and present a novel analysis showing that MORBiT converges to the first-order stationary point at a rate of $\widetilde{\mathcal{O}}(n^{1/2} K^{-2/5})$ for a class of weakly convex problems with $n$ objectives upon $K$ iterations of the algorithm. Our analysis utilizes novel results to handle the non-smooth min-max multi-objective setup and to obtain a sublinear dependence in the number of objectives $n$. Experimental results on robust representation learning and robust hyperparameter optimization showcase (i) the advantages of considering the min-max multi-objective setup, and (ii) convergence properties of the proposed \morbit.

Published as a conference paper at ICLR 2023
MIN-MAX MULTI-OBIECTIVE BILEVEL OPTIMIZATION WITH APPLICATIONS IN ROBUST MACHINE LEARNING
Alex Gu', Songtao Lu², Parikshit Ram², Tsui-Wei Weng"
"MIT CSAIL, 'IBM Research, *UCSD
gualmit.edu, *(songtao, parikshit. ramj mimn.iibm.com com, "Iweng@ucsd. edu
ABSTRACT
We consider a generic min-max multi-objective bilevel optimization problem with applications in robust machine learning such representation learning and hyper- parameter optimization. We design MORBiT, T, novel singlc-loop gradient descent- ascent bilevel optimization algorithm, to solve the generic problem and present a novel analysis showing that MORBiT converges to the first-order stationary point at rate of On K K= ) for class of weakly convex problems with objectives upon iterations of the algorithm. Our analysis utilizes novel results to handle the non-smooth min-max multi-objective setup and to obtain a sublinear dependence in the number of objectives n. Experimental results on robust representation learning and robust hyperparameter optimization showease (i) the advantages of considering the min-max multi-objective setup, and (ii) convergence properties of the proposed MORBiT. Our code is at https: / //gi thub. cwm/minniiiiccccccc.cc inniiimmaaatt ric M
INTRODUCTION
We begin by examining the classic bilevel optimization (BLO) problem as follows:
min f(a, y (x)) subject to y (x) E arg min g(z, y)
(I)
DE dd
yEY-R">
where f x R is the upper-level (UL) objective function and g X' x 3 + R is the lower-level (LL) objective function. X' and V respectively, denote the domains for the UL and LL optimization variables a and y, incorporating any respective constraints. Equation 1 is called BLO because the UL objective f depends on both x and the solution y (a:) of the LL objective BLO well-studied the optimization literature (Bard, 2013; Dempe, 2002). Recently, stochastic BLO has found various applications in machine learning (Liu et al., 2021; Chen et al., 2022a), such as hyperparameter optimization (Franceschi al., 2018), reinforcement learning or RI. (Hong et al., 2020), multi-task representation learning (Arora et al., 2020), model compression (Zhang et al., 2022), adversarial attack generation (Zhao et al., 2022) and invariant risk minimization (Zhang et al., 2023).
In this work, focus robust generalization of equation to the multi-objective setting, where there are different objective function pairs (Ji, g:). Let [n]] A (1, 2....... .n. and fi X' } + R, 9i X }1- + R denote the ith UL and LL objectives respectively. We study the following problem:
min max f1(2, v1 (x)) subject to dd tE(m)
Mi (x) E arg min gi v.), Vi [m].
=
Here, the optimization variable x is shared across all objectives f..g.. i E [n]. while the variables y: E [n] are only involved in their corresponding objectives f.. gt. The goal is to find robust solution x X' such that, the worst-case across all objectives is minimized. This is a generic problem which reduces to equation if we have single objective pair, that is n 1. Such robust optimization problem is useful in various applications, and especially necessary in any safety-critical ones. For example, in decision optimization, the different objectives (fi, gi) can correspond to different "scenarios" (such as plans for different scenarios), with x being the shared decision variable and being scenario-specift decision variables. The goal of equation 2 is find the robust shared decision & which provides robust performance across all the n considered scenarios, so that such robust assignment of decision variables will generalize well other scenarios. In machine

---

Published as a conference paper at ICLR 2023
learning, robust representation learning is important in object recognition and facial recognition where we desire robust worst-case performance across different groups of objects or different population demographics. In RL applications with multiple agents (Busoniu et al., 2006; Li et al., 2019; Gronauer & Diepold, 2022), our robust formulation in equation would generate a shared model of the world - the UL variable x: such that the worst-case utility, max; f((x; yf (3)), of the agent-specific optimal action - the LL variable y (I) - is optimized, ensuring robust performance across all agents.
An additional technical advantage of the general multi-objective problem in equation 2 is that it allows the objective-specifi variables yi E V/ to come from different domains, that is, 3 * V), i. E [n]] as stated in equation 2, this implies that the dimensionality dy, for the per-objective y4 need not be the same across all objectives. This allows for a larger class of problems where each objective can then have different number of objective specific variables but we still require a robust shared variable r. For example, in multi-agent RI,, different agents can have different action spaces because they need to operate in different mediums (land, water, air, etc).
Focusing stochastic objectives common in ML, the main contributions of this work are follows:
(New algorithm design) We present single loop Multi-Objective Robust Bilevel Two-timescale optimization algorithm, MORBiT, which uses (i) SGD for the unconstrained strongly convex LLL problem, and (ii) projected SGD for the constrained weakly convex UL problem.
(Theoretical convergence guarantees) We demonstrate that, under standard smoothness and point
regularity conditions, MORBiT T with n objectives converges to a "=="/""""""""""tiiiooo with K iterations, matching the best convergence rate for single-loop singlc-objective (n I) BLO algorithms with the constrained UL problem while using vanilla SGD for the LL problem, and providing a sublinear n' deeeddeeneeec on the number of objective pairs
(Two sets of applications) We present two applications involving min-max multi-objective bilevel problems, robust representation learning and robust hyperparameter optimization (HPO), and demonstrate the effectiveness of our proposed algorithm MORBiT.
Paper Outline. In the following section 2, we further discuss the different aspects of the problem in equation 2 and compare that to the problems and solutions considered in existing literature. We present our novel algorithm, MORBiT, and analyse its convergence properties in section 3, and empirically evaluate it in section 4. We conclude with future directions in section 5.
2 PROBLEM AND RELATED WORK
We first discuss the different aspects of the robust multi-objective BLO problem with constrained UL in equation2 2. While BLO is used in machine learning (Liu et al., 2021; I: Chen et al., 2022a), multi-objective BLO has not received much attention. In multi-task learning (MTL), the optimization problem a multi-objectivy problem in nature, but is usually solved by summing the objectives and using a single-objective solver, that is, optimizing the objective E. The robust min-max extension of MTL (Mehta et al., 2012; Collins et al., 2020) and RL (Li et al., 2019) have been shown to improve generalization performance, supporting the need for more complex multi-objective optimization problem that replaces the objective Σ, fa with the objective max; ft.
For SGD-based solutions to stochastic BLO, one critical aspect is whether the algorithm is single-toop (a single update for both x and in each iteration) or double-loop (multiple updates for the LL between each update of the UL x). Double-loop algorithms can have faster empirical convergence, but are more computationally intensive, and their performance is extremely sensitive to the step-sizes and termination criterion for the LL updates. Double-loop algorithms are not applicable when the (stochastic) gradients of the L.I. and UL problems are only provided sequentially, such as in logistics, motion planning and RL problems. Hence, we develop and analyse single-loop algorithm.
A final aspect of BLO is the constrained UL problem. When the UL variable corresponds to some decision variable in a decision optimization problem or hyperparameter in HPO, we must consider constrained form, E X C Rd. To capture more general form of the bilevel problem, we focus on the constrained UL setup. In the remainder of this section, we will review existing literature on single-objective and multi- multi-objecctiv BLO and robust optimization, especially in the context of machine learning. Table provides a snapshot of the properties of the problems and algorithms (with rigorous convergence analysis) studied in recent machine learning literature.
2

---

Published as a conference paper at ICLR 2023
Table I:
The problem studied here relative to representative related work. If studied problem not a BLO, the notion of single-loop or constrained UL (X' C Ix**) is not applicable. The 1sth row block lists general problems. The 2nd block lists algorithms with analyses. The final row is MORBiT T. This problem has been viewed both as single-level and bilevel. D: The problem can multi-objective but is treated as single-objective by summing the objectives, A In bilevel adversarial learning, the UL is unconstrained but the LL is constrained.
Problem/Methed
Bilevel
Multi-objective
Min-max
Single-loop
N' C B
Distributionally Robust Leamning
Advensarially Robust Learning
Multi-task Learning (MIL)
Robust MTL (Mchta ot al., 2012)
Meta-lcaming
BSA (Ghadimi Wang, 2018)
X
HiBSA (Lu al., 2020)
V
GDA (Lin et 2010)
X
TR TR-MAMI. (Collins al., 2020)
TTSA (Hong ct al., 2020)
X
7
StocBio et al., 2021)
X
MRBO (Yang et al, 202()
VRBO (Yang et 2021)
X
ALSET (Chen el al., 2021)
STABLE (Chen et al., 2022b)
MMB (Hu al., 2022)
X
MORRI (Ours)
Single-Objective BLO. Lately, many new algorithms have been proposed to solve the single-objective stochastic BLO problem equation 1. Ghadimi & Wang (2018) proposed the first double-loop BSA approach. StocBio (Ji et al., 2021) and VRBO (Yang et al., 2021) are double-loop schemes that improve upon the convergence rate of BSA but do not consider constrained UL problems. TTSA (Hong et al.. 2020) is a single-loop algorithm that handles UL constraints. MRBO (Yang et al., 2021) and ALSET (Chen et al., 2021) are single-loop algorithms improving TTSA's 's convergence rate but do not consider UL constraints. STABLE (Chen et al., 2022b) improves upon TTSA by leveraging an additive correction term in the LL update step (beyond a basic SGD step) while still handling UL constraints. In contrast to the above single-objective bilevel setup, our formulation in equation gives flexibility for inherently multi-objective problems in robust manner to obtain stronger guarantees, ensuring convergence of each individual objective, rather than the average objective.
Multi-Objective BLO. There has been limited number of works analyzing multi-objective BLO schemes (Sinha et al., 2015; Deb Sinha, 2009; et al., 2017). All of these works analyze the multi-objective BLO problem from game-theoretic point of view, using vector-valued objective with the notion of Pareto optimality. In contrast, we are the first to study the multi-objective BLO problem from traditional optimization perspective in terms of convergence properties and consider a min- max robust version of the multi-objectivv problem which produces a single solution that ensures the convergence of each individual objective instead of generating multiple Pareto-optimal solutions which trade-off the optimality of the different objectives. See further discussion in Appendix D.4.
Min-max Robust Optimization in Machine Learning. Min-max optimization is commonly used to achieve robustness, such as in distributionally robust learning (DRL) and adversarially robust learning (ARL). In DRI, Duchi & Namkoong (2018) and Shalev-Shwartz & Wexler (2016) showed that a min-max loss improves generalization due to variance regularization. HiBSA (I.u et al., 2020) and GDA (Lin et al., 2020) compute quasi-Nash equilibria with convergence guarantees. Robust optimization is shown to have strong generalization for new tasks in multi-task learning (Mehta et al., 2012) and meta-learning (Collins et al., 2020). While the classic MAML (Finn et al., 2017) can be formulated as a BLO problem (Rajeswaran et al., 2019), the precise problem analysed in Collins al. (2020) is a single-level one. In fact, we consider the bilevel form of the MR-MAML problem as one of our applications for empirical evaluation. In ARI., the minimum is over the loss and the maximum over the worst-case perturbation to inputs (Madry et al., 2017: Wang et al., 2019). In both DRL and ARL, the min-max objective is in the form min, max, f(z y) with oingle-objective. In contrast, we study general robust multi-objective BLO where the UL objective is dependent on the LL solutions, and where the minimization is over the variable x: shared across all objectives, and the maximization is over the multiple objectives, ensuring that each individual objective converges fast.
Closely related and concurrent work. Since our goals align with the properties of TTSA (Hong et al., 2020) - the single-loop nature and the ability to handle UL constraints our proposed MORBiT inspired by TTSA and can be viewed as robust multi-objective version. Beyond this adyancement,
3

---

Published as a conference paper at ICLR 2023
contribution also lies in the convergence analysis of MORBiT, which significantly diverges from that of TTSA. After our MORBiT was developed and released (Gu et al., 2021), STABLE (Chen et al., 2022b) was recently presented improvements to MORBiT T in future work. A very recent work (Hu et al., 2022) studies a problem that appears be quite similar equation 2, with common elements such bilevel and min-max, and proposes a single-loop multi-block min-max bilevel (MMB) algorithm. However, there are significant differences: (i) Firstly, in their setup, they consider an extension of min-max single level problem to min-max BLO, and min-max max is not meant to provide "robustness" among objectives. The applications in Hu et al. (2022) are restricted to problems such as multi-task AUC maximization instead of the common bilevel applications of representation learning and HPO. (ii) Also, Hu et al. (2022) do not consider a constrained UL problem. The problem in equation is not a generalization of their problem both our work and theirs are considering different setups with high-level commonalities. For more details, see Appendix D.2. Table shows how our setup compares to existing literature. To the best of our knowledge, the precise problem in equation 2 has not been studied in ML literature.
improvement of TTSA, and we wish to explore similar
3 ALGORITHM AND ANALYSIS
In this section, we propose a simple single-loop algorithm MORBiT to solve equation 2, and establish rigorous convergence rate and sample complexity for this algorithm. For the theoretical results, we defer the precise assumptions, statements and proofs to Appendix A and present the high-level theoretical results and critical novel proof steps here. the sequel, we will always use the subscript i [ni] to denote the objective index and the superscript (k) to denote the iteration index, with (kk)
denoting the 1ath iterate of the shared variable x E X C Rd and yi denoting the b iterate of the ---oobjecivee-peccffcc variable E Rdy, We will also use the shorthand y to denote all the per-objective variables [31,, Wu,.y2.... ......., yn], with (k)) denoting the hthh iterate of all the per-objective variables (k) (ds)
3/2 ...... Bm (k)) Given assumption that the LL objectives gi strongly convex, define y (I:) A arg min y.R (i( y), and use the shorthand (1:) A fe(x, v(
3.1 MORBiT ALGORITHM
We begin with a standard reformulation of robust min-max problems (Duchi al., 2008). We can rewrite the non-smooth min-max problem in equation 2 as
min max A. (x)) subject to ZE AEAm
y (x) = arg min gi (1, Wi), Vi E [n] B.ER
(3)
RE(m]
where A, TR the n-simplex defined as A. = E E RT: X 2 0, In),E A 1}. This problem is equivalent to the min-max problem in equation 2, but allows us to solve the problem with (projected) gradient based methods. The gradient for ys, i [xx] is the straightforward (1, yi) and we denote h4 as its stochastic estimate, with h as the shorthand for the per-objective stochastic gradient estimates [h1, h2s ... hn]. The gradient for the z-update is more involved because of the hicrarchical structure of the BLO problem. Then, we consider the following weighted objectives utilizing the simplex variable X E An to define the necessary gradients:
A(I, =
X,((z)), F(2,V,A)= Σ ig[n] >xf( v.).
(4)
Note that F(x, A) is the UL objective in equation 3, and the UL gradients can be defined as:
V=F(e, E vE(M) (x), VxF(2, A) .... (n))))
(5)
where (2((x) for any E [m] can be defined as follows utilizing the strong convexity of the L.I. problem and implicit gradients (Gould et al., 2016):
(x) Va of (4)) 2 (I,, of
(I,
(x)) V 9i,,, (x)) -V,
(6)
Note that in general, yf (1r) cannot be computed exactly. Following Ghadimi Wang (2018), we use an approximation of V (x(( as a surrogate, denoted by Vafi(z, y.). by replacing (x) in equation 6 with any y: E Rd follows:
Vali v)) Vaf( yi) V2y. (xr, 3k) [V2 9. y:)] Vy,l( v.).
4

---

Published as a conference paper at ICLR 2023
Consequently, define our approximate gradients for the UL variables (and X) as:
F2 y.
Σ im(m) (., VxF(2, y. uz), [n( Un)] (8)
We denote the (possibly biased) stochastic estimates of V. y, A) as h2 and VxF(e, y, A) as hix
Given the gradients and their stochastic estimates, we present our single-loop algorithm MORBiT in algo- rithm I, where we utilize learning rates a, 8,77 > 0 for the UL variable z, LL variables yı, E [n] and the sim- plex variable À respectively. The algorithm tracks three sets of variables 2,(8), y = [b] 3/2 (k) through a total of K iterations. The per-iterate gradient estimates h (f) of the LL variables y is defined as the collection of the per-objective gradient estimate
Algorithm 1: MORBiT with learning
rates o, 8 and y for x. y, À respectively
for ((+)) <- (x)) - Bh
1,2,- K do
(k+1) projx (300)aal ah!") ((++1) t- proja, (1 + 71)
) and
end
Sample U(f1, K}) return E <- (7) Di f- Di
h (k) evaluated (zz(*), y/ w/(h) for all i [rz]. The gra- dient estimates 1/ and (/) of the UL variables x and À are evaluated at (a:(k) y (1++1) 11++1,,00 (())). We perform standard gradient descent update for the objective specific variables y,, i [n]] from y," to y/ (k+11) For the shared UL variable z, perform projected gradient descent to satisfy the UL constraints, where (-) denotes the projection operation onto the constrained set X' We update the simplex variable via projected gradient ascent, where project the variable back onto the n--immlex after a gradient ascent step with proja., (-). Given the learning rates (a, 3, 7), MORBiT is quite straightforward in terms of implementation. When n 1, the problem reduces to single-objective BLO, (4) 1, and MORBiT reduces to TTSA (Hong et al., 2020).
X
3.2 ANALYSIS
Given the single-loop MORBiT, we establish conditions under which MORBiT has finite-horizon convergence. The coupling of the stochastic errors due the sampling process makes the convergence analysis of this three-sequence -eauenc--bbed algorithm much more challenging than existing BLO algorithms.
Assumptions. We summarize the following typical assumptions (detailed in Appendix A.1) for all objective pairs (fi.g:). E [n]. Focusing on the smoothness and regularity properties of the objectives, we assume that (i) the L.L. objective Ji is strongly convex in yú. wwice-differentiable and has sufficiently smooth first and second order gradients (Assumption 2 in Appendix A)1), ), (ii) the UL objective fi has sufficiently smooth first order gradients, and (iii) the function (x:) 2 fi(x, yf (x)) weakly convex, bounded and has bounded first-order gradients (Assumption
Appendix A.1,
also see Appendix D. Regarding the quality of the gradient estimates 4, E [n]], hgg (k) and (K)
xx we assume that, for all l: >> O, (i) 114)) is an unbiased estimate with bounded variance, (ii) (*)
is an unbiased estimate, and (iii) h.th) has bounded variance, and can be biased estimate of the ((()) term defined equation 8, but the bias norm at iteration & is bounded by by 2 O. with (be, & 2 0} forming a non-increasing sequence. These gradient estimate quality assumptions are detailed in Assumption 3 in Appendix A. While the assumptions on h, and k(()) are standard (Hong al., 2020; Lu et al., 2022), the assumption on (1) actually can be easily satisfied when a Hessian inverse approximation (HIA) based mini-batch sampling strategy is adopted, which can also avoid the matrix inversion by leveraging the Neumann series (Agarwal al., 2017;
V ((*(()) y (x++1)
Ghadimi & Wang, 2018; Hong et al., 2020).
Optimality and Stationarity of Solutions. To quantify the convergence properties of the solutions T. Di. E [n], X generated by MORBiT, we use the following optimality properties of the optimal solutions x² Vi. [m]. A' of the problem in equation 3. (i) The per-objective optimal LL variable of v, (x*) = arg min w 9:(2", vi); (ii) The optimal simplex variable X": F(z*, X*) =
maxxes, F(x A). Given the constrained UL. the first-order stationarity condition is satisfied if (V, (('''''''' X²) X*), x²) 2 OVx E X. (iii) For establishing near-stationarity of UL variable a, the proximal map (2)) X. defined below.
a(2) A arg min L x - 2112 + F(I, A), p> is a fixed constant.
(9)
5

---

Published as a conference paper at ICLR 2023
is employed (Davis & Drusvyatskiy, 2018; Hong et al., 2020) to quantify the convergence for constrained variable in the stochastic setting. If [ is small, then, near-stationarity of (#()) is achieved at iteration k. Therefore, we need to bound [[)(z) - =l| to guarantee the convergence of the UL solution ē returned by MORBiT. Given the convergence of the UL T. we also need to bound I| yt [3) for each E [n]] simultaneously to quantify the convergence of the LL variables. Finally, the convergence of requires us to bound the difference between F(x, X) and maxxeA. F(Z, A).
Theoretical Convergence Rate. Now, are ready state main theoretical result: a rigorous convergence rate for the solution returned by MORBiT (algorithm 1). We state an abbreviated version of the result, deferring details to Theorem 2 in Appendix A.2:
Theorem 1 (MORBiT convergence). Suppose that the previously stated assumptions holds and learning rates are set as = O(K- "),,, 8 = O(K- and = O(n ((nnn////p K- -/))- Then, if iy S a, the solutions こ, D.. i E [n]. À generated by algorithm 1 satisfy:
-1| - x =|| S Ō(vnk-221 2/5),
(10a)
max 10: iE(n
yg (z) S Ōθlvkk-
(10b)
max ( X)] Ō(Vnk-2/5)
(10c)
with expectation over the stochastic gradient estimates and the random index T (algorithm 1, line 6).
This result establishes the the UL and LL variables if all the assumptions are satisfied and the learning rates are selected appropriately. Note that, if the UL problem is unconstrained, that is x E X = R the definition of the proximal map (equation 9) implies that F (I, X) I Ö(n 12 K 2/), providing the convergence of E to a K[ 3/5) - - point the UL problem is unconstrained.
achieved by K iterations of MCRBiT for both
Comparison with Related Work. We would like to further highlight the differences between the convergence results of TTSA and MORBiT to highlight the major novelties in our analyses and theorem proving techniques. First, we consider more general proximal map in equation 9 involving a weighted sum of weakly convex functions & instead of a single weakly convex function in TTSA, requiring new construction of potential functions for establishing the convergence of the UL variable I in equation 10a. Secondly, even though TTSA provides a convergence rate for a single LL variable (equivalent bounding E[[l;, y (+) [|]] for single E [n]). we provide much stronger result for multiple LL optimization objectives, in the sense that simultaneously establishing convergence for all LL variables in equation 10b through measuring the convergence rate of [[[max,[[m] l1ū. y (z) This is especially challenging since a bounded [E[lg; - y (3)ll] for each i E [n] does not directly imply a bounded [[|maxe]]n] 11% yf (3) in fact this can be generally unbounded. Finally, to
satisfy the requirements of the min-max max problem in equation 3, we have to additionally establish
convergence for the simplex solution À in equation 10c while TTSA does not have any such analysis.
Given the convergence rate, another related quantity of interest is the sample complexity which pertains to the number of queries to the stochastic gradient oracle required to achieve a desired level of stationarity. For example, for an iterative algorithm that converges to a ((("">)-stttonaayy point with K iterations for some μ > 0, requiring O(1) queries to the stochastic gradient oracle in each iteration, the sample complexity to find an e-optimal solution is (/a).. The number of stochastic gradient oracle queries required directly related to the conditions the gradient estimate quality assumptions (Assumption 3 in Appendix A. in our case). While the conditions the per-iterate gradient estimates 1,(5)) (for the per-objective LL variables) and (*) (for the simplex variable A) both only require (1) stochastic gradient oracle queries from each of the objective pairs in each
iteration, the condition 1g. S a on the non-increasing squared norm of the per-iterate bias of the
gradient estimate (k) (for the UL variable) require Oloo K) stochastic gradient oracle queries for each of the n objective pairs leveraging the HIA sampling techniques in Ghadimi & Wang (2018) 8) and Hong et al. (2020) using the Neumann series (Agarwal et al., 2017). This gives us the following sample complexity bound for MORBiT T (see Appendix D.3 on potential improvements):
Corollary 1. Under the conditions of Theorem 1, MORBiT converges to stationrit with 0//%/ 5/2 log(//)) queries to the stochastic gradient oracle for each of the objective pairs.

---

Published as a conference paper at ICLR 2023
Proof Sketch of Theorem 1. We now give proof sketch of our main theorem, with constant terms abstracted away with O notation. order to show equation 10a of Theorem (convergence of æ),, we will derive a descent lemma comparing successive iterates (A) and x (x(++1), Descent lemmas often
contain a quadratic term || - =² I|? we bound the expected squared norm of the stochastic gradient estimate:
so it is natural that we must bound ||As* In Lemma 1,
(1) Lemma 1. Under our regularity assumptions, Elll 11 S O
1
x iEfn
Σ
1 (z(()) yy 1 12
Turning equation 10b of Theorem 1, we descent relation on y; (8) we would obtain a descent relation purely involving yi (k) terms themselves, the intricate coupling of the z and terms result in an extra x=(k1), - (*) term. The resulting relation is shown in Lemma 2.
V (a (x:(k-1)). While ideally
Here, we have that = 1 "g
C2 Mgg 1, where MgB 1.
Lemma 2. E
(1+1) ut (1) = S ((1
01 Billy! us (c: o2 - 111)
xx(8)
From this lemma, intuitivelyy, know that Ellly oi (1((-1))119 is decreasing as k increases, as long the [[]]=;(--1) - 122 are not too large. Therefore, it is important have another descent relation that upper bounds this quantity, which we do next in Lemma 3. The lemma naturally
involves the objective F( (a)),, which will telescope. Here, we have cg Ia T C4 4aL2.
L
As + oc, a + 0, and c3 is positive.
Lemma 3. Let ((k) A FF (x: { Eielm)
X, ((e Then, the ((k) satisfies:
((K+1) - (k)) So ( (ccgBllbx((tt) - 119 c4 max B Bllb (k+1) si (z)))) Vny + )).
sE|m]
Following the intuition previously described, we then use Lemmas and 3 to show that the (1]]x(0--1 - -x(a)113) terms are small enough and that the (k) iterates converge in Lemma 4.
Lemma 4. DK max E (k) ut (ark (=(k--))11]] O(VTK-2/0).
sE[m]
Finally, we can use the convergence of y; to prove convergence of À and æ. Theorem then follows. more detailed proof plan can be found in Appendix A.3.
EXPERIMENTAL RESULTS
In this section, we consider two applications where the min-max multi-objective bilevel formulation in equation 2 enhances robustness multi-task representation learning and hyperparameter optimization. We will highlight the advantage of the min-max formulation and the convergence of MORBi T on these applications. We use Py PyToorch (Paszke et al., 2019), and implementation details are in Appendix C. All results are aggregated over 10 trials.
Representation Learning. In this setup, each objective pair corresponds to a learning "task" i E [m]] with its own training and validation dataset pair D DV We consider a shared representation network 6e with ReLU nonlinearity (making the UL problem weakly convex) parameterized with z e Rdd and per-task linear model Wy, parameterized with yı E R Here the UT. is unconstrained. Using L(f, D) to denote the loss of model f on data D, we consider the problem in equation with
f1(x, 9:) A L(W,, o Φz. D.), 9: (x v:) A L( o Ó2. DE) + pllw,, lI2.
(11)
with p > as a regularization penalty (ensuring that the LL problem is strongly convex). We first consider multi-task setup with = 10 binary classification tasks from the FashionMNIST dataset (Xiao et al., 2017). The goal is to learn a shared representation and per-task models so that each of the tasks generalizes well. Usually, this problem is solved as single-objective BLO by minimizing 1/1 E; f.; we call this min-avg. We theoretically show that solving the min-max max multi-objective BLO in equation 2 provides fighter generalization guarantee (Proposition 2, Appendix B).
7

---

Published as a conference paper at ICLR 2023
Hoximur MSC
Tock Ralaut Standord Triekin
Val. loss(min-maxi
loec(min-sivg)
Gradient norm: [|
Highte tree (oor 2
seen RIS
unonne loss
2
Gen, lasslmin-nras
Worst case
Gen. losstmin-avgi
2.13
2-11
2:12
1003
2000
3000
4200
5000
0
2000
1208
6200
0200
10000
Iterations
Iterstions K
2000
0000
1000
ppan
(a) Quality of min-avg min-max.
(b) Convergence of V 1|?.
(c) Task-gen, of min-avg vs min-max.
Figure Numerical results for representation learning application.
Here we demonstrate the same in figure la- we plot the worst-case UL objective (the validation loss) and the worst-case generalization loss across all tasks/objectives throughout the optimization trajectory, comparing the behaviour of the solution of min-avg problem to that of the min-mar problem. The results indicate solving the min-max problem significantly reduces the worst-case validation loss and this also results in a significant reduction of the worst-case generalization loss, highlighting the utility of solving the min-max multi-objective bilevel problem in equation 2.
We study the convergence of the UL variable for the min-max problem in the form of the trajectory of the (stochastic) gradient norm | 2 in figure 1b, comparing it to the theoretical O(K-)//) rate. We see that the empirical trajectory of the gradient norm closely tracks the theoretical rate.
We also consider bilevel extension of the robust meta-learning application (Collins et al., 2020) for a sinusoid regression task, a common meta-learning application introduced by Finn et al. (2017)'. Here the goal of solving the problem in equation with the objectives defined in equation be to learn a robust representation network such that we not only improve generalization on tasks seen during the optimization but also improve generalization for related unseen tasks.
would
We theoretically show that that solving the min-max multi-objective bilevel problem in equation 2 also provides tighter generalization guarantee for the unseen tasks (Proposition 3, Appendix B). Note that these results are similar in spirit to those of Mehta et al. (2012) and Collins et al. (2020), but our results are the first for a general bilevel setup. The results in figure Ie support this, showing that solving the min-max problem not only improves the generalization on seen tasks, but significantly improves the generalization on unseen tasks when compared to solving the min-avg problem. These results are also consistent with the results for robust MTL in figure la.
Hyperparameter Optimization. In this setup, each objective pair again corresponds to a learning "task" E [n]]. each with its own d dimensional valliing/valiatiion dataset pair D Dÿ. We consider a shared hyperparameter optimization problem for kernel logistic regression (Zhu & Hastie, 2001) with K random Fourier features (RFFs) (Rahimi & Recht, 2007), where x {Ir R Io R are the regularization penalty and the bandwidth hyperparameters respectively, with Φz, denoting the RFF2. The per-task linear model Wy: on top of the RFFs are parameterized with y/ E R2K In this setup, we have a weakly convex constrained UL problem (the hyperparameters need to be positive), and an unconstrained strongly convex LL problem. Again using L(f, D) to denote the learning loss of model f on dataset D, we consider the problem in equation 2 with
f((2, y:) A L( o Φz,, D gz 34) A L(Wy, o Φz,, DF) + 1x O w.
(12)
where @ denotes the elementwise vector multiplication, and we consider a weighted regression penalty?. We generate n = 16 binary classification tasks from the Letter dataset (Frey & Slate, 1991) and compare the generalization of the min-max solution of equation to that of the min-avg.
The results in figure 2a indicate that the solution of equation provides a robust solution r (hyper- parameters), significantly improving not only the worst-case validation loss but also the worst-case generalization loss for the supervised learning problems. This result highlights the advantage of
We use the formulation of Raghu et al. (2020) to separate the representation and the model parameters.
2 For a d'dimension point p. the RFF éz, (p) = [sin(W(s (x, p)) cos(W C p))T] E R2K where W I Agxx K'dd is random normal matrix and the sin(.) and cos(:) are applied elementwise.
'The weighted regression penalty mitigates bias especially in the high- nigh-dimensional learning setting (Candes al., 2008; Gasso al., 2009; Šchié al., 2022), which is common when using RFFs.
8

---

Published as a conference paper at ICLR 2023
Val. losstmii max]]
Grad norm: [
2:40
Worst obj.
Gen. losstmin-maxi
loss(mi--aug)
2.43
3e
lozz(min aug)
2712
120
2:12
5003
10000
15000
20000
3000
10000
15000
20000
0
5000
10:00
15000
20000
Rteratioss
K
Iterations
Rteations K
(a) Quality of min-avg vs min-max.
(b) Effect of n on convergence.
(c) Effect of batch size.
Figure 2: Numerical results for hyperparameter optimization application.
solving the min-max problem in equation and the ability of the single-loop MORBiT to handle a weakly convex constrained UL problem.
We study the effect of the number of objective pairs n on the convergence. We consider n E {4, 16, 64), increasing n with factor of 4 (implying theoretical convergence slow down by factor of 2) to check how the convergence matches the /n-dependence in our theoretical result.
In this case, we consider the trajectory of the (stochastic) gradient norm V Il2 (as in figure lb). The results in figure 2b display such a behaviour - for a fixed K (outer iterations), as the number of tasks increased 4-fold, the gradient norm approximately increases 2-fold (note the -scle on the vertical axis). This validates our theoretical dependence on the number of objective pairs n.
We also study the effect of the batch size on the generalization performance of the min- max solution. In the previous experiments, we considered batch size of 8 for both the UL and L. stochastic gradients. Here, we will consider batch sizes from {8, 32, 128), using the same batch size for gradients of both levels and variables.
Note that, this problem, each of the 16 learning tasks (and hence, objective pairs) has training set size of around 900 samples (for the LL loss), with 300 samples each for the UL loss and for computing the generalization loss. Unlike figures and 2a, we only show the generalization loss (dropping the validation loss) in figure 2c. The results indicate that increasing the batch size improves the stability and reduces the variance of the overall generalization. However, the convergence follows similar trend for all batch sizes, and converges very similar level of generalization, supporting the (1) batch size requirement for convergence.
Empirical conclusion
The empirical evaluations highlight that considering the more robust min-max problem in equation 2 does provide improved generalization multiple applications (representation learning for MTL and meta-learning, and for hyperparameter optimization).
The results also highlight the validity of our theoretical convergence analysis both in terms of the number of iterations K and the number of objective pairs n.
5 CONCLUDING REMARKS
Motivated by the desiderata of robustness in bilevel learning applications, we study new min- max multi-objective BLO framework (equation 2) that provides full flexibility and generality. We propose MORBiT (algorithm I), a single-loop gradient descent-ascent based algorithm for finding an solution to our proposed min-max multi-objective framework. We establish its convergence rate (Theorem 1) and sample complexity (Corollary 1), demonstrating both the advantage of the min-max multi-objective BLO framework and the validity of our theoretical analyses on robust representation learning and hyperparameter optimization applications. We wish to explore further applications where robustness would be beneficial such as in RL, federated learning and domain generalization. On the theoretical side, we wish to develop single-loop algorithms with improved convergence rates (for example, exploring techniques in Chen et al. (2022b)) and double-loop algorithms with convergence guarantees for applications where single-lopp algorithm is not feasible (e.g., federated learning). Finally, we also wish to develop algorithms for large n (the number of objective pairs) or where MORBA is not computationally feasible.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
The formal definitions, assumptions, precise theorem statments, high level proof outline and detailed proofs for our main theoretical results are presented in Appendix A. We provide appropriate citations for the datasets used in our experiments and the experimental setup and details are presented in Ap- pendix Our implementation available at https: //g: /htbbub..gt con/ m wwwiiiaariiooo..
ACKNOWLEDGEMENTS
A.G. is supported by the National Science Foundation (NSF) Graduate Research Fellowship under Grant No. 2141064, and T.-W. Weng is supported by NSF under Grant No. 2107189. We would like thank the MIT-IBM Watson AI Lab (https: / //mit s wwwwatsonaille MIT-UROP program (https: / /urop...... mi t. edu/) /) for their support. We would also like to thank the organizers of the "Beyond First-order Methods in MI. Systems" workshop at ICML''21 and the "Bilevel Stochastic Methods for Optimization and Learning" session INFORMS 22 for giving us the opportunity to present various iterations of our work (Gu et al., 2021; 2022). Finally, we would like to thank Soumyadip Ghosh and Mark Squillante for some insightful discussions.
mit wwt.mdd.. eduu) and the
REFERENCES
Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. The Journal of Machine Learning Research, 18(1):4148 4187, 2017. 5, 6, 16, 32
Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable representation learning for imitation learning via bi-level optimization. In Proceedings of International Conference on Machine Learning (ICML), pp. 367-376. PMLR, 2020. 1
Jonathan Bard. Practical bilevel optimization: algorithms and applications, volume 30. Springer Science & Business Media, 2013.
Lucian Busoniu, Robert Babuska, and Bart De Schutter. Multi-agent reinforcement learning: A survey. In Proceedings of the 9th International Conference on Control, Automation, Robotics and Vision, pp. 1-6. IEEE, 2006. 2
Emmanuel J Candes, Michael B Wakin, and Stephen P Boyd. Enhancing sparsity by reweighted & minimization. Journal of Fourier Analysis and Applications, 14(5):877-905, 2008.
Can Chen, Xi Chen, Chen Ma, Zixuan Liu, and Xue Liu. Gradient-based bi-level optimization for deep learning: A survey. arXiv preprint arXiv: 11719, 2022a. 1,2
Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis alternating stochastic gradient methods for bilevel problems. Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 34:25294 25307, 2021. 3
Tianyi Chen, Yuejiao Sun, Quan Xiao, and Wotao Yin. single-timescale method for stochastic bilevel optimization. In Proceedings of International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 2466-2488. PMLR, 2022b. 3,4,9,35
Liam Collins, Aryan Mokhtari, and Sanjay Shakkottai. Task-robust model-agnostic meta-learning Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 33, 2020. 2, 3, 16, 30, 32, 36
Damek Davis and Dmitriy Drusvyatskiy. Stochastic subgradient method converges at the rate O(4-1/4) on weakly convex functions. arXiv preprint arXiv: 1802.0298,, 2018. 6
Kalyanmoy Deb and Ankur Sinha. Solving bilevel multi-objective optimization problems using cvolutionary algorithms. In Proceedings of International Conference on Evolutionary Multi- Criterion Optimization, pp. 110-124. Springer, 2009. 3
Stephan Dempe. Foundations of Bilevel Programming. Springer Science & Business Media, 2002. 1
10

---

Published as a conference paper at ICLR 2023
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. arXiv preprint arXiv: 1810.08750, 2018. 3
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto the bilblll for learning in high dimensions. In Proceedings of the 25th International Conference
Machine Learning, pp. 272-279, 2008. 4
Heshan Devaka Fernando, Han Shen, Miao Liu, Subhajit Chaudhury, Keerthiram Murugesan, and Tianyi Chen. Mitigating gradient bias in multi-objective learning: A provably convergent approach. In International Conference on Learning Representations, 2023. URL https:h//tpenaaaaeeeeettt /oppnrevie.....eee net / Eorum? -iteenidddddAaAGAGaa
36
Chelsea Finn, Pieter Abbeel, and Sergey Levine Model-agnostic meta-learning for fast adaptation of deep networks. arXiv preprint arXiv: 000. 770000340,, 2017. 3, 8, 32
Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riceardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In Proceedings of International Conference on Machine Learning, pp. 1568-1577. PMLR, 2018. 36
Peter W Frey and David Slate. Letter recognition using holland-style adaptive classifiers. Machine
Learning, 6(2):161-182 1991. 8
Gilles Gasso, Alain Rakotomamonjy, and Stéphane Canu. Recovering sparse signals with a certain family of nonconvex penalties and de programming. IEEE Transactions on Signal Processing, 57 (12):4686-4698, 2009. 8
Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. arXiv preprint urXiv:/802.0220 2018. 3, 4, 5, 6, 14, 16, 32
Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint arXiv: 1007:04677 2016. URL https: /awxiv.ivvi...//i//..e/aaaa org /df/ 1607. 01747.554..ppmm 4
Sven Gronauer and Klaus Diepold. Multi-agent deep reinforcement learning: a survey, Artificial Intelligence Review, 55(2):895-94) 09455-993,44 2022. 2
Alex Gu, Songtao Lu, Parikshit Ram, and Lily Weng. Nonconvex min-max bilevel optimization for task robust meta learning. In Beyond First-order Methods in ML Systems workshop at ICML 21, 2021. 4, 10
Alex Gu, Songtao Lu, Parikshit Ram, and Lily Weng. Robust multi-objective bilevel optimization with applications in machine learning. In INFORMS Annual Meeting, 2022.
Mingyi Hong. Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic. arXiv preprint 2rXiv:2007.05/70 2020. URL https: arxiv. org http://d/2200....//// pdf. 1, 3, 6, 14, 15, 16, 32, 35
Quanqi Hu, Yongjian Zhong, and Tianbao Yang. Multi-block min-max bilevel optimization with applications in multi-task deep auc maximization. arXiv preprint arXiv:2206.00120 2022. 3, 4, 35
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Provably faster algorithms for bilevel optimization and applications to meta-learning arXiv preprint arXiv:2010.07962. 2020. 33
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In Proceedings of International Conference on Machine Learning (ICML), pp. 4882-4892. PMLR, 2021. 32
Ying Ji, Shaojian Qu, and Zhensheng Yu. A method for solving multiobjective bilevel programs. Discrete Dynamics in Nature and Sociery, 2017, 2017.
Shihui Li, Yi WU, Xinyue Cui, Honghua Dong. Fei Fang. and Stuart Russell. Robust multi-agent reinforcement learning via minimax deep deterministic policy gradient. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4213-4220 2019. URL https: //ojs. http://exx..pp//a///rereeeee
4327/4205.
11

---

Published as a conference paper at ICLR 2023
Tianyi Lin, Chi Jin, and Michael Jordan. On gradient descent ascent for nonconvex-conent minimax problems. In Proceedings of International Conference on Machine Learning (ICML), pp. 6083- 6093. PMLR. 2020. 3
Risheng Liu, Jiaxin Gao, Jin Zhang, Deyu Meng, and Zhouchen Lin. Investigating bi-level opti- mization for learning and vision from unified perspective: A survey and beyond. arXiv preprint arXiv:2101. 11517, 2021. 1,2
Songtao Lu, Ioannis Tsaknakis. Mingyi Hong, and Yongxin Chen. Hybrid block successive ap- proximation for one-sided non-convex min-max problems: algorithms and applications. IEEE Transactions Signal Processing, 68:3676-360 0677666691111 2020. 3
Yucheng Lu, Yi Meng, and Christopher De Sa. general analysis of example-selection for stochas- tic gradient descent. In Proceedings of International Conference on Learning Representations (ICLR), 2022. 5, 15
Aleksander Madry. Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv: 1001 06083, 3
Nishant A. Mehta, Dongryeol Lee, and Alexander G. Gray. Minimax multi-task learning and generalized loss-compositional paradigm for mtl. In Proceedings of the 25th International Conference on Neural Information Processing Systems, pp. 2150-2158, 2012. 2,3, 8, 36
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of Machine Learning. MIT press. 2018. 31
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019. URL http: /ttt:://ppper..... neurips cc/paper///cceccpecccccccee.eeee 2015-pytorch-aa an imperative style Sectyle-hihhte wer.hhh-prfrrraancccc pdf. 7
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse?
towards understanding the effectiveness MAML. In Proceedings of International Conference on Learning Representutions (ICLR), 2020. 8
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Proceedings of the 20th International Conference on Neural Information Processing Systems, pp. 1177-1184, 2007. 8
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), pp. 113-124, 2019.
Kenan Šehic, Alexandre Gramfort, Joseph Salmon, and Luigi Nardi. Lassobench: high-dimensional hyperparameter optimization benchmark suite for lasso. In Proceedings of the First Conference on Automated Machine Learning (Main Track), 2022.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014. 31
Shai
and Yonatan Wexler. Minimizing the maximal loss: How and why. In Proceedings of International Conference on Machine Learning (ICML), pp. 793-801, PMLR, 2016.
Ankur Sinha, Pekka Malo, and Kalyanmoy Deb. Towards understanding bilevel multi-objective optimization with deterministic lower level decisions. In International Conference on Evolutionary Multi-Criterion Optimization, pp. 426-443. Springer, 2015.
12

---

Published as a conference paper at ICLR 2023
Jingkang Wang, Tianyun Zhang, Sijia Liu, Pin-Yu Chen, Jiacen Xu, Makan Fardad, and Bo Li. Towards a unified min-max framework for adversarial exploration and robustness. arXiv preprint arXiv:1906.03363. 2019. 3
Nic Wilson, Abdul Razak, and Radu Marinescu. Computing possibly optimal solutions for multi- objective constraint optimisation with tradeoffs. AAAI Press/International Joint Conferences on Artificial Intelligence, 2015. 36
Han Xiao, Kashif Rasul, and Roland Vollgraf. Rashion-mnise a novel image dataset for benchmarking machine learning algorithms. ArXiV, 2017. 7
Haibo Yang. Minghong Fang. and Jia Liu. Achieving linear speedup with partial worker participation in non-iid federated learning. arXiv preprint arXiv:2 1203, 2021. 3
Runzhe Yang, Xingyuan Sun, and Karthik Narasimhan. A generalized algorithm for multi-objective reinforcement learning and policy adaptation. Advances in neural information processing systems, 32, 2019. 36
Yihua Zhang. Yuguang Yao, Parikshit Ram, Pu Zhao, Tianlong Chen, Mingyi Hong, Yanzhi Wang, and Sijia Liu. Advancing model pruning via bi-level optimization. In Annual Conference on Neural Information Processing Systems, 2022.
Yihua Zhang, Pranay Sharma, Parikshit Ram, Mingyi Hong, Kush R. Varshney, and Sijia Liu. What is missing in IRM training and evaluation? challenges and solutions. In International Conference Learning Representations, 2023. URL https: openrevi www.wiiwwweeee.......ee........e net. forum? id-MjsDeIcDEy,
Pu Zhao, Parikshit Ram, Songtao Lu, Yuguang Yao, Djallel Bouneffouf, Xue Lin. and Sijia Liu. Learning to generate image source-agnostic universal adversarial perturbations. In International Joint Conference on Artificial Intelligence, 2022.
Ji Zhu and Trevor Hastie. Kernll logistic regression and the import vector machine. Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 14, 2001.
13

---

Published as a conference paper at ICLR 2023
A CONVERGENCE ANALYSIS OF MORBiT
A. ASSUMPTIONS
First, we begin by listing the assumptions we make:
Assumption 1 (Regularity of the outer functions). For all [n], assume that outer functions fe(2, y) and & (x) = fi(r, vt (x)) satisfy the following properties:
For any E X, fe(2z,) is Lipschitz (w.rt. y) with constant G
For any x E X, V and V Jol,,) are Lipschitz continuous (w.r... yi) with constants ,, and Lsy" >
For any E Rd,, Vu.fil.. V1) is Lipschitz continuous (w.rt. x) with constant L,, > O.
For any x E X, Vo E RA,, we have II y., vi)l S for Ciy 0.
The function (,(-) is pe weakly convex (in r), so that for all v, w E X,
((w)) 2 ((1) (Ve,)))
(VLs()), w - v) - -eellw - r1l2.
(13)
For all a e X, [[4(z)] S for Bę
For all x E X, [z)( - Cer for Ce > O.
Assumption 2 (Regularity of the inner functions). Assume that inner functions gi (z, Wi), Vi E [n]
satisfy:
For any x E X and E RRu,, ge(r, ) is twice continuously differentiable in (x.
For any x E X, (n, -) is Lipschitz continuous (w.rt. y:) with constant Lg.
For any x E X, ge(x, 1) strunnyglly convex in ys so that for all v, w E IRdv,
gi w) 2 gi(2, u) + (V Vege(I,)) w uv) +ggllu willlwwol ull?.
(14)
For any x E X. V2 ay, g( and V gi(,,)) are Lipschitz continuous (w.E. yi) with constants Lyza > 0 and Lagys > O, respectively.
For any x E and ys E Rds,, we have W.y.9((1 31)ll S Cysy for some Cgny >
For any yi E Rdd,, V2 De yi) and 2 w. wi) are Lipschitz continuous (wer.. x) with constants Lyzy > 0 and Lggyy > 0, respectively.
From these assumptions, we can show a few additional regularity- conditions. Since these conditions also be found in Ghadimi & Wang (2018, Lemma 2.2) and Hong al. (2020,
Lemma 2), we state these results without proof.
Lemma (Corollary Assumptions). Under Assumptions and stated above, for all x, æ 02
X C R da y E Rdv, [ni), we have
V (f((..)) ve,(x)) S Llbb (x:) v/1-
(15a)
(x)) vl (x2))) S Gyll ii T2 |
(15b)
[v((x)) (1,
Ly|[x] .2 ||
(15c)
where we define
L Cary
(Lyry
Lygu Cany
L 2 Ls, +
+ Cs
(16a)
Mg
Mgg
22g
(I,,, L.)
(Igzz
I JggyyCen
Lj A Ls, +
+ Csv
+
(16b)
Mg
Hg
Gy A C
(16c)
Hg
14

---

Published as a conference paper at ICLR 2023
Assumption 3 (Quality of stochastic gradient estimates). For any iteration k 0 and all i [n]]
the gradient estimates hi (k) for the LL variable y; satisfy the following for some og > 0 (Hong et al., 2020; Lu et al., 2022):
E[h (h = V v.Di(o!! ,yt
:Vii") )
(17)
E||l hg (k) v.
:3i (/ )19) (1 VV. .9 "ii )1?).
(1)
(18)
For any iteration k 0, the gradient estimate (h) for the simplex variable X satisfies:
1 = V Fx ,y gytt),, (1l ;1 ),
(n(x yn )
(19)
For any h 2 0 and a of > O, we assume that there exists a non-increasing sequence {bx fk20 such that
E 1 = (x)(x , ((++1),,11) + Bz, |Bl| S by. || (8) [[n o7.
(20)
(21)
A.2 MAIN THEOREM AND REMARKS
We now state our main theorem, in full. Recall our notation from equation 9: for fixed constant p > we defined the proximal map to be
((2) e arg min -11=- - =11² + F(x, A).
(22)
sEX
We also define the Moreau envelope as
P L11 211² + Σ Ax(s(z).
+ 11/p(2) min
(23)
2
In addition, we use the notation A
Elll (K) y (x(h-1)) 11 and j(k) A arg ma Ag
Finally, we define 0² of + 3C? (see Lemma 12). Now, ready to state the full version of our main theorem:
Theorem 2 (Convergence of MORBiT). Under Assumptions 1, 2 and 3, and the terms defined in Lemma 5, when step sizes are chosen as
K K3/55
M4 16G
a min
V
(24)
4G
4K
8 min
V,
(25)
Hg
2K -3/5
(26)
B/2
where v min (0./330+03),1/ 1/Pa), Algorithum produces F, A, D.. E [n]] satisfying:
E
max iE|n]
|| - 1/ ( A,
(27)
max E F A)]
(E(F(5,A)) S V2BBVTK-24 2/5 GjA,
(28)
16 (6 G
8(6] + L2A) (-4 + p)²
2a + 363 + 3L2A) A)
B[|((0) 2 <
(-ue+ + p)pK2/5
-4 +
(29)
where
Moo0/ /1g
(0)
16o2 1602/13
G./L
2/TB.G,/L
(b + (T2)
160
+
(30)
K3/5
K7/S
K4/5
K2/5
K2/5
K2/5
15

---

Published as a conference paper at ICLR 2023
Connection of TTSA (Hong et al., 2020). As generalize Hong et al. (2020), our proof follows similar structure. In particular, Lemma is a generalization Hong et al. (2020, Equation (14)), and our Lemma 7 combines Hong et al. (2020, Lemma and 4). Lemmas 8, 9, and 11 in our work parallel Hong et al. (2020, Lemmas 6, 5, and 7), respectively. Lemma 10 in our work deals with the maximization problem w.r.t. X, there is analogue in Hong al. (2020). However, it borrows techniques from Collins et al. (2020, Theorem 1).
We also discuss the convergence rate. Note that A in equation 30 dominated by the fourth term, /n/K2/ so it is clear that equation 27 and equation 28 converge at rate of (yykk 2/5). We give special attention to equation 29). Apart from the Stū/t-rett)* term, we see that the RHS of equation 29 converges at a rate of P(vnK-22/0 To understand the convergence of this term, we turn to (20) from Assumption 3. As discussed in section 3.2, by can be made arbitrarily small by running more iterations of the subroutine for estimating h (for example utilizing the HIA sampling scheme (Agarwal al., 2017; Ghadimi & Wang, 2018; Hong et al., 2020)). Therefore, as long as we run enough iterations (O(log K) for HIA) such that by I a, (29) will also converge at rate of O(VKK-22 2/5).
A.3 PROOF PLAN
Overall Roadmap: In what follows, we give a proof sketch of our main theorem, with constant terms abstracted away with O notation. Cı, C2, C3, are positive constants depending on Ly (defined Lemma 5), Mg (the LL objective convexity defined Assumption 2) and the learning rates a, 3 (in algorithm 1).
In order to show equation 10a of Theorem 1 (the convergence of æ), we will derive a descent lemma comparing successive iterates x*)) and ((++)), Descent lemmas often contain a quadratic term [|x(+++]) - 11 so it is natural that we will have to bound ||2 As such, in Lemma 6, we bound the averaged squared norm of the stochastic gradient estimate, B[Iln (2]:
Lemma 6. Under our regularity assumptions, the average squared norm of hz (k) can be bounded as follows, where the expectation is over the filiration Fl a {vi
(0) x0(0), ,... Mi
(k) (k), yi 1141).
E h
(k) x
50
Σ
(k) (x) yi 1
(31)
(k) yt (x-1)) While Turning to equation 10a of Theorem 1, we'll Il use a descent relation on yi
ideally would obtain a descent relation purely involving y; ((k)) terms themselves, the intricate
coupling of the x and y terms result in an extra x. (k-1) ((()) term. The resulting relation is shown in Lemma Here, we have that c = 2
3g
2 1, where 2gg < 1.
Lemma 7. The distance between the algorithm's 's iterates v satisfies the following descent equation,
(k) and the true inner optimum y (x)
Bllly (k+1) (1 s O(11 Bl)l (A) vi V. 1)1 [ - -0001111)))
(32)
From this lemma, intuitively, we know that Bllly (K) - y (x))(1 is decreasing as increases, as long as the E ([]l=(<-1) - (1) 's are not too large. Therefore, is important to have another descent relation that upper bounds this quantity, which do next in Lemma 8. The lemma naturally involves Lg 42
the objective ((k) (*), (k)),, which will telescope. Here, we have c3 As k + oc, + O. and is positive.
Lemma 8. Let ((K) & (1) (N), X1 = Σ x (1) (((,((() Then, the ((K) satisfies the descent
"
equation
C(k+1) - (())_ <o (c0))]1la((+1)
- (M) + ca max (k+1) yf (x(())) Vny +
sE|n]
(33)
16

---

Published as a conference paper at ICLR 2023
Following the intuition previously described, we then use Lemmas 7 and 8 to show that the [|||x((-1) - [[3] -[](=(6-))_(((1 terms are small enough and that the v!!) iterates converge:
Lemma 9 (Informal, see Appendix A.7 for precise statement).
K
max E [Io! vī (x((- SOlvnk-2/5) FE|M
(34)
K
Lemma then leverages the convergence ys: i [n]] to bound the convergence of X, and Lemma 11 shows the bound on By plugging in step-sizes into Lemmas 9, 10, and 11, Theorem directly follows.
Lemma 10. For any À E Aw the iterates of Algorithm I satisfy
1
E
Σ
- F(x (110)))))))))
A)
<o(vik 2/8).
(35)
K
Lemma 11. The iterates of Algorithm satisfy
E11 - -1)))) O(VKK- 2/5).
(36)
K
A.4 PROOF OF LEMMA 1 (LEMMA 6)
Stating Lemma more precisely:
Lemma 12. Under Assumptions 1, 2 and 3, the average squared norm of ha can be bounded as follows, where 03 o + 3C7
R
A (2)
36 + 3L
(1)
Ilu.(x
- v/ 1p
(37)
where the expectation is over the filtration F a {y: (0)
y (k) (1), yi (x+1))
Note that here the expectation over F/ needed in the last term.
{w, (0) (0)
yi
(k) (k) Mi (1+1)). expectation is
Proof. We can derive the following:
[2]
( (2)
E + ||
(38)
Ein! || V F (+) (k+1), ((0) + B&
(39)
(3)
03
V (z) (1) y (++1) (1)) + 1B
(40)
2
03
Σ
(8)
(z. (8) 4/ (x+1)) B
(41)
364
3
(z
fe
(k) w (1-11))
(42)
(1) is true because
E hlll (2)
11 El|lh!
+ || 11 2E(M
E (1)
|TTh
111
(43)
Elll (1)
| + 1|
(E(A
1>>
|I (2)
111
1|2
(44)
Ellla (1)
(45)
17

---

Published as a conference paper at ICLR 2023
(2) follows from (20), (3) follows from (21), (4) follows from definition of Vz, and (5) follows from Young's 's inequality, ||a + 611² I 3||a|| + [|6] Next, we bound the last term in (42). We start by using the fact that
Si
2
f f((( 3/ ((att)) - ve((2((1)))
+2
Σ X
((((x(())))
(46)
n
(8)
f (8), 2/4 ( ( 12
Vl.
(47)
2
where (1) follows from ||a b11 < 2||a|1 + 2|bb| and (2) follows from
Σ
pra
- Σ
psl
1=1
1=1
Next, we bound the first term in (47). From Lemma 5, we have
Vafi y )
(2.(2 - LI187 (ze (zx Yy; 12
(48)
Therefore, we can obtain
E 1 121 S of + 3b3 + 3L 2 n
(1)
(1)
(k)
+ y (x) - Yi ( 2 Σ
a
(49)
(=1
(2)
S 0 36 3L
y (x) y/ 1
(50)
where (1) comes from plugging (48) into (47) and (47) into (42), (2) comes from the definition of o using [[ < Of and ((k) E An.
A.5 PROOF OF LEMMA 2 (LEMMA 7)
We state the precise version of Lemma 2 here: Lemma 13. Under Assumptions 2, and when B² + oz s Bug and PGB < 1, the iterates
Vi satisfy the descent equation:
(k)
B[lly] (k+1) -v6 (x:())))1
(51)
S (1 2 |lu1" (K)
ut (2. (2:1221))122
+
( 1g,
2 - -11001100-11-- + Bog
x00
Proof. For particular (fixed) realization of the iterates z(1).... (1))) i [n], we have
.v1 (1)
(K) for some
Vi
E
Elll
1
Blll! (k) - E74, "" (k), (k) |1 | (1) , 111
(52)
(x)
(53)
0/ + +11+ +1+ +f)
(x 1
(54)
09 (1 + 03
2/ (8) V (x))) , (2
(55)
2 2 + (1 0 1 25 (1))))
(56)
where (1) follows from algebra, (2) follows from Bli V v. (z(())) -3 (b)) in equation 17 in
Assumption 3, (3) is from equation 18 in Assumption 3, (4) is from V (x(M) vi (z(())) 0 due to the optimality of yf (z.()), and (5) is due the Lipsshic continuity of V (x,
18

---

Published as a conference paper at ICLR 2023
Next, bound the difference between yi assume that x("),.... ,"",,,"""""""" ,x(,), (1) ... Mi is fixed and the expectation is over the stochasticity of the gradient estimates:
(k+1) and yf (x(()) the following, where again
E V/ (k+1) Vi (z. (12)
1 Ellby (k) /1,² (1) yi [19
(57)
l8b (k) v (a )12 + 3 | (k) |11 23( v1(())), V ( (²(
(58)
(3)
$ (1 23 )ll! (b) (2(00))) 33 11
(59)
(4)
& (1 23 1luA (1 )12 + 3o0z + 3² + 12 (2) v2 (11
(60)
(5)
S (1 1ly) (0) (z. Oo8
(61)
(
8 (1 - 3ma)
(1 (1+ c)( (1) v/ (x )1
12 + 1 (1+2) 11:1 (1(( - 21 (2) )12
+ 0
(62)
3M (1 - 3pg) (l vi (z(0-1)))
1+ 1 - Gelle(h-11 - 2000132 - + +200
(63)
where (1) is true by definition, and (2) holds by direct algebra and the unbiasedness assump-
tion B(b Vngi(a(()) (( in equation 17 in Assumption 3, (3) is from strong convex-
Bpglly (k) vI (122 (4) is from equation 56, (5) is from the assumption (² + o2
ity, B V (1,,9 ( (8) vt (x)(1)))
8p.. (6) is from the inequality ||a + b122 S
(1 + 1/c))a)la||3 + (1 + e)|bb|2, and (7) is from the Gy-lipschitzness of y (-) in Lemma 5.
Then, we choose = 2(1
fg/8 3] so that (1 (1 + c) = 1 - S/g/2 2 and + 1/c- = have e 0 because 2gg < 1. Plugging these expressions into (63), get
_ - 1. We PbB
1l (k+1) 2 (11
(64)
(1 1g
llu (k)
y (1
( 2 -1)
Gg 13 - 11 820
Mgg
which completes the proof.
A.6 PROOF OF LEMMA 3 (LEMMA 8)
We state the precise version of Lemma 3 here:
Lemma 14. Let (A)) ((F ((0)) = Σ A B(e,((()))) Under Assumptions 1, 2 and 3,
1=1
assume that the iterates (k(,), "i i A1k), Vk} are generated by MORBiT, then, ((k) satisfies the descent equation
(k++) (R)
L1 1
(65)
< 4xx max
A (1+1) (E(N) 4:
+
B[lla(tt)-
- + BnB? + 4crb3 + 2007
2
4a
Proof. First, since 6, is smoothh we know that for all E [m].
L, -0002
(((((++1)) s (((a(t)) + (1 (N), (Z,( +
(66)
2
19

---

Published as a conference paper at ICLR 2023
Taking / times the equation for i in (66) and summing, we can get
(((()tt))
(67)
(h) (x)()t) + ((ltt)) -x(0) Σ 1
Ve((v(())
Li +
2
8=1
Therefore, we have
n
(k+1) (x(x((+1)))
Σ
1 (k) (zx
==1
(k+1) (((
Σ
x (1) (((olttt)) + (x(k+1)- - ((h)
Σ 1 (2((1(())))
1=1
((1)
A(B)
+ Li -
(68)
Next, bound (A) and (B) respectively follows. First, upper bound term (A). First, from the non-expansiveness of projections and ( = proja, (2 + )))(()), we have [[ || S (1) Since (k+11) ( A.. [| - || z /2. Therefore, we know that || [A(++1)- - || I A minfv2, (1 []].
Based these facts, we can have
(A)
(K+1) E(g - Σ
x (4) ((t
(69)
(1)
(k+1) - ) (((o((E11))
1
(70)
1=1
(2)
| - -(++))>(++1] ((( (2(2( (
(71)
ABBVa @Vimin (K) 1tB,
S V |
(72)
where (1) is straightforward, (2) follows from Cauchy-Schwarz, (3) follows from the update rule for and the fact that |&(+) S B, from Assumption 1, (4) is from plugging in the definition of A, and (5) follows from >| S v2.
Then, we upper bound (B). First, from the non-expansivennss of projection and the update rule (k+1) ( and (k)) we know that
(8+11 - (N) + 11) < anl
(k) 1|?
(73)
(|x(++1) 112) 1|2 + 20 (x++1) _((),, (b)) <
(74)
[1o(0+11) - (k) + (10(k - (1),
(10)) < 0.
(75)
2a
20

---

Published as a conference paper at ICLR 2023
Therefore, can have
(B) = Σ 1 1 (((lt))))) -a(1)) _ ( (2 - -x(t)) =2l1)) (
(76)
((V2 2(6)) V. F(x((t), ((++1) ; 115)) Ba)) (zslat1) -
(x((+t))=2(h)) =z(h))
(1 y )()) + Ba) (x((h-1) x(()))
(77)
((V=F(2((4) (()) (.F(2(A) ()+ ,1(1)) BB)) (zlhit))-
- -2(ki))
(78)
+ V (z)( g(tt (()) + Bx - h!!h)² (2l(tt) - (slltt))= -2(1))
1
|
x
2a
V ((x( (k) (1))) ((r (K+)), .1(0) - -Ball2 + - (lll((tt)) - x(())
2c
1
F(x (x)(k) g(tt1) (001 + B₁ 1p + d _ (laktt)) 112
+
2d
1
(+++) - 11
(79)
2a
(6)
V F(x(()) (11))) - v. F(z!k) y ((tt1),,11 ((0)) - BBel2 + C _ (1l - 20)
2c
2
of
+ d
+
|2((+2) 11001122 (1((1 12
(80)
2d
2a
V F(
(())) - (zF(2 (k+1)), (100) B
2c
c c d
|| (0+ - - 1 1 of
+
(81)
2
2a
2d'
where (1) is by definition of ((x(K), (00)
(2) is from adding and subtracting
((VF(((()) y x()) + B&). (2(k+1) - (2))
) (3) is from adding (75) the previous
inequality, (4) is from applying the inequality (a, b) 2g |a 1 [|6]|2 both inner product terms,
(5) is from equation 20 and equation 21 in Assumption 3, and (6) is from algebra.
Plugging in our expressions for (A) and (B) into (68), we get
(1+1) (b++))
Σ
1 (M) ((((k))
1=1
1=1
<>11///
(k) || V
VV ((x())
-2 (0 (K+1), (100) -Bal2
2c
+
c
+
d
L,
1
10)
of
(82)
2
2a
2d
21

---

Published as a conference paper at ICLR 2023
Next, work bounding V ==F(x(()) yy (0+11,,(1)))) Bxll in equation 81. Observe that
V
F(x A ( V (((!! (1(+1), ,1(4) - -Ball
(1)
2 F (8), (k0) - F(o(() (k+1), (A) + 2
(83)
2 F (x (8) (10)- - -V (.F(((1) y (x+1), || +22
(84)
(3)
2 2
Σ (1
n
(Vei(z(4)) V (z 29 (+t1) })]
+ 20
(85)
( (alttt v/ (54+1111 2 +26
(4)
S 2
(86)
(5)
$ 2L max 11v2 (1)) - V/i -vi 11 + +bbb
(87)
8E|n]
where (1) comes from ||a + bl2 < 2 2| + +21bbl? (2) comes from || |& I by in Assumption 3, (3) is from expanding the definitions of ,. and V F.F(. 1 :). (4) is from Lemma 5. Therefore, plugging in equation 87 into equation 82, using d and taking expectation over the
stochasticity of the gradient estimates, we get:
(k+1) (1+t)) ((((h+)))
Σ 1=1
1 (k) ((((1))
S /nnllhh (16) (| + 4aL2
max El ElV (a:(4)) - -Mi 112
max iE(N)
+ 4ab2 +
[.g
1
E
B[[1 - 11) 22077
(88)
2
4a
Now, observe that the LHS looks like a telescoping sum. To make this more apparent, define
((k) Σ 1 Blex(((()) and (k) Ellx(( yi (1) Therefore, with the assumption
==1
that by S ty, we have
[(++1)- ((k)
S 4a max sE|ni Vi
A (k+1)
Tif
E [[[
- 1111 + /nnylh (x)) | + 4abh 200}
2 4a
Li
S 2 max A (k++1)
E [[[[2(+1])
1100 Bm + 4abh 2007.
iE(n] Vi
2
4a
(89)
A.7 PROOF OF LEMMA 9
We restate Lemma 9 in more general terms here:
Lemma 15. Assume that (k), (k)), T!()) As c1, c1, C2, do, d1, d2 are real numbers such that for all
< k < K-1,
((++1) < (k) - co (k++1)
+ max T
(k+1)
+
(90)
sc(n]
and also for all 1 1 k < K, 15is
r((+1)) < (1 d d, (M) + +
(91)
22

---

Published as a conference paper at ICLR 2023
In addition, assume that 1 do > O, do dicic > and co cidido and that Y, ,22k)
for all k, i E [ml. Then, if = arg maX ! we have
K
(r(0)
+ d,0000)
+ d2 + di
r S (du dico ())-1 EE|
200)
+ dic 02 + dz)
max
K
K
(92)
Proof. First, let ((M) & arg max T! so that
max T. seini
Summing (90) from k
se(n]
0,1, , K - 1, we get:
K
( IΩ - (2()) + c1 T(() + o2
Co
(93)
Next, we apply (91) for ((++1), Noting that 1 do > and r(()) we have
#( 12) S r r by definition of 200)
r((+1) 1(++1) S (1 - -d) ( ((+1)) + ( (k) +
(94)
S (1 do) r ((k)
+ d1 (k) + d2
(95)
Then summing for k = 1 to K, get
K
K
dg
r(k)
S (
r ((+l))
[(++)) + d1 Σ
(14)_ + dd2K.
(96)
10k)
km
kml
Now, we have
dg
K
(k)
2th) S (()) (((+1)
(1)
+ dic
22(0) (2000 + C1
K
r())
(06) + c2K ) d2K
(97)
201
k=1
(2)
K
v r 1(1) + dic
(2(0) + c1
Σ
r (R) ((k) + c2K
+ dgg
(98)
1=1
(8)
S r(
+ di cy (0 + dic
c1
Σ r((b) + dic c2K + d2
(99)
where (1) holds from plugging (93) into (96), (2) is true because T, 2 0, and (3) follows from the distributive property. We can rewrite this equation as
(do - -dd c))
r Rkk
+ d110''100) (2)00 tdicolo2 ddK S r())
(100)
K
(r.)) (0)) + ddc
(k) (0k) < (do - dic -11
+ dicō c2 d2
(101)
K
K
By plugging in
((1) (()) + d (0) + d2 into (101), we get the statement of the lemma.
Plugging the following values into Lemma 15 and utilizing Lemmas and 14 and the learning rates a, B, y from Theorem 2, we get the result in Lemma 9 in precise terms:
(() ((K)
(k)
El|lx( (k)
(k (1-1) 191
T A (k)
1
Lf
4aT22
C2 ymB + 4a0h + 200
Co = 4a
2
2
do 3/
d
-1 1 C2.
d2 PoG
(102)
3gg
23

---

Published as a conference paper at ICLR 2023
Next, recall that our step sizes were
K K-3/5
4K 4K-2/5-
2K -3/5
min
Mg
3 min
(v.
Fgg - ,,7 1/2
(103)
T T." 4G
where min( + Fs
1 Note that the choice of motivated by the conditions of Lemma
5. First, observe that 1 do > 0 is true because chose 2/ Now, observe that 3 I6G
Finally, will next show that do dco)(c))-1 > 0 and c1d1 (do) 1 > O, completing the set of conditions in Lemma By direct algebraic manipulation, we have
do - (1o)
3gg
2
MgB
1)
G2
-L
>
3g
Bg
G 4T2a
2
41
1
4
3g.3
&L
1g
64L
a
1gB
2
BPg
11
2
3
2
1g
64L
3g
ng
pg
(104)
m3
256 1
2
4
Similarly, we also have
co - (do)
Ls
4L'a
2
G2
Lj
81
MaB
1
G
4a
2
MgB
4a
2
MgB
2
2
1
Lg
8L2a a
PgB
G
2
L
16L.
16L2aG2
y
d²
4a
By
4a
2
11
32
16L2aG22
12
1
1
1
(105)
80
256 1
0
8a
16a
16a
24

---

Published as a conference paper at ICLR 2023
Now, we can bound the optimality of y by bounding the maximum difference A
go
(k)
A
(A) U:
K
Sc m
(1)
rr00) + +d10
+ d2 dicj
S (do dico (
+ dico 02 d2
(106)
K
r + +d00 + d2 + d1 co - (2))
2
+ dicj 02 + +
(107)
2g.
K
(0)
9,00
+ 08000
+
20 Maf
((0)
Bgg
K
2G
+ y
(8a) (1y + 4abj + 2007) + 8808
(108)
Mg 3
4A(0) (0)
(0)
4,9
%
6403a
64Gg
4,0
PgB
Py
(ny + 4ab2 + 2007)
+
(109)
Pg
4A (0)
k()0)
1
4 3
64G y
64G?
B
na 132
Pg
KK
1gg
12
32K
11
+ 20
a²
4a²
+
23
(110)
m
32
Pg
(6)
A
(0)
/
Gy/L
2/nB,GGgLL. K2/5
(bS + (L3)
160
160g/12
+
K3/5
KT/5
K4/5
K2/6
K2/5
(111)
Here, (1) follows directly from plugging T from (102) into Lemma 4, (2) follows from (104), (3) comes from plugging in the rest of (102), (4) is direct algebraic manipulation, (5) separates the step sizes and n, K factors from the rest of the constants, and (6) applies the definition of the step sizes. This gives us the OVynk- 225) bound Lemma
A.8 PROOF OF LEMMA 10
We present a precise form of Lemma 10 here:
Lemma 16. For any À An. under Assumptions 1, 2, and 3, assume that the iterates (x)(8) y3
(K)
(k)) Vk} generated by MORBIT, then we have
E
F(z()) (k), A) - ((a((),2(44)
K
(112)
1
G
Brynk- 3/3 Beynk-211
max A
(R)
2
K
ic(n
91
Proof. Recall that we defined
F(x, = Σ Aef((z, = Σ xe((z))
(113)
1=1
25

---

Published as a conference paper at ICLR 2023
For fixed realization of x01),.. 20(h), (1)
Mi .. Mi , [n],, we have
(8)
((x(k) A) F(x(11) (())
(A4 - 1 (k) (f((z(1)
13 (3 (2((0))
(114)
@
Σ (A: - ) .1) (z(00) fi(((1), y2 ( + fi(x!!)) y/ ((+
(115)
(3)
((> )(0)) [f( 1/ ((t1)) (n( 2/
, ]
n
Σ (Ac 11 (f( 22 (0)
fe(2 (000)). (
(116)
(4)
((^- \(11 .1((t) + Σ (1: - 11 (Fi M1 (x( (((
(117)
i=1
Y - 11 + 7 111
(k)
112 - ||1 1(8) - (2
>
27
n
+ Σ (A: 1 (00 (z)( 2/6 (x:
- So (x) ,, 1
(118)
(6) a| 112 7 (0 - a|l 11
27
+ Σ (A )x) (z 31 (z(11) fo fz(()))), .3/ (((tt11)))
(119)
1=1
a - 12 + ( 112
a| 21+
+G 0a - (2)(()) - V ((+t1)).
2
(120)
where (1) comes from the definition of F. (2) follows from adding and subtracting (A: X (k) (e(((()) Wi ((+t)) terms, (3) follows from splitting the preceding sum and writing the first term as
a dot product, (4) follows from definition of h (k) (5) uses (a, b) ||| +
2 I|al, (6) follows
from the update (1++1) proja, (2(x) - and the projection property, and (7) follows from Lipschitzness of f.
>
Therefore, applying the telescoping sum by adding the preceding inequality over it K. and taking expectation, we get:
(x) X) - F(x(1), -) (1+)))
K
(1In 12 E(IAA - 122
K + +G,
(A=
X (8)
A
(121)
24
k=1
K
BlIn 12)
1 +GG+ G Σ
max
A (1) 9.
(122)
kEjn]
(3)
nKB²
K
+ + GG
Σ
max (A) (K) VE(n]
(123)
2
(1)
V2
K
(Bivnk215 + BBeTRk3/5 + G
max A (k) yi
(124)
2
iE(m]
26

---

Published as a conference paper at ICLR 2023
where (I) follows directly from (120) and the telescoping sum, (2) follows from À E An and ||
xa) 12 & 2, (3) follows from | (x) 112 & nB², and (4) follows from selecting ? B, Brynk375
v2
Therefore, obtain
1 F(x!0), 1) () -FF(((()))
F(o(0),)))))
E K
(125)
BeVnK-33f BBV 2/5) + +G)0(VVk-211)
A.9 PROOF OF LEMMA 11
We state Lemma 11 here in precise terms:
Lemma 17. Under Assumptions 1, 2, and 3 with the iterates (x(h) Mi .iE [m], (09, Vk) generated by MORBiT, I, then we have
(x))(((())) - -1
(126)
K
1=1
4
1111(00)
2
+
bÂ+
- Me +
park
Me + p
21.
31.²2²
(
Σ
max A
(07 + 363)
- +p
K
sein
2
Proof. Recall that we defined the Moreau envelope and proximal map as follows:
11/0(2) a min 2 fir= 2112 Σ x4(x),,
argmin ((2) A arg min f1lxr 11r-20 2112+ + Σ (1(2). (127)
Therefore, we have
11 $ (1)(2((+1)
+ (llzstt1) (ll(b((((+))))))
((x((+1)))12
(128)
2
Σ (L((((((11)) (110(1111 - (11(11))12
(129)
8
Σ (L 6 (1120011) (1) +(()) - (1)(1))
(130)
Σ XL(((((())) (11x(0-1) - =11) / (1lo(0) - =(22
(((ttt))) oltt) -(0()))(()))))
(131)
where (1) is by definition of the proximal map, (2) comes from the optimality of the Moreau envelope, (3) is by adding and subtracting (k), and (4) is from expanding out ||a+bbl into [o||2+||b/'" 2 b).
27

---

Published as a conference paper at ICLR 2023
Next, from the optimality condition of the update (x(+1)) -projs., (vr(k) atls we have
(x+1)) 2(x(1)) x((+1) (k) + +ah ) 0
(132)
(x(++1 - (b) + x((h) (z((()))
x()) + ath! S
(133)
2 (x(1)) ( (h) )) - =(((()))
s --11o( - x 12 a(h?
(1) , - xb)) - an(h (1), ,,()) a(xx(1)))
(134)
@ (k+() (1,), (b) {((())))
S < (k++)) - x (6) 1| - pa(h (b), , l((+1) - x(0)) - pac(h (k), x0( =((x(1)))
(135)
(k+ (b) (b) {((((²))
S < (++1) - (80 1| + plah (k) x(() ,(0t1)) + partt , - x1)
(136)
(k+() - (1)) xI (()(()))
$ -pplz(((+1) - 1/2 + 0
- |lah 12+ + (|0(0)))))( =((+
2
+ pa (b) (x(())) - 2
(137)
@ p(z(k+1) - (1)), x1() a(x(()))
pa²
2 1 (+) 12
|11
1|2 + pa (h _(x (b))
(138)
2
where (1) is by adding and subtracting x (k), (2) is from distributing the inner product ((x((++1))
(k)) + (x)(k) - (k)( (x((+11) - 22()) + aha (k)), (3) is from multiplying both sides by p. (4) is
from simple algebra, (5) is from rewriting (a, b) S +|o||*+|00] and (6) is from combining terms.
2
Therefore, substituting (138) into (131), we get
11 (()(((()11)
S Σ (1) + (1x(() (11000)111 + pa²
||
- (8)) (139)
2
+ pa²
= (1)()()))))
|1 + phx( ;(((())) - (8))
(140)
2
where the second equality is by definition of the Moreau envelope. Now, we bound the last term in (140):
(1000 (0x - (8))
(x)) (8), (b(()) (t y (K+1), ,110))
+ (x)) 11) V F.F(1()) (g (00) (zF(((t), 2111)
+ (i())())) (k), V F (x.()) (=((0),(4))) 2((1))
(141)
(2) (k)) - 11k) BB) + (z)) - (()), V. (x(()) y (16++1,,(11)) V Fzo(()) (10( (10))
(B)
(b(x())) (k) V.F( ((o)))
(142)
where (1) follows from adding and subtracting V, (((()) ((++1), 2(0)), V. (xx((())) (()) terms and
(2) is from splitting the inner product and applying h - (z(((((), y 1k+ (()) = B, from
equation 20 in Assumption 3. To bound (A) and (B), we simply apply (a, b) &||aa12 + 1 |8|2 to
28

---

Published as a conference paper at ICLR 2023
both inner products:
(4) ( - x1(), BBx - 1 [1 - =2/ -0%
(143)
(B) ( (5(2())) - -2(h) VoF(x(1) (1++2), 1((t)),(6b) (zF((((), (1(((())11))))
1 [lV yl \10) - VaF(z(t)) {()))1 + (1(((o(() - 112
(144)
We proceed to bound (C'). First, from weak convexity of pez, we have that for all i E [m].
(((((())) 2 (x))) + (Ve(x(")), ((x((k)) - (k)) Me
[| - I. 112
(145)
2
Taking x; times the ith of these equations, get
Σ (ch(((((()))
(e -(()))
Me [,(
10.2
2
(146)
By definition of the Moreau envelope, we also have
Σ ((
(x)(((())))))) + P [((()())) - 1p.
(147)
Adding (146) and (147), have
14
P [|((z((a))
11 2 ((.
- (K))
(148)
If we
we let c= =mcte in (143) and (144), we can rewrite (142) as
(h) (k) ((x((1)) - (
$ 5 [1((x((0) - 1110 -ba IV F(x y 1(tt1), v ((x( ((2
V
ue + Tette1laoa (1(((o((0) - - 12
(149)
2
$ (1)())(t) - 1/ + 1 1/ + I
Σ (z)))
- -1 1
c
Me P (1(((x(()
- 112
(150)
2 29
2L2
m Σ
Aellyr (z) - y/ 11
-Me +
(x) -
5 -4e +
P
ue + P
4
(151)
Taking the full expectation Fi a {u;
(0) (0)
(K) (.)7, we have
g( (k) (x)) [())))]
2
6
+
2L
E
Auller (2((0) 3i 1atti112
3e +
[10(z11)) - -x (8)
Me
+
-M + p
4
(152)
2
82 +
21
A (2 w/ (22t)111))
14 +
(1:((z - (x)
1e + P
11 + P
(153)
2
62
212
max E Elllz (z((k) - vi 1(1111133) (E(m)
-14 1t +
11(((2(1)) 20(
14 + P
- Pt + P
4
(154)
2
62
2L
max A (t)
- +
(2(()(())
(155)
-Pe +
Pa + p
[E(m) 31
4
29

---

Published as a conference paper at ICLR 2023
Therefore, rewriting everything into (146) and taking the full expectation over Fs have (recall that the definition of A %. includes an expectation):
(A)
E| E[01/0 1 (xx(+))))
(1)
pa²
E 1 (x()) +
1 [13] paBB(ih p ,olB((h1),,x((00) xx(1/)]
(156)
2
2
E 1/101/1000 (x)))
2pa
b2
2T.² pa
max
0a01-12 -10 p)
Bl1 (((()(x(()) ,x(1) ||
-Me +
-He +
iE|m 4/1
4
pa 2
2
of
+
3b
+ 3L
max A sE|n|
U1
(1)
(157)
3
E 1 1/p (z
(1))
Me + b3 +
2pa
22 pac
Mr + +
pa²
max A (k) (E(m]
2
palue p) (1)) - (100)
pa²
+
(0)
+3bb),
(158)
2
where (1) is copy of (140) and (2) is from (158), plugging in ||/g
(1) ||2 from Lemma 1, and doing the same expectation calculation from (152) to (155). Finally, (3) combining terms via algebra. Summing up from h = 0, 1, K 1, we get the following bound:
K
(k))
11
K
4
(211p(00)
2 by
jeg +
pack
1Me + p
+
2L2
+
3L
Σ max
a (k) m
a
()) + +bb))
(159)
114 +
2
x (r)
2
4
①
1110(00)
2
1ie + P
pacK
tee + p
K-2
212
1
max
A M
We + P
K
EE(n]
Vik 2/0
0f + 360 + 2L
+
max (k) kE(m) 91
(160)
2
K
K-3/5+
B GENERALIZATION BOUNDS
In addition to convergence, we also show the generalization abilities of the bilevel optimizer. The theorem in this section is inspired by Collins et al. (2020, Theorem 4), but our results hold for the fully general min-max multi-objective BLO setup while Collins et al. (2020) study a min-max multi-objective single-level problem. Assume that for learning task i, observe m, batches of train/test data, Di and j E [me]. Assume that each train and test batch has K and J input- output pairs, respectively, so that sets D DV, are drawn from a common distribution D,. Also, let yf (z; D: arg miny, I. E gi(x, Vi, Di be the value of y; that minimizes the empirical inner loss on some dataset DJ. For outer and inner objectives fi, Di. consider the following
30

---

Published as a conference paper at ICLR 2023
function class Fi, where E X' is the set of optimization parameters introduced in equation 2 and (D; D)) are any train and test datasets sampled from D;:
Fi { (I, 17 (1; Dj), DT) E x
We use and g to denote the empirical function values evaluated the points in Dj and Dj, and so the empirical Rademacher complexity of F; on my samples D; - {(Dig Dig m ~ (D,)
my (z.00
R (F)) ED, Bey
sup ZEX m
where E,S are Rademacher random variables (=1/2 with equal probability). The empirical loss for fixed samples Dij and Di.
1 m mh E (z (i,);
A())
Similarly, define (x & ED, [x)]. First, from classical generalization results such as Shalev- Shwartz & Ben-David (2014, Theorem 26.5) Mohri al. (2018, Theorem 3.3), directly conclude the following proposition, which bounds the true loss of the classifier as a function of the empirical loss.
Proposition Assume the regularity assumptions considered in Appendix 3, specifically that the function is Br-bounded. Then, with probability at least 1 - 8, we have
log 1/6
F1(x) S F((x) + 2R (Fi) + Bu
(161)
2m,
This extends to the following in a straightforward manner, providing guarantee for the worst-case generalization for any learning task
Proposition 2. Assume the regularity assumptions given in Section 5, specifically that the function E is By-bounded. Then, with probubility at least 8, we have
logn/5
max Fx(x) < max (2) + 2R (F) + Bey
(162)
*E|n]
iE[n]
2m
Here we assume that F; F and m; m for all E [n] and hence Ri (Fi) Rm (F) for all i [n].
Next, we proceed to bound the generalization on unseen tasks. Consider a new task with distribution Dn-1= a;D; for some E A,, meaning that the distribution of the new task is anywhere in the convex hull of the distribution of the old tasks. We make this assumption because if the new task very dissimilar to the existing tasks, there is no reason expect good generalization in the first place. We then show the following proposition:
Proposition 3. For all X. with probability at least 1 - 8, we have
log(n/0)
Fx+1(x) max p. Fxx + PEAn
Σ a Ri, (Fi)
a, B,
(163)
2m;
Notice that while Proposition holds true for all xx the tighest upper bound found when r minimizes maXpEA, pt F.(r), which is precisely when x = x=,, the optimal solution to problem we study in equation 2. This highlights another advantage of our formulation over TTSA: when we use the solution obtained by the single averaged objective E, fi in equation 1, min-ooy: We will have a looser upper bound for Fx+1 (wiin..yg) compared to Ex+ (x*)) ), showing that our formulation gives tighter robust (or worst case) generalization guarantees and this behaviour has been demonstrated empirically in section 4.
31

---

Published as a conference paper at ICLR 2023
Proof. (of Proposition 3) First, by definition, for all 2, we have that
(2) Bepra
1. ) DrvD
Si
xr arg min
Σ (x, Vvs Digs) Dia
(164)
ag
Din n+
So arg min
Σ
(a, Vi. Dig Dieg
(165)
i=1
Σ (x).))).
(166)
Therefore, for all x, using a union bound over have
< i & n, we get that with probability at least - n28,
log 1/0
F2+ Σ (z)) F.(x + Σ
R (F)
E B,
B.
(167)
2m4
Letting no' in (167), have
log(n/0)
P (a) S max R(() + - R (F) + Σ
BB
(168)
rEA,
2m,
2m1
Therefore, plugging in xx gives us
log(n/8)
(x+1(**) min max pt P.()))))))))))))))())) + +2 _ #EX PEd,
0g R (F) + Σ a Be
(169)
2m,
C IMPLEMENTATION DETAILS AND COMPUTE RESOURCES
We perform our experiments in Python 3.7.10 and Py PyToch 1.8.1 with Intel(R) Core(TM) i5 U5-8265U CPU @ 1.60GHz. The code is available bilevel. 1. For our empirical evaluation, we first select o; B that give good performance/converrgnnee for the min-avg problem (the baseline). We do hyperparameter search to choose these parame- ters, specifically the initial learning rates. Then we fix a, 3 and only select 7 that provides good convergence for the min-max problem (our proposed scheme).
repository https //gi thub com mi n ma of/
Hypergradient computation. We would like to note that the analysis does not require the actual hypergradient but rather stochastic estimate with bounded bias. The standard Hessian inverse approximation using the Neumann series (Agarwal et al., 2017: Ghadimi & Wang, 2018; Hong al., 2020) is one way of computing this estimate (as we have discussed in section 3.2 preceding Corollary 1). Since we are considering single-loop algorithm, even straightforward iterative differentiation (Ji et al., 2021) can provide an sufficiently useful estimate the hypergradient. We utilize this for our empirical evaluations.
C. SINUSOID REGRESSION TASK
We consider the sinusoid regression experiment (Finn et al., 2017), multi-task representation learning problem where each task T, is a regression problem y = 6, (xx = a, aisin - Φ.). We uniformly sample the amplitude E 5j. frequency and phase d; [0. for each task. We use n training tasks and testing tasks, with "easy tasks" (a, E [0.1, 1.05)) and one "hard tasks" (a, E [4.95, 5]) for each set. We easy and hard tasks following the setup in (Collins et al., 2020). During training, for each task i, the learner is given samples (æ, y). x E [-5, 5). The goal is to learn function approximating ( as best possible in the mean squared error sense.
32

---

Published as a conference paper at ICLR 2023
neural network divided into two pieces, i.e., embedding network and a task-specific network. The embedding network f R - R!0 consists of two hidden ReLU layers of size 80 and final fully connected layer of size 10. Each task-specifie network : R200 + R, E [m] is a one-layer linear layer. Therefore, the loss on an input a E R and y (x) for task i is (=:() y)² and the true loss of the network with parameters f,
As described in Section 2,
4ff:g g)) (( - y) The embedding network is as follows:
Input (R)
Linear FC Layer (output in R*0)
ReLU
Linear FC Layer (output in R0))
ReLU
Linear FC Layer (output in !T")
Training: At each iteration, first perform the inner loop optimization step (meta-training) by sampling shots from each of the tasks in order to update each of the task-specific network weights. We use just inner loop step. Pseudocode for the inner loop is shown below in PyTorch-style:
for task id in task list:
xs, ys hample_battee (task_id, n_shots) embedding for in range inner):
head optimizers [tasssstskiid id zero zero-grade () total_ = get_loss (task_id, xs, ys)
total backward () head.optimizerss task_ step))
For each outer loop optimization, of the tasks. We then take an outer-loop step. optimizing the embedding weights using the results of the meta-validation batch. The meta-validation batch is sampled in the exact same way as the meta-training batch shown above.
meta- validation batch again containing 10 shots from each
Regularization: First, as in (Ji al., 2020), add weight regualarization during inner loop training of the form e Ewew [|w]] where W denotes the set of weight parameters, where = 0.01. In PyTorch, this expressed as
12_reg constant 0.01 for p in heads task ]. parameters (): 12_reg += p. norm (2) total_loss += 12_reg * 12_reg-constant
Next, for the inner À updates, we add a regularization term to the overall loss, -c >i (A; - ) where EΛ =3, which pulls the As closer uniform. In PyTorch, the update expressed
task gradient task losses [i] reg gr-ddrddent -m lambda (lambdas [ [aak]] 1/n) lambdas task] +=
reg update) gamma
Parameters: For the Task-Robust version of the algorithm, we use a = 0.007, 8 = 0,005, y = 0.003. For the standard version of the algorithm,
0.007, 3 0.011, 0.003.
Loss curves: To approximate the true loss for measurement purposes, we use 100 equally-spaced samples from ( After each iteration, we calculated the maximum loss among all the tasks. In 1c, we show the minimum of these maximum losses up until each epoch.
Results with more tasks: Finally, show another figure similar Figure 1, but with 20 training tasks and 20 test tasks. It can be observed that both the task-robust training loss and the task-robust testing loss greatly outperform their respective standard losses.
33

---

Published as a conference paper at ICLR 2023
Maximum MSE Loss with Task-Roben vs. Standard Training
Standard Train Loss
Standard Test Loss
Task-Robust Train Loss
log MSF Loss Until Epoch)
Task- Rak-RRouut Test Loss
1-
2500
5000
7500
10000
12500
15000
17500
20000
Epoch
Figure 3: Comparison of standard (min-avg) training and robust (min-max) training using 20 tasks
C.2 NONLINEAR REPRESENTATION LEARNING
We consider binary classification tasks generated from the FashionMNIST data where we select "easy" tasks (lowest log loss ~ 0.3 from independent training) and 2 "hard" tasks (lowest loss ~ 0.45 from independent training). We learn shared representation network that maps the 784 dimensional (vectorized 28 x28 images) to a 100 dimensional space. Each tasks then learns binary classifier on top of this representation. The task specific objective yi for task corresponds to the cross-entropy loss on the training set, while the upper level objective fa corresponds to the loss of the v (ar) with the learned representation z on a validation set. also maintain a heldout test set which we use to evaluate the generalization of the learned representation and per-task models.
For our data, we had IR 784x100 and y E 1000x We used step sizes a 0.01, 8 = 0.01, and y 0.3. We used batch sizes of 8 and 128 compute 9, for each inner step and f. for each outer iteration, respectively. In addition, we included & Ezregulaarizaton of y with regularization penalty 0.0005. We used vanilla SGD with learning rate scheduler MeeduceLROnPltteee 100 outer iterations, with patience of 10. Each optimization was executed for 10000 outer iterations. The results generated by aggregation over with 10 different seeds.
invoked every
C.3 HYPERPARAMETER OPTIMIZATION
In this application, iterations. We use a batch size of for both the inner and outer steps for each i E [16] for the initial experiment in figure 2a. The optimizer was vanilla SGD with learning rate scheduler (ReduceLROnPlatcau) invoked every 100 outer iterations, with patience of 30. The results are generated by aggregating over 10 runs with different seeds. For the other HPO experiments, the number of tasks n and the batch sizes are discussed in the main text.
learning rates a 0.0001, 0.001, 7 0.001 and 20000 outer
D ADDITIONAL TECHNICAL DETAILS
Here we provide further discussion on some technical aspects of the problem we are studying in this paper.
34

---

Published as a conference paper at ICLR 2023
D.1 WEAK-CONVEXITT AND NON-CONVEXITY
We consider weakly convex UL objective, and here we discuss how it is related to non-convexity. Weak convexity captures class of non-convex problems. Weakly convex functions are not convex note difference in the following definitions (also in Appendix A.I, Assumptions and 2). For any convex function 7, there exists μ 2 0 such that, for any 2, x' (x + x')
T(x') (2) + (VaT(z), x) ul| =11²
(170)
whereas, for weakly-convex function there exists > 0 such that, for any x, x' (x # x')
<(x') 2 x<x) + (V=<(1), x' - z) - 11|2 - x11²
(171)
Note the "--" for a weakly-convex <<)) instead of the "+,"" " for convex 7(.) in the third term on the right hand side of the above two inequalities. So x(-) clearly not convex. Moreover, note that the || ==|² term on the right-hand side of the inequality for the weakly-convex function is strictly positive, implying that, for large enough v, the inequality will be true for any function. We provide convergence results which depend on the coefficient of weak-convexity (for our UL function in question, it is denoted as pę), with slower rates for larger coefficients.
D.2 COMPARISON WITH HU ET AL. (2022)
Hu al. (2022) may appear similar our work glance, but we would like to clarify that the differences are nontrivial as we are solving a different problem. We address this briefly in section (Closely related and Concurrent Work), but will elaborate further here make the distinction clearer.
At high level, the problem in Hu al. (2022) not multi-objective: the authors explicitly call it multi-block. They are still solving the single-objective min-max problem min,, max, f(z, a). Hence the problem setup in Hu et al. (2022) cannot solve standard bilevel learning applications such as representation learning and HPO; they choose AUC maximization as their motivating example instead.
Now, we explain what may be source of confusion: why it seems like they are solving multi- objective problem. Hu et al. (2022) start with the min-max problem min, max, fax, or) with strong concavity in a, such as in AUC maximization. Then they make bilevel to min, max, f f(r, y (x), subject to (z) = arg miny (Ix y, a) by splitting the x variable and then further splitting into multi- block to min, max,. Σ, (z, y ax). Here, each fi is strongly concave in og. This is different problem setup than ours and does not include our problem formulation.
Therefore, the crucial difference is this: they study single-objective problem min, max, f(z a). and we consider the robust multi-objective bilevel problem min, max, fi y (z)). Their approach seems similar at first glance because they solving the single-objective problem in bilevel, multi-block way, but their problem class does not encompass the multi-objective one we consider.
D.3 IMPROVING THE SAMPLE COMPLEXITY OF MORBiT
There is a potential room for improvement in the sample complexity of MORBiT. In the n = 1 case, algorithm builds off of TTSA (Hong al., 2020) with O(1/225) complexity. The only existing work in the n = 1 case with a better sample complexity in a single-loop constrained UL case the extremely recent STABLE (Chen et al., 2022b), achieving ((1/²).. STABLE, has much more complex LL update than TTSA using variance reduction techniques. We are optimistic that more complex algorithms like STABLE can be extended the robust multi-objective bilevel optimization setting with improved sample complexity.
D.4
WHY ROBUST min max INSTEAD OF PARETO
OPTIMIZATION?
Bilevel optimization problems are ubiquitous in machine learning applications such as representation learning and hyperparameter optimization, which is difficult to formulate a single-level problem. We consider standard stochastic bilevel problems such as these, formulating a natural robust multi- objective version of these problems inspired by the benefits of robust multi-objective learning
35

---

Published as a conference paper at ICLR 2023
highlighted in Mehta et al. (2012) and Collins et al. (2020). These papers consider the robust multi- objective view but do not study stochastic bilevel learning problems, which we do. Existing bilevel optimization problems, however, are all single-objective rather than multi-objective.
The advantages of taking single objective problems and formulating them as robust multi Multi-objecctve the literature cited in section 2 (Min-max
problems have been highlighted in various works Robust Optimization in Machine Learning). To summarize, the main advantage is that we can get guarantees on the worst-case performance instead of the usual average case performance (see for example our generalization guarantees in Appendix B). If we just summed the objectives and solved a single-objective problem, we would only be able to establish guarantees for the average-case performance: maybe we would find a solution that is good for most tasks, but might do extremely poorly on some. Moreover, at the lower level (I.I.) problem, there are different objectives for the learners as the individual problem structures and data distributions are different, again forming a natural multi-objective optimization (MOO) problem.
Much like our motivating existing literature on robust multi-objective learning, we focus on single robust solution instead of set of Pareto optimal solutions since, in various applications, we finally need select a single solution, and the robust (ınin max) solution provides stronger worst-case guarantees than any Pareto-optimal solution, which is our main motivation.
Pareto frontiers can be very useful and informative, potentially allowing us to understand the tradeoff between the multiple objectives. However, we would like to note that there are various forms of solutions in multi-objective optimization. There Pareto optimal solutions, but also "possibly optimal" solutions (Wilson et al., 2015), convex coverage set of solutions (Yang et al., 2019), and min max robust solution (that we consider). The appropriate form of solution(s) would depend on the application, and we are focusing on min max applications, motivated by existing work such as Mehta et al. (2012) and Collins et al. (2020), since a min max solution can be shown to have good generalization guarantees (as we have also shown in Appendix B).
Furthermore, while the Pareto frontier can be more informative and the Pareto curves better demon- strate tradeoff between the objectives, it is important to note that, this curve is mostly intuitive with obvious tradeoffs for n = 2 objectives. With n > objectives, the Pareto frontier cannot even be visualized, and one has to resort to pairwise comparisons, making it hard to reason about the tradeoffs between objectives even for moderately high n since we will have to consider n² such comparisons (for example n ~ O(10). Therefore, given a Pareto front of solutions, it is not clear which of the Pareto optimal solutions we should select.
One advantage of the min max formulation (equation 2) is that it tries to seek single solution instead of a set of solutions. This allows us to use the solution for a new related problem (like for a new related task in representation learning application or hyperparameter optimization application in Franceschi et al. (2018)), we can use the robust min max solution - we select the robust solution for the shared UL variable (the representation network the hyperparameter configuration). With a Pareto front, it is not clear which solution to pick for a new task since we would have a set of solutions, without the knowledge of which one would be useful for new task/objective.
Furthermore, while a solution on the Pareto frontier implies that there is no other solution that "dominates" it, to the best of our knowledge, there is no guarantee that some solution on the obtained Pareto frontier achieves the optimal value for the robust min max objective max; min, f((x) unless the Pareto frontier is completely dense, which is never the case. Multi-objective optimizers can return a set of solutions on the Pareto frontier, but even uniformly covering the Pareto frontier requires the size of the solution set to grow exponentially in the number of objectives n.
Finally, for nonconvex objective functions, the Pareto frontier refers to the Pareto stationarity rather than Pareto optimality. Our considered first-order stationarity condition defined on the weighted average of the objective value, while the classical Pareto stationarity (please see Fernando et al. (2023, equation (2)) and references therein) is measured on the size of the weighted average of the gradients. The weighting vector in both of these two notations is optimized over simplex. Therefore, the stationarity condition of our proposed min max formulation can be considered as one variant of Pareto stationarity for nonconvex problems.
36


Reviews:
[{"review_id": "hUIiYF7hbf8", "reviewer": "Reviewer_UZe1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies robust solution for bilevel optimization under min-max multi-objective form. The main idea is build the problem by a min-max upper-level objective function and using weighted sum of weakly convex functions li instead of a single weakly convex function. To solve such a problem, the SGD(projected SGD) to iterate $y$, $x$ and $\\lambda$. Some applications such as robust representation learning and robust hyperparameter selection  are provided to evaluate the performance. ", "summary_of_the_review": "Overall, I feel this work does a good job in designing a novel from of bilevel problem. It provided a precious frame for further and finer solution. I tend to accept this work but with a conservative score given the above questions. However, I am open to change my mind based on other reviewer’s comments and the authors’ feedbacks.  ", "strength_and_weaknesses": "Pros:\n1. This paper builds a frame for such a specific bilevel optimization problem, which has not been considered before. The proposed algorithm is simple and easy to implement with a single-loop structure. \n\n2. In the reformulated problem in (3), the upper level problem turns to optimization a maximization problem over a simplex. This new component requires some new treatments.  Comprehensive convergence rate analysis is provided to justify the algorithmic designs. \n\n3. The algorithm itself is easy to understand and apply. Experiments seem to support the proposed algorithm well. \n\nCons:\n1. Excluding the nicely designed problem itself, the algorithm is not that surprising given existing single-loop algorithms such as Hong et al., 2020 in bilevel optimization. It looks like applying SGD (projected-SGD) to a specific problem. I wonder how the authors implement the hypergradient in practice? Is there any instructions here? \n\n2. For the experiments, e.g., in Fig. 5, why does the variance look so large? \n\n3. The analysis seems to follow from TTSA (Hong et al., 2020), and yields a similar rate of $K^{-2/5}$. It seems that some other SGD types of bilevel methods like stocBiO (Ji et al., 2021) have better complexity. Can it be used here for improved complexity?\n\n4. It seems to me that the complexity still has space to be improved using some strategies like variance reduction. Can the authors elaborate more on this? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "VE9PMfsSxO", "reviewer": "Reviewer_XgUo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors formulate a new min-max multi-objective bilevel optimization problem and applied it in robust ML and HPO.\nThey also proposed a new algorithm to find a solution to the proposed new problem and establish its convergence rate and computational complexity.", "summary_of_the_review": "The framework and algorithms are novel enough but the benefit of the way formulating the problem compared with the existing MOO formulation is skeptical. The strong assumption w.r.t. the leader and follower problems is another concern. Considering this, it is hard to recommend its publication in its present form.\n\n------------------------\nI have read the authors' replies. The response has not significantly changed my view of the paper, and thus I would like to keep the same score.", "strength_and_weaknesses": "Strength\n1. The authors did a thorough literature review w.r.t. robust learning and multi-task learning.\n2. The authors created a new MOO bilevel problem, proposed a solution to it, and established its convergence rate.\n\nWeaknesses\n1. The main issue with this work is that the paper is trying to create an (unnecessary) multi-objective optimization (MOO) problem and address it. It is challenging to understand the benefit of formulating a MOO problem as a bilevel one instead of summing the objectives and using a single-objective solver or generating multiple Pareto-optimal solutions. Specifically, it is worth comparing the proposed formulation with the Pareto-optimal solutions considering that the proposed one only gives one robust solution and the Pareto front gives a full set of solutions. What is the complexity difference? What is the convergence difference? Besides, the Parent solution can be generalized easily to nonconvex use cases but can the proposed algorithm be extended to the nonconvex cases? How can the proposed one be justified to be beneficial? \n2. It is also interesting to convert problem (2) to (3) by replacing the objective in the leader problem with a weighted sum of multiple objectives. Though the equivalence might be straight for the authors, it seems not the case for the majority of readers. It will be good to prove why these two formulations are equivalent.\n3. Equation (2) has a quite strong assumption that there are n different objective function pairs (fi, gi), i.e, the number of objectives in the leader problem and the follower problem is the same. This is usually not true considering that the number of leaders and the number of followers are usually not of the same magnitude.\n4. The examples given in the introduction part, e.g., safety-critical, scenario-specific, etc are very unclear. What is a concrete example whose formulation is like (2)? This formulation is also different from the existing MOO works in Sinha et al., 2015; Deb & Sinha, 2009; Ji et al., 2017.\n5. The authors assume the LL problem is unconstrained strongly convex and the UL problem is constrained weakly convex. This is again a very strong assumption and hardly satisfied in reality considering the main application is in RL and HPO.\n6. For this iterative algorithm, how to choose the three learning rates alpha, beta and gamma are critical. The authors only give the chosen numbers at the end of the appendix but did not give any intuition on how they are selected especially how this matches what Theorem 1 suggested regarding these three learning rates.\n7. In Theorem 1, the bias norm is bounded by a sequence bk. How to verify that b_k^2 <= alpha in practice?\n8. The writing needs to be polished.\na. At the end of the first paragraph, there is typo 'Similar'.\nb. In Table 1, it is not fair to list the problem solution and solving strategy together. Specifically, the single loop should not be listed with other columns.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "P3ZvNOlylzQ", "reviewer": "Reviewer_fwCP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work, the authors consider a class of bilevel optimization problems in which he objective function of the upper level is a maximum over $n$ different functions for each one a separate lower level problem is given. In particular, they propose a single-loop algorithm in which a gradient descent update is performed for both upper and lower level decision variables per iteration. The authors then provide simple complexities of finding a stationary point of the upper level problem and optimal solutions of the lower level ones.", "summary_of_the_review": "The paper needs more discussion and comparison with the existing works.", "strength_and_weaknesses": "The problem under consideration is interesting, however, a similar problem already studied by (Hu et al., 2022) with minor differences. Moreover, their sample complexity is in the order of ${\\cal O}(1/\\epsilon^2)$ which is better than the one in this work ${\\cal O}(1/\\epsilon^{2.5})$ when $n=1$. Thus, the authors need to discuss why their result is worse for the case of $n=1$. Moreover, a table is required to compare the proposed complexity bounds with the existing ones.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster