Provably Faster Algorithms for Bilevel Optimization via Without-Replacement Sampling

Thanks for spending your valuable time reviewing our manuscript and provide insightful feedbacks. Our responses are provided as below:

Convergence Rate of SOBA. We first want to clarify that SOBA [1] has convergence rate of $O(\epsilon^{-4})$ under the notation of our manuscript. In fact, the $\epsilon$-stationary point in our manuscript is defined as $\|\nabla h(x)\| \leq \epsilon$. Note that the SOBA paper uses an alternative definition of $\|\nabla h(x)\|^2 \leq \epsilon$.

The value of $p$. Our WiOR-BO has convergence rate of $O(max(m,n)^p\epsilon^{-3})$, note that this rate is strictly better than the $O(\epsilon^{-4})$ rate of stocBiO, BSA and SOBA. As for the value of $p$, it relates to the specific example order we choose. More specifically, in Assumption 4.5 (i.e. how well the consecutive example gradients approximate the true gradients), we have the constant $C = \max(m,n)^pA$. For shuffle-once, we have $p=1$, for random-reshuffling, we can show $p=0.5$, while for the GraB algorithm~[2], we can reach $p=0$. In our experiments, we consider the random-reshuffling and shuffle-once, as for the GraB algorithm, although it has favorable property that $p=0$, we need to store at least one gradient (also Hessian in the bilevel case), and this leads to $O(d^2)$ extra space, which is expensive for large scale neural networks.

Comparison with other Methods. In Table 1 of the manuscript, we compare WiOR-BO with methods such as stocBiO, BSA, and SOBA, which all adopt SGD-like updates. Our algorithm differs from these methods in the sampling strategy. The faster convergence rate of our algorithms demonstrates the efficacy of using without-replacement sampling. Then in Table 2 of the Appendix, we include more algorithms, some of them obtain faster convergence rate than WiOR-BO, yet this is achieved by using some form of variance reduction technique. For example, SABA ($O(max(m,n)^{2/3}\epsilon^{-2})$) is based on SAGA, SRBA ($O(max(m,n)^{1/2}\epsilon^{-2})$) is based on SARAH, while MRBO ($O(\epsilon^{-3})$) is based on STORM. However our algorithm still demonstrates favorable empirical advantages over these variance-reduction enabled methods. In fact, WIOR-BO has fewer hyper-parameters (learning rates for inner and outer problems) to tune compared to MRBO, which has six independent hyper-parameters, and the optimal theoretical convergence rate is achieved only if the complicated conditions among hyper-parameters are satisfied in Theorem 1 of the MRBO paper are satisfied, and this requires significant tuning effort in practice. Next, SRBO/SABA evaluates the full hyper-gradient at the start of each outer loop, where we need to evaluate the first and second order derivatives over all samples in one step, which is very expensive for modern ML models (such as the transformer model with billions of parameters). In contrast, our WiOR-BO never evaluates the full gradient. Note that the inner loop length $I = lcm(m,n)$ is analogous to the concept of "epoch" in single level optimization: where we go over the data samples following a given order, but at each step we evaluate gradient over a mini-batch of samples. The practical advantage of our WiOR-BO is further verified by the superior performance over MRBO and VRBO in the Hyper-Data Cleaning Task.

Novelty: Without-replacement sampling is widely used in practice in model training including bilevel optimization. However, most existing literature for bilevel optimization assumes the independent sampling assumption to simplify the analysis. To the best of our knowledge, our WiOR-BO (WiOR-CBO) is the first without-replacement sampling bilevel algorithm with theoretical guarantee.

Responses to Reviewer Questions:

Q1: The paper regards the special cases as one of the main contributions. Can the special cases be fully covered by the general cases? Are convergence results or algorithms for these cases different from what we could achieve by applying the general case?

R1: Due to space limitation, the convergence results of the minimax and compositional cases are deferred to Corollary B.20 and Corollary B.21 in Appendix. For the compositional case, we can directly apply the analysis of bilevel optimization and get the same convergence rate. For the minimax case, we need to remove the bounded gradient assumption, i.e. $\|\nabla_y f(x,y, \xi) \| \leq C_f$ as this assumption conflicts with the assumption of $\mu$-strong convexity with respect to y. Once this assumption is removed, we can follow the same analysis used in bilevel optimization to establish the convergence rate for the minimax case.

