Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization
Authors: Qianli Shen
Yezhen Wang
Zhouhao Yang
Xiang Li
Haonan Wang
Yang Zhang
Jonathan Scarlett
Zhanxing Zhu
Kenji Kawaguchi
Venue: NeurIPS 2024 poster (2024)
Keywords: bi-level optimization
large-scale optimization
meta learning
Primary area: optimization_for_deep_networks
TL;DR: This paper introduces a new method for bi-level optimization that offers unbiased meta-gradient approximations with enhanced memory efficiency, and demonstrates superior performance in large-scale applications.

Abstract: Bi-level optimizaiton (BO) has become a fundamental mathematical framework for addressing hierarchical machine learning problems.
As deep learning models continue to grow in size, the demand for scalable bi-level optimization has become increasingly critical.
Traditional gradient-based bi-level optimizaiton algorithms, due to their inherent characteristics, are ill-suited to meet the demands of large-scale applications.
In this paper, we introduce **F**orward **G**radient **U**nrolling with **F**orward **G**radient, abbreviated as **$($FG$)^2$U**, which achieves an unbiased stochastic approximation of the meta gradient for bi-level optimizaiton.
$($FG$)^2$U circumvents the memory and approximation issues associated with classical bi-level optimizaiton approaches, and delivers significantly more accurate gradient estimates than existing large-scale bi-level optimizaiton approaches.
Additionally, $($FG$)^2$U is inherently designed to support parallel computing, enabling it to effectively leverage large-scale distributed computing systems to achieve significant computational efficiency.
In practice, $($FG$)^2$U and other methods can be strategically placed at different stages of the training process to achieve a more cost-effective two-phase paradigm.
Further, $($FG$)^2$U is easy to implement within popular deep learning frameworks, and can be conveniently adapted to address more challenging zeroth-order bi-level optimizaiton scenarios.
We provide a thorough convergence analysis and a comprehensive practical discussion for $($FG$)^2$U, complemented by extensive empirical evaluations, showcasing its superior performance in diverse large-scale bi-level optimizaiton tasks.

