Fast Linear Interpolation for Piecewise-Linear Functions, GAMs, and Deep Lattice Networks

We present fast implementations of linear interpolation operators for piecewise linear functions and multi-dimensional look-up tables . We use a compiler-based solution using MLIR to accelerate this family of workloads . On two-layer and deep lattice networks , we show these strategies run 5 − 10× faster on a standard CPU compared to a C++ interpreter implementation that uses prior techniques , producing runtimes that are 1000s of times faster than TensorFlow 2.0 for single evaluations . 1 INTRODUCTION . Linearly-interpolated look-up tables ( LUTs ) are a core operation of many machine-learned models , ranging from simple one-dimensional piecewise-linear ( PWL ) functions ( see Fig . 1 ) , to Ddimensional generalized additive models ( Hastie & Tibshirani , 1986 ) ( GAMs ) that are a linear combination of a PWL on each input variable , to D-dimensional interpolated LUTs also known as lattices . Interpolated LUTs can be used to efficiently express arbitrarily flexible functions over an arbitrary number of inputs D by increasing the size of the underlying LUTs , or by using ensembles or multiple layers of lattices ( Canini et al. , 2016 ; You et al. , 2017 ; Cotter et al. , 2019 ) . Interpolated LUTs have long been a standard choice for applications where fast evaluation and flexible models are needed . For example , two-layer models formed by PWLs followed by lattices are part of the International Color Consortium standard for colorspace transformations ( Sharma & Bala , 2002 ) . Garcia & Gupta ( 2009 ) showed that lattices can be trained by empirical risk minimization . As concerns grow about the black-box nature of AI , the highly-structured nature of interpolated LUTs has made them a popular choice for producing AI with safety guarantees such as bounded outputs and satisfying shape constraints ( Groeneboom & Jongbloed , 2014 ; Chetverikov et al. , 2018 ) , such as monotonicity which guarantee that a model only responds positively to increases in designated inputs ( Barlow et al. , 1972 ; Howard & Jebara , 2007 ; Gupta et al. , 2016 ; Canini et al. , 2016 ; You et al. , 2017 ) , and other shape constraints ( Pya & Wood , 2015 ; Chen & Samworth , 2016 ; Gupta et al. , 2018 ; Cotter et al. , 2019 ) . Fast evaluation of models is paramount for real-time estimation needs , and cheap evaluation is key to lowering the costs of AI . In theory , what makes interpolating LUTs so fast is that only a small subset of parameters is needed for any given evaluation , in contrast to models like DNNs where every parameter may be needed for every evaluation . In this paper we investigate how fast one can evaluate interpolated LUTs on standard CPUs . The main contributions of this paper are : ( i ) we give two new strategies to reduce the runtime of PWLs , and ( ii ) we show how to optimize the kernels for two types of multi-dimensional linear interpolation ( simplex and multilinear ) . Overall we show a 5-10x speed-up for the runtime of real-world lattice models , compared to a C++ interpreter implementation that uses prior techniques , producing runtimes that are 1000s of times faster than TensorFlow 2.0 for single evaluations . All comparisons are on the same standard CPUs . 2 BACKGROUND ON LINEAR INTERPOLATION OPERATORS AND SOME INITIAL ANALYSIS . We review PWLs and the two standard types of multi-dimensional interpolation , and what the challenges are in making them run fast . 2.1 PIECEWISE-LINEAR FUNCTIONS ( PWLS ) . PWLs are an old technique for efficiently representing functions ( Sang , 1875 ; Perry , 1899 ; Farr , 1860 ) . Multi-layer machine-learned models may use thousands of PWLs to transform the inputs in their first-layer , which makes the overall performance of a lattice model highly sensitive to the runtime of the PWL kernel . As illustrated in Fig . 1 , we define a PWL by N key-value pairs ( ki , vi ) N i=1 , where the keys are sorted ki < ki+1 , and the values are interpolated to evaluate any input x ∈ [ k1 , kN ] as per the formula in Fig . 1 . The key challenge to efficiently evaluating PWLs is to quickly find the index j in the linear interpolation formula in Fig . 1 needed for a given x . We estimate the cost of a memory load and compare at 2 cycles if speculative memory loads are allowed , and 4 cycles if they are not , due to the 2-3 cycle latency of a memory load . Each branch misprediction costs an expected 15-20 cycles depending on architecture ( Fog , 2018 ) . As a result , we estimate the overall expected cost to be E [ # Cycles ] = 4 × E [ # Comparisons ] + 17 × E [ # Mispredictions ] . Here , we assume that all parameters are within the L1-cache because there are few enough parameters to be loaded together . 2.1.1 LINEAR SEARCH FOR PWLS IS SLOW , EVEN FOR SMALL PWLS . Linear search over the N keypoints requires E [ # Comparisons ] = N+12 . However , as we show in this section , it is the branch mispredictions that are the bigger problem . Recent recommendations for machine-learned lattice models , e.g . ( Gupta et al. , 2016 ; 2018 ) , are to set the keypoints based on the quantiles of the training examples for that input : assign keypoint k1 to the minimum possible value of the PWL ’ s domain , assign the last keypoint kN to the maximum possible value of the PWL ’ s domain , and assign the remainingN−2 keypoints to equally-spaced quantiles of the training examples ’ inputs . We assume the keypoints stay fixed at these values and are not trained ( though their corresponding values { vi } are trained ) . Quantile keypoints are good for machine-learning because each keypoint sees roughly 1/N of the training examples , reducing the chance of overfitting any of the trained PWL values { vi } . Quantile keypoints also aid in interpretability because the PWLs can be visualized , and quantile keypoints show how the training data was distributed over the input domain . However , we next show that quantile keypoints are terrible for PWL runtimes because each keypoint is equally likely to be the left-keypoint for a random input x . Recall that a branch predictor predicts whether or not a given branch is taken , and an optimal branch predictor always predicts the outcome associated with the highest probability . During the linear search , a branch prediction will predict whether the for-loop over the keypoints will stop , for each i = 1 , . . . , N . Typically , a branch predictor is able to access a summary of its history , and any static information the compiler may be able to provide . For quantile keypoints , each of the first N − 1 keypoints is equally likely to be the correct index . Thus the optimal branch prediction is to continue unless the linear search has reached i = N − 2 , in which case there is a 50-50 chance of either of the remaining two keypoints being the right one . However , ( N − 2 ) / ( N − 1 ) of the examples will find their correct index before the linear search reaches i = N − 2 , which means the branch prediction will be wrong once with probability ( N − 2 ) / ( N − 1 ) , producing E [ # Mispredictions ] = ( N − 2 ) / ( N − 1 ) . 2.1.2 BINARY SEARCH FOR PWLS . In a branch-free implementation of binary search with known depth , the compiler is able to fully unroll the structure and thus avoid branch mispredictions . Additionally , a well-optimized binary search implementation is able to perform each step of the binary search in approximately 6 cycles ( Khuong , 2012 ) . It thus takes 6 × dlog2 ( N ) e cycles to find the appropriate location . This makes binary search roughly 2× more efficient than linear search even for N = 3 to 10 , and is the baseline that any proposed indexing must beat . 2.1.3 A MAP-TO-INDEX FUNCTION FOR PWLS . More abstractly , the goal is to construct an efficient function that can map an input x to the correct index j . An old trick is to build an auxiliary LUT over [ k1 , kN ] with B uniformly-spaced buckets , use that to map x to a bucket , and then linearly-search through all the keypoints that fell in that bucket . However , with irregularly-spaced keypoints , a uniform bucket can still have O ( N ) keypoints to search through . Aus & Korn ( 1969 ) proposed constructing a hierarchy of such auxiliary uniformly-spaced LUTs to better cover irregular keypoints . O ’ Grady & Young ( 1991 ) proposed using a sufficiently large B such that no uniform bucket contains more than one keypoint , but at the cost of potentially large B . An analogous problem arises in database indexing , where recent work has proposed machine-learning a two-layer DNN to produce the map-to-index function ( Kraska et al. , 2018 ) . Our proposed solution will be in a similar spirit but lighter-weight . 2.2 MULTILINEAR INTERPOLATION . Consider one cell of a D-dimensional LUT , which without loss of generality is a D-dimensional unit hypercube parameterized by LUT values v ∈ R2D corresponding to the O ( 2D ) vertices of the hypercube . For an input x ∈ [ 0 , 1 ] D , multilinear interpolation outputs f ( x ) = ∑2D i=1 viwi ( x ) , where vi is the stored multi-d LUT value for the ith vertex in the D-dimensional unit hypercube , and wi ( x ) is the multilinear interpolation weight on the ith unit hypercube vertex ξi ∈ [ 0 , 1 ] D taken in lexicographical order , computed from x ∈ [ 0 , 1 ] D as : wi ( x ) = D∏ d=1 x [ d ] ξi [ d ] ( 1− xd ) 1−ξd for all i = 1 , . . . , 2D ( 1 ) For example , for D = 1 , ξ1 = 0 , ξ2 = 1 , and thus w1 ( x ) = 1 − x and w2 ( x ) = x , the same as in the PWL formula in Fig . 1 given that there are only two keypoints k1 = 0 and k2 = 1 . For D = 2 , this is standard bilinear interpolation . Gupta et al . ( 2016 ) gave an O ( 2D ) dynamic programming algorithm for computing ( 1 ) . We note that while asymptotically efficient , the dynamic programming solution introduces a loop-carry dependency , and thus prevents critical compiler optimizations , a problem we show how to avoid . 2.3 SIMPLEX INTERPOLATION . Simplex interpolation is a more efficient linear interpolation of a D-dimensional LUT cell that produces a locally linear surface comprised ofD ! pieces . For each input x ∈ [ 0 , 1 ] D , theD components of x are sorted , the resulting sort order determines a set of D + 1 vertices whose simplex is guaranteed to contain x , and then a sparse inner product is taken with the corresponding D + 1 LUT values to produce f ( x ) ( Rovatti et al. , 1998 ; Weiser & Zarantonello , 1988 ; Gupta et al. , 2016 ) . This algorithm is the same as the Lovász extension in submodularity ( Bach , 2013 ) . Gupta et al . ( 2016 ) gave runtimes for single-layer models that only did either multilinear and simplex interpolation and were implemented in C++ on a single-threaded 3.5GHz Intel Ivy Bridge proces- sor : for D = 4 inputs both interpolations ran in about 50 nanoseconds , but for D = 20 , simplex interpolation ran in 750 nanoseconds , and multilinear interpolation ran in 12 milliseconds , around 15,000× slower . Despite the slower runtime for D > 4 , multilinear interpolation might be preferred because it produces a smoother surface and might produce more accurate models for some problems . Furthermore , the large-scale multiplications needed for multilinear are a better match than the sorting needed for simplex for machine learning libraries such as TensorFlow . Like with PWLs , branch prediction poses a significant challenge when implementing the simplex interpolation kernel . For example , we found that LLVM defaults to using either a hard-coded insertion sort or quicksort depending on the input size , which is determined at runtime . However , we note that this run-time decision is not needed for machine learning models , because the number of inputs D is fixed and known . Using std : :sort < std : :pair < double , int > > , we found that the sorting operation accounts for approximately 70 % of the simplex kernel ’ s overall runtime .

Summary: This paper proposes several low-level code optimizations aimed at speeding up evaluation time for linearly interpolated look-up tables, a method often used in situations where fast evaluation times are required. The paper is very practically oriented in the sense that it does not aim for improving asymptotic speed-ups, but rather real speed-ups expressible in terms of CPU cycles that are achievable by exploiting compiler optimizations such as branch prediction and loop unrolling. The focus is on unbatched computations which can arise in many real-time situations. The proposed implementation techniques are as follows: