Nonsmooth Implicit Differentiation for Machine Learning and Optimization
Year: 2021

Abstract:
In view of training increasingly complex learning architectures, we establish a nonsmooth implicit function theorem with an operational calculus. Our result applies to most practical problems (i.e., definable problems) provided that a nonsmooth form of the classical invertibility condition is fulfilled. This approach allows for formal subdifferentiation: for instance, replacing derivatives by Clarke Jacobians in the usual differentiation formulas is fully justified for a wide class of nonsmooth problems. Moreover this calculus is entirely compatible with algorithmic differentiation (e.g., backpropagation). We provide several applications such as training deep equilibrium networks, training neural nets with conic optimization layers, or hyperparameter-tuning for nonsmooth Lasso-type models. To show the sharpness of our assumptions, we present numerical experiments showcasing the extremely pathological gradient dynamics one can encounter when applying implicit algorithmic differentiation without any hypothesis.

text: In view of training increasingly complex learning architectures, we establish a nonsmooth implicit function theorem with an operational calculus. Our result applies to most practical problems (i.e., definable problems) provided that a nonsmooth form of the classical invertibility condition is fulfilled. This approach allows for formal subdifferentiation: for instance, replacing derivatives by Clarke Jacobians in the usual differentiation formulas is fully justified for a wide class of nonsmooth problems. Moreover this calculus is entirely compatible with algorithmic differentiation (e.g., backpropagation). We provide several applications such as training deep equilibrium networks, training neural nets with conic optimization layers, or hyperparameter-tuning for nonsmooth Lasso-type models. To show the sharpness of our assumptions, we present numerical experiments showcasing the extremely pathological gradient dynamics one can encounter when applying implicit algorithmic differentiation without any hypothesis.

