Fine-grained Analysis of Stability and Generalization for Stochastic Bilevel Optimization

1 Introduction

In this paper, we focus on establishing stability and generalization analysis for the stochastic bilevel optimization (SBO) (Bracken and McGill, 1973; Ji et al., 2021; Bao et al., 2021) defined as follows:

∗Corresponding author.

where d1, d2 ∈ N+, the outer objective function f and the inner objective function g are both continuous and differentiable, ξ, ζ are samples drawn from the validation set and training set respectively.

For this bilevel optimization scheme, we often call miny∈Rd2 Eζ[g(x, y; ζ)] as the inner (or lower-level) problem, and name minx∈Rd1 Eξ [f (x, y∗ (x); ξ)] as the outer (or upper-level) problem. The goal of (1) is to minimize the outer objective function R(x) (also F(x, y∗ (x))) with respect to (w.r.t.) the model parameter x, where parameter y∗ (x) is derived from the inner minimization formulation.

The SBO formulation in (1), stemming from (Bracken and McGill, 1973), has attracted increasing attention in many machine learning applications including hyper-parameter optimization (Franceschi et al., 2017, 2018; Lorraine and Duvenaud, 2018; MacKay et al., 2019; Okuno et al., 2021; Zhang et al., 2023), generative adversarial learning (Pfau and Vinyals, 2016), meta learning (Franceschi et al., 2018; Bertinetto et al., 2018; Z¨ugner and G¨unnemann, 2019; Rajeswaran et al., 2019; Ji et al., 2020), and reinforcement learning (Tschiatschek et al., 2019). Indeed, there are rich computing methods to implement this bilevel optimization scheme, as well as theoretical works on optimization convergence analysis (Li et al., 2020; Ji et al., 2021). However, the generalization analysis of SBO is still far less understood from the viewpoint of statistical learning theory (STL), e.g., algorithmic stability and generalization analysis (Hardt et al., 2016; Lei and Ying, 2020; Bao et al., 2021).

Stability-based generalization analysis can be traced back to the 1970s (Rogers and Wagner, 1978) and has achieved rapid developments in STL, see e.g., (Bousquet and Elisseeff, 2002; Elisseeff et al., 2005; Hardt et al., 2016; Liu et al., 2017; Lei and Ying, 2020; Lei et al., 2021a; Deng et al., 2021; Kuzborskij and Lampert, 2018). To match the characterizations of various algorithms, different definitions of algorithmic stability have been formulated (including the uniform stability (Bousquet and Elisseeff, 2002), uniform argument stability (Liu et al., 2017), locally elastic stability (Deng et al., 2021), on-average stability (Kuzborskij and Lampert, 2018) and on-average argument stability (Lei and Ying, 2020)) to better investigate their generalization bounds. The on-average argument stability was proposed in (Lei and Ying, 2020) to establish the fine-grained generalization analysis of singlelevel pointwise stochastic gradient descent (SGD). Subsequently, Lei et al. extended the stability-based generalization assessment to the pairwise SGD (Lei et al., 2021a), where systematic strategies have been provided to make a better balance between generalization error and optimization error. As far as we know, there is only one study exploring the generalization analysis of SBO (Bao et al., 2021), which presents an expectation generalization bound w.r.t. the validation set via the uniform stability approach. However, the theoretical analysis of (Bao et al., 2021) is limited to unrolled differentiation (UD) based algorithms with re-initialization in inner-level for hyper-parameter optimization, which may not be applicable to other commonly used optimization algorithms, e.g., single timescale SGD (SSGD) (Zhou et al., 2022a; Liu et al., 2022; Chen et al., 2022) and two timescale SGD (TSGD) (Zhou et al., 2022a; Liu et al., 2022; Hong et al., 2023). Therefore, it is important to further investigate the generalization guarantees for general SBO formulation to cover wider bilevel optimization algorithms.

To address the aforementioned issue, this paper establishes the fine-grained stability and generalization analysis for general first-order bilevel optimization methods. Our main contributions are summarized as follows:

