Paper: Bilevel Optimization under Unbounded Smoothness: A New Algorithm and Convergence Analysis (ICLR-2024)

Reviewer comment (Evaluation):
Too many loops in Algorithm 1,2 which may lead to time-consuming.

Author response:
**A**: There is a minor issue in the pseudo-code for Algorithm 2 in the original version of the paper and we are sorry for the ambiguity and confusion brought to the readers. The outer loop in Algorithm 2 with the iteration variable $k$ spanning from $0$ to $K-1$ is indeed repeatedly defined, and we already removed it in the new version of our paper. Algorithm 1 is actually a double-loop algorithm instead of a triple-loop one. Notably, we only need to execute the inner loop occasionally since updating for the lower-level variable is periodic rather than in each iteration, which helps reduce the running time.

Reviewer comment (Theory):
How will the iteration numbers affect the convergence performance? Is the inner loop necessary? Because this method seems to be a momentum method based on SOBA, why author use an inner loop? why not just update $y$ for a single step?

Author response:
**A**: Thank you for your insightful question. The inner loop for periodic updates is quite important for convergence analysis. We will explain why we cannot use single-loop based on our current proof strategy. The main difficulty of the paper comes from the bias term $\\|\mathbb{E}\_k[\widehat{\nabla}\Phi(x_k)]-\nabla\Phi(x_k)\\|$ of the hypergradient estimator, whose upper bound depends on both $L_{x,1}\\|y_k-y_k^*\\|\\|\nabla\Phi(x_k)\\|$ and $\\|y_k-y_k^*\\|$. Please check the proof of Lemma 4 in Appendix D.6 (in the middle of page 33, between inequality (ii) and (iii)) for details. Due to the fact that both $\\|y_k-y_k^*\\|$ and $\\|\nabla \Phi(x_k)\\|$ are measurable with respect to $\mathcal{F}_{k-1}$, we have to use high probability analysis for $y_k$ to disentangle these two terms. 

If we use single-loop to solve our algorithm, then for every iteration we update $y_k$ just via one step of stochastic gradient descent. By the high probability analysis of SGD for the lower-level problem (i.e., Lemma 12 in Appendix D.2), if we choose $
\lambda = \max\left(\sqrt{3\ln\left(\frac{2K}{\delta}\right)}, \ln\left(\frac{2K}{\delta}\right)\right)$, $\alpha_s=\gamma$ and $T_s=1$ (i.e., single loop) in formula (19) and (20) in Lemma 12, then with probability at least $1-\delta/K$, we have 

$$     
 \frac{\mu}{2}\\|y_1-y_0^*\\|^2 \leq g(x_0,y_1)-g(x_0,y_0^*)  \leq \frac{1}{\gamma}D^2+2\sigma_{g,1}^2\gamma+\lambda\left(2\sqrt{2}D\sigma_{g,1}+2\sigma_{g,1}^2\gamma\right) \\
$$
$$
=\frac{1}{\gamma}D^2+2\sigma_{g,1}^2(\lambda+1)\gamma+2\sqrt{2}\lambda\sigma_{g,1}D = 2\sqrt{2(\lambda+1)}\sigma_{g,1}D+2\sqrt{2}\lambda\sigma_{g,1}D \\
$$
$$
= 2\sqrt{2}(\lambda+\sqrt{\lambda+1})\sigma_{g,1}D,
$$
where $D=\epsilon/8K_0$ is the upper bound for $\|y_0-y_0^*\|$, and we choose $\gamma=\frac{D}{\sigma\sqrt{2(\lambda+1)}}$ that minimizes the value of the right-hand-side of the inequality. Then we have 
$$
\\|y_1-y_1^*\\| \leq \\|y_1-y_0^*\\| + \\|y_0^*-y_1^*\\| \leq \sqrt{4\sqrt{2}(\lambda+\sqrt{\lambda+1})\sigma_{g,1}/\mu}\sqrt{D} + \frac{C_{g_{xy}}}{\mu}\eta
$$
$$
 = \sqrt{\sqrt{2}(\lambda+\sqrt{\lambda+1})\sigma_{g,1}/2K_0\mu}\sqrt{\epsilon} + \frac{C_{g_{xy}}}{\mu}\eta = \widetilde{\mathcal{O}}(\epsilon^{1/2}),
