M IN -M AX M ULTI - OBJECTIVE B ILEVEL O PTIMIZATION WITH A PPLICATIONS IN R OBUST M ACHINE L EARNING
A BSTRACT
We consider a generic min-max multi-objective bilevel optimization problem with applications in robust machine learning such as representation learning and hyper - parameter 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 (cid:101) 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
Our code is at https://github com/ minimario
/MORBiT
I NTRODUCTION
We begin by examining the classic bilevel optimization (BLO) problem as follows:
[Equation]
subject to
[Equation]
where f : X × Y → R is the upper-level ( UL ) objective function and g : X × Y → R is the lower-level ( LL ) objective function.
X and Y , respectively, denote the domains for the UL and LL optimization variables x and y , incorporating any respective constraints.
Equation 1 is called BLO because the UL objective f depends on both x and the solution y ⋆ ( x ) of the LL objective g .
BLO has been well-studied in the optimization literature (Bard, 2013; Dempe, 2002).
Recently, BLO has found various applications in modern machine learning (Liu et al., 2021; Chen et al., 2022a), such as hyperparameter optimization (HPO) (Franceschi et al., 2018), reinforcement learning (Hong et al., 2020) (RL), multi-task learning (MTL) and representation learning (Arora et al., 2020).
In this work, we focus on a robust generalization of equation 1 to the multi-objective setting, where there are n different objective function pairs ( f i , g i ) .
Let [ n ] ≜ { 1 , 2 , · · · , n } and f i : X
× Y i
→ R , g i : X × Y i
→ R denote the i th UL and LL objectives respectively.
We study the following problem:
[Equation]
Here, the optimization variable x is shared across all objectives f
i , g i , i ∈ [ n ] , while the variables
y
i , i ∈ [ n ] are only involved in their corresponding objectives f
i , g i .
The goal is to find a robust solution x ∈ X such that the worst-case across all objectives is minimized.
This is a generic problem which reduces to equation 1 if we have a single objective pair, that is n = 1 .
Such a robust optimization problem is useful in various applications, and especially necessary in any safety-critical ones.
For example, in decision optimization, the different objectives ( f i , g i ) can correspond to different “scenarios” (such as plans for different scenarios), with x being the shared decision variable and y i ’s being scenario-specific decision variables.
The goal of equation 2 is to find the robust shared decision x which provides robust performance across all the n considered scenarios, so that such a robust assignment of decision variables will generalize well on other scenarios.
In machine 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 2 would generate a shared model of the world – the UL variable x – such that the worst-case utility, max i f i ( x, y ⋆i ( x )) , of the agent-specific optimal action – the LL variable y ⋆i
( x ) – 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-specific variables y i ∈ Y i to come from different domains,
that is, Y i ̸
=
Y j , i, j ∈ [ n ] ; as stated in equation 2, this implies that the dimensionality d y i for the per-objective y
i 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 x .
For example, in multi-agent RL, different agents can have different action spaces because they need to operate in different mediums (land, water, air, etc).
The main contributions of this work are highlighted as follows:
▶ (New algorithm design)
We present a single loop M ulti-
O bjective R obust Bi level T wo-timescale optimization algorithm, MORBiT , which uses (i) SGD for the unconstrained strongly convex LL problem, and (ii) projected SGD for the constrained weakly convex UL problem.
▶ (Theoretical convergence guarantees)
We demonstrate that, under standard smoothness and regularity conditions, MORBiT with n objectives converges to a (cid:101) O ( n 1 / 2 K − 5 / 2 ) -stationary point with K iterations, matching the best convergence rate for single-loop single-objective ( n = 1 ) BLO algorithms with the constrained UL problem while using vanilla SGD for the LL problem, and providing a sublinear n 1 / 2 -dependence on the number of objective pairs n .
▶ ( 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.
P ROBLEM AND R ELATED W ORK
We first discuss the different aspects of the robust multi-objective BLO problem with constrained
UL in equation 2.
While BLO is used in machine learning (Liu et al., 2021; Chen et al., 2022a), multi-objective BLO has not received much attention.
In multi-task learning (MTL), the optimization problem is a multi-objective problem in nature, but is usually solved by summing the objectives and using a single-objective solver, that is, optimizing the objective (cid:80) i f i .
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 a more complex multi-objective optimization problem that replaces the objective (cid:80) i f
i with the objective max i f i .
For SGD-based algorithms solving BLO, one critical aspect is whether the algorithm is single-loop (a single SGD step for both x and y in each iteration) or double-loop (the LL variable y taking multiple
SGD steps between each step of the UL variable 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 LL iterations.
Double-loop algorithms are not applicable when the (stochastic) gradients of the LL and UL problems are only provided sequentially, such as in logistics, motion planning and RL problems.
Hence, we develop and analyse a single-loop algorithm.
A final aspect of BLO is the constrained UL problem .
When the UL variable x corresponds to some decision variable in a decision optimization problem or a hyperparameter in HPO, we must consider a constrained form, x ∈ X ⊂ R d x .
To capture a 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-objective BLO and robust optimization, especially in the context of machine learning.
Table 1 provides a snapshot of the properties of the problems and algorithms (with rigorous convergence analysis) studied in recent machine learning literature.
[Table]
Lately, many new algorithms have been proposed to solve the singleobjective BLO problem in equation 1.
BSA (Ghadimi & Wang, 2018) proposed the first double-loop 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 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 2 gives flexibility for inherently multi-objective problems in a robust manner to obtain stronger guarantees, ensuring convergence of each individual objective, rather than the average objective.
Multi-Objective BLO There has been a limited number of works analyzing multi-objective BLO schemes (Sinha et al., 2015; Deb & Sinha, 2009; Ji et al., 2017).
All of these works analyze the multi-objective BLO problem from a game-theoretic point of view, using a vector-valued objective with the notion of Pareto optimality.
In contrast, we are the first to study the multi-objective BLO problem from a traditional optimization perspective in terms of convergence properties and consider a min-max robust version of the multi-objective 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.
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 DRL, Duchi & Namkoong (2018) and Shalev-Shwartz & Wexler (2016) showed that a min-max loss improves generalization due to variance regularization.
HiBSA (Lu 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).
Note that the classic MAML (Finn et al., 2017) can be formulated as a BLO problem (Rajeswaran et al., 2019), and the precise problem analysed in Collins et al.
(2020) is a single-level one as well.
In fact, we consider the bilevel form of the TR-MAML problem as one of our applications for empirical evaluation.
In ARL, the minimum is over the loss and the maximum is 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 x max y f
( x, y ) with a single-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 is inspired by TTSA and can be viewed as a robust multi-objective version.
Beyond this advancement, our contribution also lies in the convergence analysis of MORBiT , which significantly diverges from that of TTSA .
After our MORBiT was developed and released, STABLE (Chen et al., 2022b) was recently presented as an improvement of TTSA , and we wish to explore similar improvements to MORBiT in future work.
A very recent work (Hu et al., 2022) studies a problem that appears to be quite similar to equation 2, with common elements such as 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 a min-max single level problem to a min-max BLO, and min-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 2 is not a generalization of their problem – both our work and theirs are considering different setups with high-level commonalities.
Table 1 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 .
A LGORITHM AND A NALYSIS
In this section, we propose a simple single-loop algorithm MORBiT to solve equation 2, and establish a 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.
In the sequel, we will always use the subscript i ∈ [ n ] to denote the objective index and the superscript ( k ) to denote the iteration index, with x ( k ) denoting the k th iterate of the shared variable x ∈ X ⊆ R d x and y ( k ) i denoting the k th iterate of the i th -objective-specific variable y i ∈
R d yi .
We will also use the shorthand y to denote all the per-objective variables [ y 1 , y 2 , . . .
, y n ] , with y ( k ) denoting the k th iterate of all the per-objective variables [ y ( k ) 1 , y ( k ) 2 , . .
, y ( k ) n ] .
Given our assumption that the LL objectives g i are strongly convex, we define y ⋆i ( x ) ≜ arg min y i ∈
R dyi g i ( x, y ) , and use the shorthand ℓ i ( x ) ≜
f i ( x, y ⋆i ( x )) .
3. MORBiT A LGORITHM
We begin with a standard reformulation of robust min-max problems (Duchi et al., 2008).
We can rewrite the non-smooth min-max problem in equation 2 as
[Equation]
subject to
[Equation]
where ∆ n ∈ R n + is the n -simplex defined as ∆ n := { λ ∈ R n + : λ i ≥ 0 , ∀ i ∈ [ n ] , (cid:80) i ∈ [ n ] λ i = 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 y i , i ∈ [ n ] is the straightforward ∇ y
i g i ( x, y i )
and we denote h i as its stochastic estimate, with h as the shorthand for the per-objective stochastic gradient estimates [ h 1 , h 2 , . . .
, h n ] .
The gradient for the x -update is more involved because of the hierarchical structure of the BLO problem.
Then, we consider the following weighted objectives utilizing the simplex variable λ ∈ ∆ n to define the necessary gradients:
[Equation]
Note that F ( x, λ ) is the UL objective in equation [Equation] and the UL gradients can be defined as:
[Equation]
where ∇ x ℓ i ( x ) for any i ∈ [ n ] can be defined as follows utilizing the strong convexity of the LL problem and implicit gradients (Gould et al., 2016):
[Equation]
Note that in general, y ⋆i ( x ) cannot be computed exactly.
Following Ghadimi & Wang (2018), we use an approximation of ∇ x ℓ i ( x ) as a surrogate, denoted by ∇ x f i ( x, y i ) , by replacing y ⋆i ( x ) in equation 6 with any y i ∈
R d yi as follows:
[Equation]
Consequently, we define our approximate gradients for the UL variables x (and λ ) as:
[Equation]
We denote the (possibly biased) stochastic estimates of ∇ x F ( x, y , λ ) as h x and ∇ λ F ( x, y , λ ) as h λ .
1 for k = 1 , 2 , · · · , K do y
( k +1)
← y ( k ) − β h ( k ) x
( k +1) ← proj X ( x ( k ) − αh ( k ) x ) λ ( k +1)
← proj ∆ n ( λ ( k ) + γh ( k ) λ ) 5 end 6 Sample τ ∼ U ( { 1 , · · · , K } ) 7 return ¯ x ← x ( τ ) , ¯ y i ← y ( τ − 1) i , ¯ λ ← λ ( τ )
Given the gradients and their stochastic estimates, we present our single-loop algorithm MORBiT in algorithm 1, where we utilize learning rates α, β, γ > 0 for the UL variable x , LL variables y i , i ∈ [ n ] and the simplex variable λ respectively.
The algorithm tracks three sets of variables x ( k ) , y ( k ) =
[ y ( k ) 1 , y ( k ) 2 , . .
, y ( k ) n ] and λ ( k ) through a total of K iterations.
The per-iterate gradient estimates h ( k ) of the LL variables y is defined as the collection of the per-objective gradient estimate h ( k ) i evaluated at ( x ( k ) , y ( k ) i ) for all i ∈ [ n ] .
The gradient estimates h ( k ) x
and h ( k ) λ of the UL variables x and λ are evaluated at ( x ( k ) , y ( k +1) , λ ( k ) ) .
We perform a standard gradient descent update for the objective specific variables
y i , i ∈ [ n ] from y ( k ) i to y ( k +1)
i .
For the shared UL variable x , we perform a projected gradient descent to satisfy the
UL constraints, where proj X ( · ) denotes the projection operation onto the constrained set X .
We update the simplex variable λ via projected gradient ascent , where we project the variable back onto the n -simplex after a gradient ascent step with proj ∆ n ( · ) .
Given the learning rates ( α, β, γ ) ,
MORBiT is quite straightforward in terms of implementation.
When n = 1 , the problem reduces to single-objective BLO, λ ( k ) = 1 , and MORBiT reduces to TTSA (Hong et al., 2020).
3. A NALYSIS
Given the single-loop MORBiT , we establish conditions under which MORBiT has finite-horizon convergence.
The coupling of the stochastic errors due to the sampling process makes the convergence analysis of this three-sequence-based algorithm much more challenging than existing BLO algorithms.
Assumptions .
We summarize the following typical assumptions (detailed in Appendix A.1) for all objective pairs ( f i , g i ) , i ∈ [ n ] .
Focusing on the smoothness and regularity properties of the objectives, we assume that (i) the LL objective g i is strongly convex in y i , twice-differentiable, and has sufficiently smooth first and second order gradients (Assumption 2 in Appendix A.1), (ii) the UL objective f
i has sufficiently smooth first order gradients, and (iii) the function ℓ i ( x ) ≜
f i ( x, y ⋆i ( x )) is weakly convex, bounded and has bounded first-order gradients (Assumption 1 in Appendix A.1).
Regarding the quality of the gradient estimates h ( k ) i , i ∈ [ n ] , h ( k ) x
and h ( k ) λ , we assume that, for all k > 0 , (i) h ( k ) i is an unbiased estimate with bounded variance, (ii) h ( k ) λ is an unbiased estimate, and (iii) h ( k ) x has bounded variance, and can be a biased estimate of the ∇ x F ( x ( k ) , y ( k +1) , λ ( k ) ) term defined in equation 8, but the bias norm at iteration k is bounded by b k ≥ 0 , with { b k , k ≥ 0 } forming a non-increasing sequence.
These gradient estimate quality assumptions are detailed in Assumptionin Appendix A.1.
While the assumptions on h ( k ) i and h ( k ) λ are standard (Hong et al., 2020; Lu et al., 2022), the assumption on h ( k ) x 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 et al., 2017; Ghadimi & Wang, 2018;
Hong et al., 2020).
Optimality and Stationarity of Solutions .
To quantify the convergence properties of the solutions ¯ x, ¯ y i , i ∈ [ n ] , ¯ λ generated by MORBiT , we use the following optimality properties of the optimal solutions x ⋆ , y ⋆i , i ∈ [ n ] , λ ⋆ of the problem in equation 3.
(i)
The per-objective optimal LL variable y ⋆ i =
y ⋆ i
( x ⋆ ) =
arg min
y i ∈
R dy
i g i ( x ⋆ , y i ) ; (ii)
The optimal simplex variable λ
⋆ : F ( x ⋆ , λ ⋆ ) = max λ ∈ ∆ n F ( x ⋆ , λ ) .
Given the constrained UL , the first-order stationarity condition is satisfied if ⟨∇ x F ( x ⋆ , λ ⋆ ) , x − x ⋆ ⟩ ≥ 0 ∀ x ∈ X .
(iii)
For establishing near-stationarity of UL variable x , the proximal map ˆ x ( z ) ∈ X , defined below,
[Equation]
is a fixed constant.
[Equation]
) is employed (Davis & Drusvyatskiy, 2018; Hong et al., 2020) to quantify the convergence for a constrained variable x in the stochastic setting.
If ∥ ˆ x ( x ( k ) )
− x ( k ) ∥ 2 is small, then, near-stationarity of x ( k ) is achieved at iteration k .
Therefore, we need to bound ∥
ˆ x (¯ x ) − ¯ x ∥ to guarantee the convergence of the UL solution ¯ x returned by MORBiT .
Given the convergence of the UL ¯ x , we also need to bound ∥
¯ y i − y ⋆i
(¯ x ) ∥ 2 for each i ∈ [ n ] simultaneously to quantify the convergence of the LL variables.
Finally, the convergence of ¯ λ requires us to bound the difference between F (¯ x, ¯ λ ) and max λ ∈ ∆ n F (¯ x, λ ) .
Theoretical Convergence Rate .
Now, we are ready to state our 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 − 3 / 5 ) , β
=
O
( K − 2 / 5 ) and γ = O ( n − 1 / 2 K − 3 / 5 ) .
Then, if b 2 k ≤ α , the solutions ¯ x, ¯ y i , i ∈ [ n ] , ¯ λ generated by algorithm 1 satisfy:
[Equation]
with expectation over the stochastic gradient estimates and the random index τ (algorithm 1, line 6).
This result establishes the (cid:101) O ( n 1 / 2 K − 2 / 5 ) -stationarity achieved by K iterations of MORBiT for both 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 ∈ X =
R d x , the definition of the proximal map (equation 9) implies that E ∥∇ x F (¯ x, ¯ λ ) ∥ ≤ (cid:101) O ( n 1 / 2 K − 2 / 5 ) , providing the convergence of ¯ x to a (cid:101) O ( n 1 / 2 K − 2 / 5 )
-stationary point if the UL problem is unconstrained.
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 a more general proximal map in equation 9 involving a weighted sum of weakly convex functions ℓ i instead of a single weakly convex function in TTSA , requiring new construction of potential functions for establishing the convergence of the UL variable
¯ x in equation 10a.
Secondly, even though TTSA provides a convergence rate for a single LL variable (equivalent to bounding E [ ∥ ¯ y i − y ⋆i
(¯ x ) ∥ 2 ] for a single i ∈ [ n ] ), we provide a 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 E [max i ∈ [ n ] ∥ ¯ y
i − y ⋆i
(¯ x ) ∥ 2 ] .
This is especially challenging since a bounded E
[ ∥ ¯ y i − y ⋆i
(¯ x ) ∥ ] for each i ∈ [ n ] does not directly imply a bounded E [max i ∈ [ n ] ∥ ¯ y
i − y ⋆i
(¯ x ) ∥ 2 ] ; in fact this can be generally unbounded.
Finally, to satisfy the requirements of the min-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 O ( K − µ ) -stationary 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 ϵ -optimal solution is O ( ϵ − 1 / µ ) .
The number of stochastic gradient oracle queries required is directly related to the conditions in the gradient estimate quality assumptions (Assumption 3 in Appendix A.1 in our case).
While the conditions on the per-iterate gradient estimates h ( k ) i (for the per-objective LL variables) and h ( k ) λ (for the simplex variable λ ) both only require O (1) stochastic gradient oracle queries from each of the n objective pairs in each iteration, the condition b 2 k ≤ α on the non-increasing squared norm of the per-iterate bias of the gradient estimate h ( k ) x (for the UL variable) require O (log K ) stochastic gradient oracle queries for each of the n objective pairs leveraging the HIA sampling techniques in Ghadimi & Wang (2018) and Hong et al.
(2020) using the Neumann series (Agarwal et al., 2017).
This gives us the following sample complexity bound for MORBiT :
Corollary 1.
Under the conditions in Theorem 1, MORBiT converges to ϵ -(near)-stationarity with
O ( n 5 / 4 ϵ − 5 / 2 log( 1 / ϵ )) queries to the stochastic gradient oracle for each of the n objective pairs.
Proof Sketch of Theorem 1.
We now give a proof sketch of our main theorem, with constant terms abstracted away with O notation.
In order to show equation 10a of Theorem 1 (convergence of x ), we will derive a descent lemma comparing successive iterates x ( k ) and x ( k +1) .
Descent lemmas often contain a quadratic term ∥ x ( k +1)
− x ( k ) ∥ 2 , so it is natural that we must bound ∥ h ( k )
x
∥ 2 .
In Lemma 1, we bound the expected squared norm of the stochastic gradient estimate:
Lemma 1.
Under our regularity assumptions, E [ ∥ h ( k )
[Equation]
∥
[Equation]
O
[Equation]
Turning to equation 10b of Theorem 1, we use a descent relation on y ( k ) i − y
⋆i
( x ( k − 1) ) .
While ideally we would obtain a descent relation purely involving y ( k ) i terms themselves, the intricate coupling of the x and y
i terms result in an extra x ( k − 1) − x ( k ) term.
The resulting relation is shown in Lemma 2.
Here, we have that c 1 = 1 − µ g β 2 , c 2 = 2 µ g β − 1 , where µ g β < 1 .
Lemma 2.
E [ ∥ y
[Equation]
i
[Equation]
From this lemma, intuitively, we know that E [ ∥ y ( k ) i − y
⋆i
( x ( k − 1) )
∥ 2 ] is decreasing as k increases, as long as the E [ ∥ x ( k − 1) − x ( k ) ∥ 2 ]
’s are not too large .
Therefore, it is important to have another descent relation that upper bounds this quantity, which we do next in Lemma 3.
The lemma naturally involves the objective F ( x ( k ) , λ ( k ) ) , which will telescope.
Here, we have c 3 = 14 α − L f 2 , c 4 = 4 αL 2 .
As k → ∞ , α → 0 , and c 3 is positive.
Lemma 3.
Let L ( k ) ≜
E [ F ( x ( k ) , λ ( k ) )]
= n (cid:80)
[Equation]
E [ ℓ i ( x ( k ) )] .
Then, the L ( k ) satisfies:
[Equation]
Following the intuition previously described, we then use Lemmas 2 and 3 to show that
the
E [ ∥ x ( k − 1) − x ( k ) ∥ 2 ] terms are small enough and that the y ( k ) i iterates converge in Lemma 4.
Lemma
[Equation]
Finally, we can use the convergence of y i to prove convergence of λ and x .
Theorem 1 then follows.
A more detailed proof plan can be found in Appendix A.3.
E XPERIMENTAL R ESULTS
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 MORBiT on these applications.
We use PyTorch (Paszke et al., 2019), and implementation details are in Appendix C.
Representation Learning.
In this setup, each objective pair corresponds to a learning “task” i ∈ [ n ] , with its own training/validation dataset pair D t i , D v i .
We consider a shared representation network ϕ x with ReLU nonlinearity (making the UL problem weakly convex) parameterized with x ∈ R d x and a per-task linear model w y i parameterized with
y i ∈
R d yi .
Here we consider an unconstrained UL .
Using L ( f, D ) to denote the loss of a model f on data D , we consider the problem in equation 2 with
[Equation]
with ρ > 0 as a regularization penalty (ensuring that the LL problem is strongly convex).
We first consider a multi-task setup with n = 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 a single-objective bilevel problem by minimizing 1 / n (cid:80) i f
i ; we call this min-avg .
We theoretically show that solving the min-max multi-objective bilevel problem in equation 2 provides a tighter generalization guarantee in
Proposition 2 in Appendix B.
Here we demonstrate the same in figure 1 – 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-max problem.
The results indicate solving the minmax 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 also show the convergence of the UL variable for the above minmax problem in the form of the trajectory of (stochastic) gradient norm ∥∇ x ∥ 2 in figure 2, comparing it to the theoretical (cid:101)
O
( K − 2 / 5 ) rate.
The plot shows that the empirical trajectory of the gradient norm closely tracks the theoretical rate .
We also consider a bilevel extension of the robust meta-learning application (Collins et al., 2020) for a sinusoid regression task (a commonly used meta-learning application introduced by Finn et al.
Here the goal of solving the problem in equation 2 with the objectives defined in equationwould 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.
We theoretically show that that solving the min-max multiobjective bilevel problem in equation 2 also provides a tighter generalization guarantee for the unseen tasks in Proposition 3 in
Appendix B. Note that these results are similar in spirit to those of
Mehta et al.
(2012) and Collins et al.
but our results are the first for a general bilevel setup.
The results in figure 3 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 1.
Hyperparameter Optimization.
In this setup, each objective pair again corresponds to a learning “task” i ∈
[ n ] , each with its own d dimensional training/validation dataset pair D t i , D v i .
We consider a shared hyperparameter optimization problem for a kernel logistic regression (Zhu & Hastie, 2001) with K random Fourier features (RFFs) (Rahimi & Recht, 2007), where x = { x ρ ∈ R 2 K + , x σ ∈ R d + } are the regularization penalty and the bandwidth hyperparameters respectively, with ϕ x σ denoting the RFF.
The per-task linear model w y i on top of the RFF are parameterized with y i ∈ R 2 K .
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 a model f on a dataset D , we consider the problem in equation 2 with
[Equation]
(12) where ⊙ denotes the elementwise vector multiplication, and we consider a weighted regression penalty.
3 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 2 to that of the min-avg.
, cos( W ( x σ ⊙ p )) ⊤ ] ⊤ ∈ R 2 K , where
W ∈ R K × d is a random normal matrix and the sin( · ) and cos( · ) are applied elementwise.
The weighted regression penalty mitigates bias especially in the high-dimensional learning setting (Candes et al., 2008; Gasso et al., 2009; ˇ Sehi ´ c et al., 2022) which is common when using RFFs.
The results in figure 4 indicate that the solution of equation 2 provides a robust solution x (hyperparameters), 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 solving the min-max problem in equation 2 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 ∈ { 4 , 16 , 64 } , increasing n with a factor of 4 (implying theoretically that the convergence should slow down with a factor of 2) to study how the convergence matches the √ n -dependence in our theoretical result.
In this case, we consider the trajectory of the (stochastic) gradient norm ∥∇ x ∥ 2 (as in figure 2).
The results in figure 5 display such a behaviour – for a fixed K (outer iterations), as the number of tasks is increased 4-fold, the gradient norm approximately increases 2-fold (note the log 2 -scale 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 a batch size of 8 for both the UL and LL 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, in this problem, each of the 16 learning tasks (and hence, objective pairs) has a 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 1 and 4, we only show the generalization loss (dropping the validation loss) in figure 6.
The results indicate that increasing the batch size improves the stability and reduces the variance of the overall generalization.
However, the convergence follows a similar trend for all batch sizes, and converges to a very similar level of generalization, supporting the (cid:101) O (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 in 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 .
C ONCLUDING R EMARKS
Motivated by the desiderata of robustness in bilevel learning applications, we study a new minmax multi-objective BLO framework (equation 2) that provides full flexibility and generality.
We propose MORBiT (algorithm 1), 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 a single-loop algorithm is not feasible (e.g., federated learning).
Finally, we also wish to develop algorithms for large n (the number of objective pairs) or even n → ∞ where MORBiT is not computationally feasible.
R EPRODUCIBILITY S TATEMENT
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
Appendix C
A C ONVERGENCE A NALYSIS OF MORBiT A. A SSUMPTIONS
First, we begin by listing the assumptions we make:
Assumption 1 (Regularity of the outer functions) .
For all i ∈ [ n ] , assume that outer functions f i ( x, y ) and ℓ i ( x ) =
f i ( x, y ⋆i ( x )) satisfy the following properties:
▶ For any x ∈ X , f i ( x, · ) is Lipschitz (w.r
t. y ) with constant G f > 0 .
▶
For any x ∈ X , ∇ x f i ( x, · ) and ∇ y i f i ( x, · ) are Lipschitz continuous (w.r.t.
y i ) with constants
L f x > 0 and L f y > 0 .
▶
For any y i ∈ R d yi , ∇ y i f i ( · , y i ) is Lipschitz continuous (w.r.t.
x ) with constant ¯ L f y > 0 .
▶
For any x ∈ X , y i ∈ R d yi , we have ∥∇ y
i f i ( x, y i )
∥ ≤ C fy for C fy > 0 .
▶
The function ℓ i ( · ) is µ ℓ weakly convex (in x ), so that for all v, w ∈ X , [Equation] i
[Equation]
▶ For all x ∈ X , ∥ ℓ i ( x ) ∥ ≤ B ℓ for B ℓ > 0 .
▶
For all x ∈ X , ∥∇ ℓ i ( x ) ∥ ≤ C ℓ , for C ℓ > 0 .
Assumption 2 (Regularity of the inner functions) .
Assume that inner functions
g i ( x, y i ) , ∀ i ∈ [ n ] satisfy:
▶ For any x ∈ X and y i ∈ R d yi , g i ( x, y i ) is twice continuously differentiable in ( x, y i ) .
▶
For any x ∈ X , ∇ y
i g i ( x, · ) is Lipschitz continuous (w.r.t.
y i ) with constant L g .
▶
For any x ∈ X , g i ( x, · ) is µ g -strongly convex in y i , so that for all v, w ∈ R d yi ,
[Equation]
▶ For any x ∈ X , ∇ 2 xy
i g i ( x, · ) and ∇ 2 y
i g i ( x, · ) are Lipschitz continuous (w
r
t. y i ) with constants L gxy > 0 and L gyy > 0 , respectively.
▶
For any x ∈ X and y i ∈ R d yi , we have ∥∇ 2 xy
i g i ( x, y i )
∥ ≤ C gxy for some C gxy > 0 .
▶
For any y i ∈ R d yi , ∇ 2 xy
i g i ( · , y i ) and ∇ 2 y
i g i ( · , y i ) are Lipschitz continuous (w.r
t. x ) with constants ¯ L gxy > 0 and ¯ L gyy > 0 , respectively.
From these assumptions, we can show a few additional regularity-type conditions.
Since these conditions can also be found in Ghadimi & Wang (2018, Lemma 2.2) and Hong et al.
Lemma 2), we state these results without proof.
Lemma 5 (Corollary of Assumptions) .
Under Assumptions 1 and 2 stated above, for all x, x 1 , x 2 ∈
X ⊆
R d x , y ∈ R d yi , i ∈ [ n ] , we have
[Equation]
where we define
[Equation]
Assumption 3 (Quality of stochastic gradient estimates) .
For any iteration k > 0 and all i ∈ [ n ] , the gradient estimates h ( k ) i for the LL variable y
i satisfy the following for some σ g >
(Hong et al., 2020; Lu et al., 2022):
[Equation]
[Equation]
For any iteration k >
[Equation]
the gradient estimate h ( k ) λ for the simplex variable λ
satisfies:
[Equation]
[Equation]
For any k ≥ 0 and a σ f >
[Equation]
we assume that there exists a non-increasing sequence { b k } k ≥ 0 such that
[Equation]
[Equation]
A. M AIN T HEOREM AND R EMARKS
We now state our main theorem, in full.
Recall our notation from equation 9: for a fixed constant ρ > 0 , we defined the proximal map to be [Equation]
[Equation]
We also define the Moreau envelope as
[Equation]
In addition, we use the notation ∆ ( k ) y i ≜ E [ ∥ y ( k ) i − y
⋆i
( x ( k − 1) )
∥ 2 ] and i ( k ) ≜ arg max i ∈ [ n ] ∆ ( k ) y i .
Finally, we define ˜ σ 2 f = σ 2 f + 3 C 2 ℓ (see Lemma 12).
Now, we are 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
[Equation]
where ν = min ( µ g / L 2 g (1+ σ 2 g ) , 1 / µ g ) , Algorithm 1 produces ¯ x, ¯ λ, ¯ y i , i ∈ [ n ] satisfying:
[Equation] [Equation]
where
[Equation]
Connection of TTSA (Hong et al., 2020)
As we generalize Hong et al.
(2020), our proof follows a similar structure.
In particular, Lemma 6 is a generalization of Hong et al.
(2020, Equation (14)), and our Lemma 7 combines Hong et al.
(2020, Lemma 3 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.
λ , so there is no analogue in Hong et al.
However, it borrows techniques from Collins et al.
(2020, Theorem 1).
We also discuss the convergence rate.
Note that A in equation 30 is dominated by the fourth term, √ n/K 2 / 5 , so it is clear that equation 27 and equation 28 converge at a rate of O ( √ nK − 2 / 5 ) .
We give special attention to equation 29).
Apart from the 8 b 20 / ( − µ ℓ + ρ ) 2 term, we see that the RHS of equation 29 converges at a rate of O ( √ nK − 2 / 5 ) .
To understand the convergence of this term, we turn to (20) from Assumption 3.
As discussed in section 3.2, b k can be made arbitrarily small by running more iterations of the subroutine for estimating h ( k )
x
(for example utilizing the HIA sampling scheme (Agarwal et al., 2017; Ghadimi & Wang, 2018; Hong et al., 2020)).
Therefore, as long as we run enough iterations ( O (log K ) for HIA) such that b 2 k ≤ α , (29) will also converge at a rate of
O ( √ nK − 2 / 5 ) . [Equation]
Overall Roadmap:
In what follows, we give a proof sketch of our main theorem, with constant terms abstracted away with O notation.
c 1 , c 2 , c 3 , c 4 are positive constants depending on L f (defined in Lemma 5), µ g (the LL objective convexity defined in Assumption 2) and the learning rates α, β (in algorithm 1).
In order to show equation 10a of Theorem 1 (the convergence of x ), we will derive a descent lemma comparing successive iterates x ( k ) and x ( k +1) .
Descent lemmas often contain a quadratic term ∥ x ( k +1)
− x ( k ) ∥ 2 , so it is natural that we will have to bound ∥
h ( k ) x
∥ 2 .
As such, in Lemma 6, we bound the averaged squared norm of the stochastic gradient estimate, E [ ∥ h ( k ) x
∥ 2 ] :
Lemma 6.
Under our regularity assumptions, the average squared norm of h ( k ) x can be bounded as follows, where the expectation is over the filtration F ′ i ≜
[Equation]
y
[Equation]
[Equation]
Turning to equation 10a of Theorem 1, we’ll use a descent relation on y ( k ) i − y
⋆i
( x ( k − 1) ) .
While ideally we would obtain a descent relation purely involving y (( k ))
i terms themselves, the intricate coupling of the x and y terms result in an extra x ( k − 1) − x ( k ) term.
The resulting relation is shown in Lemma 7.
Here, we have that c 1 = 1 − µ g β 2 , c
2 = 2 µ g β − 1 , where µ g β < 1 .
Lemma 7.
The distance between the algorithm’s iterates y ( k ) i and the true inner optimum y ⋆i
( x ( k ) ) satisfies the following descent equation,
[Equation]
From this lemma, intuitively, we know that E [ ∥ y ( k ) i − y
⋆i
( x ( k − 1) )
∥ 2 ] is decreasing as k increases, as long as the E [ ∥ x ( k − 1) − x ( k ) ∥ 2 ]
’s are not too large .
Therefore, it is important to have another descent relation that upper bounds this quantity, which we do next in Lemma 8.
The lemma naturally involves the objective L ( k )
=
F ( x ( k ) , λ ( k ) ) , which will telescope.
Here, we have
[Equation]
As k → ∞ , α → 0 , and c 3 is positive.
Lemma 8.
Let L ( k ) ≜
E [ F ( x ( k ) , λ ( k ) )]
= n (cid:80)
[Equation]
E [ ℓ i ( x ( k ) )] .
Then, the L ( k ) satisfies the descent equation
[Equation]
Following the intuition previously described, we then use Lemmas 7 and 8 to show that
the
E [ ∥ x ( k − 1) − x ( k ) ∥ 2 ] terms are small enough and that the y ( k ) i iterates converge:
Lemma 9 (Informal, see Appendix A.7 for precise statement) .
[Equation]
Lemma 10 then leverages the convergence of y i , i ∈ [ n ] to bound the convergence of λ , and Lemmashows the bound on x .
By plugging in our step-sizes into Lemmas 9, 10, and 11, Theorem 1 directly follows.
Lemma 10.
For any λ ∈ ∆ n , the iterates of Algorithm 1 satisfy
[Equation]
Lemma 11.
The iterates of Algorithm 1 satisfy
[Equation]
A. P ROOF OF L EMMA 1 (L EMMA 6)
Stating Lemma 1 more precisely:
Lemma 12.
Under Assumptions 1, 2 and 3, the average squared norm of h ( k ) x can be bounded as follows, where ˜ σ 2 f = σ 2 f + 3 C 2 ℓ :
[Equation]
where the expectation is over the filtration F ′ i
[Equation]
y
[Equation]
[Equation]
x
[Equation]
x [Equation] [Equation] Note that here the expectation is over F ′ i ≜ { y (0) i , x (0) , · · · , y ( k ) i , x ( k ) , y ( k +1) i } , so no expectation is needed in the last term.
Proof.
We can derive the following:
[Equation]
(1) is true because
[Equation]
(2) follows from (20), (3) follows from (21), (4) follows from definition of ∇ x , and (5) follows from
Young’s inequality, ∥ a + b ∥ 2 ≤ 3 ∥ a ∥ 2 + 32 ∥ b ∥ 2 .
Next, we bound the last term in (42).
We start by using the fact that
[Equation]
where (1) follows from ∥ a + b ∥ 2 ≤ 2 ∥ a ∥ 2 + 2 ∥ b ∥ 2 , and (2) follows from (cid:13)(cid:13)(cid:13) n (cid:80)
[Equation]
Next, we bound the first term in (47).
From Lemma
[Equation]
we have
[Equation]
Therefore
[Equation]
we can obtain
[Equation]
where (1) comes from plugging (48) into (47) and (47) into (42), (2) comes from the definition of ˜ σ 2 f using ∥∇ ℓ i ( x ( k ) )
∥ 2 ≤ C 2 ℓ and λ ( k ) ∈ ∆ n .
A. P ROOF OF L EMMA 2 (L EMMA 7)
We state the precise version of Lemma 2 here:
Lemma 13.
Under Assumptions 1, 2, and [Equation] when β 2 (1 + σ 2 g )
L 2 g ≤ βµ g and µ g β < 1 , the iterates y ( k )
i satisfy the descent equation:
[Equation]
Proof.
For a particular (fixed) realization of the iterates x (1) , · · · , x ( k ) , y (1) i , · · · , y ( k ) i for some i ∈ [ n ] , we have
[Equation]
where (1) follows from algebra, (2) follows from E [ h ( k ) i ]
= ∇ y
i g i ( x ( k ) , y ( k ) i ) in equation 17 in
Assumption 3, (3) is from equation 18 in Assumption 3, (4) is from ∇
y
i g i ( x ( k ) ,
y ⋆i
( x ( k ) ))
=
due to the optimality of y ⋆ i ( x ( k ) ) , and (5) is due to the L g -Lipschitz continuity of ∇
y
i g i ( x, · ) .
Next, we can bound the difference between y ( k +1)
i and y ⋆i ( x ( k ) ) as the following, where again we assume that x (1) , · · · , x ( k ) , y (1) , · · · , y ( k ) i is fixed and the expectation is over the stochasticity of the gradient estimates: [Equation]
[Equation]
where (1) is true by definition, and (2) holds by direct algebra and the unbiasedness assumption
E [ h ( k ) i ]
= ∇ y
i g i ( x ( k ) , y ( k ) i ) in equation 17 in Assumption 3, (3) is from strong convexity, β
[Equation]
(4) is from equation 56, (5) is from the assumption β 2 (1 + σ 2 g )
L 2 g ≤ βµ g , (6) is from the inequality ∥ a + b ∥ 2 ≤ (1 + 1 /c )
∥ a ∥ 2 + (1 + c )
∥ b ∥ 2 , and (7) is from the G
y -lipschitzness of y ⋆i ( · ) in Lemma 5.
Then, we choose c = µ g β 2(1 − µ g β ) , so that (1 − βµ g )(1 + c ) = 1 − βµ g / 2 and 1 + 1 /c
= 2 µ g β − 1 .
We have c > 0 because µ g β < 1 .
Plugging these expressions into (63), we get
[Equation]
[Equation]
which completes the proof.
A. P ROOF OF L EMMA 3 (L EMMA 8)
We state the precise version of Lemma 3 here:
Lemma 14.
Let L ( k ) ≜
E [ F [Equation] x ( k ) , λ ( k ) [Equation]
] = n (cid:80)
[Equation]
=
[Equation]
( k ) i E [ ℓ i ( x ( k ) )] .
Under Assumptions 1, 2 and 3, assume that the iterates { x ( k ) , y ( k ) i , i ∈ [ n ] , λ ( k ) , ∀ k } are generated by MORBiT , then, L ( k ) satisfies the descent equation [Equation]
[Equation]
Proof.
First, since ℓ i is L f -smooth, we know that for all i ∈ [ n ] ,
[Equation]
Taking λ ( k ) i times the equation for i in (66) and summing, we can get n
[Equation]
Therefore, we have
[Equation]
Next, we bound ( A ) and ( B ) respectively as follows.
First, we upper bound term ( A ) .
First, from the non-expansiveness of projections and λ ( k +1) = proj ∆ n ( λ ( k ) + γh ( k ) λ ) , we have ∥ λ ( k +1) − λ ( k ) ∥
≤
∥
γh ( k ) λ ∥ .
Since λ ( k +1) , λ ( k ) ∈ ∆ n , ∥ λ ( k +1)
− λ ( k ) ∥ ≤ √ 2 .
Therefore, we know that ∥ λ ( k +1) − λ ( k )
∥
≤ Λ ≜ min {√ 2 , γ ∥ h ( k ) λ ∥} .
Based on these facts, we can have
[Equation]
where (1) is straightforward, (2) follows from Cauchy-Schwarz, (3) follows from the update rule for λ and the fact that | ℓ i ( · ) | ≤ B ℓ from Assumption 1, (4) is from plugging in the definition of Λ , and (5) follows from γ ∥ h kλ ∥ ≤ √ 2 .
Then, we upper bound ( B ) .
First, from the non-expansiveness of projection and the update rule x ( k +1) =
proj
X
( x ( k ) − αh ( k ) x ) , we know that
[Equation]
[Equation]
[Equation]
[Equation]
Therefore
[Equation]
we can have
[Equation]
where (1) is by definition of F ( x ( k ) , λ ( k ) ) , (2) is from adding and subtracting (cid:10)(cid:0) ∇ x F ( x ( k ) , y ( k +1) , λ ( k ) )
+ B k (cid:1) , (cid:0) x ( k +1)
− x ( k ) (cid:1)(cid:11) , (3) is from adding (75) to the previous inequality, (4) is from applying the inequality ⟨ a, b ⟩ ≤ 12 c ∥ a ∥ 2 + c 2 ∥
b ∥ 2 to 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
[Equation]
Next, we work on bounding ∥∇ x F ( x ( k ) , λ ( k ) )
− ∇ x F ( x ( k ) , y ( k +1) , λ ( k ) )
− [Equation] k ∥ 2 in equation 81.
Observe that
[Equation]
where (1) comes from ∥ a + b ∥ 2 ≤ 2 ∥ a ∥ 2 + 2 ∥ b ∥ 2 , (2) comes from ∥ B k ∥ ≤ b k in Assumption 3, (3) is from expanding the definitions of ∇ x F ( · , · ) and ∇ x F ( · , · , · ) , (4) is from Lemma 5.
Therefore, plugging in equation 87 into equation 82, using c = d = 14 α , and taking expectation over the stochasticity of the gradient estimates, we get: [Equation] Now
[Equation]
observe that the LHS looks like a telescoping sum.
To make this more apparent [Equation] define
L ( k ) =
n (cid:80) i =
1 λ ( k ) i E [ ℓ i ( x ( k ) )] and ∆ ( k ) y
i
=
E [ ∥ y ( k )
i − y
⋆i
( x ( k − 1) )
∥ 2 ] .
Therefore, with the assumption that b k ≤ b 0 , we have
[Equation]
A. P ROOF OF L EMMA
We restate Lemma 9 in more general terms here:
Lemma 15.
Assume that Ω ( k ) , Θ ( k ) , Υ ( k ) i , λ i , c 0 , c 1 , c 2 , d 0 , d 1 , d 2 are real numbers such that for all 0 ≤ k ≤ K − 1 ,
[Equation]
and also for all 1 ≤ k ≤ K, 1 ≤ i ≤ N ,
[Equation]
In addition, assume that 1 − d 0 > 0 , d 0 − d 1 c 1 c − 1 0 > 0 and c 0 − c 1 d 1 d − 1 0 , and that Υ ( k ) i , Ω ( k ) ≥for all k, i ∈ [ n ] .
Then, if i (0) =
arg max i ∈ [ n ] Υ (0)
i , we have
[Equation]
Proof.
First, let i ( k ) ≜ arg max i
[Equation]
) i , so that Υ ( k ) i ( k ) =
max i ∈ [ n ] Υ ( k ) i .
Summing (90) from k = 0 , 1 , · · · , K − 1 , we get:
[Equation]
Next, we apply (91) for i
=
i ( k +1) .
Noting that 1 − d 0 > 0 and Υ ( k ) i ( k +1) ≤ Υ ( k ) i ( k ) by definition of i ( k ) , we have
[Equation]
Then summing for k = 1 to K , we get
[Equation]
Now
[Equation]
we have
[Equation]
where (1) holds from plugging (93) into (96) [Equation] (2) is true because Υ , Ω ≥
[Equation]
and (3) follows from the distributive property.
We can rewrite this equation as
[Equation]
By plugging in Υ (1) i (1) =
Υ (0) i (0) + d 1 Θ (0) + d 2 into (101), we get the statement of the lemma.
Plugging the following values into Lemma 15 and utilizing Lemmas 13 and 14 and the learning rates α, β, γ from Theorem 2, we get the result in Lemma 9 in precise terms:
[Equation]
Next, recall that our step sizes were
[Equation]
where ν = min( µ g L 2 g (1+ σ 2 g ) , 1 µ g ) .
Note that the choice of ν was motivated by the conditions of Lemma 5.
First, observe that 1 − d 0 > 0 is true because we chose β < 2 /µ g .
Now, observe that αβ ≤ µ g 16 G y L .
Finally, will next show that d 0 − d 1 c 1 ( c 0 ) − 1 > 0 and c 0 − c 1 d 1 ( d 0 ) − 1 > 0 , completing the set of conditions in Lemma 5.
By direct algebraic manipulation, we have
[Equation]
Similarly, we also have
[Equation]
Now
[Equation]
we can bound the optimality of y by bounding the maximum difference ∆ ( k ) y i :
[Equation]
Here, (1) follows directly from plugging Υ ( k ) i 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 O ( √ nK − 2 / 5 ) bound in Lemma 9.
A. P ROOF OF L EMMA
We present a precise form of Lemma 10 here:
Lemma 16.
For any λ ∈ ∆ n , under Assumptions 1, 2, and 3, assume that the iterates { x ( k ) , y ( k ) i , i ∈ [ n ] λ ( k ) , ∀ k } generated by MORBiT, then we have
[Equation]
Proof.
Recall that we defined
F ( x, λ ) :
= n (cid:88) i = λ
i f i ( x,
y ⋆i ( x )) =
n (cid:88) i = λ i ℓ i ( x ) .
For a fixed realization of x (1) , · · · , x ( k ) , y (1) i , · · · , y ( k ) i , i ∈
[ n ] , we have
[Equation] [Equation]
where (1) comes from the definition of F , (2) follows from adding and subtracting ( λ i − λ
( k ) i )
f i ( x ( k ) , y ( k +1) i ) 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 ⟩ = ∥ a ∥ 2 +
∥
b
∥
2 −∥ a − b ∥ 2, (6) follows from the update λ ( k +1) =
proj ∆ n ( λ ( k ) − γh ( k ) λ ) and the projection property, and (7) follows from
Lipschitzness of f .
Therefore, applying the telescoping sum by adding the preceding inequality over k = 1 , 2 , . . . , K , and taking expectation, we get:
[Equation]
where (1) follows directly from (120) and the telescoping sum, (2) follows from λ ∈ ∆ n and ∥ λ − λ (1) ∥ 2 ≤ 2 , (3) follows from ∥ h ( k ) λ
∥ 2 ≤ nB 2 , and (4) follows from selecting γ = √
[Equation]
ℓ √ nK 3 /
[Equation]
Therefore, we obtain
[Equation]
A. P ROOF OF L EMMA
We state Lemma 11 here in precise terms:
Lemma 17.
Under Assumptions 1, [Equation] and 3 with the iterates { x ( k ) , y ( k ) i , i ∈ [ n ] , λ ( k ) , ∀ k } generated by MORBiT , then we have
[Equation]
Proof.
Recall that we defined the Moreau envelope and proximal map as follows:
[Equation]
Therefore
[Equation]
we have
[Equation]
where (1) is by definition of the proximal map, (2) comes from the optimality of the Moreau envelope, (3) is by adding and subtracting x ( k ) , and (4) is from expanding out ∥ a + b ∥ 2 into ∥ a ∥ 2 + ∥ b ∥ 2 +2 ⟨ a, b ⟩ .
Next, from the optimality condition of the update x ( k +1) = proj ∆ n ( x ( k ) − αh ( k ) x )
[Equation]
we have
[Equation]
where (1) is by adding and subtracting x ( k ) , (2) is from distributing the inner product ⟨ ( x ( k +1)
− x ( k ) )
+ ( x ( k ) − ˆ x
( x ( k ) )) , ( x ( k +1) − x ( k ) )
+ αh ( k ) x ⟩ , (3) is from multiplying both sides by ρ , (4) is from simple algebra, (5) is from rewriting ⟨ a, b ⟩ ≤ ∥ a ∥ 2 + ∥ b ∥ 2, and (6) is from combining terms.
Therefore, substituting (138) into (131), we get
[Equation]
where the second equality is by definition of the Moreau envelope.
Now, we bound the last term in (140):
[Equation]
[Equation]
142) where (1) follows from adding and subtracting ∇ x F ( x ( k ) , y ( k +1) , λ ( k ) ) , ∇ x F ( x ( k ) , λ ( k ) )
terms and (2) is from splitting the inner product and applying h ( k ) x − ∇ x F ( x ( k ) , y ( k +1) , λ ( k ) )
=
B k from equation 20 in Assumption 3.
To bound ( A ) and ( B ) , we simply apply ⟨ a, b ⟩ ≤ c 4 ∥ a ∥ 2 + 1 c ∥ b ∥ 2 to both inner products:
[Equation]
We proceed to bound ( C ) .
First, from weak convexity of µ ℓ , we have that for all i ∈ [ n ] ,
[Equation]
Taking λ
i times the i th of these equations
[Equation]
we get n (cid:88) i =
[Equation]
By definition of the Moreau envelope, we also have
[Equation]
Adding (146) and (147), we have
[Equation]
If we let c = − µ ℓ + ρin (143) and (144), we can rewrite (142) as ⟨ h ( k ) x ,
[Equation]
x
[Equation]
x ( k )
[Equation] [Equation]
Taking the full expectation F i ≜ { y (0) i , x (0) , · · · , y ( k ) i , x ( k ) } , we have
E [ [Equation] h ( k ) x , ˆ x ( x ( k )
[Equation]
( k )
[Equation] [Equation]
Therefore, rewriting everything into (146) and taking the full expectation over F i , we have (recall that the definition of ∆ ( k ) y i includes an expectation):
[Equation]
where (1) is a copy of (140) and (2) is from (158), plugging in ∥ h ( k ) x
∥ 2 from Lemma
[Equation]
and doing the same expectation calculation from (152) to (155).
Finally, (3) is combining terms via algebra.
Summing up from k = 0 , 1 , · · · , K − 1 , we get the following bound:
[Equation]
B G ENERALIZATION B OUNDS
In addition to convergence [Equation] we also show the generalization abilities of the bilevel optimizer.
The theorem in this section is inspired by Collins et al.
(2020, Theorem 4)
[Equation]
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 a learning task i , we observe m i batches of train/test data, D t i,j and D v i,j , j ∈ [ m i ] .
Assume that each train and test batch has K and J inputoutput pairs, respectively, so that sets D t i,j , D v i,j are drawn from a common distribution D i .
Also, let y ⋆i ( x ; D t i ) =
arg min y i
m i (cid:80) m i j ′
=
1 g i ( x, y i , D t i,j ′ ) be the value of y i that minimizes the empirical inner loss on some dataset D t i .
For outer and inner objectives f i , g i , we consider the following function class F i , where x ∈ X is the set of optimization parameters introduced in equation 2 and ( D t i , D v i ) are any train and test datasets sampled from D i :
[Equation]
We use f and g to denote the empirical function values evaluated at the points in D t
i and D v i , and so the empirical Rademacher complexity of F i on m i samples D i ≜ { ( D t i,j , D v i,j } m i j =1 ∼
( D i ) m i is
[Equation]
where ϵ j s are Rademacher random variables ( ± 1 / 2 with equal probability).
The empirical loss for fixed samples D t
i,j and D v i,j ,
[Equation]
Similarly, define F i ( x ) ≜
E D i
[ ˆ F i ( x )] .
First, from classical generalization results such as ShalevShwartz & Ben-David (2014, Theorem 26.5) or Mohri et al.
(2018, Theorem 3.3), we directly conclude the following proposition, which bounds the true loss of the classifier as a function of the empirical loss.
Proposition 1.
Assume the regularity assumptions considered in Appendix 3, specifically that the function ℓ is B ℓ -bounded.
Then, with probability at least 1 − δ , we have
[Equation]
This extends to the following in a straightforward manner, providing a guarantee for the worst-case generalization for any learning task i :
Proposition 2.
Assume the regularity assumptions given in Section 5, specifically that the function ℓ is B ℓ -bounded.
Then, with probability at least 1 − δ , we have
[Equation]
Here we assume that F i = F and m i
= m
for all i ∈ [ n ] and hence R im i ( F i )
=
R m ( F ) for all i ∈ [ n ] .
Next, we proceed to bound the generalization on unseen tasks.
Consider a new task with distribution
D n +1 = (cid:80) ni
=1 a i D i , for some a ∈ ∆ n , 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 is very dissimilar to the existing tasks, there is no reason to expect good generalization in the first place.
We then show the following proposition:
Proposition 3.
For all x ∈ X , with probability at least 1 − δ , we have
[Equation]
Notice that while Proposition 3 holds true for all x , the tighest upper bound is found when x minimizes max p ∈ ∆ n p i
ˆ F i ( x ) , 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 (cid:80) i f i in equation 1, x ⋆ min-avg , we will have a looser upper bound for F n +1 ( x ⋆ min-avg ) compared to F n +1 ( x ⋆ ) , showing that our formulation gives tighter robust (or worst case) generalization guarantees and this behaviour has been demonstrated empirically in section 4.
Proof.
(of Proposition 3)
First, by definition, for all x , we have that
[Equation]
Therefore, for all x
[Equation]
using a union bound over 1 ≤ i ≤ n
[Equation]
we get that with probability at least 1 − nδ ′ , we have
[Equation]
Letting δ = nδ ′ in (167) [Equation] we have
[Equation] [Equation]
Therefore, plugging in x ⋆ gives us
[Equation]
≤ min x ∈X max p ∈ ∆ n p i
ˆ F i ( x ) +n (cid:88)
i = a i R im i
( F ) + n (cid:88) i = a i B ℓ (cid:115) log( n/δ ) 2 m i .
C I MPLEMENTATION DETAILS AND C OMPUTE R ESOURCES
We perform our experiments in Python 3.7.10 and PyTorch 1.8.1 with Intel(R) Core(TM) i5-8265U
CPU @ 1.60GHz.
C. S INUSOID R EGRESSION T ASK
We consider the sinusoid regression experiment (Finn et al., 2017), a multi-task representation learning problem where each task T i is a regression problem y
= t i ( x ) =
a i sin( x − ϕ i ) .
We uniformly sample the amplitude a i ∈ [0 . 1 , 5] , frequency and phase ϕ i ∈ [0 , π ] for each task.
We use n = 3 training tasks and 3 testing tasks, with 2 ”easy tasks” ( a i ∈ [0 . 1 , 1 . 05]) and one ”hard tasks” ( a i ∈ [4 . 95 , 5]) for each set.
We use easy and hard tasks following the setup in (Collins et al., 2020).
During training, for each task i , the learner is given samples ( x, y ) , x ∈ [ − 5 , 5] .
The goal is to learn a function approximating t i as best as possible in the mean squared error sense.
As described in Section 2, we use a neural network divided into two pieces, i.e., an embedding network and a task-specific network.
The embedding network f : R → R 10 consists of two hidden
ReLU layers of size 80 and a final fully connected layer of size 10.
Each task-specific network g i : R 10 → R , i ∈ [ n ] is a one-layer linear layer.
Therefore, the loss on an input x ∈ R and y =
t i ( x ) for task i is ( g i ( f ( x )) − y ) 2 , and the true loss of the network with parameters f, g
i are ℓ
i ( f ; g i ) =
E ( x,y ) [( g i ( f ( x )) − y ) 2 ] .
The embedding network is as follows:
Training: At each iteration, we first perform the inner loop optimization step (meta-training) by sampling 10 shots from each of the tasks in order to update each of the task-specific network weights.
We use just 1 inner loop step.
Pseudocode for the inner loop is shown below in PyTorch-style: f o r t
a s k i d in t a s k l i s t :
Input ( R )
Linear FC Layer (output in R 80 )
ReLU
Linear FC Layer (output in R 80 )
ReLU
Linear FC Layer (output in R 10 ) xs , ys = sample batch ( t a s k i d , n s h o t s ) embedding =
embedding network ( xs )
f
o r in range ( n i n n e r ) : h e a d o p t i m i
z e r s [ t a s k i d ] .
z e r o g r a d ( ) t o t a l l o s s = g e t l o s s ( t a s k i d , xs , ys ) t o t a l l o s s .
backward ( ) h e a d o p t
i m i z e r s [ t a s k i d ] .
s t e p ( )
For each outer loop optimization, we run a meta-validation batch again containing 10 shots from each 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.
Regularization: First, as in (Ji et al., 2020), we add weight regualarization during inner loop training of the form ϵ w (cid:80) w ∈W ∥ w ∥ , where W denotes the set of weight parameters, where ϵ w = 0 .
In
PyTorch, this is expressed as
[Equation]
f [Equation] p in heads [ t a s k ] .
parameters ( ) :
[Equation]
Next, for the inner λ updates, we add a regularization term to the overall loss, − ϵ λ (cid:80) ni
=
( λ i − 1 n ) 2 , where ϵ λ = 3 , which pulls the λ s closer to uniform.
In PyTorch, the λ update is expressed as t a s k g r a d i e n t =
t a s k l o s s e s [ i ] r e g g r a d i e n t = −mu lambda * ( lambdas [ t a s k ] − 1/ n ) lambdas [ t a s k ]
+=
( t a s k u p d a t e + r e g
u p d a t e ) * gamma
Parameters: For the Task-Robust version of the algorithm, we use α = 0 .
007 , β = 0 .
005 , γ = 0 .
For the standard version of the algorithm, we use α = 0 .
007 , β = 0 .
011 , γ = 0 .
Loss curves: To approximate the true loss for measurement purposes, we use 100 equally-spaced samples from [ − 5 , 5] .
After each iteration, we calculated the maximum loss among all the tasks.
In 3, we show the minimum of these maximum losses up until each epoch.
Results with more tasks: Finally, we show another figure similar to 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.
C. N ONLINEAR R EPRESENTATION L EARNING
We consider binary classification tasks generated from the FashionMNIST data set where we select“easy” tasks (lowest log loss ∼ 0 .
3 from independent training) and 2 “hard” tasks (lowest loss ∼
.from independent training).
We learn a shared representation network that maps the 784 dimensional (vectorized 28 × 28 images) to a 100 dimensional space.
Each tasks then learns a binary classifier on top of this representation.
The task specific objective g i for task i corresponds to the cross-entropy loss on the training set, while the upper level objective f i corresponds to the loss of the y ⋆i ( x ) with the learned representation x on a validation set.
We 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 x ∈ R 784 × 100 and y ∈ R 100 × 2 .
We used step sizes α = 0 .
01 , β = 0 .
01 , and γ = 0 .
We used batch sizes of 8 and 128 to compute g i for each inner step and f i for each outer iteration, respectively.
In addition, we included ℓ 2 -regularization of y with regularization penalty 0.0005.
We used vanilla SGD with a learning rate scheduler (ReduceLROnPlateau), invoked every 100 outer iterations, with patience of 10.
Each optimization was executed for 10000 outer iterations.
The results are generated by aggregation over runs with 10 different seeds.
C. H YPERPARAMETER OPTIMIZATION
In this application, we use learning rates α = 0 . 0001 , β = 0 .
001 , γ = 0 .
001 and 20000 outer iterations.
We use a batch size of 8 for both the inner and outer steps for each i ∈ [16] for the initial experiment in figure 4.
The optimizer was vanilla SGD with a learning rate scheduler (ReduceLROnPlateau), 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.