Efficient Bilevel Optimization with KFAC-Based Hypergradients

Disen Liao; Felix Dangel; Yaoliang Yu

cs.LGMar 31, 2026

Efficient Bilevel Optimization with KFAC-Based Hypergradients

Authors: Disen Liao, Felix Dangel, Yaoliang Yu

Organizations: University of Waterloo · Vector Institute

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.

Figures & tables

Panel of Figure 1 — Figure 1: Top : KFAC balances accuracy and efficiency better than Hessian/GGN-based IHVP approximations. $T$ is the number of CG iterations and $K$ the number of truncated Neumann terms. Bottom : On BERT, KFAC incurs only a small memory overhead compared to the curvature-free SAMA, while using less memory than CG/Neumann.

Figure 2: Results on data hypercleaning. Left two panels: existing BO algorithms improve when their IHVP solvers are replaced with IKVPs. Right two panels: KFAC-based methods compared against other non-IFT baselines, showing faster convergence in both time and iterations. Curves are truncated at the lowest test loss reached. Solid black line is trained on the validation dataset, while dashed blackline is trained on the validation dataset and the noisy dataset together.

Method	$d=10$	$d=100$	$d=500$
KFAC	$7.62\text{\times}{10}^{-3}$	$7.62\text{\times}{10}^{-3}$	$7.62\text{\times}{10}^{-3}$
Neu-3	$1.12\text{\times}{10}^{-1}$	$2.45\text{\times}{10}^{-1}$	$8.01\text{\times}{10}^{-1}$
Neu-20	$2.22\text{\times}{10}^{-2}$	$4.46\text{\times}{10}^{-2}$	$2.41\text{\times}{10}^{-1}$
Neu-50	$1.01\text{\times}{10}^{-4}$	$5.62\text{\times}{10}^{-3}$	$1.39\text{\times}{10}^{-1}$
CG-3	$1.05\text{\times}{10}^{-3}$	$7.49\text{\times}{10}^{-2}$	$3.02\text{\times}{10}^{-1}$
CG-5	$6.31\text{\times}{10}^{-4}$	$7.74\text{\times}{10}^{-3}$	$1.96\text{\times}{10}^{-1}$

Table 1: Relative error for different IHVP approximations as the parameter dimension

d

increases. Neu-

K

denotes a truncated Neumann series with

K

terms. CG-

T

denotes conjugate gradient with

T

iterations. Identity corresponds to the identity approximation used in one-step unrolling.

Figure 3: Batch size effects on data hypercleaning. KFAC consistently achieves lower minimum test loss than CG across models and batch sizes. For linear model, CG will converge early to a worse point than KFAC.

$\tau$	Best Test Loss		Avg. iter Time (s)
$\tau$	KFAC	EKFAC	KFAC	EKFAC
1	1.24	1.21	0.69	1.83
5	1.33	1.24	0.28	0.62
10	1.35	1.30	0.26	0.36
50	1.36	1.33	0.22	0.28

Table 2: Best test loss and average per-iteration runtime for KFAC and EKFAC with varying update intervals (

\tau

) on ResNet-18. Amortization significantly reduces computational cost while maintaining competitive performance.

Dataset	CIFAR-10					CIFAR-100
IF	200	100	50	20	10	200	100	50	20	10
Base Model	65.3	70.0	74.2	81.8	86.1	34.4	38.5	43.4	50.7	55.5
Focal Loss	65.2	69.9	74.3	81.5	85.9	34.9	38.4	43.7	51.8	55.6
MWN + $T_{1}$ - $T_{2}$	66.3	72.8	78.6	84.8	87.5	36.0	40.2	46.0	52.0	58.6
MWN + SAMA	66.3	72.9	79.5	85.4	87.6	35.8	40.0	45.7	52.5	59.7
MWN + KFAC (ours)	67.0	73.1	78.9	85.6	88.4	35.6	40.9	46.2	53.2	60.8

Table 3: ResNet-32 accuracy (%) on long-tailed CIFAR-10/100 with different imbalanced factors (IFs), best is in bold .

