Preconditioning Neural Tangent Kernel for Adaptive Optimization

Xiyuan Yang, Wenxuan Bao, Katherine Tieu, Jingrui He
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:144211-144226, 2026.

Abstract

The Neural Tangent Kernel is a theoretical framework for understanding the training dynamics of neural networks. However, standard NTK and its variants fail to properly depict the finetuning of foundation models, as they neglect the preconditioning effects of adaptive gradients. To bridge this gap, we propose the Optimizer Aware Kernel (OAK), which incorporates the optimizer’s influence into standard NTK framework by a preconditioner estimation technique. Furthermore, we conduct an analysis to answer: when and why kernel regime fails in finetuning. We derive explicit error bounds showing that the collapse of kernel regime is primarily due to the cumulative training effects and the task discrepancy between pretraining and finetuning. Theoretically, we justify OAK’s preconditioner estimation by bounding its error term. Empirically, experiments on various model architectures show both the effectiveness of the OAK method and validity of our arguments on kernel regime collapse.

Cite this Paper


BibTeX
@InProceedings{pmlr-v306-yang26b, title = {Preconditioning Neural Tangent Kernel for Adaptive Optimization}, author = {Yang, Xiyuan and Bao, Wenxuan and Tieu, Katherine and He, Jingrui}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {144211--144226}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/yang26b/yang26b.pdf}, url = {https://proceedings.mlr.press/v306/yang26b.html}, abstract = {The Neural Tangent Kernel is a theoretical framework for understanding the training dynamics of neural networks. However, standard NTK and its variants fail to properly depict the finetuning of foundation models, as they neglect the preconditioning effects of adaptive gradients. To bridge this gap, we propose the Optimizer Aware Kernel (OAK), which incorporates the optimizer’s influence into standard NTK framework by a preconditioner estimation technique. Furthermore, we conduct an analysis to answer: when and why kernel regime fails in finetuning. We derive explicit error bounds showing that the collapse of kernel regime is primarily due to the cumulative training effects and the task discrepancy between pretraining and finetuning. Theoretically, we justify OAK’s preconditioner estimation by bounding its error term. Empirically, experiments on various model architectures show both the effectiveness of the OAK method and validity of our arguments on kernel regime collapse.} }
Endnote
%0 Conference Paper %T Preconditioning Neural Tangent Kernel for Adaptive Optimization %A Xiyuan Yang %A Wenxuan Bao %A Katherine Tieu %A Jingrui He %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-yang26b %I PMLR %P 144211--144226 %U https://proceedings.mlr.press/v306/yang26b.html %V 306 %X The Neural Tangent Kernel is a theoretical framework for understanding the training dynamics of neural networks. However, standard NTK and its variants fail to properly depict the finetuning of foundation models, as they neglect the preconditioning effects of adaptive gradients. To bridge this gap, we propose the Optimizer Aware Kernel (OAK), which incorporates the optimizer’s influence into standard NTK framework by a preconditioner estimation technique. Furthermore, we conduct an analysis to answer: when and why kernel regime fails in finetuning. We derive explicit error bounds showing that the collapse of kernel regime is primarily due to the cumulative training effects and the task discrepancy between pretraining and finetuning. Theoretically, we justify OAK’s preconditioner estimation by bounding its error term. Empirically, experiments on various model architectures show both the effectiveness of the OAK method and validity of our arguments on kernel regime collapse.
APA
Yang, X., Bao, W., Tieu, K. & He, J.. (2026). Preconditioning Neural Tangent Kernel for Adaptive Optimization. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:144211-144226 Available from https://proceedings.mlr.press/v306/yang26b.html.

Related Material