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

We thank the reviewer for the feedback. We will try to address the concerns point by point.

Regarding the computational complexity

We direct the reviewer to our general response. We place a computational cost analysis in A.1 and a more detailed discussion on practical computation in A.2.

Regarding the convergence of problem (1)

To extend the convergence of optimization problem (2) into (1), we need to assume that the lower-level objective function $g$ is strongly convex w.r.t. $\theta$ as commonly done by previous works [2][3]. From the strong convexity and first-order smoothness (Assumption 3.2) of $g$, we have 1) the zeroth and first-order smoothness of $\theta^*(\phi)$; 2) $\|\theta_T(\phi')-\theta^*(\phi')\|\to 0$ as $T\to +\infty$. Then the inequality $\|\theta_T(\phi)-\theta_T(\phi')\|\leq \|\theta_T(\phi)-\theta^*(\phi)\|+\|\theta_T(\phi')-\theta^*(\phi')\|+\|\theta^*(\phi)-\theta^*(\phi')\|$ implies the the zeroth and first-order smoothness of $\theta_T(\phi)$. Following the same line of proof as presented in our paper, we derive the smoothness of $f(\theta^*(\phi),\phi)$ and $f(\theta_T(\phi),\phi)$, and subsequently the convergence of either problem (1) or (2).

However, as discussed in lines 73 to 77, given that the scope of this paper is large-scale BO, the inner optimization typically involves deep neural networks. Therefore, the optimal parameters are not explicitly accessible and can only be estimated through iterative procedures. Most related works [1][2][3] are implicitly or explicitly solving (2) instead of (1).

Additionally, it is important to acknowledge that achieving strong convexity is often unfeasible in practical applications. Consequently, in this paper, we aim to present a more practical convergence theory that proves the effectiveness of our method.

Regarding larger-scale datasets

As discussed in Appendix H, this work serves as the first attempt to apply FG/ZO in BO. While the scale of cases studied in this paper is relatively small compared to the most powerful generative models in the industry, it is comparable to recent large-scale BO works such as [1] and significantly larger than traditional BO works. We believe the empirical studies conducted in this paper are sufficient to demonstrate the potential of $\text{(FG)}^2$U in large-scale BO. We anticipate that the effectiveness of $\text{(FG)}^2$U will be further validated in larger scales.

[1] Making Scalable Meta Learning Practical, https://arxiv.org/abs/2310.05674

[2] Truncated Back-propagation for Bilevel Optimization, https://arxiv.org/abs/1810.10667

[3] Bilevel Optimization: Convergence Analysis and Enhanced Design, https://arxiv.org/abs/2010.07962

We thank the reviewer for the overall positive evaluation and will try to address the concerns raised point by point.

Regarding the complexity

We direct the reviewer to the general response for computational cost analysis and discussions regarding the computation in practice, including the empirical success of FG/ZO in large-scale applications with constant batch size, tricks to reduce the variance, how the parallelizable nature of $\text{(FG)}^2$U fits modern AI computation, and the more cost-effective two-phase methodology in practice.

Regarding the strategy to choose $b$

We follow the methodology [2][3] developed for FG/ZO in large-scale applications to choose $b=\mathcal{O}(1)$ rather than $\mathcal{O}(N)$. We select the largest possible $b$ that does not exceed the GPU memory limit to fully utilize the GPU. See Tables C and D in the attached PDF in the general response for the exact memory consumption. Notice that we additionally perform gradient accumulation through iterative/multi-GPU parallel computation, which will also influence the variance. The number of accumulation steps is tuned in our initial attempts according to the wall-clock efficiency and stability.

Regarding the IF methods

Although the memory efficiency of IF methods makes them suitable for large-scale BO to some extent, there are several inherent limitations associated with the approximation of these methods. As discussed in Appendix B.2, the approximation bias due to unsatisfied KKT conditions cannot be controlled by hyperparameter tuning. Additionally, the challenge of stochastic approximation in IHVP is highlighted in [1]. Specifically, while the Hessian $H$ is approximated via mini-batch sampling, it holds that $H^{-1} = (E[H])^{-1} \neq E[H^{-1}]$. This corresponding bias also cannot be mitigated through hyperparameter tuning.

In addition, following the reviewer's suggestion, we conducted additional experiments to complement the findings in Table 2, focusing on the comparison between IF methods and other approaches.

We considered two IF methods: Hessian-free and Neumann. The Hessian-free method approximates the inverse Hessian as an identity matrix scaled by a scalar. The Neumann method approximates the inverse Hessian vector product by utilizing the Neumann Series (see Equation 25, noting that the exponent $k$ is missing for $(I - \alpha H)$).

Due to the limited time available during the rebuttal period, we focused solely on the StreamingQA dataset. We fixed the base model as GPT2 and the unrolled depth at 48.

For the Hessian-free method, the scalar can be merged into the learning rate since the explicit gradient defined in Equation 3 is zero in this case. We selected the learning rate from $[1E-4, 5E-5, 1E-5, 5E-6, 2.5E-6]$, with $5E-6$ yielding the best validation loss. For the Neumann method, we tuned the hyperparameters $\alpha$ and $K$ in Equation 25. We conducted a grid search with $(\alpha, K) \in [1, 0.1, 0.01, 0.001] \times [10, 20, 40]$ and a learning rate of $2.5E-6$. The optimal combination was $\alpha = 0.01$ and $K = 40$. No further performance improvement was observed by tuning the learning rate or increasing $K$ to 80.

