Nystrom Method for Accurate and Scalable Implicit Differentiation

TL;DR

This paper introduces a Nystrom-based method for implicit differentiation in bilevel optimization, offering a stable, efficient alternative to iterative approaches like conjugate gradient, with demonstrated superior performance in large-scale tasks.

Contribution

It proposes a novel Nystrom and Woodbury matrix identity-based approach for inverse Hessian vector products, improving stability and efficiency over existing iterative methods.

Findings

Achieves stable and faster inverse Hessian computations.

Performs comparably or better in hyperparameter optimization and meta learning.

Reduces numerical instability compared to iterative methods.

Abstract

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.