$$
where the last equation holds because $\eta=\mathcal{O}(\epsilon^3)$ and hence the second term is dominated by the first term when $\epsilon$ is small. Then repeat the same process, with probability at least $1-2\delta/K$, we have $\\|y_2-y_2^*\\|=\widetilde{\mathcal{O}}(\epsilon^{1/4})$. Generally speaking, we have with probability at least $1-k\delta/K$, $\\|y_k-y_k^*\\|=\widetilde{\mathcal{O}}(\epsilon^{1/2^k})$. Since for $\epsilon<1$, we have $\epsilon < \epsilon^{1/2} < \epsilon^{1/2^2} \leq \cdots \leq \epsilon^{1/2^K}$, this means that with high probability, $\frac{1}{K}\sum_{k=0}^{K-1}\\|y_k-y_k^*\\| = \frac{1}{K}\sum_{k=0}^{K-1}\mathcal{O}(\epsilon^{1/2^k})$, and we cannot conclude that $\frac{1}{K}\sum_{k=0}^{K-1}\\|y_k-y_k^*\\| = \mathcal{O}(\epsilon)$ when $K$ is large (note that $\mathcal{O}(\epsilon)$ bound for this quantity is required in our proof to obtain the $\widetilde{\mathcal{O}}(\epsilon^{-4})$ complexity). Therefore the single loop algorithm does not work if we follow our current proof strategy.

We also conduct the ablation study to figure out how the iteration numbers affect the performance. The empirical result is shown in Figure 6 (b), where $N$ is the iteration numbers and $I$ is the update period. The ablation experiment is implemented on the hyper-representation task. The blue line (with $N=1, I=1$) denotes the result of the single update for $y$. Then we fix the update period $I=2$ (which is fixed as $2$ in Section 4) and increase the number of inner loops. The result shows that the performance of BO-REP algorithm would increase as the number of inner loops. The algorithm can achieve the best performance when $N\geq 5$. In our experiments, it is good enough to choose $N=3$ to achieve good performance universally for all tasks.

Reviewer comment (Experiments):
Why only use the SOBA in [1]. SABA also proposed in [1]. I think it should be compared.

Author response:
**A**:
We have added the SABA baseline in our experiments, please refer to Figure 1 in the new version. SABA slightly improves over SOBA but it is still much worse than our algorithm BO-REP.

Reviewer comment (Theory):
How to ensure the problem considered in experiments is unbounded smooth? Because the applications considered in this paper are usually considered in other papers with bound smoothness assumption.

Author response:
**A**: First, we want to emphasize that the unbounded smoothness assumptions were empirically observed in many prior works [3, 2] for RNNs, LSTMs and Transformers. Second, we have empirically verified that the Recurrent Neural Network model satisfies Assumption 1 in the revised version. In Figure 4, we plot the estimated smoothness at different iterations during training neural networks. In particular, we conduct the hyper-representation experiments to verify Assumption 1. We adopt the same recurrent neural network as that in section 4.1. The lower-level variable is the model parameter of the last linear layer, and the upper-level variable is the model parameter of the previous layers (2 hidden layers). In each training iteration, we calculate the gradient norm w.r.t. the upper variable $\\|\nabla_x f(x, y)\\|$ and the lower variable $\\|\nabla_y f(x, y)\\|$, and use same the method as described in Appendix H.3 in [3] to estimate the smoothness of $x$ and $y$. We can get consistent results which are shown in Figure 1 in [3] and Figure 1 in [2]. In particular, we can find that the smoothness parameter scales linearly in terms of gradient norm for both layers. This verifies the Assumption 1 empirically.

Reviewer comment (Theory):
The theoretical complexity result of the proposed algorithm, Theorem 1, lacks persuasiveness. Specifically, the choice of the iteration number $N$ in Algorithm 2 for lower-level variable updates is left unspecified. Furthermore, the selection of the parameter $I$ appears improper, particularly in comparison to the selection of the parameter $K$. This peculiarity arises because, according to Theorem 1, $I = O(1/\epsilon^2)$, while $K = O(1/\epsilon)$. Consequently, for sufficiently small values of $\epsilon$, $I$ is significantly larger than $K$. However, both the implementation of the proposed algorithm and the theoretical analysis, as exemplified by Lemma 2, necessitate $I$ to be smaller than $K$.