• Firstly, we establish the quantitative connection between the generalization gap of bilevel optimization methods and on-average argument stability. Especially for the l2 on-average argument stability, the derived stabilitybased generalization bounds involve the empirical risks, which is consistent with the previous analysis for singelevel optimization (Lei and Ying, 2020; Lei et al., 2021a).
• Secondly, this paper provides several stability bounds of bilevel optimization methods associated with both SSGD and TSGD algorithms, where different conditions of objective functions (i.e., SC-SC, C-C, NC-NC) are considered. Moreover, we extend the results to the more general setting by relaxing the restriction (e.g., Lipschitz continuity and smoothness assumptions) of the optimization objective. As far as we know, this is the first systemic generalization analysis for first-order SGD-based bilevel optimization under the low-noise setting.
• Finally, we conduct experimental evaluations for bilevel optimization methods including hyperparameter optimization. Empirical results validate our theoretical findings about the relationship between the generalization gap and the size of the validation set as well as the maximum value of inner (outer) iterations.

To better evaluate our results, we compare them with the most related work on stability and generalization analysis (Bao et al., 2021) from the following perspectives:

• Optimization strategy. The previous UD-based hyperparameter optimization (Algorithm 1 in (Bao et al., 2021)) requires reinitialization in the inner-level parameters before each iteration. Different from this special case, this
• paper considers the SBO algorithms where the parameters in inner-level and outer-level are both updated continuously (e.g., SSGD (Zhou et al., 2022a; Chen et al., 2022) and TSGD (Zhou et al., 2022a; Liu et al., 2022; Hong et al., 2023)). The iteration strategy matching our analysis has been used extensively in practice (Ji et al., 2021; Liu et al., 2022; Ghadimi and Wang, 2018). Especially for theoretical analysis of TSGD, it is challenging to deal with the gradient summation during the inner iterations and the previous analysis technique (Bao et al., 2021) can not be extended to this case directly.
• Analysis tool. Different from uniform stability used in (Bao et al., 2021), this paper develops the analysis technique of on-average argument stability to provide the fine-gained generalization bounds under low noise settings, where the stability bounds involve a weighted sum of empirical risks instead of the uniform Lipschitz constants.
• Conditions of objective functions. Similar to the previous stability analysis in (Lei et al., 2020; Shen et al., 2020; Zhou et al., 2022b), the objective functions in (Bao et al., 2021) are assumed to be bounded, third order continuously differentiable and smooth. Here, we merely need the bilevel objective functions to be nonnegative, smooth and Lipschitz continuous, where the last condition for the outer-level function can be further removed by the l2 on-average argument stability. Detailed stability results have been derived for both SSGD and TSGD algorithms under NC-NC, C-C and SC-SC settings. In addition, we also establish generalization bounds by replacing the smooth condition with the weaker H¨ older continuous assumption.

2 ProblemFormulation

Given distributions D1, D2, we get the validation set

and the training setDm2 := {ζi}m2

by independent sampling, where m1 and m2 are the sample sizes. This paper focuses on the outer-level population risk w.r.t D1 and empirical risk w.r.t Dm1 1 , which are defined respectively as

where f : Rd1 × Rd2 → R is an objective function and y(x) is the inner model parameter given the outer model parameter x (also see (1)).

Let (x, y(x)) in (1) be estimated by a stochastic algorithm A with data Dm1 ,Dm2 , i.e. A(Dm1 ,Dm2 ). Similar to the previous works (Bao et al., 2021; Hoffer et al., 2017; Keskar et al., 2017), in order to evaluate the approximated searching of hyperparameters, we define

1 Thus we consider adding corruptions to Dm1 to access the generalization behavior of the meta-learner (Thrun, 1998) at upper level.

in the upper (outer) level as the generalization gap of A, which measures the difference between the population risk R(A) and the empirical risk RDm1 (A).

The following conditions have been used to characterize the theoretical properties of objective functions in (1).

Definition 1. (Joint Lipschitz Continuity (Ji et al., 2021; Liu et al., 2022)). An objective function f is jointly Lf -Lipschitz over Rd1 ×Rd2 , if there holds

Definition 2. (Joint Smoothness (Lei et al., 2021b)). An objective function f is ℓf -smooth over Rd1 ×Rd2 , if

Definition 3. (Strong Convexity). A function ψ is µ-stronglyconvex over a set X, if ∀t, t′ ∈ X,

Definition 4. (H¨ older Continuity). Let τ > 0, α ∈ [0, 1]. Gradient ∇f is (α, τ)-H¨ older continuous over Rd1 × Rd2 , if there holds

