Functional Bilevel Optimization for Machine Learning
Authors: Ieva Petrulionytė
Julien Mairal
Michael Arbel
Venue: NeurIPS 2024 spotlight (2024)
Keywords: bilevel optimization
functional optimization
adjoint method
neural networks
Primary area: optimization
TL;DR: A functional perspective on bilevel optimization for deep learning applications.

Abstract: In this paper, we introduce a new functional point of view on bilevel optimization problems for machine learning, where the inner objective is minimized over a function space. These types of problems are most often solved by using methods developed in the parametric setting, where the inner objective is strongly convex with respect to the parameters of the prediction function. The functional point of view does not rely on this assumption and notably allows using over-parameterized neural networks as the inner prediction function. We propose scalable and efficient algorithms for the functional bilevel optimization problem and illustrate the benefits of our approach on instrumental regression and reinforcement learning tasks.

Functional Bilevel Optimization for Machine Learning
Ieva Petrulionyte, Julien Mairal, Michael Arbel
Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK, 38000 Grenoble, France f firstname. wlstname@iinii..c..
Abstract
In this paper, we introduce a new functional point of view on bilevel optimization problems for machine learning, where the inner objective is minimized over a function space. These types of problems are most often solved by using methods developed in the parametric setting, where the inner objective is strongly convex with respect to the parameters of the prediction function. The functional point view does not rely this assumption and notably allows using over- wver-parameterized neural networks as the inner prediction function. We propose scalable and efficient algorithms for the functional bilevel optimization problem and illustrate the benefits of our approach on instrumental regression and reinforcement learning tasks.
Introduction
Bilevel optimization methods solve problems with hierarchical structures, optimizing two interdepen- dent objectives: an inner-level objective and an outer-level one, Initially used in machine learning for model selection [Bennett et al.. 2006] and sparse feature learning [Mairal et al., 2012], these methods gained popularity as efficient alternatives to grid search for byyer-parameter tuning [Feurer and Hutter, 2019, Lorraine et al., 2019, Franceschi al., 2017]. Applications of bilevel optimization include meta-learning [Bertinetto et al., 2019], auxiliary task learning [Navon et al., 2021], ]. reinforcement learning [Hong et al., 2023, Liu al., 2021a, Nikishin et al., 2022], inverse problems [Holler et al., 2018] and invariant risk minimization [Arjovsky et al., 2019, Ahuja et al., 2020).
Bilevel problems are challenging solve, even in the well-defined bilevel setting with unique inner-level solution. This difficulty stems from approximating both the inner-level solution and its sensitivity to the outer-level variable during gradient-based optimization. Methods like Iterative Differentiation (ITD. Baydin et al., 2017) and Approximate Implicit Differentiation (AID, Ghadimi and Wang, 2018) were designed to address these challenges in the well-defined setting, resulting in scalable algorithms with strong convergence guarantees [Domke. 2012. Gould et al.. 2016. Ablin al., 2020, Arbel and Mairal, 2022a, Blondel et al., 2022, Liao al., 2018, Liu et al., 2022b, Shaban et al., 2019]. These guarantees usually require the inner-level objective to be strongly convex. However, when the inner-level variables are neural network parameters, the lower-level problem becomes non-convex and may have multiple solutions due to over- parameterization. While non- convexity considered "benign" in this setting [Allen-Zhr et al., 2019, Liu et al., 2022a], multiplicity of inner-level solutions makes their dependence on the outer-level variable ambiguous [Liu et al., 2021b), posing a major challenge bilevel optimization for modern machine learning applications.
We identify common functional structure in bilevel machine learning problems to address the ambiguity challenge that arises with flexible models like neural networks. Specifically, we consider prediction function h optimized by the inner-level problem over a Hilbert space H. This space consists of functions defined over an input space X' and taking values in a finite dimensional vector space V. The optimal prediction function is then evaluated in the luter-level to optimize an luter- parameter w in a finite dimensional space Q Rdd resulting in a functional bilevel problem:
F(w): Lout (w, ht) s.t. h wEn
min
arg min Lin (w,h).
(FBO)
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Functional approach
Parametric approach
Step 1:
Step 2:
Step 1:
Step 2:
Implicit differentiation
Function approximation
Function approximation
Implicit differentiation
Unique solution
Multiple solutions
liw> Liin (wo, h)
R
Oun Join (wo To)
V is well-defined
Prediction function
Model with parameters
Prediction function
Model with parameters 1
V V is ill-defined
Figure I: Parametric vs functional approaches for solving FBO by implicit differentiation.
In contrast to classical bilevel formulations involving neural networks, where the inner objective is non-convex with respect to the network parameters, the inner objective in (FBO) defines an optimization problem over a prediction function h in a functional vector space H.
A crucial consequence of adopting this new viewpoint is that renders the strong convexity of the inner objective with respect to h a mild assumption, which ensures the uniqueness of the solution h for any outer parameter value w. Strong convexity with respect to the prediction function is indeed much weaker than strong convexity with respect model parameters and often holds in practice. For instance, a supervised prediction task with pairs of features/labels (at, y) drawn from some training data distribution is formulated as regularized empirical minimization problem:
min Lin (wi h) := E.,y [llv - h h( NCH
w ||A
(1)
where H is the Le space of square integrable functions w.r.t the distribution of a, and w > 0 controls the amount of regularization. The inner objective is strongly convex in h, even though the optimal prediction function hn can be highly nonlinear in æ. The function h may then be approximated. e.g., by an overparameterized deep neural network, used here function approximation tool.
Although appealing, the (FBO) formulation necessitates the development of corresponding theory and algorithms, which is the aim of our paper. To the best of our knowledge, this is the first work to propose a functional bilevel point of view that can leverage deep networks for function approximation. The closest works are either restricted to kernel methods [Rosset, 2008, Kunapuli et al., 2008] and thus cannot be used for deep learning models, or propose abstract algorithms that can only be implemented for finite Hilbert spaces [Suonperä and Valkonen, 2024].
We introduce in Section 2 theoretical framework for functional implicit differentiation in an abstract Hilbert space H that allows computing the total gradient F(w)) using a functional version of the implicit function theorem [loffe and Tihomirov, 1979] and the adjoint sensitivity method [Pontryagin, 2018]. This involves solving well-conditionel functional linear system, equivalently formulated as a regression problem in H. to find an adjoint function a used for computing the total gradiens VF(w). We then specialize this framework to the common scenario where H is an L2 space and objectives are expectations of point- (see Sections 1. 4.2, and Appendix A). In Section 3, we propose an efficient algorithm where the prediction and adjoint functions can be approximated using parametric models, like neural networks, learned with standard stochastic optimization tools. We further study its convergence using analysis for biased stochastic gradient descent [Demidovich et al., 2024].
losses. This setting covers many machine learning problems
The proposed method, called finctional implicit differentiation (FuncID), adopts a novel "differentiate implicitly, then parameterize" approach (left Figure 1): functional strong convexity is first exploited to derive an unambiguous implicit gradient in function space using well-defined adjoint function. Then, both the lover-level solution and adjoint function are approximated using neural networks. This contrasts with traditional AID/TTD approaches (right Figure I), which parameterize the inner-level solution as a neural network, leading to a non-convex 'parametric" bilevel problem in the network's parameters. An ambiguous implicit gradient is then computed by approximately solving unstable or ill-posed linear system [Arbel and Mairal, 2022b]. Consequently, FuncID addresses the ambiguity challenge by exploiting the functional perspective and results in a stable algorithm with reduced time and memory costs. In Section 4, we demonstrate the benefits of our approach in instrumental regression and reinforcement learning tasks, which admit natural functional bilevel structure.
Related work on bilevel optimization with non-convex inner objectives. In principle, considering amended versions of the bilevel problem can resolve the ambiguity arising from non-convex inner

---

objectives. This is the case of optmmssticppssimisii versions of the problem, often considered in the literature on mathematical optimization, where the outer-level objective optimized over both outer and inner variables, under the optimality constraint of the inner-level variable [Dempe et al., 2007, Ye and Ye, 1997, Ye and Zhu, 1995, Ye et al., 1997]. While tractable methods were recently proposed to solve them [Liu et al., 2021a,b, 2023, Kwon et al., 2024, Shen and Chen, 2023], it is unclear how well the resulting solutions behave on unseen data in the context of machine learning. For instance, when using wer-pparamteriee models for the inner-level problem, their parameters must be further optimized for the outer-level objective, possibly resulting in over-fitting [Vicol et al., 2022]. More recently, Arbel and Mairal [2022b] proposed a game formulation involving selection map to deal with multiple inner-level solutions. Such a formulation justifies the use of ITD/AID outside the well-defined bilevel setting, by viewing those methods approximating the Jacobian of the selection map. However, the justification only hold under rather strong geometric assumptions. Additional related work is discussed in Appendix B on bilevel optimization with strongly-convex inner objectives, the adjoint sensitivity method that is often used in the context of ordinary differential equations, and amortization techniques [Amos al., 2023] that have been also exploited for approximately solving bilevel optimization problems [MacKay et al., 2019, Bac and Grosse, 2020].
2 A Theoretical Framework for Functional Bilevel Optimization
The functional bilevel problem (FBO) involves an optimal prediction function hig for each value of the outer-level parameter w. Solving (FBO) by using a first-order method then requires characterizing the implicit dependence of h on the luter-level parameter i to evaluate the total gradient VF(w)) in Rd Indeed, assuming that hi and Lout are Fréchet differentiable (this assumption will be discussed later), the gradient (W() may be obtained by an application of the chain rule:
VF(ww) gw + duhidww with gw == a Loat (w, h and d :- Lout (w) h+))
(2)
The Fréchet derivative 0,, : H. + Rd is linear operator acting on functions in H. and measures the sensitivity of the optimal solution on the outer variable. We will refer to this quantity as the "Jacobian" in the rest of the paper. While the expression of the gradient in Equation (2) might seem intractable in general, we will see in Section 3 class of practical algorithms to estimate it.
2.1 Functional implicit differentiation
Our starting point is to characterize the dependence of h on the outer variable. To this end, we rely on the following implicit differentiation theorem (proven in Appendix C) which can be seen functional version of the one used in AID [Domke, 2012, Pedregosa, 2016], albeit, under a much weaker strong convexity assumption that holds in most practical cases of interest.
Theorem 2.1 (Functional implicit differentiation). Consider problem (FBO) and assume that:
For any w E Ω there exists p>0 0 such that h + Lin(w', h) is p-strongly convex for any w' near w.
+ Lin(w, h) has finite values and is Fréchet differentiable on H for all w E St
-p is Hadamard differentiable on Q X H (in the sense of Definition C.I in Appendix C).
Then, + h% is uniquely defined and is Fréchet differentiable with Jacobian auhi given by:
B. + O. = 0, with B Ow, (wi, hw), and C= Li (w> h.).
The strong convexity assumption on the inner-level objective ensures the existence and uniqueness of the solution 15, while differentiability assumptions on Lin and Ör ensure Fréchet differentiability of the map w + h Though the implicit function theorem for Banach spaces [loffe and Tihomirov, 1979] could yield similar conclusions, it demands the stronger assumption that 2, is continuously Fréchet differentiable, which is quite restrictive in our setting of interest: for instance, when H an Lz-space and Lin is an integral functional of the form Lin (w, = f (w, h(x)) dx, with li defined on Ω x V and satisfying mild smoothness and growth assumptions on lin. then h + Ōn Lin Lan h)) cannot be Fréchet differentiable with uniformly continuous differential on bounded sets, unless + li (w,, v) is polynomial of degree at most 2 (see INemioovski and Semenov, 1973, Corollary 2, p 276) and discussions in [Noll, 1993, Goodman, 1971].. Instead, Theorem 2.1 employs the weaker notion of Hadamard diffcrentiability for 2 Lin, widely used in statistics, particularly for
3

---

deriving the delta-method [van der Vaart and Wellner, 1996, Chapter 3.9]. Consequently, Theorem 2.1 allows us cover a broader class of functional bilevel problems, as we see in Section 2.2.
Similarly to AID, only a dacobian-vector product is needed when computing the total gradient V.))). The result in Proposition 2.2 below, relies on the adjoint sensitivity method [Pontryagin, 2018] to provide a more convenient expression for F(W)) and is proven in Appendix C.2.
Proposition 2.2 (Functional adjoint sensitivity). Under the same assumption on Lin as in Theo- rem 2 and further assuming that Lout is jointly differentiable in and h, the total objective F is differentiable with VF()) given by: VF(ww = Iw + B.w
where the adjoint function at -w dw is an element of that minimizes the quadratic objective:
a arg min Lads (w,a) = (a, Cua) + (a, dw)).
aEH
Equation (4) indicates that computing the total gradient requires optimizing the quadratic objective (5) to find the adjoint function a The strong convexity of the adjoint objective Ladj ensures the existence of a unique minimizer, and stems from the positive definiteness of its Hessian operator Cu due to the inner-objective's 's strong convexity. Both adjoint and inner-level problems occur in the same function space H and are equivalent in terms of conditioning, as the adjoint Hessian operator equals the inner-level Hessian optimum. The strong convexity H of the adjoint objective guarantees well-defined solutions and holds in many practical cases, as opposed to classical parametric bilevel formulations which require the more restrictive strong convexity condition in the model's s parameters, and without which instabilities may arise due to ill-conditioned linear systems (see Appendix F.1).
2.2 Functional bilevel optimization in L2 spaces
Specializing the abstract results from Section 2. to common scenario in machine learning, we consider both inner and outer level objectives of (FBO) expectations of point-wisc functions over observed data. Specifically, we have two data distributions IP and Q defined over a product space X' x ) Rd, x 12 and denote by H the Hilbert space of functions h X + V, where V Rd.
Given an outer parameter space f2 we address the following functional bilevel problem:
min Lout (w, hw)) == E2 Lows (,,h((()) 3]] wEa
(6)
s.t. h arg min Lin (w, h) == Ep [ein (w, ((r), xr, v)]. kEH
where louts lin point-wai wise loss functions defined on Ω V x X' y. This setting encompasses various deep learning problems discussed in Sections 4.1 and 4.2, and in Appendix A, representing specific instance of FBO. The Hilbert space H of square-integrable functions not only models broad range of prediction functions but also facilitates obtaining concrete expressions for the total gradient (Ww).. enabling the derivation of practical algorithms in Section
The following proposition, proved in Appendix D, makes mild technical assumptions on probability distributions P. Q and the point-wise losses lins lout- It gives expression for the total gradient in the form of expectations under P and Q
Proposition 2.3 (Functional Adjoint sensitivity in L2 spaces.). Under the technical Assumptions (A) (L) stated in Appendix D.I, the conditions Lin and Louu in Proposition 2.2 hold and the total FF(w) = Iw + Bua with o E H being the minimizer of the objective Ludj in Equation (5). Moreover, Ladj. Iw and B are given by:
gradient F(w) of F is expressed as
Lodi (w. a) = EE> [a(x))oRenn ol in (w, (x),, x, x) a(x))
Eg [a(x) (lx Oolout (w, h
9w R lott (w, (x), Bo = Ep (0w, (w, hi, (x), v) (x)),,
(8)
where bont Oykoue, pin and 8lemn are partial first and second order derivatives of lont and lin with respect to their first and second arguments w and v.

---

Assumptions (A) and (B) on P and Q ensure finite second moments and bounded Radon-Nikodym
derivatives, maintaining square integrability under both distributions in Equation (6). Assumptions (C) to (L) on li and lout primarily involve integrabilityy differentiability, Lipschitz continuity, and strong convexity of lin in its second argument, typically satisfied by objectives like mean squared error or cross entropy (see Proposition D.1 in Appendix D.1). Next, by leveraging Proposition 2.3, we derive practical algorithms for solving FBO using function approximation tools like neural networks.
Methods for Functional Bilevel Optimization in L2 Spaces
We propose Functional Implicis Differentiation (FuncID), flexible class of algorithms for solving the functional bilevel problem L2 spaces described in Section 2.2 when samples from distributions P and Q are available.
FuncID relies on three main components detailed in the next subsections:
1. Empirical objectives. These approximate the objectives Lout, Lin and Lody empirical expectations over samples from inner and outer datasets Din and Dout-
2. Function approximation. The search space for both the prediction and adjoint functions is restricted to parametric spaces with dimee-dimensional parameters and &. Approximate solutions
hw and بي to the optimal functions h and o are obtained by minimizing the empirical objectives. 3. Total gradient approximation. FuncID estimates the total gradient VF(w) using the empirical objectives, and the approximations h.s and س the prediction and adjoint functions.
Algorithm provides outlines of FuncID which has nested structure similar to AID: (1) inner- level optimizations (InnerUpt and Adjoint0pt) to update the prediction and adjoint models using scalable algorithms such stochastic gradient descent [Robbins and Monro, 1951], and (2) an outer-level optimization to update the parameter w using a total gradient approximation TotalGrad. An optional warm-start allows initializing the parameters of both the prediction and adjoint models for the current outer-level iteration with those obtained from the previous one.
Algorithm FuncID
Input: initial outer, inner, and adjoint parameter wp. Gu. Eo: warm-start option WS. for O,...,N do
# Optional warm-start
if WS=True then (0n, En) +- En)) end if
Inner-level optimization
h h.w..0m 0u+1 + InnerOpt(w.. Oo Din)
Adjoint optimization
aw Sa+1 f- AdjointOpt(wn Eor Ien. D)
# Outer gradient estimation
Sample mini-batch B (Bouts Bio) from D = (Dout, Din)
Jont f TotalGrad(us) Newns awa. B)
Wa+1 update wn using Jouti
end for
3.1 From population losses empirical objectives
We assume access to two datasets Din and Douts comprising i.i.d. samples from IP and Q. respectively. This assumption can be relaxed, such as when using samples from a Markov chain or Markov Decision Process to approximate population objectives. For scalability, we operate in mini- batch setting, where batches = (Bout: B.,) are sub-sampled from datasets := (Douts D..). Approximating both inner and outer level objectives in (6) can be achieved using empirical versions:
Lout (w, h, Bout) = BBuuel
Σ (R.g)
Lout (w, h()), a, 9)
Lin (w. h, Bin)
Σ
ei (w:)(x) I,
[B
(a.
5

---

Adjoint objective. Using the expression of Lodj from Proposition 2.3, we derive finite-sample approximation of the adjoint loss by replacing the population expectations by their empirical counter- parts. More precisely, assuming we have access to an approximation i., to the inner-level prediction function, we consider the following empirical version of the adjoint objective:
Lads (w,a, BB = 18,m
Σ (2.g)EBI.
a(x) T Cnlia hes x. a(x)
alx) Uplout ( (w),
(9)
B
EBot
The adjoint objective in Equation (9) requires computing Hessian-vector product with respect to the output v in Rd of the prediction function hw, which typically of reasonably small dimension, unlike traditional AID methods that necessitate a Hessian-vector product with respect to some model parameters. Importantly, compared to AID, FuncID does not requires differentiating twice w.rt model's 's parameters r r0) which results in memory and time savings discussed in Appendix F.2.
3.2 Approximate prediction and adjoint functions
To find approximate solutions to the prediction and adjoint functions we rely on three steps: 1) speci- fying parametric search spaces for both functions, 2) introducing optional regularization prevent overfitting and, 3) defining gradient- -based optimization procedure on the empirical objectives.
Parametric search spaces. We approximate both prediction and adjoint functions using parametric search spaces. A parametric family of functions defined by map θ + H. over parameters C -Pin constrains the prediction function as /(x) = (z) We only require r to be continuous and differentiable almost everywhere such that baok-propagation Notably, unlike AID, we do not need 7 to be twice differentiable, as functional implicit differentiation computes the Hessian w.r.t. the output of т, not w.r.t. its parameters 0. For flexibility, can consider a different parameterized model parameters =\t IRP, constraining the adjoint similarly т. In practice, we often the same parameterization, typically neural network, for both the inner-level and the adjoint models.
be applied [Bolte et al., 2021].
- H. for approximating the adjoint function, defined over
Regularization. With empirical objectives and parametric search spaces, we can directly optimize parameters of both the inner-level model and the adjoint model However, to address finite sample issues, regularization may be introduced to these empirical objectives for better generalization. The method allows flexibility in regularization choice, accommodating functions 0 - R (0) and g +> Radj (E), such as ridge penalty or other commonly used regularization techniques
Optimization. The function InnerOpt (defined in Algorithm 2) optimizes inner model parameters for given w, initialization Og. and data Din. using M gradient updates. It returns optimized parameters 8M and the corresponding inner model h. T(0M). approximating the inner-level solution. Similarly, AdjointOpt (defined in Algorithm 3) optimizes adjoint model parameters with K gradient updates, producing the approximate adjoint function á_, v(Ex). ). Other optimization procedures may also be used, especially when closed-form solutions are available, as exploited in some experiments in Section 4. Operations requiring differentiation can be implemented using standard optimization procedures with automatic differentiation packages like PyTorch [Paszke et al., 2019] or Jax [Bradbury al., 2018].
Algorithm 2 InnerOpt(w, 0b. D...)
Algorithm AdjointOpt(w,, Eo. hw: D)
for m O...., M do Sample batch Bin from Din
for k O.... K do
Sample batch B from D
+ V Lin (w, T(mm), B.n) + HRn (9m)
Jadj + V LLadt (w, v(Ee), hw, B)d (Ex))
0m+11 + Update 9m using Iin
Ex+1 + Update En using Indj
end for
end for
Return ((0M)) 0M
Return v(EK),

