大型矩阵函数的梯度
机器学习
2024-10-28 v2 数值分析
数值分析
机器学习
摘要
调节科学和概率机器学习模型——例如偏微分方程、高斯过程或贝叶斯神经网络——往往依赖于随数据集大小或参数数量而增长的矩阵函数的评估。虽然评估这些量的现有方法几乎总是基于 Lanczos 和 Arnoldi 迭代,但本文首次解释了如何高效地对这些数值线性代数的基石进行微分。为此,我们推导了先前未知的 Lanczos 和 Arnoldi 迭代的 adjoint 系统,实现了它们在 JAX 中的应用,表明所得代码在对偏微分方程进行微分、选择高斯过程模型方面能够与 Diffrax 竞争,并且在校准贝叶斯神经网络时优于标准分解方法。所有这一切都通过无需任何问题特定代码优化即可实现。代码可在 https://github.com/pnkraemer/experiments-lanczos-adjoints 获取,通过 pip install matfree 安装该库。
引用
@article{arxiv.2405.17277,
title = {Gradients of Functions of Large Matrices},
author = {Nicholas Krämer and Pablo Moreno-Muñoz and Hrittik Roy and Søren Hauberg},
journal= {arXiv preprint arXiv:2405.17277},
year = {2024}
}