Nonsmooth Implicit Differentiation for Machine Learning and Optimization

Introduction

The recent introduction of deep equilibrium networks , the increasing importance of bilevel programming (e.g., hyperparameter optimization) and the ubiquity of differentiable programming (e.g., TensorFlow , PyTorch , JAX ) in modern optimization call for the development of a versatile theory of nonsmooth differentiation. Our focus is on nonsmooth implicit differentiation. There are currently two practices lying at the crossroads of mathematics and computer science: on the one hand the use of the standard smooth implicit function theorem “almost everywhere” and on the other hand the development of algorithmic differentiation tools . The empirical use of the latter in the nonsmooth world has shown surprisingly efficient results , but the current theories cannot explain this success. We bridge this gap by providing nonsmooth implicit differentiation results and illustrating their impact on the training of neural networks and hyperparameter optimization.

Let us consider $z$ implicitly defined through $F(z(x))=h(x)$ where $F$ and $h$ have full domain and adequate dimensions. How does autograd apply to evaluating the “derivative” of the implicitly defined function $z$? Regardless of differentiability or nonsmoothness, and provided that inversion is possible, one commonly uses (or dynamically approximates) this derivative by

applying the implicit differentiation framework of using JAX library, as presented in , provides inconsistency of the derivative at the origin, see Figure 1. As mentioned above, despite these unpredictable outputs, propagating derivatives leads to an undeniable efficiency. But can we parallel these propagation ideas with a simple mathematical counterpart? Is there a rigorous theory backing up formal (sub)differentiation or formal propagation? The answer is positive and was initiated in through conservative Jacobians (see also ).

This property exactly parallels the idea of “propagating derivatives” in practice. It gives a strong meaning to the formal use of Jacobians proposed in , and many empirical approaches .

— We establish a nonsmooth conservative implicit function theorem that comes with an implicit calculus which is the central focus of this paper. Our calculus amounts somehow to formal subdifferentiation with Clarke Jacobians. This approach cannot rely on classical tools like the inverse of a Clarke Jacobian or a composition of Clarke Jacobians, which are not in general Clarke Jacobians. Indeed, a surprising example (Example 1) shows that an “inverse function theorem with Clarke calculus” is not possible.

— We study a wide range of applications of our implicit differentiation theorem, covering deep equilibrium problems , conic optimization layers , and hyperparameter optimization for the Lasso . Each case is detailed and its specificities are discussed.

— As a consequence, we obtain convergence guarantees for mini-batched stochastic algorithms with vanishing step size for training wide classes of Neural Nets, or for Lasso hyperparameter selection. The assumptions needed for our results are mild and fulfilled by most losses occurring in ML in the spirit of : elementary log-exp functions , semialgebraic functions , all being subclasses of definable functions . The use of such structural classes has become standard in nonsmooth optimization and is more and more common in ML (see, e.g., ).

— As in the smooth implicit function theorem, the invertibility condition is not avoidable in general. We provide various examples for which the assumption is not satisfied; this results in severe failures for the corresponding gradient methods. In Figure 1, one sees how lack of invertibility on an otherwise ordinary problem may provide totally unpredictable behavior for smooth quadratic optimization.

Implicit Differentiation with Conservative Jacobians

A locally Lipschitz function is called path differentiable if it has a conservative Jacobian. Recall that the Clarke Jacobian is defined as

Let $J_{F}$ be a conservative Jacobian for $F$, then there is a residual $R$ such that

Propagating derivatives within a nonsmooth function finds its justification in the following:

There is already a long tradition of nonsmooth implicit function theorems, e.g., . What makes the following theorem useful is that it comes with a qualification-free calculus. The proofs are given in Appendix B.

Furthermore, assume that for each $x\in\mathcal{U}$, for each $[A\ B]\in J_{F}(x,G(x))$, the matrix $B$ is invertible where $J_{F}$ is a conservative Jacobian for $F$. Then, $G:\mathcal{U}\to\mathcal{V}$ is path differentiable with conservative Jacobian given, for each $x\in\mathcal{U}$, by

(a) (On the necessity of conservativity) Example 1 in Appendix B shows that one cannot hope for the formulas in Corollaries 1 & 2 to provide Clarke Jacobians in general, even if the input(s) are Clarke Jacobians themselves. (b) (Lipschitz definable implicit and inverse function theorems) See Theorem 4 and 5 in the appendix

Nonsmooth implicit differentiation in Machine Learning

Detailed proof arguments for all considered models are given in Appendix C.

Deep Equilibrium Networks (DEQs) are specific neural network architectures including layers whose input-output relation is implicitly defined through a fixed point equation of the form

has a unique output $z(W,b)$ [50, Theorem 2]. The transformation $(W,b)\mapsto z(W,b)$ is a monotone implicit layer.

The set-valued mapping obtained from Theorem 2 provides a conservative Jacobian for $(W,z)\mapsto z(W,z)$. A similar expression was described in [50, Theorem 2], without using conservativity and using the Clarke Jacobian formally as a classical Jacobian. The proposition below provides a full justification of this heuristic and ensures convergence of algorithmic differentiation based training.

Convexity and invertibility assumptions are satisfied when $J_{\sigma}$ is the Clarke Jacobian .

Optimization layers in deep learning may take many forms; we consider here those based on conic programming . We follow , simplifying the analysis by ignoring infeasability certificates, which correspond to the absence of a primal-dual solution , in line with the implementation described in [2, Appendix B]. Consider a conic problem (P) and its dual (D):

The mapping $\mathcal{N}$ is a synthetic form of optimality measure for (P) and (D), capturing KKT conditions. To simplify the presentation, we ignore the extreme cases of infeasibility and unboundedness which correspond to an absence of solution in .

The following result extends the discussion in , limited to situations where $\Pi$ is differentiable at the proposed solution $z$, to a fully nonsmooth setting; its proof is postponed to Appendix C.2.

In practice, the path differentiability of conic projections is pervasive since they are generally semialgebraic (orthant, second-order cone, PSD cone). See for the computations of the corresponding Clarke Jacobians (which are conservative). Note that a conservative Jacobian for $\mathcal{N}$ may be obtained from $J_{P_{\mathcal{K}^{*}}}$ using Proposition 2.

Implicit differentiation can be used to tune hyperparameters via first-order methods optimizing some measure of task performance, see and references therein. In a nonsmooth context, we recall the formulation in of the general hyperparameter optimization problem as a bi-level optimization problem:

Optimizing implicit problems with gradient descent

We establish the convergence of gradient descent algorithms for compositional learning problems involving implicitly defined functions. The result follows from the previous section and the general convergence results of .

The applications considered in the previous section all yield minimization problems of the type

For $i\in\{1,\ldots,N\}$ and $j\in\{1,\ldots,L\}$, the function $g_{i,j}$ is locally Lipschitz with conservative Jacobian $J_{i,j}$ and one of the following holds

$\bullet$ $g_{i,j}$ and $J_{i,j}$ are semialgebraic (or, more generally, definable).

$\bullet$ $g_{i,j}$ is defined as $G$ in Theorem 2, with $F$ and $J_{F}$ semialgebraic (or, more generally, definable).

Actually, in Assumption 1 the second point implies the first point; we list both for clarity. More details on semialgebraicity and definability are given in Appendix A.2. Let us stress that virtually all elements entering the definition of neural networks are semialgebraic or, more generally, definable, see for example for a constructive model. In particular, beyond classical networks with usual nonlinearities (e.g., relu, sigmoid, max pooling …), this setting encompasses (through Corollary 1):

(a) Deep equilibrium networks: each $g_{i,j}$ may correspond to usual explicit layers or an implicit layer involving a fixed point mapping and a learning sample $i$ as in (4) or (3).