heading: 1 Introduction
text: Differentiable programming. The recent introduction of deep equilibrium networks [7], the increasing importance of bilevel programming (e.g., hyperparameter optimization) [48] and the ubiquity of differentiable programming (e.g., TensorFlow [1], PyTorch [47], JAX [16]) 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” [31, 30] and on the other hand the development of algorithmic differentiation tools [2, 3, 59]. The empirical use of the latter in the nonsmooth world has shown surprisingly efficient results [59], 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.
Backpropagation: a formal differentiation approach. 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
(backpropF (z(x))) −1 backprophx,
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
where backprop outputs the result of formal backpropagation, see e.g., [50]. This identity1 is used to provide efficient training despite the fact that the rules of classical nonsmooth calculus are transgressed [7, 59]. Note that spurious outputs may be created by this approach, but on a negligible set. Consider for example the simple implicit problem x = f(z(x)) where f(z) := tanh(z) + relu(−z) + z − relu(z), whose solution is z(x) = tanhx. Yet applying the implicit differentiation framework of [7] using JAX library, as presented in [59], 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 [14, 15] through conservative Jacobians (see also [41, 24]).
A mathematical model for propagating derivatives. Conservative calculus models nonsmooth algorithmic differentiation faithfully and allows for a sharp study of training methods in Deep Learning [14, 15]. It involves a new class of derivatives, generalizing Clarke Jacobians [20]. A distinctive feature of conservative calculus is that it is preserved by Jacobian multiplication. Consider for example a feed forward network combining analytic or relu activations and max pooling. A conservative Jacobian for this network can be obtained by using Clarke Jacobians formally as classical Jacobians, regardless of qualification conditions, compared to other approaches, e.g. [37], for which stricter qualification conditions are imposed to ensure that an element of the Clarke Jacobian itself can be computed using algorithmic differentiation. For instance, Figure 1 depicts a selection in a conservative Jacobian. This approach is general enough to handle spurious points such as in Figure 1 while keeping the essence of the properties one expects from a derivative. It was proved in [14] that backprop, applied to any reasonable program of a function, is a conservative Jacobian for this function; in contrast, backprop cannot be modelled by some subdifferential operator. For instance for the fixed point problem above, given conservative Jacobians JF and Jh (e.g., Clarke Jacobians) for F and h one obtains a new conservative Jacobian Jz implicitly defined through
JF (z(x))Jz(x) = Jh(x).
This property exactly parallels the idea of “propagating derivatives” in practice. It gives a strong meaning to the formal use of Jacobians proposed in [7], and many empirical approaches [32, 2, 31, 30].
Main contributions: — 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 [7], conic optimization layers [2], and hyperparameter optimization for the Lasso [9]. 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 [15, 40]: elementary log-exp functions [15], semialgebraic functions [12], all being subclasses of definable functions [23, 55]. The use of such structural classes has become standard in nonsmooth optimization and is more and more common in ML (see, e.g., [18, 15, 40, 35]).
— 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.
Related work on implicit differentiation: The classical implicit function theorem has two parts: an existence and regularity part, and a calculus part, which is called implicit differentiation. Nons-
1The notation backpropz instead of backprop(z) is indicative of the fact that backprop is an operator that does not act on functions themselves but rather on the program used to represent them, see [15].
mooth generalizations have mostly focused on the existence part. We focus on Lipschitz equations, for which existence of an implicit functional relation is due to [34], based on Clarke’s inverse mapping theorem [19], all these elements being summarized in [20]. Various extensions of this result have been proposed, let us mention a result in an approximation context in [36], a semismooth extension in [52] and a tame extension in [29]; however, all these extensions lack a calculus amenable to implicit differentiation in practice. There have been extensions of the nonsmooth implicit function theorem with a calculus limited to directional derivatives and lexicographic derivatives (a sort of directional derivative introduced in [45]) in [49, Corollary 3.4] and [38, Theorem 2.3], respectively; the former requiring stronger assumptions than just the invertiblity condition. The limitation of the calculus of these extensions to directional derivatives is furthermore incompatible with the most common numerical libraries used for algorithmic differentiation in practice, as opposed to the framework we present.
Definitions and Notations. A function F : Rn → Rm is locally Lipschitz if, for each x ∈ Rn, there exists a neighborhood U of x such that F is Lipschitz on U . Given matrices A ∈ Rn×m and B ∈ Rn×p, [A B] ∈ Rn×(m+p) denotes their concatenation; Idn denotes the n× n identity matrix. For q ∈ Rn, diag (q) ∈ Rn×n denotes the diagonal matrix whose diagonal entries are given by the qi; sign (q) ∈ {−1, 0, 1}n denotes the componentwise sign function. The convex hull of U is denoted conv U . The projection onto a closed convex set C ∈ Rn is given, for each x ∈ Rn, by PC (x) := argmin{ 12 ‖u− x‖ 2 : u ∈ C}. Given a convex proper lower semicontinuous function f : Rn → R∪{+∞}, we define its proximal operator through x ∈ Rn, proxf (x) := argmin{f (u)+ 1 2 ‖u− x‖ 2 u ∈ Rn}. Set-valued maps are denoted by ⇒, for example the subgradient ∂f : Rn ⇒ Rn. Additional details and notations are provided in Appendix A.

