Efficient Bilevel Optimization with KFAC-Based Hypergradients
This addresses the scaling challenge in bilevel optimization for machine learning practitioners, offering a more efficient method with modest overhead.
The paper tackles the computational bottleneck of computing hypergradients in bilevel optimization by incorporating Kronecker-factored approximate curvature (KFAC), which provides curvature-aware hypergradients with better performance-efficiency trade-offs than existing methods like Conjugate Gradient or Neumann expansions, and consistently outperforms unrolling across tasks including meta-learning and AI safety on models up to BERT.
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.