Memory-Effiiien Gradient Unrolling for Large-Scale Bi-level Optimization
Qianli Shen
Yezhen Wang'
Zhouhao Yang'
Xiang Li
Haonan Wang'
Yang Zhang'
Jonathan Scarlett'
Zhanxing Zhu?
Kenji Kawaguchi
National University of Singapore
University of Southampton, UK
Abstract
Bi-level optimization (BO) has become fundamental mathematical framework for addressing hierarchical machine learning problems. As deep learning models continue to grow size, the demand for scalable bi-level optimization has become increasingly critical. Traditional gradient-based bi-level optimization algorithms, due their inherent characteristics, are ill-suited to meet the demands of large-scale applications. In this paper, we introduce Forward Gradient Unrolling with Forward Gradient, abbreviated as (FG)UU, U, which achieves an unbiased stochastic approx- imation of the meta gradient for bi-level optimization. (FG)²U circumvents the memory and approximation issues associated with classical bi-level optimization ap- proaches, and delivers significantly more accurate gradient estimates than existing large-scale bi-level optimization approaches. Additionally, (FG)*U is inherently de- signed support parallel computing, enabling it effectively leverage large-scale distributed computing systems to achieve significant computational efficiency. practice, (FG)²U U and other methods can be strategically placed different stages of the training process to achieve more cost-effective two-phase paradigm. Further, (FG)²U is easy to implement within popular deep learning frameworks, and can be conveniently adapted to address more challenging black-box bi-level optimization scenarios. We provide thorough convergence analysis and a comprehensive practical discussion for U, complemented by extensive empirical evaluations, showcasing its superior performance in diverse large-scale bi-level optimization tasks. Code is available at /ttpt::///// /tte::///iiii /ggthuu...... com/ wttm/Shnnianni/ieenn
Introduction
Bi-level optimization is mathematical framework with long history of research 1T dealing with hierarchical optimization problems where one problem is nested within the other. A bi-level optimization problem can be formulated as:
65
min f(0" (())) ②) s.t. 0* (Φ) min ((,, +).
(1)
where 0 E θ C RM denotes the inner parameter, & O R denotes the meta parameter, and f. g are called the meta objective function and inner objective function, respectively.
Recently, with the rise deep learning, bi-level optimization has regained attention as a theoret- ical framework covering wide range of machine learning problems, including hyperparameter neural architecture search [78] 38 14), robust machine learn-
optimization [6 43
16 45
ing 779 76 71 26 meta learning [3]
(33 49 2 and physics- ahysics-informed machine learning 23] 62)
In these scenarios, the inner problem often pertains to the optimization of neural networks, thereby
"Correspondence: shengianli@u edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Category
Method
Constant Memory
Hessian Free
Stochasci Optinization
Grad. Approx.
PERFORMANCE
EFFICIENT GRAD APPROX
GUI VF
Classical
30
k7
33
X
TRGU
60
Efficiency
GRAD APPROX
Large-scale Hessian- Froe
26
75
Efficieccy
IPG) (ours)
Accuracy
TRAINING
Learning Online Adastation (DSA/OPP2)
Outa Condersation ICISAR IPC-50)
Condensation (OFAR (PC+58)
30
Hesdor Iree
CTU (G
Efficenccy (uscroou)))
Memary toage
21
GPUs
(FOPU GRUS
12
07107
15
TMGJ
10
155
35
MethadUroolled Depth
Method
GPU Homory Usage 1G8B1
Figure 1: Top Left: A comparison of bi-level optimization methods. (FG) circumvents the large-scale challenges inherent in classical bi-level optimization techniques. Within large-scale bi-level optimization, (FG)"U prioritizes the accuracy of gradient approximation over efficiency. Top Right: An overview of the cost-effective two-phase paradigm. (FG)²U is ideally positioned in Phase to enhance performance after an approximate solution has been obtained using other efficient methods. Bottom Left: GPU Memory Usage and Performance on Meta Learning Online Adaptation experiment. (FG)*U can effectively address the memory issue of RGU when both the inner model and the unrolled depth are large. Bottom Center: GPU Memory Usage and Performance on Data Condensation experiments. The performance of (FG)U surpasses that of other large-scale bi-level optimization methods, owing to its accurate gradient approximation, while demonstrating better memory efficiency. Bottom Right: Efficiency tradeoff of (FG)²U on Data Condensation experiments. The efficiency U can be well enhanced via intra/inter-GPU parallelism.
precipitating challenges associated with gradient- based bi-level optimization. Consequently, various gradient-baal -based bi-level optimization algorithms have been developed 773.. These algorithms typically employ iterative solution 0r obtained by executing multiple inner optimization steps to approxi- mate the meta gradient, and provide different tradeoffs between computational costs and performance for meta gradient approximation.
However, as the scale of deep learning models continues to expand, the requirements for scalability in bi-level optimization correspondingly increase. Existing gradient-basee bi-level optimization algorithms, due their inherent characteristics, are ill-suited meet the demands of large- scale applications. Concretely, gradient unrolling (GU) methods [17] 16 40 60 are bottlenecked by
the memory overhead associated with either the dimension of the inner parameter or the number
of iterative steps for the inner problem. Implicit Function (IF) approaches [43
19
64
76 are
compromised by approximation errors, which stem from the iterative estimation of inner solutions and computations that involve the Hessian matrix. Value Function (VF) based strategies 39
37 although exhibit commendable theoretical properties [8] for deterministic bi-level optimization, have yet to gain traction in practical applications, predominantly due to their limitations in addressing large- scale stochastic challenges [73]. Recent advancements algorithms [60 9 have been specifically tailored for large-scale bi-level optimization. Although these methodologies facilitate efficient gradient approximation by compromising accuracy, they may result in significantly suboptimal performance due biased gradient approximations. Additionally, these methods struggle in more complex scenarios, such as when inner problems are addressed through black-box optimization.
61 33
In this paper, propose novel method called Forward Gradient Unrolling with Forward Gradient, abbreviated as (FG)²U, which achieves an unbiased stochastic approximation of the meta gradient for bi-level optimization. (FG)2U U circumvents the memory issues associated with GU-based approaches and approximation issues associated with IF-based approaches. Compared recently developed large-scale bi-level optimization approaches, (FG)²U delivers significantly more accurate gradient estimates. Additionally, (FG)²U is inherently designed to support parallel computing, enabling it to effectively leverage large-scal distributed computing systems to achieve significant computational efficiency. In practice, cost-cffective two-phase paradigm can be achieved by strategically placing

---

(FG)'U and other methods at different stages of the training process to balance efficiency and performance. Further, (FG))U is easy to implement within popular deep learning frameworks, and can be conveniently adapted to address more challenging zeroth-order bi-level optimization scenarios.
We provide an overview of (FG)'U in Figure to illustrate its strengths and role in large- scgee bi-level optimization. The rest of the paper organized as follows. Firstly, in Section 2 we provide summaries of existing bi-level optimization algorithms and discuss their limitations in large-scale
contexts. Next, in Section 3 we introduce the propased method, (FG)²U, followed by a convergence analysis Section (.11 and detailed discussion of the practical considerations in Section 3.2) Further, in Section 4 we conduct extensive empirical studies covering large-scale bi-level optimization in
computer vision, natural language processing, and physics- informed machine learning to demonstrate the efficacy of (FG)*U U in large-scale bi-level optimization scenarios.
2 Background
Gradient-based Bi-level Optimization. Within deep learning applications, the model concerned with optimizing over as presented in ( typically constitutes deep neural networks. The optimal pa- rameters of such networks are not explicitly accessible and are estimated through iterative procedures. Consequently, the primal problem of bi-level optimization (T) is approximately reformulated as follows:
min ((6) =- ((0r(0),
ΦE
where Ou(Φ) (2,(d). (Φ) (,(00-1 Φ) E E t= 1=
where & C RN,, C RM are the parameter spaces; T, commonly called the unrolled depth, denotes the number of inner optimization steps for approximating 0" (Φ)); 82o IR - R- specifies the initialization of the inner optimization, and 52 θ x + - & delineates the transition dynamics of the inner optimization 0,-1 Nh V eg (0-10.10 Φ). where denotes the step size at timestep L.
timestep In particular, for gradient descent, (2, (Ø),
To optimize & using first-order method, it is necessary estimate the meta gradient Vah, which can be further decomposed according to the chain rule:
V @f(Or()), )) d0 d( 0f(0((0).)) 00,
ds
QΦ
meta gradient
implicit gradient
explicit gradient
The computation of meta meta-gradient poses significant challenge, primarily due to the need for efficient approximation of the implicit gradient. This task is complicated by the recursive dependency of Or on Φ. To surmount this challenge, a variety of gradient-based bi-level optimization algorithms have been developed, as extensively reviewed recently in [73]. These algorithms can be fundamentally categorized into three types based on their approach to meta- meta-gradient approximation: Gradient Unrolling (GU), Implicit Function (IF), and Value Function (VF). Recent innovations such as truncated RGU (TRGU) (60) and Hessian-Free approaches (76 75 9 which are predicated on GU and IF methodologies respectively, have introduced significant biases in their approximations to accommodate the computational constraints of large-scale scenarios. In the subsequent paragraph, we furnish a concise overview of GU-based approaches, addressing their non-constant memory issues in large-scale applications. Extended discussions on the remaining methods are reserved for Appendix B
Gradient Unrolling. The core idea behind GU 11 16 40 60 entails unrolling the inner optimiza- tion into an expansive computational graph, followed by the employment of automatic differentiation (AD) techniques for the iterative computation of gradients.
Forward Gradient Unrolling (FGU) 17 16] computes the meta gradient using the following forward recursive formula, starting from Zo
dgf(
de
d0,())
an, (0, (0), ))
(0)1(()) do
an, (0-10, Φ)) t= 1, T,
(4)
dΦ
00
dΦ
Z:
A.
Z
B
3

---

Reverse Gradient Unrolling (RGU) 46 16], instead of the employment of explict reccursive formulas of TT, focuses on the implicit reccursive formulas of Vah:
V Vah(Ø)= @f(Or(Ø), Φ) d0r(Ø)
of(Or(0).)) dΦ
d0r
dΦ
Zr
Cr
(5)
=d Z CT
dr Ar Z -r +
Br =
di C
dr
cr-1
The corresponding reverse recursive formulas can thus be summarized C-1 = + d,B,, d =-11 L=T,....
(6)
Weakness (GU): Non-Constant Memory. Both GU approaches exhibit non-constant memory overhead, which constrains their utility in large-scale scenarios. The forward reccursive formulas in (4) revolve around the Jacobian matrix product, demanding ((MN) space consumption. The reverse recursive formulas in (6) necessitate the storage of the entire trajectory of the inner optimization 0..r for backward computation, thereby imposing memory requirement of O (M). These requirements are often impractical for aarge-scale bi-le optimization, when & and @ are of high dimension and significant unrolled depth is required.
Forward mode automatic differentiation (forward mode AD) has been applied variety of research fields, including the training of recurrent neural networks (70), the computation of Hessian vector products [50], etc. However, the computation of the true gradient via fordard-mode AD requires the full Jacobian, which is typically costly to compute.
Forward Gradient.
To solve this, forward gradient learning (69 4
56], built upon forward-mode AD, was proposed. Forward gradient methods update parameters based on the directional gradient along random perturbation direction for backpropagation- -ree training. More formally, given a differentiable function h I I, the gradient for given input RVV can be approximated as
(M( = Vh(a)vo".
where e ~ p(t)) is N-dimensional multivariate random variable, satisfying B(vw?) = I. Common choices of the distribution include Rademacher ~ Unif([-1, ), Gaussian v ~ No I), and uniform distribution over a set of normalized orthogonal coordinates v ~ Unif((//e.h.w)) For any given Φ, (h(Φ) is an unbiased estimator of Vh()), as [T)] = [[V)(())0=)) Vh(o)E(ve71 = Vh( Vh(Ø). Despite the unbiasedness of Vh, the dimension- dependent variance of the estimated gradient with a single direction impedes the scaling-up to high-dimensional problems. In practice, Monte Carlo gradient estimation can be used via averaged forward gradients over multiple random directions to reduce the variance.
3 (: Forward Gradient Unrolling with Forward Gradient
We aim to circumvent the memory overhead issues associated with forward gradient unrolling (FGU) as discussed in Section 2
We begin by examining the forward gradient of h at Φ.
v1(() = Vl( (d (dZZTu+cTe)e"
(8)
where a ~ p(v) is N-dimensional multivariate random variable, satisfying E(voT] =I We follow the idea of FGU introduced in Section 2 to compute Zrv. v. By multiplying both sides of by on the right, we can obtain the recursive formulas for Z, as
Zo Bov; Z,U 1.Z Biv, t= 1=
The revised recursive formulas in 19 facilitate the tracking of a M-dimensionnal vector Z,U, rather than full Jacobian Z, of size M N, throughout the forward pass. The stochastic estimation in (8) is unbiased, adhering to the properties of forward gradient methods. To reduce the variance,

---

Monte Carlo estimate via averaged forward gradients over b i..dd. random directions:
Vh(σ) 1 Vh(o)u,u, 1 6 E d Zr, cruu)u."
(10)
We call this algorithm (FG)UU, abbreviation of Forward Gradient Unrolling with Forward Gradient. The algorithm is summarized in Appendix A as Algorithm
Compared to GU-based methods, as discussed in Section 2 (FG) eliminates the dependency on the meta parameter dimension N and the depth of unrolling I' without introducing bias, significantly enhancing memory efficiency. Unlike ba-dd methods, as discussed in Appendix B.2 (FG)* U
overcomes the approximation issues associated with them while maintaining a constant memory overhead, thus providing superior gradient approximation. Compared to TRGU and Hessian-Free methods, which compromise approximation accuracy for efficiency, (FG) consistently delivers accurate gradient approximations. The computational efficiency of (FG))U U can be further enhanced by leveraging large-scale distributed computing resources, capitalizing on its inherently parallelizable formulation as presented in (10). In practice, a more cost-effective two-phase paradigm can be achieved by strategically placing UU and other methods at different stages of the training process. we will diseuss in Section 3.2 For an illustration of the role of (FG)*U in large-scale bi-level optimization, please refer to Figure
3.1 Convergence
In this section, we provide convergence analysis for (FG) The proofs can be found in Appendix C First, establish bound the variance of the estimated gradient, when employing random vectors whose entries follow the Rademacher distribution.
Lemma 3.1. For any & ④, v, Unif([-1, 1), the gradient estimation in (10). satisfies
vh( =
where p = - E (0, 1] as the sample size is selected 1, , 1.
The resultant error is bounded by O (امل), where represents the sample size used for computing the forward gradient, and N is the dimensionality of the gradient itself. This bound demonstrates how the error scales inversely with the sample size while also being influenced by the gradient's dimensionality.
Next, we lay down the following assumptions, on which our main theorems are based. Let &b
(0, Φ) E θ x & denote the combination of the lower-level parameter O and the meta parameter Φb Following existing papers on the theory of bilevel optimization [55) 60 28), Assumption 3.2 we adopt some standard assumptions over the smoothness of the objective functions f and 9.
The meta objective function f(bb) and the lower-level objective function g(bb) are both C-Lipschitz and L-smooth, i.e.,for any 1/, r/² e x b
Assumption 3.2.
(f(b) - f((b') < cl 2b'll. [V - f((bb)) s L| - 2b11. (g( g (b)) - cl 2b'l.
(11)
[Vg({b) - Vg(hb))l I| - L| ه'1l
(12)
The next assumption regulates that the transition functions Ω satisfy similar smoothness conditions. Assumption 3.3. The transition functions Ωa:r are Cn-llipschitz and La-smooth, i.e., for any Φ. Ø'
[ (b) 2 < Call Φ'1. [V) (L. < For any wb 26 θ x Φ ,,...,T, T, |2 (+) - (2((")))))) || < Cn Cn A'1. (b) - (2)'')) اn
Tall -'ll-
(13)
(14)
Assumption 3.3 made ensure the generality our analysis over different optimizers. Note that Ω is scheme-dependent w.r.t. the gradient-based optimizer we adopt for lower-level problems. In many cases, such as gradient descent where Ω2 (1/r1)) = 0,- ?) Vog(u,
Assumption 3.3
direct consequence of Assumption 3.2
5

---

We propose the following theorem and remark for convergence analysis of (FG)³U on problem (2) Notice the convergence result can be extended to the primal BO problem with some further assumptions. We place proof scratch and some discussions in Appendix
Theorem 3.4 (Convergence). Suppose that Asumption 3.2 and Assumption 3.3 hold. Setting the learning rate 3 To+iyL for gradient descent over the hyperparameter Φ. then there exists constant La (depending on C, L, Cr. La, and T, and formally in the proof) such that
4Ln (E[h(0o)]
- h(Φ))
E [IVA((W))
(15)
K
pK
Remark 3.5. Theorem 3.4 shows that Algorithm Converges to an e-accurate stationary point with convergence rate of Ole p
Recall that p -- which indicates that the convergence rate is linearly dependent on N, which poses a significant challenge when managing high- high-dimensional meta-parameters Φ. However, is important note that the cemeenionnddeendeee scalability has been found to be feasible with several practical considerations, as discussed in the following subsection.
convergence rate represents an upper bound, and
3.2 Practical Considerations
Choice of b. According to the convergence analysis in Section 3.1 a sample size of b - O(N) is required to achieve a convergence rate of O(E-1). However, it has been widely observed that forward gradient and zeroth-orde optimization, despite having dimension-dependee well empirically with b = (1) in large-scale scenarios, such as in LLM fine-tuning 47 74). In this paper, we select the largest possible b that does not exceed the GPU memory limit for our empirical study. Additionally, gradient accumulation is utilized to further control variance and stabilize the training process.
convergence rates, work
Cost-Effective Two-Phase Paradigm. It is important to note that the upper bound delineated in (15) linearly depends on the performance discrepancy between the initialized meta parameter Φ and the optimal. This dependence motivates the adoption of more cost-effective two-phase paradigm for large-scale bi-level optimization. In the initial phase, we utilize efficient yet less accurate gradient approximation methods, such as TRGU [60] and Hessian-Free [9), to efficiently establish an initial Φ, that surpasses random initialization, while keeping computational overhead manageable.
Subsequently, in the second phase, ( is utilized for a more accurate, albeit less efficient, gradient approximation further elevate the performance, leveraging extensive computational resources.
The technique employed in computing Vh())v is identified as forward-mode automatic differentiation (forward- AD). In advanced automatic differentiation libraries, such as JAX (5] and PyTorch B. forward-mode AD is efficiently implemented as aecobian-vector product (jvp), without the necessity of explicitly computing the Jacobian matrix. The FL.OP cost of jvp is approximately three times that of a standard forward pass, while the memory overhead is doubled. In practice, is only necessary to define the forward computational graph of inner optimization and invoke forward-mod AD, which simplifies the implementation process significantly. Regarding distributed training, JAX offers the vmap interface for efficient unU parallelism and the pmap interface for effective inter-GPU parallelism.
Implementation.
Zeroth-order Bi-level optimization. In certain applications of bi-level optimization, the inner problem is approached as a black box, where the gradient of S2 is inaccessible, rendering the analytical gradient unrolling unfeasible. For example, in PDE-constrained optimization [3] 62). in which the inner problem entails solving a Partial Differential Equation (PDE) using a non-difierentiable solver. In such scenarios, rather than employing forward-mode Automatic Differentiation (AD), can resort to Finite Difference methods to approximate the directional gradient Vh(Φ)v by
h(Φ + pv) - h h())) h + jiv) h(())
Vh(Φ)n
= lim p-0
(16)
P
with sufficiently small positive >
O. We refer to this zeroth-order variant of (FG)²U as (FG)²U- 70, noting that the computation solely encompasses two forward passes and does not involve the utilization of any first-order information. The memory complexity is the same forward-mi mode AD

---

and the actual computation time will be slightly less than forward-mooe AD, at the cost of introducing an approximation bias. We give a more detailed discussion within the context of zeroth-order optimization 41 in Appendix D and empirically study corresponding case in Section
Experiments
We conduct experiments across various contexts, as detailed in the respective subsections. Initially, we engage in an image data condensation task, where we focus on comprehensive performance comparison between (FG) and both classical and large-scale bi-level optimization algorithms. Subsequently, we investigate meta- meta-learning for the online adaptation of language models, employing GPT model as the inner model, to illustrate how (FG)²U effectively circumvents the non-constant memory issue associated with RGU. Finally, we address a physics-informed bi-level optimization problem, where gradient-based inner solvers are ineffective, to demonstrate the efficacy of combining U-Z), the zeroth-orde variant of (FG)UU discussed in Section 3.2 with non- con-differentiable numerical solvers.
To overcome the challenges posed by large-scale datasets, a line of works known data condensation 668 72 has been proposed. The main idea is generate compact, synthesized dataset, designed to elicit similar behaviors machine learning models as those trained with the original, massive dataset. The objective of the mainstream principles 722 designed for data condensation can be naturally formulated as bi-level optimization problem. We focus on the best-known principle performance matching 72] classification tasks, which can be formulated as
Data Condensation.
min [(Or; where 0g = 0, - () t= 1..,..., D
(17)
where Do. De respectively denote the original and condensed dataset, 9 denotes the model parameter, denotes the cross-entropy loss function, and 7) represents the step-size for inner optimization.
Approaches
For Reference
Dataset
IPC
Ratio (%)
TRGU
Hessian-Free
Neumann
(
RGU WHOLE
1
0.017
1..76
63.98
6..37
82.44
92.32+0.3
MNIST
10
0.17
94.05
94.97-1044
96.12 95.75
2022
96.79
99.6=a.0
50
0.83
96. +0.41
96.34+0a1
96.78
97.01 +0.19
97.72+003
0.02
20.78
19.72
21.33
29.37 20:75
34.08
CIFAR-10
0.2
44.01 ::0.7
1.002:+102
4..67
50.10
0.556
53.15
84.8=a.0
50
1
49. 2044
48.73
50.02+t0.9
51.98 00.44
56.37+4037
0.2
3.96
3.14#04:
4.52+:0.56
8.22
15.61
CIFAR-100
10
2
20.20 20.6
19.01
20.87:2082
23.38=0.33
25.42:005
56.2#030
50
10
22.33
23.59
24.52.10.7
25.84_0031
28.52:t0.5
Table 1: The performance (testing accuracy %) comparison among various bilevel optimization methods on the data condensation task over three datasets. All the datasets are condensed using a 3-layer ConvNet. IPC: image(s) per class. Ratio (%): the ratio of condensed examples to the whole training set.
We conducted our experiments following the standard data condensation setting established by [68] 77 677. more detailed task description is given in Appendix E. and implementation details given in Appendix F.
The condensed datasets are evaluated using 3- 3-layer convolutional networks with randomly initialized parameters, and the average accuracies on test datasets are summarized in Table [ Compared to large-scale bi-level optimization methods like TRGU and Hessian-Free, which prioritize efficiency at the expense of approximation accuracy, (FG)²U exhibits significantly better performance, due to more accurate gradient approximation as explained in Appendix B Additionally, we assessed Neumann Series (denoted Neumann in Table [. an IF-based method that mitigates gradient approximation enrors through extended computations, as introduced Appendix B.2 While it demonstrates perfor- mance enhancements over the Hessian- Free method, Neumann still yields suboptimal performance
7

---

compared to UF owing to the inherent bias of the IF-based method. Further discussions and supporting evidence are available Appendix (B.2 B.2
The results of RGU, which represent the upper performance bound for both TRGU and (FG)UU, are provided for reference, along with the results from training on the entire dataset (denoted as WHOLE in Table [ representing the upper performance bound for all approaches. However, it is crucial to acknowledge that RGU is not practical in large-scale bi-level optimization scenarios due to its non-constant memory requirements, as discussed in Section [ This limitation will be further exemplified in the subsequent, where the inner model is significantly larger. In principle, the performance of ( U can be further improved to approach that of RGU by increasing the number of random directions for gradient approximation.
The memory and computational efficiencies of TRGU, Hessian-Free, and (FG) in the most challenging case (CIFAR-100, IPC=50) are reported in I (Bottom Right), demonstrating that the efficiency of UU can be significantly enhanced through intra/inter-GPU parallelism.
Meta Learning Online Adaptation of Language Models. The online adaptation of language models (LM) has been studied recently to keep the knowledge of LM current [34 (2].. However, trivial auto-regressive fine-tuning the LM, which applies uniform weights to all tokens, often results in suboptimal performance in downstream tasks. This issue stems from the default average negative
log- eigllikelhoodd (NLL) loss, which fails to capture the significance of tokens (25). To overcome this limitation, [25 proposed Context-aware Meta-learned Loss Scaling (CaMeLS), a strategy that employs meta-learning to adjust token weights for more effective online adaptation. Specifically, they meta train a weight model to reweight the auto-regressive loss during online fine-tuning, aiming to enhance LM performance on downstream preestion-answenigg tasks. A comprehensive task description and the mathematical formulation of the objectives are detailed in Appendix E.2
The trained weight model subsequently fine-tuned on unseen online documents and evaluated on corresponding question -nntwernn tasks. In [25 RGU is utilized for meta gradient approximation. To mitigate the non-constannt memory issue associated with RGU, a DistilGPT2 model [59] is chosen as the surrogate base model for training the weight model, instead of larger models typically employed for online adaptation. Additionally, a very limited unrolled depth of 6 utilized within a 40 GiB
GPU memory budget. In our experiments, since (FG)'U has circumvented the non-constant memory issue associated with RGU, we are able to increase the unrolled depth and upscale the base model for training the weight model. Empirical evaluations are conducted on two datasets, StreamingQA (36
and SQuAD-Seq 54
StreamingQA
SQuAD-Seq
Model (# params)
Method
EM (1)
Fl (t)
EM (t)
(t)
CaMeLS + RGU 25
66
1.62
5.79
1.45
3.08
DistilGPT2 (82M)
CaMeLS + RGU (impl.)
2.04
5.53
1.52
3.16
CaMeLS + (FG)²U (ours)
2.22
6.37
1.72
3.50
CaMeLS + RGU (25 66
5.35
10.60
4.97
8.63
GPT2-Large (774M)
CaMeLS + RGU (impl.)
7.02
12.19
4.86
8.57
CaMel + (FG)²U (ours)
7.21
12.50
5.56
8.99
CaMeL RGU 25 66
6.55
11.67
6.70
10.15
GPT2-XL (1.5B)
CaMel + RGU (impl.)
7.93
12.94
6.71
9.65
CaMeLS (FG) (ours)
8.89
14.42
7.37
10.37
Table 2: Comparison of the online adaptation performance. The reported evaluation metrics include
the exact match (EM) and Fl scores. For vanilla CaMeLS (25). RGU is conducted with unrolled depth 6, using DistilGPT2 as the base model. We present both the results reported by [66] and those from our implementation (denoted as impl.). For CaMeL.S + (FG)²U, we select unrolled depths from (24, 48), and the base model from (DistilGPT2, GPT2]. We report the results for the combination that yields the best Fl score. Additional details and ablation studies are documented in Appendix G.I
Firstly, we increased the unrolled depth while maintaining the base model DistilGPT2. We plotted the Fl scores and GPU memory usages for RGU with unrolled depths of {1, 2, 4. 6) and (FG)²U with unrolled depths of (24, 48) on StreamingQA in Figure (Bottom Left). The performance of the weight model is positively correlated with the unrolled depth. substantiating the benefits of
8

---

training with larger unrolled depths. The non-constant memory issue associated with RGU can be observed when the unrolled depth increases, while (FG)"U maintains constant memory even with large unrolled depth. Subsequently, we endeavored to upscale the base model GPT2 to reduce the disparity between training and evaluation. The performances are summarized in Table 2 with detailed
ablation studies on unrolled depths and base model variants documented in Table G.I and Table G.2
Data-driven Discovery of Partial Differential Equations (PDEs). Let us consider the following general forms of parametrized and nonlinear PDEs:
Ut + Muu Φ) = 0, x V E,
(18)
where denotes the space-time coordinate, v denotes bounded domain with boundary, [o, T] x - IR denotes the latent solution, 4, represents the first-order derivative of u with respect to t, and N is general differential operator parameterized by Φ, acting V This setup encompasses broad spectrum of problems in physics For example, the one-dimensional Burgers' equation is defined by N iu Φ) = ,uuu, Vulaa where (H,)) v) IR?, and Ugy Uzx represent the first and second-order derivatives of with respect to x, respectively.
Burgers: PDE Salver Efficiency
Method
(FG)UU
(FG)"U-Z0
PINN (Adam)
Relative L2 Error (log)
MNN 15CD)
Inner solver
PINN [2
Numerical
16-1
(E 4)
2143.58 1855..6
0.94
Burgers
Effcient
e. (B-3,
336.062:6.99
0.63
1F-2
Allen-Cahn
(F-2,44
438. 13
2.34
e.. (F-3,4)
133.61
0.97+0.4
E (E-2,4)
94.40+4.31
0.72=0.677
Kdv
50
150
250
e. (E-33,,
837.81167.01
167.01
2.72:1.55
Speed (second per solution)
Figure 2: Left: Comparison of efficiency between the PINN solver and the numerical solver. We evaluated Adam [29] and SGD as the inner optimizers for the PINN solver, with steps ranging from 100 to 50,000. The results demonstrate that the numerical solver is significantly more efficient. Right: Comparison of relative L2 errors in the prediction of and u. EΦ = 1l@ /1 ||| 6u ||Upred u/|2/||u||/
The problem of data-driven discovery of PDEs 52 can be framed as follows: given a set of scattered observations of the latent solution u(x)) what are the parameters most accurately describing the observed data? The problem can be formulated a PDE-constrained optimization problem (PDECO):
min By,u~D [u(x; cb) - ul2 s.l. de + Nlu(s Φ); Φ] = 0, V, +
(19)
where D = { (Ix, 4,)): denotes the observed data. In cases where the closed-for form solutions of the nonlinear PDEs are intractable, parametric solutions ue are used to approximate the latent solution u for given Φb The PDECO in (19 is then reformulated into a bi-level optimization problem:
min Er (1) Φ)- ul? s.t. ( =
Φ). = I.
(20)
Employing gradient-based PDE solvers, such as physics-informed neural networks (PINN) 52
facilitates the direct application of (FG)²U. However, demonstrated in Figure 2 (Left), the accuracy and efficiency of PINNs fall short the rigorous demands of scientific computing. This limitation has prompted us to integrate faster and more accurate traditional solvers like the spectral method [I] (see also Appendix E.3.4 to tackle the inner problem. Given these solvers are non non-differentiable, we employ (FG))U-ZO, the zeroth-order variant of (FG)2U introduced in Section 3.2 to solve the problem.
We conduct experiments on three non-linear PDEs: Burgers, Allen-Cahn, and KdV, with a more de- tailed task description available in Appendix E.3 The results are summarized in Figure F We can observe that the combination of U-Z) and the numerical solver significantly outperforms (FG)UU and the PINN solver, in terms of both the prediction on & and u. The implementation details documented in Appendix

---

Conclusion
In this work, we propose a novel algorithm Forward Gradient Unrolling with Forward Gradient, abbreviated as (FG)*U, designed tackle the challenges associated with large-scale bi-level op- timization. We conduct convergence analysis of (FG)',, perform extensive comparisons with existing methods, and provide detailed discussions on its practical applications. Additionally, we undertake empirical evaluation across series of large-scale bi-level optimization tasks. Our findings indicate that (FG)*U effectively complements existing bi-level optimization algorithms, addressing gaps in large-scale bi-level optimization scenarios.
Limitations and future works. The experiments conducted in this paper are of relatively small scale, with the largest inner model being GPT-2 model. We look forward validating its effectiveness on larger-scale bi-level optimization tasks. Additionally, the application of black-box bi-level optimization and the potential of (FG)² remain underexplored, considering the prevalent black- box interaction between users and models today. We hope our work will inspire further development of large-scale bi-level optimization algorithms and their application in corresponding scenarios. Furthermore, we have not specifically addressed the efficiency issues inherited by (FG)²U from the forward gradient method. Enhancing the efficiency of (FG)²U while maintaining its gradient estimation accuracy will be an important direction for future research.
6 Acknowledgements
This research supported by the National Research Foundation Singapore under the AI Singapore Programme (AISG Award No: AISG2-TCC-003-000--2 Academic Research Fund Tier (Award No: T1 251RES2207, T1 251RES2218). The computational work for this article was partially performed on resources of the National Supercomputing Centre, Singapore (https:/wwww.sccc.s)
and the Singapore Ministry of Education
10

---

References
[1] William F Ames. Numerical methods for partial differential equations. Academic press, 2014.
[2] Marcin Andrychowicz. Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
[3] Jason Ansel. Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky. Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary De Vito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalamharkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang. Yinghai Lu, CK Luk, Bert Maher. Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet. Eikan Wang. Xiaodong Wang. William Wen, Shunting Zhang. Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '24). ACM, April 2024.
[4] Atilm Güneş Baydin, Barak A Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation. arXiv preprint arXiv:2202. 085887 2022.
[5] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:180 187-1999,, 2020.
[7] Claudio Canuto, M Yousuff Hussaini, Alfio Quarteroni, and Thomas A Zang. Spectral methods: funda- mentals in single domains. Springer Science & Business Media, 2007.
[8] Lesi Chen, Yaohua Ma, and Jingzhao Zhang. Near-optimal fully first-order algorithms for finding stationary points in bilevel optimization. arXiv preprint arXiv:23066. 14853, 2023.
[9] Sang Choe, Sanket Vaibhav Mehta, Hwijeen Ahn, Willie Neiswanger, Pengtao Xie, Emma Strubell, and Eric Xing. Making scalable meta learning practical. Advances in neural information processing systems.
36, 2024.
[10] Benoît Colson, Patrice Marcotte, and Gilles Savard. An overview of bilevel optimization. Annals of operations research, 153:235-256, 2007.
[11] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina N. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint (rXiv:1810.0480) 2018.
[12] Tian Dong, Bo Zhao, and Lingjuan Lyu. Privacy for free: How does dataset condensation help privacy? In International Conference on Machine Learning, pages 5378-5396. PMLR, 2022.
[13] John C. Duchi, Michael 1 Jordan, Martin J Wainwright, and Andre Wibisono. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions on Information Theory, 61(5):2788-288 2806, 2015.
[14] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: survey. Journal of Machine Learning Research, 20(55):1 21, 2019.
[15] Chelsea Finn, Picter Abbeel, and Sergey Levine. Model-agnostic meta-lcarning for fast adaptation of deep networks. In International conference on machine learning, pages 1126-1135. PMIRR 2017.
[16] Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient- based hyperparameter optimization. In International Conference on Machine Learning, pages 1165 1173. PMLR, 2017.
[17] Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning International conference on machine learning, pages 1568-1577. PMLR, 2018.
[18] Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better fow-shot leamerr. arXiv preprint arXiv:2012.1572. 2020.
[19] Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. arXiv preprint arXiv:1802. 1802.02246. 2018
[20] Ian J Goodfellow, Mehdi Mirza, Da Xiao, Aaron Courville, and Yoshua Bengio. An empirical investigation of catastrophic forgetting in gradient- neural networks. arXiv preprint arXiv: 13Xi::3112.62111 2013.
[21] David Gottlieb and Steven A Orszag. Numerical analysis of spectral methods: theory and applications. SIAM, 1977.
11

---

[22] Evan Greensmith, Peter L Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates reinforcement learning. Journal of Machine Learning Research, S(9), 2004.
[23] Zhongkai Hao, Chengyang Ying, Hang Su, Jun Zhu, Jian Song, and Ze Cheng. Bi-level physics- informed neural networks for pde constrained optimization using broyden's s hypergradients. arXiv preprint arXiv:2209.07D 2022.
[24] Ryuichiro Hataya and Makoto Yamada. Nyström method for accurate and scalable implicit diffcrentiation. In International Conference Artificial Intelligence and Statistics, pages 4643-4654, PMLR, 2023.
[25] Nathan Hu, Eric Mitchell, Christopher D Manning, and Chelsea Finn. Meta-learning online adaptation of language models. arXiv preprint arXiv:2305. 230.:35555577 2023.
[26] Ronny Huang, Jonas Geiping, Liam Fowl, Gavin Taylor, and Tom Goldstein. Metapoison: Practical general-purpose clean-label data poisoning. Advances in Neural Information Processing Systems, 33: 13:12001 20800 12091, 2020.
[27] Joel Jang, Seonghyeon Ye, Sohee Yang, Joongbo Shin, Janghoon Han, Gyeonghun Kim, Stanley Jungkyu Choi, and Minjoon Seo. Towards continual knowledge learning of language models. arXiv preprint arXiv:2110.03215.22 2021.
[28] Kaiyi Ji, Junjic Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In International conference on machine learning, pages 4882-4892. PMLR, 2021.
[29] Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6900 2014.
[30] Steven George Krantz and Harold R Parks. The fuplicitfunction function theorem: History, theory, and applications. Springer Science Business Media, 2002.
[31] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[32] Harshat Kumar, Dionysios S Kalogerias, George J Pappas, and Alejandro Ribciro. Zeroth-order determin- istic policy gradient. arXiv preprint arXiv:2006.077
2020.
[33] 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, pages 9083318813. PMLR, 2023.
[34] Angeliki Lazaridou, Adhi Kuncoro, Elena Gribovskaya, Devang Agrawal, Adam Liska, Tayfun Terzi, Mai Gimenez, Cyprien de Masson d' Autumm, Tomas Kocisky, Sebastian Ruder, et al. Mind the gap: Assessing temporal generalization in neural language models. Advances in Neural Information Processing Systems, 34:29348 29363, 2021.
[35] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 1 1):2278-2324, 1998.
[36] Adam Liska, Tomas Kocisky, Elena Gribovskaya, Tayfun Terzi, Eren Sezener, Devang Agrawal, D' Autume Cyprien De Masson, Tim Scholtes, Manzil Zaheer, Susannah Young, et al. Streamingqa: A benchmark for adaptation to new knowledge over time in question answering models. In International Conference on Machine Learning, pages 13604-13622. PMLR, 2022.
[37] Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and Qiang Liu. Bome! bilevel optimization made casy: A simple first-order approach. Advances in neural information processing systems, 35:17248 17262, 2022.
[38] Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806. 09055. 2018
[39] Risheng Liu, Xuan Liu. Xiaoming Yuan, Shangzhi Zeng. and Jin Zhang, buauto--nnni interior- point method for non-convex bi-level optimization. In International conference on machine learning, pages 6882-6892. PMLR, 2021.
[40] Risheng Liu, Pan Mu, Xiaoming Yuan, Shangzhi Zeng, and Jin Zhang. generic first-order algorithmic framework for bi-level programming beyond lower-level singleton. In International conference on machine learning, pages 6305-6315. PMLR, 2020.
[41] Sijia Liu, Pin- Chen, Bhavya Kailkhura, Ganyuan Zhang, Alfred O Hero II, and Pramod K Varshney. primer on zeroth-order optimization in signal processing and machine learning: Principals, recent advances, and applications. IEFE Signal Processing Magazine, 37-5),33-54, 2020.
[42] Sijia Liu, Bhavya Kailkhura, Pin- Yu Chen, Paishun Ting, Shiyu Chang, and Lisa Amini. Zoroth-order stochastic variance reduction for nonconvex optimization. Advances in Neural Information Processing Systems, 31, 2018
[43] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International conference on artificial intelligence and statistics, pages 1540- 1552. PMLR, 2020.
12

---

[44] Lu Lu, Raphael Pestourie, Wenjie Yao, Zhicheng Wang. Francesc Verdugo, and Steven G Johnson. Physics- informed neural networks with hard constraints for inverse design. SIAM Journal on Scientific Computing. 43(6):B1 (16)):B105--BD 100--B113, 2021.
[45] Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable gradient-based tuning of continuous regularization hyperparameters. In International conference machine learning, pages 2952- 2960. PMLR, 2016.
[46] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International conference on machine learning, pages 21 2112-2122. PMLR. 2015.
[47] Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine- tunint language models with just forward passes. Advances in Neural Information Processing Systems, 36:53038- 53075, 2023.
[48] John L Nazareth. Conjugate gradient method. Wiley Interdisciplinary Reviews: Computational Statistics. 133):388-355 2009.
[49] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-lcaming algorithms. arXiv preprint arXiv: pet.i 2018.
[50] Barak A. Pearlmutter. Fast exact multiplication by the hessian. Neural Comput., 6(1): 160, 1994.
[51] George F Pinder. Numerical methods for solving partial differential equations: a comprehensive introduc- tion for scientists and engineers. John Wiley & Sons, 2018
[52] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 778:686-707, 2019.
[53] Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. Advances in neural information processing systems, 32, 2019.
[54] Pranav Rajpurkar, Jian Zhang. Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprini arXiv:1606.0525. 2016.
[55] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International conference on machine learning. pages 8821 18311. Pmlr, 2021.
[56] Mengye Ren, Simon Korblithh, Renjie Liao, and Geoffrey Hinton. Scaling forward gradient with local losses. arXiv preprint arXiv:2210.03310. 2022.
[57] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611. 16l.01661..7776.. 2016.
[58] Levent Sagun, Utku Evci, V. Ugur Güney, Yann N. Dauphin, and Léon Bottou. Empirical analysis of the hessian of over- parametrized neural networks. In 6th International Conference on Learning Representations, ICLR 2018 Vancouver, BC, Canada, April 30- May 3, 2018, Workshop Track Proceedings, 2018.
[59] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert. a distilled version of bert: smaller. faster, cheaper and lighter. arXiv preprint arXiv: 1910001 10000000000000000000000000000000 2019.
[60] Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated back-propagation for bilevel optimization. In The International Conference on Artificial Intelligence and Statistics, pages 1723-1732. PMLR, 2019.
[61] Han Shen and Tianyi Chen. On penalty- based bilevel gradient descent method. In International Conference on Machine Learning, pages 30992-3 10015 PMLR, 2023.
[62] Qianli Shen, Wai Hoh Tang, Zhun Deng, Apostolos Psaros, and Kenji Kawaguchi. Picprop: Physics- informed confidence propagation for uncertainty quantification. Advances in Neural Information Processing Systems, 36, 2024.
[63] David Silver, Anirudh Goyal, Ivo Danihelka, Matteo Hessel, and Hado van Hasselt. Learning by directional gradient descent. In International Conference on Learning Representations, 2021.
[64] Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. Advances in Neural Information Processing Systems. 1301888188009, 2020.
[65] Ankur Sinha, Pekka Malo, and Kalyanmoy Deb. A review on bilevel optimization: From classical to cvolutionary approaches and applications. IEEE Transactions on Evolutionary Computation, 22(2):276- 295, 2017.
[66] Jihoon Tack, Jachyung Kim, Eric Mitchell, Jinwoo Shin, Yee whye Teh, and Jonathan Richard Schwarz. On- line adaptation of language models with memory of amortized contexts. arXiv preprint 2433w:2403.04317, 2024.
13

---

[67] Kai Wang, Bo Zhao, Xiangyu Peng. Zheng Zhu, Shuo Yang, Shuo Wang. Guan Huang, Hakan Bilen, Xinchao Wang, and Yang You. Cafe: Learning condense dataset by aligning features. Proceedings of the IEEE/CVP Conference on Computer Vision Puttern Recognition, pages 12196-12205, 2022.
[68] Tongzhou Wang, Jun Yan Zhu, Antonio Torralba, and Alexci A Efros. Dataset distillation. arXiv preprint arX1v.1811.1095,, 2018.
[69] R. E. Wengert. A simple automatic derivative evaluation program. Commun. ACM, 7(8):463-465 464, 1964.
[70] Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural Comput., 1(2):270-288) 1989.
[71] Peng Yang. Yingje Lao, and Ping Li. Robust watermarking for deep neural networks via bi-level optimization. In Proceedings the IEEE/CVF International Conference Computer Vision, pages 14841-14850, 2021.
[72] Ruonan Yu, Songhua Liu, and Xinchao Wang. Dataset distillation: A comprehensive review. IEEE
Transactions on Pattern Analysis and Machine Intelligence, 2023.
[73] Yihua Zhang. Prashant Khanduri, Joannis Tsaknakis, Yuguang Yao, Mingyi Hong, and Sijia Liu. An introduction to bi-level optimization: Foundations and applications in signal processing and machine learning. arXiv preprint S.fiv:2308.0078885 2023.
[74] Yihua Zhang, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin- Chen, Jason D Lee, Wotao Yin. Mingyi Hong. et al. Revisiting zeroth-order optimization for memory-efficient IIm fine-tuning: A benchmark. arXiv preprint arXiv:2402. 1192, 2024.
[75] Yihua Zhang, Yuguang Yao, Parikshit Ram, Pu Zhao, Tianlong Chen, Mingyi Hong, Yanzhi Wang, and Sijia Liu, Advancing model pruning via bi-level optimization. Advances in Neural Information Processing Systems, 35:18309-1888 18326, 2022.
[76] Yihua Zhang, Guanhua Zhang, Prashant Khanduri. Mingyi Hong, Shiyu Chang, and Sijia Liu. Revisiting and advancing fast adversarial training through the lens bi-level optimization. In International Conference on Machine Learning, pages 26693-26712. PMLR, 2022.
[77] Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen. Dataset condensation with gradient matching. arXiv preprint arXiv:2006. 05929, 2020.
[78] Barret Zoph and Quoc V Le.. Neural architecture search with reinforcement learning. arXiv preprint 1.X1v:6611.0177 2016.
[79] Simiao Zuo, Chen Liang, Haoming Jiang, Xiaodong Liu, Pengeheng Hc, Jianfeng Gao, Weizhu Chen, and Tuo Zhao. Adversarial regularization stackelberg game: An unrolled optimization approach. arXiv preprint 1rXiv:2104.04500 2021.
14

---

A Algorithm
Algorithm 1 (F:) Forward Gradient Unrolling with Forward Gradient
Require: Initial inner parameters 0.. initial meta parameter Φo. random direction distribution number of random directions b, total meta steps K, meta update mappings L I.Kk
1: 0 <- 0o. Φ Φn
2: for ...... K do
3: for
1,.... do
4: Sample v, p() and initialize y, (0,(0..) U; 86
5: end for
6: for t :=:...... I,.... T do
7:
0 < Q0, Φ). A +
(0.10. 80
B (0.(0..)
8:
for i
1..... do
9:
y; + Ay; Bv,
10:
end for
11:
end for
12: for 1,. ..... b do
W;=- (5(0.0)
80 y; + 01(0.0)
8Φ
14: end for
1 Σ wins) 15: & +- 4 (Ø)
&
end for
17: return
B Extended Discussion on Bi-level Optimization
B.1 Truncated Reverse Gradient Unrolling (TRGU)
To address the memory issue of GU methods, truncated Reverse Gradient Unrolling (TRGU) 60 proposed to reduce the memory usage by preserving only the last K steps the inner optimization trajectory. However, this introduces a significant bias large-sealc scenarios, particularly when the permissible
small.
Recall 5 and (6), where the conventional RGU method computes the hypergradient by fully unrolling the Tssppp inner optimization into a computational graph. Instead, TRGU performs s-step truncated sack-propagation and approximates the gradient with the intermediate term -T
T
CT
Σ t=T-s+1
B,A++1 Ardr-
(21)
According to Proposition 3.1 in 60). if the inner-level objective function is L-smooth, twice-differentiable and globally a-strongly convex, the gradient update rule writes 0:-1 VVVg Φ). then the bias s-step TRGU would be bounded by
I cT < (1-
na)" |dT]|
max .......T--.
(22)
7/a
The bound (22) demonstrates exponentially decaying gets smaller, which means that we truncate computational graph heavier in pursuit of lower memory cost, bias would grow exponentially. This would result in an inaccurate calculation of the hypergradient. Contrastively, our (UG)²U is an unbiased estimator of the hypergradient, while still keeping high memory efficiency with
the bias of s-step TRGU. However, when
small sample size of forward gradient as in 10
B.2 Implicit Function (IF)
Another idea for computing the implicit gradient is to utilize the implicit function theorem (IFT) 30) Suppose that the inner optimality Veg(Oz (Φ), Φ) s 0 is approximately achieved by sufficient inner optimization steps. If g is second-order differentiable, by applying the implicit function theorem and taking the first-order derivative of Φb
8 g( Φ) d0r(#) 8 g Φ)
0.
(23)
00%
dΦ
r0
15

---

Then, if the Hessian is further assumed to be invertible, the meta gradient can be approximated as
Vh(o) ~ 0f(Or(0), +) 80r
3 gg0r(()) -
8 ((r((() Φ)
0f(Or()).
(24)
00%
00,00
0Φ
The main challenge lies in computation the inverse Hessian matrix H which is intractable when 0 of high dimensionality. Fortunately, several iterative inverse Hessian vector product (1hvp) approximators requiring
only Hessian vector product (hvp) and O (M) space be employed produce dH for approximating
dH based Conjugate Gradient 48
53 Neumann Series 19 28
and low-rank approximation 644 24
Neumann Series. The inverse Hessian vector product approximated with truncated sum of Neumann series 19 28
K > d aH) dH k=0
HH
E diI d aH)" k=
(25)
where a is a hyperparameter to ensure the convergence, and the number truncated steps. Compared to other IF-besed methods, the Neumann Series has demonstrated good empirical performance and stability [9] and its stochastic variant has been well studied (28)
Weakness (IF); Approximation Errors. The errors of IF emanate from two distinct sources Firstly, IFT presupposes that the Karush-Kuhnn Kuhn-Tucker (KKT) conditions of the inner problem satisfied, leading an approximation error in (23) when iterative approximations of inner solutions are used. Secondly, singular nature of the Hessian within neural network training [57] leads to costly and unstable inverse Hessian
approximation in practical applications, with heavy reliance on engineering efforts
[ More formally, recall
Vi(Φ) (f( Φ) d0r(Φ)
@f(Or()).c
(26)
00r
dø
0Φ
The approximation error can be decomposed into
Vh(o) - v = d(
(d)
dH ")Y
(27)
Cin
To reduce the computational cost, Hessian-free approaches
propose approximating the Hessian as an identity matrix, incorporating additional assumptions about the inner model and objective.
76 75
dH
adI,
(28)
where a > 0 is hyperparameter to control the magnitude. However, these numerous assumptions often diverge from practical scenarios, resulting in significant approximation errors and consequently inducing suboptimal outcomes.
0.0
inner Lo55
L2 Grod.
605
1.5
Itoration
(a). Inner v.s. Iterations.
Gradient v.s. Iterations.
Figure B.1: CIFAR100, IPC=50: Inner Loss and gradient norm for Neumann
In Figure B.I it evident that the inner optimization has not converged by the unrolled step 100, as indicated by both inner loss and gradient norm. This observation implies that the Karush-Kulle satisfied, leading the conclusion that the approximation used in 23 introduces bias.
(KKT) conditions
16

---

B.3
Value Function (VF)
The VF-based methodology
39
37
61
33
considers equivalent reformulation of the original optimization
problem as outlined in 2
min f(0, ④) s.t. g(0,((( () ((0r(()) g(rr)()
(29)
0.4
This reformulation casts the standard hi-level optimization challenge into a constrained single-level optimization framework. VF-based methods circumvent the need for second-order computations and have demonstrated near-optimal complexity. comparable to second-order methodologies in deterministic settings, reported in [8]
Weakmess (VF): stochastic optimization. VF-based strategies have yet to gain widespread acceptance in practical ML applications. This limited adoption primarily altributed the challenges these methods face in addressing large-scale stochastic problems, where the complexity significantly impedes their performance 73
C Proofs of Theoretical Results
In this section, detail the proof's of Lemma 3.1 and Theorem 3.4 Our approach proving Theorem 3.4
follows similar high-level approach [19 28 with some important distinctions. Initially, Lemma B.2 we extend the smoothness properties of the objective functions f and g, and the transition functions $2, to the 7'th iteration lower-level parameter or. Following this, Lemma [33 establishes the smoothness of the meta-learning objective f(Or. +). incorporating results from the inner-loop computations. Building the demonstrated smoothness of the meta objective function and the variance of the forward gradient method (as shown Lemma 3.I. we then validate the convergence properties of Algorithm
The novelty in our proof Theorem 3.4 lies two primary aspects. Firstly, our analysis does not presume that the lower-level optimization yields an optimal solution 0" instead, more realistically assumes the use of or, which is derived from finite number of iterations. This assumption aligns more closely with the computational constraints encountered in real-word world scenarios. Secondly, convergence analysis explicitly accounts for the variance of unbiased gradient estimator, achieving convergence rate O O(c-tp-1) ). This demonstrates that utilizing the forward gradient method, while significantly reducing memory requirements, does not adversely affect the algorithm's 's convergence rate, underscoring the practical viability and efficiency of our approach even with memory constraints.
In Appendix C.3 discuss how to extend the convergence of optimization problem into with additional assumptions.
C.1 Proof of Lemma 3.1
For convenience, the lemma is restated as follows.
Lemma 3.1. For any & Φ the gradient estimation with forward gradient method:
Vi(()) Vh(c)o.e
where Unif([-1,1))) and b denotes the sample size, satifies
1 [h(T)]
B( vh( =
where p:= N-I E (0, 1] as b selected from 1,
Proof. We start by computing the variance of one-sample estimation, Vh(Φ) = Vh(o)vv" Since E(vv") ]
I, we know that [ = (b))). Consequently.
E(( - ÆVト()) 2 = E (ur I>
= B(Vh(() (wuT T 1) (vr" I))()(()))
(vr (ND TVh()) + (Vh(#) Tvh()) (30)
= [[(()(()) (ww") T) T or (7(())2(V(h())) Twor"
= BB(V)(()) ou 1|2 V [|
Since v an N N-dimensional Rademacher random variable, we have E| [/))e [[V(()]] and
Ell 1|2 N. Then,
30
1)]
17

---

For the sampee estimation Vi(() = - Vh(o)ore; Sincev; are i..d. sampled, and E[Vh(6)]]
Vi(o), we have
E [))
BV = E Vi( b (w: T -D -D
E
(h(o))o;; 1)1
[
b
C.2 Proof of Theorem 3.4
To prove our main result (Theorem 3.4 we first establish useful smoothness properties of the hyperparameter
learned from solving the lover-level optimization problem. Subsequently, establish the smoothness of the meta objective function examined the approximated lower-level parameters in Lemma B.3
Regarding the lower-level parameter (()). we present the following lemma, which is based Assumption 3.3 and establishes that (()) inherits similar Lipschitz continuity and smoothness properties as Z.
Lemma B.2. Under Assumptions 3.2 and 3.3 (r()) is Cz-Iipchirz and I.z -smoohh, i.e., for any Φ &' E +
and
[10r()) 0r(4')ll S Cz Φ
d'1.
(v)
v0r(6') S Lz|l<b l'1l.
where Cz
G. c--2 and Lz = Ln C +
Cp-1
cx+2T
Cn-
(Cn-1)²
C -1
Prouf. We start with the proof of Lipshitz continuity. For any pair of Φb o' t. using and Assumption 3.3 we have
(+) 0.(6')l |[ 020 (Φb) か')l
(31)
< CaIl0..:(0)) - 0. (x)) || Cn c l'|l.
Applying recursively over = 1..... gives
[[)( - (.(())) < ( - 00(20)))) |1 Σ chll - ø'1l
(32)
++
Ch+²2
Ca
Σ
cal 001
|
l'll
Cn
where the last inequality holds from the fact that 0, = Zo well as Assumption equality follows from the geometric series summation formula.
3.3 and the subsequent
cf
Substituting t T, we get Cz = Cz(T)) = Therefore, (e(Φ) is Cz (itphiit,,,,,,, where Cz = Cp-1
cT+2 -a- We now proceed with the proof of Lz smoothness. For simplicity of notation, we follow and denote
Zz(Φ)= (0:(0) A: (0-() 6) B((Φ)= (0
00,-1
Subsequently, considering the update rule Z:(Ø) = (.()) + B(( of Forward Gradient Unrolling
we have
(v) (0.(d')l
[Ze(Ø) - Z((o')l
[A:()) + (.( -
(B') II
(A(())Z)-1()) ((0))) B,('))
B(()')
(A(())._.()) ((0)) +
B.(() B.(d')ll
< (4)))))._1()) ( + (
+IB.(0)- B.(d')l
< [|A((0)] (z) Z (0') (A.( A. I| 12. (b - Bi(
< (2)) - Z.-1(6))l + (Cz(() + 1)La 0'11
(33)
18

---

where the last inequality follows from Assumption 3.3 that (a(()) and l1rr(bb) are Cn-Lipshiz and Ln- smooth, and the previously proved result that (x(() (6) Cz C2((2) (1)-Licchit.
Noting that (o( (0.(Ø) (lo(d)) applying
33 recursively over I,.... gives
(0r(6)- v0r( - Zo
c (Cz(1)+ + 1)Lanl I'll
Cz(() +1
= -] - vn
+ C
Ln LaIg 2'1
C
2++
I L 0
4'1 + C
n
La 1Φ 0'1
(Ca - 1)
T
Cp
C
T
= La
Cn
C
| Φ11
Cn
Cn 1=1
C
[=1
d (1-
C
Ln
n
|6 w''l
Cn 1
-n-1
1 c
2T
ci
Lo
C
| 2
Co
(Cn 1)
C -
where the third line follows from Assumption 3.3 that (o) is Ln smooth and the choice Cz(()
On
C
(which gives (t) +
C ) and the fifth line again uses the geometric series summation formula.
Hence, (x)) is Lz-smooth with Lz La C
C
1y²
Next, we provide lemma establishing that the upper-level objective f. evaluated the learned parameter (0r(Φ), 6). also adheres to certain smoothness properties.
Lemma B.3. Define (((()) = f(0 (Φ), ②). Under Assumptions 3.2 and
3.3
h(() is La-smooth, ie., for any
Φ 6' E Q.
([) - Vh((')Il S 1 &'1l where Ln = (C2 + CLz, with Cz and Lz defined in Lemma B.2
Proof. For simplicity of notation, we follow (5) and denote
Zi(0)= = (0) cr( (f)
(f(Or(0). Φ)
dr
dΦ
00,
For any Φ o' E 4, following similar proof as Lemma B.2 we have
[Vh(d) Vh(d')l
< (011)))22-000--((000000000(()))1
(34)
[(d(())ZT()) (111::/0000011011110000000000011
cr(3))l
(er()) 2. |rt -o(())
$
(0) Zr(
Subsequently, we deduce that
S Cz)(dr( dr( ||| CLZ - w'1l + ||T - cT''ll
(1 (f(Or( x) 0f(Or( x)
(f(0r( Φ)
0f (f(Or Φ')
S2²
00, (10F(0r(0))) Φ) (f(Or(
00,
80,
80,
0f(Or
(0 ) Φ6)
d0
dΦ
dΦ
OΦb
+CZZIl - l'1l
$ (22 (0'))l + C2L +'11+110 ( + L Φ'11
+CL2110- -6'll
3C2CC21l0 - 0'1 + C2 - 601 + LC21 w11 + 0 - 2'1 + CL2ll0 - 6'1
= =[C2+1) + CLz]Il@ 0'1b
19

---

where the first, third, and fourth lines all follow directly from Lemma B.2 and Assumption 3.3 (recall that the latter states that is L-Lipschitz and C-smooth).
Therefore, l(() is Lm-smooth with La. (Cz + 1) CLz.
Now based on the aforementioned lemmas. we put forward the proof our main theorem: the convergence analysis for our bilevel optimization method (FG)
Theorem 3.4 (Convergence). Suppose that Asumption 3.2 and Assumption 3.3 hold. Setting the learning rate
(o+1) for gradient descent over the hyperparamette ④. have
4L (E)((,)) min h(0))
E [IV
(35)
K
p
Proof. We have
1((b41)) 1((b)))
$ (Vh(#t) PA+1 - (b)) + La 11 - 2xl2
LL 1V
=-B(VH(OL) Vl(d.)) (.VK(()).V((d))))
(36)
= -B(Vh((O), Vh(ow) + LL [hh(d)) + V)(o)) vh(oa)11
2
LA [V(( + (0 LA 3)()(( VM(D)) 8L (Vh(d,))
-
2
where the second line is well-known inequality for smooth functions with the Sm-ooothsness itself following from Lemma B.3 and the third line uses the gradient descent rule Φ++1 Øx vV(()))
By Lemma 3.1 and the fact that V is unbiased (see the proof of Lemma 3.1
we know that
(((V)( (I(A MMA. = (h(Ø,)) )
(((bb( V(((,))" (1bb _ [((()))))
Therefore, taking the conditional expectation E[- - Φz over 36 gives
8 Ln
[ ((()))
B
[h((a)))
(37)
Furthermore, taking the full expectation and telescoping
37 over & form 0 to K - 1 yields
K-
(0 + Ln
g2
[[ E((( ) - [(
K
2p
k=
(38)
E(((@)) m min h()))
Choosing 9
we have
Tp+ 13T."
2(p 1)La (B[h(db)) ming 1(()))
Σ [1
<
K
pK
(39)
4hw (B[)())) min& h()))
pK Hence, Algorithm requires O 11) steps attain an e-accurate stationary point.
C.3 Extended Discussions
Convergence of Problem (T). To extend the convergence of optimization problem into T we need to assume that the lever-level objective function is strongly convex w.r.t. 0 as commonly done by previous works (60 28). From the strong convexity and first-order smoothness (Assumption 3.2) of g. we have 1) the
zeroth and first-order smoothness of 0' (+): 2) [[0((6') 0" ( 0 T + oo. Then the inequality
|0r (Φ) (r''lll < [0r(Φ) 0* (+) |r((') ( + (Φ) (
(40)
20

---

implies the the zeroth and first-order smoothness of (r()). Following the same line of proof as presented in our paper, derive the smoothness of (()) Φ) and (Ø),, +), and subsequently the convergence of either
problem or
However, as discussed in Section 2 given that the scope this paper 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 60] 28 are implicitly or explicitly solving (2) instead of (I. Additionally, it is important to acknowledge that achieving strong convexity is often unteasible in practical applications. Consequently, we focus on aiming to present convergence theory that proves the effectiveness of our method.
practical
D Zeroth-Order Derivative Estimator
In this section, we give a more detailed introduction to zeroth-order (ZO) derivative estimators. These estimators pivotal in scenarios where the computation of exact derivatives is either infeasible due to memory constraints computationally prohibitive. Apart from the forward gradient method employed in (FG)²U, randomized smoothing (RS) is another widely-used derivative estimator, both Reinforcement Learning Language Models
and Large
18
For a function F
I
R. gradient estimation via RS can he mathematically formulated as:
F'x + cc) 0
Fix + cea) F(x) 0,
F(w)
Va F 2 E. A((.I))))))
(41)
6
6
where b is the number of random samples, is the smoothing parameter, V are samples drawn from a standard Gaussian distribution. Regarding the accuracy of estimation, it has been shown in [13 42 that the variance of
RS is roughly in the order of (NN/b) which is the same FG proved in Lemma 3.1
RS stands out particularly in its ability to estimate gradients of functions evaluated through black-box systems, where internal operations are inaccessible or highly complex. This characteristic makes RS exceptionally valuable in practical applications such as adversarial rohustness and black-box optimization, where obtaining direct gradients might not be possible. Another advantage of RS is its robustness against noise and discontinuities in the function landscape. Unlike deterministic methods. the stochastic nature of RS allows it to approximate the gradient over smoothed version of the function, providing stability scenarios where slight perturbations lead to substantial changes in the output.
While RS provides robust gradient estimates across various scenarios, it is critical recognize that RS inherently introduces bias if the expectation not computed during inference. In many CV and NLP applications, the computational expense of Monte Carlo sampling at the evaluation stage is prohibitive, leading to a biased estimation when using RS. However, in the context of inverse PDE problems, where the inner-loop solvers are non-differentiable numerical solvers, employ RS zeroth-order derivative estimator.
E
Detailed Task Description
E.1
Data Condensation
In the era of rapid advancement in machine learning, multitude of foundation models [1 6 SS has benefited from training large-scale datasets, exhibiting formidable performance that models trained small-scale data cannot match. However, exponential growth of data also presents challenges: (1) Models updated with only new data are prone to catastrophic forgetting 200 while retaining all historical data for subsequent training imposes significant storage and computational burdens. (2) Applications within the realm of meta- learning, such as hyperparameter tuning [46 4 and neural architecture search [7 38), necessitate multiple training iterations
over datasets. The computational cost of these operations scales dramatically with the size of the datasets, posing a bottleneck for efficiency and scalability. (3) The widespread dissemination and utilization of datasets have raised significant concerns regarding privacy and copyright [2).
To overcome the challenges posed by large-scale datasets, line of work known as data condensation (68 72 has been proposed. with the idea to generate compact, synthesized dataset, designed to elicit similar behaviors in machine learning models those trained with the original, massive dataset. The objectives of the mainstream principles 772 designed for data condensation can naturally formulated a bi-level optimization problem. We focus on the best-known principle performance matching 722 on classification task, which can be formulated as,
min L(Or; D.), where 0, 0 100-1-1200001100 (VC00. 1 I....
(42)
We conduct experiments condense the following image datasets:
21

---

handwritten digits dataset containing 60, 000 training images and 10, 000 testing images with the size of 28 x 28 from 10 categories.
MNIST [5];
colored natural images datasets contraining 50, 000 training images and 10, 000 testing images from 10/100 categories, respectively.
CIFAR 10/100 [1:
The scale of the condensed dataset will fundamentally impact the results. Therefore, we consider different scales for each dataset, with images per class set to I, 10, and 50. The condensed dataset will be used to train random initialized models, and evaluated on test dataset.
E.2
Meta Learning Online Adaptation of Language Models
The online adaptation of language models (LM) 34 27 has been studied recently to keep the knowledge of LM updated to date. However, trivial auto-regressive fine-tuning the LM with uniform weights for all tokens results in poor performance
downstream tasks, as the default average negative log- likelihood (NLL) loss does not accurately reflect the importance of tokens (25). To address the issue, [25] proposed Context-awarc Meta-learned Loss Scaling (CaMeLS) to meta learning the weights of tokens for effective online adaption. More formally, let 0 denote the parameter of the hase model for adaptation, denote the parameter of parametric weight model to assign weights for each token, the meta-learning online adaption of I.M can be formulated as the following bi-level optimization,
min Limeta (rr (Φ), +) s.t. (.( 0--1( V Vooorrin (0:-1, Wo) t = 1,....,, T.
(43)
We follow the setting studied by 25 where the downstream task is question-answering. The meta-objective
consists of a anesttoon-nsweere term measuring the performance gained from adaptation, and a locality term that prevents the updated base model parameters from excessively changing the base model's behavior. Let Dq denotes question- caation-ansween answering dataset. Pioe denotes the locality dataset, and E 12t denotes the weight of the locality term, then the meta objective is formally defined as
Lmett (Ø)) x) := Ra - Pog (alg) + cRauDDue Σ kI, (+Ia:) || po. (-Ixs))
(44)
The inner objective defined as weighted NLL loss, where the weights are determined by the weight model Wib.
Cirnin (0, ug) == Dewaen Σ
-w (1; æ²) log pe (r.(1r.)
(45)
The trained weight model is then fine-tuned unseen online documents and evaluated on corresponding question- answering tasks.
E.3
Data-driven Discovery of Partial Differential Equations (PDEs)
Burgers
AllenCahnn
KdV
0.4
3A
0.2
02-
03-
Figure E.1: Visualization of the 2D latent solutions for the Burgers, Allen-Cahn, and KdV equations. The observed data are sampled on an 8 x 8 grid, denoted by white points.
We conducted experiments on three non-linear PDEs, with the latent solutions visualized Figure E.1 The PDE structures (46) (47) (48) are assumed to be known while the PDE parameters v in (46) (47) (48 are assumed to be unknown. For each equation, 64 observed data points are sampled on an 8 x 8 grid. The objective is to predict the unknown PDE parameters using the observed data. The predicted PDE parameters evaluated by comparing the error with the ground truth, well the error in the corresponding prediction of the latent solution.
E.3.1 Burgers Equation
The nonlinear viscous Burgers equation is pivotal partial differential equation arising in diverse domains of applied mathematics such as fluid mechanies, nonlinear acoustics, and traffic flow. This equation can be deduced
22

---

from the Vavier-Stokes equations for the velocity field by omitting the pressure gradient term. In our experiment, the equation along with Dirichlet boundary conditions, expressed follows:
We +
= 0, = [-,,1], [0, 11. 0,
u x) sin(TT),
(46)
utt, -1] u(t, - O,
with actual viscosity v p.01 = 0.0031831. For PINN, following the output by choosing surrogate model of the solution as
we enforced the initial condition into
i(x)= = (1 - e1-eep(----epppp (N; -sin (TT),
where (a; 0) is neural network.
E.3.2 Allen-Cahn Equation
The Calen-Cahn equation is a reaction-diffusion equation of mathematical physics describing the process of phase separation in multi-componen alloy systems, including aider-dioorde transitions. In our experiment, it is expressed as follows:
d Vlls 5(u - v²), 11,, te e [0,1],0>00 0,
w00, x) x² cos(TTx).
(47)
utt, -1) = u(t, 1) -1,,
with the actual diffusion coefficient v = 0.001. For PINN, we enforced the initial condition into the output by choosing surrogate model of the solution as
i(x) (1 exp(-t)) 0) + + cos(TT))
where NN(x; 0) is neural network.
E.3.3 Korteweg-De Vries (KdV) Equation
The Korteweg-de Vries (KdV) equation serves a mathematical model for waves shallow water surfaces. This equation is distinguished prototypical example of an integrable PDE. is characterized by features typical of integrable systems, including a plethora explicit solutions, notably soliton solutions, and an infinite number of conserved quantities. These properties are particularly noteworthy given the inherent nonlinearity of the equation, which generally complicates the solvability of PDEs. In specific, we consider:
Ut + uu, + Pugas 0,
x E (-1,,,), [0, 1), * 0,
(48)
(<00 .r) cos(
with the actual coefficient of dispersion equal to 0.0025. For PINN, we enforced the initial condition into the output by choosing surrogate model of the solution
à (x)) (1 (xp((()) 0) cos(TTe)
where NN(x:; 0) is neural network.
E.3.4 Numerical PDE solver
Numerical solvers play critical role in the study and application of PDEs, enabling the simulation and analysis of complex physical phenomena that cannot be addressed analytically [5]]. These solvers convert PDEs into form that can be handled computationally, typically by discretizing the domain into finite of points or elements and approximating the derivatives. Conventional numerical methods include finite difference methods, finite element methods, and spectral methods M1
Among the various numerical methods for solving PDEs, spectral methods stand out for their ability to deliver highly accurate solutions, particularly for problems with smooth solutions 21 Spectral methods involve representing the solution globally defined over the domain. This approach contrasts finite difference finite element methods, where the solution is localized to the grid points or elements. In this paper, we mainly adopt spectral methods, as we focus on the Burgers, Allen-Cale Cahn, and KdV equations. All these three equations can be efficiently and accurately resolved by spectral techniques.
PDE
of basis functions, such as trigonometric polynomials, which
23

---

Implementation Details
F.1 Data Condensation
We conducted experiments following the standard data condensation setting established by 68 77
67
condensation and evaluation are both performed on depth-3 convolutional neural network [8]. The hyperparameters we used for (FG)"U are summarized in Appendix F.1
All experiments are conducted on
NVIDIA-L40S (40G).
Datasets IPC
MNIST
CIFARIO
CIFAR100
10
50
10
50
10
50
Unrolled Depth
100
100
100
100
100
100
100
100
200
Random Directions
32
32
32
32
32
32
32
# Inner Batch Size
Full
Full
Full
Full
Full
Full
Full
Full
100
Hessian-Free Pretraining
Gradient Accumulate
16
32
32
32
32
32
64
64
64
Outer Steps
10000
10000
10000
10000
10000
10000
10000
10000
10000
Outer Step Size
le-2
5e-4
5e-4
le-2
5e-4
5e-4
1e-2
5e-4
5e-4
Evaluation Steps
1000
10000
10000
1000
10000
10000
1000
10000
10000
Table F.1: (FG)²U hyperparameters for data condensation experiments.
F.2
Meta Learning Online Adaptation of Language Models
We adhered the standard settings of CaMeLS 25 and adapted their official code for our implementation. The only modification made replacing the meta gradient approximation module with (FG) U. It is important note base models used for meta-learning were initially pre-trained on a split QA-paired set. While the official codebase provided the script for pretraining, it did not include the exact base model (weights) they used. We executed the official script to generate the pre-traine pre-trained base models and observed that meta- meta-learning performance is sensitive the choice of base models. For fair comparison, [66] (where CaMeLS 25 presented performance improvements over baselines using bar plots without specific metric values) and results with our best custom pre-trained base models. Following the two-phase training paradigm introduced in Section 3.2 we performed training of (FG) on RGU (DistilGPT2, unrolled depth 6) results. The hyperparameters used for (FG)U are summarized in Appendix F22 while all remaining hyperparameters were kept the same in 25 (80G).
reported both the results from
All experiments conducted
NVIDIA A100 GPU
base model
DistilGPT2
GPT2
Unrolled Depth
24/48
24/48
# of Random Directions
12
8
Gradient Accumulate
32
32
Outer Optimizer
Adam
Adam
Outer Step Size
2.5E-6
2.5E-6
Table F.2: (FG)²U hyperparameters for CaMeLS experiments.
F33 Data-driven Discovery of Partial Differential Equations (PDEs)
The hyperparameters used for this experiment are summarized in Appendix FK
All experiments are conducted NVIDIA-LAOS (40G). The structure for PINN is depth-9 and width-20 MLP with tanh activations.
G Additional Experimental Results
G.1 Meta Learning Online Adaptation of Language Models
Ablation results on the unrolled depth and the base model are summarized in Table G.T and Table G.2
24

---

PDEs
Burgers
Allen AllenCahn
KdV
Inner Solvers
PINN
Numerical
PINN
Numerical
PINN
Numerical
Directional Grad. Calculation
FAD
Zo
FAD
ZO
FAD
ZO
Random Directions
Outer Steps
5000
5000
5000
5000
5000
5000
Unrolling Depth
1000
1000
1000
Grid Size for Numerical Method
256x512
256x512
256x512
Range of initial
(0, lel]
(0, lell
(0, le-1]
(0, le-1]
le-21
(0, 1e-2]
Outer Optimizer
Adam
Adam
Adam
Adam
Adam
Adam
Outer Slep Size
le-2
le-2
lc-2
le-2
le-3
lc-3
Inner Oplimizer
SGD
SGD
SGD
Inner Batch Size
5000
5000
5000
Inner Step Size
le-3
le-3
le-3
μ for Finite Difference
le.4
le44
le-4
Table F3: ( hyperparameters for discovery of PDEs experiments. v denotes the unknown PDE parameters.
Unrolled Steps (#)
DistilGPT2 EM (t)
GPT2
Model (# params)
Method
(1)
EM (f)
FI (t)
RGU (impl.)
2.04
5.53
OOM
DistilGPT2 (82M)
2.10
5.59
2.22
6.37
(FG) (ours)
48
2.10
6.25
2.16
6.32
RGU (impl.)
7.02
12.19
OOM
GPT2-Large (774M)
(FG)²U (ours)
24
6.91
12.12
7.21
12.50
48
7.03
12.31
7.27
12.45
RGU (impl.)
7.93
12.94
OOM
GPT2-XL (I.5B)
(FG)²U (ours)
8.34
13.46
8.89
14.42
48
8.23
13.70
8.65
13.91
Table G.1: StreamingQA: Ablation results on the unrolled depth and the base model.
Unrolled Steps (#)
DistilGPT2
GPT2
Model (# params)
Method
EM (t)
F1 (t)
EM (1)
FI (1)
RGU (impl.)
1.52
3.16
OOM
DistilGPT2 (82M)
24
1.72
3.49
1.72
3,50
U (ours)
48
1.75
3.47
1.73
3.49
RGU (impl.)
4.86
8.57
0OM
GPT2-Large (774M)
24
5.49
8.88
5.56
8.99
(FG) (ours)
5.45
8.90
5.32
8.97
RGU (impl.)
6.71
9.65
OOM
GPT2-XL (1.5B)
7.00
10.13
7.27
10.33
(FG) (ours)
48
7.37
10.37
7.25
10.32
Table G.2: SQuAD: Ablation results on the unrolled depth and the base model.
25

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant
information is Not Available.
Please provide a short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked "[Yes] is generally preferable to "[No] ". is perfectly acceptable to answer "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or used"). In general, answering "[No] oo "[NA] is not
the checklist as one of the factors in their evaluation. While
"we were unable to find the license for the dataset grounds for rejection. While the questions are phrased in binary way. acknowledge that the true answer often more nuanced, so please just use your best judgment and write a justification elaborate. evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
supporting
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can he expected to generalize other setrings.
It is fine to include aspirational goals motivation long it is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of the work in Section i
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged create separate "Limitations" section in their paper.
26

---

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model sell-sppeiiicattonn. asympiotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tes -oxttttt system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If l'applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve integrity the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The assumptions included in Section 3.1 and the complete proof is placed in Ap- pendix [
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should he numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: To the best of our knowledge, all the information needed to reproduce the main experi- mental results included in Section 4 Appendix B and Appendix F
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of model checkpoint, or other means that appropriate to the research performed.
27

---

While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(h) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. large language model), then there should either be way
this model for reproducing the results way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct
dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors welcome to deseribe the particular way they provide for reproducibility. In the of cosed-source models, it may he that access to the model is limited some way (e.g., to registered users), hut it should be possible for other rescarchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does paper provide open access to the data code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The data used in this paper is public. We will release the code the acceptance of the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines htttpe::/nipa.cc//eeee
Muides/CodeSubenn soioooliir for more details.
While we encourage release of code and data, we understand that this might not possible,
so "No" an acceptable answer. Papers cannot rejected simply not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines /ottps://nips... cc/public/ for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for new proposed method and baselines. If only a subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should relcase anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Deeting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: To hest of our knowledge, necessary information been included Sec- tion 4 Appendix E and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that necessary appreciate the results and make sense them.
The full details can be provided either with the code, in appendix, or as supplemental material.
Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined other appropriate informa- tion about the statistical significance of the experiments?
Question:
Answer: [Yes]
28

---

Justification: We have reported the error bars in Data Condensation experiments and Duta-driven Discovery PDE experiments. For Meta Learning Online Adoptation LM experiments, adhere to the standard evaluation protocol employed [25 66), specifically, using the identical evaluation script with the same random seed, to ensure a fair comparison, hence, without error bars. Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, statistical significance tests, least for the experiments that support the main claims the paper.
factors variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given
experimental conditions).
The method for calculating the enror bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It OK report I-sigma error hars, but should state it. The authors should preferably report a 2-sigma error than state they not verified.
96% CL, if the hypothesis of Normality of errors
For asymmetric distributions, the authors should be careful not to show in tables or figures
symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: each experiment, does paper provide sufficient information on computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification:
We provide the information of computer resources in Appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httts:://eerrpp::: wttpeiipppccpppuaaee http://///pcccccccco.oo..eeeeee.
Answer:
[Yes]
The rescarch conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from Code of Ethics.
authors should make sure preserve anonymity (e.g.. if there special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts the work performed?
Answer: [NA]
29

---

The research presented in this paper focuses on fundamental algorithms rather than specific applications. Consequently, is challenging predict the potential social impacts of this work.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance). fairness considerations (e.g.- deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will he foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it not needed point out that generic algorithm for optimizing neural networks
could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not have plans to release any new data or models in conjunction with this work.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recngnize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All usages of the related assets are accompanied by formal citations and comply with the respective licensing terms and conditions.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papprrswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide help determine the license of a dataset.
30

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's
13. New Assets
Question: Are introduced in the paper well documented and is the documentation provided alongside assets?
Answer: [NA]
Justification: not have plans to release any new assets models conjunction with this work. The code, which will be released upon acceptance, will be well documented, and the documentation will also be made publicly available.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cole/mmoel as part of their sub- missions via structured templates. This includes details about training, license, limitations,
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least minimum wage in country data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurlPS Code of Ethics the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
31


Reviews:
[{"review_id": "r2qHigPxX0", "reviewer": "Reviewer_cwe5", "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.", "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}$?", "limitations": "The paper appears to have no potential negative societal impact. The authors discussed the limitations in the Appendix H.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed method (FG)^2U is simple and easy to implement. \n- The writing is clear and easy to follow.\n- The theoretical results clearly demonstrate the relationship between convergence, sampling batch size, and parameter space size.\n- 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\n- With a fixed batch size, convergence on large-scale applications would be slow (vanilla (FG)^2U without additional techniques).\n- With an $O(N)$ batch size, convergence is satisfactory, but the memory cost grows as $O(MN)$, similar to FGU.\n- 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$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the overall positive evaluation and will try to address the concerns raised point by point.\n\n**Regarding the complexity**\n\nThe reviewer's understanding of Theorem 3.4 is generally correct. We would like to highlight the following points:\n\nFirstly, 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.\n\nSecondly, 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.\n\nThirdly, 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$. \nAs 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])$. \n\nThe differences are:\n\n(1) Plan A maintains the rows of $Z$, while Plan B maintains the columns.\n\n(2) Plan B utilizes randomization, while Plan A does not.\n\nWe argue that both choices of Plan B are better:\n\n(1) Considering the update rule in Equation 4, row-wise updates require communication among threads, whereas column-wise updates do not.\n\n(2) Plan A requires exactly $M$ threads, whereas Plan B provides smooth trade-offs via randomization.\n\n**Regarding the formula for data condensation**\n\nYes, the reviewer is correct. We will fix the typo in the revised manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3xT8OP1jev", "reviewer": "Reviewer_4Mx8", "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", "questions": "see weakness", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The method significantly reduces memory overhead compared to traditional gradient unrolling methods, making it suitable for large-scale applications.\n2. 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?\n2. The convergences analysis is for the problem (2) not the original problem (1). \n3. More large scale datasets are needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. We will try to address the concerns point by point.\n\n**Regarding the computational complexity**\n\nWe 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.\n\n**Regarding the convergence of problem (1)**\n\nTo 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$. \nThen the inequality \n$$\\|\\theta_T(\\phi)-\\theta_T(\\phi')\\|\\leq \\|\\theta_T(\\phi)-\\theta^*(\\phi)\\|+\\|\\theta_T(\\phi')-\\theta^*(\\phi')\\|+\\|\\theta^*(\\phi)-\\theta^*(\\phi')\\|$$\nimplies 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).\n\nHowever, 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. \nTherefore, the optimal parameters are not explicitly accessible and can only be estimated through iterative procedures. \nMost related works [1][2][3] are implicitly or explicitly solving (2) instead of (1).\n\nAdditionally, it is important to acknowledge that achieving strong convexity is often unfeasible in practical applications. \nConsequently, in this paper, we aim to present a more practical convergence theory that proves the effectiveness of our method.\n\n**Regarding larger-scale datasets**\n\nAs discussed in Appendix H, this work serves as the first attempt to apply FG/ZO in BO. \nWhile 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. \nWe believe the empirical studies conducted in this paper are sufficient to demonstrate the potential of $\\text{(FG)}^2$U in large-scale BO. \nWe anticipate that the effectiveness of $\\text{(FG)}^2$U will be further validated in larger scales.\n\n[1] Making Scalable Meta Learning Practical, https://arxiv.org/abs/2310.05674\n\n[2] Truncated Back-propagation for Bilevel Optimization, https://arxiv.org/abs/1810.10667\n\n[3] Bilevel Optimization: Convergence Analysis and Enhanced Design, https://arxiv.org/abs/2010.07962", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fta6PnKuaH", "reviewer": "Reviewer_J9FE", "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.", "questions": "See weakness part.", "limitations": "See weakness part.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1.\tThe presentation is great and I feel relaxed to read this paper.\n2.\tThe experiment is extensive including both small and large scale settings.", "weaknesses": "1.\tThe 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. \n2.\tBased 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?\n3.\tIn 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?\n4.\tThe 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? \n5.\tCan authors provide real space consumption in the comparison (Table 1, Table 2)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the overall positive evaluation and will try to address the concerns raised point by point.\n\n**Regarding the complexity**\n\nWe 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.\n\n**Regarding the strategy to choose $b$**\n\nWe follow the methodology [2][3] developed for FG/ZO in large-scale applications to choose $b=\\mathcal{O}(1)$ rather than $\\mathcal{O}(N)$. \nWe 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.\nNotice that we additionally perform gradient accumulation through iterative/multi-GPU parallel computation, which will also influence the variance.\nThe number of accumulation steps is tuned in our initial attempts according to the wall-clock efficiency and stability.\n\n**Regarding the IF methods**\n\nAlthough 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.\n\nIn 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. \n\nWe considered two IF methods: Hessian-free and Neumann.\nThe 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)$).\n\nDue 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. \n\nFor 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.\nFor 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.\n\nPlease refer to Table A in the attached PDF in the general response for results. We observe the following:\n\n(1) Both IF methods yield improvements on RGU (with an unrolled depth of 6).\n\n(2) By carefully tuning the hyperparameters, the Neumann method achieves further improvements over the Hessian-Free method, albeit with additional computational cost.\n\n(3) Despite the careful hyperparameter tuning for Neumann, its performance remains inferior to $\\text{(FG)}^2$U.\n\n\n**Regarding the stochastic $\\text{(FG)}^2$U**\n\nIn the stochastic context, the bilevel optimization problem is formulated as:\n\n$$\\min_{\\phi\\in \\Phi}\\  h(\\phi):=f(\\theta_T(\\phi), \\phi)=E_{\\xi}[F(\\theta_T(\\phi), \\phi; \\xi)]$$\n\n$$where\\ \\ \\ \\theta_0(\\phi) = \\Omega_0(\\phi), \\theta_t(\\phi) = \\Omega_t(\\theta_{t-1}(\\phi), \\phi; \\zeta)\\in\\Theta, t = 1,...,T, $$\n\nwhere $\\xi$, $\\zeta$ are random variables.\n\nAssume 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:\n\n $$\\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}}.$$\nWith 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\n$$E[\\nabla_\\phi ' h(\\phi)] = E_{\\xi,\\zeta}[E_{v}[\\nabla_\\phi ' h(\\phi)vv^{\\top}]]= E_{\\xi,\\zeta}[\\nabla_\\phi ' h(\\phi)]$$\n\n$$=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]$$\n\n$$= \\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),$$\nwhich gives the unbiasedness of $\\text{(FG)}^2$U for stochastic bilevel optimization problems.\n\n**Regarding the real space consumption**\n\nWe direct the reviewer to Tables B, C, and D in the attached PDF in the general response.\n\n[1] Making Scalable Meta Learning Practical, https://arxiv.org/abs/2310.05674\n\n[2] Fine-Tuning Language Models with Just Forward Passes, https://arxiv.org/abs/2305.17333\n\n[3] Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark, https://arxiv.org/abs/2402.11592", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## (A) General Response

We thank all reviewers for their constructive feedback. 
In this response, we will address some common concerns raised by the reviewers.
We will revise the manuscript covering the following discussions.

All reviewers raised concerns about the computation. 
More specifically, Reviewer 4Mx8 specifically requested an analysis of the complexity. Reviewer cwe5 and Reviewer J9FE raised concerns about the dimension-dependent convergence rate of forward gradient (FG). 
We will begin with a theoretical computation cost analysis of GU-based methods, then move to A.2 for discussions on practical computation.

### (A.1) Theoretical Computational Cost

We follow existing works [4][5] in treating the transitions $\Omega_t: \theta_{t-1} \rightarrow \theta_{t}$ as indivisible computational units and conduct the computational cost analysis by focusing on Jacobian (jac), Jacobian-vector product (jvp), Hessian-vector product (hvp), vector-Jacobian product (vjp), vector-Hessian product (vhp), and Hessian-Jacobian product (hjp) operations around $\Omega$.

The gradient computation costs are as follows:

(i) FGU involves $T$ hjp and $T$ jac operations (Equation 4).

(ii) RGU involves $T$ vjp and $T$ vhp operations (Equation 6).

(iii) $\text{(FG)}^2$U involves $bT$ hvp and $bT$ jvp operations (Equation 9).

Overall, the computational complexities of FGU, RGU, and $\text{(FG)}^2$U are $\mathcal{O}(MNT)$, $\mathcal{O}(MT)$, and $\mathcal{O}(bMT)$, respectively.

Here are some discussions:

(a) $b=\mathcal{O}(N)$ or $\frac{N}{b}$ times updates are required to achieve convergence for $\text{(FG)}^2$U. The total computation will be $\mathcal{O}(MNT)$.

(b) Notice (a) is about the theoretical upper bound, we will further discuss the computation cost in practice in (A.2).

(c) RGU is the most computationally efficient. However, as discussed in lines 105-111, memory issues impede RGU in large-scale scenarios. $\text{(FG)}^2$U improves memory performance at the cost of increased computation.

$ $
### (A.2) Practical Computational Cost

**Firstly, the dimension-dependent convergence rate is an upper bound and the scalability has proven acceptable in practice**.
As evidence, zeroth-order optimization (ZO), which can be regarded as a finite difference approximation of FG, has been used as a standard technique for fine-tuning large language models (LLMs) [1][2]. 
The size of the LLMs studied is up to 66B, and the fine-tuning performance is competitive with full-parameter fine-tuning, with $b=\mathcal{O}(1)$ ($b=1$ in [1] and $b=16$ in [2]).

**Secondly, the variance of the FG/ZO can be reduced by various tricks.**
[2] explored ZO + gradient pruning, which reduced the effective number of dimensions based on the lottery hyperthesis.
[3] propose ZO + SVRG to reduce the variance.
We didn't explore these tricks in our initial attempt to apply FG/ZO in BO, but we believe these are promising directions for our future works to scale up $\text{(FG)}^2$U and $\text{(FG)}^2$U-ZO, as we have discussed in Appendix H.

**Further, $\text{(FG)}^2$U is suitable for modern AI computation due to its parallelizable nature.**
With the practice of scaling laws and the development of AI infrastructure, the computational cost concerns of large-scale models are mitigated when efficient parallelization is available. 
A vivid example is Transformer, whose quadratic complexity in sequence length is less favorable compared to RNNs, yet it has achieved impressive empirical success. 
The key to scaling up the Transformer lies in the parallelizable nature of attention.

Similarly, the inherently parallelizable and hardware-friendly nature of $\text{(FG)}^2$U enables it to leverage large-scale distributed computing resources (as mentioned in lines 146-148) within popular deep learning frameworks (as discussed in lines 196-204), with minimal engineering effort.

**Furthermore, the more cost-effective two-phase methodology can be utilized to reduce the overall computational expense.**
Considering the computational cost, we do not recommend using $\text{(FG)}^2$U from scratch, as discussed in Sec. 3.2. 
Instead, we advocate for a two-phase methodology: initially, employing efficient yet less accurate gradient approximation methods, such as TRGU or Hessian-Free, and subsequently using $\text{(FG)}^2$U for more accurate, albeit less efficient, gradient approximation to further enhance performance.

Finally, we want to emphasize the scope of this work: $\text{(FG)}^2$U is intended to complement, rather than overturn, the existing methodology of BO. 
The trade-off between computation and performance is a perpetual theme in computer science. 
Within the BO community, previous works have tended to sacrifice performance for reduced computational costs. 
We believe $\text{(FG)}^2$U will bring new insights to the BO community, prompting a reconsideration of this trade-off, especially in the current era of rapidly developing AI infrastructure and scaling law methodologies.

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

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

[3] Variance-reduced Zeroth-Order Methods for Fine-Tuning Language Models, https://arxiv.org/abs/2404.08080v1

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

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

## (B) Additional Experiments

Please refer the the attached PDF for additional results, including comparisons to additional baselines and real memory consumptions.

Meta-review: This paper introduces a Forward Gradient Unrolling with Forward Gradient, abbreviated as (FG)2U, achieving an unbiased stochastic approximation of the meta gradient. All the reviewers agree to accept this paper based on reasons of (1) simple and easy to implement, (2) well linked theoretical contribution, (3) comprehensive experiments, etc. I mostly agree.

Some concerns are also raised: (1) Theorem 3.4 indicates that the convergence rate depends linearly on N and states the various useful implications about it; (2) the computational cost takes the order of O(N) and it is suggested to add discussion on the zero-order BO.

Thus, I recommend an accept subject to addressing the mentioned issues in the final version
Decision: Accept (poster)