Please refer to Table A in the attached PDF in the general response for results. We observe the following:

(1) Both IF methods yield improvements on RGU (with an unrolled depth of 6).

(2) By carefully tuning the hyperparameters, the Neumann method achieves further improvements over the Hessian-Free method, albeit with additional computational cost.

(3) Despite the careful hyperparameter tuning for Neumann, its performance remains inferior to $\text{(FG)}^2$U.

Regarding the stochastic $\text{(FG)}^2$U

In the stochastic context, the bilevel optimization problem is formulated as:

$\min_{\phi\in \Phi}\ h(\phi):=f(\theta_T(\phi), \phi)=E_{\xi}[F(\theta_T(\phi), \phi; \xi)]$

$where\ \ \ \theta_0(\phi) = \Omega_0(\phi), \theta_t(\phi) = \Omega_t(\theta_{t-1}(\phi), \phi; \zeta)\in\Theta, t = 1,...,T,$

where $\xi$, $\zeta$ are random variables.

Assume that the sampled dataset for the meta-objective and the lower-level objective are respectively $D_{F}$ and $D_{G}$, related to random variables $\xi$ and $\zeta$. We have the hypergradient under stochastic context:

$\nabla_\phi ' h(\phi) = \frac{\partial F(\theta_T(\phi;D_G), \phi;D_F)}{\partial \theta_T} \frac{d\theta_T(\phi,D_G)}{d\phi} + {\frac{\partial F(\theta_T(\phi;D_G), \phi;D_F)}{\partial \phi}}.$ With Forward Gradient, the estimation of hypergradient would be $\nabla_\phi ' h(\phi) vv^{\top}$. Subsequently, from the independence between the forward gradient vectors $v$ and the batch sampling, we have $E[\nabla_\phi ' h(\phi)] = E_{\xi,\zeta}[E_{v}[\nabla_\phi ' h(\phi)vv^{\top}]]= E_{\xi,\zeta}[\nabla_\phi ' h(\phi)]$

$=E_{\xi,\zeta}\left[\frac{\partial F(\theta_T(\phi;D_G), \phi;D_F)}{\partial \theta_T} \frac{d\theta_T(\phi,D_G)}{d\phi} + {\frac{\partial F(\theta_T(\phi;D_G), \phi;D_F)}{\partial \phi}}\right]$

$= \frac{\partial f(\theta_T(\phi), \phi)}{\partial \theta_T} \frac{d\theta_T(\phi)}{d\phi}+ \frac{\partial f(\theta_T(\phi), \phi)}{\partial \phi}=\nabla_\phi h(\phi),$ which gives the unbiasedness of $\text{(FG)}^2$U for stochastic bilevel optimization problems.

Regarding the real space consumption

We direct the reviewer to Tables B, C, and D in the attached PDF in the general response.

[1] Making Scalable Meta Learning Practical, https://arxiv.org/abs/2310.05674

[2] Fine-Tuning Language Models with Just Forward Passes, https://arxiv.org/abs/2305.17333

[3] Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark, https://arxiv.org/abs/2402.11592

We thank the reviewer for the overall positive evaluation and will try to address the concerns raised point by point.

Regarding the complexity

The reviewer's understanding of Theorem 3.4 is generally correct. We would like to highlight the following points:

Firstly, the convergence rate in Theorem 3.4 is an upper bound. We direct the reviewer to our general response (A.2) for a more detailed discussion on the choice of $b$ and the computation in practice, covering the empirical success of FG/ZO in large-scale applications with constant batch sizes, techniques to reduce variance, the parallelizable nature of $\text{(FG)}^2$U which aligns well with modern AI computation, and the more cost-effective two-phase methodology in practice.

Secondly, the $\mathcal{O}(bM)$ memory cost with batch size $b$ can be reduced to $\mathcal{O}(M)$ through serial/parallel accumulation. In other words, we do not have to compute the full batch on a single GPU simultaneously. In practice, the memory cost of $\text{(FG)}^2\text{U}$ can be more manageable than FGU and RGU by choosing $b$ according to the hardware. It is important to note that both serial and parallel accumulation are not trivial for FGU and RGU.

Thirdly, the parallelization method for FGU (Plan A) proposed by the reviewer can be expressed as $d\theta / d\phi = Z = I Z = \sum_i^M \text{diag}(e_i) Z = \sum_i^M e_ie_i^T Z$. As mentioned in Line 122, a special choice of $v$ is the normalized orthogonal coordinates (Plan B), which leads to $Z vv^T = N Z e_j e^T_j$, where $j$ is a random index from $Unif([0,\ldots,N])$.

The differences are:

(1) Plan A maintains the rows of $Z$, while Plan B maintains the columns.

(2) Plan B utilizes randomization, while Plan A does not.

We argue that both choices of Plan B are better:

(1) Considering the update rule in Equation 4, row-wise updates require communication among threads, whereas column-wise updates do not.

(2) Plan A requires exactly $M$ threads, whereas Plan B provides smooth trade-offs via randomization.

Regarding the formula for data condensation

Yes, the reviewer is correct. We will fix the typo in the revised manuscript.