heading: 2 Implicit Differentiation with Conservative Jacobians
text: Definitions and conservativity. Conservative Jacobians are generalized forms of Jacobians well suited for automatic differentiation, introduced in [14]. Given a locally Lipschitz continuous function F : Rn → Rm, we say that JF : Rn ⇒ Rn×m is a conservative mapping or a conservative Jacobian for F if JF has a closed graph, is locally bounded, and is nonempty with
d dt F (γ(t)) ∈ JF (γ(t))γ̇(t) a.e. (1)
whenever γ is an absolutely continuous curve in Rn. When m = 1, the corresponding vectors are called conservative gradient fields. Note that when JF is conservative, so is its pointwise convexified extension conv JF .
A locally Lipschitz function is called path differentiable if it has a conservative Jacobian. Recall that the Clarke Jacobian is defined as
Jacc F (x) = conv { lim
k→+∞ JacF (xk) : xk ∈ diffF , xk −→ k→+∞ x } where diffF is the full measure set of points where F is differentiable and JacF is the standard Jacobian of F . Path differentiability is equivalent to having a chain rule as in (1) for the Clarke subdifferential, see [14, 25].
Examples of path differentiable functions and conservative Jacobians. (a) Convex functions and concave functions are path differentiable, see [14]. This implies that their subdifferential in the sense of convex analysis is a conservative field. (b) The vast class of definable functions are path differentiable [25, 14]. As a result, the Clarke Jacobian of a Lipschitz definable mapping is a conservative Jacobian. Definable functions (see [5, 25, 18, 14] for an optimization context and [23] for a foundational work) encompass semialgebraic functions [12], elementary log-exp selection [14], PAP [40] (restricted to analytic functions with full domain), and many others, see [55] and references therein. This includes networks with common nonlinearities: for example analytic with full domain (e.g., square, exponential, logistic loss, hyperbolic tangent, sigmoid), relu, max pooling, sort, (see Appendix A.2 for more detail). (c) The backpropagation can be seen as an oracle (in the optimization sense) for a conservative Jacobian. Let PF be a numerical program for a function F , aggregating elementary functions, for instance, relu, max pooling, affine mappings, polynomials (in general, any definable function). Then the backpropagation algorithm applied to PF , which we denote (abusively) by backpropPF := backpropF , outputs an element of a conservative Jacobian [14, Theorem 8] which depends on PF and can be constructed by a closure procedure [15, definition 5]. As described in [15], due to spurious behaviors, backpropF is not in general an element of the Clarke Jacobian of F .
The structure of conservative Jacobians. As established in [41] in a semialgebraic context, the discrepancy between conservative gradients and Clarke subdifferentials is somehow negligible. Let us provide a version of that result matching our concerns. We call conservative mappings of the null function residual or residual conservative. Such a mapping R has the property that R(x+ tv)v = 0 for almost all t in R and all x, v in Rn × Rn. The following theorem and proposition (partially) extend results from [14] and [41], their proof is given in Appendix B.
Theorem 1 (The Clarke Jacobian is a minimal conservative Jacobian) Given a nonempty open subset U of Rn and F : U ⊂ Rn → Rm locally Lipschitz, let JF be a convex-valued conservative Jacobian for F . Then for almost all x ∈ U , JF (x) = {JacF} and, for all x ∈ U , Jac cF (x) ⊂ JF (x).
Proposition 1 (Decomposition of conservative fields) Let JF be a conservative Jacobian for F , then there is a residual R such that
JF ⊂ Jac cF +R.
Note that the above may not hold with equality. Consider F (x) = |x| and JF (0) = [−1, 1] ∪ [2, 3], JF (x) = sign (x) otherwise. One cannot write JF = Jac cF +R with a residual operator R.
Formal subdifferentiation in a nonsmooth setting. Propagating derivatives within a nonsmooth function finds its justification in the following:
Proposition 2 (Stability by composition, [14]) Let F : Rn → Rm and G : Rm → Rl be two locally path differentiable functions having respective conservative Jacobians JF and JG. Then F ◦G is path differentiable and the point-to-set matrix-valued x⇒ JF (G(x))JG(x) is conservative.
A conservative Implicit Function Theorem. There is already a long tradition of nonsmooth implicit function theorems, e.g., [20, 6, 49, 26]. What makes the following theorem useful is that it comes with a qualification-free calculus. The proofs are given in Appendix B.
Theorem 2 (Implicit differentiation) Let F : Rn × Rm → Rm be path differentiable on U × V ⊂ Rn × Rm an open set and G : U → V a locally Lipschitz function such that, for each x ∈ U ,
F (x,G(x)) = 0. (2)
Furthermore, assume that for each x ∈ U , for each [A B] ∈ JF (x,G(x)), the matrix B is invertible where JF is a conservative Jacobian for F . Then, G : U → V is path differentiable with conservative Jacobian given, for each x ∈ U , by
JG : x⇒ { −B−1A : [A B] ∈ JF (x,G(x)) } .
Corollary 1 (Path differentiable implicit function theorem) Let F : Rn × Rm → Rm be path differentiable with conservative Jacobian JF . Let (x̂, ŷ) ∈ Rn × Rm be such that F (x̂, ŷ) = 0. Assume that JF (x̂, ŷ) is convex and that, for each [A B] ∈ JF (x̂, ŷ), the matrix B is invertible. Then, there exists an open neighborhood U × V ⊂ Rn × Rm of (x̂, ŷ) and a path differentiable function G : U → V such that the conclusion of Theorem 2 holds.
Corollary 2 (Path differentiable inverse function theorem) Let U and V be open neighborhoods of 0 in Rn and Φ : U → V path differentiable with Φ(0) = 0. Assume that Φ has a conservative Jacobian JΦ such that JΦ(0) contains only invertible matrices. Then, locally, Φ has a path differentiable inverse Ψ with a conservative Jacobian given by
JΨ(y) = { A−1 : A ∈ JΦ(Ψ(y)) } .
Remark 1 (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. This example is bivariate and piece-wise linear, illustrating that even in this simple situation, implicit differentiation fails and produces artifacts when formaly applied to Clarke Jacobians. (b) (Lipschitz definable implicit and inverse function theorems) See Theorem 4 and 5 in the appendix

heading: 3 Nonsmooth implicit differentiation in Machine Learning
text: Detailed proof arguments for all considered models are given in Appendix C.
Monotone deep equilibrium networks. Deep Equilibrium Networks (DEQs) [7] are specific neural network architectures including layers whose input-output relation is implicitly defined through a fixed point equation of the form
z = f(z, x) (3)
where x ∈ Rp is a given input and z ∈ Rm is the corresponding output. We may consider that the variable x represents both the input layer and layer parameters. Assuming that, for each x ∈ Rp, there is a unique z ∈ Rm satisfying the relation (3), this defines an input-output relation z : Rp → Rm. Furthermore, if f is path differentiable with convex-valued conservative Jacobian Jf : Rm × Rp ⇒ Rm×(m+p) whose projection on the first m columns are all invertible, then the function z itself admits a conservative Jacobian which can be computed from Theorem 2.
We now focus on monotone operator implicit layers [58] for which assumptions are easily stated. Our method applies to other similar architectures, e.g., DEQs [7] or implicit graph neural networks [32]. Let σ : Rm → Rm be the proximal operator of a convex function and assume σ is path differentiable with conservative Jacobian Jσ : Rm ⇒ Rm×m, assumed to be convex-valued. This encompasses the majority of activation functions used in practice [21]. Let W ∈ Rm×m be a matrix such that W +WT 2θI with θ > 0. Under these assumptions the implicit equation
z = σ(Wz + b) (4)
has a unique output z(W, b) [58, Theorem 2]. The transformation (W, b) 7→ z(W, b) is a monotone implicit layer.
The set-valued mapping obtained from Theorem 2 provides a conservative Jacobian for (W, z) 7→ z(W, z). A similar expression was described in [58, 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.
Proposition 3 (Path differentiation through monotone layers) Assume that Jσ is convex-valued and that, for all J ∈ Jσ(Wz(W, b) + b), the matrix (Idm − JW ) is invertible. Consider a loss-like function ` : Rm → R with conservative gradient D` : Rm ⇒ Rm, then g : (W, z) 7→ `(z(W, b)) is path differentiable and has a conservative gradient Dg defined through
Dg : (W, b) ⇒ { JT (Idm − JW )−T vzT , JT (Idm − JW )−T v) : J ∈ Jσ(Wz + b), v ∈ D`(z) } .
Remark 2 Convexity and invertibility assumptions are satisfied when Jσ is the Clarke Jacobian [58].
Optimization layers: the conic program case. Optimization layers in deep learning may take many forms; we consider here those based on conic programming [17, 3, 2, 4]. We follow [3], simplifying the analysis by ignoring infeasability certificates, which correspond to the absence of a primal-dual solution [17], in line with the implementation described in [2, Appendix B]. Consider a conic problem (P) and its dual (D):
(P) inf cTx subject to Ax+ s = b
s ∈ K
(D) inf bT y subject to AT y + c = 0
y ∈ K∗, (5)
with primal variable x ∈ Rn, dual variable y ∈ Rm, and primal slack variable s ∈ Rm. The set K ⊂ Rm is a nonempty closed convex cone and K∗ ⊂ Rm is its dual cone. The problem parameters are the matrix A ∈ Rm×n and the vectors b ∈ Rm and c ∈ Rn; the cone K is fixed. Under the assumption that there is a unique primal-dual solution (x, y, s), we study the path differentiability of the solution mapping as a function of its parameters:
(A, b, c) 7→ sol(A, b, c) = (x, y, s).
For this, let us interpret the solution mapping as a composition mapping involving equation-like implicit formulations. Set N = n+m, given A, b, c ∈ Rm×n × Rm × Rn, define
Q(A, b, c) =
[ 0 AT
−A 0
] ∈ RN×N V (b, c) = [ c b ] ∈ RN .
Consider a vector z = (u, v) ∈ Rn × Rm, denote by Π the projection onto Rn ×K∗ and define the residual map N : RN × Rm×n × Rm × Rn → RN as
N (z,A, b, c) = (Q(A, b, c)− IdN )Πz + V (b, c) + z.
The mapping 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 [17].
Define the function φ : RN → Rn × Rm × Rm through φ(u, v) := (u, PK∗(v), PK∗(v) − v). As shown in Appendix C.2, φ(u, v) provides a primal-dual KKT solution of problems (P) and (D) if and only ifN (z,A, b, c) = 0. When we assume that, for fixed A, b, and c, there is a unique z ∈ RN such that N (z,A, b, c) = 0, we have an implicitly defined a function z = ν(A, b, c), such that
sol(A, b, c) = [φ ◦ ν] (A, b, c). (6)
The following result extends the discussion in [17, 3], limited to situations where Π is differentiable at the proposed solution z, to a fully nonsmooth setting; its proof is postponed to Appendix C.2.
Proposition 4 (Path differentiation through cone programming layers) Assume that PK∗ , N are path differentiable, denote respectively by JPK∗ , JN corresponding convex-valued conservative Jacobians. Assume that, for allA, b, c ∈ Rm×n×Rm×Rn, z = ν(A, b, c) ∈ Rn×Rm is the unique solution to N (z,A, b, c) = 0 and that all matrices formed from the N first columns of JN (z,A, b, c) are invertible. Then, φ, ν, and sol are path differentiable functions with conservative Jacobians:
Jν(A, b, c) := { −U−1V : [U V ] ∈ JN (ν(A, b, c), A, b, c) } ,
Jφ(z) := [ Idn 0 0 JPK∗ (v) 0 (JPK∗ (v)− Idm) ] ,
Jsol(A, b, c) := Jφ(ν(A, b, c))Jν(A, b, c).
In practice, the path differentiability of conic projections is pervasive since they are generally semialgebraic (orthant, second-order cone, PSD cone). See [33, 43, 39, 43] for the computations of the corresponding Clarke Jacobians (which are conservative). Note that a conservative Jacobian for N may be obtained from JPK∗ using Proposition 2.
Hyperparameter selection for Lasso type problems. Implicit differentiation can be used to tune hyperparameters via first-order methods optimizing some measure of task performance, see [10] and references therein. In a nonsmooth context, we recall the formulation in [9] of the general hyperparameter optimization problem as a bi-level optimization problem:
min λ∈Rm C(β̂(λ)) such that β̂(λ) ∈ argmin β∈Rp ψ(β, λ)
where C : Rp → R is continuously differentiable (e.g., test loss) and ψ : Rp ×Rm → R is a possibly nonsmooth training loss, convex in β, with hyperparameter λ ∈ Rm. We seek a subgradient type method for this problem with convergence guaranties; our nonsmooth implicit differentiation results can be used for this purpose. We demonstrate this approach on the Lasso problem [53]
β̂ (λ) ∈ argmin { 1
2 ‖y −Xβ‖22 + e λ ‖β‖1 : β ∈ R p
} (7)
where y ∈ Rn is the vector of observations, X = [X1, . . . , Xp] ∈ Rn×p is the design matrix with columns Xj ∈ Rn, j ∈ {1, . . . p}, and λ ∈ R is the hyperparameter. Define F : R× Rp → Rp to be
F (λ, β) := β − proxeλ‖·‖1 ( β −XT (Xβ − y) ) and recall that, for each i ∈ {1, . . . , p}, [proxeλ‖·‖1(β)]i = sign(βi) max{|βi|−e
λ, 0}. The function F (λ, β) is thus nonsmooth but locally Lipschitz on R× Rp. An optimal β̂(λ) for (7) must satisfy F (λ, β̂(λ)) = 0 [22, Prop. 3.1]. For a given solution β̂(λ), we introduce the equicorrelation set by E := {j ∈ {1, . . . , p} : |XTj (y −Xβ̂(λ))| = eλ} which contains the support set supp β̂ := {i ∈ {1, . . . , p} : β̂i 6= 0}. In fact, E does not depend on the choice of the solution β̂, see [54, Lemma 1]. The proof of the following result is given in Appendix C.3.
Proposition 5 (Conservative Jacobian for the solution mapping) For all λ ∈ R, assume XTE XE is invertible where XE is the submatrix of X formed by taking the columns indexed by E . Then β̂(λ) is single-valued, path differentiable with conservative Jacobian, Jβ̂ (λ), given for all λ as{[ −eλ ( Idp − diag (q) ( Idp −XTX ))−1 diag (q) sign ( β̂ −XT ( Xβ̂ − y ))] : q ∈M(λ)
} where M(λ) ⊂ Rp is the set of vectors q such that qi = 1 if i ∈ supp β̂, qi = 0 if i 6∈ E and qi ∈ [0, 1] if i ∈ E \ supp β̂.
Taking, in Proposition 5, qi = 1 for all i ∈ E corresponds to the directional derivative given by LARS algorithm [28], see also [42]. Alternatively, taking qi = 0 for i 6∈ supp β̂ gives the weak derivative described by [9]. Both are particular selections in Jβ̂ , which is the underlying conservative field.

heading: 4 Optimizing implicit problems with gradient descent
text: 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 [15].
The minimization problem. The applications considered in the previous section all yield minimization problems of the type
min w∈Rp
`(w) := 1
N N∑ i=1 `i(w) with `i = gi,L ◦ gi,L−1 ◦ . . . ◦ gi,1 (8)
where, for each i ∈ {1, . . . , N}, `i : Rp → R is a composition of functions having appropriate input and output dimensions. The indices i correspond in practice to learning samples while the loss ` embodies an empirical expectation, as for instance in deep learning. We will enforce the following structural condition.
Assumption 1 For i ∈ {1, . . . , N} and j ∈ {1, . . . , L}, the function gi,j is locally Lipschitz with conservative Jacobian Ji,j and one of the following holds
• gi,j and Ji,j are semialgebraic (or, more generally, definable). • gi,j is defined as G in Theorem 2, with F and JF 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 [15] 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 gi,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).
(b) Training with optimization layers: similarly, the inner maps gi,j may also be solution mapping to convex conic programs and related to the sol function (6) of conic problems.
(c) One may assume that N = 1, L = 2 and retrieve the hyperparameter tuning for Lasso in its implicit formulation.
SGD with backpropagation. Algorithmic differentiation (AD) is an automated application of the chain rule of differential calculus. When applied to `i, it amounts to computing one element of the product Ji := ∏L j=1 Ji,j by choosing one element in each Ji,j with appropriate inputs given by intermediate results kept in memory during a forward computation of the composition.
In this context AD stochastic gradient descent requires an initial w0 ∈ Rp and a sequence of i.i.d. random indices uniform in {1, . . . , N}, (Ik)k∈N. It gives:
wk+1 = wk − sαkvk (9) vk ∈ JIk(wk), (given by backprop), (10)
where (αk)k∈N is a sequence of positive step sizes and s ∈ (smin, smax) is a scaling factor where smax > smin > 0. A simpler choice could be vk ∈ ∂c`Ik(wk), however, the chain rule used within algorithmic differentiation routines does not produce subgradients (see, e.g., Figure 1). In contrast, conservative Jacobians are faithful models of AD outputs. The asymptotic behavior of the above algorithm depends on the variational properties of the conservative Jacobian J := 1N ∑N i=1 Ji.
Theorem 3 (Convergence result) Consider minimizing ` given in (8) using algorithm (10) under Assumption 1. Assume furthermore the following
• Step size: ∑+∞ k=1 αk = +∞ and αk = o(1/ log(k)).
• Boundedness: there exists M > 0, and K ⊂ Rp open and bounded, such that, for all s ∈ (smin, smax) and w0 ∈ cl K, ‖wk‖ ≤M almost surely.
For almost all w0 ∈ K and s ∈ (smin, smax), the objective value `(wk) converges and all accumulation points w̄ of wk are Clarke-critical in the sense that 0 ∈ ∂c`(w̄).
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 [14, 11]; details are given in Appendix D.

heading: 5 Numerical experiments
text: Using implicit differentiation when the invertibility condition in Theorem 2 does not hold can result in absurd training dynamics.
A cyclic gradient dynamics via fixed-point/optimization layer. Consider the bilevel problem: min x,y,s `(x, y, s) := (x− s1)2 + 4(y − s2)2 (11)
s.t. s ∈ s(x, y) := arg max {(a+ b)(−3x+ y + 2) : a ∈ [0, 3], b ∈ [0, 5]} .
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:
∇(x,y)`(x, y, s(x)) + J̃s(x, y)T∇s`(x, y, s(x)) (12)
where J̃s is piecewise derivative. We implement gradient descent for (11), evaluating (12) either using cvxpylayers [2] or the JAX tutorial [59] 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 2a); see Appendix E.1 for details.
Persistence under small perturbations: For different initial points the gradient flow converges to the same limit cycle (Figure 2a). The cycle persists even if we perturb the coefficients in the problem (11) ( see Appendix E.1.2 for more details) .
A Lorenz-like dynamics via implicit differentiation. The Lorenz Ordinary Differential Equation (ODE) writes:
ẋ = σ(y − x), ẏ = x(ρ− z)− y, and ż = xy − βz. (13) It is well-known that taking (σ, ρ, β) = (10, 28, 8/3), and (x(0), y(0), z(0)) = (0, 1, 1.05) gives a chaotic trajectory, displayed in Figure 3a. Denoting F : (x, y, z) 7→ (σ(y−x), x(ρ−z)−y, xy−βz) the vector field of the Lorenz system (13), consider the optimization problem:
max u∈R3 uT z s.t. z ∈ arg min s∈R3 ‖s− F (u)‖4 (14)
which is obviously equivalent to max u∈R3 uTF (u). (15)
The function g : u 7→ uTF (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 saddlepoint. 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 3b) resembles the Lorenz attractor (Figure 3a), in stark contrast to the conventional method (Figure 3c). 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.

heading: 6 Conclusion and future work
text: 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.
Acknowledgments and Disclosure of Funding
The authors acknowledge the support of AI Interdisciplinary Institute ANITI funding, through the French “Investing for the Future – PIA3” program under the Grant agreement ANR-19-PI3A0004, Air Force Office of Scientific Research, Air Force Material Command, USAF, under grant numbers FA9550-19-1-7026, FA9550-18-1-0226, and ANR MaSDOL 19-CE23-0017-01. J. Bolte also acknowledges the support of ANR Chess, grant ANR-17-EURE-0010 and TSE-P.

Review summaries:
This paper surveys and summarizes the theoretical support for performing implicit differentiation in the broader, not-necessarily smooth settings. The derived formulas and theories, importantly, are compatible with AD. The authors have established the generalization of their conservative-Jacobian-based approach to multiple scenarios, including equilibrium networks, optimization layer, etc.

Conservative Jacobians provide a formalism that justifies the application of standard AD techniques even when the computational graph contains nonsmooth Lipschitz elemental functions such as relu. An element of the conservative Jacobian can be computed by picking any element of the Clarke Jacobian for the elemental functions and applying the chain rule. In that sense the formalism of conservative Jacobians has a significant practical benefit compared to previously proposed approaches for nonsmooth AD, such as e.g. [1] and [2], because no directional information needs to be propagated (there is no qualification).  In this paper the authors present the natural application of the conservative Jacobian formalism to implicit differentiation and highlight the necessity of invertibility as analogous to the classical implicit function theorem.  The authors highlight some applications of implicit differentiation in the area of machine learning:  - Monotone implicit layers  - Embedded optimization problems such as optimization layers and hyperparam optimization for L1 regularized optimization problems  If the optimality condition for an embedded optimization problem can be formulated as a nonsmooth implicit function with the assumed Lipschitz composed structure then the authors' implicit function theorem applies.  In some cases it can be possible to prove that the invertibility condition is guaranteed.  The authors give a theorem showing that gradient descent with a sufficient step size decay converges to Clarke stationary points when conservative gradients are used. This theorem is not related to the implicit function theorem and also holds for non-implicit conservative differentiation.  The authors show that conservative gradient dynamics for nonsmooth functions can be cycling and even chaotic in the case the invertibility assumption is not satisfied and demonstrate the resulting discrete gradient dynamics in numerical experiments to highlight the relevance of the invertibility assumption of their implicit function theorem in practical examples.  [1] Branch-locking AD techniques for nonsmooth composite functions and nonsmooth implicit functions [2] Provably Correct Automatic Subdifferentiationfor Qualified Programs

This work provides theoretical foundations for nonsmooth implicit differentiation based on conservative Jacobians. Specifically, it provides conditions under which implicit differentiation is compatible with backpropagation and first-order optimization algorithms. The authors then study the applications of the implicit differentiation algorithm, including deep equilibrium problems, conic optimization layers, and hyperparameter optimization problems for Lasso. The specific conditions and the forms of the Jacobians are also provided. After that, the authors obtain the convergence guarantees for first-order optimization algorithms. Finally, two examples -- a cyclic gradient dynamic via optimization layer and a Lorenz-like attractor via implicit differentiation – are used to demonstrated that the optimization may fail when the conditions are not satisfied.