Functional Bilevel Optimization for Machine Learning

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.

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. 1. 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. 2. 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). 3. 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. 2. 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). 3. 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.) 4. 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. 5. FuncID seems to require more outer iterations to converge than baselines (Figure 1b), which I don't see being discussed by the authors.

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. 2. Why is the prediction error of FuncID worse than MLE in Figure 7b? Notes: 1. In (FBO), there's a space between the colon and the equals (use \coloneqq). 2. In line 94, link to where the assumption is discussed. 3. Font size in the figures is very small and hard to read.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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.

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.

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. - 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. - The assumptions necessary for the theoretical framework might not be easily verifiable in practical scenarios, potentially limiting the method's applicability.

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? 2. Please explain how does the choice of function space impact the stability and convergence of the FuncID method?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Authors addressed their work limitation.

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.

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. 2. 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]. [1] K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design. ICML, 2021. [2] T. Chen, Y. Sun, and W. Yin. Closing the Gap: Tighter Analysis of Alternating Stochastic Gradient Methods for Bilevel Problems. NeurIPS, 2021. [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. 2. 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. [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. [5] H. Shen, and T. Chen. On Penalty-based Bilevel Gradient Descent Method. ICML 2023.

Questions

Same as weakness.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Same as weakness.

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.

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

Questions

1. On Line 126, there is a typo: "elemen" should be replaced by "element". 2. On Line 629, capitalize the word "euclidean". 3. 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. 4. What is the rationale for the label "FuncID linear" in Figure 1? Why do you refer to it as "linear"? 5. Why do you feel "FuncID linear" converge faster in terms of outer iterations than "FuncID" in Figure 1?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

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.

Discussion period

Dear Reviewers, Thanks for the contributions in the review period and looking forward to a fruitful discussion on Submission 6898 in the next few days. Kind regards, Your AC

Reviewer response?

Reviewer MXrv, could you please review the authors' response and see whether it addresses your questions? Please acknowledge having done so in a comment. Thanks.

Discussion period authors' comment

Thank you to all reviewers for responding to our rebuttal, and to the area chair for encouraging the discussion on our submission. We will ensure that the changes mentioned in the rebuttal and confirmed by the reviewers are integrated into the final version of the paper, if it is accepted.