The above conditions for objective functions have been used extensively in convergence analysis for bilevel optimization (Ji et al., 2021; Ghadimi and Wang, 2018; Liu et al., 2022) and stability-based generalization analysis for singlelevel optimization methods (Hardt et al., 2016; Lei et al., 2021b). Moreover, theH¨ older continuity is much weaker than the Lipschitz continuity and smoothness (Lei and Ying, 2020; Nesterov, 2015). If Definition 4 holds with α = 1, then f is smooth (see Definition 2). And if Definition 4 holds with α = 0, f becomes Lipschitz continuous as in Definition 1 and can be non-differentiable (Lei and Ying, 2020). The objective functions satisfying Definition 4 include the mean absolute function, the hinge function and some of their variants (Lei and Ying, 2020; Steinwart and Christmann, 2008).

Definition 5. (On-average Argument Stability (Lei and Ying, 2020)). Let Dm1 = {z1, . . . , zm1 } and ˜Dm1 = {˜ z1, . . . , ˜ zm1 } be two sets drawn independently from distribution Dm1 1 . For any i = 1, . . . ,m1, define D(i) = {z1, . . . , zi-1, ˜ zi, zi+1, . . . , zm1 }. Denote the E as the expectation of EDm1 ,Dm2 , ˜Dm1 ,A. We say a randomized algorithm A is l1(β) on-average argument stable if

and l2(β2 ) on-average argument stable if

Remark 1. The on-average argument stability measures the average sensitivity (stability) of output parameters of the learning algorithm when at most one validation sample is changed. Definition 5 is different from Definition 1 in (Bao et al., 2021), where the uniform stability is evaluated by the drift of prediction error of hyperparameter optimization algorithm and the boundedness of loss function often is required.

Based on the above definitions, we introduce the requirements of f, g in our analysis.

Assumption 1. (Outer Function Assumption). Assume that the outer objective function f in (1) satisfies

(I) f is jointly Lf -Lipschitz.

(II) f is nonnegative, continuously differentiable and ℓf - smooth.

Assumption 2. (Inner Function Assumption). Assume that the inner objective function g in (1) satisfies

(I) g is jointly Lg-Lipschitz.

(II) g is continuously differentiable and ℓg-smooth.

3 Quantitative Relationship between Generalization and Stability

This section states that the generalization gap of (1) can be bounded by the on-average argument stability. Before providing the detailed conclusion of Theorem 1, we first introduce the self-bounding property definition.

Lemma 1. (Self-bounding property). Assume that for all z ∈ D, the map w ↦→ f(w; z) is nonnegative, and w ↦→ ∂f(w; z) is (α, τ)-H¨ older continuous with α ∈ [0, 1]. Then we have

The self-bounding property of f with (α, τ)-H¨ older continuous (sub)gradient contains the specific Lipschitz continuous (α = 0) and smoothness (α = 1) conditions (Lei et al., 2021b).

Theorem 1. (I) If algorithmA is l1(β) on-average argument stable in expectation and the outer-level function f is Lf - Lipschitz continuous w.r.t. (x, y) ∈ Rd1 ×Rd2 , denote E as EA,Dm1 ,Dm2 , there holds

(II) If algorithm A is l2(β2 ) on-average argument stable in expectation and f is nonnegative and ℓf -smooth w.r.t. (x, y) ∈ Rd1 × Rd2 , denote E as EA,Dm1 ,Dm2 , then

Algorithm 1 Computing algorithm of SSGD

Input: Validation data Dm1 = {ξi}m1 i=1 and training set Dm2 = {ζi}m2 i=1 , the total number of iterations K, step sizes ηx, ηy.

Initialization: x0 and y0.

• 1: for k = 1 to K -1 do
• 2: Uniformly sample ξi ∈ Dm1 and ζi ∈ Dm2 :
• 3: yk+1 = yk - ηy∇yg (xk, yk (xk) ; ζi)
• 4: xk+1 = xk - ηx∇xf (xk, yk (xk) ; ξi)
• 5: end for

Output: xK and yK.

Algorithm 2 Computing algorithm of TSGD

Input: Validation data Dm1 = {ξi}m1 i=1 and training set Dm2 = {ζi}m2 i=1 , the total number of inner iterations T and outer iterations K, step sizes ηx and ηy.

