ASDL: A Unified Interface for Gradient Preconditioning in PyTorch

TL;DR

ASDL is a PyTorch extension library that provides a unified interface for various gradient preconditioning methods, facilitating easier implementation, comparison, and study of second-order optimization techniques in deep learning.

Contribution

The paper introduces ASDL, a library that standardizes and simplifies the integration of diverse gradient preconditioning methods in PyTorch, enabling structured analysis and comparison.

Findings

ASDL supports multiple gradient preconditioning algorithms.

It simplifies implementation and experimentation in deep learning.

The library promotes systematic evaluation of second-order methods.

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.