ASDL: A Unified Interface for Gradient Preconditioning in PyTorch
Kazuki Osawa Satoki Ishikawa Rio Yokota Shigang Li Torsten Hoefler
Gradient preconditioning is a key technique to integrate the second-order
information into gradients for improving and extending gradient-based learning
algorithms. In deep learning, stochasticity, nonconvexity, and high
dimensionality lead to a wide variety of gradient preconditioning methods, with
implementation complexity and inconsistent performance and feasibility. We
propose the Automatic Second-order Differentiation Library (ASDL), an extension
library for PyTorch, which offers various implementations and a plug-and-play
unified interface for gradient preconditioning. ASDL enables the study and
structured comparison of a range of gradient preconditioning methods.