Initialization: x0 and y0 0 .

• 1: for k = 0 to K -1 do
• 2: for t = 0 to T - 1 do
• 3: Uniformly sample ζi ∈ Dm2 :
• 4: yt+1 k = ytk - ηy∇yg (xk, ytk (xk) ; ζi)
• 5: end for
• 6: Uniformly sample ξi ∈ Dm1 :
• 7: xk+1 = xk - ηx∇xf (xk, yT k (xk) ; ξi )
• 8: y0 k+1 = yT k
• 9: end for

Output: xK and y0K .

where the constant γ > 0.

(III) If algorithm A is l2(β2 ) on-average argument stable in expectation, f is nonnegative and (α, τ)-H¨ older continuous w.r.t. (x, y) ∈ Rd1 ×Rd2 with α ∈ [0, 1], then

2 , where the constant γ > 0.

Remark 2. Theorem 1 validates the connection between onaverage argument stability and the generalization gap. Especially, the smoothness assumption is further relaxed by the H¨ older continuity in Theorem 1(III).

Remark 3. Different from the uniform stability technique employed in (Bao et al., 2021), the on-average argument stability further exploits the Lipschitz continuous (Lf ) or smooth properties (ℓf ) of the objective function as well as the stability parameter (β) to bound the algorithmic generalization gap. Especially, there is a trade-off between the empirical risk and the algorithmic stability bound.

Remark 4. There are several advantages of l2 on-average argument stability in Theorem 1 (II), where Assumption 1(I) is removed and the low noise assumption can be used to obtain a fine-gained result instead of the Lipschitz constant (Lei and Ying, 2020). If algorithm A is l2(β2 ) on-average argument stable, then we derive the upper bound of generalization gap

Algorithm 3 Computing algorithm of UD (Bao et al., 2021)

Input: Validation data Dm1 = {ξi}m1 i=1 and training set Dm2 = {ζi}m2 i=1 , the total number of inner iterations T and outer iterations K, step sizes ηx and ηy.

Initialization: x0 and y0 .

for t = 0 to T - 1 do

Uniform sampling ζi ∈ Dm2 :

yt+1

end for

k = ytk - ηy∇yg (xk, ytk (xk) ; ζi)

Uniform sampling ξi ∈ Dm1 :

xk+1 = xk - ηx∇xf (xk, yT

y0

end for

k+1 = yT

k

Output: xK and y0K .

k (xk) ; ξi )

4 Stability Analysis for Stochastic Bilevel Optimization

To solve bilevel optimization formulation (1), some gradientbased algorithms are designed based on the single timescale or two timescale strategies (Ji et al., 2021; Chen et al., 2022; Liu et al., 2020, 2022; Zhou et al., 2022a). In the following, we introduce the computing approaches for (1) (SSGD in Algorithm 1 and TSGD in Algorithm 2), and then establish their generalization assessments by presenting their algorithmic stability bounds.

4.1 Stability and Generalization Analysis for SSGD

Let ηx and ηy be the step sizes for updating x and y. According to Theorem 1, the on-average argument stable metrics in Definition 5 for SSGD algorithm A with K iterations

Now we state the upper bounds of on-average argument stability for SSGD in Algorithm 1.

Theorem 2. Suppose that Assumptions 1, 2 hold and Algorithm A is SSGD with K iterations. Denote ℓ = max {ℓf , ℓg}, η = max{ηx, ηy}.

(I) Assume that the bilevel optimization problem (1) is SCSC with strong convexity parameters µf and µg. Let the step sizes satisfy that 2(µf+µg)-√4(µf+µg)2-2(ℓ2f +ℓ2g) 2(ℓ2f +ℓ2g ) ≤ ηx =

average argument-stable in expectation with

and l2(β2 ) on-average argument-stable in expectation with

(II) Assume that the bilevel optimization problem (1) is CC. If η ≤ c1 ln(K) √2Kℓ for some c1 > 0, then A is l1(β) onaverage argument-stable in expectation with β =

And A is l2(β2 ) on-average argument-stable in expectation, where β2 =

