Paper: Exploring the Generalization Capabilities of AID-based Bi-level Optimization (ICLR-2024)

Reviewer comment (Theory):
Even though the algorithm is complex with several intermediate variables involved in the analysis, the proof techniques for each are analogous to [1] and some existing stability-based work [3-5]. The employed argument-based stability tool has also been used previously [4]. The authors are suggested to further add some descriptions of the employed assumptions (e.g., the e-accuracy from [2], Assumption 3 from …), existing stability-based work and provide some discussions on the differences with them. Besides, it is kindly suggested to highlight the challenges and difficulties of the proof so as to help the readers understand the contributions of this work.

Author response:
First, our stability analysis is far different from [1], where [1] reduces the bi-level problem into a single-level problem. More specifically, in [1], $y_t$ can be written as a function of $x_t$, i.e. $y_t = F(x_t)$. By analyzing the property of function $F$, the problem can be reduced to a single-level problem. However, as $y_t$ depends on all previous iterations in the AID-based algorithm, we use a very different way, by making the connection of successive iterates on all $x,y$ and $m$.
 We do borrow some proof from single-level stability analysis from [3-4]. However, we have to deal with the case from a single variable to 3 correlated variables, where they use different update rules and even based on different functions.

Reviewer comment (Writing):
As described in the contribution, clarification is needed regarding from which perspective in this paper the robustness of the AID algorithm can be demonstrated. Additionally, is the so-called robustness referring to adversary robustness or some form of statistical robustness [1]?

Author response:
Sorry for the confusion, we use robustness to show more stable on changing one data sample in the dataset. Since it will cause confusion, we deleted the word in the contribution.

Reviewer comment (Theory):
In Definition 1, only perturbation in the validation (or outer) dataset is considered. Why not consider perturbation in the inner-level training dataset? Some discussions of the reasons for this definition are needed. In fact, there are prior works that consider perturbation on both datasets for a special case of bilevel learning, meta learning (see [1]).

Author response:
As we discuss in Section 3.2, usually we have much more data in the lower-level problem than the upper-level problem. It is important to first consider the bottleneck, the upper-level problem. [1] gives an example of how the 'inner loop' will affect the difference of iterates, but to consider the inner loop, we have to first define a proper decomposition rule like Section 3.2. The most difficult part is that (III) in Section 3.2 may not be non-positive when changing the constraints for constrained optimization. Thus, we decided to analyze the generalization of the inner level to future work.

Reviewer comment (Novelty):
Algorithms 1 & 2 are not the same as the previous works [1-3]. While it is understandable that approximation algorithms are not unique for AID algorithms (see Table 1 in [2]), the authors should indicate in detail which related works have proposed or used these algorithms.

Author response:
For AID-based algorithm, there are different ways to approximate $\nabla_{yy} g(x,y)^{-1}\nabla_y f(x,y)$. In algorithm 1, we use gradient descent to approximate it. In Table 1 in [2], our algorithm belongs to CG row instead of taking conjugate gradient descent, we use gradient descent for simplicity.

Reviewer comment (Theory):
Although the authors provide an example of a non-convex upper problem, it would be interesting and necessary to further analyze other conditions. For example, in convergence analysis, some works conduct asymptotic analysis under assumptions that the lower and upper functions are convex and strongly convex [1,2]. [3] studied the setting where the average version Φ() is strongly convex or convex, and g(x,·) is strongly convex. It would be better if this issue were considered, though it is not necessarily required.

Author response:
Thank you for the suggestion. The difficulty of the analysis is that the property on $\Phi$ only controls the function property of $f(x,y)$ when y is near to $y*(x)$. However, the stability analysis needs to hold on arbitrary (x,y). Thus, the condition that $\Phi$ is strongly convex is not enough to give a better stability bound. We will give some more conditions for strongly convex cases in future work.

Reviewer comment (Theory):
In Corollary 1, the constants α and β are involved in the learning rate without detailed definition. Is the β here the same as the stability definition in Definition 1 or 2?

Author response:
The $\alpha$ and $\beta$ are two constants in the learning rate, which will affect the value of $q = \frac{2C_m\alpha + L_1\beta}{2C_m\alpha + L_1\beta+1}$ in Corollary 1.

Reviewer comment (Theory):
In Appendix A, the conclusion of Lemma 8 seems to be wrongly used in the following lemmas and Theorem 1. Dz could only bound z^K, which does not hold for z^0 or z^k where k < K. This could be the most serious weakness and may render some arguments in the paper invalid.