Author response:
**A**:  Thank you for pointing it out. We actually state the choice of the iteration number $N$ in the proof of Lemma 2, but without explicitly pointing it out in statement of Lemma 2.  We have fixed it in the new version of the paper.
For the exact value of $N$, you can check equation (45) and (50) in proof of Lemma 2 in Appendix D.3. The iteration complexity $N$ in Algorithm 2 for lower-level variable updates (within a single subroutine) should be (by eq. (45) and (50)) $N = 64^2\sigma_{g,1}^2K_0^2\left(\lambda+\sqrt{\lambda+1}\right)^2/\mu^2\epsilon^2$, where $\lambda$ is some logarithmic factors. Thus we conclude $N=\widetilde{\mathcal{O}}(\sigma_{g,1}^2K_0^2/\mu^2\epsilon^2)$ in Algorithm 2.  

 For Theorem 1, $I$ is indeed much smaller than $K$ for small enough $\epsilon$. Note that we choose update period $I$ to be $I = \sigma_{g,1}^2K_0^2/\mu^2\epsilon^2 = \mathcal{O}(1/\epsilon^2)$, and there is a learning rate dependency (i.e. $1/\eta$) in $K$: $K=\frac{4\Delta}{\eta\epsilon}=\mathcal{O}(1/\epsilon^4)$ because $\eta=\mathcal{O}(\epsilon^3)$. In this case, $I$ is much smaller than $K$ for small enough $\epsilon$.

Reviewer comment (Theory):
Despite the authors' assertion that they do not require the boundedness of the norm of the gradient of the upper-level objective function, $ || \nabla_y f(x, y) ||$, as stipulated in Assumption 2 (i), it is worth noting that they still demand the boundedness of $|| \nabla_y f(x, y^*(x)) ||$. The rationale behind the relative ease of achieving this condition in practical scenarios, as opposed to the former requirement, remains unclear. Moreover, it remains unaddressed whether the problems examined in the numerical experiments exhibit unbounded smoothness yet satisfy Assumption 2 (i).

Author response:
**A**: Thank you for this insightful question. First of all, to the best of our knowledge, this is an assumption (i.e., $\\|\nabla_yf(x,y^*(x))\\|\leq M$) which is required by all bilevel optimization papers [1, 2, 3]. Also, the boundedness of $\\|\nabla_y f(x, y^*(x))\\|$ does not contradict with relaxed smoothness assumption, i.e., Assumption 1. We will illustrate it in the following. As can be seen from Figure 4 in Appendix G, we show that the smoothness parameter scales linearly in terms of the gradient norm of $x$ (parameter of the recurrent layers) and $y$ (parameter of the last linear layer).

In addition, we show that if we denote $x$ as previous layers of recurrent neural networks and $y$ as the last layer (i.e., classification layer), the condition 
$\\|\nabla_yf(x,y^*(x))\\|\leq M$ indeed holds both theoretically and experimentally. Figure 4 (b) in Appendix G shows that the gradient norm in terms of $y$ is bounded, i.e., the maximal value of the horizontal axis can be bounded by $M=10$. Besides, when considering the setting of hyper-representation, we have included a proof in Appendix H, indicating that this condition holds true for any $(x, y)$ and thus also holds for $(x, y^*(x))$.

[1] *Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. arXiv preprint arXiv:1802.02246, 2018.*

[2] *Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In International conference on machine learning, pp. 4882–4892. PMLR, 2021.*

[3] *Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert D Nowak. A fully first-order method for stochastic bilevel optimization. In International Conference on Machine Learning, pp. 18083-18113. PMLR, 2023.*

Reviewer comment (Experiments):
The selection of the parameters $I$ and $N$ in the implementation of the proposed algorithm in the numerical experiments deviates from the parameters required by the theoretical analysis. Specifically, the theoretical analysis mandates that both $I$ and $N$ should be large as $\epsilon$ becomes small. However, in the context of the numerical experiments, the authors have set the values for $I$ and $N$ as merely 2 or 3, which contradicts the theoretical analysis.

