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