Author response:
Thanks for pointing it out, we have updated Lemma 8. Now, it holds for all k that $||z_t^k|| \leq D_z :=||z_0|| + \frac{L_0}{\mu}$. Since $\|z_0\|$ and $\frac{L_0}{\mu}$ are constants, it will not affect the latter proof.

Reviewer comment (Theory):
In the proof of Lemma 11 (Appendix, p.14), the rules for updating variable ‘m’ seem contradictory with Algorithm 1. Additionally, the mixed usage of y and y*(x) causes confusion (see Lemma 16 and Theorem 2).

Author response:
Sorry for the typos in the algorithm, we have updated Algorithm 1.

Reviewer comment (Writing):
Potential typos or mistakes have been identified as follows, and the authors are asked to check these issues during revision:  
- In Remark 3, change \Theta -> O  
- In the proof of Lemma 8 (p.12), A is undefined  
- In the proof of Lemma 9 (p.12), the symbol ‘+’ before \eta_z is missing; on p.13 line 3, the ‘=’ symbol is missing  
- In the proof of Lemma 10 (p.13), the symbol ‘<=’ is missing; the description “By selecting \eta_y such that \mu \eta_y ≥ L_2 \eta_y” seems incorrect; on p.14, x should lie in [0,0.5] in the first line; ‘2’ should be replaced by ‘K+1’ in the fourth line  
- In the proof of Lemma 11 (p.14), rules for updating m seem contradictory with Algorithm 1  
- In the proof of Lemma 15 (p.18), distinguish between ‘X’ and ‘x’  
- In the proof of Lemma 16 (p.18), replace ‘\nabla_t’ with ‘\nabla_y’; replace single ‘\nabla’ with ‘\nabla^2_yy’

Author response:
Thanks for you pointing them out, we have updated the paper accordingly.

Reviewer comment (Theory):
The influence of the inner iteration K is not discussed. In Theorem 2, the paper only gives 'K is large enough'. However, in the experiments, the author only uses K=2. I think this is confumsing. Is there any relation between K and T? [1] proposed a method where K=1, so may be 'K is large enough' is not needed.

Author response:
In Therorem 2, we have added the theorectical results for $K = \Theta(\log T)$, and we increase the K in both experiments. The results are shown in Appendix, the results indicate that increasing K can help get smaller error, while the diminishing learning rate is indeed better than constant learning rate even K is large.  
[1] gives the results when K = 1, because they assume higher order smothness like $\nabla^3 g$ is Lipschitz smooth and $\nabla^2 f$ is Lipschitz smooth. With weaker condition, we can only show the K with a large value.

Reviewer comment (Experiments):
If K is needed to be large, i think the performance with different K should be given.

Author response:
In Therorem 2, we have added the theorectical results for $K = \Theta(\log T)$, and we increase the K in both experiments. The results are shown in Appendix, the results indicate that increasing K can help get smaller error, while the diminishing learning rate is indeed better than constant learning rate even K is large.  
[1] gives the results when K = 1, because they assume higher order smothness like $\nabla^3 g$ is Lipschitz smooth and $\nabla^2 f$ is Lipschitz smooth. With weaker condition, we can only show the K with a large value.

Reviewer comment (Theory):
The influence of the inner iteration K is not discussed. In Theorem 2, the paper only gives 'K is large enough'. However, in the experiments, the author only uses K=2. I think this is confumsing. Is there any relation between K and T? [1] proposed a method where K=1, so may be 'K is large enough' is not needed.

Author response:
In Therorem 2, we have added the theorectical results for $K = \Theta(\log T)$, and we increase the K in both experiments. The results are shown in Appendix, the results indicate that increasing K can help get smaller error, while the diminishing learning rate is indeed better than constant learning rate even K is large.  
[1] gives the results when K = 1, because they assume higher order smothness like $\nabla^3 g$ is Lipschitz smooth and $\nabla^2 f$ is Lipschitz smooth. With weaker condition, we can only show the K with a large value.

Reviewer comment (Theory):
This analysis in this paper appears heavily reliant on smoothness assumptions for the inner and outer functions, the theoretical results may be not suited to loss function with L1 penalty. Therefore, the results in this paper do not yet seem generalizable to the broader field of artificial intelligence.