Author response:
**A**:  We respectfully disagree. Indeed, by Lemma 2 and Theorem 1, we know that $I = \sigma_{g,1}^2K_0^2/\mu^2\epsilon^2 = \mathcal{O}(1/\epsilon^2)$ and $N = 64^2\sigma_{g,1}^2K_0^2\left(\lambda+\sqrt{\lambda+1}\right)^2/\mu^2\epsilon^2 = \widetilde{\mathcal{O}}(1/\epsilon^2)$. However, we do not think the practical values of $I$ and $N$ contradict the theoretical analysis. Big $\mathcal{O}$ notation only gives a conservative upper bound which makes sure that the algorithm can converge even we are given the worst function instance from certain function class, but in practice, the function is never adversarial such that we can choose smaller $I$ and $N$ and hence the algorithm in fact enjoys low complexity. This is not uncommon in the literature of computer science. For example, the simplex method is very efficient in solving linear programming in practice, but it takes exponential time in the worst case from complexity point of view: this is not a contradiction, and it also inspires the seminal smoothed analysis framework [4].

[4] *Spielman, Daniel A., and Shang-Hua Teng. "Smoothed analysis of algorithms: Why the simplex algorithm usually takes polynomial time." Journal of the ACM (JACM) 51, no. 3 (2004): 385-463.*

Reviewer comment (Evaluation):
The proposed algorithm incorporates an initialization refinement procedure for the lower-level variable. I am curious about whether the authors have taken into account the computational cost associated with this initialization refinement procedure when presenting the numerical results in Figure 1.

Author response:
**A**: For a fair comparison between different baselines, we have included loss curves in terms of running time in Figure 6 in Appendix I, where the loss curve of BO-REP takes the time cost of initialization refinement into consideration. Our proposed BO-REP achieves the best performance among all the baselines and shows the fastest convergence in various experiments, including Hyper-Representation, hyperparameter optimization and Data Hyper-Cleaning.

Reviewer comment (Evaluation):
While the proposed algorithm exhibits advantages over its competitors in terms of the number of epochs required, how about its performance concerning computation time?

Author response:
**A**:   We compare our proposed algorithm with other single-loop and double-loop algorithms in terms of running time, the result is shown in Figure 6 in Appendix I. To accurately evaluate each algorithm, we use the machine learning framework PyTorch 1.13 to run each algorithm individually on an NVIDIA RTX A6000 graphics card, and record its training and test loss. As we can observe from the figure, our algorithm (BO-REP) is much faster than all other baselines in the experiment of Hyper-representation (Figure 5 (a)) and Data Hyper-Cleaning (Figure 5 (c)). For the hyperparameter optimization experiment (Figure 5 (b)), our algorithm is slightly slower at the very beginning, but quickly outperforms all other baselines. This means that our algorithm indeed has better runtime performance than the existing bilevel baselines.

Reviewer comment (Theory):
This paper relaxes the bounded assumptions but is not based unbounded assumption as stated in the abstract. The example pf RNN and LSTMs are misleading unless the authors are able to prove that these networks satisfy the Assumption 1.

Author response:
**A**: Thank you for your suggestion. First, we want to emphasize that the unbounded smoothness assumptions were empirically observed in many prior works ([1, 2]) for RNNs, LSTMs and Transformers. Second, we have empirically verified that the Recurrent Neural Network model satisfies Assumption 1 in the revised version. In Figure 4 in Appendix G, we plot the estimated smoothness at different iterations during training neural networks. In particular, we conduct the hyper-representation experiments to verify Assumption 1. We adopt the same recurrent neural network as that in Section 4.1. The lower-level variable is the parameter of the last linear layer, and the upper-level variable is the parameter of the previous layers (2 hidden layers). In each training iteration, we calculate the gradient norm w.r.t. the upper variable $\\|\nabla_x f(x,y)\\|$ and the lower variable $\\|\nabla_y f(x,y)\\|$, and use the same method as described in Appendix H.3 in [1] to estimate the smoothness of $x$ and $y$. From Figure 4 in Appendix G, we can find that the smoothness parameter scales linearly in terms of gradient norm for both layers. This verifies the Assumption 1 empirically. Note that these results are consistent with the results in the literature, e.g., Figure 1 in [1] and Figure 1 in [2].