(c) One may assume that $N=1,L=2$ and retrieve the hyperparameter tuning for Lasso in its implicit formulation.

Step size: $\sum_{k=1}^{+\infty}\alpha_{k}=+\infty$ and $\alpha_{k}=o(1/\log(k))$.

This result shows that AD SGD may be applied successfully to all problems described in Section 3, combining algorithmic differentiation with implicit differentiation. Its proof may be adapted directly from ; details are given in Appendix D.

Numerical experiments

Using implicit differentiation when the invertibility condition in Theorem 2 does not hold can result in absurd training dynamics.

Problem (11) has an equivalent fixed-point formulation using projected gradient descent on the inner problem (Appendix E.1.1). Backpropagation applied to (11) associates to $(x,y)$ the following:

We implement gradient descent for (11), evaluating (12) either using cvxpylayers or the JAX tutorial for fixed-point layers. In both cases, the invertibility condition in Theorem 2 fails when $-3x+y+2=0$, resulting in discontinuity of $s$, affecting the dynamics globally: the gradient trajectory converges to a limit cycle of non critical points (Figure 2(a)); see Appendix E.1 for details.

Persistence under small perturbations: The limit cycle remains (Figure 2(b)) when running the same experiment on a slightly perturbed version of Problem (11) with perturbed initializations as well (Appendix E.1.2).

The Lorenz Ordinary Differential Equation (ODE) writes:

It is well-known that taking $(\sigma,\rho,\beta)=(10,28,8/3)$, and $(x(0),y(0),z(0))=(0,1,1.05)$ gives a chaotic trajectory, displayed in Figure 3(a). Denoting $F:(x,y,z)\mapsto(\sigma(y-x),x(\rho-z)-y,xy-\beta z)$ the vector field of the Lorenz system (13), consider the optimization problem:

The function $g:u\mapsto u^{T}F(u)$ is a nondegenerate quadratic function whose expression can be found in Appendix E.2.1. The function $g$ has for unique critical point $(0,0,0)$ which is a strict saddle-point. We perform gradient ascent with implicit differentiation using cvxpylayers on (14), and the classical gradient ascent on the equivalent problem (15). The path obtained by implicit differentiation (Figure 3(b)) resembles the Lorenz attractor (Figure 3(a)), in stark contrast to the conventional method (Figure 3(c)). The chaotic dynamics are a consequence of the lack of invertibility, due to the power $4$ in (14), and various numerical approximations related to optimization and implicit differentiation.

Conclusion and future work

This article provides a rigorous framework and calculus rules for nonsmooth implicit differentiation using the theory of conservative Jacobians. In particular, it describes precise conditions under which implicit differentiation can be used, in a way that is compatible with backpropagation and first-order algorithms.

We show the applicability of our results on practical machine learning problems including training of neural networks involving layers with implicitly defined outputs (deep equilibrium nets, networks with optimization layers) and nonsmooth hyperparameter optimization (Lasso-type models).

Finally, we demonstrate the necessity of a rigorous theory of nonsmooth implicit differentiation through multiple numerical experiments. These illustrate the range of extremely pathological gradient dynamics that can occur when algorithmic differentiation is combined with nonsmooth implicit differentiation outside the scope of our theorem, i.e., without satisfying the invertibility condition we specify.

References

Appendix A Lexicon

$D$ has a closed graph and is locally bounded.

Although conservative fields are not assumed to be locally bounded in , we add this restriction here to ensure they are upper semicontinuous. This will allow us to use a nonsmooth Lyapunov method to prove convergence of first-order algorithms.

A.2 A simpler and more operational view on definability

We recall basic definitions and results on definable sets and functions used in this work. More details on this theory can be found in .

We make a specific attempt to provide a new simple view on this subject by using dictionaries, in the hope that machine learning users consider utilizing these wonderful tools.