---

3.3 Total gradient estimation
We exploit Proposition 2.3 derive Algorithm which allows us to approximate the total gradi- ent V.F(w) after computing the approximate solutions has and uw.. There, we decompose the gradient into two terms: 9Ezp. empirical approximation of Iw in Equation (8) representing the explicit dependence Lout on the outer variable w, and gimps an approximation to the implicit gradient term B in Equation (8). Both terms are obtained by replacing the expectations by empirical averages batches Bin. and Bout, and using the approximations h and dh instead of the exact solutions.
Algorithm 4 TotalCrad(w; Pww, iw 13)
gEep O.Loout (w, hus Bout)
91mp 5 Ba ElayEEB. Ow, (ws thw (x)) (z),, x.1) (1)
Return !Ezp + 91mp
3.4 Convergence Guarantees
Convergence of Algorithm to stationary points of depends on approximation errors, which result from sub-optimal inner and adjoint solutions, as shown by the convergence result below.
Theorem 3.1. Assume that F is L-smooth and admits a finite lower bound F* Use an update rule Wn+11 wn "[oout with step size O 1j S Th, in Algorithm 1. Under Assumption (a) on sub-optimality of hws and â Assumptions (b) ) to (h) on the smoothness of lin and lowr all stated in Appendix E.1, and the assumptions in Proposition 2.3, the iterates (waIn20 of Algorithm 1 satisfy:
min E [[I S ( - F*) + +nnooe1f + (Cit + (2tadz), OSREN-1
N
where C1, 022 6211 are positive constants, and Eins Eadj are sub-optimality errors that result from the inner and adjoint optimization procedures.
Theorem 3.1 is proven in Appendix and relies on the general convergence result in Demidovich et al. [2024, Theorem for stochastic biased gradient methods. The key idea is to control both bias and variance of the gradient estimator in terms of generalization errors enn and Eady when approximating the inner and adjoint solutions. These generalization errors can, in turn, made smaller in the case of over- weer-aaameterized networks, by increasing network capacity, number of steps and sample size [Allen-Zht et al., 2019, Du al., 2019, Zou al., 2020].
Applications
We consider two applications of the functional bilevel optimization problem: Two-stage least squares regression (2SLS) and Model-based reinforcement learning. To illustrate its effective- ness we compare it with other bilevel optimization approaches like AID or ITD, as well as state-of-the-art methods for each application. We provide a versatile implementation of FuncID http:://gghhub..co//iiiatttooooo
in PyTorch [Paszke et al., 2019], compatible with standard optimizers (e.g., Adam [Kingma and Ba, 2015)). and supports common regularization techniques. For the reinforcement learning application, we extend existing JAX [Bradbury et al., 201 8] im- plementation of model-based RL from Nikishin et al. [2022] to apply FuncID. To ensure fairness, experiments conducted with comparable computational budgets for hyperparameter tuning using the MLXP experiment management tool [Arbel and Zouaoui, 2024]. Additionally, we maintain consistency by employing identical neural network architectures across all methods and repeating experiments multiple times with different random seeds.
4.1 Two-stage least squares regression (2SI.S)
Two-stage least squares regression (2SLS) is commonly used in causal representation learning, including instrumental regression or proxy causal learning [Stock and Trebbi, 2003]. Recent studies have applied bilevel optimization approaches to address 2SLS, yiclding promising results [Xu et al.,
7

---

34x30
Out- Out-ot-Sampli MSE
15:30
60
Olieer Loss
Vel. penv. DPY
Inner Loss
Ag
panal
23x30
22x30
Fure
GD penai
AD
Val. paral
OPl
PandD
2439
606
100 150
1030
175
nD DFV fundb In. Fundo
Outer Iterations
Innar Iterations
Figure 2: Performance metrics for Instrumental Variable (IV) regression. All results are averaged over 20 rums with 5000 training samples and 588 test samples. (Left) box plot of the test loss, with the dashed black line indicating the mean test error. (Middle) outer loss vs training iterations, (Right) inner loss vs training iterations. The bold lines in the middle and right plots indicate the mean loss, the shaded area corresponds to standard deviation.
2021b,a, Hu et al., 2023]. We particularly focus 2SLS for Instrumental Variable (IV) regression, widely-used statistical framework for mitigating endogeneity in econometrics [Blundell et al., 2007, 2012], medical economics [Cawley and Meyerhoefer, 2012], sociology [Bollen, 2012], and more recently, for handling confounders in off-line reinforcement learning [Fu et al., 2022].
Problem formulation. In an IV problem, the objective is to model fw + o that approximates the structural function fstruct using independent samples (o, f.): from data distribution PP, where is an instrumental variable. The structural function fetruct delineates the true effect of a treatment t on an outcome A significant challenge in IV is the presence of an unobserved confounder e, which influences both l and o additively, rendering standard regression ineffective for recovering fa- However, if the instrumental variable æ solely impacts the outcome through the treatment and is independent from the confounder e, it can be employed to elucidate the direct relationship between the treatment and the outcome using the 2SLS framework, under mild assumptions on the confounder [Singh et al., 2019]. This adaptation replaces the regression problem with a variant that averages the effect of the treatment conditionally on
min Ee WEn
[10 - E. 1fw (t) [+)]=
(10)
Directly estimating the conditional expectation Er- [z] is hard in general. Instead, it is easier express it, equivalently, as the solution of another regression problem predicting (t) from z:
to
h Ep [1)az) = arg min Ep [1],..() h(<)))²]
(11)
NER
Both equations result in the bilevel formulation in Equation (6) with y (t, 0). Q P and the wiset-wise losses fn and lout given by (w, 2, 8) = (w,, v, x, o)) flf o and lomt(as (w, u, v. = Lovt (w, v, x. (t, 0)) = 1|o - w|l It is, therefore, possible to directly apply Algorithm to learn for as we illustrate below.
Experimental setup. We study the IV problem using the dsprites dataset [Matthey et al., 2017]. comprising synthetic images representing single objects generated from five latent parameters: shape, scale, rotation, and posX, posY positions on image coordinates. Here, the treatment variable / is the images, the hidden confounder is the posY coordinate, and the other four latent variables form the instrumental variable The outcome o is an unknown structural function Sstruct of t, contaminated by confounder as detailed in Appendix G.I. We follow the setup of the Deep Feature Instrumental Variable Regression (DFIV) dsprites experiment by Xu et al. [2021a, Section 4.2], which achieves -fth--of- performance. structural model, optimized to solve the bilevel problem in Equations (10) and (11). We explore two versions of our method: FuncID, which optimizes all adjoint network parameters, and FuncID linear, which learns only the last layer in closed-fome while inheriting hidden layer parameters from the inner prediction function. We compare our method with DFIV, AID, ITD, and Penalty- -bsed methods: gradient penalty (GD penal.) and value function penalty (Val penal.) [Shen and Chen, 2023]. using identical network architectures and computational budgets for hyperparameter selection. Full details on network architectures, hyperparameters, and training settings are provided in Appendix G.2.
this setup, neural networks serve as the prediction function and
8

---

Results. Figure 2 compares structural models learned by different methods using 5K training samples (refer Figure in Appendix G.3 for 10K sample results). The left subplot illustrates out-of-sample mean squared error of learned structural models compared to ground truth outcomes (uncontaminated by noise <), while the middle and right subplots show the evolution of outer and inner objectives over iterations. For the 5K dataset, FuncID) outperforms DFIV (p-value=0.003, one-sided paired t-test), while showing comparable performance on the 10K dataset (Figure 6). AID and ITD perform notably worse, indicating their parametric approach fails to fully leverage the functional structure. FuncID outperforms the gradient penalty method and performs either better comparably the value function penalty method, though the latter shows higher variance with some particularly bad outliers. While all methods achieve similar outer losses, this criterion alone is only reliable as indicator of convergence when evaluated near the "exact' inner-level solution corresponding to the lowest inner-loss values. Interestingly, FuncID obtains the lowest inner-loss values, suggesting its outer-loss is more reliable indicator of convergence.
4.2 Model-based reinforcement learning
Model-based reinforcement learning (RL) naturally yields bilevel optimization formulations, since several components of an RL agent need to be learned using different objectives. Recently, Nikishin et al. [2022] showed that casting model-based RL as a bilevel problem can result in better tolerance to model- model-miisspecifccttion Our experiments show that the functional bilevel framework yields improved results even when the model is well-specified, suggesting a broader use of the bilevel formulation.
Problem formulation. In model-based RL, the Markov Decision Process (MDP) is approximated by a probabilistic model d* with parameters w that can predict the next state 8 (x) and reward (x:), given a pair == (s, a) where is the current environment state and is the action of an agent. A second model h can be used to approximate the action-value function (xx) that computes the expected cumulative reward given the current state-action pair. Traditionally, the action-value function is learned using the current MDP model, while the latter is learned independently from the action-value function using Maximum Likelihood Estimation (MLE) [Sutton, 1991].
In the bilevel formulation of model- based RL by Nikishin et al. [2022]. the inner-level problem in- volves learning the optimal vetion-value function h. with the current MDP model q. and minimizing the Bellman error. The inner-level objective can be expressed as an expectation of a point-wise loss f with samples (xr, r", ss) P, derived from the agent-environment interaction:
h arg min Ex [f(h(x), rw()), (x))) NEH
(12)
Here, the future state and reward (r', s') are replaced by the MDP model predictions rw (x) and (z). In practice, samples from IP are obtained using a replay buffer. The buffer accumulates data over several episodes of interactions with the environment, and can therefore be considered independent of the agent's ' policy. The point-wise loss function f represents the error between the action-value function prediction and the expected cumulative reward given the current state-action pair:
[(, 8¹) 1 - r
- log Eas clee .a)
with h lagged version of (exponentially averaged network) and y a discount factor. The MDP model is learned implicitly using the optimal function h3 by minimizing the Bellman error w.r.t. wt
min Ep [f ,')). EE
(13)
Equations (12) and (13) define a bilevel problem as in Equation (6), where Q = P. = (r', s'), and the point-wise losses Ly and lout are given by: ey (w, v, y) I (v, (x), sw (x)) and lout (w, v, 8) = f (v, r', s'). Therefore, we can directly apply Algorithm I to learn both the MDP model qu and the optimal action-value function h
Experimental setup. We apply FuncID to the CartPole control problem. classic benchmark in reinforcement learning [Brockman et al., 2016, Nagendra et al., 2017]. The goal is to balance pole attached to a cart by moving the cart horizontally. Following Nikishin et al. [2022], we use model- based approach and consider two scenarios: one with well-specified network accurately representing the MDP, and another with misspecified model having fewer hidden layer units, limiting its capacity.

---

162
Reward
Reward
Algorithms AID (OMD)
Algorithms
AID (OMB)
0.00
0.25
1.00
1.25
1.50
1.75
2,00
Steps
Steps
Figure 3: Average reward on evaluation environment vs. training iterations the CartPole task. (Left) Well-specified model with 32 hidden units. (Right) Misspecified model with hidden units. Both plots show mean reward over 10 runs where the shaded region is the 95% confidence interval.
Using the bilevel formulation in Equations (12) and (13), we compare FuncID with the Optimal Model Design (OMD) algorithm [Nikishin et al., 2022], variant of AID. Additionally, we compare against a standard single-level model-based RI. formulation using Maximum Likelihood Estimation (MLE) [Sutton, 1991]. For the adjoint function in FuncID,, we derive simple closed-for expression based on the structure of the adjoint objective (see Appendix H.1). We follow the experimental setup of Nikishin et al. [2022], providing full details and hyperparameters in Appendix H.2.
Results. Figure 3 illustrates the training reward evolution for FuncID, OMD, and MLE in both well-specified and misspecified scenarios. FuncID consistently performs well across settings. In the well-specified case, where OMD achieves a reward of 4, FuncID reaches the maximum reward of 5, matching MLE (left Figure 3). In the misspecified scenario, FuncID performs comparably to OMD and significantly outperforms MLE (right Figure 3). Moreover, FuncID tends to converge faster than MIE (see Figure in Appendix H.3) and yields consistently better prediction (see Figure 8 in Appendix H.3). These findings align with Nikishin et al. [2022], suggesting that MLE may prioritize minimizing prediction errors, potentially leading overfitting irrelevant features. In contrast, OMD and FuncID focus on maximizing expected returns, especially in the presence of model misspecification. Our results highlight the effectiveness of (FBO) even in well-specified settings, suggesting, for future work, further investigations for more general RL tasks.
than OMI)
5 Discussion and concluding remarks
This paper introduced functional paradigm for bilevel optimization in machine learning, shifting focus from parameter space function space. The proposed approach specifically addresses the ambiguity challenge arising from using deep networks in bilevel optimization. The paper establishes the validity of the functional framework by developing theory of functional implicit differenntation. proving convergence for the proposed FuncID method, and numerically comparing it with other bilevel optimization methods.
The theoretical foundations of our work rely on several key assumptions worth examining. While our convergence guarantees assume both inner and adjoint optimization problems are solved to some optimality, this assumption is supported by recent results on global convergence in over-parameterized networks [Allen-Zhu et al., 2019, Liu et al., 2022a). However, quantifying these optimality errors more precisely and understanding their relationship to optimization procedures remains an open challenge. Additionally, like other bilevel methods, our approach requires careful hyperparameter selection, which can impact practical implementation.
Several promising directions emerge for future research. While we focus on L2 spaces, exploring alternative function spaces (such Reproducing Kernel Hilbert Spaces or Sobolev spaces) could reveal additional advantages for specific applications. Furthermore, extending our framework to non-smooth objectives or constrained optimization problems, potentially building on existing work in non-smooth implicit differentiation [Bolte et al., 2022], would broaden its applicability.
10

---

Acknowledgments
This work was supported by the ERC grant number 101087696 (APHELAIA project) and by ANR 3IA MIAI@ Grenoble Alpes CE23-0012-01). We thank Edouard Pauwels and Samuel Vaiter for their insightful discussions.
and the ANR project BONSAI (grant ANR-23-
References
P. Ablin, G. Peyré, and T. Moreau. Super-efficiency of automatic differentiation for functions defined as minimum. International Conference on Machine Learning (ICML), 2020.
K. Ahuja, K. Shanmugam, K. Varshney, and A. Dhurandhar. Invariant risk minimization games. International Conference on Machine Learning (ICML), 2020.
Z. Allen-Zhu Y. Li, and Z. Song. A convergence theory for deep learning via over- prraaaacttttiiiiionnn In International Conference on Machine Learning (ICML). 2019.
B. Amos et al. Tutorial on amortized optimization. Foundations and Trends@ in Machine Learning. 16(5):592-732, 2023.
M. Arbel and J. Mairal. Amortized implicit differentiation for stochastic bilevel optimization. International Conference on Learning Representations (ICLR), 2022a.
M. Arbel and Mairal. Non-convex bilevel games with critical point selection maps. Advances in Neural Information Processing Systems (NeurIPS), 2022b.
M. Arbel and A. Zouaoui. Mlxp: A framework for conducting replicable machine learning ex- periments in python. arXiv preprint arXiv:2402.13831 2024. URL https://githir com/
Anria-thoth/mine
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint 1907.0289) 2019.
Bae and R. B. Grosse, Delta-stn: Efficient bilevel optimization for neural networks using structured response jacobians. Advances in Neural Information Processing Systems (NeurIPS), 2020.
D. Bansal, R. T. Chen, M. Mukadam, and B. Amos. Taskmet: Task-driven metric learning for model learning. Advances in Neural Information Processing Systems (NeurIPS), 2023.
H. H. Bauschke and P. L. Combettes. Convex Analysis and Monotone Operator Theory in Hilbert Spaces. Springer Publishing Company, Incorporated, 1st edition, 2011. ISBN 1441994661.
A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind. Automatic differentiation in
machine learning: A survey. Journal of Machine Learning Research (JMLR), 18(153): 1-43, 2017.
Y. Bengio, Simard, and P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions Neural Networks, 5(2):157-166, 1994.
K. P. Bennett, Hu, X. Ji, G. Kunapuli, and J.-S. Pang. Model selection via bilevel optimization. IEEE International Joint Conference on Neural Network Proceedings, 2006.
L. Bertinetto, J. F. Henriques, P. H. Torr, and A. Vedaldi. Meta-learning with differentiable closed- form solvers. International Conference on Learning Representations (ICLR), 2019.
M. Binikowki, D. J. Sutherland, M. Arbel, and Gretton. Demystifying mmd gans. arXiv preprint 1801.01807.01444 2018.
M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer, F. Llinares-López, F. Pedregosa, and J.-P. Vert. Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems (NeurIPS), 2022.
R. Blundell, X. Chen, and D. Kristensen. Semi-nooparametric iv estimation of shape-invariant engel curves. Econometrica. 75(6):6613-166, 2007.
R. Blundell, J. L. Horowitz, and M. Parey. Measuring the price responsiveness of gasoline demand: Economic shape restrictions and nonparametric demand estimation. Quantitative Economics, 3(1): 29-51, 2012.
K. Bollen. Instrumental variables in sociology and the social sciences. Annual Review of Sociology. 38:37-72, 2012.

---

