ASDL: A Unified Interface for Gradient Preconditioning in PyTorch
This work addresses the implementation complexity for researchers and practitioners in machine learning, but it is incremental as it focuses on software tooling rather than algorithmic breakthroughs.
The authors tackled the problem of inconsistent implementation and performance of gradient preconditioning methods in deep learning by proposing ASDL, a PyTorch extension library that provides a unified interface and various implementations, enabling structured comparison and study of these methods.
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.