Nystrom Method for Accurate and Scalable Implicit Differentiation
This work addresses a computational bottleneck in gradient-based bilevel optimization for machine learning practitioners, offering a more stable and efficient alternative to existing methods.
The paper tackles the challenge of efficiently computing inverse Hessian vector products in bilevel optimization by proposing a Nyström method that exploits Hessian low-rankness, achieving stable performance and faster computation than iterative approximations across hyperparameter optimization and meta-learning tasks.
The essential difficulty of gradient-based bilevel optimization using implicit differentiation is to estimate the inverse Hessian vector product with respect to neural network parameters. This paper proposes to tackle this problem by the Nystrom method and the Woodbury matrix identity, exploiting the low-rankness of the Hessian. Compared to existing methods using iterative approximation, such as conjugate gradient and the Neumann series approximation, the proposed method avoids numerical instability and can be efficiently computed in matrix operations without iterations. As a result, the proposed method works stably in various tasks and is faster than iterative approximations. Throughout experiments including large-scale hyperparameter optimization and meta learning, we demonstrate that the Nystrom method consistently achieves comparable or even superior performance to other approaches. The source code is available from https://github.com/moskomule/hypergrad.