Differentiable Spline Approximations

1 Introduction . Motivation : Differentiable programming has been a paradigm shift in algorithm design . The main idea is to leverage gradient-based optimization to optimize the parameters of the algorithm , allowing for end-to-end trainable systems ( such as deep neural networks ) to exploit structure in data and achieve better performance . This approach has found use in a large variety of applications such as scientific computing [ Innes , 2020 ; Innes et al. , 2019 ; Schafer et al. , 2020 ] , image processing [ Li et al. , 2018a ] , physics engines [ Degrave et al. , 2017 ] , computational simulations [ Alnæs et al. , 2015 ] , and graphics [ Li et al. , 2018b ; Chen et al. , 2019 ] . One way to leverage differentiable programming modules is to encode additional structural priors as “ layers ” in a larger machine learning model . Inherent structural constraints such as monotonicity , or piecewise constancy , are particularly prevalent in applications such as physics simulations , graphics rendering , and network engineering . In such applications , it may be beneficial to build models that obey such priors by design . Challenges : For differentiable programming to work , all layers within the model must admit simple gradient calculations ; however , this poses a major limitation in many settings . For example , consider computer graphics applications for rendering 3D objects [ Kindlmann et al. , 2003 ; Gross et al. , 1995 ; Loop and Blinn , 2006 ] . A common primitive in such cases is a spline ( or a piecewise polynomial ) function which either exactly or approximately interpolates between a discrete set of points to produce a continuous shape or surface . Similar spline ( or other piecewise polynomial ) approximations arise in ⇤†Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . partial differential equation ( PDE ) solvers [ Hughes et al. , 2005 ] , network flow problems [ Balakrishnan and Graves , 1989 ] , and other applications . For such problems , we would like to compute gradients “ through ” operations involving spline approximation . However , algorithms for spline approximation often involve discontinuous ( or even discrete ) co-domains and may introduce undefined ( or even zero ) gradients . Generally , embedding such functions as layers in a differentiable program , and running automatic differentiation on this program , requires special care . A popular solution is to relax these non-differentiable , discrete components into continuous approximations for which gradients exist . This has led to recent advances in differentiable sorting [ Blondel et al. , 2020 ; Cuturi et al. , 2019 ] , dynamic programming [ Mensch and Blondel , 2018 ] , and optimization [ Djolonga and Krause , 2017 ; Agrawal et al. , 2019 ; Deng et al. , 2020 ] . Our contributions : We propose a principled approach for differentiable programming for spline functions without the use of continuous relaxation2 . For the forward pass , we leverage fast algorithms for computing the optimal projection of any given input onto the space of piecewise polynomial functions . For the backward pass , we leverage a fundamental locality property in splines that every piece ( or basis function ) in the output approximation only interacts with a few other elements . Using this , we derive a weak form of the Jacobian for the spline operation and show that it exhibits a particular block-structured form . While we focus on spline approximation in this paper , our approach can be generalized to any algorithmic module with piecewise smooth outputs . Our specific contributions are as follows : 1 . We propose the use of spline function approximations as “ layers ” in differentiable programs . 2 . We derive efficient ( nearly-linear time ) methods for computing forward and backward passes for various spline approximation problems , showing that the ( weak ) Jacobian in each case can be represented using a block sparse matrix that can be efficiently used for backpropagation . 3 . We show applications of our approach in three stylized applications : image segmentation , 3D point cloud reconstruction , and finite element analysis for the solution of partial differential equations . Related Work Before proceeding , we briefly review related work . Extensions of autodiff : Automatic differentiation ( autodiff ) algorithms enable gradient computations over basic algorithmic primitives such as loops , recursion , and branch conditions [ Baydin et al. , 2018 ] . However , introducing more complex non-differentiable components requires careful treatment due to undefined or badly behaved gradients . For example , in the case of sorting and ranking operators , it can be shown that the corresponding gradients are either uninformative or downright pathological , and it is imperative the operators obey a ‘ soft ’ differentiable form . Cuturi et al . [ 2019 ] propose a differentiable proxy for sorting based on optimal transport . Blondel et al . [ 2020 ] improve this by proposing a more efficient differentiable sorting/ranking operator by appealing to isotonic regression . Berthet et al . [ 2020 ] introduce the use of stochastic perturbations to construct smooth approximations to discrete functions , and other researchers have used similar approaches to implement end-to-end trainable top-k ranking systems [ Xie et al. , 2020 ; Lee et al. , 2020 ] . Several approaches for enabling autodiff in optimization have also been researched [ Pogančić et al. , 2020 ; Amos and Kolter , 2019 ; Agrawal et al. , 2019 ; Mensch and Blondel , 2018 ] . Structured priors as neural “ layers ” : As mentioned above , one motivation for our approach arises from the need for enforcing structural priors for scientific computing applications . Encoding non-differentiable priors such as the solutions to specific partial differential equations [ Sheriffdeen et al. , 2019 ] , geometrical constraints [ Joshi et al. , 2020 ; Chen et al. , 2019 ] , and spatial consistency measures [ Djolonga and Krause , 2017 ] perform well but typically require massive amounts of structured training examples . Spline approximation : Non-Uniform Rational B-splines ( NURBS ) are commonly used for defining spline surfaces for geometric modeling [ Piegl and Tiller , 1997 ] . NURBS surfaces offer a high level of control and versatility ; they can also compactly represent the surface geometry . The versatility of NURBS surfaces enables them to represent more complex shapes than Bèzier or B-splines . Several frameworks that leverage deep learning are beginning to use NURBS representations . Minto et al . [ 2018 ] use NURBS surfaces fitted over the 3D geometry as an input representation for the object 2While tricks such as straight-through gradient estimation [ Bengio , 2013 ] also avoid continuous relaxation , they are heuristic in nature and may be inaccurate for specific problem instances [ Yin et al. , 2019 ] . classification task of ModelNet10 and ModelNet40 datasets . Erwinski et al . [ 2016 ] presented a neural-network-based contour error prediction method for NURBS paths . Fey et al . [ 2018 ] present a new convolution operator based on B-splines for irregular structured and geometric input , e.g. , graphs or meshes . Very recently , Sharma et al . [ 2020 ] perform point cloud reconstruction to predict a B-spline surface , which is later processed to obtain a complete CAD model with other primitives “ stitched ” together . Differentiable PDE solvers : With the advent of deep learning , there has been a recent rise in the development of differentiable programming libraries for physics simulations [ Hu et al. , 2019 ; Qiao et al. , 2020 ] . Most often , the physics phenomena are represented using partial differential equations ( PDEs ) [ Sanchez-Gonzalez et al. , 2020 ; Holl et al. , 2020 ] . Considerable effort has gone into designing physics-informed loss functions [ Raissi et al. , 2019 ; Raissi and Karniadakis , 2018 ; Kharazmi et al. , 2021 ] whose optimization leads to desired solutions for PDEs . Due to space limitations , we defer to a detailed survey of this ( vast ) area by Cai et al . [ 2021 ] . 2 Differentiable Spline Approximation . We now introduce our framework , Differentiable Spline Approximation ( DSA ) , as an approach to estimate gradients over piecewise polynomial operations . Our main goal will be to estimate easy-to-compute forms of the ( weak ) Jacobian for several spline approximation problems , enabling their use within backward passes in general differentiable programs . Setup . We begin with some basic definitions and notation . Let f 2 Rn be a vector where the ith element is denoted as fi . Let us use [ n ] = { 1 , 2 , . . . , n } to denote the set of all coordinate indices . For a vector f 2 Rn and an index set I ✓ [ n ] , let fI be the restriction of f to I , i.e. , for i 2 I , we have fI ( i ) : = fi , and fI ( i ) : = 0 for i /2 I . Now , consider any fixed partition of [ n ] into a set of disjoint intervals I = { I1 , . . . , Ik } where the number of intervals |I| = k. The ` 2-norm of f is written as kfk2 : = pPn i=1 f 2 i while the ` 2 distance between f , g is written as kf gk2 . We first define the notion of a discretized k-spline . Note that the use of “ spline '' here is non-standard and somewhat more general than what is typically encountered in the literature . ( Indeed , the spline concept used in computer graphics is a special instance of this definition ; we explain further below . ) Def . 2.1 ( Discretized k-spline ) . A vector h 2 Rn is called a discretized k-spline with degree d if : ( i ) there exists a partition of [ n ] into k disjoint intervals I1 , . . . , Ik ; ( ii ) within each interval Ii , the coefficients of hj , j 2 Ii , can be perfectly interpolated by some polynomial function of degree d. Let us illustrate this by an example . Suppose that d = 1 and k = 5 . Then , h is a discretized k-spline with degree d if , in a “ line plot ” of the vector h ( i.e. , we interpolate the 2D points ( j , hj ) for all j 2 [ n ] ) , we see up to k = 5 distinct linear pieces . A different way to interpret this definition is that we start with a piecewise degree-d polynomial function H : R ! R with k = 5 pieces ( with suitably defined knot points , which are the location of the intervals I ) , and evaluate H at any n equally spaced points in its domain . This gives us a vector h 2 Rn , which we call a discretized k-spline . In contrast with traditional splines , we allow H to be arbitrarily defined at the knot points and require no specific continuity or differentiability properties . Therefore , our definition encompasses all standard spline families ( including interpolating/approximating splines such as smoothing- , cubic- , and B-splines ) . 2.1 Spline Approximation . Our focus in this paper is the problem of computing the best possible spline fit to a given set of data points ( where both the parameters of the spline as well as the knot vectors are allowed to be variable ) . We provide an algebraic interpretation of this problem . For a given vector space Rn , consider Skd , the set of all discretized k-splines with degree d. Since ( standard ) splines are vector spaces for a fixed set of knots , one can easily see that for any fixed partition of [ n ] into k subsets , the family of discretized k-splines is a k ( d+ 1 ) -dimensional subspace of Rn . Now suppose that the knot indices are allowed to vary . The number of possible partitions is finite ( of the order of n k ) , and therefore the set Skd is a finite union of subspaces , or a nonlinear submanifold , embedded in Rn . Therefore , the problem of discretized k-spline approximation can be viewed as an orthogonal projection onto this nonlinear manifold . Consider any arbitrary vector x 2 Rn ( we can think of ( i , xi ) as a set of n data points to which we are trying to fit a k-spline ) . Then , the best k-spline fit to x ( in the sense of ` 2 distance ) amounts to solving the optimization problem : F ( x ) = argmin h 1 2 kx hk22 = 1 2 nX i=1 ( xi hi ) 2 s.t . h 2 Skd ( 1 ) This operation resembles standard spline regression . But it is strictly more general since this requires not only optimizing piecewise spline parameters but also the knot indices . Crucially , we note that F is both a non-differentiable and a non-convex map . Nevertheless , such an orthogonal projection can be computed in polynomial ( in fact , nearly-linear ) time [ Jagadish et al. , 1998 ; Acharya et al. , 2015 ] using many different techniques , including dynamic programming . This forms the forward pass of our DSA “ layer ” . Our first main conceptual contribution is a formal derivation of the backward pass of the orthogonal projection operation . Strictly speaking , the Jacobian is not well-defined due to the non-differentiable nature of the forward pass ( owing to the non-differentiability built into the definition of the kspline ) . Therefore , we will instead be deriving the so-called “ weak ” form of the Jacobian ( borrowing terminology from Blondel et al . [ 2020 ] ) . We leverage two properties of the projection operation : ( 1 ) the output of the forward pass h corresponds to a partition of [ n ] , that is , each element of hj corresponds to a single interval , Ij , and ( 2 ) within each interval , the least-squares operation is continuous and differentiable . The first property ensures that every element xi contributes to only a single piece in the output h. Given that the sub-functions from the piecewise partitioning function are smooth , we also observe that the size of each block corresponds to the size of the partition , Ii . Using this observation , we get : Theorem 1 . The Jacobian of the operation F with respect to x 2 Rn can be expressed as a block diagonal matrix , J 2 Rn⇥n , whose ( s , t ) th entry obeys : Jx ( F ( x ) ) ( s , t ) = @ h ( x ) s @ xt = ( @ hIi ( x ) s @ xt if s , t 2 Ii 0 otherwise ( 2 ) As a concrete instantiation of this result , consider the case d = 0 . This is the case where we wish to best approximate the entries of x with at most k “ horizontal ” pieces , where the break-points are obtained during the forward pass3 . Call this approximation h. Then , the Jacobian of h with respect to x forms the block-diagonal matrix J 2 Rn⇥n : J = 2 664 J1 0 . . . 0 0 J2 . . . 0 ... ... . . . ... 0 0 . . . Jk 3 775 ( 3 ) where all entries of each block , Ji 2 R|Ii|⇥|Ii| are constant and equal to 1/|Ii| , i.e. , they are row/column-stochastic . Note that the sparse structure of the Jacobian allows for fast computation and that computing the Jacobian vector product JT ⌫ for any input ⌫ requires O ( n ) running time . As an additional benefit , the decoupling induced by the partition enables further speed up in computation via parallelization . See the Appendix for proofs , as well as derivations of similar Jacobians for k-spline approximation of any degree d 1 , and generalization to 2D domains ( surface approximation ) . In Section 3 we demonstrate the utility of this approach for a 2D segmentation ( i.e. , piecewise constant approximation ) problem , similar to the setting studied in Djolonga and Krause [ 2017 ] . 3In the data summarization literature , this class of functions is sometimes called k-histograms [ Jagadish et al. , 1998 ]

Summary: The paper presents sparse Jacobians/backward algorithms for several spline formulations and motivates their use in differentiable algorithms. The formulations include general splines, NURBS, and basis functions for FEM solvers.  The resulting differentiable modules are applied to the three proof of concept tasks of image segmentation, point cloud reconstruction, and solving Poisson equations, where they show promising results.