Nonsmooth Implicit Differentiation for Machine-Learning and Optimization

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 . 1 Introduction . 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 can not 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 can not 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 can not 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 . 2 Implicit Differentiation with Conservative Jacobians . 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 can not 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 can not 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

Summary: 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