An Accelerated Algorithm for Stochastic Bilevel Optimization under Unbounded Smoothness

Summary

This work develops the first algorithm that achieves the near-optimal oracle complexity $\widetilde{O}(\epsilon^{-3})$ (the number of evaluations to gradient or Hessian/Jacobian vector product) to achieve an $\epsilon$-stationary point of a stochastic bilevel optimization problem with unbounded smooth nonconvex upper-level function and strongly-convex lower-level function. The algorithm design is novel which updates the upper-level variable by normalized stochastic gradient descent with recursive momentum and updates the lower-level variable by the stochastic Nesterov accelerated gradient (SNAG) method.

Strengths

This work studies a specific stochastic bilevel optimization problem with significant applications in sequential data learning. The upper-level function has unbounded smoothness which is more general than the widely used Lipscthiz smoothness assumption. The oracle complexity is improved on this specific bilevel optimization problem compared with existing works. The algorithm design not only combines but also adjusts STORM variance reduction and Nesterov acceleration. The 2 real-world experiments look convincing. The presentation is clean and clear.

Weaknesses

The major weakness is Assumption 4.2 as mentioned in the question (1) below. Some points could be clarified as mentioned in other questions below.

Questions

(1) **My major concern** is that Assumption 4.2 is made on algorithm-generated variables not on problem setting, which cannot be verified. Is it possible to bound $\Delta_t$ in terms of $\|x_{t+1}-x_t\|$ to get rid of Assumption 4.2? If not, at least Assumption 4.2 should also be mentioned in the main result (Theorem 4.1). **I would like to raise my rating if all the assumptions become problem-related.** (2) In line 48 "each realization of the stochastic oracle calls is Lipschitz with respect to its argument", do you mean $F(x,y,\xi)$ and $G(x,y,\xi)$ are Lipschitz continuous functions of $(x,y)$ for each $\xi$? You might express in formula to make it more clear. (3) You could also cite [1] as a paper on relaxed smoothness: [1] Convex and Non-convex Optimization Under Generalized Smoothness (4) In line 127, do you mean "... is a good approximation of $\nabla \Phi(x)$ **if $y\approx y^*(x)$**"? (5) Is there an intuition why line 21 (averaging of y) is needed in Algorithm 2? Is it possible to implement $N=\mathcal{O}(1)$ steps of SNAG such that $\mathbb{E}\|y_{t+1}-y^*(x_t)\|\le \frac{1}{2}\mathbb{E}\|y_t-y^*(x_t)\|$ and remove averaging of y for every outer iteration t, and obtain convergence result purely in expectation instead of high probability? This yields the same oracle complexity per round since $I,N=\widetilde{\mathcal{O}}(\epsilon^{-1})$ are of the same order, and may get rid of Assumption 4.2 by bounding $\Delta_t$ in terms of $\|x_{t+1}-x_t\|$. (6) You said your upper-level algorithm design is inspired from [17,49]. However, in line 23 (STORM variance reduction) of Algorithm 2, your second and third additive terms use the same samples $\overline{\xi}_t$ while [49] uses different samples. The algorithm design differs from [17] more. Such difference still exists even if Algorithm 2 is reduced to single-level. Why do you use such different design? (7) Also in line 23, does $\overline{\nabla}f(x_t,\hat{y}_t;\overline{\xi}_t)$ denote the average of the expression right below Eq. (2) over samples $s=1,\ldots,S$? You could define more clearly. (8) At the beginning of Section 5 experiments, $ph(w; x)I_{[y=-1]}$ appears in your AIC maximization. Do you mean $rh(w; x)I_{[c=-1]}$? If not, what do $p$ and $y$ mean?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The conclusion section mentions the limitation that the convergence analysis for the Option I of Algorithm 2 relies on the lower-level problem being a one-dimensional quadratic function. Actually I think this is fine since Option II works for more general strongly convex lower-level function. The major limitation from my perspective in mentioned in my above question (1) about the undesired Assumption 4.2. The authors said in the checklist that they do not see any negative societal impacts in their theoretical work, which I agree.

Summary

The authors consider bilevel optimization problems under the unbounded smoothness assumption of the outer function. They propose AccBO, an AID-based method that uses Stochastic Nesterov Accelerated Gradient to approximate the lower-level solution and Neumann approximations to estimate the inverse Hessian-Vector product involved in the hypergradient expression. Then, a normalized STORM-like update is performed for the outer variable. They show a $\tilde{\mathcal{O}}(\epsilon^{-3})$ complexity of their algorithm. Finally, numerical experiments are provided on the deep AUC maximization and data hypercleaning problems.

Strengths

- **S1**: The paper is well-written - **S2**: The algorithms proposed by the authors achieve SOTA complexity for stochastic bilevel optimization under unbounded smoothness assumption.

Weaknesses

