Efficient Bilevel Optimization with KFAC-Based Hypergradients

TL;DR

This paper introduces a KFAC-based method for efficient hypergradient computation in bilevel optimization, improving scalability and performance over traditional approximations, especially on large models like BERT.

Contribution

It proposes integrating KFAC into implicit function theorem-based algorithms for curvature-aware hypergradients, outperforming existing methods in efficiency and accuracy.

Findings

KFAC-based hypergradients outperform unrolling and Neumann methods.

Curvature information is beneficial at large scale, including BERT models.

Implementation available at https://github.com/liaodisen/NeuralBo.

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 https://github.com/liaodisen/NeuralBo.