Figure 4: Results on BERT data hypercleaning. We report test accuracy (y-axis) versus total training time (x-axis) on three datasets when fine-tuning 1, 7, 12 encoder layers (columns). Each marker corresponds to one hypergradient solver. The gray region highlights the Pareto frontier.

	ChemProt	HyperPartisan	ACL-ARC	SciERC
Baseline	82.54 ±0.32	90.05 ±1.24	67.85 ±2.03	79.87 ±0.59
TARTAN-MT	83.97 ±0.25	94.64 ±0.91	72.86 ±3.01	80.49 ±0.82
SAMA	83.75 ±0.24	95.18 ±0.03	71.75 ±1.65	81.06 ±0.09
KFAC (ours)	85.04 ±0.23	95.13 ±0.05	72.98 ±1.58	81.22 ±0.14

Table 4: Experiment results for auxiliary learning with the continued pretraining task. Following Choe et al. (2023b) , we report test micro-F1 for ChemProt and macro-F1 for other datasets; the result is averaged over 3 runs.

Model	Dataset	Clean	TGDA Hessian $\to$ KFAC (Ours)		BOME (VF)		Label Flip
Model	Dataset	Acc	Acc Drop	time	Acc Drop	time	Acc Drop	time
LR	MNIST	92.34	$2.78_{\pm 0.04}\to\textbf{5.40}_{\pm 0.06}$	0.7 $\to$ 0.4 hrs	$2.22_{\pm 0.04}$	0.4 hrs	$1.52_{\pm 0.03}$	-
MLP	MNIST	98.02	$1.48_{\pm 0.04}\to\textbf{3.99}_{\pm 0.05}$	7.3 $\to$ 3.2 hrs	$1.63_{\pm 0.05}$	3.0 hrs	$0.05_{\pm 0.01}$	-
CNN	CIFAR-10	64.30	$2.04_{\pm 0.03}\to\textbf{3.44}_{\pm 0.07}$	21.5 $\to$ 6.2 hrs	$2.15_{\pm 0.03}$	5.9 hrs	$0.17_{\pm 0.02}$	-

Table 5: Accuracy drop (%) and training time of different data poisoning algorithms on LR, MLP, and CNN. Results are averaged over 5 runs. Label flipping requires no training, so time is omitted.

	ResNet18	CNN	Time	Bilevel?
None (clean)	92.36	86.89	-	-
AID-CG	78.23	57.32	11.5h	✓
BOME	48.09	26.34	2.2h	✓
stocBiO	84.92	75.13	6.7h	✓
AID-KFAC (ours)	42.33	22.96	3.8h	✓
EM	32.71	21.83	$<$ 0.5h	✗

Table 6: Unlearnable examples with different BO algorithms and non-bilevel algorithms. The unlearnable examples trained on CNN and tested on ResNet18 and CNN.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

1 Layer (6.5% Params)
Method	Trec		SemEval		Youtube
	Acc(%)	Time (s)	Acc(%)	Time (s)	Acc(%)	Time (s)
Baseline	60.12 ±0.43	73.3 ±0.1	78.94 ±0.87	121.2 ±0.3	68.53 ±1.24	124.6 ±1.2
KFAC	77.62 ±1.28	772.8 ±3.4	83.83 ±0.47	840.5 ±3.8	81.33 ±1.55	955.7 ±4.4
KFAC EMP	74.25 ±2.57	791.5 ±3.5	81.61 ±0.75	838.6 ±4.5	80.53 ±1.32	962.8 ±3.2
SAMA	62.93 ±0.84	602.7 ±2.2	81.61 ±1.46	628.0 ±5.7	77.86 ±1.15	630.1 ±5.3

Appendix

Table 7: Comparison of Methods on Trec , SemEval , and Youtube at Different Tuning Levels. Each entry shows Accuracy and Training Time. Each result is averaged over 5 runs. These are the same results as Figure 4 but in numbers.

Figure 5: Data hypercleaning results for CNN on CIFAR-10 and linear on MNIST. For the linear model, Solid black line is the test loss reached on the cleaned dataset, while the dashed black line is trained only on validation set. For CNN, solid black line is test loss reached on validation dataset, while dashed line is trained on the validation dataset and noisy dataset.