### Major - **W1**: Some errors in the code may invalidate the experimental results. See the comment **C2** in the Code section for details. - **W2**: AccBO comes with many hyperparameters, which can be a concern for practical applications. - **W3**: l.219: "is nearly optimal in terms of the dependency of $\epsilon$". The authors should be careful when talking about optimality of an algorithm. In particular, the algorithm AccBO does not belong to the algorithm class considered in [1] which does not take into account the biasedness of the hypergradient estimate due to the approximation of the inverse Hessian-vector product and the solution of the inner optimization problem. To my knowledge, the only works considering lower bounds for bilevel problems are [2, 5]. - **W4**: The works [2, 3, 4] also use Nesterov acceleration for the inner optimization problem, but in the deterministic setting. They should be mentioned in the paper. ### Code Even though it is a theoretical paper, I have several comments regarding the experiments code: - **C1**: In the checklist, the authors indicate "The code and data are attached as a supplement with instructions for reproducibility." However, the code documentation is insufficient to run the experiments. The data are not present in the folders and there is no instruction indicating where to download them. As a consequence, when I launch the commmand `python main.py --methods accbo ` as indicated in the `README.md`, I get an error. - **C2**: It seems that the implementation of the different methods do not correspond to the one described in their respective original paper, which might invalidate the empirical results: - For `TTSA`, `SUSTAIN` and `AccBO`, the estimation of the inverse Hessian-vector product is computed by the formula $$ v = \eta\sum_{q=-1}^Q\prod_{j=Q-q}^Q(I-\eta \nabla^2_{yy}G(x,y;\zeta^{(j)}))\nabla_y F(x,y;\xi) $$ while in the original paper (and even the present paper for AccBO), the authors use the formula $$ v = \eta\prod_{j=1}^{q(Q)}(I-\eta \nabla^2_{yy}G(x,y;\zeta^{(j)}))\nabla_y F(x,y;\xi) $$ with $q(Q)\sim\mathcal{U}{\{1,\dots,Q-1\}}$. - For `SABA`, the SAGA-like variance reduction is absent in the code while being the core of the algorithm. Moreover, `SABA` uses constant step sizes. - For `VRBO`, the oracles in the case `step % self.args.update_interval == 0` should be computed with a larger batch-size (or even with full batches) than in the inner for-loop. - **C3**: The code of `VRBO` for the datacleaning task is missing. The method to compute the hypergradient in `ma_soba.py` is missing as well for the datacleaning task (but not AUC experiment, thus I assume it is almost the same). ### Minor - l.41: "transforms" -> "transformers" - In `README.md` (both): `python main.py --medthods [algorithm] ` -> `python main.py --methods [algorithm] ` I am inclined to raise my score if my concerns, particularly regarding the experiments, are addressed.

Questions

N/A

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This paper investigates a class of stochastic bilevel optimization problems where the upper-level function is nonconvex with potentially unbounded smoothness, and the lower-level problem is strongly convex. To improve the convergence rate, the authors propose a new Accelerated Bilevel Optimization algorithm named AccBO. This algorithm updates the upper-level variable using normalized stochastic gradient descent and the lower-level variable using the stochastic Nesterov accelerated gradient descent algorithm with averaging. The proof shows that AccBO has an oracle complexity of $ \tilde{O}(1/\epsilon^3) $, relying on a novel lemma describing the dynamics of the stochastic Nesterov accelerated gradient descent algorithm under distribution drift. Experimental results demonstrate that AccBO significantly outperforms baseline algorithms in various tasks, achieving the predicted theoretical acceleration.

Strengths

1.The paper introduces AccBO, a novel algorithm that leverages normalized recursive momentum for the upper-level variable and Nesterov momentum for the lower-level variable. This dual approach to acceleration in bilevel optimization under stochastic settings is a significant advancement. 2.AccBO achieves an oracle complexity of $ \tilde{O}(1/\epsilon^3) $ for finding an $\epsilon$-stationary point. This is a substantial improvement over the previously best-known complexity of $ \tilde{O}(1/\epsilon^4) $ for similar problems.The paper provides new proof techniques, particularly in analyzing the dynamics of stochastic Nesterov accelerated gradient descent under distribution drift. 3.The effectiveness of AccBO is empirically verified across various tasks, including deep AUC maximization and data hypercleaning. The results show that AccBO achieves the predicted theoretical acceleration and significantly outperforms baseline algorithms. 4.The paper extends the bilevel optimization framework to handle upper-level functions with unbounded smoothness, which is a realistic scenario in many neural network applications. This makes the proposed algorithm applicable to a broader range of practical problems.

Weaknesses

The convergence analysis for Option I of the algorithm is restricted to the lower-level problem being a one-dimensional quadratic function. This limitation reduces the generality and applicability of the theoretical results for more complex lower-level problems.

Questions

In reference [33], Assumption 2 states $|| \nabla_y f(x, y^*(x)) || \leq M$, while in this paper, Assumption 3.2 states$ || \nabla_y f(x, y) || \leq l_{f,0} $ . Can this assumption be relaxed to a less stringent form?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The convergence analysis for Option I of the AccBO algorithm assumes that the lower-level problem is a one-dimensional quadratic function. This limits the algorithm's applicability to scenarios where the lower-level function deviates significantly from this simple form, especially in higher dimensions or more complex optimization landscapes.

Summary

This paper considers a class of stochastic bilevel optimization problems where the upper-level function is nonconvex with potentially unbounded smoothness, and the lower-level problem is strongly convex. Their novel algorithms achieve an oracle complexity of $O(1/\epsilon^3)$ to find an $\epsilon$-stationary point.

Strengths

1. The setting is novel: they consider problems with potentially unbounded smoothness. 2. Their complexity strictly improves the state-of-the-art oracle complexity for unbounded smooth nonconvex upper-level problems and strongly convex lower-level problems.

Weaknesses

1. In the algorithm, for the update of the normalized stochastic gradient, it seems that it can be reformulated to tune the learning rate by dividing by the norm of the stochastic gradient. Therefore, the novelty and reasonableness of this step may be unclear. 2. In the experiments, it is common to use the same learning rate for baselines and the proposed algorithms for a fair comparison. However, in the experimental details, the authors provide the best learning rate pairs for different baselines, which seems unusual.

Questions

1. Regarding the main challenge, the authors state why previous algorithms and analyses do not work. However, this explanation is not clear to me. Can you elaborate on how your work addresses the problem of hypergradient bias? Also, how does your potential function argument differ from previous work?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

Reviewer 2WtS agrees with the authors and will keep rating 5.

Reviewer 2WtS agrees with the authors and will keep rating 5.

Thank you!

Dear Reviewer 2WtS, Thank you for raising the rating. Best regards, Authors