Efficient Bilevel Optimization with KFAC-Based Hypergradients

Disen Liao, Felix Dangel, Yaoliang Yu
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:4042-4050, 2026.

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

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-liao26b, title = { Efficient Bilevel Optimization with KFAC-Based Hypergradients }, author = {Liao, Disen and Dangel, Felix and Yu, Yaoliang}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {4042--4050}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/liao26b/liao26b.pdf}, url = {https://proceedings.mlr.press/v300/liao26b.html}, 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}. } }
Endnote
%0 Conference Paper %T Efficient Bilevel Optimization with KFAC-Based Hypergradients %A Disen Liao %A Felix Dangel %A Yaoliang Yu %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-liao26b %I PMLR %P 4042--4050 %U https://proceedings.mlr.press/v300/liao26b.html %V 300 %X 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}.
APA
Liao, D., Dangel, F. & Yu, Y.. (2026). Efficient Bilevel Optimization with KFAC-Based Hypergradients . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:4042-4050 Available from https://proceedings.mlr.press/v300/liao26b.html.

Related Material