Reviewer comment (Novelty):
Momentum technique has been incorporated in the bilevel optimization area before. For example, MRBO [1] achieves the $\mathcal{\widetilde O}(1/\epsilon^3)$ after momentum introduction. The authors are expected to improve the convergence rate of proposed algorithm.

Author response:
**A**: Please note that these two momentum techniques are completely different. Our momentum is simply doing moving-average estimate for history stochastic gradients, but the momentum used in [3] is a type of recursive momentum which is based on the STORM estimator [7]. The STORM estimator requires the mean-squared smoothness assumption while our momentum does not. Without this assumption, our complexity result is already optimal up to logarithmic factors and it is impossible to achieve the $\widetilde{\mathcal{O}}(1/\epsilon^3)$ rate as proved in [6]. We could also improve our complexity result to $\widetilde{\mathcal{O}}(1/\epsilon^3)$ by imposing the mean-squared smoothness assumption and adopting the STORM technique, but it is out of the scope of our current paper.

Reviewer comment (Evaluation):
The authors are suggested to show the experimental results w.r.t. time rather than only epochs. The proposed algorithm involves three loops per epoch, which might introduce extra computational cost. It will be more fair to compare proposed algorithm with other single-loop or double-loop based algorithms in terms of running time.

Author response:
**A**: We have to clarify a typo in our manuscript. Please check the algorithm framework in the new version of our paper, we fix the typo in the pseudo-code for Algorithm 2 (remove the unnecessary loop in line 1). Therefore our algorithm is a double-loop algorithm instead of a triple-loop one. Notably, we only need to execute the second loop occasionally since updating for the lower-level variable is periodic rather than in each iteration. 

We compare our proposed algorithm with other single-loop and double-loop algorithms in terms of running time, the result is shown in Figure 6 in Appendix I. To accurately evaluate each algorithm, we use the machine learning framework PyTorch 1.13 to run each algorithm individually on an NVIDIA RTX A6000 graphics card, and record its training and test loss. As we can observe from the figure, our algorithm (BO-REP) is much faster than all other baselines in the experiment of Hyper-representation (Figure 5 (a)) and Data Hyper-Cleaning (Figure 5 (c)). For the hyperparameter optimization experiment (Figure 5 (b)), our algorithm is slightly slower at the very beginning, but quickly outperforms all other baselines. This means that our algorithm indeed has better runtime performance than the existing bilevel baselines.

Reviewer comment (Experiments):
The main drawback of the proposed algorithm is that the structure is complex, and each subroutine contains additional hyperparameters that need to be tuned.

Author response:
**A**: Our problem setting is much more challenging to solve than the standard bilevel optimization for smooth objectives. The reason is that, under Assumption 1, i.e., $(L_{x,0},L_{x,1}, L_{y,0}, L_{y,1})$-smoothness assumption, we have to carefully update both upper-level and lower-level variables. Therefore our proposed algorithm has more complex structure compared to other bilevel algorithms. However, we want to emphasize that our algorithm is not sensitive to additional hyperparameters such as $I$ and $N$ (please see our ablation study in Appendix J). In fact, for each subroutine of the periodic updates, we do not need to carefully tune these values in practice:  some default values of $I$ and $N$ (e.g., $I=2, N=3$) do not seem to hurt the performance for a wide range of bilevel optimization tasks.

Reviewer comment (Experiments):
In the experiments, how is the ball radius R chosen?

Author response:
**A**: The ball radius $R$ is set as $0.5$ in the experiments. We experimentally explore how the ball radius $R$ affects the algorithm's performance. We set the ball radius as $\{0.001, 0.005, 0.01, 0.05, 0.10, 0.50, 1.00, 5.00\}$ respectively, and then conduct the bilevel optimization on Hyper-representation tasks. We keep all the other hyperparameters the same as in Section 4.1. The result is shown in Table 1.  When the ball radius $R$ is too small (i.e., $R<0.10$), the performance significantly drops, possibly due to the overly restricted search space for the lower-level variable.  When the ball radius $R \geq 0.50$, the performance becomes good and stable. In our experiments in the main text, we choose $R=0.50$.