Q2: Can you compare the emprical performance with SOBA?

R2: In the Hyper-Data Cleaning task, we have compared with SOBA in the name of WiR-BO. We use the name WiR-BO as there are several single loop bilevel algorithms such as AmIGO [3] and FSLA [4] in the literature. We also compare with variance-reduction enable methods: MRBO and VRBO.

Finally, we want to thank the reviewer once again. If anything remains unclear, please do not hesitate to let us know.

References

[1]. Dagréou, Mathieu, et al. "A framework for bilevel optimization that enables stochastic and global variance reduction algorithms." Advances in Neural Information Processing Systems 35 (2022): 26698-26710.

[2]. Y. Lu, W. Guo, and C. M. De Sa. Grab: Finding provably better data permutations than random reshuffling. Advances in Neural Information Processing Systems, 35:8969–8981, 2022.

[3]. M. Arbel and J. Mairal. Amortized implicit differentiation for stochastic bilevel optimization,2021.

[4]. J. Li, B. Gu, and H. Huang. A fully single loop algorithm for bilevel optimization without hessian inverse. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36,pages 7426–7434, 2022

Sorry for the confusion! Since in your original review, it reads "SOBA with moving average and with-replacement sampling", we interpret this as meaning the standard stochastic case (which involves sampling).

Firstly, we agree with the reviewer that the value of $p$ is important and we indeed include it in Table 1 of the manuscript and also discuss them in the main text. As stated in the rebuttal, the dependence over $max(m,n)^p$ comes from the sample gradient error $||\frac{1}{k} \sum_{\tau = t}^{t+k-1} \nabla f (x_t,y_t; \xi_{\tau}^{\pi}) - \nabla f (x_t,y_t)||^2$, while the value of $p$ depends on specific sampling orders. In particular, random-reshuffling and the GraB method can force $p < 1$, which has a better dependence over $max(m,n)$ compared to deterministic algorithms. Additionally, we want to add that similar dependence over $max(m,n)^p$ exists for single level algorithms that adopt without-replacement sampling.

Next, we believe our work makes a valuable contribution to the study of bilevel optimization. Specifically, we demonstrate that the independent sampling assumption, commonly used in SGD-type bilevel algorithms, can be replaced by without-replacement sampling, leading to a reduced per-iteration cost (less number of backward passes) and an improved convergence rate (with the improvement depending on specific example orders). Furthermore, the new algorithms (WiOR-BO/WiOR-CBO) still perform SGD-type updates, making them practical for modern machine learning, in contrast to other variance-reduction-based or deterministic algorithms.

Thanks for spending your valuable time in reviewing our manuscript and provide insightful comments.

Firstly, about the upper bound in Assumption 4.3, 4.4 and 4.7, we assume a common bound due to the following reasons. Note that these assumptions measure similar properties, namely the sample estimation error of first or second order derivatives, using a unified bound can simplify the analysis and makes the final convergence bound obtain simpler constant factors. In fact, we can assume different upper bounds, but the convergence rates are not affected.

Next, We want to thank you for your advice and will use $\tilde{O}(\cdot)$ to indicate that logarithmic terms are hidden. Furthermore, we will adjust the equations so that they can fit into the line width.

Thanks for spending your valuable time reviewing our manuscript and provide insightful feedbacks. Our responses are provided as below:

W1: The paper claimed that "Compared to independent sampling-based algorithm such as stocBiO, ... computational cost of training large scale models." Is there any empirical result to show the memory costs?

R1: In our algorithms, we maintain the permutation of example indices. For example, in the Hyper-Data Cleaning Task: we have 40000 training examples and 5000 validation examples and we store the permutations with numpy.uint16 data type and this leads to memory cost of 90KB. Meanwhile, the GPU memory related to training neural network is around 20.07MB, which is roughly 200 times greater than the memory used to store the index permutations. Note that for large scale neural networks, this ratio will be larger. Additionally, the permutation process is executed on the CPU and does not consume expensive GPU memory.