J. Bolte, T. Le, E. Pauwels, and T. Silveti-Falls. Nonsmooth implicit differentiation for machine- learning and optimization. Advances in Neural Information Processing Systems (NeurIPS), 2021.
J. Bolte, E. Pauwels, and S. Vaiter. Automatic differentiation of nonsmooth iterative algorithms. Advances in Neural Information Processing Systems (NeurIPS), 2022.
J. Bolte, E. Pauwels, and S. Vaiter. One-step differentiation of iterative algorithms. Advances in Neural Information Processing Systems (NeurIPS), 2024.
J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, VanderPlas, Wanderman-MMlne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL. http: / //ggtttttb/.... com/ hoan/gooole/aaa
A. Brock, T. Lim, J. Ritchie, and N. Weston. SMASH: One-shot model architecture search through hypernetworks. International Conference on Learning Representations (ICLR), 2018.
G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai
gym. arXiv preprint 1606.01540, 2016.
Cawley and C. Meyerhoefer. The medical care costs of obesity: an instrumental variables approach. Journal of Health Economics, 31(1):219 2300 2012.
R. Т. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems (NeurIPS), 2018.
L. Debnath and P. Mikusinski. Introduction to Hilbert spaces with applications. Academic press, 2005.
Y. Demidovich, G. Malinovsky, I. Sokolov, and P. Richtárik. A guide through the zoo of biased sgd. Advances in Neural Information Processing Systems (NeurIPS), 2024.
S. Dempe, J. Dutta, and B. Mordukhovich, New necessary optimality conditions in optimistic bilevel programming. Optimization, 56(5-6))577-604, 2007.
J. Domke. Generic methods for optimization-base modeling. International Conference on Artificial Intelligence and Statisties (AISTATS), 2012.
S. Du, J. Lee, H. Li, L. Wang, and X. Zhai Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning (ICML), 2019.
Z. Fang and A. Santos. Inference on Directionally Differentiable Functions. The Review of Economic Studies, 86(1 (6(1):377-412, 2018
M. Feurer and F. Hutter. Hyperparameter optimization. Springer International Publishing, 2019.
L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. Forward and reverse gradient- hyperpa- rameter optimization. International Conference on Machine Learning (ICML), 2017.
Z. Fu, Z. Qi, Z. Wang, Z. Yang, Y. Xu, and M. R. Kosorok. Offline reinforcement learning with instrumental variables in confounded markov decision processes. arXiv preprint 2209.08666 2022.
S. Ghadimi and M. Wang. Approximation methods for bilevel programming. Optimization and Control, 2018
V. Goodman. Quasi-differentiable functions of banach spaces. Proceedings of the American Mathe- matical Society, 30(2):367-370, 1971.
S. Gould, B. Fernando, A. Cherian, P. Anderson, R. S. Cruz, and E. Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint 1607.05447, 2016.
R. Grazzi, L.. Franceschi, M. Pontil, and S. Salzo. On the iteration complexity of hypergradient computation. International Conference on Machine Learning (ICMI.), 2020.
D. Ha, A. M. Dai, and Q Le. Hypernetorrss. International Conference on Learning Representa- tions (ICLR), 2017.
G. Holler, K. Kunisch, and R. C. Barnard. A bilevel approach for parameter learning in inverse problems. Inverse Problems, 34(11):115012, 2018.
M. Hong, H. Wai, Z. Wang, and Z. Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application actor-critic, SIAM Journal on Optimization, 33(1):147-180, 2023.
12

---

Y. Hu, J. Wang, Y. Xie, A. Krause, and D. Kuhn. Contextual stochastic bilevel optimization. Advances in Neural Information Processing Systems (NeurIPS), 2023.
A. D. loffe and V. M. Tihomirov. Theory of Extremal Problems. Series: Studies in Mathematics and its Applications 6. Elsevier, 1979.
K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design. International Conference on Machine Learning (ICML), 2021.
J. Jia and A. Benson. Neural jump stochastic differential equations. Advances in Neural Information Processing Systems (NeurIPS), 2019.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), 2015.
D. P. Kingma and M. Welling. Auto-encoding variational bayes. International Conference on Learning Representations (ICLR), 2014.
G. Kunapuli, K. Bennett, J. Hu, and J.-S. Pang. Bilevel model selection for support vector machines. CRM Proceedings and Lecture Notes, 45:129-158, 2008.
J. Kwon, D. Kwon, S. Wright, and R. D. Nowak. On penalty methods for nonconvex bilevel optimization and first-order stochastic approximation. In International Conference on Learning Representations (ICLR), 2024.
S. Li, Z. Wang, A. Narayan, R. Kirby, and S. Zhe. Meta-learning with adjoint methods. International Conference on Artificial Intelligence and Statistics (AISTATS), 2023.
X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
R. Liao, Y. Xiong, E. Fetaya, L. Zhang, K. Yoon, X. Pitkow, R. Urtasun, and R. Zemel. Reviving and improving recurrent back- propaapppaga International Conference on Machine Learning (ICML), 2018.
C. Liu, L. Zhu, and M. Belkin. Loss landscapes and optimization in over wver-prraetteiieed non-lineat systems and neural networks. Applied and Computational Harmonic Analysis, 59:85-116, 2022a.
R. Liu, X. Liu, S. Zeng. J. Zhang, and Y. Zhang. Value-function dunstion-based sequential minimization for bi-level optimization. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 45:15930-15948, 2021a.
R. Liu, Liu, S. Zeng, and J. Zhang. Towards gradient- based bilevel optimization with non-convex followers and beyond. Advances in Neural Information Processing Systems (NeurIPS), 2021b.
R. Liu, J. Gao. J. Zhang. D. Meng, and Z. Lin. Investigating bi-level optimization for learning and vision from a unified perspective: a survey and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 10455-00100077 2022b.
R. Liu, Y. Liu, W. Yao, S. Zeng, and J. Zhang. Averaged method of multipliers for bi-level optimization without lower-level strong convexity. In International Conference on Machine Learning (ICML), pages 21839-21886 PMLR, 2023.
J. Lorraine, P. Vicol, and D. K. Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
M. MacKay, P. Vicol, Lorraine, D. Duvenaud, and R. B. Grosse. Self-tuning networks: Bilevel op- timization of hyperparameters using structured best-response functions. International Conference on Learning Representations (ICLR), 2019.
Mairal, F. Bach, and Ponce. Task-driven dictionary learning. IEEE Transactions on Puttern Analysis and Machine Intelligence (PAMI), 34(4):791 8844, 2012.
C. C. Margossian and M. Betancourt. Efficient automatic differentiation of implicit functions. arXiv preprint 2/ 17, 2021.
Marion, A. Korba, Bartlett, M. Blondel, V. De Bortoli, A. Doucet, F. Llinares-Lópezz C. Paquette, and Q Berthet. Implicit diffusion: Efficient optimization through stochastic sampling. arXiv preprint arXiv:2402.05466; 2024.
L. Matthey, I. Higgins, D. Hassabis, and A. Lerchner. dsprites: Disentanglement testing sprites dataset, 2017.
13

---

S. Nagendra, N. Podila, R. Ugarakhod, and K. George. Comparison of reinforcement learning algorithms applied the cart-pole problem. International Conference on Advances in Computing. Communications and Informatics (ICACCI). 2017.
A. Navon, I. Achituve, H. Maron, G. Chechik, and E. Fetaya. Auxiliary learning by implicit differentiaaion. International Conference on Learning Representations (ICLR). 2021.
A. Nemirovski and S. Semenov. On polynomial approximation of functions on hilbert space. Mathematics of the USSRRSSborin 21(2):255, 1973.
E. Nikishin, R. Abachi, R. Agarwal, and P.-L.. Bacon. Control-oriented model-based reinforcement learning with implicit differentiation. AAAI Conference on Artificial Intelligence, 2022.
D. Noll. Second order differentiabilityy of integral functionals on Sobolev spaces and L2- Walter de Gruyter, Berlin/New York Berlin, New York, 1993.
R. Pascanu, T. Mikolov, and Y. Bengio. On the difficulty of training recurrent neural networks. International Conference on Machine Learning (ICML., 2013.
A. Paszke, S. Gross, Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang. J. Bai, and S. Chintala. Pytorch: An imperative style, high-peeformance deep
learning library. Advances in Neural Information Processing Systems (NeurIPS), 2019.
F. Pedregosa. Hyperparameter optimization with approximate gradient. International Conference on Machine Learning (ICML., 2016.
S. Pontryagin. Mathematical Theory of Optimal Processes. Routledge, 2018.
J. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference
in deep generative models. International Conference on Machine Learning (ICML). 2014.
H. Robbins and S. Monro. A stochastic approximation method. The Annals of Mathematical Statistics, pages 400-407, 1951.
S. Rosset. Bi-level path following for cross validated solution of kernel quantile regression. Interna- tional Conference on Machine Learning (ICML), 2008.
A. Shaban, C.-A. Cheng, N. Hatch, and B. Boots. Truncated packkpropagation for bilevel optimization. International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
H. Shen and T. Chen. On penalty-based -baee bilevel gradient descent method. International Conference on Machine Learning (ICMI), 2023.
R. Singh, M. Sahani, and A. Gretton. Kernel instrumental variable regression. Advances in Neural Information Processing Systems (NIPS), 2019.
J. H. Stock and F. Trebbi. Retrospectives: Who invented instrumental variable regression? Journal of Economic Perspectives, 1733)1177-14, 2003.
E. Suonperä and T. Valkonen. Linearly convergent bilevel optimization with single-step inner methods. Computational Optimization and Applications, 87(2):571-610 2024.
R. S. Sutton, Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 164):160-163, 1991.
A. der Vaart. Efficiency and hadamard differentiabilityy Scandinavian Journal of Statistics, 18(1): 63-75, 1991.
A. W. van der Vaart and J. A. Wellner. Weak Convergence and Empirical Processes with Applications to Statistics, pages 16-28. Springer New York, 1996.
P. Vicol, J. Lorraine, D. Duvenaud, and R. Grosse. Implicit regularization in overparameterized bilevel optimization. International Conference on Machine Learning (ICMI.), 2022.
C. Williams and M. Seeger. Using the nyström method speed up kernel machines. Advances in Neural Information Processing Systems (NIPS), 2000.
L. Xu, Y. Chen, S. Srinivasan, N. de Freitas, A. Doucet, and A. Gretton. Learning deep features in instrumental variable regression. International Conference on Learning Representations (ICLR), 2021a.
I. Xu, H. Kanagawa, and Gretton. Deep proxy causal learning and its application to confounded bandit policy evaluation. Advances in Neural Information Processing Systems (NeurIPS), 2021b.
14

---

J. Ye and X. Ye. Necessary optimality conditions for optimization problems with variational inequality constraints. Mathematics of Operations Research, 22(4):977-9977, 1997.
J. Ye and D. Zhu, Optimality conditions for bilevel programming problems. Optimization, 33(1): 9-27, 1995.
Ye, D. Zhu, and Q. J. Zhu. Exact penalization and necessary optimality conditions for generalized bilevel programming problems. SIAM Journal on Optimization, 7(2):481-507, 1997.
C. Zhang, M. Ren, and R. Urtasun. Graph hypernetworks for neural architecture search. International Conference on Learning Representations (ICLR), 2019.
Y. D. Zhong, B. Dey, and A. Chakraborty. Symplectic ode-net: Learning hamiltonian dynamics with control. arXiv preprint 1909.12077, 2019.
D. Zou, Y. Cao, D. Zhou, and Q. Gu. Gradient descent optimizes over- serepprramterized deep relu networks. Machine learning, 109:467-492, 2020.
15

---

A Examples of FBO formulations
The functional bilevel setting applies to various practical bilevel problems where objectives depend solely on model predictions, not their parameterization. Below, we discuss a few examples.
Auxiliary task learning. As in Equation (1), consider main prediction task with features x and labels y, equipped with a loss function f(y, ). The goal of auxiliary task learning is to learn how set of auxiliary tasks represented by vector fouz (y, h(z)) could help solve the main task. This problem is formulated by Navon et al. [2021] as a bilevel problem, which can be written as (FBO) with
Lout (ws, h) = B(yra))Duua [f h(x))): where the loss is evaluated over a validation dataset Doal, and
Lin (w, h) Deynen (f h((x)) + 9w (faua (y, h(x)))),
where an independent training dataset Derain is used, and In is a function that combines the auxiliary losses into scalar value.
Task-driven metric learning. Considering now a regression problem with features and labels y, the goal of taskedriven metric learning formulated by Bansal et al. [2023] is to learn metric parameterized by w for the regression task such that the corresponding predictor ht performs well on a downstream task Ltask. This can be formulated as (FBO) with Loue h) (h) and
Lin(w, h) = B(u..) [1l - ((x))
where 11 is the squared Mahalanobis norm with parameters w and (z) is a data data-dependent metric that allows emphasizing features that are more important for the downstream task.
B Additional Related Work
Bilevel optimization in machine learning. Two families of bilevel methods are prevalent in machine learning due to their scalability: iterative (or 'unrolled') differentiation (TTD, Baydin et al., 2017) and Approximate Implicit Differentiation (AID, Ghadimi and Wang, 2018). ITD approximates the optimal inner-level solution using an 'unrolled' function from a sequence of differentiable optimization steps, optimizing the outer variable via back-propagation [Shaban et al., 2019, Bolte et al., 2024]. The gradient approximation error decreases linearly with the number of steps when the inner-level is strongly convex, though at increased computational and memory costs [Grazzi et al., 2020, Theorem 2.1]. ITD is popular for its simplicity and availability in major deep learning libraries [Bradbury al., 2018], but can be unstable, especially with non-convex inner objectives [Pascanu et al., 2013, Bengio et al., 1994, Arbel and Mairal, 2022b]. AID uses the Implicit Function Theorem (IFT) derive the Jacobian of the inner-level solution with respect the outer variable [Lorraine et al., 2019, Pedregosa, 2016), solving a finite-dimensional linear system for an adjoint vector representing optimality constraints. AID offers strong convergence guarantees for smooth, strongly convex inner objectives [Ji et al., 2021, Arbel and Mairal, 2022a]. However, without strong convexity, the linear system can become ill ill-posed due to a degenerate Hessian, leading instabilities, especially with overparameterized deep neural networks. Our proposed approach avoids this issue, even when using deep networks for function approximation.
Adjoint sensitivity method. The adjoint sensitivity method [Pontryagin, 2018] efficiently dif- ferentiates a controlled variable with respect to a control parameter. In bilevel optimization, AID applies a finite-dimensional version of this method [Margossian and Betancourt, 2021, Section 2]. Infiniteddimensional versions have differentiated solutions of ordinary differential equations (ODEs) with respect to defining parameters [Margossian and Betancourt, 2021, Section 3]. and have been used in machine learning to optimize parameters of vector field describing an ODE [Chen et al., 2018]. Here, the ODE's vector field, parameterized by neural network. is optimized to match observations. The adjoint sensitivity method offers an efficient alternative to the unstable process of back-propagatis through ODE solvers, requiring only the solution of an adjoint ODE to compute gradient updates, improving performance [Jia and Benson, 2019, Zhong et al., 2019, Li et al., 2020]. This method has been adapted to meta-lcarning [Li et al., 2023], viewing the inner optimization as an
16

---

ODE evolution with gradients obtained via the adjoint ODE. Recently. Marion et al. [2024] employ the adjoint sensitivity method for optimizing diffusion models. where an adjoint SDE solved to compute the total gradient. Unlike these works, which use the adjoint method for ODE/SDE solutions as functions of time, our work applies an infinie-dimmnnsioall version of the adjoint sensitivity method to general learning problems, where solutions are functions of input data.
Amortization. Recently several methods have used amortization approximately solve bilevel problems [MacKay et al., 2019, Bae and Grosse, 2020]. These methods employ a parametric model called hypernetwork [Ha al., 2017, Brock et al., 2018, Zhang et al., 2019], optimized directly predict the inner-level solution given the outer-level parameter w. Amortized methods treat the two levels as independent optimization problems: (1) learning the hypernetwork for a range of w, and (2) performing first-order descent on using the hypernetwork
proxy for the inner-level solution. Unlike ITD, AID, or our functional implicit differentiation method, amortized methods do not fully exploit the implicit dependence between the two levels. They are similar to amortized variational inference [Kingma and Welling, 2014, Rezende et al., 2014], where a parametric model produces approximate samples from a posterior distribution. Amortization methods perform well when the inner solution's dependence on w is simple but may fail otherwise [Amos et al., 2023, pages 71-72). In contrast, our functional implicit differentiation framework adapts to complex implicit dependencies between the inner solution and w.
C Theory for Functional Implicit Differentiation
C.1 Preliminary results
We recall the definition of Hadamard differentiability and provide in Proposition C.2 general property for Hadamard differentiable maps that we will exploit to prove Theorem 2.1 in Appendix C.2.
Definition C.1. Hadamard differrentiability. Let A and B be two separable Banach spaces. A function L A + B is said to be Hadamard differentiable [van der Vaart, 1991, Fang and Santos, 2018] if for any a A. there exist a continuous linear map d,, L(a) : A + B that for any sequence (un)n21 in A converging to an element E A, and any real valued and non-vanishing sequence (tn)n21 converging to 0, it holds that:
_ (I(a + (nun)) - -da) La)y
- d.
a-+++- O.
(14)
in
Proposition C.2. Let A and B be two separable Banach spaces. Let L A B be Hadamard differentiable map with differential da L. at point Consider bounded linear map defined over a euclidean space Th of finite dimension n and taking values in A, i.e J R - A. Then, the following holds:
Lia Ju) = L(a) + de La))u +((lu))
Proof. Consider sequence (ux) )=2 (wx)<21 in R that us converges to 0 with |ull > 0 for all k 2 1 and define the first order error Eg as follows:
Ex
I[(a + Juk) - L(a) dd daL(a)ukll
The goal is to show that Eg converges to O. We can write ug as us twū with t || and
|ل = 1, so that:
F ||tzoll
(I(a + (ūx) - I(a)) de Majaus
If Es were unbounded, then, by contradiction, there must exist a subsequence (E0())22) converging to +oo, with o(k) increasing and ((k) + +oc. Moreover, since üg is bounded, one can further choose the subsequence Eω(E) that ( converges to some element ů. We can the following upper-bound:
Eg S
(L(a + tūx) L(a)) - d La)ūū + ||d (())l || - ū|| (15)
||tx||
E2
17

---

where we used that du ((a) is bounded. Since L is Hadamard differentiable, Ea(s) converges to O. Moreover, | - ã| also converges to O. Hence, Bş( converges to 0 which contradicts E Therefore, Eg is bounded.
Consider now any convergent subsequence of (Ex)a21. Then, it can be written as (Eag)) with (kk) increasing and ((R) +o0. We then have Flg(k)
+00 by construction. Since üg is bounded, one can further choose the subsequence Eœ(R) so that (o(K)) converges to some element deduce that Ex(k) must converge 0, and by definition of (o(k). that (( - =i|l converges to O. Therefore, follows that Eœ(N) - 0, that = O. We then have shown that (Ex)x21 is bounded sequence and every subsequence of it converges to 0. Therefore, Ea must converge to O, which concludes the proof. D
ü. Using again Equation (15) and the fact that L is Hadamard differentiable,
C.2 Proof of the Functional implicit differentiation theorem
Proof of Theorem 2.1. The proof strategy consists in establishing the existence and uniqueness of the solution map w + hw deriving candidate Jacobian for it, then proving that + h is differentiable.
Existence and uniqueness of solution map w h Let in l be fixed. The map h Lin (w, h) lower remi-continuou since it is Fréchet differentiable by assumption. It is also strongly convex. Therefore, it admits unique minimier 16 [Bauschke and Combettes, 201 Corollary 11.17]. We then conclude that the map w + h is well-defined on Ω.
Strong convexity inequalities. We provide two inequalities that will be used for proving differ- entiability of the map w + ht The map + Lin (w, h) is Fréchet differentiable on H and μ-strongly convex (with positive by assumption). Hence, for all h1, h2 in H the following quadratic lower-bound holds:
Lon(w,h2) 2 Lin(w, + (i Lin (w, h1), (h2 h11) 2 +
11 || hu 1
(16)
2
From the inequality above, we can also deduce that h + On Lin (ui, h) is j-strroggly monotone operator:
LinLin(w, h1) - On Toin Lus 122), In - h2) 2 |// - h2ll Finally, note that, since h Lin(w,)) is Fréchet differentiable, its gradient must vanish at the optimum he ie:
(17)
Og (w, k()) 0.
(18)
Candidate Jacobian for - No. Let be in Ω. Using Equation (17) with yn h to and h2 h for some h, v E H. and a non-zeros real number f. we get:
(O, L (w, h + tv) - 0g Lin h), u)n 2 p ||=]|
(19)
By assumption, + Ōn Lin (w, h) is Hadamard differentiable and, a fortiori, directionally differen- tiable. Thus, by taking the limit when O, it follows that:
( (w., h)e,,)) 2 / ||0||
(20)
Hence, Lin (w, (w, h) H - H. defines coercive quadratic form. By definition of Hadamard differentiability, is also bounded. Therefore, it follows from Lix-Malgram's theorem [Debnath and Mikusinski, 2005, Theorem 4.3.16), that 0% Lin (w, /t) is invertible with a bounded inverse. Moreover, recalling that B_ = Li (w, h(t) is a bounded operator, its adjoint (Ba)* is also a bounded operator from to H. Therefore, we can define J -- (Bw)) which is a bounded linear map from Ω to H and will be our candidate Jacobian.
Differentiability of r hig By the strong convexity assumption (locally in w), there exists an open ball B centered at the origin 0 that is small enough so that we can ensure the existence of ft > 0 for which h + Lin (w + h, p-strongly convex for all E B. For a given e B, we use the f--stonng monotonicity of h + 0 Li (w + e, (17) at points hC + Je and hi to get:
Alll - 1 - Jell? S Lin (w + e, hi+ee) - O2 Lin (w + e, + Je)) (hate - n1 -Je)
= Lod (w - e, hi + Je), (hi-e h" -JE)
< Li (w + e, h+ + Je) I h - JEll
18

---

where the second line follows from optimality of hate (Equation (18)), and the last line uses Cauchy-Schwarz's inequality. The above inequality allows us deduce that:
1 1ih (w e, h + Je) /
h hu Jell
(21)
Moreover, since O Lin is Hadamard differentiable, by Proposition C.2 it follows that:
OnLin (w + e, thw + Je) = =ū, =ipLin (w, (w) +dl,,)) Lin (w, h (e. + o(lle)),
(22)
where the first term vanishes consequence of Equation (18), since hy is minimizer of h Lin (co, ft). Additionally, note that the differential d(w,,) Ôn (w h) 2 x H + H acts on elements (c.g) E Q x as follows:
dn On (e)))) 0 Lwn h)g (Ow,h hw)"
(23)
where 0f. (u, hw H + H. and Dn h)) H Z are bounded operators and (Ow, Lii (w, h))* denotes the adjoint of Ow. Lin (w, h). By definition of J, and using Equation (23), follows that:
(lrn)OnLise (w. h)) Je) = CwJe + B = O.
Therefore, combining Equation (22) with the above equality yields:
Lin (w + e, hy Je) o(Ilel)-
(24)
Finally, combining Equation (21) with the above equality directly shows that | h Jell s # [ We have shown that w h is differentiable with Jacobian map h given by J = -B
C.3 Proof of the functional adjoint sensitivity in Proposition 2.2
Proof of Proposition 2.2. We use the assumptions and definitions from Proposition 2.2 and express the gradient F(w) using the chain rule:
VF(w)= = Lou (w, h + [awhi]]
hū).
The Jacobian 0,h is the solution of a linear system obtained by applying Theorem 2.1
duhi" - CB
We note = Ow (w, and du = Or Loue ht). It follows that the gradient VF((w) can be expressed as:
VF(w) = 9w + 1w dw = Iw + B
a = - 1
In other words, the implicit gradient F(w) can be expressed using the adjoint function o which is an element of H and can defined as the solution of the following functional regression problem:
-
= arg min Lady (w, a) = (a, Caa) + (a, d=)
aEH
D Functional Adjoint Sensitivity Results in Lg Spaces
In this section we provides full proofs of Proposition 2.3. We start by stating the assumptions needed on the data distributions and the point-wic bint losses in Appendix D.1, then provide some differentiation results in Appendix D.2 and conclude with the main proofs in Appendix D.3.
19

---

D.1 Assumptions
Assumption on IP and Q
(A) IP and Q admit finite second moments.
(B) The marginal of X w.r.t. Q admits a Radon-Nikodym derivative r(x) w.r.t. the marginal of X w.r.t. P. i.e. dQ(x, 3) - r(z) dP,(,, J). Additionally, r(x) is upper-bounded by positive constant M.
Assumptions on iin-
(C) For any S2 there exists positive constant i and neighborhood B of for which lin is j-stroogly convex in its second argument for all (w', xx, E B x X' x V
(D) For any E Ω Ep
[km[[, O, v)1 + 110.kn (w, 0, x, u)] +00.
EE) u - (w,, v, Ft v) is continuously differentiable for all (w, v) E Ω x X' x y
(F) For any fixed w E Ω2 there exists a constant L and a neighborhood B of w s.t. v + lan (w', U, 22yyy y) is L-smooth for all w' x. B X y.
(G) (w, u) + (wl v, v, y) is continuously differentiable on Ω V for all x, y X x >
(H) For any w Ω2 there exists a positive constant C and neighborhood B of s.t. for all (w/, x. v) E B x x y:
eln (w'0, w)ll S C |||||| + llyl)
(25)
(I For any w E Ω. there exists a positive constant C' and a neighborhood B of w s.t. for all (wl, Ui, U2, U2,
x V x X x y we have:
1i (w', ,1, v) Sw,nlin (w/, 02, 22, v)ll S C Cle -2ll-
(26)
Assumptions on Lout.
(J) For any w E Ω, EN (w, y] < +00.
(K) (w. v) + low U, x, y) is jointly continuously differentiable on Ω x V for all (z,y) e
X x V.
(L) For any w E Ωt there exits neighborhood B of w and a positive constant C' s.t. for all (w', u, x, B
X we have:
[ (w' - Botu (w/, ,'' v) < (1 + || + ||el + || + IIvll) /ee w'll.
|10goout (w', (,,,, doluu (w',v', < C || v'll,
10 (W' < C (1 + ||인||| + ||||||| + l|>/l) Lott (w',u,z, u. v) s C +++ ||a|| + | [lell")
Example. Here consider the squared error between two vectors v, = in V. Given map (w, xr y) fz y) defined over Ω x X' x y and taking values in V. we define the following point-wise objective:
1 |00 211
f(w, v, v, =
(x.
(27)
2
We assume that for any w E Ω, there exists constant C > such that for all w/ in a neighborhood of w and all z, X x V. the following growth assumption holds:
fl (z, r) + [ow (x. v\ll s C(1 + ||^||| + llvll).
(28)
This growth assumption is weak in the context of neural networks with smooth activations as discussed by Binkowski al. [2018, Appendix C.4].
Proposition D.I. Assume that the map + (a y) is continuously differentiable for any x. E X x Y, and that Equation (28) holds. Additionally, assume that P and Q admit finite second order moments. Then the wisnt-wise objective E in Equation (27) satisfies Assumptions (C) to (L).
20

---

Proof. We show that each of the assumptions are satisfied by the classical squared error objective.
Assumption (C); the squared emror is 1-strongly convex in v, since 0ee = I. Hence, the strong convexity assumption holds with /4 =.
Assumption (D): For any w E S2 we have
Ep [ O, x, w) + | (w, O, æ. w)ll = 1PP Ilf (2, + [I w)]l < t00,
which holds by the growth assumption on fww (zx, y), and having finite second moments.
Assumption (E): With a perturbation u E we have:
1 ||e =||² + (v u) + o((lul)), = fw,
flw, + u, v)
2
with o[[ = || 1|2 The mapping e + e continuous, thus the assumption holds.
Assumption (F): For any two points 01, 0/ E V using the expression of a, v, v, v 2 with (x, y) we have:
[0,6(w, u1 x. v) O (ww 22 v)ll = [[(r - 2) - (u2 - 2)1l
101 021l,
We see that e is L-smooth with 1 and the assumption holds.
Assumption (K): By the differentiation assumption fa v). with perturbation Z we can write:
fatela, y) = fa, y) (z. u)c o(e).
With perturbation 9 and substituting fwte (2, y) with the expression above we have:
2 (v + - + +L o(e)]]
||u fu( (c, f (a,8) (f
+ (u, v (z, 3) o(| + ||u])),
which allows conclude that (w, v) f(w, U, x. y) is continuously differentiable
Ω x for all x, y E X' x y and the assumption holds.
Assumption (G): With a perturbation x u E Ω x V using the expression of 0,E(w, v, 2, write: 0.l(w + t, u. v. = (u + u) - (x,
(v + u) - (fu(x,)) + fa +) = (v fw (z, v)) + ( y) o(c),
by continuously differentiable fx, y). we have that the assumption holds. Assumptions (H) and (I): From the expression of f.e(w, v, u, 2. y):
o l (w, u, x. v) O. (v (x) fw y);
then using the expression above and the growth assumption for (x, y) we have that the two assumptions hold.
Assumption (J): For any w E Ω we have:
Eg [e[(w, 0, vill Eg [l [I w|1l < +00,
by the growth assumption on [w y), and P having finite second moments, thus the assumption is verified.
21

---

Assumption (L): Using the growth assumption on fwr (z, we have the following inequali-
ties:
1 (w', U, x, - d.e O (w', 1,, v)1l
for (z, v) (v' -m)l
- for (I, 3) 7 + --lle' =
< C (1 + ||||| + |1000l + ||||| + l1lll) 110 uoll
10,,('', e, x. v)ll = - (2,
$ || (x,,-(z,,)))
s || + CC +
llyl)
| (w', 2,,,
@awf (2, (fw-(x,)) v))
- 1 (2, v) [[] v) +)ll
< / (x, w)1 (Ilfwr (x, y)ll + [[lll)
SC (1 [1+ ||a|| + ||=|| + -llll")
combining the above with L-smoothness of we can conclude that the assumption holds.
D.2
Differentiability results
The next lemmas show differentiability of Lout, Lin and Li and will be used to prove Proposi- tion 2.3.
Lemma D.2 (Differeetiiability of Lin in its second argument). Under Assumptions (D) to (F), the function h + Lin(w, h) is differentiable in H with partial derivative vector 2 Lin (w, h) E H. given by:
Or (w, h):
+ Ep Jolin (w, h(x),(())))))))),) x, 8) [z]
Proof. We decompose the proof into three parts: verifying that Lin is well-defined, identifying a bounded map as candidate for the differential and showing that it is the Fréchet differential of Lin.
Well-defined objective. Consider (w, h) in Q2 H. To show that [nn(w, h) well-defined, we need to prove that (n,(w, h(x), v, is integrable under IP. We use the following inequalities to control lin (wo, hxx), x, y): (a, h(x), v, v)1 Cm h(z), la (w, 0, 8) + 1 (w. 0, 8)
L' at
+ (w, 0, v))l
||((2)|| //' dt (11000mm (w, th(x), 9) 0glin (w, 0, 0,
+ [[((x)] In (w. 0. v) + |n (w. 0, x, x)
2 [|(())]] + (11(=))1 | (w, O, x, w)ll") ee (w. O, x. v)l
L
where the first line follows by triangular inequality, the second follows by application of the fundamen- tal theorem of calculus since fm is differentiable by Assumption (E). The third uses Cauchy-Schwar inequality along with triangular inequality. Finally, the last line follows using that ei is L-smooth in its second argument, locally in and uniformly in x and y by Assumption (F). Taking the expectation under yields:
Lin (w, Ab) SEp [lem (w, h(xx,(())))))))). 2)
L+1
w01l + (w, 0,2, < too, || + Ep [[ 0, (w,
2
22

---

where 1 is finite since h E H. and expectations under IP of [|0gesn (w. 0, 0.z.y)] y)ll and (ein(w,,,) 0, ar [.) are finite by Assumption (D). This shows that Lin (w. ft) is well defined on 52 H.
Candidate differential. Fix (w, h) in Ω x H and consider the following linear form dia in H:
d == Ee [xx) 0. (w. h(x), xr 2))
Vg H.
We need to show that it is a bounded form. To this end, we will show that dın is a scalar product with some vector Din in H. The following equalities hold:
dmn = Ep [g(x) O, (wo, h(x), 2, y)]
Er [o(x) Ep (0, Lin (w, ((z),
Er [o(x)) Dm](1]]
where the second line follows by the "tower" property conditional expectations and where we define Din : Ep (n (w, h(a)), x. y)]2 in the last line. Din is a the candidate representation of din in H. We simply need check that Din is element of H. To see this, we use the following upper-bouue opper-bounds:
Ep [x)][
s EE>
Ez [10 (iOpein (w, 2,
E. [10 h(x) v|
s 2Fp [10glin M(), v) 0, (w, 0, w)]l 2Rp (w, v))l
s 22L [1b(z)1 + 21p (w,
+00.
The first inequality is an application of Jensen's inequality by convexity the squared norm. The second line follows by the "tower" property for conditional probability distributions while the third follows by triangular incquality and Jensen's 's inequality applied to the square function. The last line uses that ln is L-smooth in its second argument, locally in and uniformly in x, y by Assumption (F), Since h is square integrable under P by construction and [i (w, O, x. y) is also square integrable by Assumption (D), deduce from the above upper-bounds that D (x:) must also be square integrable and thus an element of H. Therefore, we have shown that din is continuous linear form admitting the following representation:
dnn (Dins 9)n.
(29)
Differentiability of h Lin (w, h). To prove differentiability, we simply control the first order error E(g) defined as:
F(g)) TTin (w, h + 9) Trin [in(w,,)) dingl-
(30)
For given H. the following incqualities hold:
E(g) = Ee L/ dt (g(x)T (Wols ((g)
(x, On ((2))) v)))
Lo Io(x))
Ep
(0,lin h(x) tg(x),,,,)) v) Jolin (w, (xr), v)) dt
T. Br [|[[[))]]]
I 2 lla1
5
where the first inequality follows by application of the fundamental theorem of calculus since lir is differentiable in its second argument by Assumption (E). The second line follows by Jensen's inequality while the last line uses that / de., (w, v, x. y) is L-Lipschitz locally in x and by Assumption (F). Therefore, we have shown that E(g) = of o(119ll which precisely means that + L,n (w, h) is differentiable with differential d Moreover, Dan is the partial gradient of Lin (wo, h) in the second variable:
and uniformly
(, h) Di = Ez 10 (w. h(z), v)(x]
23

---

Lemma D.3 (Differentiability of Lout). Under Assumptions (A), (B) and (J) to (L), Lout is jointly differentiable w and h. Moreover, its partial derivatives d.L h) and On (w, h) are elements in St and H given by:
Tout h.) =Bo (ww (w, h(x)), 2, Ö (w) h) =x + r(x)Eg (0 Lont (wo, h(x)) 2, y)[=].
(31)
Proof. We follow a similar procedure as in Lemma D.2, where we decompose the proof into three steps: verifying that the objective Lout is well-defined, identifying a candidate for the differential and proving that it is the differential of Lout.
Well-definiteness of the objective. Let h) be in S H. First, note that by Assumption (B), have that
Ee [|(0a] = Ep [In(a)) r(()) s M |1ilii < +00.
(32)
The next inequalities control the growth of lout: |oau 1(xx), v\ < (ws 0, v) + Mou (w, (x)),,,))) lonelw, (, Hoo (w, 0, [. dt (allo))olllllll (w, th()),
oou (w, O,,,0,,.))) +,),0, [|(())] /-" dt [10, th(x)) y>ll
< oot (w, 0, + [|x(=)]] (1 [[+())) II + ||=||| + |!vl))
3 1 (w, 0, + C (1 + ||<=)]1 + ||0|| + ||vl
The first line is due the triangular inequality while the second line follows by differentiability of Oolout in its second argument (Assumption (K)). The third line follows by
inequality wile the fourth line uses that lout has at most linear growth in its last three arguments by Assumption (L). Using the above inequalities, we get the following upper-bound on Louti
Loou b)]
Eq [lout (w, 0, x. v) C (1+ + 3BQ [[((=)]]] + Be ||==|| + I1be1()) +00.
(33)
In the above upper-bound, Eq [[ is finite by Equation (32). Additionally, Eq ||=]|] + [mvl]]
finite since Q has finite second moments by Assumption (A) while Ie [looue (w, O, xe, y)]] is also finite by Assumption (J). Therefore, Toau is well defined over Ω x H.
Candidate differential. Fix (w, in Ω x and define the following linear form: doue (c,e((,,0) g) := c =E loueout (w, h(x), x. v)1 + Eq [9(x)) Opfoue(w, h(x),
Define Dont (Dw, Dn)) be:
Dw == Eq (8 Lost (w, (()),,,1)
Dr = x + ((x) 1ol (w, h(z),
By an argument similar to the one in Lemma D.2, we see that doue (e, g) = (g, DRY + E ED.. We now need to show that D., and D;, are well defined elements of Ω and H.
Square integrability of Dr. We use the following upper-bounds:
Ex [
< Ep
[e [11, (w, h(x), x. =)
< Ep
TE [1 Mxx), v)
< MEp
([ [10,four(w h(xx)) w)
ME4
[1a [10. (w, h( x,, m)
(1 + te [11(0a)11]] + Be ||=0|/ + [I]l(])
< AMC
24

---

The first inequality is an application of Jensen's inequality by convexity of the norm, while the second one is an application of Cauchy-Schwaad -Schwarz inequality. The third line uses that r(z) is upper- by constant M by Assumption (B), and the fourth line follows from the "tower" property for conditional probability distributions. Finally, the last line follows by Assumption (L) which ensures that 0, has at most a linear growth in its last three arguments. By Equation (32), we have that
Eq [[b(z)]1]]
< +00. Moreover, since Q has finite second order moment by Assumption (A), we
also have that E2 ||0=|| + Ilvll"] < +oc. We therefore conclude that Ep [x)]( is finite which ensure that Dj. belongs H.
Well-definiteness of Du. To show that D. is well defined, we need to prove that (xr, y) + Eot (we h(r),, x. v) is integrable under Q. By Assumption (L), know that Oshout has
most a quadratic growth in it last three arguments so that the following inequality holds.
[or (w, h(x), x, w)ll S C [|{)))]] + ||:|| [lvl*]
We can directly conclude by taking the expectation under Q in the above inequality and recall- ing that 12 [|](=))1]] is finite by Equation (32), and that Q has finite second-order moments by Assumption (A).
Differentiaaility of Lour. Since differentiability is local notion, we may assume without loss of generality that [|e|] +1ggl+ I 1. Introduce the functions A and Ag defined over H2
as follows:
(1 9, v, Oylout (w te, h(x)) tg(x), v) Oplout (w te, ((x),
(i g, x. v. Oplout (w+ te, h(z),,,)) 0ylout h(x), 2,1)) A2 9, v. U. v,,, 1) = lone (w Le, h(x) (g(())((()))))))(())))))))))[[ 2, O lous (w + Le, 2,,) A 9, v. 1) = lone (w + Le, (t)))(())).))) (r, Eat (w, (1x), x,
We consider the first-order error F(e, g) which admits the following upper-bounds:
F(c,g) = Trout(w (w + t. g) Trout (w, h)) dont (5,
EQQ 1 dt (g(x) + Ai)e [A22 &.(e,
<Ea Lh di
+ Eo Lo'
<Fa [10(0))1]]
EQ Jf di (I
+ +Ee [[. di 1 J. (1
Ai(e.))
Asle.))
EQ 1/' dt (.A2(c,g.2.,0 v1
+ +Eg I// dt (lllA2(e,...
+ +|llll
Asle.g)
Asle.))
<M b1ll (A(cc.))+AA(s + ||l|| (As(c,g) + As(c,g))
The second line uses differentiabilityy of eut (Assumption (K)). The third uses the triangular inequality, while the fourth line uses Cauchy-Schward inequality. Finally, the last line uses Equation (32). We simply need to show that each of the terms A1, A2, Ag and As converge to 0 as and y converge to O. We treat each term separately.
25

---

Controlling A1 and A3. For e small enough so that Assumption (L) holds, the following upper- bounds on A, and A2 hold:
A1( SCIE [1
SCM² & l1al
9s(r,g) a [11 + [1b(x)) + tg(x))] + [|(xa)] + ||->|| + lllll) [I
sCe [1l((>21]] $ (1+ [B 2Bg [1((=)]]]
+ Bo [I
+ Eq ||0l| -Iwll]]
M SCM² 1|all (1+M²= + 28 [
+ Fe [||||| [[]]]]]
For Ai, we used that Sylout has is Lipschitz continuous in its second argument for any x, y E X x y and locally in by Assumption (L). The second upper-bound A1 uses Equation (32). For A2, used the locally Lipschitz property of 0,t from Assumption (L), followed by Cauchy- Schhw- inequality and Equation (32). For the last line, also used that ||||| I 1 by assumption. The above upper-bounds on A and A3 ensure that these quantities converge to 0 as e and approach O.
Controlling A2 and A2. To show that A2 and Aa converge to O, we will use the dominated convergence theorem. It is easy to see that A (c, g. 2. y, and D/2 (c, I. y, converge point-wise to 0 when and 9 converge to since (w, v) + Lo U, v) and (w, v) + 0u (wi v, x, v) are continuous by Assumption (K). remains to dominate these functions. For small enough so that Assumption (L) holds, we have that:
A (e, S16C2 (1 [|((2)]] + ||0|| + Ilull")
(c <2C [1(x)()) + ||>|| + -1e11)))
Both upper-hounds are integrable under Q since F. [I)()))]] < +0 by Equation (32) and Q has finite second-orde moment by Assumption (A). Therefore, by the dominated convergence theorem, we deduce that Ag and A₁ converge to as and g approach 0.
Finally, we have shown that F(c, = o (Ilcll + (1gllx) which allows to conclude that Lout is differentiable with the partial derivatives given by Equation (31).
Lemma D.4 (Differentiabilityy of On Lin). Under Assumptions (A) and (E) to (I), the differential map (w, h)) + On Lin(w, h) defined in Lemma D.2 is differentiable on Q x H. in the sense of Definition C.I. Its differential di by Lin (w, h) Q x H H acts on elements (e,g) Q H as follows:
dis On Lin (w, (), = Lan(ww,h) + (0w, (in,
(34)
where Li (w. h) 71- H is linear symmetric operator representing the partial derivative of On Lin(w, h) werl h and (a.. Trin hi)" is the adjoint of Qw., Tun (w.,)) h) H + Ω which represents the partial derivative of Ūg Lin (wo, h)) w.nt w. Moreover, L2 Lin (w. h) and 0w. Lanhin (w,h) h) are given by:
Lg (w, h)g =x + Ee [olla h(x))))))))))))))) 9(x) O., Lin (wi hig EE lin (w, ((x), v))()))
(35)
(36)
Proof. Let (w,, fi) be in Ω x H. To show that a., Lin Hadamard differentiable, we proceed two steps: we first identify candidate differential and show that is a bounded operator, then prove Hadamard differentiability.
Candidate differential. For given (w, h) E Ω x H, we consider the following linear operators Ω:
C..ph H - H and Bu,
Cw,rg= = Ex [n ((x), x, wr|=| g()),
Bu,ng = Ex Veu, (wi h(xx)))))))))()z)) v)9())}
where the expectations are over y conditionally on æ. Next, we show that Cw,b and Bag are well-defined and bounded.
26

