Efficient Bilevel Optimization with KFAC-Based Hypergradients

Disen Liao, Felix Dangel, Yaoliang Yu

TL;DR: We scale bilevel optimization by using KFAC for curvature-aware hypergradients, outperforming unrolling and Neumann/CG, and showing curvature is valuable at scale with modest overhead.

Abstract:

Bilevel optimization (BO) is widely applicable to many machine learning problems.
Scaling BO, however, requires repeatedly computing hypergradients, which involves solving inverse Hessian-vector products (IHVPs). In practice, these operations are often approximated using crude surrogates such as one-step gradient unrolling or identity/short Neumann expansions, which discard curvature information.
We build on implicit function theorem-based algorithms and propose to incorporate Kronecker-factored approximate curvature (KFAC), yielding curvature-aware hypergradients with a better performance efficiency trade-off than Conjugate Gradient (CG) or Neumann methods and consistently outperforming unrolling. We evaluate this approach across diverse tasks, including meta-learning and AI safety problems.
On models up to BERT, we show that curvature information is valuable at scale, and KFAC can provide it with only modest memory and runtime overhead. Our implementation is available at \url{https://github.com/liaodisen/NeuralBo}.

Code Dataset Promise: Yes

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 1745