### **Table 1: Test accuracy vs. Projection radius $R$.**
| **Radius**   | $R=0.001$ | $R=0.005$ | $R=0.01$ | $R=0.05$ | $R=0.10$ | $R=0.50$ | $R=1.00$ | $R=5.00$ |
|-------------|----------------|-----------------|---------------|--------------|---------------|---------------|---------------|--------------|
| **Test ACC**| $65.06\\%$ | $67.64\\%$  | $70.76\\%$| $81.22\\%$| $86.84\\%$ | $88.84\\%$ | $88.84\\%$ | $88.84\\%$ |

Reviewer comment (Presentation):
Caption of table: "$\mathcal{C}_L^{a, a}$ denotes a-times differentiability with Lipschitz k-th order derivatives." I think it should be $\mathcal{C}_L^{a, k}$ here.

Author response:
**A**: We have fixed this issue in the new version.

Reviewer comment (Presentation):
In the pseudo-code for Algorithm 2, the outer loop with the iteration variable k spanning from 0 to K-1 is repeatedly defined, which may lead to ambiguity or confusion for the readers. To enhance clarity, it's recommended to pass the current iteration k as an input to the UpdateLower subroutine (which is a part of the BO-REP algorithm) and subsequently remove the loop from the algorithm.

Author response:
**A**: Thank you for catching this! It is indeed a typo in our original manuscript. Please check the algorithm framework in the new version of our paper, we fix the typo in the pseudo-code for Algorithm 2 (remove the unnecessary loop in line 1).

Reviewer comment (Theory):
Lack of lower bounds does not allow to understand optimality of the proposed method under given assumption.

Author response:
**A**: Thank you for your insightful question. We want to emphasize that our oracle complexity of finding an $\epsilon$-stationary point is indeed optimal up to logarithmic factors by directly applying the lower bound results in [1]. Let us give more details here. For the upper-level problem, we assume the stochastic gradient estimator is unbiased and has bounded variance. Under this stochastic gradient oracle, the lower bound for finding $\epsilon$-stationary points in non-convex smooth single-level stochastic optimization with first-order oracle is at least $\Omega(1/\epsilon^4)$ (Theorem 1 of [1]). Since the single-level optimization is a special case of bilevel optimization (e.g., bilevel optimization becomes a single-level optimization when the upper-level function does not depend on the optimal solution of the lower-level problem) and standard smoothness is a special case of relaxed smoothness (e.g., relaxed smoothness with $L_1=0$ correspond to standard smoothness), our $\widetilde{\mathcal{O}}(1/\epsilon^4)$ complexity result is optimal up to logarithmic factors. We have added descriptions in the main paper to make it more clear.

Reviewer comment (Theory):
Proposed guarantees holds only for finding an $\epsilon$-stationary point in expectation.

Author response:
**A**: Please note that we do not assume light tail assumption for upper-level problems in our paper, so our convergence guarantees hold only in expectation. Note that the high probability result even in the single-level stochastic nonconvex optimization setting requires the light tail assumption (e.g., Corollary 2.5 in [2]).

Reviewer comment (Theory):
Is it possible to relax the strong convexity assumption of the lower function to log-concavity?

Author response:
**A**: Strong convexity guarantees that the optimal solution $y^*(x)$ is unique for the lower-level problem: this is very important for our analysis. However, this property does not always hold for log-concave functions. For example, let $g(x,y)=(\sin{x}+2)\exp(-y^2/2)$, then $g(x,y)$ is log-concave in $y$ for any $x$, and $g(x,y)$ also satisfies other assumptions in Assumption 2. For this specific function, $y^*(x)$ does not even exist, since for any fixed $x$, we could not find a minimum for $y$ if $y\in\mathbb{R}^{d_y}$.

Reviewer comment (Theory):
What are challenges to provide high-probability guarantees for $\epsilon$-stationary point for the given algorithm?

Author response:
**A**: For Corollary 2.5 in [2], it requires stochastic estimators to be unbiased and have bounded variance, and also the light-tail assumption for stochastic gradient noise to provide high probability bound for single-level non-convex problems. The main challenge for bilevel optimization is the following: even if we have all of these assumptions in bilevel problems, the hypergradient estimator is still biased  (check Lemma 4 in our paper for details) and therefore the techniques in the single-level stochastic optimization may not apply.