A framework for bilevel optimization that enables stochastic and global variance reduction algorithms
Anonymous Author(s)
AffiliationAddress email
Abstract
Bilevel optimization, the problem of minimizing a value function which involvesthe arg-minimum of another function, appears in many areas of machine learning.
In a large scale empirical risk minimization setting where the number of samples ishuge, it is crucial to develop stochastic methods, which only use a few samples at atime to progress.
However, computing the gradient of the value function involvessolving a linear system, which makes it difficult to derive unbiased stochasticestimates.
To overcome this problem we introduce a novel framework, in whichthe solution of the inner problem, the solution of the linear system, and the mainvariable evolve at the same time.
These directions are written as a sum, making itstraightforward to derive unbiased estimates.
The simplicity of our approach allowsus to develop global variance reduction algorithms, where the dynamics of allvariables is subject to variance reduction.
We demonstrate that SABA, an adaptationof the celebrated SAGA algorithm in our framework, has O ( 1 T ) convergence rate,and that it achieves linear convergence under Polyak-Łojasciewicz assumption.
This is the first stochastic algorithm for bilevel optimization that verifies either ofthese properties.
Numerical experiments validate the usefulness of our method.
Introduction
Bilevel optimization is attracting more and more attention in the machine learning community thanksto its wide range of applications.
Typical examples are hyperparameters selection [5, 35, 16, 6],data augmentation [11, 39], implicit deep learning [3] or neural architecture search [31].
Bileveloptimization aims at minimizing a function whose value depends on the result of another optimizationproblem:
[Equation]
that
[Equation]
where F and G are two real valued functions defined on R p × R d .
G is called the inner function , Fis the outer function and h is the value function .
Similarly, z is the inner variable and x is the outervariable .
In most cases, the function z ∗ can only be approximated by an optimization algorithm,which makes bilevel optimization problems challenging.
Under appropriate hypotheses, the functionh is differentiable, and the chain rule and implicit function theorem give for any x ∈ R d
[Equation]
where v ∗ ( x ) ∈ R p is the solution of a linear system
[Equation]
In the light of (2) and (3), it turns out that the derivation of the gradient of h at each iteration iscumbersome because it involves two subproblems: the resolution of the inner problem to find an approximation of z ∗ ( x ) and the resolution of a linear system to find an approximation of v ∗ ( x ) .
Itmakes the practical implementation of first order methods like gradient descent for (1) challenging.
[Figure]
As is the case in many machine learning problems, we suppose in this paper that F and G are empirical means:
[Equation]
This structure suggests the use of stochastic methods tosolve (1).
For single-level problems (that is, classicaloptimization problems where one function should be minimized), using Stochastic Gradient Descent (SGD; [38, 7])and variants is natural because individual gradients arestraightforward unbiased estimators of the gradient.
Inthe bilevel framework, we want to develop algorithms thatmake progress on problem (1) by using only a few functions F j and G i at a time.
However, since ∇ h involvesthe inverse of the Hessian of G , building such stochastic algorithms is quite challenging, one of thedifficulties being that there is no straightforward unbiased estimator of ∇ h .
Still, in settings where mor n are large, where computing even a single evaluation of F or G is extremely expensive, stochasticmethods are the only scalable algorithms.
Variance reduction [25, 13, 40, 12] is a popular technique to obtain fast stochastic algorithms.
In asingle-level setting, these methods build an approximation of the gradient of the objective functionusing only stochastic gradients.
Contrary to SGD, the variance of the approximation goes to 0 as thealgorithm progresses, allowing for faster convergence.
For instance, the SAGA method [13] achieveslinear convergence if the objective function satisfies a Polyak-Łojasciewicz inequality, and O ( 1 T )convergence rate on smooth non-convex functions [37].
The extension of these methods to bileveloptimization is a natural idea to develop faster algorithms.
However, this idea is hard to implementbecause it is hard to derive unbiased estimators of ∇ h , let alone variance reduction ones.
Contributions.
We introduce a novel framework for bilevel optimization in Section 2, wherethe inner variable, the solution of the linear system (3) and the outer variable evolve jointly.
Theevolution directions are written as sums of derivatives of F j and G i , which allows us to derivesimple unbiased stochastic estimators.
In this framework, we propose SOBA, an extension of
SGD (Section 2.1), and SABA (Section 2.2), an extension of the variance reduction algorithm SAGA [13].
In Section 3 we analyse the convergence of our methods.
SOBA is shown to achieveinf t ≤ T E [ ∥∇ h ( x t ) ∥ [-
-]2 ]
= O (log( T ) T − 12 ) with decreasing step sizes.
We prove that SABA withfixed step sizes achieves 1 T (cid:80)
Tt =1 E [ ∥∇ h ( x t ) ∥ 2 ] = O ( 1 T ) .
SABA is therefore, to the best of ourknowledge, the first stochastic bilevel algorithm that matches the convergence rate of gradientdescent on h .
We also prove that SABA achieves linear convergence under the assumption thath satisfies a Polyak-Łojasciewicz inequality.
To the best of our knowledge, SABA is also the firststochastic bilevel algorithm to feature such a property.
Importantly, these rates match the rates ofthe single level counterparts of each algorithm in non-convex setting (SGD for SOBA and SAGAfor SABA).
Finally, in Section 4, we provide an extensive benchmark of many stochastic bilevelmethods on hyperparameters selection and data hyper-cleaning, and illustrate the usefulness of ourapproach.
Related work.
The bilevel optimization problem has a strong history in the optimization community,taking root in game theory [42].
Gradient-based algorithms to solve (1) can be mainly classified intwo different categories depending on how ∇ h is computed, by automatic or implicit differentiation .
Since the solution of the inner problem z ∗ ( x ) is approximated by the output of an iterative algorithm, it is possible to use automatic differentiation [43, 29] to approximate ∇ h ( x ) .
It consists indifferentiating the different steps of the inner optimization algorithm – see [4] for a review – andhas been applied successfully to several bilevel problems arising in machine learning [14, 15].
Oneof the main drawbacks of this approach is that it requires to store in memory each iterate of the inner optimization algorithm, although this problem can sometimes be overcome using invertibleoptimization algorithms [32] or truncated backpropagation [41].
The use of the implicit function theorem to obtain (2) and (3) is known as implicit differentiation [5].
While the cost of computing exactly (2) can be prohibitive for large scale problems, Pedregosa [35]showed that we can still converge to a stationary point of the problem by using approximate solutionsof the inner problem and linear system (3), if the approximation error goes to 0 sufficiently quickly.
The complexity of approximate implicit differentiation has been studied in [18].
Ramzi et al.
[36]propose to reuse the computations done in the forward pass to approximate the solution of the linearsystem (3) when the inner problem is solved thanks to a quasi-Newton method.
In the last few years, several works have proposed different strategies to solve (1) in a stochasticfashion.
A first set of methods relies on two nested loops : one inner loop to solve the inner problemwith a stochastic method, and one outer loop to update the outer variable with an approximate gradientdirection.
In [17, 24, 9] the authors use several SGD iterations for the inner problem and then usestochastic Neumann approximations to get an estimate solution of the linear system, which providesthem with an approximation of ∇ h used to update x . The analysis of this kind of method wasrefined by Chen et al.
[9], allowing to achieve the same convergence rates as those of SGD.
Theconvergence of the hypergradient when using stochastic solvers for the inner problem and the linearsystem has been studied in [19].
Arbel and Mairal [2] replace the Neumann approximation by SGDsteps to estimate (3).
Other authors have proposed single loop algorithms, alternating steps in theinner and the outer problem.
Hong et al.
[22] propose to perform Neumann approximations of theinverse Hessian and use a single SGD step for the inner problem.
It was refined in [21] and [44]where the optimization procedure uses a momentum acceleration.
Other variations around this ideainclude [23, 26, 10, 20, 28].
We refer to Table 1 in appendix for a detailed comparison of thesemethods.
Notation.
The set of integers between 1 and n (included) is denoted [ n ] .
For f : R p × R d → R ,we denote ∇ i f
( z, x ) its gradient w.r.t.
the i th variable.
The Hessian of f with respect to the firstvariable is denoted ∇ 211 f ( z, x ) ∈ R p × p , and the cross-derivatives matrix is ∇ 221 f ( z, x ) ∈ R d × p .
If vis a vector, ∥ v ∥ is its Euclidean norm.
If M is a matrix, ∥ M ∥ is its spectral norm.
A function is saidto be L -smooth, for L > 0 , if it is differentiable, and its gradient is L -Lipschitz.
Proposed framework
Algorithm Input: initializations z 0 ∈ R p , x 0 ∈ R d , v 0 ∈ R p , number of iterations T , step size sequences ( ρ t ) t<T and ( γ t ) t<T .
for t = 0 , . . .
, T − 1 do
Update z : z t +1 =
z t − ρ t D tz ,
Update v : v t +1 = v t − ρ t D tv ,
Update x : x t +1 = x
t − γ t D tx , where D tz , D tv and D tx are unbiased estimators of D z ( z t , v t , x t ) , D v ( z t , v t , v t ) and D x ( z t , v t , x t ) .
end for
In this section, we introduce our framework in whichthe solution of the inner problem, the solution of thelinear system (3) and the outer variable all evolve atthe same time, following directions that are writtenas a sum of derivatives of F j and G i .
We define
[Equation]
These directions are motivated by the fact that wehave ∇ h ( x ) =
D x
( z ∗ ( x ) , v ∗ ( x ) , x ) , with z ∗ ( x )the minimizer of G ( · , x ) and v ∗ ( x ) the solution of∇ 211 G ( z ∗ ( x ) , x ) v = −∇ 1 F ( z ∗ ( x ) , x ) .
When x is fixed, we approximate z ∗ by doing a gradient descent on G , following the direction
− D z
( z, v, x ) .
Finally, when z and x are fixed, wefind v ∗ by following the direction − D v ( z, v, x ) , which corresponds to a gradient descent onv (cid:55)→ 1 2 ⟨∇ 2 11 G ( z, x ) v, v ⟩ + ⟨∇ 1 F ( z, x ) , v ⟩ .
The rest of the paper is devoted to the study of theglobal dynamics where the three variables z, v and x evolve at the same time, following stochasticapproximations of D z , D v and D x .
The next proposition motivates the choice of these directions.
Proposition 2.1.
Assume that for all x ∈ R d , G ( · , x ) is strongly convex.
If ( z, v, x ) is a zero of( D z , D v , D x ) , then z
= z ∗ ( x ) , v = v ∗ ( x ) and ∇ h ( x ) = 0 .
We also note that the computation of these directions does not require to compute the matrices∇ 211 G ( z, x ) and ∇ 221 G ( z, x ) : we only need to compute their product with a vector, which can becomputed at a cost similar to that of computing a gradient.
The framework we propose is summarized in Algorithm 1.
It consists in following a joint updaterule in ( z, v, x ) that follows directions D tz , D tv and D tx that are unbiased estimators of D z , D v , D x .
The first and most important remark is that whereas ∇ h cannot be written as a sum over samples, thedirections D z , D v and D x involve only simple sums, since their expressions are “linear” in F and G :
[Equation]
It is therefore straightforward to derive unbiased estimators of these directions.
In [28], the authorsconsidered one particular case of our framework, where each direction is estimated by using the
STORM variance reduction technique (see [12]).
Taking a step back by proposing the framework summarized in Algorithm 1 opens the way to potential new algorithms that implement other techniquesthat exist in stochastic single level optimization.
In what follows, we study two of them.
2.1 First example: the SOBA algorithm
The simplest unbiased estimator is obtained by replacing each mean by one of its terms chosenuniformly at random, akin to what is done in classical single-level SGD.
We call the resultingalgorithm SOBA (StOchastic Bilevel Algorithm).
To do so, we choose two independent randomindices i ∈ [ n ] and j ∈ [ m ] uniformly and estimate each term coming from G using G i and eachterm coming from F using F j .
This gives the unbiased SOBA directions
[Equation]
This provides us with a first algorithm, SOBA, where we plug Equations (10a) to (10c) in Algorithm 1.
We defer its analysis to the next section.
Importantly, we use different step sizes for the update in( z, v ) and for the update in x .
We use the same step size in z and in v since the inner problem and thelinear system have similar conditioning, which is that of ∇ 211 G ( z t , x t ) .
The need for a different stepsize for the outer and inner problem is clear: both problems can have a different conditioning.
An important remark for SOBA is that all the stochastic directions used are computed at the samepoint z t , v t and x t with the same indices ( i, j ) .
The update of z , v and x can thus be performedin parallel instead of sequentially, benefiting from hardware parallelism.
Moreover, this enablesto share the computations between the different directions.
This is the case in hyperparametersselection where G i ( z, x ) =
ℓ i ( ⟨ z, d i ⟩ ) + x 2 ∥ z ∥ 2 , with d i a training sample, and ℓ i that measureshow good is the prediction ⟨ z, d i ⟩ .
In this setting, we have ∇ 1 G i ( z, x ) =
ℓ ′ i ( ⟨ z, d i ⟩ ) d
i + xz and∇
G i ( z, x ) v = ℓ ′′ i ( ⟨ z, d i ⟩ ) ⟨ v, d i ⟩ d i .
The prediction ⟨ z, d i ⟩ can thus be computed only once toobtain both quantities.
For more complicated models, where automatic differentiation is used tocompute the different derivatives and Jacobian-vector products, we can store the computational graphonly once to compute at the same time ∇ 1 G i ( z, x ) , ∇ 211 G i ( z, x ) v and ∇ 221 G i ( z, x ) v , requiring onlyone backward pass, thanks to the R technique [34].
Finally, like all single loop bilevel algorithms, our method updates at the same time the inner and outervariable, avoiding unnecessary optimization of the inner problem when x is far from the optimum.
2.2 Global variance reduction with the SABA algorithm
In classical optimization, SGD fails to reach optimal rates because of the variance of the gradientestimator.
Variance reduction algorithms aim at reducing this variance, in order to follow directionsthat are closer to the true gradient, and to achieve superior practical and theoretical convergence.
In our framework, since the directions D z , D v and D x are all written as sums of derivatives of F j and
G i , it is easy to adapt most classical variance reduction algorithms.
We focus on the celebrated SAGA algorithm [13].
The extension we propose is called SABA (Stochastic Average Bilevel Algorithm).
The general idea is to replace each sum in the directions D by a sum over a memory, updating onlyone term at each iteration.
To help the exposition, we denote y =
( z, x, v ) the vector of joint variables.
Since we have sums over i and over j , we have two memories for each variable: w ti for i ∈ [ n ] and˜ w tj for j ∈ [ m ] , which keep track of the previous values of the variable y .
At each iteration t , we draw two random independent indices i ∈ [ n ] and j ∈ [ m ] uniformly andupdate the memories.
To do so, we put w t
+1 i =
y t and w t
+1 i ′ =
w ti ′ for
i ′
̸
= i , and ˜ w t +1 j = y t and˜ w t +1 j ′ =
˜ w tj ′ for j ′ ̸ = j .
Each sum in the directions D is then approximated using SAGA-like rules:given n functions ϕ
i ′ for i ′ ∈ [ n ] , we define S [ ϕ, w ] ti = ϕ
i ( w t {++
+} i
[Equation]
This is an unbiased estimators of the average of the ϕ ’s since E i (cid:104) S [
[Equation]
[-=
 ϕ i ( y t ) .-][Equation]
With a slight abuse of notation, we call ∇
Gv the sequence of functions ( y (cid:55)→ ∇ 211 G i ( z, x ) v )
i ∈
[ n ]and ∇ 221 Gv the sequence of functions ( y (cid:55)→ ∇ 221 G i ( z, x ) v ) i ∈ [ n ] .
We define the SABA directions as
[Equation]
[Equation]
11a
[Equation]
11b
[Equation]
11c
[Equation]
These estimators are unbiased estimators of the directions
D z , D v and D x .
The SABA algorithmcorresponds to Algorithm 1 where we use Equations (11a) to (11c) as update directions.
When takinga step size γ t = 0 in the outer problem, hereby stopping progress in x , we recover the iterationsof the SAGA algorithm on the inner problem.
In practice, the sum in S is computed by doing arolling average (see Appendix B for precision), and the quantities ϕ
i ( w ti ) are stored rather thanrecomputed: the cost of computing the SABA directions is the same as that of SGD.
It requires anadditional memory for the five quantities, of total size n × p + ( n + m ) ×
( p + d ) floats that can bereduced by using larger batch sizes.
Indeed, if b in and b out are respectively the inner and the outerbatch sizes, the memory load is reduced to n b × p + ( n b + m b ) × ( p × d ) with n b = ⌈ nb inn ⌉ andm b = ⌈ mb out ⌉ which are smaller than the number of samples.
This memory load can also be reducedin specific cases, for instance when G and F correspond to linear models, where the individualgradients and Hessian-vector products are proportional to the samples.
In this case, we only store theproportionality ratio, reducing the memory load to 3 n + 2 m floats.
Like for SOBA, the computationsof the new quantities ϕ i ( w t +1 i ) are done in parallel, thus benefiting from hardware acceleration andshared computations.
Despite this memory load, using SAGA-like variance reduction instead of STORM as done in [28, 44, 26] has the advantage to bring the variance of the estimate directions tozero, enabling faster O ( 1 T ) convergence.
In the next section, we show that SABA is fast.
It essentially has the same properties as SAGA:despite being stochastic, it converges with fixed step sizes, and reaches the same rate of convergenceas gradient descent on h .
Theoretical analysis
In this section, we provide convergence rates of SOBA and SABA under some classical assumptions.
Note that, unlike most of the stochastic bilevel optimization papers, we work in finite sample settingrather than the more general expectation setting.
Actually, SABA does not make any sense forfunctions that don’t have a finite sum structure.
However, we stress that SOBA could be studied in amore general setting to obtain the same bounds as here.
Also, the finite sum setting is still interestingsince doing empirical risk minimization is very common in practice in machine learning.
The proofsand the constants in big-
O are deferred in Appendix C.
3.1 Background and assumptions
We start by stating some regularity assumptions on the functions F and G .
Assumption 3.1.
The function F is twice differentiable.
The derivatives ∇ F and ∇ 2 F are Lipschitzcontinuous in ( z, x ) with respective Lipschitz constants L F 1 and L F 2 .
Note that the above assumption is typically verified in the machine learning context, e.g., when F isthe ordinary least squares (OLS) loss or the logistic loss.
Assumption 3.2.
The function G is three times continuously differentiable on R p × R d .
Forany x ∈ R d , G ( · , x ) is µ G -strongly convex.
The derivatives ∇ G , ∇ 2 G and ∇ 3 G are Lipschitzcontinuous in ( z, x ) with respective Lipschitz constants L G 1 , L G 2 and L G 3 . We denote by κ = L G 1 µ G thecondition number of G .
Strong convexity and smoothness with respect to z of G are verified when G is a regularized leastsquares/logistic regression with a full rank design matrix, when the data is not separable for thelogistic regression.
Moreover, the strong convexity ensures the existence and uniqueness of the inneroptimization problem for any x ∈ R d .
Assumption 3.3.
There exists C F > 0 such that for any x we have ∥∇ 1 F ( z ∗ ( x ) , x ) ∥ ≤ C F .
This assumption, combined with the strong convexity of G ( · , x ) , shows boundedness of v ∗ .
Thisassumption holds, for instance, in the case of hyperparameters selection for a Ridge regressionproblem.
Note that in Assumptions 3.1 and 3.2, we assume more regularity of F and G than instochastic bilevel optimization literature (see for instance [17, 22, 24, 2]).
It is necessary to get thesmoothness of v ∗ which will allow to adapt the proof of Chen et al.
[9] and get tight convergencerates.
The following lemma gives us some smoothness properties of the considered directions thatwill be useful to derive convergence rates of our methods.
Lemma 3.4.
Under the Assumptions 3.1 to 3.3, there exist constants L z , L v and L x such that
[Equation]
x ) ∥ 2 + ∥ v − v ∗ ( x ) ∥ 2 ) and∥
D x
( z, v, x ) − ∇ h ( x ) ∥ 2 ≤ L 2 x ( ∥ z − z ∗ ( x ) ∥ 2 + ∥ v − v ∗ ( x ) ∥ 2 ) .
In first order optimization, a fundamental assumption on the objective function is the smoothnessassumption.
In the case of vanilla gradient descent applied to a function f , it allows to get aconvergence rate of ∥∇ f ( x t ) ∥ 2 in O (1 /T ) , i.e. convergence to a stationary point [33].
The followinglemma proved by Ghadimi and Wang [17, Lemma 2.2] ensures the smoothness of h .
Lemma 3.5.
Under the Assumptions 3.1 to 3.3, the function h is L h -smooth for some L h > 0 .
The constant L h is specified in Appendix C.3.
As usual with the analysis of stochastic methods, wedefine the expected norms of the directions V tz = E
[ ∥ D tz ∥ 2 ] , V tv =
E [ ∥ D tv ∥ 2 ] and V tx =
E
[ ∥ D tx ∥ 2 ] ,where the expectation is taken over the past.
Thanks to variance-bias decomposition, they are thesum of the variance of the stochastic direction and the squared-norm of the unbiased direction.
For SOBA, we use classical bounds on variances like those found for instance in [22]:
Assumption 3.6.
There exists B z , B v and B x such that for all t , 241 E t [ ∥ D tz ∥ 2 ] ≤
B 2 z (1 + ∥ D z ( z t , v t , x t ) ∥ 2 ) and E t [ ∥ D tv ∥ 2 ] ≤ B 2 v (1 + ∥ D
v ( z t , v t , x t ) ∥ 2 )where E t denotes the expectation conditionally to ( z t , v t , x t ) . For SOBA and SABA, we need to bound the expected norm of D t x . For SABA, this assumptionallows to get a the same sample complexity as SAGA for single level problems. Assumption 3.7.
There exists B x such that for all t , E t [ ∥ D tx ∥ 2 ] ≤ B 2 x .
Assumptions 3.6 and 3.7 are verified for instance, if all the G i and ∇ 1 G i have at most quadraticgrowth, and if F has bounded gradients.
They are also verified if the iterates remain in a compact set.
Note that we do not assume that G has bounded gradients, as this would contradict its strong-convexity.
Finally, for the analysis of SABA, we need regularity on each G i and F j :
Assumption 3.8.
For all i ∈ [ n ] and j ∈ [ m ] , the functions ∇ G i , ∇ F j , ∇ 211 G i and ∇ 221 G i are
Lipschitz continuous in ( z, x ) .
3.2 Fundamental descent lemmas
Our analysis for SOBA and SABA is based on the control of both δ tz = E
[ ∥ z t − z ∗ ( x t ) ∥ 2 ] andδ tv =
E [ ∥ v t − v ∗ ( x t ) ∥ 2 ] , Strong convexity of G and smoothness of z ∗ ( x ) and v ∗ ( x ) allow to obtainthe following lemma by adapting the proof of Chen et al. In what follows, we drop the dependencyof the step sizes ρ and γ in t for clarity.
Lemma 3.9.
Assume that γ 2 ≤ min (cid:16) µ G L 2 ∗ 4 B 2 x L 2 zx , [Equation] G L [Equation]
8 B 2 x L 2 vx (cid:17) ρ .
We have:
[Equation]
where β zx = β vx = 3 L 2 ∗ , β zx = 8 L 2 ∗ µ G , β vx = L 2 ∗ µ G , L ∗ is the maximum between the Lipschitzconstants of z ∗ and v ∗ (see Lemma {+C.1), β vz = 
+}
µ 3 G ( L F 1 µ G + L G 2 ) 2 , L zx and L vx are respectivelythe smoothness constants of z ∗ and v ∗ .
We insist that this result is obtained in general for Algorithm 1 with arbitrary unbiased directions.
Wecan therefore invoke this lemma for the analysis of both SOBA and SABA.
We use the smoothness ofh to get the following lemma, which is similar to [9, Lemma 1].
Lemma 3.10.
Let h t =
E
[ h ( x t )] and g t = E [ ∥∇ h ( x t ) ∥ 2 ] .
We have
[Equation]
If z t = z ∗ ( x t ) , v t
= v ∗ ( x t ) , that is δ z , δ v both cancel and D x ( z t , v t , x t ) = ∇ h ( x t ) , we get aninequality reminiscent of the smoothness inequality for SGD on h .
3.3 Analysis of SOBA
The analysis of SOBA is based on Lemmas 3.5 and 3.9.
We have the following theorem, with fixedstep sizes depending on the number of iterations:
Theorem 1 (Convergence of SOBA, fixed step size) .
Fix an iteration T > 1 and assume that
Assumptions 3.1 to 3.7 hold.
We consider fixed steps ρ t = ρ √ T and γ t = ξρ t with ρ and ξ are precised in the appendix.
Let ( x t ) t ≥ 1 the sequence of outer iterates for SOBA.
Then, 1 T (cid:80) Tt =1 E
[ ∥∇ h ( x t ) ∥ 2 ]
=
O ( T − 12 ) .
As opposed to [22], we do not need that the ratio γρ goes to 0, which allows to get a complexity (thatis, the number of call to oracles to have an ϵ -stationary soluion) in O ( ϵ − 2 ) better than the ˜ O
( ϵ − 52 )they have.
Also, note that this rate is the same as the one of SGD for non-convex and smooth objective[8].
We obtain a similar rate using decreasing step sizes:
Theorem 2 (Convergence of SOBA, decreasing step size) .
Assume that Assumptions 3.1 to 3. hold. We consider steps ρ t = ρt − 12 and γ t = ξρ .
Let x t the sequence of outer iterates for SOBA.
Then, inf t ≤ T E [ ∥∇ h ( x t ) ∥ 2 ]
=
O (log( T ) T − 12 ) .
As for SGD, SOBA suffers from the need of decreasing step sizes to get actual convergence because ofthe variance of the estimation on each directions.
On the other hand, the analysis of SABA leveragesthe dynamic of all three variables, resulting in fast convergence with fixed step sizes.
3.4 SABA: a stochastic method with optimal rates
In what follows, we denote N = n + m the total number of samples.
The following theoremshows O ( N 23 T − 1 ) convergence for the SABA algorithm in the general case where we only assumesmoothness of h .
Our analysis of SABA is inspired by the analysis of single-level SAGA by Reddiet al.
Theorem 3 (Convergence of SABA, smooth case) .
Assume that Assumptions 3.1 to 3.3 and 3.7 to 3.8 hold.
We suppose ρ = ρ ′ N − 2 3 and γ = ξρ , where ρ ′ and ξ depend only on F and G and are specified in appendix. Let x t the iterates of SABA.
Then, 1 T (cid:80) Tt =1 E
[ ∥∇ h ( x t ) ∥ 2 ]
= O (cid:16) N 23 T − 1 (cid:17) .
To prove the theorem, the idea is to control the distance from the memory to the current variables.
Wedefine S t = 1 n (cid:80) ni
=1 ∥
y t − w ti ∥ 2 + 1 m (cid:80) mj
=1 ∥
y t − ˜ w tj ∥ 2 .
In appendix, we show that we can find scalars ϕ s , ϕ z , ϕ v > 0 such that the quantity L t = h t
+ ϕ s S t + ϕ z δ tz + ϕ v δ tv satisfies L t +1 ≤ L t − γ 2 g t .
Summing these inequalities for t = 1 . . .
T and using the fact that L t is lower bounded demonstratesthe theorem.
Note that the step sizes are constant with respect to the time, but they scale with N − 23 .
As aconsequence, the sample complexity is O ( N 23 ϵ − 1 ) which is analogous of the one of SAGA fornon-convex single level problems [37].
This is better than the sample complexity of Algorithm 1 withfull batch directions, which is O
( Nϵ − 1 ) .
Hence, with SABA, we get the best of both worlds: thestochasticity makes the scaling in N of the sample complexity goes from N in full batch mode
to
N 23 for SABA, and the variance reduction makes the scaling in ϵ goes from ϵ − 2 for SOBA to ϵ − 1 for SABA. Our experiments in Section 4 confirm this gain.
Furthermore, if we assume that h satisfies a Polyak-Łojasiewicz (PL) inequality, we recover linearconvergence.
Recall that h has the PL property if there exists µ h > 0 such that for all x ∈ R d ,12 ∥∇ h ( x ) ∥ 2 ≥ µ h ( h ( x ) − h ∗ ) with h ∗ the minimum of h .
Theorem 4 (Convergence of SABA, PL case) .
Assume that h satisfies the PL inequality and that
Assumptions 3.1 to 3.3 and 3.7 to 3.8 hold. We suppose ρ = ρ ′ N − 23 and γ = ξρ ′ N − 1 , where ρ ′ and ξ depend only on F and G and are specified in appendix. Let x t the iterates of SABA and c ′ ≜ min (cid:0) µ h , 1 16 P ′ (cid:1) with P ′ specified in the appendix.
Then, E [ h T ] − h ∗ = (1 − c ′ γ ) T ( h 0 − h ∗ + C 0 ) where C 0 is a constant specified in appendix that depends on the initialization of z, v, x and memory.
The proof is similar to that of the previous theorem: we find coefficients ϕ s , ϕ z , ϕ v such that
L t = h t + ϕ s S t + ϕ z δ tz + ϕ v δ tv satisfies the inequality L t +1 ≤ (1 − c ′ γ )
L t , which is thenunrolled.
Note that in the case where we initialize z and v with z
= z ∗ ( x 0 ) , v 0 = v ∗ ( x 0 ) , andthe memories w
i = w 0 , ˜ w 0 j = w 0 for all i, j , the constant C 0 cancels and the bound simplifies
to
E [ h ( x T )] − h ∗ ≤ (1 − γµ h 4 ) T ( h ( x 0 ) − h ∗ ) .
Just like classical variance reduction methods in single-level optimization, this theorem shows thatour method achieves linear convergence under PL assumption on the value function.
To the bestof our knowledge, our method is the first stochastic bilevel optimization method that enjoys suchproperty.
We note that the PL hypothesis is more general than µ h -strong convexity of h – it is anecessary condition for strong convexity.
We see here the importance of global variance reduction.
Indeed, using variance reduction onlyon z and SGD on x would lead to sub-linear convergence in x .
This would be the case even witha perfect estimation of z ∗ ( x ) .
Similarly, using variance reduction only on x and SGD on z wouldlead to sub-linear convergence in z , and hence in x . Using global variance reduction with respectto each variable as we propose here is the only way to achieve linear convergence.
We now turn toexperiments, where we find that our method is also promising from a practical point of view.
Experiments
Here we compare the performances of SOBA and SABA with competitor methods on different tasks.
The different methods being compared are stocBiO [24], AmiGO [2], FSLA [28], MRBO [44], TTSA[22], BSA [17] and SUSTAIN [26].
A detailed account of the experiments is provided in Appendix B.
4.1 Hyperparameters selection
The first task we perform is hyperparameters selection to choose regularization parameters onℓ 2 logistic regression.
Let us denote (( d train i , y train i ))
1 ≤ i ≤ n and (( d val i , y val i ))
1 ≤ i ≤ m the trainingand the validation sets.
In this case, the inner variable θ corresponds to the parameters of themodel, and the outer variable λ to the regularization.
The functions F and G of the problem (1)are the logistic loss, with ℓ 2 penalty for G , that is to say F ( θ, λ ) = 1 m (cid:80) mi =1 φ ( y val i ⟨ d val i , θ ⟩ ) and
G ( θ, λ ) = 1 n (cid:80) n i =1 φ ( y train i ⟨ d train i , θ ⟩ ) + 1 2 (cid:80) p k =1 e λ k θ 2 k where φ ( u ) =
log(1 + e − u ) .
We fit abinary classification model on the IJCNN1 1 dataset.
Here, n = 49 990 , m = 91 701 and p = 22 .
The suboptimality gap is plotted in Figure 2a for each method.
The lowest values are reached by
AmIGO and SABA, but much quicker for SABA.
Moreover, SABA is the only single-loop method
[Figure]
that reaches a suboptimality below 10 − 3 . Among all methods, SOBA is the first to reach it bestvalue, but this value is also the highest.
The gap between SOBA and SABA highlights the benefits ofvariance reduction: it gives us a lower plateau and the fixed step sizes enable faster convergence.
4.2 Data hyper-cleaning
The second task we perform is data hyper-cleaning introduced in [15] on the MNIST 2 dataset.
The datais patitioned into a training set ( d train i , y train i ) , a validation set ( d val i , y val i ) , and a test set.
The trainingset contains 20000 samples, the validation set 5000 samples and the test set 10000 samples.
Thetargets y take values in { 0 , . . .
, 9 } and the samples x are in dimension 784 .
Each sample in the trainingset is corrupted with probability p : a sample is corrupted when we replace its label
y
i by a randomlabel in { 0 , . . .
Samples in the validation and test sets are not corrupted.
The goal of datacleaningis to train a multinomial logistic regression on the train set and learn a weight per training sample,that should go to 0 for corrupted samples.
This is formalized by the bilevel optimization problem(1) with F ( θ, λ ) =
1 m (cid:80) mi =1 ℓ ( θd val i , y val i ) and G ( θ, λ ) = 1 n (cid:80) ni =1 σ ( λ i ) ℓ ( θd train i , y train i )
+ C r
∥ θ ∥ 2where ℓ is the cross entropy loss and σ is the sigmoid function.
The inner variable θ is a matrix ofsize 10 × 784 , and the outer variable λ is a vector in dimension n train = 20000 .
For the estimated parameters θ during optimization, we report in Figure 2b the test error, i.e. , thepercent of wrong predictions on the testing data.
We use for this experiment a corruption probabilityp = 0 .
In general, the error decreases quickly until it reaches a final value.
We observe thatour method SABA outperforms all the other methods by reaching faster its smallest error, which issmaller than the ones of the other methods.
For SOBA, it reaches a lower final error than stocBiOand BSA.
In appendix, we provide other convergence curves, and find that for higher values of p ,
SABA is still the fastest algorithm to reach its final accuracy, but the algorithm MRBO can reach alower test accuracy.
Overall, we find that among all methods, even those that implement variancereduction (that is FSLA, MRBO, SUSTAIN, SABA), SABA is the one that demonstrates the bestempirical performance.
Conclusion
In this paper, we have presented a framework for bilevel optimization that enables the straightforwarddevelopment of stochastic algorithms.
The gist of our framework is that the directions in Equations (4)to (6) are all written as simple sums of samples derivatives.
We leveraged this fact to propose SOBA,an extension of SGD to our framework, and SABA, an extension of SAGA to our framework, whichachieves similar convergence rates as gradient descent on the value function.
Finally, we think thatour framework opens a large panel of potential methods for stochastic bilevel optimization involvingtechniques of extrapolation, variance reduction, momentum and so on. Checklist For all
authors.
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes] See Section 3 and Section (b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix C
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See supplementary material
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix B
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix B
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]