where, for $i\in\{1,...,I\}$ and $j\in\{1,...,J\}$, $P_{ij}$ and $Q_{ij}$ are polynomials. The stability properties of semialgebraic sets may be axiomatized to give rise to the general notion of an o-minimal structure:

Denoting by $\pi$ the projection on the $p$ first coordinates, if $A\in\mathcal{O}_{p+1}$ then $\pi(A)\in\mathcal{O}_{p}$.

The elements of $\mathcal{O}_{1}$ are exactly the finite unions of intervals.

Note that the collection of semialgebraic sets verifies 3 in Definition 6 according to the Tarski-Seidenberg theorem.

There are several major structures which have been explored . But rather than relying on traditional description of these structures, we provide instead classes of functions that are contained in an o-minimal structure. The goals achieved are twofold:

The classes we provide are o-minimal and thus all the results provided in the main text apply to functions in these classes.

It is very easy to verify that a function belongs to one of the classes. Everything boils down to checking that the problem under consideration can be expressed in one of the dictionaries we provide.

Note however that we do not aim at providing neither a comprehensive nor a sharp picture of what could be done with o-minimal structures.

We consider first a collection of functions which will serve to establish dictionaries:

Analytic functions restricted to semialgebraic compact domains (contained in their natural open domain), examples are $\cos$ and $\sin$ restricted to compact intervals.

“Globally subanalytic functions”: $\arctan,\tan_{|]-\pi/2,\pi/2[}$ or any functions in (a) (see for a precise definition of global subanalyticity).

With this collection of functions we may build elementary dictionaries. To demonstrate, we consider the following dictionaries

The last dictionary describes a larger class of functions, we shall come back on this later on.

Consider the dictionary $\mathcal{D}={\rm Dic}(\cdot)$ based on the properties (a)-(e) described above.

Then, in the spirit of , we can extend the idea of piecewise selection functions with the following three definitions.

An elementary $\mathcal{D}$-function is a $C^{2}$ function described by a finite compositional expression involving the basic operations $\times,+,/$, multiplication by a constant, and the functions of $\mathcal{D}$ inside their domain of definition.

$(\beta,\lambda)\mapsto\|X\beta-Y\|_{2}+e^{\lambda}\|\beta\|_{1}$.

where, for $i\in\{1,\ldots,I\}$ and $j\in\{1,\ldots,J\}$, the $g_{ij}$ and $h_{ij}$ are elementary $\mathcal{D}$-functions.

We denote $\mathcal{P}\mathcal{D}$ the set of piecewise $\mathcal{D}$-functions. With the assumptions we have on the dictionary $\mathcal{D}$, the piecewise selections we consider are all definable (it ’s not always the case in general). Notice that piecewise log-exp functions are a specific case of $\mathcal{D}$-functions with the dictionary $\mathcal{D}={\rm Dic}(\mbox{\ref*{list:logexp}})=\left\{\log,\exp\right\}$. It is easy to see that the following functions are in $\mathcal{P}\mathcal{D}$ and thus definable:

$x\mapsto\begin{cases}\frac{1}{2}{x^{2}}&\text{for }|x|\leq\delta,\\ \delta(|x|-\frac{1}{2}\delta),&\text{otherwise,}\end{cases}$ with $\delta>0$ (Huber loss).

Appendix B Results from Section 2

for each $x\in\mathcal{U}$. Finally, by graph closedness and convexity of $J_{F}$ we get, for each $x\in\mathcal{U}$,

Let $J_{F}$ be a conservative Jacobian for $F$, then there is a residual $R$ such that

Furthermore, assume that for each $x\in\mathcal{U}$, for each $[A\ B]\in J_{F}(x,G(x))$, the matrix $B$ is invertible where $J_{F}$ is a conservative Jacobian for $F$. Then, $G:\mathcal{U}\to\mathcal{V}$ is path differentiable with conservative Jacobian given, for each $x\in\mathcal{U}$, by