Remark 5. Theorem 2 demonstrates that the algorithmic stability can be improved when the model can achieve a relatively small optimization error. In addition, the ℓf -smooth assumption also can be replaced by H¨ older continuous condition. In order to obtain tighter bounds for the SSGD algorithm, we further derive its algorithmic stability with refined step sizes in Proposition 1 in Appendix C.

Combining Theorems 1 and 2, the algorithmic generalization bounds of SSGD are further summarized in Table 1 under the low noise settings (small empirical risk). As shown in Table 1, the generalization bounds of some SSGD algorithms achieve the rate of O(m-1 1 ) under the limitations of step sizes in Theorem 2. From Table 1, one can easily find that objective functions with better (convexity) properties usually lead to better algorithmic stability and generalization performance, which is consistent with the existing stability and generalization analysis for single-level problems (Lei and Ying, 2020; Kuzborskij and Lampert, 2018; Lei et al., 2021b).

4.2 Stability and Generalization Analysis for TSGD

Now we turn to establish the stability bounds of the TSGD algorithm with different inner and outer functions (i.e., NCNC, C-C and SC-SC).

Assume that f is ℓf -smooth and g is ℓg-smooth. Let ηx and ηy be the step sizes for updating x and y, respectively. Denote ∇yg(x, y) as the partial derivative of function g over variable y. ytK represents the inner parameter y in K-th outer loop and t-th inner loop. For the TSGD algorithm A with K outer iterations and T inner iterations, the argument sta-

Remark 6. Analogous to TSGDalgorithm, theUDalgorithm employed in (Bao et al., 2021) (see Algorithm 3) also involves two layers of nested loops but requires re-initialization in the inner level before each new outer loop. In their stability analysis, the inner-level parameter updates are not considered, but used to determine the constants of Lipschitz continuity and smoothness of the outer-level function. This paper considers the general TSGD algorithms where both innerlevel and outer-level parameters are updated continuously, e.g. yTK-1 = y0K . The gradient summation of the inner-level parameter is relatively complex and makes it difficult to utilize the (smooth or convex) properties as (Bao et al., 2021;

Hardt et al., 2016) directly, which brings challenges to the stability analysis.

Theorem 3. Suppose that Assumptions 1 and 2 hold and algorithm A is TSGD with T inner loops and K outer loops. Denote ℓ = max{ℓf , ℓg}, η = max{ηx, ηy}, E[RDm1 ] = EA,Dm1 [RDm1 (xk, yk)].

Let the step sizes satisfy that ρ1 - ρ2 ≤ η ≤ ρ1 + ρ2 = C1 for some positive constant c1, C1. Then A is l1(β) on-average argument-stable in expectation with

and l2(β2 ) on-average argument-stable with

(II) Assume that the bilevel optimization problem is C-C. When η ≤ c2 ln(T) √1+T 2Kℓ for some c2 > 0, A is l1(β) on-average argument-stable in expectation with

and is l2(β2 ) on-average argument-stable with β2 =

(III) Assume that the bilevel optimization problem is NCNC. Denote ηy,t as the inner step size in t-th inner loop, denote ηk = max{ηx,k, ηy,k} as the outer step size in k-th outer loop. Let ηy,t ≤ c3 ℓg(t+1) , ηk ≤ c4 ℓk and ∑bx=a f(x) ≤ c5 ∫ b a f(x)dx for some positive constants c3, c4, c5, then A is l1(β) on-average argument-stable in expectation with

and l2(β2 ) on-average argument-stable with β2 =

Notice that T c6 √1+T 2 with ηy,t ≤ c3 ℓg(t+1) is obtained from Lemma 6 for NC-NC in Appendix D, where the original form is T c0c1 √1+T 2 with ηy,t ≤ c0 ℓg(t+1) .

Remark 7. After integrating Theorems 1 and 3, we summarize the generalization bounds of Algorithm 2 in Table 1. Similar to Theorem 2, the results of Theorem 3 also demonstrate that the total numbers of the validation samples m1 (↑), the inner iterations T (↓) and outer iterations K (↓) directly affect the generalization performance (↑) of TSGD algorithms. We also observe that the impacts of K and T on generalization are suppressed for Algorithm 2 with SC-SC (or C-C) with a small enough step size. In order to obtain tighter bounds w.r.t. Theorems 2 and 3, we further derive the corresponding results with refined step sizes in Propositions 1 and 2 in Appendix C, D. The results shown in Table 1 are comparable to Bao et al. (2021) with the bound of O( Kc m ) where 0 < c < 1. Relaxing the stepsize limitations, especially for SC-SC, is a meaningful direction, which is left for future work.