W2:Projection is applied to $u$, but I found other bilevel papers (such as StocBio, BSA) don't use that. Is there any intuition or explanation about that?

R2: Note that our algorithms perform single loop update, i.e. update $x_t$, $y_t$ and $u_t$ alternatively; in contrast, stocBiO and BSA are double loop algorithms, where for each $x_t$, they use a loop to estimate $y_{x_t}$ and the corresponding hyper-gradient, and there is no variable $u_t$ involved. Recall that $u_t$ is used to estimate the solution of the quadratic problem in hyper-gradient computation, i.e. $u_x = \nabla_{y^2} g(x, y_x)^{-1} \nabla_y f(x, y_x)$ and we only apply projection operation to $u_t$, therefore, stocBiO and BSA does not need the projection. In fact, this projection operation guarantee $u_t$ is always bounded during the update and is essential in the analysis. The projection operation can be found in recent single loop bilevel algorithms such as SRBA in [1].

W3: In Algorithm 2, the generation of is not very clear. It seems that there are samples for inner loops given a outer sample. How do you execute the permutation for the inner dataset.

R3: As we state in line 149-150 of main text, for the inner dataset, we concatenate $S$ random permutations of the inner dataset $D_{l,\xi_t^{\pi}}$ to get {$\zeta_{\xi_{t}^{\pi}, t_l}^{\pi} \in D_{l,\xi_t^{\pi}}, t_l \in [T_l]$} with a sequence length $T_l = S\times n_{\xi_i}$. To put it simply, for each given outer example $\xi_{t}^{\pi}$, we go through its corresponding inner dataset $S$ rounds, and for each round we go through the inner dataset following a random permutation. Finally, we correct a typo in the Line 7 of Algorithm 2: it should be {$\zeta_{\xi_i, j}^{\pi}, j\in[n_{i}]$} instead of {$\zeta_{\xi_i, j}^{\pi}, i\in[n_{i}]$}.

W4:In experiments, what is the best hyper-parameters for each baselines, such inner and outer learning rates?

R4: For the invariant risk minimization task: the inner and outer learning rates are 0.001; for the hyper-data cleaning task: reverse, stocBiO, BSA, AID-CG, WiR-BO, WiOR-BO-SO, WiOR-BO-RR use inner learning rate 0.1 and outer learning rate 1000, WiR-BO, WiOR-BO-SO, WiOR-BO-RR set $\rho = 0.1$; MRBO uses $d=10$, $m=500$, $c_1=0.9$, $c_2=0.9$, $\gamma=500$, $\lambda=0.2$; VRBO uses $q=5$, $\alpha=1000$, $\beta=0.2$; for the hyper-representation learning task: for the Omniglot dataset, the outer learning rate is 0.1, the inner learning rate is 0.4, $\rho = 0.002$; for the MiniImageNet dataset, the outer learning rate is 0.05, the inner learning rate is 0.01, $\rho = 0.002$, for the RL-MLMC baseline, we set $K_{max} = 6$.

W5: What are the details of data sampling for other baselines? For other baselines, they usually need the independent training and validation set for upper and lower-level update, how did you construct data for them?

R5: In our empirical implementation, we use a dataloader class to handle example sampling. For our algorithms, the dataloader returns examples following the order of a permutation, while for other baselines that require independent sampling, the dataloader will randomly sample a mini-batch of samples (sample indices) from the whole dataset.

W6: Typos: Line 107 $w_x$, should be $u_x$.

R6: We will correct this typo and carefully proofreading our manuscript.

References

[1]. Dagréou, Mathieu, et al. "A lower bound and a near-optimal algorithm for bilevel empirical risk minimization." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

Thanks for spending your valuable time reviewing our manuscript and provide insightful feedbacks. Our responses are provided as below:

Technical contributions. We provide the first theoretical analysis of bilevel optimization that uses without-replacement sampling. A key insight in our analysis is to analyze the descent of potential function over an interval of length k instead of one step. This is important as only the amortized gradient estimation error ($|\frac{1}{k}\sum_{\tau=t}^{t+k-1}\nabla f(x,y,\xi_\tau) - \nabla f(x,y)|$) of without-sampling exhibits improvement compared to independent sampling. Meanwhile, we need to carefully balance the intertwined hyper-gradient estimation, inner variable estimation error and the estimation error to $u_x$ in the analysis.

Fair-Classification Task: we provide further experimental results over a fair-classification task [1,2], which can be formulated as a minimax optimization problem. More specifically, for a classification task with $K$ categories, the objective minimizes the maximum loss over all categories, this improves the fairness among different categories. The objective is as follows:

$\min_{w}\max_{u \in U} \sum_{i=1}^K u_i L_i(w) - \lambda ||u - \frac{1}{K}||^2$

where $u_i$ is the weight for category $i$ and $U$ denotes the probability simplex, $L_i(w)$ denotes the classification loss over category $i$ and $w$ is the model parameter. Follow the setting in [2], we consider T-shirt/top, Coat and Shirt categories in the Fashion-MNIST dataset (K = 3) and solve the objective with SGDA [3], WiOR-MiniMax-SO and WiOR-MiniMax-RR. Note that SGDA performs the minimization and maximization alternatively and differs our WiOR-MiniMax-SO(-RR) only in the sampling strategy. Be We show results in the table below, we show the training loss across different epochs:

As shown in the table, our algorithms outperforms SGDA, which demonstrate the efficacy of using without-replacement sampling.

Responses to Reviewer Questions

Q1: How was the bias due to the without-replacement sampling handled in the analysis? Using window averaging? In this case, did you need to store the gradients of previous samples?

A1: We do not need to store the gradients of previous samples, and the window averaging is solely for analysis purpose. For example to bound the estimation error of $\nabla_x f(x_t, y_t; \xi_{t})$ to $\nabla_x f(x_t, y_t)$, we bound the following term $||\sum_{\tau=t}^{t+k} (\nabla_x f(x_\tau, y_\tau; \xi_{\tau}) - \nabla_x f(x_t, y_t))||^2$ and this term can be further broken into two sub-terms: $||\sum_{\tau=t}^{t+k} (\nabla_x f(x_t, y_t; \xi_{\tau}) - \nabla_x f(x_t, y_t))||^2$ and $||\sum_{\tau=t}^{t+k} (\nabla_x f(x_\tau, y_\tau; \xi_{\tau}) - \nabla_x f(x_t, y_t; \xi_{\tau}))||^2$. For the first term, the error decrease as $k$ increase (as we will go through more samples due to without replacement sampling), while for the second term, it is equivalent to bound $||(x_\tau, y_\tau) - (x_t, y_t)||^2$ by the smoothness assumption and this variable drift term can be bounded by tuning the value of learning rates.

Q2: When you compared the convergence rate of algorithm 1 with baselines, e.g., stocBiO, how did you justify that the faster convergence is because of the sampling but not the single-loop update? Given that stocBiO is not a single-loop algorithm.

A2: The rate improvement is indeed due to the without-replacement sampling strategy. SOBA[4] also adopts a single loop update rule, but this algorithm also obtains the same $O(\epsilon^{-4})$ rate as stocBiO.

Finally, thanks for your review again. We will also carefully polish our presentation, including the font sizes in the figures suggested in your review.

References:

[1] Nouiehed, Maher, et al. "Solving a class of non-convex min-max games using iterative first order methods." Advances in Neural Information Processing Systems 32 (2019).

[2]. Huang, Feihu, Xidong Wu, and Zhengmian Hu. "Adagda: Faster adaptive gradient descent ascent methods for minimax optimization." International Conference on Artificial Intelligence and Statistics. PMLR, 2023.

[3]. Lin, T., Jin, C., and Jordan, M. (2020a). On gradient descent ascent for nonconvex-concave minimax problems. In International Conference on Machine Learning, pages 6083–6093. PMLR.

[4]. M. Dagréou, P. Ablin, S. Vaiter, and T. Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. arXiv preprint arXiv:2201.13409,2022.

