ASDL: A Unified Interface for Gradient Preconditioning in PyTorch
Machine Learning
2023-05-09 v1
Abstract
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.
Cite
@article{arxiv.2305.04684,
title = {ASDL: A Unified Interface for Gradient Preconditioning in PyTorch},
author = {Kazuki Osawa and Satoki Ishikawa and Rio Yokota and Shigang Li and Torsten Hoefler},
journal= {arXiv preprint arXiv:2305.04684},
year = {2023}
}