5 Empirical Evaluations

This section empirically validates our theoretical findings on two real-world datasets. We consider Algorithm 2 here, since it is equal to Algorithm 1 as T = 1. The distributions of the testing samples and validation samples are assumed to be same, but can differ from the training data (Ren et al., 2018; Bao et al., 2021). Similar to (Bao et al., 2021), we focus on evaluating the generalization behavior of outer-level problems based on the validation set. All experiments are implemented in Python on an Intel Core i7 with 32 GB memory. Implemented codes (including (Bao et al., 2021) for hyperparameter optimization) and data sets (including the MNIST data (LeCun, 1998) and the Omnilot data (Lake et al., 2015)) are from publicly available sources.

This section considers the general hyperparameter optimization formulation (Ji et al., 2021). Given the training set Dtrain and the validation set Dval, the hyperparameter optimization scheme can be formulated as

where h is the loss function, Ωy,x is the regularizer and |Dtrain | represents the size of training data.

5.1 Experiment Settings

We evaluate the impact of several factors on the generalization gap (2) based on the famous MNIST data (LeCun, 1998), which totally consists of more than 6 × 105 handwritten figures with the size of 28 × 28. Following the same task of data reweighting in (Bao et al., 2021), we corrupt the labels of training samples randomly with the probability of 50% and employ a fully connected network (with size of 784/256/10) with cross-entropy loss for classification. Initially, we randomly select 2000, 2000, and 1000 figures for training, validation and testing, respectively. Meanwhile, set the initial batch size as 8, the maximum number of inner iterations as T = 5000, and the number of outer iterations as K = 5000. The initial step sizes for inner and outer minimization problems are 0.01 and 5, respectively. For the given parameter settings, each experiment is randomly repeated five times on one GeForce GTX 1660 SUPER GPU, and the average results are reported.

5.2 Experimental Results

The generalization gap defined in (2) is estimated by the divergence between the validation error and the testing error.

Impact of iteration numbers K and T. Now we evaluate the impact of parameters (e.g., the numbers of validation samples m1, inner iteration T and outer iteration K) on the generalization performance. Figure 1 shows the curves of validation error, testing error and the generalization gap under different settings of maximum inner iteration T and maximum outer loop K. Figures 1(a) and 1(b) imply that the classification model might be overfitting with increasing testing errors as K > 3000 and T = 32. Besides, Figure 1(c) demonstrates that too large K and T may reduce the generalization ability of the hyperparameter optimization method due to overfitting. This empirical finding is consistent with our theoretical results and the previous related analysis (Franceschi et al., 2018; Bao et al., 2021).

Impact of sample size m1 with T = 1. Figure 2 presents the results of SSGD (Algorithm 1) under different choices of K and m1. From Figure 2(c), we observe that a small sample size m1 = 500 leads to an increase in validation error and testing error. This indicates that the larger number of validation samples is beneficial to reduce the generalization gap. The above empirical findings match our theoretical results, see e.g., Theorem 3 and Table 1.

Based on theoretical analysis and empirical evaluations, we can get some understanding of the generalization performance of bilevel optimization. Explicitly, the generalization ability of SBO often can be improved with the increase of m1 and proper iteration numbers K, T , where too small iterations may cause underfitting and large ones can lead to overfitting. Usually, it is beneficial for generalization through setting appropriate learning rates, especially for NC-NC. For real applications, the trade-off between m1, T, and K is of great importance to guarantee the effectiveness of SBO methods.

6 Conclusion

This paper established the stability and generalization analysis for stochastic bilevel optimization with first-order gradient-based approximate algorithms. Our theoretical results are obtained by developing the analysis technique associated with the on-average argument stability, and can cover wider bilevel optimization algorithms under low noise settings. Compared with the state-of-the-art analysis (Bao et al., 2021), our theoretical results do not require reinitializing the inner-level parameter before each iteration and suit for objective functions under milder conditions.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (Nos. 62376104 and 12071166), the Fundamental Research Funds for the Central Universities of China (No. 2662023LXPY005), and HZAU-AGIS Cooperation Fund (No. SZYJY2023010).