Proof: Let $\gamma:\to\mathcal{U}$ be absolutely continuous, then the composition $G\circ\gamma$ is also absolutely continuous since $G$ is locally Lipschitz. By (16) we have, for all $t\in$,

which we can differentiate almost everywhere; for almost every $t\in$, for any $[A\ B]\in J_{F}(\gamma(t),G(\gamma(t)))$,

Since $B$ is assumed to be invertible, we have, for almost every $t\in$,

The set-valued mapping $J_{G}\colon x\rightrightarrows\left\{-B^{-1}A:[A\ B]\in J_{F}(x,G(x))\right\}$ is nonempty, locally bounded, and has a closed graph for each $x\in\mathcal{U}$ since $J_{F}(x,G(x))$ is a conservative Jacobian and $B$ is invertible . We conclude that $G$ is path differentiable on $\mathcal{U}$ with conservative Jacobian $J_{G}$. $\Box$

Proof: Consider the function $F(x,y)=x-\Phi(y)$ and observe that it satisfies the assumptions of Corollary 1, so that we obtain a function $G$ which is exactly the desired inverse. $\Box$

It is tempting to think that Corollary 2 should come with a formula of the type

for all $z$ in a neighborhood of . This happens to be false, making the use of the notion of conservativity necessary to catpure the artifacts resulting from application of ordinary calculus rules to nonsmooth inverse functions. Note that since the inverse function theorem is a special case of the implicit function theorem, this also rules out a Clarke calculus for implicit functions.

It is locally Lipschitz and semialgebraic and thus path differentiable with its Clarke Jacobian a conservative Jacobian. We have the following explicit piecewise linear representation

from which we deduce that the Clarke Jacobian of $\Phi$ has the following structure

Using the above, one can verify that $\Phi$ is a homeomorphism whose inverse is also piecewise linear. We set $\Psi=\Phi^{-1}$; it is given by

A graphical representation of these sets is given in Figure 4.

From this explicit piecewise linear representation of $\Psi$, we deduce that its Clarke Jacobian at is the following

In the definable (e.g. semialgebraic case) our results have a remarkably simple expression that we give below.

Proof: It suffices to use the fact that definable mappings are path differentiable, see , and that the the graph of $\Psi$ is given by a first-order formula. $\Box$

Moreover, for each $x\in\mathcal{U}$, the mapping $J_{G}\colon x\rightrightarrows\left\{-B^{-1}A:[A\ B]\in J_{F}(x,G(x))\right\}$ is conservative for $G$.

Appendix C Results from Section 3

Proof: The quantity $z(W,b)$ is defined implicitly by the relation

so that $\mathcal{B}$ is infinitely differentiable. Equation (18) is then equivalent to

C.2 Optimization layers: the conic program case

Let us first expand on the link between zeros of the residual map and KKT solutions. We provide a simplified view of , ignoring cases of infeasibility and unboundedness. Note that this corresponds to enforcing $w=1$ as done in .

$v=s+y$, $s\in\mathcal{K}$, $y\in\mathcal{K}^{\circ}$, $s^{T}y=0$.

$s=P_{\mathcal{K}}(v)$, $y=P_{\mathcal{K}^{\circ}}(v)$.

We may reformulate this equivalence as follows, using changes of signs on $y$ and $v$, noticing that $-P_{\mathcal{K}^{\circ}}(-\cdot)=P_{\mathcal{K}^{*}}(\cdot)$ since $\mathcal{K}^{*}=-\mathcal{K}^{\circ}$,

$v=y-s$, $s\in\mathcal{K}$, $y\in\mathcal{K}^{*}$, $s^{T}y=0$.

$s=P_{\mathcal{K}^{*}}(v)-v$, $y=P_{\mathcal{K}^{*}}(v)$.

Now the KKT system in $(x,y,s)$ for problem (P) and (D) can be written as follows (see, for example, ),