---

Well-definiteness of the operator Cw.h. The first step is to show that the image C of any element g E by Cuuh is also an element in H. To this end, we simply need find a finite upper-bound on ||Cu,nolxx for given g E H:
|[C Rt = Ee
Ee [ (w. h(x), 2, v) ((x))
< Ep
Ep [o (w. M(x), x v) [x] I. [1(())11²
Ep
Ep |1 (w,
1lop xr 11
2 Ep
[|06 (w. h(z)) x, 2,,) v [=(
1 |191
The second line follows using the operator norm inequality, the third line follows by Jensen's inequality applied to the norm, while the fourth uses the "tower" property for conditional distributions. Finally, the last line uses that of is upper-bounded bounded uniformly in x and y by Assumption (F). Therefore, we conclude that Cr. belongs to H. Moreover, the inequality || |CC,,lll S L l1g1l also establishes the continuity of the operator Cw.h-
Well-definiteness of the operator Bu,k. We first show that the image Bu, is bounded. For a given g in H. write:
[Bw.ng] = ||Bp [w,, (w. 2()),,,)))(
<Ep [1I e Ca (w, v))x)il
< Ep 110,
< 11 1E>
10
v)
s ||ll| E8>
10
Ou.ulin( 0,,,,
+ -1lll Er
110 (w, 0,
(1 + ||~|||| + l1el1)]+))
<C11ll
Ep
[11(())1 + Egg
SC 11
(1llllee + 2
[1 ||a|12 -lvll) < +00.
In the above expression, the second line is due to Jensen's inequality applied to the norm function, the third line follows from the operator norm inequality, while the fourth follows by Cauchy-Schwarz. The fifth line is due to the triangular inequality. Finally, the sixth line relies on two facts: 1) that ++ dn (w. v, xr 3) is Lipschitz uniformly in x and y and locally in w by Assumption (I), and, 2) that ean (w. 0, x. y)ll has at most linear growth in z and locally in w by Assumption (H). Since has finite second order moments by Assumption (A) and both h and g are square integrable, we conclude that the constant || Baa ll is finite. Moreover, the last inequality establishes that B a continuous linear operator from to Ω. One can then see that the adjoint of Bach admits representation of the form:
(Bw,a)* (o, e x + Ep (oouulin (wo, h(x), x v)) v)) [
Therefore, we can consider the following candidate operator de, for the differential of 0 T.mn:
de (e. g) Cu,ng + (B.,
Differentiablity of On Lin. We will show that On Lin jointly Hadamard differentiable (w, h)) with differential operator given by:
Ow (ww (t)(c, Cw,ng + (B..
(37)
27

---

To this end, we consider a sequence (ck, 9s 221 converging in Ω x H. towards an element (c, g) E Q x H. and non-vanishing real valued sequence t, converging to Define the first-order error Ex as follows:
Es - -/222 (OnLin(w + LEEN h + lagk) - On Lin (w; - Cw. - (Bu,a) 1
1x
Introduce the functions P1, P2, A and Ag defined over N X x y/ [0, 1] as follows:
Cnlin + slEEk, h(xx) + slagk (l (,,,, 0glin (w> ,,,),,,
A: 2
A: = 0
(Ow, (w + styck, h(x) P2, (Ow,olin (w, h(x),
k 1 k= 0.
(,, P (K,,,, v, P1 x. 8) A2 2, Pa(k, x, 2, 8) P2(0, x, y, 8)
By joint differentiability of (w, + Oplan (w. v, w. y) (Assumption (E)), we use the fundamental theorem of calculus to express Eg in terms of Ai and Ag:
Fg
Er
[a 2.l.::/00000.0000....00--------0 x
Ep 1/" dt
< 1Ep
L/ dt (kk
= -Ep
Pi 20. + Ps (k. 2,y, P2
%/' || (k, x, y, 1) ||
+4 +4||e|| Ep Lo di 11 x, v. 1)
<4 1p
B
+ Ep [1" di [|| (k. vs 11 (g(x) (x))
ex EEp [ di || v. 11
+4 +4
The second line uses Jensen's 's inequality applied to the squared norm, the fourth line results from the "tower" property of conditional distributions. The fifth line uses Jensen's 's inequality for the square function followed by the operator norm inequality. It remains to show that a and A(
converge to 0 and that (22) bounded.
Upper-bound on A We will use the dominated convergence theorem. Assumption (F) ensures the existence of positive constant 7, and neighborhood B of w so that + || (w', v, v) op bounded by for any ,,, E B x x y. Since w + 1ĄεR - wi then there exists some Ko so that, for any k 2 Ko. we can ensure that w leeR E B. This allows to deduce that:
| (k, x. y, t) 2 op [=)())) < 4L 11
(38)
for any l: 2 Ko and any x, E X x V, with |= Moreover, we also have the following point-wise convergence for Paalmost all x E X::
2 being integrable under P.
A x. y, 11 [9(x))] O.
(39)
Equation (39) follows by noting that w + thEk + and that h(x) + t&gu (z) + h(x) for Plalmost all E X, since ta converges to 0, e converges and 9k converges g in H. (a fortiori converges point-wise for PP-almost all x E X²). Additionally, the map (w:,) + ||aßen fn (w, v, 21,,) is
28

---

continuous by Assumption (G), which allows to establish Equation (39). From Equations (38) and (39) we can apply the dominated convergence theorem which allows to deduce that A + O.
Upper-bound on A By a similar argument as for A and using Assumption (F), we know that there exists Kg > 0 so that for any k 2 Ko:
(K y, lop - L2
(40)
Therefore, we directly get that:
2 S 1,2 11g())))) (x(()) + +0,
(41)
where we used that gk 7 9 by construction.
Upper-bound on B We will show that || at, y, 1)lp upper-bounded by square integrable function under P. By Assumptions (H) and (I), there exists neighborhood B and a positive constant
C' such that, for all w', u1, 02 V V E B x V x V x X x y:
10 lan (w', 0, w)ll - C (1 + |||||| + Illl
(42)
D (w', v. v) - ein (w', 02 n. w)ll & C [|e] - 12ll
(43)
By a similar argument as for (1] there exists Ko so that for any k 2 Ko. the above inequalities hold when choosing w' w + tEER. Using this fact, we obtain the following upper-bound on PP (k, a, y, t)p for k 2 Ko:
i
Il (,, 8) S @ww (w + StrEk, h(x) stuggu (x), v) Ow.eeilww + strEk. x, v)
+ 100, (w + StREK, 0, x, 1>,,,,,,,
50 (1 + [1(x)) + st.0 (x) + |||||| + Ilvll)
00 (1 + [=<]=|]]||]]0]>>> + tu [x) + ||^||| + llvll)
Therefore, by taking expectations and integrating t, it follows:
B < C
(1 + [A(x)] + 11 (z) + |||||| + [1xl))]]
s Er (1 + || + 2 (x)] ||0|| + llvvl*)]
By construction 12 (x)] + 0 and is therefore a bounded sequence. Moreover, E3 [1<(z)) I1] is finite since belongs to H. Finally, Ep
||<|| + [|vl["] < +oo by Assumption (A). Therefore, we
have shown that (222 is bounded.
Upper-bound on B By similar argument as for B(a) and using again Assumptions (H) and (I), there exists Ko so that for any & 2 Ko:
|| x. y, olop S lo stuck, h(x) + stagk v) Ow (w stack, h(x), x, v)
+ en + styck, h(z), x, dr, (w, h(x)))()))))))))) x, y)pp
<Ct/ (1l + 110w, (w st&ck h(z)) x, Ow M(x), x, 4) 1)pp
where used Equation (43) get upper-bound on the first terms. By squaring the above inequality and taking the expectation under P we get:
D 2Ctx 11
2E
[las,lim(w + st (x),
M(x), x. v) 1
(44)
(x(2))
We only need to show that Ep [ex (ar, y)] converges to 0 since the first term 2C1s |1 l already con- verges by construction la and gk. To achieve this, we will use the dominated convergence theo- rem. It is easy to that en (z, y) converges to 0 point-wise by continuity of + Ou (w, v, 2,, v)
29

---

(Assumption (G)). Therefore, we only need to show that (x y) is dominated by integrable func- tion. Provided that k 2 Ko. we can use Equations (42) and (43) to get the following upper-bounds:
1
e (z, v) - ( stack, h(x), 2. v) ww + StyEk, O, 2, 1l 4
+ 110 ((2))(()))))))))) 2, y) O (w, 0, 2, 2, y)
+ 11 (w, O1 + 100nriin StrER, 0, x, x) 1)
$2C² (1 + [|][[[](]] ||||| + llvl²))
The Lh.s. of the last line is an integrable function that is independent of k, since is square integrable by definition and ||x/|| + |lyl1 are integrable by Assumption (A). Therefore, by application of the dominated convergence theorem, it follows that Ep (2, v)] 0, we have shown that B - 0.
To conclude, we have shown that the first-order error E. converges which means that (w. h) H> Ön Lin (w, h) is jointly differentiable on f2 x H. with differential given by Equations (34) and (35).
D.3 Proof of Proposition 2.3
Proof. The strategy is to show that the conditions on Lin and Lout stated in Proposition 2.2 hold. By Assumption (C), for any w Ω there exists a positive constant and neighborhood B of w which the function Lin (w/, U, y) is j-sttroggyy convex in for any (w',x,)) B x X V. Therefore, by integration, we directly deduce that f + Lin(w', ,h) is μ strongly convex H for all w E Q and θ₂ Lin
for any
w' B. By Lemmas D.2 and D.4, h + (w, h)) is differentiable Hadamard differentiable on x H. Additionally, Lout jointly differentiable in w and h by Lemma D.3. Therefore, the conditions notations from Proposition 2.2, we have that the total gradient V.W)) can be expressed as:
Lin and Lout for applying Proposition 2.2 hold. Using the
VF(w) = Iw B
(45)
where 9u- = Low ht), B., Owin Lin (w, hi) and where at is the minimizer of the adjoint objective Ladj:
Ladj (w, a) = a C a dw:
with = Li h and du = 2g Lout (w, h.). Recalling the expressions of the first and second order differential operators from Lemmas D.2 and D.4, we deduce the expression of the adjoint objective as a sum of two expectations under P and Q given the optimal prediction function
Ladj (w, a) = ((
(x)) y) a(x)]
E [a(z) Oolout (w, (z) x y)]
Furthermore, the vectors Iu and B appearing in Equation (45) can also be expressed as expecta- tions:
I- (a (Oslout (w, (x), xr, y)]
a E 10,, (w, hi (z), x, v) a (z)]
E Convergence Analysis
We provide a convergence result of Algorithm to stationary points of F. Our analysis uses the framework of biased stochastic gradient descent (Biased SGD) [Demidovich et al., 2024] where the bias arises from suboptimality errors when solving the inner-level and adjoint problems.
30