Author response:
Thank you for pointing it out. Since as far as we know, this is the first paper that analyzes the generalization of AID-based bilevel algorithms, we start from the simple case where the function has nice smooth properties. We have added the suggestion in the future work in conclusion that there is still a mystery for the proper choice of stepsize, and for the weaker conditions of the upper-level and lower-level objective function, we will leave for future work.

Reviewer comment (Evaluation):
The paper does not seem to consider the generalization ability of the inner problem, but sometimes the solution to the inner problem is equally important. Therefore, it would be better if the generalization ability of the lower problem could also be analyzed.

Author response:
As the inner problem will have a much larger dataset than the upper level, it is more important to analyze the upper level first. After that, we can further analyze the inner-level problem. Different from the upper-level objective, the inner-level objective is considered as a constraint, such the generalization needs to be defined first on the constraints for constrained optimization problems. We will consider it in future work.

Reviewer comment (Theory):
The paper considers the trade-off between convergence and stability of AID based algorithms, but only gives the cases of constant learning rate and diminishing learning rate, without discussing the impact of decaying learning rate at different rates to choose better learning rate.

Author response:
For stability results and convergence results, the formulations are complicated, where statbility is in order of $O(\sum_{t=1}^T \Pi_{k=t+1}^T (1+ \eta_{x_k}\eta_{m_k}+\eta_{m_k} + \eta_{y_k} L_1)(1+\eta_{x_t})\eta_{m_t}/n)$ and convergence is in order of $\mathcal{O}(\frac{1+\sum_{k=1}^T {\eta_{y_k}\eta_{m_k} + \eta_{m_k}^2}}{\sum_{k=1}^T\eta_{m_k}})$. Thus, we just analyze some popular learning rate choices instead of getting an optimal schedule. We will leave it to future work.

Reviewer comment (Experiments):
Might it be feasible to incorporate more challenging datasets and tasks in order to validate the theoretical analysis, such as CIFAR?

Author response:
Due to the limited time, we added an experiment in the Appendix. We use dataset CIFAR10 and the rest settings are the same as MNIST. The results show that with diminishing step size, the algorithm can perform better than constant step size.

Reviewer comment (Theory):
Will there be more significant challenges in extending the conclusions of this paper to AID-based bi-level optimization variants?

Author response:
We are sorry that we get confused about the AID-based bi-level optimization variants. However, we think for any small modification of the algorithm, the whole proof needs to be updated with at least 2-3 lemmas.

Reviewer comment (Novelty):
Why do the authors focus on the specific double-loop AID-based bi-level optimization algorithm in Algorithm 1? There are several single-loop AID-based bi-level optimization algorithms available, as demonstrated in Dagréou et al.'s work ``A framework for bilevel optimization that enables stochastic and global variance reduction algorithms" in NeurIPS (2022), as well as the related references.

Author response:
1. Why do the authors focus on the specific double-loop AID-based bi-level optimization algorithm in Algorithm 1? There are several single-loop AID-based bi-level optimization algorithms available, as demonstrated in Dagréou et al.'s work ``A framework for bilevel optimization that enables stochastic and global variance reduction algorithms" in NeurIPS (2022), as well as the related references.  
> Different from ``A framework for bilevel optimization that enables stochastic and global variance reduction algorithms", we do not assume the Lipschitz continuity of $\nabla^2_x f(x,y)$ and $\nabla^3_y g(x,y)$. With weaker assumptions, it is hard to prove the convergence of a single-loop algorithm.  We have added the discussion to the related work that  Dagréou et al. (2022)  show the convergence under higher-order smoothness setting.

Reviewer comment (Theory):
Why are the solutions of bi-level optimization problems in Section 3.2 unique, considering the definitions of $(\bar{x}, \bar{y})$ and $(x^*, y^*)$? Note that Remark 1 mentions that the bi-level optimization problem is likely to have a nonconvex objective with respect to $x$.

Author response:
2. Why are the solutions of bi-level optimization problems in Section 3.2 unique, considering the definitions of $(\bar{x}, \bar{y})$ and $(x^*,y^*)$? Note that Remark 1 mentions that the bi-level optimization problem is likely to have a nonconvex objective with respect to x.  
> Sorry for this confusion, we define $(\bar{x},\bar{y})$ and $(x^*,y^*)$ as one of the optimal solution, where we change $=$ into $\in$. Because any optimal solution in the set shares the same function value, the decomposition will remain the same in Section 3.2.

Reviewer comment (Theory):
Why can we use the same sample for different datasets? Note that there is an internal randomness of algorithm $\mathcal{A}$ (i.e., Algorithm 1). The proof of Theorem 1 should be more rigorous. The proof of Theorem 2.2 in Hardt et al. (2016) may provide valuable insights.