Thanks for spending your valuable time reviewing our manuscript and provide insightful feedbacks. We provide our responses to your comments below:

Hyper-parameter Tuning. Our algorithms include three hyper-parameters: the outer learning rate $\eta$, the inner learning rate $\gamma$ and the learning rate of variable $u$: $\rho$, and they are easy to tune and robust to specific choices of hyper-parameters. More specifically, we use grid search (as stated in the manuscript) to perform hyper-parameter tuning, which is a standard practice used by optimization literature [1,2], furthermore, we use the loss over a hold-out validation set as the metric for tuning and adopt the early stoping technique. Take the hyper-data cleaning task as an example, we search $\eta$ in {1, 10, 100, 1000, 10000}, $\gamma$ in {0.001, 0.01, 0.1, 0.5, 1} and $\rho$ in {0.001, 0.01, 0.02, 0.05, 0.1}. Similar to the behavior of the learning rate in SGD for single-level problems, the performance initially improves as we increase the learning rate, and it will start to diverge beyond a certain point. In the following experiments, we perform ablation studies to the learning rates for our WiOR-BO algorithm over the hyper-data cleaning task, where show the validation loss vs Number of Hyper-iterations:

Table 1: Ablation Study over $\gamma$ ($\eta = 1000$ and $\rho = 0.1$)

Table 2: Ablation Study over $\eta$ ($\gamma = 0.1$ and $\rho = 0.1$)

Table 3: Ablation Study over $\rho$ ($\gamma = 0.1$ and $\eta = 1000$)

As shown by the results in Tables 1 to 3, the best performance is achieved with $\eta=1000$, $\gamma=0.1$, and $\rho=0.1$. The performance with other hyperparameter settings is also comparable.

Computational Efficiency. We can discuss the influence of dataset size and model size based on the convergence rate. As stated in Table-1 of the manuscript, WiOR-BO has convergence rate of $max(m,n)^q\epsilon^{-3}$ and WiOR-CBO has rate of $max(m,n)^{2q}\epsilon^{-4}$. Note that the convergence rate is measured by the number of gradient queries and our convergence results show that this number is not affected by the model size, however, the cost of one gradient query indeed increases as the model size increase; as for the dataset size, the convergence rate of our model has a dependence of $max(m,n)^q$, where $q$ is a constant that relates to the without-replacement sampling strategy we use. For shuffle-once, we have $q=1$; for random-reshuffling, we have $q=0.5$, and if use the GraB algorithm in [3] (which is based on a herding and balancing), $q=0$.

Dynamic Environments. In our manuscript, we focus on the finite-sum setting (i.e. over a static dataset), which is a standard setting in modern machine learning [4], and we also believe it is a suitable setting to study the effect of example correlation effect (due to without-replacement sampling) to the convergence of bilevel optimization. While the dynamic environment setting is an interesting extension, it is beyond the scope of this manuscript and warrants a separate, dedicated study.

Finally, in terms of the limitation, since the convergence results are based on various regularity assumptions to the outer and inner problems, real world tasks might not hold these assumptions. As for the future direction, we can study the effect of without-replacement sampling under other assumptions, in this manuscript, we consider the nonconvex-strongly-convex class of functions, and we could extend our analysis to other classes of functions, furthermore, the reviewer's advice to consider dynamic environments is also possible.

References:

[1]. Yang, Junjie, Kaiyi Ji, and Yingbin Liang. "Provably faster algorithms for bilevel optimization." Advances in Neural Information Processing Systems 34 (2021): 13670-13682.

[2]. Ji, Kaiyi, Junjie Yang, and Yingbin Liang. "Bilevel optimization: Convergence analysis and enhanced design." International conference on machine learning. PMLR, 2021.

[3]. Lu, Yucheng, Wentao Guo, and Christopher M. De Sa. "Grab: Finding provably better data permutations than random reshuffling." Advances in Neural Information Processing Systems 35 (2022): 8969-8981.

[4]. Johnson, Rie, and Tong Zhang. "Accelerating stochastic gradient descent using predictive variance reduction." Advances in neural information processing systems 26 (2013).