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.

From functional to parametric

Our 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.

A. Challenges of bilevel optimization with deep networks

• 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).
• 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.

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.

• 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.

• 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.

Other comments

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).

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. We 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.

Experimental results:

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.

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.

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).

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).

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.

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.

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).

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.

[Bach 2017] Bach, F. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19), 1-53 2017.

[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

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.

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.

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.

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.

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).

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.

Typos. We will fix the typos in the final version.

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.

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.

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.

Thank you for your responses to my questions and for the clarifications! Good to hear that the typos will be fixed in the final draft and that the missing details will be added to the final draft! I am also glad that you are adding a detailed discussion/conclusion at the end as it is important to leave the reader with some take-away points from your paper. Good luck!