Author response:
3.  Why can we use the same sample for different datasets? Note that there is an internal randomness of algorithm $\mathcal{A}$((i.e., Algorithm 1). The proof of Theorem 1 should be more rigorous. The proof of Theorem 2.2 in Hardt et al. (2016) may provide valuable insights.  
> Based on the definition, we want to bound $\mathbb{E}_\mathcal{A} ||\mathcal{A}(D_t,D_v) - \mathcal{A}(D_t,D_v')||$, where for each instance of algorithm $\mathcal{A}$, two trajectories share the same randomness. In practice, the randomness of $\mathcal{A}$ is the sample index. Thus, two trajectories either sample the same data point or sample the only different data point in the dataset. This technique is identical to the proof of Hardt et al.(2016).

Reviewer comment (Theory):
Where is the proof of Proposition 3? According to Definition 2, it's not immediately obvious, as there are different samples in $f$, and the assumptions on $f$ are made for the same sample.

Author response:
(4) Where is the proof of Proposition 3? According to Definition 2, it's not immediately obvious, as there are different samples in $f$, and the assumptions on $f$ are made for the same sample.  
> We have updated Definition 2, it is a typo to use $D_v$ and $D_v'$ in f. The correct version is $|| E _{\mathcal{A}, D _v \sim P _{D _v} ^n, D' _v \sim P _{D _v}^n}[ f(\mathcal{A}(D _t,D _v),z) -  f(\mathcal{A}(D _t,D _v'),z)]||\leq\beta,\ \forall D _t \in \mathcal{Z} _t^q, z \in Z _v.$ Based on the new definition 2, it is obvious that Proposition 3 is correct.

Reviewer comment (Writing):
Equation (2) in Theorem 2: $L_{\Phi}$ is not defined before, so clarify its meaning.

Author response:
(5) Equation (2) in Theorem 2: $L_\Phi$ is not defined before, so clarify its meaning.  
> Similar to the gradient descent, we have a necessary condition to ensure descent that the step size is less than or equal to $1/L$, where $L$ is the gradient Lipschitz constant of the function. Here, we prove that the gradient of function $\Phi$ is Lipschitz with constant $L_\Phi$. We have updated the value of $L_\Phi$ in Thereom 2.

Reviewer comment (Theory):
Where is the proof of Corollary 2? At the end of the proof of Corollary 3, how does $e^{\frac{\alpha}{\epsilon \gamma}}$ become $\mathcal{O}(e^{1/\epsilon})$? Is this statement correct?

Author response:
(7) Where is the proof of Corollary 2? At the end of the proof of Corollary 3, how does $e ^{\alpha/(\gamma\epsilon)}$ become $O(e ^{1/\epsilon})$? Is this statement correct?  
> We have updated the statement in Corollary 2 and Corollary 3. Instead of giving the order of $\epsilon _{stab}$, we give the order of $\log \epsilon _{stab}$. The proof of Corollary 2 has been added to the Appendix. It is correct that $\log e ^{\frac{\alpha}{\gamma\epsilon}}$ becomes $O(1/\epsilon)$.

Reviewer comment (Theory):
In the toy example in Section 5.1, why is $(\hat{X}, \hat{y})$ considered the ground truth? Please note that there is an $L_2$ regularization term in the context.

Author response:
(8)Toy example in Section 5.1: Why is $(\hat{X},\hat{y})$ considered the ground truth? Please note that there is an $L_2$ regularization term in the context.  
> As the goal is to recover $\hat{X}$ and $\hat{y}$, we consider $(\hat{X}, \hat{y}))$ as the ground truth. $L_2$ regularization in the lower-level problem is for generalization as a prior because of limited data. It is a well-known trick to improve generalization which helps $y$ close to $\hat{y}$. $L_2$ regularization in the upper-level problem is for softening the constraint that $X$ is orthogonal, which helps $X$ close $\hat{X}$. Thus, $L_2$ regularization helps to find the ground truth.

Reviewer comment (Theory):
Does the lower-level objective in Section 5.2 is strongly convex for $y$?

Author response:
(9)Does the lower-level objective in Section 5.2 strongly convex for $y$?  
> No, the lower-level objective in Section 5.2 is not strongly convex for $y$. The experiment is to show that even without the strong convexity in the lower-level objective, our conclusion still holds.