which is equivalent, by setting $v=y-s$ and $u=x$, to

The system (19) is equivalent to $\mathcal{N}(z,A,b,c)=0$ with $z=(u,v)$. We have shown that $(x,y,s)$ is a KKT solution to the system if and only if $(x,y,s)=(u,P_{\mathcal{K}^{*}}(v),P_{\mathcal{K}^{*}}(v)-v)=\phi(z)$ for $z=(x,y-s)$ such that $\mathcal{N}(z,A,b,c)=0$.

Let us now turn to $\phi$. Since $P_{\mathcal{K}^{*}}$ has for conservative Jacobian $J_{P_{\mathcal{K}^{*}}}$, we may construct a conservative Jacobian for the function $\phi$ as follows using [14, Lemmas 3, 4, and 5]:

C.3 Hyperparameter selection for nonsmooth Lasso-type model

Thus a conservative Jacobian for $F$ at $(\lambda,\beta)$ is given by

with $\mathcal{C}:=\{q:q_{i}\in\mathds{1}_{e^{\lambda}}\left(\beta_{i}-X_{i}^{T}\left(X\beta-y\right)\right)\}$. Let us estimate the factors $q_{i}$ above in terms of the equicorrelation set $\mathcal{E}$. Recall the KKT conditions for the Lasso problem; a solution $\hat{\beta}$ must satisfy

Taking $q_{i}=1$ for all $i\in\mathcal{E}$ gives a selection of the conservative Jacobian for $\hat{\beta}$ in Proposition 5, for all $j\in\left\{1,\ldots,p\right\}$,

Appendix D Results from Section 4

Step size: $\sum_{k=1}^{+\infty}\alpha_{k}=+\infty$ and $\alpha_{k}=o(1/\log(k))$.

Appendix E Results from Section 5

Where $P_{\mathcal{U}}$ is the projection on the set $\mathcal{U}:=\left[0,3\right]\times\left[0,5\right]$ which can be implemented as a difference of relu functions

E.1.2 Perturbed experiments

Perturbed experiments are done on the following perturbed loss function

with $\varepsilon_{1},\ldots,\varepsilon_{6}$ the perturbations. In Figure 2(b), we consider several realizations of independent Gaussian variables $\varepsilon_{1},\ldots,\varepsilon_{6}\sim\mathcal{N}(0,\sigma^{2})$ with $\sigma^{2}=0.05$; despite this added noise, the unwanted dynamics persist.

E.1.3 Conic canonicalization

It can be formulated as a cone program (P) and its dual (D):

Let $(x,y,s)$ be a solution to the cone program (24) where $x$ is the primal variable, $y$ is the dual variable, and $s$ the primal slack variable. Then it follows from (6) that a solution $z$ to $\mathcal{N}(z,c)=0$ is obtained by $z=(x,y-s)$. For $c=(0,0)$, the solutions are $x\in\left[0,3\right]\times\left[0,5\right]$, $s=b-Ax$, and $y=(0,0,0,0)$, hence the uniqueness assumption for Proposition 4 is not satisfied.

This will combine the two cycles but the parameter $\eta$ will make one cycle “faster” than the other. Projecting the path of the gradient descent on the variables $(y,z)$ we obtain a chaotic dynamics filling the space as the number of iterations increases.

E.2 Lorenz-like attractor

where $H=\begin{bmatrix}-2\sigma&\sigma+\rho&0\\ \sigma+\rho&-2&0\\ 0&0&-2\beta\end{bmatrix}$.

For $(\sigma,\rho,\beta)=(10,28,\frac{8}{3})$, $g$ has for unique critical point $(0,0,0)$ which is a strict saddle-point.

E.3 License of assets used

All assets used: cvxpy, cvxpylayers, and JAX were released under the Apache License, Version 2.0, January 2004, http://www.apache.org/licenses/.