Panel of Figure 6 — Figure 6: Effect of EMA coefficient $\beta$ on convergence for data hypercleaning with logistic regression on MNIST. Each curve corresponds to a different $\beta\in\{0,0.3,0.5,0.9\}$ . Left column : batch size 1000. Right column : batch size 50000. In this setting, EMA is most helpful for small damping and small batch sizes, where curvature estimates are noisier, while its effect is modest for larger batches or stronger damping.

Figure 7: Convergence of AID-CG and AID-KFAC on data hypercleaning for different batch sizes. KFAC achieves faster and more stable convergence, particularly for smaller batches, where curvature estimates are noisier and the inner problem becomes ill-conditioned. Compared with Figure 3 , this figure shows the evolution of the test loss for different batch sizes.

Kronecker-Factored Approximation

See all 14 papers

Newton-Based Methods for Bilevel Optimization Beyond Lower-Level Strong Convexity
Oct 19, 2026Youran Dong, Junfeng Yang, Wei Yao +1Bilevel Optimization Second-Order Optimization
School of Mathematics, Nanjing University, Nanjing, P. R. China. · School of Artificial Intelligence, Southern University of Science and Technology, Shenzhen, P. R. China. · Department of Mathematics and National Center for Applied Mathematics Shenzhen, Southern University of Science and Technology, Shenzhen, P. R. China.
Sharper Analysis of Single-Loop Methods for Bilevel Optimization
Jul 11, 2026Yubo Zhou, Jun Shu, Luo Luo +4Implicit Differentiation Optimization Convergence Analysis
School of Mathematics and Statistics, Xi’an Jiaotong University · School of Data Science, Fudan University · 3SGIT AI Lab, State Grid Corporation of China +1
To Solve Bilevel Optimization with Nonconvex Lower Levels, We Need Second-Order Stationarity
Sep 24, 2026Zhiyao Zhang, Menglu Yu, Alvaro Velasquez +2Bilevel Optimization Second-Order Optimization
The Ohio State University · Meta · University of Colorado Boulder +1

Oct 19, 2026J/K move · Enter open · S save

Newton-Based Methods for Bilevel Optimization Beyond Lower-Level Strong Convexity

Youran Dong, Junfeng Yang, Wei Yao, Jin Zhang

School of Mathematics, Nanjing University, Nanjing, P. R. China. · School of Artificial Intelligence, Southern University of Science and Technology, Shenzhen, P. R. China. · Department of Mathematics and National Center for Applied Mathematics Shenzhen, Southern University of Science and Technology, Shenzhen, P. R. China.

Bilevel Optimization Second-Order Optimization Hypergradient Optimization

Open paper AskPDF arXiv

Bilevel optimization provides a framework for machine learning tasks such as hyperparameter optimization, data reweighting, and meta-learning. Existing Newton-based bilevel methods typically rely on the classical hypergradient formula, which is generally unavailable when lower-level solutions are nonunique and the lower-level Hessian is singular. We develop two Newton-based methods beyond lower-level strong convexity by exploiting curvature shared by the lower-level and multiplier updates. Newton-based GALET handles potentially singular lower-level curvature under a positive spectral-gap condition, exploiting its range structure to avoid the squared-Hessian multiplier system used in GALET. Newton-based BAMM uses a strongly convex aggregation model when the lower-level objective is convex and the upper-level objective is strongly convex in the lower-level variable. We establish convergence guarantees and improved complexity bounds relative to GALET and sl-BAMM under the corresponding assumptions. We further introduce GNBO, a generalized Newton-based framework that captures the shared-curvature structure of the proposed methods and several existing Newton and quasi-Newton bilevel methods. Experiments on synthetic problems, data hyper-cleaning, and GPT-2 data reweighting evaluate the proposed methods across different bilevel learning tasks.

Abstract

Kronecker-Factored Approximation

Explore similar work

Newton-Based Methods for Bilevel Optimization Beyond Lower-Level Strong Convexity

Sharper Analysis of Single-Loop Methods for Bilevel Optimization

To Solve Bilevel Optimization with Nonconvex Lower Levels, We Need Second-Order Stationarity