[Paper Review] Provably Correct Automatic Subdifferentiation for Qualified Programs

Sham M. Kakade, Jason D. Lee|arXiv (Cornell University)|Sep 23, 2018

Numerical Methods and Algorithms12 references4 citations

TL;DR

This paper proposes a provably correct automatic subdifferentiation method for a class of nonsmooth, piecewise-defined functions—commonly used in machine learning—by extending automatic differentiation to compute generalized (Clarke) subgradients. It establishes a 'Cheap Subgradient Principle' showing that subgradients can be computed at most 6 times the cost of evaluating the function itself, under mild constraint qualifications, ensuring correctness even at nondifferentiable points where existing libraries fail.

ABSTRACT

The Cheap Gradient Principle (Griewank 2008) --- the computational cost of computing the gradient of a scalar-valued function is nearly the same (often within a factor of $5$) as that of simply computing the function itself --- is of central importance in optimization; it allows us to quickly obtain (high dimensional) gradients of scalar loss functions which are subsequently used in black box gradient-based optimization procedures. The current state of affairs is markedly different with regards to computing subderivatives: widely used ML libraries, including TensorFlow and PyTorch, do not correctly compute (generalized) subderivatives even on simple examples. This work considers the question: is there a Cheap Subgradient Principle? Our main result shows that, under certain restrictions on our library of nonsmooth functions (standard in nonlinear programming), provably correct generalized subderivatives can be computed at a computational cost that is within a (dimension-free) factor of $6$ of the cost of computing the scalar function itself.

Motivation & Objective

• To address the lack of correctness in subgradient computation by popular ML libraries like TensorFlow and PyTorch, especially at nondifferentiable points.
• To investigate whether a 'Cheap Subgradient Principle'—analogous to the Cheap Gradient Principle for smooth functions—can be established for nonsmooth, locally Lipschitz functions.
• To develop a computational model and algorithmic framework that provably computes generalized (Clarke) subgradients for qualified programs, ensuring correctness across equivalent program representations.
• To extend automatic differentiation to handle nonsmooth functions such as ReLU, absolute value, and linear algebra operations (e.g., SVD, eigenvalues) with guaranteed subgradient correctness.
• To lay the foundation for numerically stable and correct subgradient computation in safety-critical applications like physical simulators and online learning systems.

Proposed method

• The method extends reverse-mode automatic differentiation to compute limiting total derivatives via a chain rule that respects the structure of piecewise functions.
• It introduces a new computational model where each function is represented as a finite sum over regions defined by linear inequalities, with smooth components (polynomials) on each region.
• For each function, the algorithm determines the active region based on the direction of the input perturbation, using a limiting process to select the correct subgradient.
• The core algorithm, Algorithm 6, computes the subgradient by evaluating the gradient of the active polynomial component in the direction of the input perturbation.
• Library functions are overloaded via Algorithm 7, which returns the function value, directional derivative, and a subgradient by identifying the region that is 'activated' in the direction of the perturbation.
• The method relies on constraint qualifications (Assumption 3.1) to ensure that the active region is well-defined and that the limiting subgradient is consistent across equivalent program representations.

Experimental results

Research questions

• RQ1 Can a provably correct automatic subdifferentiation method be designed for a broad class of nonsmooth, piecewise-defined functions used in machine learning?
• RQ2 Is there a 'Cheap Subgradient Principle' such that computing the generalized subgradient costs at most a constant factor (independent of dimension) more than computing the function itself?
• RQ3 Why do existing libraries like TensorFlow and PyTorch produce inconsistent subgradients for mathematically equivalent programs, and can this be formally resolved?
• RQ4 Can the framework be extended to handle nonsmooth linear algebra operations such as SVD, eigenvalues, and QR decompositions with correct subgradient computation?
• RQ5 What conditions ensure numerical stability and correctness in subgradient computation, especially in the presence of roundoff errors and near-degenerate cases?

Key findings

• The paper establishes a 'Cheap Subgradient Principle' showing that generalized subgradients can be computed at most 6 times the cost of evaluating the function, independent of input dimension.
• The method correctly computes subgradients even at nondifferentiable points—such as the origin in ReLU(x) or |x|—where current libraries like TensorFlow and PyTorch fail.
• For mathematically equivalent programs (e.g., f1(x)=x, f2(x)=ReLU(x)−ReLU(−x), f3(x)=10f1(x)−9f2(x)), the method produces consistent subgradients, unlike existing libraries.
• The algorithm guarantees correctness by using a limiting process based on the direction of perturbation, ensuring that the subgradient corresponds to the correct active region.
• The framework is extendable to univariate piecewise polynomials and can be generalized to more complex functions under constraint qualifications.
• The approach provides a foundation for developing numerically stable, provably correct subdifferentiation in safety-critical applications such as physical simulation and online learning.