---

E.1
Setup and assumptions
Gradient estimators. Recall that the total gradient V.Fww) admits the following expression under Assumptions (A) to (L):
VF(ww) = Eg (d, lout (w, (2), x, y)l + Ee Ouwpiin (w, hw), , x, 8) (2)),
We denote by ((w) the gradient estimator, i.e., the mapping 9 Ω computed by Algorithm and which admits the following expression:
((w)
Σ
Ouloue (w). +
Σ (r.
Ou,ulen (w, hw, (z), (x), 2. a (x),
Bouel
(Bot)
where Bin and Bout are samples from P and Q independent from hi and å and independent from each other (i.e. Bin 1 Bout). Here, there are three independent sources of randomness when computing g(w)): estimation of hnw and a., in Algorithms 2 and 3, as well as random batches B and Bott- We denote by E[-] the expectation with respect all random variables appearing in the expression of ((w), and by [8]]kh]] and E[- hhw d the conditional expectations knowing hh only and both hi and âw- j(w) is a biased estimator of V.F(w) (i.e., R[((w))] is not equal to V.F(w)), as the bias is due to using sub-optimal solutions hi and i instead of h and a in the expression of ((w). Furthermore, we define G(w) :- E ((w)))/ aw to be the conditional expectation of g(w)
given estimates hiw and aw of htt and a By independence of Bin and Bout from h and a.. G(w) admits the following expression:
G(w) =BQ (oufoue (w,) 22 v) Ep (ww,eln (w, (w), 8) a (v(
where the expectation is taken w.r.t. (x.v) Q. Approximate adjoint objective.
We introduce the approximate adjoint objective Tnaet (wo a) where h is replaced by Awi
1 Las (w, a) za 00, To Lii (w, hw)a + a Or Lout (wo, hww.
(46)
By independence the estimator h and the samples used for computing Lady (w, a, har B) it is
casy to see that IE [Lay(w, (w, a, hww, B3) Pew = Lads (w,a), a). Hence, it is natural to think of aw as an approximation to the minimizer ae of Lady (w, a) in H. Sub-optimality assumption.
The following assumption quantifies the sub-optimality errors made by hw and aw.
(a) For some positive constants Gin and Eadj and for all S2 h and â satisfy:
E [wn hw) Lin (wn hi) S Gin
E
[oad (w, aw) Lods (w, aw)] I Cadj-
Assumptions on Lm.
(b) (Strong convexity) l. (wo, v, x, v) is p-strongly convex in E V. (c) v + 8. (w, u, v, y) is Ci-Lipschitz on St x X' x V.
(d) V + au lan (w, U. x, y) is differentiable and Cg-Lipschitz and bounded by B₂ E R.
Assumptions on Lowt.
| Lout (we, h Mh is bounded by a positive constant B3-
(f) r Opbout (w. U. x. y) is Ca-Lipschitz for all (w, xr, y) Z X x y
+ Lod (w, v. is differentiable and Cy-Lipschitz for all (w, xr v) E 9 X X
31

---

(h) (z. v) + Owtout (w,, (2),,, 2, y) has variance bounded by out for all w E Q.
Smoothness of the total objective.
(i) Function F(w)) is C-smooth for all E Ω and bounded from below by F ER.
Remark E.1. Assumption (a) reflects the generalization errors made when optimizing Lin and Luadt using Algorithms 2 and 3. In the case of aver-parameterized networks, these errors can be made smaller by increasing network capacity, number of steps and the number of samples JAllen-Zhi et al., 2019, Du al., 2019, Zou al., 2020].
Remark E.2. Assumptions (b) (h) are similar in spirit to those used for analyzing bi-level optimiza-
tion algorithms (ex: [Arbel and Mairal, 2022a, Assumptions to 5]). In particular, Assumption (e) even weaker than [Arbel and Mairal, 2022a, Assumptions 2] where Ō₁ Loue (w, h)) needs to be bounded for all w, h in Ω x H. For instance. Assumption (e) trivially holds when Lo (w, h) is linear transformation of 0n Lin (w, f)), as the case for min-max problems. There Lp (w, (0, so that 0 out (w, kw) = 0 is bounded, while Oo, h) might not be bounded in general.
E.2 Proof of the main result
The general strategy of the proof is to first show that the conditions for applying the general conver- gence result for biased SGD of Demidovich et al. [2024, Theorem 3] hold. We start with Proposi- tion E.3 which shows that the biased gradient estimate ((w) satisfies the conditions of Demidovich al. (2024, Assumption 9].
Proposition E.3. Let Assumptions (a) to (h) and Assumptions (A) to (L) hold.
1 ||F[(W))
VF(w) TE[((W]
2
2 (CI Gin + C26out)
(47)
2
E [[b(W)]]] S Lo + (c(rin + c2Cad) +2 [V
(48)
where and are constants defined in Equation (50) and og is a positive constant defined in Equation (51).
Proof. We prove each bound separately.
Lower-bound on V.F(w) V.F(W) [[](]]] Fist note that [[ğ(w)] B[G(w)), Hence, by direct calculation,
we have that:
VF(ww) [[j(w)] =vw( T
= (IVF(w)) E [[G(W)]]] - 1e [iG(w] V=(WI)I] [| (ciein + o2fout)
where we use Lemma E.4 to get the last bound. Upper-bound on 1[
||( = [10(w] G( + + [[G(w) - VF(WW) + V=(W)]]]
+ 2E [[G[w) [F(W))1 + + E [1((W) - G(w)]|]
oog + (cIem o2cad)) + [
where the first line uses that B[g(w)] E[G(ww) and the last line uses Lemmas E.4 and E.5.
We can now directly use Proposition E.3 and Assumption (i) on F to prove the Theorem 3.1 using the biased SGD convergence result [Demidovich al., 2024, Theorem 3].
Proof of Theorem 3.1. The proof is direct application of [Demidovich et al., 2024, Theorem 3] given that the variance and bias conditions on the estimator (w)) are satisfied by Proposition E.3 and that F is C-smooth and has finite lower-bound F R by Assumption (i).
D
32

---

E.3
Bias-variance decomposition.
Lemma E.4 (Bias control). Let Assumptions (a) to (g) and Assumptions (A) to (L) hold.
E
[[G[[]) ((G(W)-V[(W)II" S CiEin + C2Cadj
(49)
where 6 and 02 are non-negative constants defined in Equation (50).
Proof.
G(w) - VF(w)) = Ba loue (ws, huw (1) - Owlour 2,
A1
+ Ep [owplin (w, (w))), x., v) @w(xx) Owrolin in (w, (x), x, v) a1
- + Ep (ounnlin (w. (x)), x. Ownnei (4s, (1(),,,,)))) oz (2)
+ Ep /owolin (ws, hu(x),, x. v) (aw az (1(2))
We have the following:
|| SMCss M hi l
|| <C² ||ailxx how hi hill S C2ByNr &C2Bs/+ Il
h1 #
||A < Bp [llāw (x) - (1)()]l]
<2 (Er [1lāw(()) (x))l + EE (1lad (x)
<B (1 āwll + lق at (,,).,
The first inequality holds since + Owlout (w, u, v) is C3- Ca-Lipschitz by Assumption (g) and ()) admits a density w.r.t (P)())) bounded by a positive constant M by Assumption (B). The second inequality holds since u + Ow (w: x, z) is Cz-Lipschitz by Assumption (d) and | H is bounded by μ=1 wu'Bs "3 as consequence of Lemma E.6. Finally, the inequality on || holds since vu (wo v, 2, y) bounded by constant Bg by Assumption (d). Therefore, holds that the difference between G(w) and V satisfies:
C2B3)) P
ll
| H + +3B; 1ق 2w1l + 3B² l1ā atll
[[G(w) [ 3 MC3+ +
h
Taking the expectation over hnw and ân and using the bounds in Lemma E.6 yields:
e [IG[(w) [=(W))]]] <6 (MCss C2BsH 1)2 p
ein + +6B2/ Codj
+ +B ( + ,C Ein.
Finally, the upper bound on the bias holds with e₁ and 02 defined
= 6 (MC's + CBB p- + +6B3p (H-'C, C,M +
BBC))² 2
= 6B3,
(50)
Lemma E.5 (Variance control), Let Assumptions (a) to (h) and Assumptions (A) (L) hold, then the variance is upper-bounded as follows:
E [1ا(w) G(w))l]] < Sog
where dn is a positive constant given by:
4B3
0 = Boutl
(2C 1C - Ein + )ou)) +
(p- E Eadj + + 30G22 em + x=² (51)
B
33

---

Proof. By definition of ((w) and G(w) we have that:
[l((w) [(w)ll
1
E
@ouloue (ws (u))))))) Ee
@wlous (ws, huule))) hrs
Bout
vi @uroem (wo hull)), x, v) )m(سس) |
1
E
Bin
De Owrulin (wi (x), 2) ( aw ]1
Bin
V1
1/2
+
Bout
B
[Binl
Where the first line is direct consequence of the independence of Bin and Bout. Moreover, we can bound V1 and 1/2 as follows:
| (w. (w, x.), Oulout (w, (w,n (x)) v) [1]
Vi 2
[1a (w, Mo (2), ar 2,, E2 lowooue (wi 1o (2'), of // )
+ 2
||| [ - -hh--h% /| 20out I 403// (03/1-'co can + 20
<20
1/2 <BEE [| 112c] $2B? E |1 20llit] + IE ||قق 11
<4B² (p 1 Eadii 2p 3 30 M² Ein 14 (=3)
For V1, we used that v + boue (wl, v. v. y) is C3-Lipschitz uniformly in w, x and y by Assump- tion (g), and that (x. y) Oyhoue (w, (z), (,,, y) has variance uniformly bounded by a²
consequence of Assumption (h). For V2, we use Assumption (d) where we have that a.., u en is
uniformly bounded by a constant B2 and apply the bounds on || 1 and E || au 7 from Lemma E.6.
Lemma E.6. For the optimal prediction and adjoint functions h a defined in Equations (6) and (7), as well as their estimated versions hw åw given in Algorithm Under Assumptions (a) (f) and Assumptions (A) to (L), the following holds:
E how
S 2' Eins
E [1ā0 aw 1116
s 2u -
E
[lā atll,]] 2 22- -C 4 C.M) pr p
E
|لااا [1]]]] < 44 Cc 24 B
|aill S 11
Proof. We show each of the upper bounds separately.
Upper-bound on E I - n: |
We use the strong-convexity Assumption (b) and Assump-
tion (a) to show the first bound:
E New n: 1011 74
(Tan(ws, hww) Tein kr))]]
T
(52)
Upper-bound on
The second bound can be proven in the same way as the first, using that, by definition, Ladi is continuous and μ-stroogly convex in together with Assumption (a).
||@w - aw 11
34

---

Upper-bound on E |1āw a atlle.] J.
We exploit the closed form expressions of aw and az:
aw -
Lil hws)
Onlour (w, (w) o -
0f
Lin (wi, hit)
On Loue (wi mw)
Hw
bw
Bw
b
By standard linear algehra, the difference āي - a can be expressed as:
aw - a = = (Hl Ho) H - -1 (bw-
By taking the Li(P) norm of the above and using the upper-bounds from Lemma E.7, we can write:
1/ -ill, 1 H (ow
Hw)
H
bb
+ +1
(w
8
1.1
111
[ - Ha)
A 1 bw
+ +14 -=
hw
b
L1
L1
S a
Pim h
1 bw
+ --
bbw
8
21
n
72
s c how - h
|bwll 7 + 1
C
Now
1
71
s (w Cs - CcM) - n:ll1x
The first line used the triangular inequality, the second line follows from Equation (53) from Lemma E.7. The third line applies Equation (54) from Lemma E.7 to the first term and uses
that bew 1l.-ban Ba L1 S [bw - b 6011 by Cauchy-Schwae inequality. The fourth line uses that / 11 111 11w for the first term since | S ~ー by Assumption (b) and uses Equation (55) from Lemma E.7 for the second term. The final bound is obtained using Assumption (e) to upper-bound ||6, by B3. By taking the expectation w.r.t. ho and using Equation (52), we get:
[| ni 1122
E سلاق atll..]
s (
$ (M C 1 tcn.
Upper-bound on E
سسقق 1
We use the closed-fomm expression of aw
| 1l = |
/,. |
[ ll
$ -1 (1|x0- |bol + -(bullze)
S p -1 C2M | hir 11 + Bs)
where the first line uses that A S 1-1
1 by Assumption (b), the second line follows by triangular inequality while the last line uses Equation (55) from Lemma E.7 for the first term and Assumption (e) for the second terms. By squaring the above bound and taking expectation w.r.t h we get:
(1=
+ + -ta Bss) 27
E |لwااا
112d
S E
(CMM
C,M
27
2 22
+ 221 2-
$ 4p 30 ci
+ 211 -B
35

---

where the last line uses Equation (52). Upper-bound ||a Using the closed-form expression of a", it holds that:
| 21 = 4 bwwl - -u-B.....
where we used that | 1lop S xー by Assumption (b) and that 11lllx S B2 by Assumption (d).
Lemma E.7. Consider the operators H., Hu, bw by defined Lemma E.6. Under Assumptions (b), (c) and (f) and Assumptions (A) (L), the following holds for any (w, 8) E 52 x H:
H s112, Sp |
$
S 5 |1
(53)
( Hu)s|) s SC |iou hiu 1|sllze
(54)
lbw - bull S CCM now - hio hillin
(55)
Proof. We show each of the upper bounds separately.
Upper-bound on || i|l and
H s Using strong convexity Assumption (b), we have the
following inequality for the Hessian operator H., acting function H (P)):
Hadll,, Eg [[ [0] h (w), x. w)(a] c(=) 2 p [
by positive-definiteeesss of Has we take c Hwt for some & E H:
H 11 S ||s||E..
By the same arguments, the above bound applies ||
Upper-bound on
( Hw) dle.- We can express the operator (Aw Hw) acting on some
8 E H as follows:
Hu. - H=)s= (-)==BBrr Ep [(B Of (w, Pew (x), Pi (,))),, xr u1 w) s(xr)]
B
Using Assumption (c) we upper-bound the L2 norm the above quantity as follows:
[ - Ho) oll, S. (())())) ng (=)) ||ا())||]
s Ci li 1i 11112 l|
where we used Cauchy-Sear -Shhhar inequality to get the last inequality.
Upper-bound on bu bu 60112 Using Lemma D.3 to get expression of b. and b... we obtain the following upper-boute bound:
10w 1111
1 Ex ||Be (Onfoouelus 1 x. y) Opkane(w, Bo (w) 2, v.1)) |o r(xc)
S c2 12 | will
where we used that the density r(x) upper-bounded by positive constant M by Assumption (B) and that Eou is C2 Lipschitz in its second argument by Assumption (f).
36

---

F Connection with Parametric Implicit Differentiation
E.1 Parametric approximation of the functional bilevel problem
In this section, we approximate the functional problem in Equation (FBO) with a parametric bilevel problem where inner-level functions are parametrized as h(x)) (1) with parameters 0. Here, the inner-level variable is 0 instead of the function h. Standard bilevel optimization algorithms like AID can be applied, which involve differentiating twice with respect to the parametric model. However, for models like deep neural networks, the inner objective may not be strongly convex in 0, leading to a non-positive or degenerate Hessian (Proposition F.I). This can cause numerical instabilities and divergence from the gradient in Equation (4) (Proposition F.2), especially when using AID, which relies on solving a quadratic problem defined by the Hessian.
If the model has multiple solutions, the Hessian may be degenerate, making the implicit function theorem inapplicable. In contrast, functional implicit differentiation requires solving positive definite quadratic problem in H. to find an adjoint function o ensuring a solution even when h is sub-optimal, due the strong convexity of Lin (w, ht. This stability with sub-optimal solutions is crucial for practical algorithms like the one in Section 3, where the optimal prediction function is approximated within parametric family, such neural networks.
Formally, to establish a connection with parametric implicit differentiation, let us consider : + + H to be map from a finite dimensional set of parameters θ to the functional Hilbert space H. and define a parametric version of the outer and inner objectives in Equation (FBO) restricted to functions in He == (r(0) A θ}:
Gour 0) = Loue (w, (0))
Gi((w,
0) = Lin (0)).
(56)
The map 7 can typically be a neural network parameterization and allows to obtain a "more tractable" approximation to the abstract solution 4g in H where the function space H is often too large to perform optimization. This is typically the case when H is an L2-space of functions as we discuss in more details in Section 3. When H. is Reproducing Kernel Hilbert Space (RKHS), may also correspond to the Nyström approximation [Williams and Sceger, 2000], which performs the optimization on miiee-dimenioonl subspace of RKHS spanned by few data points.
The corresponding parametric version of the problem (FBO) is then formally defined as:
min (wo == Goue(w, (w,, 0%) WES
(PBO)
s.1. 0% E arg min Gmn(w,
Oee
The resulting bilevel problem in Equation (PBO) often arises in machine learning but is generally ambiguously defined without further assumptions on the map as the inner-level problem might admit multiple solutions [Arbel and Mairal, 2022b]. Under the assumption that twice continuously differentiable and the rather strong assumption that the parametric Hessian G&G.. (w, 6%) is invertible for a given w, the expression for the total gradient (wo follows by direct application of the parametric implicit function theorem [Pedregosa, 2016]:
V Grot (w) = Gut 0 + Gin 6.)
(57)
W = G.. (w, 0 Og Gowelw, (w),
where nt, is the adjoint vector in Equation (57) is generally different from the one obtained in Proposition 2.2 using the functional point of view. Nevertheless, a precise connection between the functional and parametric implicit gradients can be obtained under expressiveness assumptions on the parameterization 7, as discussed in the next two propositions.
Ô. Without further assumptions, the expression of the gradient in
Proposition F.1. Under the same assumptions as in Proposition 2.2 and assuming that is twice continouslly differentiable, the following expression holds for any (w, 0) E S x O:
GG 0) := (,,, (0)) 033 10n (w, (0)),,
(58)
where oz is a linear operator measuring the distortion induced by the purameterization and acts on functions in H by mapping them to matrix p p where p the dimension of the parameter space e. If. in addition, T is expressive enough so that T TBt)) = h" then the above expression simplifies to:
GgG (w, ( == (0,
(59)
37

---

Proposition F.1 follows by direct application of the chain rule, noting that the distortion term on the right of (58) vanishes when 0 = 8 since 0₁ L., (w, ((z)) = Ô (w, = 0 by optimality of h A consequence is that, for an optimal parameter 0 the parametric Hessian is necessarily symmetric positive semi-definite. However, for an arbitrary parameter 0, the distortion does not vanish in general, making the Hessian possibly non-positive. This can result in numerical instability when using algorithms such as AID for which an adjoint vector is obtained by solving a quadratic problem defined by the Hessian matrix 0 G.n evaluated on approximate minimizers of the inner-level problem. Moreover, if the model admits multiple solutions C- the Hessian is likely to be degenerate making the implicit function theorem inapplicable and the bilevel problem in Equation (PBO) ambiguously defined'. On the other hand, the functional implicit differentiation requires finding an adjoint function a by solving a positive definite quadratic problem in H. which is always guaranteed to have solution even when the inner-level prediction function is only approximately optimal.
Proposition F.2. Assuming that T is twice continuously differentiable and that for a fixed w E S we have T r(%) hy, and Jw (gT( has a full rank, then, under the same assumptions in Proposition 2.2. Grol(( is given by:
T 9w Bu.
(60)
where Pw: H. + H. is a projection operator of rank dim(@). If. in addition, the equality r(0,)) hh holds for all w' in a neighborhood of w. then Wror == VF(w)= 9w + B
Proposition F.2, which is proven below, shows that, even when the parametric family is expressive enough to recover the optimal prediction function 4" at single value w, the expression of the total gradient in Equation (60) using parametric implicit differentiation might generally differ from the one obtained using its functional counterpart. Indeed the projector P.,. which has rank equal to dim, biases the adjoint function by projecting it into finite dimensional space before applying the cross derivative operator. Only under much stronger assumption T, requiring to recover the optimal prediction function hü in a neighborhood of the outer-level variable w, both parametric and functional implicit differentiation recover the same expression for the total gradient. In this case, the projector operator aligns with the cross-derivativy operator so that B Pa = B Finally, note that the expressiveness assumptions made in Propositions F.1 and F.2 only used here to discuss the connection with the parametric implicit gradient and are not required by the method we introduce in Section 3.
Proof of Proposition F22. Here we want to show the connection between the parametric gradient of the outer variable (w) usually used in approximate differentiation methods and the functional gradient of the outer variable V.F(w) derived from the functional bilevel problem definition in Equation (FBO). Recall the definition of the parametric inner objective G (w. 0) == Lin (w, (()).
According to Proposition F.1, we have the following relation
on [I with == (67(
By assumption, Ju has a full rank which matches the dimension of the parameter space O Recall from the assumptions of Theorem 2.1 that the Hessian operator C., is positive definite by the strong convexity of the inner-objective Lin in the second argument. We deduce that Gg (%) must be invertible, since, by construction, the dimension of θ is smaller than that of the Hilbert space H which has possibly infinite dimension. Recall from Theorem 2.1, B_, := dw, h. Lan (w, h&) and the assumption that T(t)) = h We apply the parametric implicit function theorem to get the following expression of the Jacobian O=:
0 = -Bo (JaCCrl))
Hence, differentiating the total objective (wo == Gout (wi 0%) = Loutt (w, ((t))) and applying the chain rule directly results in the following expression:
V.oGol( Iw BJ (JoCwrlt)) Jodws
(61)
with previously defined gw := Lo htt) and du == Ūr Lout (w, ht).
although a generalized version of such theorem was recently provided under restrictive assumptions [Arbel and Mairal, 2022b].
38

---

We now introduce the operator Pw: := Jh satisfies P Pa. Hence, using the fact that the Hessian operator is invertible, and recalling that the adjoint function is given by a = CD du, we directly get form Equation (61) that:
JuCw- The operator Pa is projector it
V Groe(w) :- 9n + B. Puo
we further assume that =((",) ) = h" holds for all w' in a neighborhood of w, then differentiating with respect to w results in the following identity:
Ou0 Ow hw-
Using the expression of 0 from Equation (3), we have the following identity:
- Bw.
In other words, B_ is of the form B = DJ for some finite dimensional matrix D of size dim(St) x d.). Recalling the expression of the total gradient, can deduce the equality between parametric and functional gradients:
Gror (w) g- - BuJe (JoCodl)) Judh
9w DJJOOwN (JoCwJw) Judw
gw -DJudw
9w - -JJ.CwCold dw
gw B VF(Ww.
The first equality follows from the general expression of the total gradient V. Gtor(w). In the second line use the expression of B. which then allows to simplify the expression in the third line. Then, recalling that the Hessian operator C,,, is invertible, we get the fourth line. Finally, the result follows by using again the expression of B., and recalling the definition of the adjoint function a
F.2 Computational Cost and Scalability
The optimization of the prediction function hh in the inner-level optimization loop is similar to AID, although the total gradient computation differs significantly. Unlike AID, Algorithm does not require differentiating through the parameters of the prediction model when estimating the total gradient F.(w). This property results in an improved cost in time and memory in most practical cases shown in Table and Figure 4. More precisely, AID requires computing Hessian-vector products size Pin, which corresponds to the number of hidden layer weights of the neural network hw While FuncID only requires Hessian-vector products of size d,, i.c. the output dimension of h In many practical cases, the network's parameter dimension pin is much larger than its output size du, which results in considerable benefits in terms of memory when using FuncID) rather than AID, as shown in Figure 4 (left). Furthermore, unlike AID, the overhead of evaluating Hessian-vector products in FuncID is not affected by the time cost for evaluating the prediction network. When hww is a deep network, such an overhead increases significantly with the network size, making AID significantly slower (Figure 4 (right)).
Method AID
Time cost (T + TT,
Memory cost
3pm Mn. _d, + Ma
FuncID
TTL.. + (2+ 8)T, +Th
Table I: Cost in time and memory for performing a single total gradient estimation using either AID FuncID and assuming the prediction model is learned. Time cost: T₁, and T2 represent the time cost of evaluating both prediction and adjoint models h and a, while Tin is the time cost for evaluating the inner objective once the outputs of h are computed. The factors y and are multiplicative overheads for evaluating hessian- -vector products and gradient. Memory cost: Mr. and M,, represent the memory cost of storing the intermediate outputs and a, Pin and d, are the memory costs of storing the Hessian- vector product for AID) and FuncID) respectively and B is multiplicative constant that depends on particular implementation.
39

---

10
65%
10x
82%
100%
72%
Seme Pemertoolccgn
Vencyy Picaan
526
40%
10" 527
30"
Figure 4: Memory and time comparison of single total gradient approximation using FuncID vs AID. (Left) Memory usage ratio of FuncID over AID vs inner model parameter dimension Pins for various values of the output dimension d,. (Right) Time ratio of FuncID over AID inner model parameter dimension Pin averaged over several values of d, and 10* evaluations. The continuous lines experimental results obtained using JAX implementation [Bradbury al., 2018] running on a GPU. The dashed lines correspond to theoretical estimates obtained using the algorithmic costs given in Table with 12, 6 2 for time, and the constant factors in the memory cost fitted the data.
G Additional Details about 2SLS Experiments
We closely follow the experimental setting of the state-of-the-an method DFIV [Xu et al., 2021a]. The goal of this experiment is to learn model fw approximating the structural function fstract that accurately describes the effect of the treatment t on the outcome o with the help of an instrument æ, as illustrated in Figure 5.
fstruct
Figure 5: The causal relationships between all variables in an Instrumental Variable (IV) causal graph, where is the treatment variable (dsprites image), is the outcome (label in R), x is the instrument and is the unobserved confounder
G.1 Dsprites data.
We follow the exact same data generation procedure as in Xu et al. [2021a, Appendix E.3]. From the dsprites dataset [Matthey et al., 2017], we generate the treatment and outcome o as follows:
1. Uniformly sample latent parameters scale, rotation, posX, posY from dsprites.
2. Generate treatment variable as
Fig(scale, rotation, posX, posY) + X.
3. Generate outcome variable as
At | 5000
+ 32 - 0..5) +
1000
Here, function Fig returns the corresponding image to the latent parameters, and X, are noise variables generated from (0.0, 0.17) and 5 (0000, 0.5). Each element of the matrix A E 10 is generated from Unif(0.0, 1.0) and fixed throughout the experiment. From the data generation process, we can see that and are confounded by posY. We use the instrumental variable =(scale, rotation, PosX) 123 and figures with random noise treatment variable The variable
40

---

posY is not revealed to the model, and there is no observable confounder. The structural function for this setting is
|| 5000
fstruct (t)
1000
Test data points are generated from grid points of latent variables. The grid consist of 7 evenly spaced values for posX, posY, evenly spaced values for scale, and 4 evenly spaced values for orientation.
G.2 Experimental details
All results are reported over an average of 20 runs with different seeds on 24GB NVIDIA RTX A5000 GPUs.
Feature maps. As in the DFIV setting, we approximate the true structural function fstruct with La - u (1) where 1ix is a feature map of the treatment I, u is vector in Rdde and fw is parameterized by = x). To solve the inner-problem of the bilevel formulation in Section 4.1, the inner prediction function h is optimized over functions of the form h(x) = VΦ(ω) where denote Φ the feature map of the instrument and V is matrix in Xd1 The feature maps th and á are neural networks (Table 2) that are optimized using empirical objectives from Section 3.1 and synthetic dsprites data, the linear weights V and are fitted exactly each iteration.
Choice of the adjoint function in FuncID. In the dsprites experiment, we call linear FuncID the functional implicit diff. method with linear choice of the adjoint function. Linear FuncID uses an adjoint function of the form a (x) = W ((z) with W E R xd1, In other words, to find a the features & are fixed and only the optimal linear weight W computed in closed-form. In the F FuncID method, the adjoint function lives in the same function space as h This is achieved by approximating a with a separate neural network with the same architecture as hw-
Layer
instrument feature map Ø
Layer
treatment feature map 2'x Input (t) FC(4096, 1024), SN, ReLU FC(1024, 512), SN, ReLU, LN FC(512, 128). SN, ReLU FC(128, 32). SN, LN, Tanh
Input(a:) 2 FC(3, 256). SN, ReLU
2
FC(656, 128), SN, ReLU, LN FC(128, 128), SN, ReLU, LN FC(128, 32), SN, LN, ReLU
3
4
4
5
Table Neural network architectures used in the dsprites experiment for all models. The FuncID model has an extra fully-connected layer FC(32, 1) in both networks. LN corresponds to LayerNorm and SN to SpectralNorm.
Hlpper-parameter tuning. As in the setup of DFIV, for training all methods, we use 100 outer iterations (N in Algorithm 1), and 20 inner iterations (M in Algorithm 1) per outer iteration with full-batch. We select the hyper-parameters based the best validation loss, which we obtain using validation set with instances of all three variables (t, o, x) [Xu et al., 202 la, Appendix A]. Because of the number of linear solvers, the grid search performed for AID very large, so we only run with one seed. For other methods, we run the grid search on 4 different seeds and take the ones with the highest average validation loss. Additionally, for the parameesrsseters that not tuned, we take the ones reported in Xu et al. [2021a].
Deep Feature Instrumental Variable Regression: All DFIV hyper-parameters are set based on the best ones reported in Xu et al. [2021a].
Approximate Implicit Differentiation: We perform grid search over 5 linear solvers (two variants of gradient descent, two variants of conjugate gradient and an identity heuristic solver), linear solver learning rate 10-" with E {3, 4, 5). linear solver number of iterations {2, 10, 20), inner optimizer learning rate 10 n with n E (2, 3, 4), inner optimizer weight decay 10 with (1.2.3} and outer optimizer learning rate 10 with (2,3, 4).
Iterative Differentiation: We perform grid search over number of "unrolled" inner iterations {2, 5) (this chosen because of memory constraints since "unrolling" an iteration is meaory- number of warm-start inner iterations {18, 15), inner optimizer learning
41

---

rate 10Th with (2.3, inner optimizer weight decay 10 with E (1,2,3) and outer optimizer learning rate 10 " with n E 3,.3,41
Gradient Penalty: The method is based on Eq. 5.1 in Shen and Chen [2023]. for this single- level method we perform grid search on the learning rate 10 with n E (,.3, 4, 6). weight decay 10-" with
{1,2, 3), and the penalty weight 10-" with n E 2,
Since the method has only a single optimization loop, we increase the number total itera- tions to 2000 compared to the other methods (100 outer-iterations and 20 inner iterations).
The method based on Eq. 3.2 in Shen and Chen [2023], for this method perform the same grid search for the Gradient Penalty method. However, since this method has an inner loop, we perform 100 outer iterations and perform grid search on the number of inner iterations with 10 and 20.
Value Function Penalty:
We perform a grid search over the number of iterations for learning the adjoint network {10, 20), adjoint optimizer learning rate 10 with {2,3, 4,5, 6} and adjoint optimizer weight decay 10 with n E 2,,2,3 3). The rest of the parameters are the same for DFIV since the inner and outer models almost equivalent to the treatment and instrumental networks used their experiments.
FuncID:
G.3 Additional results
We additional experiment with 10k training points using the same setting described above to illustrate the effect of the sample size on the methods. Figure 6 shows that a similar conclusion can be drawn when increasing the training sample size from 5k to 10k, thus illustrating the robustness of the obtained results.
CLL- of Sampii MSE
ITD
W0. pona DETV
Locs
AG
2.0
Vol 201t
42x1
370
Parco
DFV
Furei
ReCO
aaad Wereesa DRV PanlDla PealO
iterations
130
iterations
100 2012
Figure 6: Performance metrics for Instrumental Variable (IV) regression. (Left) final test loss. (Middle) outer loss vs training iterations, (Right) inner loss vs training iterations. All results are averaged over 20 runs with 10000 training samples and 588 test samples.
H
Additional Details about Model-Based RL Experiments
H.1
Closed-form expression for the adjoint function
For the FuneID method, we exploit the structure of the adjoint objective to obtain a closed-form expression of the adjoint function a In the model-based RL setting, the unregularized adjoint objective has a simple expression the form:
1
Lodj (w, a, Rw.
Σ
||a(x)]1
(62)
2 B
+
alx) 0.f (h.
(63)
Bout
The key observation here that the same batches of data are used for both the inner and outer problems, i.c. Bin Bout. Therefore, only need to evaluate the function where (x. y) B.n Without restricting the solution set of optimal solution at simply matches (h y) on the set of points x s.t. (x, y) Bin Our implementation directly exploits this observation and uses the following expression for the total
finite set of points adding regularization to Ladj. the
42

---

gradient estimation:
Dout
Σ (I.
(64)
H.2 Experimental details
As in the experiments of Nikishin et al. [2022], we use the CartPole environment with 2 actions, 4-dimensional continuous state space, and optimal returns of 500. For evaluation, copy of the environment. The reported return is an average of 10 runs with different seeds.
separate
Networks. We us the same neural network architectures that are used in the CartPole experiment of Nikishin et al. [2022, Appendix DJ. All networks have two hidden layers and ReLU activations. Both hidden layers in all networks have dimension 32. In the misspecified setting with the limited model class capacity, we set the hidden layer dimension to for the dynamics and reward networks.
Hyper-parameters. We perform 200000 environment steps (outer-level steps) and set the number of inner-level iterations M 1 for both OMD and FuncID. for MLE, perform single update to the state-value function for each update to the model. For training, we use a replay buffer with batch size of 256, and the discount factor 7 to 0.99. When sampling actions, parameter = 0.01 as in Nikishin et al. [2022]. The learning rate for outer parameters w is set to 10 For the learning
temperature
of the inner neural network and the moving average coefficient т, we perform a grid search over 110 10-3,010000 3. 10 31 and {5 10 10-3 10 1002} as in Nikishin et al. [2022].
H.3 Additional results
Time comparison. Figure 7 shows the average reward on the evaluation environment as a function of training time in seconds. We observe that our model is the fastest reach best performance both in the well-specified and misspecified settings.
1e2
Reward
Reward
Algorithms funcID
Algorithms
AID (OMD)
funciD
MLE
MLE
Time
163
Time (s)
Figure 7: Average Reward on an evaluation environment vs. time in seconds on the CartPole task. (Left) Well-specified predictive model with 32 hidden units to capture the variability in the states dynamics. (Right) misspecified predictive model with only hidden states.
MDP model comparison. Figure 8 shows the average prediction error of different methods during training. The differences in average prediction error between the bilevel approaches (OMD, FuncID) and MLE reflect their distinct optimization objectives and trade-offs. OMD and FuncID focus on maximizing performance
task environment, while MLE emphasizes accurate representation of all aspects of the environment, which can lead to smaller prediction errors but may not necessarily correlate with superior evaluation performance. We also observe that FuncID has stable prediction error in both settings meanwhile OMD and MLE exhibit some instability.
43

---

10a
10
Model prediction error
10:
Model prediction error
10"
10"1
10
10
10-1
Aggrithmms
Algorithms
www
10-3
(OMD)
TunciD
funcID
10"
1.800 1.825 1,850
1.875
1,900 1.925 Steps
1.950
1.975
2:000
1.800 1.825 1.350 1.675 1.900 1.925 1. 950 1.975 2.000 Steps
1e5
1e5
Figure 8: Average MDP model prediction error in the training environment vs. inner optimization steps on the CartPole task. (Left) Well-specified predictive model with 32 hidden units to capture the variability the states dynamics. (Right) misspecified predictive model with only 3 hidden states.
44

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All the claims made in the paper are either rigorously proven, shown experi- mentally or describe well-established facts in the literature.
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 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 be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: All strong assumptions used in the theoretical results are clearly stated, dis- cussed, and put into perspective with other published work.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic 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, e.g., if the approach was only tested on few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 an impor- tant role in developing norms that preserve the integrity of 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 set of assumptions and complete (and correct) proof?
45

---

Answer: [Yes]
Justification: All theoretical claims are rigorously proven. We provide structured list of all assumptions used before the statement of result and refer specific assumption whenever it is employed in the proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix 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 ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide the code to reproduce our experiments, which includes a README file with instructions. We also include the information about tuning, datasets and other experimental detail in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If 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 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 to the model. In general. releasing code and data is 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
46

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide anonymized repository with extensively commented code and instructions on how to reproduce our main experiment.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for 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 (https:
ctpps/cccppbbbccguuu getii//ggge///eeeeeeeeeiieeessee /Co///oooeeseeeeeeeeeeee simmmmmssioonl ssisssoliiyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Yes, we specify and discuss all experimental details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In all of our experiments we run multiple times and report the error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
47

---

The method for calculating the error bars should be explained (closed form formula, call a 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 or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that 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 or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report which GPUs we use for larger experiments and show how our method compares to similar methods in the section about computational cost.
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 compute required for each of the individual experimental runs well 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 or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: https::/neurbba //tp:://neuupspeeeeeee wttp/e/ppiii...c.c.e httee/Etthiiiiieeeeeeeeeeeee ines?
Answer: [Yes]
Justification: We reviewed the NeurIPS Code of Ethics and attest that our work conforms with
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We present new bilevel optimization method, which does have any immediate societal impacts.
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 societal impact why the paper does not address societal impact.
48

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used 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 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 monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our new optimization method does not have an identifiable risk for misuse.
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 or restrictions to access the model or 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 recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Our code includes an open access licence in the README file.
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 dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a 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 package should be provided. For popular datasets, paperswi twwode..coo..... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
49

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not present new assets in this work.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details the dataset/code/model as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
We do have crowdsourcing nor human subjects in our work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not have human subjects in work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
50


Reviews:
[{"review_id": "uyVzZLwAIx", "reviewer": "Reviewer_p8LC", "summary": "This paper offers a new functional point of view for bilevel optimization problems in machine learning. This functional approach allows the use of an overparameterized neural network as inner prediction function while previous works have used an inner objective that is strongly convex with respect to the parameters of the prediction function. For the inner problem, the prediction function is a function that lies in a Hilbert space of square integrable functions ($L^2$). The authors develop the theory of Functional Implicit Differentiation, which is a flexible class of algorithms to do functional bilevel optimization over $L^2$ spaces. First, they show that strong convexity assumption with respect to the prediction function as opposed to the model parameters ensures the existence and uniqueness of the optimal prediction function $h_\\omega^\\ast$ for the inner-level objective. Second, they show that differentiability assumptions on the inner-level objective and its Fr\\'echet derivative with respect to prediction function $h$ ensure the differentiability of the map $\\omega \\rightarrow h_\\omega^\\ast$. Given further assumption about joint differentiability of outer objective, they show that it is possible to compute total objective $\\mathcal{F}$ using the adjoint function that minimizes a quadratic objective over a Hilbert space. Finally, assuming that the inner objective and outer objective are defined over the distributions and that a batch of data samples from each distribution can be sampled, an iterative algorithm is proposed that has the following three steps: (1) approximation of inner objective, outer objective and quadratic objective for adjoint function using the batch samples, (2) do a gradient-based update for the parameters of the prediction function and adjoint function assuming they have a fixed parametric form and (3) total gradient approximation and gradient-based update of the functional parameter $\\omega$ obtained by solving the outer objective. They show that these class of algorithms converge to a stationary point at $\\mathcal{O}(1/N)$ rate. Experiments are performed using Two-stage least squares regression and Model-based reinforcement learning as use cases.", "questions": "1. On Line 126, there is a typo: \"elemen\" should be replaced by \"element\". \n2. On Line 629, capitalize the word \"euclidean\".\n3. What purpose does the variable $y$ and space $\\mathcal{Y}$ serve in general? Is it a variable or a parameter? Your prediction function $h_\\omega$ is a function of $x$. In the case of 2SLS, you have $y = (t, o)$ and $x$ is the instrumental variable. The function you are interested in is $f_\\omega(t)$ and is a function of $t$, which is a subvariable of $y$ and $h_\\omega(x) = \\mathbb{E}[f_\\omega(t) | x]$. Maybe it will help to clarify what $y$ and $\\mathcal{Y}$ are in the theoretical section. \n4. What is the rationale for the label \"FuncID linear\" in Figure 1? Why do you refer to it as \"linear\"?\n5. Why do you feel \"FuncID linear\" converge faster in terms of outer iterations than \"FuncID\" in Figure 1?", "limitations": "The authors didn't go into a discussion about the limitations nor did they provide a conclusion section and simply concluded with a results section.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper makes a substantial contribution by developing the theory of functional bilevel optimization with less restrictive assumptions of strong convexity with respect to the prediction function as opposed to model parameters. This is useful because it allows the prediction function $h_\\omega$ to be modeled by deep neural network, which has a non-convex training objective with respect to the model parameters. This paper could lead to more research into optimization over function spaces as strong convexity over model parameters is a restrictive assumption in practice.", "weaknesses": "The paper is very technical and requires a good understanding of monotone operator theory and theory of Fr\\'echet and Hadamard differentiability to understand it fully. Still, the technical details about the existence and uniqueness of the prediction function and the map $\\omega \\rightarrow h_\\omega^\\ast$ are deferred to the appendix for interested readers. I would recommend the authors simplify the notation in the main paper a bit to reduce the clutter. For instance, $h_\\omega^\\ast(x)$ can be represented simply as $h_\\omega$. It may help to have a table that shows the functional arguments (e.g. $h_\\omega$, $a_\\omega$, etc.), functional parameters $\\theta$, $\\xi$ and $\\omega$ and then function variables $x$ and $y$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for thoroughly reading our work and giving us helpful feedback. Taking your feedback into account, we will simplify the notation and include a notation table. We agree that this could help the reader get a quick grasp of the mathematical objects considered. We will also include a discussion section (presented in the general response) with a paragraph on the limitations of our work. Below, we address your comments in detail.\n\n**Typos**. We will fix the typos in the final version.\n\n**Purpose of $y$ and $Y$**. Data is represented by pairs $(x,y)$, noting that the function $h$ takes only $x$ as input. Thus, $x$ and $y$ are both random variables. In the simplest supervised learning setup (see *Eq. 1* for instance), $y$ is a label living in a space $Y$ and $h(x)$ tries to predict $y$. The theory of Section 3 is however more general, and $y$ can serve other purposes. For instance, in 2SLS, $y$ is made of two variables $t$ and $o$, whose causal relationship is described in Fig. 4 (App. H). We admit that the setup of 2SLS Instrumental Variable regression is a bit particular and can be confusing. We will clarify these points in the final version of the paper.\n\n**Difference between *FuncID* and *FuncID linear***. *FuncID linear* uses a linear model to approximate the adjoint function, while *FuncID* uses a trainable neural network. The linear model is obtained by learning a linear combination of the frozen penultimate layer of the current prediction function $h$ (modelled as a neural network). We will clarify this in the text.\n\n**Why *FuncID linear* converges faster**. Following the previous explanation, *FuncID* optimizes all adjoint network parameters, while *FuncID linear* learns only the last layer in closed form. \"FuncID linear\" converges faster because solving a linear system is computationally faster than iteratively approximating the full adjoint function using a neural network. However, *FuncID linear* is less expressive for approximating the adjoint, which may result in suboptimal solutions that could explain the performance gap in terms of test error.", "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": "YrCO4sm24S", "reviewer": "Reviewer_MXrv", "summary": "This paper introduces a novel functional perspective on bilevel optimization, where the inner objective is defined over a function space. The authors developed functional implicit differentiation and functional adjoint sensitivity, which together facilitate the establishment of a gradient-based algorithm in the functional space. They also analyze the convergence rate of the proposed algorithm and apply it to two-stage least squares regression and model-based reinforcement learning. Experimental results validate the effectiveness of the proposed method.", "questions": "Same as weakness.", "limitations": "Same as weakness.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method offers a new insight into solving bilevel optimization problems with nonconvex lower-level objectives by leveraging their strong convexity in the functional space. This is particularly noteworthy because, although neural networks are nonconvex, the loss function in model training can be convex or strongly convex.\n2. This paper provides a heuristic approach with both theoretical and practical impact. The proposed method not only has a convergence guarantee but is also implementable in real-world applications. The two applications chosen in this paper are also novel: the first has potential impacts on causal representation learning, while the second provides a new perspective on model-based reinforcement learning.", "weaknesses": "1. The convergence analysis in this paper is based on the stochastic biased gradient descent framework, making the results explicitly dependent on the sub-optimality constant $\\epsilon_{in}$ and $\\epsilon_{adj}$. However, it is unclear how these errors relate to the inner loop $M$ and $K$. It might be beneficial to leverage the strong convexity of the inner and adjusted objective functions to clarify these dependencies. See similar techniques used in [1]-[3]. \n\n[1] K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design. ICML, 2021. \n\n[2] T. Chen, Y. Sun, and W. Yin. Closing the Gap: Tighter Analysis of Alternating Stochastic Gradient Methods for Bilevel Problems. NeurIPS, 2021. \n\n[3] M. Dagréou, P. Ablin, S. Vaiter, T. Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. NeurIPS, 2022. \n\n\n2. Minor issues: experimental baselines. It might be better to also compare with those bilevel methods that are capable to solve nonconvex lower-level problem [4]-[5] as they can also potentially solve the two-stage least squares regression and model-based reinforcement learning problem. \n\n[4] J. Kwon, D. Kwon, S. Wright, and R. D. Nowak. On penalty methods for nonconvex bilevel optimization and first-order stochastic approximation. ICLR 2024. \n\n[5] H. Shen, and T. Chen. On Penalty-based Bilevel Gradient Descent Method. ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Error analysis**. Thank you for pointing out these references, we will make sure to discuss them. We agree that the result of Thm. 3.1 does not provide an explicit dependence of the errors on the inner-level optimization. Providing such dependence would require introducing another level of technical complexity, beyond the techniques used in *[1]-[3]* which are tailored for the strongly convex case in a parametric setting. In our case, one would instead need to use quantitative approximation results of functions in $L_2$ spaces by NNs [Bach 2017], as well as global convergence results for NNs [Allen-Zhu et al., 2019, Liu et al., 2022]. Such analysis would require substantial effort that is best suited for a separate future work. We discuss this in the future work section (see the general response).\n\n**Additional comparison**. Thank you for suggesting these methods. We performed an additional comparison on the Instrumental Variable (IV) problem (see Fig. 1 in the pdf file). These methods handle non-convexity by considering an optimistic version of the bilevel problem and turning it into a penalized single-level problem. However, they do not exploit the functional strong convexity of the IV problem. Consequently, the new results, on average, still favor *FuncID*, which exploits functional strong convexity as shown in the general response.\n\n[Bach 2017] Bach, F. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19), 1-53 2017.\n\n[Allen-Zhu et al. 2019] Zeyuan Allen-Zhu. Yuanzhi Li. Zhao Song. A Convergence Theory for Deep Learning via Over-Parameterization. ICML 2019.\n\n[Liu et al. 2022] Chaoyue Liu, Libin Zhu, Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis. 2022", "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": "zf5i3Lg1y6", "reviewer": "Reviewer_98ub", "summary": "The paper proposes a functional approach to bilevel optimization for machine learning, focusing on inner-level problems defined over function spaces rather than traditional parametric settings. This allows the application of the proposed method to machine learning tasks without requiring the strong convexity typically assumed in bilevel optimization.", "questions": "1. Can the proposed functional framework be generalized to other types of bilevel problems that involve constraints on the lower-level problem? How would you address potential non-smoothness or non-convexity in these extended settings?\n\n2. Please explain how does the choice of function space impact the stability and convergence of the FuncID method?", "limitations": "Authors addressed their work limitation.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper introduces a functional perspective to bilevel optimization, extending its applicability to settings where traditional assumptions (like strong convexity) do not hold.", "weaknesses": "- The paper does not sufficiently compare the proposed methods against a broad spectrum of existing algorithms, particularly the latest advancements in the field of bilevel optimization problem. This lack of comprehensive benchmarking restricts the ability to fully evaluate the performance enhancements or potential drawbacks of the proposed methods relative to the state-of-the-art.\n\n- The paper lacks a thorough analysis of how the proposed methods perform across varied settings and parameter configurations. It is not clear if there are specific scenarios where the methods might underperform or fail to converge. Additional details on the robustness of the methods in diverse operational environments would be beneficial.\n\n- The assumptions necessary for the theoretical framework might not be easily verifiable in practical scenarios, potentially limiting the method's applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Additional comparisons**. As suggested by reviewer *MXrv* and in addition to the comparisons already made with most widely-used bilevel algorithms (AID, ITD, and variants) and SoTA methods for each problem (DFIV for the IV problem and MLE for model-based RL), we now additionally include a comparison with a recent approach for solving non-convex bilevel problems based on penalty methods (see general response). The new results are, on average, consistent with the previous ones and are still in favour of *FuncID* which exploits functional strong convexity. We would be happy to include additional methods if you think they would be relevant.\n\n**Relevant settings for FuncID**. We expect *FuncID* to outperform AID in settings where the bilevel problem has a ‘hidden’ strong convexity in functional space. That is simply because AID does not exploit such strong convexity, while *FuncID* does. While this setting covers many practical scenarios (such as those considered in the paper (two-stage least squares regression and model-based RL), we do not expect particular improvements in the absence of such a structure. We will make that clear in the text and limitation section. \n\n**Verifiable assumptions**. We agree that the assumptions in Prop. 2.3 might seem complex, but they are easily verifiable through standard calculations. In Proposition E.1 of the appendix, we verify that these assumptions hold for regression problems involving feedforward networks when using quadratic objectives. The verification only requires computing derivatives and upper-bounding them, and could be applied to other problems similarly. We will make sure this is clear in the text.\n\n**Extensions**. Thank you for raising these points, we discuss them in the future work section (presented in the general response). Extending the framework to a constrained inner problem or non-smooth setting should be possible, if the uniqueness of the solutions in functional space is preserved. However, this would require introducing additional tools to handle non-smoothness/constraints such as those from the recent works on non-smooth bilevel optimization [Bolte et al., 2022] and would be an interesting future work direction.\n\n**Choice of the function space**. The function space we consider is motivated by existing bilevel problems that are already formulated in such spaces (of which we consider 2 examples in the applications). Using different spaces would require a different analysis which is beyond the scope of this work but would certainly be interesting for future work as discussed in the new limitation/future work section (provided in the general response).", "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": "PkdZleoWYu", "reviewer": "Reviewer_QfTd", "summary": "The authors propose a functional view of bilevel optimization in machine learning, in which the inner objective is often strongly convex for many problems of interest. The authors prove the resulting optimization algorithm's convergence and benchmark their approach for regression and model-based RL tasks.", "questions": "1. Can you elaborate on $\\epsilon_{in}$ and $\\epsilon_{adjoint}$ in Theorem 3.1? I'm not sure what the expectation is taken over in the definition of line 1100.\n2. Why is the prediction error of FuncID worse than MLE in Figure 7b?\n\nNotes:\n1. In (FBO), there's a space between the colon and the equals (use \\coloneqq).\n2. In line 94, link to where the assumption is discussed.\n3. Font size in the figures is very small and hard to read.", "limitations": "I'd like to see an explicit limitations section. The authors' response to the limitations question in the checklist is that they state theorem assumptions -- this is not a comprehensive discussion of limitations, which should include practical experimental considerations as well.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I'm borderline on this submission as the authors have invested substantial effort, but I'm not fully convinced. I am interested to see what other reviewers think.\n\n1. The authors address an important problem with a range of applications. I'd additionally suggest adversarial robustness / adversarial games in general as another important bilevel optimization problem area.\n2. The authors have a good balance of theory and experiments, with particularly extensive theoretical contributions (although I haven't had the bandwidth to check the proofs).\n3. The writing and logical flow of the paper are solid, and the authors do a good job of making complicated theory comprehensible.", "weaknesses": "1. The authors seem to do a \"bait-and-switch\" when going from function space to parameter space. Namely, the authors repeatedly emphasize that in function space, the inner objective is generally strongly convex. However, as the authors point out, this is no longer true when moving to function parameterizations -- and all the experiments concern concrete parameterizations. I'm fairly lost here as to how the authors theoretically handle this jump, as all the theory seems to assume that we are operating in function space.\n2. The assumptions in Theorem 3.1 seem very strong. Namely, it is assumed that the inner and adjoint problems are solved close to optimality; however, the inner optimization is nonconvex in the parameters (as discussed above).\n3. The authors repeatedly scatter in comparisons to AID throughout the text. I think it would be good to have these all summarized in one table in the appendix (runtime, convergence guarantees, assumptions, etc.)\n4. The experimental results are quite weak. Namely, buried in H.3 the authors experiment with boosting the training sample size by a factor of two, and in this setting FuncID seems to underperform DFIV (although FuncID linear is now better). It seems that the improvement over the state of the art is not very robust.\n5. FuncID seems to require more outer iterations to converge than baselines (Figure 1b), which I don't see being discussed by the authors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback. We will correct the typos, increase the font size in the figure captions, and add a link to the discussion of the assumptions as suggested. Additionally, we will include a discussion section in the main paper (provided in the general response) to address limitations and perspectives. We address other points below.\n\n### From functional to parametric\n\nOur approach addresses one of the two main challenges below (C1 and C2) in using deep networks for bilevel optimization. We outline these challenges and show how our method overcomes C2 using strong convexity in function space. These clarifications will be added to the text.\n\n**A. Challenges of bilevel optimization with deep networks** \n\n- **C1- Non-convexity of the lower-level optimization problem** is unavoidable with deep networks. However, several studies show this non-convexity to be 'benign' and ensure global convergence when optimising deep networks using gradient methods in the overparameterized regime (see general comments). \n- **C2- Ambiguity of the bilevel problem (Arbel & Mairal 2022)**. Exact solutions at the lower level can be multiple due to over-parameterization. Thus, no 'implicit function' links the inner-level solution to the outer-level parameter, making it impossible to use the implicit function theorem to compute the total gradient.\n\n\n**B. Algorithm derivation**. From a functional bilevel problem, our approach is to \"Differentiate implicitly first, then parameterize,\" leading to funcID. The alternative, \"Parameterize first, then differentiate implicitly,\" results in AID, as described below.\n \n- **Differentiate implicitly first, then parameterize (FuncID)**. Functional strong convexity is used to apply the implicit function theorem in function space and derive the implicit gradient. This gradient is then estimated by approximating the lower-level solution and adjoint function using neural networks. While optimising these networks involves a non-convex problem (Challenge C1), the approach avoids Challenge C2 since the neural networks merely estimate quantities in the well-defined implicit gradient expression.\n\n- **Parametrize first, then differentiate implicitly (AID)**. The inner-level function is constrained to be a NN, converting the problem into a non-convex ‘parametric’ bilevel problem where the lower-level variables are the NN parameters. Computing the implicit gradient requires viewing the optimal network parameter as an implicit function of the outer parameter. However, this implicit function does not exist due to multiple global solutions to the inner problem for a given upper variable $\\omega$. Thus, this approach faces both challenges C1 and C2.\n\n### Other comments\n\n**Optimality assumption in Thm 3.1**. The assumptions in Thm 3.1 align with recent findings that non-convexity is 'benign' when optimizing over-parametrized NN ensuring that gradient methods linearly converge to global solutions. Thus, *in such a regime*, one can reduce the errors $\\varepsilon_{\\text{in}}$​ and $\\varepsilon_{\\text{adj}}$ by optimising the inner and adjoint functions using gradient descent. We will clarify these points (see proposed limitation section).\n\n**Comparison table between AID and FuncID**. For more clarity, we propose to merge section D and F of the appendix which both compare AID and FuncID from different aspects. \nWe will also summarize the comparisons already present throughout the paper between AID and *FuncID* into a single table that will be included in the paper. \n\n**Experimental results**: \n\n*FuncID vs. DFIV for instrumental regression*: to reach a more rigorous conclusion, we have performed statistical tests. For the 5K dataset, funcID outperforms DFIV (p-value=0.003, one-sided paired t-test), but for the 10K dataset in H.5, the difference between both approaches was not statistically significant. Overall, FuncID performs in the same ballpark as the state-of-the-art approach DFIV, which is specifically designed for the IV problem. We will add these observations in the discussion of the results.\n\n*Robust/consistent improvement over AID/ITD*: all results, including those in Appendix H.3, show that FuncID outperforms commonly used bilevel optimization algorithms in ML (namely AID and ITD). These results were obtained by fairly allocating a budget for selecting the best hyperparameters for AID and ITD.\n\n*Additional comparisons*: we have compared our method with a recent approach that handles non-convexity by considering an optimistic version of the bilevel problem and turning it to a penalised single-level problem (see general response for more details).\n\n**Number of outer iterations**. Convergence in Fig 1 can be assessed by monitoring both inner-level and outer-level losses. Outer-level loss alone does not indeed indicate convergence unless the outer loss is evaluated at the ‘exact’ inner-level solution, which is generally inaccessible before convergence. We agree that this is a source of confusion. The part of Fig 1 from which it is easier to draw conclusion is the out-of-sample MSE, which reflects generalization (see comment above). \n\n**Q 1**. $\\mathbb{E}$ denotes the expectation with respect to the random data samples. The quantities $\\epsilon_{in}$ and $\\epsilon_{adj}$ represent a bound on the optimality error of the approximate solutions $\\hat{h}$ and $\\hat{a}$.  For instance, $\\mathbb{E}[L_{in}(w,\\hat{h_w})- L_{in}(w,h_w^{\\star})]  \\leq \\epsilon_{in}$. The expectation accounts for the fact that these approximations are based on random samples. We agree that this deserves some clarifications.\n\n**Q 2**. It is expected that MLE has a smaller prediction error because it explicitly minimizes the prediction error in its objective. In contrast, the bilevel formulations (FuncID and AID) learn an MDP whose state-value function aligns with the true MDP.", "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:
# General comments

We thank the reviewers for their useful feedback. We now list the main changes made to the paper.

## Discussion section (limitations and perspectives). 

We agree that such a section is important, and propose to include the following discussion:

### Discussion and concluding remarks
This paper introduces a functional paradigm for bilevel optimization in machine learning, shifting the optimization focus from the parameter space to the function space. This new approach addresses the limitations of traditional bilevel optimization methods when dealing with over-parameterized neural networks. Specifically, the proposed method exploits the functional strong convexity of certain bilevel problems to derive an abstract, yet approximable, expression for the implicit gradient that requires solving both an inner and an adjoint optimization problem in functional space. Approximation is achieved by restricting both problems to a flexible class of functions, such as neural networks. The paper establishes the validity of this approach by developing a theory of functional implicit differentiation and providing a general convergence result for the proposed method. Despite these contributions, we discuss several limitations of our work and highlight potential research directions.

**Hyperparameter selection**. One notable limitation is the presence of multiple hyperparameters in the proposed algorithms. This is a common challenge shared by all bilevel optimization methods, complicating the practical implementation and tuning of these algorithms. Selecting and optimizing these hyperparameters can be time-consuming and may require extensive experimentation to achieve optimal performance.

**Convergence guarantees:** The result in Theorem 3.1 relies on the assumption that both the inner and adjoint optimization problems are solved up to some optimality errors. This assumption is motivated by recent global convergence results for over-parameterized networks [Allen-Zhu et al., 2019, Liu et al., 2022]. Although over-parameterized networks are ubiquitous in the machine learning literature, it is unclear to what extent this optimality assumption remains realistic beyond such settings. Moreover, the result in Theorem 3.1 does not explicitly relate these optimality errors to the optimization procedure used for the inner and adjoint problems. A precise quantification of these errors would be valuable to strengthen the theoretical foundations of the proposed methods and provide principled guidelines for the choice of hyperparameters. These theoretical considerations do not prevent from applying the method even beyond the setting where the convergence results hold, much like with popular bilevel algorithms such as AID or ITD.
	
**Choice of the function space**. Another important consideration is the choice of the function space in the functional bilevel optimization framework. While we primarily focus on $L_2$ spaces, there are numerous other function spaces that could be explored, such as Reproducing Kernel Hilbert Spaces and Sobolev spaces. Investigating these alternative spaces may reveal additional advantages and open the way for a broader class of machine learning applications where higher-order derivatives of a prediction function appear naturally in the objectives.

**Non-smooth/constrained setting**. The proposed method primarily focuses on smooth and unconstrained problems, but many practical machine learning applications involve non-smooth objectives or constraints. Extending the proposed framework to handle these scenarios would significantly broaden its applicability. Notably, the works [Bolte et al.] on non-smooth implicit differentiation could perhaps be leveraged to adapt our methods to non-smooth settings. Future work should explore these opportunities to further enhance the flexibility and applicability of the functional bilevel optimization approach.

## Additional experiment (see attached PDF). 
As suggested by reviewer *MXrv*, we have compared to a bilevel penalty-based method for the Instrumental Variable application. We use two variants and perform an extensive grid search to adjust their hyper-parameters:

1.  **Gradient penalty** (Eq. 5.1 in [5] (Shen et. al. 2023)): we perform a grid search with the following hyper-parameters: learning rate [0.01, 0.001, 1e-4, 1e-5, 1e-6]; weight decay [0., 0.1, 0.01, 0.001]; penalty constant [1, 0.1, 0.01, 0.001]. Since the method has only a single optimization loop, we increase the number of total iterations to 2000 compared to the other methods (100 outer-iterations). The rest of the parameters are the same for all methods.

2. **Value function penalty** (Eq. 3.2a in [5]). We use the same grid search and test number of inner steps [10,20]. Since the method has a double optimization  loop, we then use 100 outer iterations. The rest of the parameters are the same for all methods.

This method handles non-convexity by considering an optimistic version of the bilevel problem and turning it into a penalised single-level problem. *funcID* performs significantly better than the gradient penalty, whereas it is in the same ballpark as the value function penalty method or better (lower mean on 5k/10k, but higher median for 5k). Notably, the value function penalty seems to have a high variance, with some particularly bad outliers, despite the extensive grid search for tuning hyper-parameters. These conclusions will be added to the paper.

[Allen-Zhu et al. 2019] Zeyuan Allen-Zhu. Yuanzhi Li. Zhao Song. A Convergence Theory for Deep Learning via Over-Parameterization. ICML 2019.

[Liu et al. 2022] Chaoyue Liu, Libin Zhu, Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis. 2022

[Bolte et al., 2022]  Jérôme Bolte, Edouard Pauwels, Samuel Vaiter.  Automatic differentiation of nonsmooth iterative algorithms. Neurips. 2022.

Meta-review: This paper develops a new functional point of view for bilevel optimization problems in machine learning that provides a new insight into solving bilevel optimization problems with nonconvex lower-level objectives by leveraging their strong convexity in the functional space. This development is noteworthy because, although neural networks are nonconvex, the loss function in model training can be convex or strongly convex. The authors have developed a complete paper with theory, practical insights, and plenty of computational tests. Bilevel optimization is a tough topic, but the authors have also made the paper notably accessible.
Decision: Accept (spotlight)