Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization

Summary

This paper studies scalable bi-level optimization problems and points out the limitations of most traditional methods. Among all, to mitigate the high memory cost issue of the GU method, it proposes (FG)2U such that the space consumption is reduced to $\mathcal{O}(M)$ from $\mathcal{O}(MN)$. Convergence analysis is provided and extensive experiments show the benefit of the algorithm.

Strengths

1. The presentation is great and I feel relaxed to read this paper. 2. The experiment is extensive including both small and large scale settings.

Weaknesses

1. The memory cost is reduced but the computational cost seems to increase significantly. The computational cost takes the order of $\mathcal{O}(KTN)$ since $b=\mathcal{O}(N)$. This feels unrealistic as well in the large-scale application when $N$ is large. 2. Based on the previous question, I checked the choice of $b$ in the experiment, which is small. How do authors tune this parameter? Is there any ablation study on it? 3. In terms of the discussion in Appendix B.2, IF methods also require $\mathcal{O}(M)$ space using some approximation tricks. There indeed exist approximation errors in IF methods. However, normally, the approximation errors can be controlled to be very small based on the hyperparameter selections. Thus, IF methods seem suitable for scalable bi-level optimization problems to some extent. Do authors have time to include IF methods in Table 2? 4. The algorithm is deterministic without consideration of the batch data. Can authors provide some insights in stochastic settings? For example, can the unbiasedness of the hypergradient estimator still be satisfied, or is there any requirement for the batch size? 5. Can authors provide real space consumption in the comparison (Table 1, Table 2)?

Questions

See weakness part.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

See weakness part.

Summary

The paper introduces a method called Forward Gradient Unrolling with Forward Gradient, abbreviated as (FG)²U, which is designed to address the large memory requirements of forward method in bi-level optimization in large-scale machine learning model

Strengths

1. The method significantly reduces memory overhead compared to traditional gradient unrolling methods, making it suitable for large-scale applications. 2. Can be easily implemented within popular deep learning frameworks and adapted to various optimization problems.

Weaknesses

1. The proposed method introduces additional computational complexity. Can the author give an analysis of the complexity? 2. The convergences analysis is for the problem (2) not the original problem (1). 3. More large scale datasets are needed.

Questions

see weakness

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

NA

Summary

This paper presents a novel gradient unrolling algorithm for bi-level optimization. The authors highlight that existing methods for calculating meta gradients in the literature are not memory efficient. They propose a sampling-based method, (FG)^2U, to estimate the meta gradient. This approach approximates the meta gradient by multiplying it with a random rank-1 matrix, thereby simplifying the Forward Gradient Unrolling (FGU) scheme. The paper includes discussions on sampling efficiency, convergence analysis, and numerical experiments for (FG)^2U.

Strengths

- The proposed method (FG)^2U is simple and easy to implement. - The writing is clear and easy to follow. - The theoretical results clearly demonstrate the relationship between convergence, sampling batch size, and parameter space size. - The three experiments, data condensation, LM fine-tuning, and PDE recovery, show the wide range of applications of bi-level optimization and the proposed method (FG)^2U.

Weaknesses

A major concern is the theoretical dependence on the parameter dimension $N$. Theorem 3.4 indicates that the convergence rate depends linearly on N. It means that - With a fixed batch size, convergence on large-scale applications would be slow (vanilla (FG)^2U without additional techniques). - With an $O(N)$ batch size, convergence is satisfactory, but the memory cost grows as $O(MN)$, similar to FGU. - While the authors mention that (FG)^2U allows parallelization to mitigate computational overhead, it seems that the calculations of FGU (4) also permit parallelization, correct? Specifically, $d \theta / d \phi$ can be computed in parallel as $d \theta[1] / d \phi, d \theta[2] / d \phi, d \theta[3] / d \phi, \cdots$

Questions

Regarding the formula for data condensation (17) and related formulas in the appendix: should it be minimizing over ${\mathcal{D}_c}$ instead of minimizing over ${\mathcal{D}_o}$?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The paper appears to have no potential negative societal impact. The authors discussed the limitations in the Appendix H.

Response to author

Thank you for the rebuttal, the authors have addressed all my concerns. I will increase my score.

Thanks for the response

I thank the authors for their response. Some of my concerns have been addressed. However, I am still slightly concerned about the theoretical guarantee on the problem dimension. There are also some works on bilevel optimization using zeroth-order types of methods. Maybe the authors would like to include them and provide a comparison. I keep my current score. Best, Reviewer

Response to authors

I greatly appreciate the efforts made by the authors during the rebuttal phase, including their responses to my questions and the inclusion of additional experimental results. However, my primary theoretical concern, which I raised in the initial review, remains only partially addressed. The statement "Plan A requires exactly $M$ threads" may not be entirely accurate. It is possible to adaptively allocate the $M$ calculation tasks based on the constraints of the hardware. For instance, if $M=5$, one could opt for 3 threads and distribute the tasks as $[1,1,3]$, rather than being limited to $[1,1,1,1,1]$. This flexibility also casts doubt on the assertion that "the memory cost of $\text{(FG)}^2\text{U}$ can be more manageable than FGU and RGU," as one could potentially implement simple variants of FGU or RGU. Of course, as the authors have noted, these analyses are merely upper bounds. The actual performance in the experiments presented by the authors appears promising. Therefore, I will keep my positive score.