Giga-scale Kernel Matrix-Vector Multiplication on GPU
Anonymous Author(s)
Abstract
Kernel matrix-vector multiplication (KMVM) is a foundational operation in machine learning and scientific computing.
However, as KMVM tends to scalequadratically in both memory and time, applications are often limited by these computational constraints.
In this paper, we propose a novel approximation procedurecoined Faster-Fast and Free Memory Method (F 3 M) to address these scaling issuesof KMVM for tall ( 10 8 ∼ 10 9 ) and skinny ( D ≤ 7 ) data.
Extensive experimentsdemonstrate that F 3 M has empirical linear time and memory complexity with arelative error of order 10 − 3 and can compute a full KMVM for a billion points inunder a minute on a high-end GPU, leading to a significant speed-up in comparisonto existing CPU methods.
We demonstrate the utility of our procedure by applyingit as a drop-in for the state-of-the-art GPU-based linear solver FALKON, improvingspeed 1.5-5.5 times at the cost of < 1% drop in accuracy.
We further demonstrate competitive results on Gaussian Process regression coupled with significantspeedups on a variety of real-world datasets.
Introduction
Kernel matrix-vector multiplication (KMVM) is one of the most important operations needed inscientific computing with core applications in diffeomorphic registration, geometric learning [11],[31], numerical analysis [28], fluid dynamics [6] and machine learning [27].
For a dataset of sizen , KMVM using direct computation has complexity and memory footprint O ( n 2 ) , both unfeasiblefor modern large scale applications where n ≈ 10 9 is becoming increasingly common.
Pioneeringcontributions presented in the Fast Multipole Method (FMM) [10] amend the complexity of theseproblems to O ( n log ( ϵ − 1 )) , where ϵ is the chosen error tolerance, with varying reductions in memoryfootprint for data restricted to dimension D = 2 .
Subsequent developments in [8, 17] mainly focusedon extending approximations for a broader set of kernels for a fixed dimensionality D ≤ 3 , tailoredfor problems in physics with narrow data such as electrostatics, stellar dynamics, Stokes flow, andacoustic problems, amongst others.
In this paper, we introduce Faster-Fast and Free Memory Method (F 3 M), a novel algorithm builtupon the FFM [3] framework to perform KMVM on a GPU for tall and skinny ( D ≤ 7 ) data of ordern ∼ 10 9 in under a minute with user-specified error tolerance, providing between 2 − 8500 timesspeed-up over existing methods.
Notations.
We use capital and lower case bold letters to represent matrices and vectors, respectively.
In this paper, we will work with matrices X ∈ R n x × D , Y ∈ R n y × D and vector b ∈ R n y .
For akernel k , the goal for KMVM is to compute v :=
K · b , where K :=
k ( X , Y ) =
{ k ( x i , y j ) } n x ,n y i =1 ,j =1 ,and x i , y j denote the i th , j th row of X , Y respectively.
[Figure]
Motivation and Related Work
Kernel methods are often limited by their O ( n 2 ) memory footprint and computational complexity for
KMVM.
These constraints make scaling beyond n = 10 6 challenging.
Many recent developmentshave been made to improve both of these constraints, ranging from hardware acceleration using GPUsin KeOps [11], to various approximation techniques proposed in [36, 33, 34, 3, 9].
In this work, wefocus our attention on kernel independent KMVM methods.
KeOps.
Charlier et al.
[11] proposes a map-reduce scheme to compute kernels using exactly O ( n )memory and O ( n 2 ) complexity on GPU.
This is achieved by computing the full KMVM producton-the-fly by summing v i = (cid:80) nj =1 k
( x i , y j ) b
j directly, without ever storing the kernel matrix
K ∈ R n × n explicitly.
Extensive experiments show that this method is practical when n ≤ 10 6 , asthe GPU hardware acceleration allows the KMVM product to be computed in less than a second on aconventional GPU.
Moreover, the method places no constraint on the number of features D it can beapplied to, making it favourable for KMVM on medium size datasets.
In application contexts, KeOpsis currently adopted into conjugate gradient solver FALKON [24, 26] as part of the default pipeline.
[Table]
The Fast and Free Memory Method (FFM).
While KeOps can theoretically scale to a billion points, it becomes practically infeasible asthe O ( n 2 ) complexity would imply a computational time of 10 6 seconds, or roughly 11 days.
To overcome this billion points barrier, Aussalet al.
[3] deploys a geometric space partitioningscheme, and proposed the Fast and Free Memory Method (FFM), a KMVM approach that extends the FMM [10] family of algorithms.
In contrastto traditional FMM methods, which require specific series expansion of the kernel, FFM deploys Lagrange interpolations to approximate them instead.
This allows FFM to be applied to almost anyconventional kernel and further enables the user to trade off accuracy with computational efficienciesby controlling the order of the approximating polynomial [19].
Compared to KeOps, FFM demonstrates both linear memory and time complexity in experiments and scales to compute a billion-points KMVM on a smaller CPU cluster under 4 hours, outscaling the GPU implementation of FMM [21].
While 4 hours is a significant improvement compared to 11 days from KeOPS, it still renders manymachine learning techniques infeasible.
Further, as recursive partitioning of the data space scalespoorly with D [5], both FMM and FFM can only be applied to D ≤ 3 data, a price to pay for thespeed-up of KMVM operations when n = 10 9 .
Furthermore, we show in our experiments that a direct FFM port to GPU gives unstable results for n = 10 9 , D = 3 for non-trivial data simulations (bottomrow in Appendix 9).
Our contribution.
To surpass the billion point barrier while maintaining high-speed and stablecomputation, we propose F 2 .
5 M and our main algorithm F 3 M, the first pair of KMVM algorithmsthat can reliably scale to n = 10 9 on skinny data using a single GPU.
We build F 2 .
5 M on top of
FFM by introducing non-trivial GPU parallelisation and low-level optimisations.
We further stabilizeand improve the original geometric partitioning scheme in FFM to significantly reduce memoryconstraints, leading to a relaxation of dimensionality constraints from 3 to 7 .
At last, we introduce anadaptive far-field and smooth field approximation scheme for kernel interpolation, resulting in ourmain algorithm F 3 M, which runs 2 .
− 33 .
3 times quicker and more stable than a direct port of FFMon GPU.
See Fig.
1 and Table 1 for an overview and comparisons of the methods.
We summarise ourcontribution as follows:
We propose Faster-Fast and Free Memory Method ( F 3 M ) , a KMVM algorithm building ontop of FFM by applying multiple low-level enhancements, GPU parallelisation, and algorithmiccomputational and memory enhancements, allowing for KMVM operations on n ≤ 10 9 data in undera minute.
Codebase is released here [14].2.
We characterize theoretical time and memory complexity of F 3 M.3.
We run extensive KMVM experiments of F 3 M on a variety of tall and skinny data with n ≤ 10 9 ,demonstrating empirical linear time and memory scaling, and achieving speedups between 2–8500times when compared to FFM (GPU and CPU) and KeOps.4.
We run a practical application of F 3 M as a drop-in replacement for KeOps in conjugate gradientsolver FALKON [24, 26] for kernel ridge regression and classification (KRR) on giga-scale data,obtaining a solution 3.4 times faster with <1% drop in accuracy.
We further demonstrate competitiveresults on Gaussian process regression against KISS-GP [35], SVGP [18] and SVGR [30] withsignificant speed-ups.
Background
The FFM method considers KMVM for a kernel k evaluated on two data matrices X , Y andb are weights associated with Y .
The KMVM is expressed as v :
=
k
( X , Y ) · b = K · b .
For example, b could be the weights in a KRR or the strength of electronic charges.
As n xand n y are taken to be very large, a full computation is unfeasible.
In this section, we illustrate and detail the main steps of FFM, before presenting our improvements in Section 4.
For illustration purposes, we first consider a simple 2D
KMVM.
Our goal is to calculate k ( X , Y ) · b for X , Y in
Figure 2.
The intuition behind FFM is to reduce the complexity of calculating the full KMVM by partitioning Xand Y such that certain calculations can be approximatedin a fast manner, based on the pairwise distances betweenpartitions.
Enclosing and partitioning the data.
The first step is to partition the data.
To begin, we find a large enough box that can just enclose X or Y .
The edge length of this box is calculated as
[Equation]
where x ( d )max , x ( d )min denotes the largest value and the smallest value along the d -dimension in X andsimilarly for Y .
Figure 2 illustrates this enclosing procedure.
Defining near and far-field In FMM, an octree [23] is applied to recursively partition data into smaller boxes B Xp ⊂ X , B Yq ⊂ Y , with p, q denoting box indices.
Here each box correspondsto a subset of rows in the data matrix.
Let usalso denote b q as the partition of b j ’s groupedwith the same indices as B Yq .
To calculate the
KMVM between two boxes B Xp , B Yq with thegrouped vector b q , for each x i ∈ B Xq , we compute
[Equation]
with v p,q =
[ v p,q 1 . . .
v p,q n x ] .
Now the target v can be computed as v = σ ([ v p =1 , . .
. , v p = P ] ⊤ ) , wherev p = (cid:80) Qq =1 v p,q , P, Q denote the total number of boxes and σ ( · ) a permutation such that v pi appearin the same order as x i appears in X .
Figure 3 shows how boxes are recursively partitioned.
Far and near-field interactions FFM relies on a divide-and-conquer strategy to effectively compute a KMVM product; data is partitioned into boxes and then separated into far-field and near-field interactions, where near-field interactions are computed exactly and far-field interactions are approximated using Lagrange interpolation for speed, explained in the paragraph below.
The partitioningprocedure in FFM is recursive, where the recursion depth tree_depth controls the size of the edgel =
E 2 tree_depth of the box.
An interaction is defined to be in the far-field if the distance between the twocenter points of the boxes exceeds 2 l , i.e. IsFarField :
=
∥ x center − y center ∥ ≥ 2 l .
While l for eachbox will decrease with the number of divisions, this rule ensures a fixed minimal distance for a givendepth for far-field interactions.
Figure 3 illustrates how far(green) and near(orange)-field interactionsarise between X and Y when tree_depth increases.
Lagrange interpolation
We review Lagrange interpolation used for far-field approximations in
FFM.
Given a function f ( x ) : [ − 1 , 1] → R and r +1 unique points s i ∈ [ − 1 , 1] , i = 0 , . . .
, r , thereexists a unique polynomial p r ( x ) of degree ≤ r that interpolates f at p r
( s i ) =
f ( s i ) .
The Lagrangepolynomial is given by p r ( t ) =
(cid:80) ri =
f ( s i )
L i ( t ) , where L i ( t ) = (cid:81) rj
=0 ,j ̸
= i ( t − s j ) (cid:81) rj =0 ,j ̸
= i (
s i − s j ) , i = 0 , . .
r.
We are free to chose the degree r as well as the points s i to interpolate through.
The choice of s i isespecially important in minimizing large oscillations around the edges of the interpolation interval(Runge’s phenomenon [12]).
For this reason, Chebyshev nodes of the second kind are used [7]s i =
cos θ i , where θ i = iπr , i = 0 , . . .
, r .
Interpolating k ( x , y )
By noticing that k ( x , y ) is a bivariate function, we can apply Lagrangeinterpolation twice, thus interpolating k ( x , y ) as k ( x , y )
≈ (cid:80) r
X i =
1 L i ( x ) (cid:80) r Y j =1 k (cid:0) s x i , s y j (cid:1)
L j ( y ) .
Here r X , r Y denotes the number of the interpolation nodes and s xi , s yj ∈ R D denotes thegrid of interpolation nodes for B Xp and B Yq .
Note that since x ∈ R D , we take L i ( x ) :=
[Equation]
i = 0 , . .
r. These operations can be vectorized and computed sequentially on-the-fly with linear memory footprint v ≈ L TX ·
[Equation]
) , which is done by firstcomputing v 1 , then v 2 and lastly v .
A far-field KMVM between two boxes v p,q = k ( B Xp , B Yq ) · b qis then approximated by using double Lagrange interpolation according to Figure 4.
[Figure]
Faster-FFM (F 2 . 5 M and F 3 M)
To fully leverage the port of FFM to GPU, we enhance FFM with novel approximation procedures forimproved complexity and memory optimizations to scale to n = 10 9 .
We coin this improved version
Faster-FFM (F 3 M).
The capabilities of F 3 M against previous methods are summarized in Table 1.
4.1 CPU to GPU optimizations
In FFM, every computation is serial and on CPU.
When moving to GPU, we have parallelizedall major computations.
These parallelizations are non-trivial and require low-level algorithmicoptimizations, with challenges such as:
Box-to-threadblock alignment – A major challenge in the implementation of both the parallelfar-field and near-field computations was correctly aligning thread blocks to boxes.
This aligning requirement imposed non-trivial boundary conditions on data indexing when using shared memory.
We present how the parallelization is done for calculating near-field interactions in Appendix I.
No native sorting methods – We found that LibTorch [25] sorting methods often led to out-ofmemory (OOM) due to allocation of large long-type vectors on GPU.
When n = 10 9 , this impliesallocating 8GB of memory, 25% of the 32GB card used, making it a necessity to avoid native sortingmethods.
Ensuring interactions are sorted – To avoid any unnecessary sorting, we ensure that thematrix containing interactions is always sorted.
To do this without using sorting algorithms,we had to design a non-trivial algorithm that would recursively retain order during division.
We refer to the get_new_interactions function in n_tree.cu for exact implementation.
In-place grouping data on boxes – Due to infeasible LibTorch sorting methods, we additionally had to design an algorithm that found a permutation that would group X into its corresponding boxes in linear time and memory.
We useda non-trivial count and increment-based strategy.
We refer to the box_division_cum_hashand box_division_assign_hash function inn_tree.cu for exact details.
However, we found that these optimizations andporting alone were not enough to scale to n =10 9 on 3D datasets, as Figure 5 demonstrates.
FFM doesn’t remove empty boxes or handle boxes with few points in them and keeps exponentiallycreating new empty boxes and interactions, thus leading to out-of-memory (OOM) errors on nonuniform data (see Appendix 9).
4.2 Scaling to n = 10 9 on GPU (F 2 . 5 M)
In this section, we detail the memory enhancements that allow F 3 M to consistently scale to n = 10 9 .
Removing empty boxes with hash list indexing To ensure linear memory on GPU, we only keepa reindexing vector σ of size n x (resp.
n y )
in memory during the computation of the algorithm inaddition to a list of interactions and box centers.
This reindexing vector rearranges the data pointsso they appear in the order of the box they belong to.
We optimize both the computation and thememory footprint of these objects by avoiding recursive formulas and hash lists.
Naively, points can be assigned to boxes by direct comparison to all existing box centers.
Asthe number of centers grow exponentially with depth tree_depth , this method quickly becomespathological.
To amend this, we propose a linear complexity formula to retrieve the box index
β
i apoint x ∈ X ⊂ R D belongs to β i =
[Figure]
, where α d denotesthe minimum value of X in dimension D and x d is the value of x in dimension D .
To preventthe number of boxes from growing exponentially, we remove empty boxes with each division.
Toassign points to the corresponding boxes, we use a hash list to store β
i and the order i .
We can thengroup points { x i } i =1 to their respective ordering i using the hash list in O ( n ) time in contrast to
O ( n · 2 D · tree_depth ) by direct computation.
Handling boxes with few points with small field
In cases when the number of points in each boxcan vary greatly, we separately consider the interactions where the number of points in boxes is small.
Hence, we say that there is a small field interaction between boxes B Xp , B Xq if both have a smallnumber of points, i.e. if
[Equation] | B Xq | ≤ ρ , for some threshold number ρ .
To minimize the computationsneeded, ρ can be set to ρ =
r X + r Y .
This intuitively allows F 3 M to directlycompute interactions that are too small to benefit from interpolation savings(i.e.
| B Xp | + | B Xq | ≤ r X + r Y ), thus limiting memory usage by stoppingpartitions from dividing further than necessary.
In higher dimensions where the division rate is faster, ρ can be set to a higher value to limit memory usageat the expense of more direct computations which are slower.
Sparse grids As the number of Lagrange polynomials increases exponentially with dimension, weimplement Sparse grids [29] to allow for a finer selection of interpolation nodes.
With sparse grids,the number of nodes needed grows slower [20], thus saving memory.
We give an example of a sparsegrid versus a full grid in 2D in Figure 6.
4.3 Speeding up F 2 . 5 M (cid:0) F 2 . 5 M → F 3 M (cid:1)
Smoothness criteria FFM speeds up its computations with minimal loss in accuracy by selectivelyinterpolating interactions that are far apart.
To improve speed, we introduce the smoothness criterionto widen the selection of interactions that can be interpolated with minimal loss in accuracy.
For a
Gaussian Kernel k ( x , y ) = exp (cid:16) ∥ x − y ∥ {+
 [Equation]+} 2 (cid:17) , with lengthscale γ , the smoothness criteria is definedas is_smooth :=
[Equation]
≤ η betweenan adjacent interaction of boxes B Xp , B Yq .
The quantity computed can be understood as “Effective
Variance” (EV), as it considers total variation in the exponent of the Gaussian kernel.
We justify thesmoothness criteria with the following proposition.
Proposition 1.
Consider x , y ∈ X
⊂ R d such that d ( x , y ) :
= ∥ x − y ∥ 2 2 γ= 12 γ 2 (cid:80) d
i ( x ( i ) − y ( i ) )
2 ≤η < 1 for all x , y .
When interpolating k ( x , y ) =
exp ( − d ( x , y )) using bivariate Lagrange interpolation L r ( x , y ) :
= L TX · K · ( L Y · b ) with degree r = 2 p , for any p ∈ N > 0 there exist nodes s x , s yfor L r ( x , y ) such that the pointwise interpolation error is bounded by O ( η p +1 ) .
See Appendix D for proof.
Hence for small η < 1 , we see that the error becomes smallfor well specified L r ( x , y ) .
To avoid calculating the sample variance during computationswhich costs O ( n ) , we exploit that data is partitioned into hypercubes with a known edge Eand take the upper bound of the variance in each cube as E 2 4 along a dimension.
A prooffor this bound is provided in Appendix B. Adjacent interactions are then classified as smoothwhen (cid:80) Dd E 2 4 · 2 γ ( EV of B Xp ) + E 2 4 · 2 γ ( EV of B Yq ) =
D E 2 γ 2 · 4 ≤ η which only costs O (1) to compute.
[Figure]
Adaptive far-field approximation
To further improve speed we introduce an adaptive rule to select thenumber of interpolation nodes usedwhen calculating far-field interactions.
Error bounds for multidimensional Lagrange interpolation have been proposed in [22], however, these boundscannot be directly used to create anadaptive interpolation rule.
We thussimulate KMVM errors for k ( X , Y ) ·b where X , Y are uniformly distributed and b is normally distributed.
We fix a distance between X and Yand vary the squared of this distancebetween boxes against nodes in Figure 7.
We use a Gaussian Kernel with γ = 1 √ 2 .
Based on Figure 7, we use the following rule for selecting the number of interpolation nodes forfar-field interactions
[Equation]
where r is the number of nodes chosen to interpolate with in the general case.
Barycentric lagrange interpolation We slightly improve the complexity further by implementingbarycentric Lagrange interpolation [7] evaluated at the chebyshev nodes of the second kind.
As thisis a well-known technique, we refer to the appendix for more details.
4.4 Complexity
The time complexity of FFM is O ( n log ( n ))
[3] and we use a similar derivation strategy for F 3 M toobtain a complexity that is dependent on the effective variance limit η (chosen parameter) and thebox width E (data).
We first present two propositions needed to derive the complexity of F 3 M.
Proposition 2.
A far-field interaction between two boxes containing n x and n y points respectivelyhas time complexity O ( n ) , where n = max( n x , n y ) .
Proposition 3.
Given n data points in dimension D , the maximum number of divisions
Tree max divisions is given by
[Equation]
With the above results, the complexity of FFM is taken as the maximum number of divisionsmultiplied by the complexity of far-field interactions at each division which yields O ( n log ( n )) .
We remark that near-field interactions between boxes containing only 1 data point have linear timecomplexity, hence the results hold.
Theorem 1.
Given a KMVM with edge E (dependent on data X , Y ), lengthscale γ , effective variancelimit η , n data points and data dimension D , F 3 M has time complexity O ( n · log 2 (cid:16) D ·E 2 γ 2 · 4 · η (cid:17) ) , whichcan be taken as O ( n · log 2 (cid:16) Cη (cid:17) ) where C ∝ D ·E 2 γ 2 .
Memory footprint As our implementation uses the same partitioning strategy as FFM, the theoretical memory complexity remains O ( n ) for F 3 M (see [3] for proof).
However, this does notaccurately reflect the memory footprint of the actual implementations, whose memory mostly dependson the number of interactions stored.
We summarize these memory footprints for FFM and F 3 M in
Theorem 2 below.
Theorem 2.
The number of interactions M i against tree depth i of FFM and F 3 M grows as
O (cid:16) M i − 1 2 2 · D − m far i (cid:17) and
O
[Equation]
− ( m empty i ) 2 − m far i − m smooth i − m small i (cid:17) respectively.
Here M − 1 = 12 2 D and m far= m smooth= m small= m empty= 0 andm far i , m smooth i , m small i , m empty i denotes the number of far-field, smooth field, small field interactions andthe number of empty boxes respectively at depth i > 0 .
We see that the additional approximations presented in F 3 M also impacts memory footprint, as theadditional ( m empty i ) 2 , m smooth i , m small i terms removes a substantial amount of interactions at each i ,significantly slowing down the growth of interactions, reducing memory growth. The efficacy of( m empty i ) 2 , m smooth i , m small i is widely dependent on data.
As an example, data where points are veryclose to eachother would significantly benefit m smooth i more, as the closeness of points would implymore smooth interactions.
If points are sparsely spread out ( m empty i ) 2 , m small i would provide the mostbenefit as they remove empty boxes and stops boxes with few points to divide unnecessarily.
Allproofs can be found in Appendix E.
Experiments
We demonstrate the utility of F 3 M over a variety of experiments using the Gaussian kernel k ( x , y ) =
exp − ( ∥ x − y ∥ 2 2 γ 2 ) .
2 We generate data such that the EV (see section 4.3) varies between 0 . 1 , 1 , 10 fordata of sizes n = 10 6 , 10 7 , 10 8 , 10 9 .
The parameters used for F 3 M are η = 0 .
5 andr = 2 D , 3 D , 4 D with a cap at r = 2048 .
The error for the approximated KMVM product ˆ v iscalculated as Relative error :
=
∥ ˆ v − v ∥ 2 ∥ v ∥ 2 , where the true KMVM product v is obtained by calculatingthe full KMVM on a subset X ′ consisting of the first 5000 points in X against the entire dataset indouble precision, i.e. v = k
( X ′ , X ) · b , where we fix b ∼ N (0 , I n ) .
All experiments were run on
NVIDIA V100-32GB cards, where the data is fitted entirely on the GPU.
These cards were chosensince the extra graphic memory is necessary to fit the data on one card when n = 10 9 .
It should benoted that n = 10 9 can only be run up to D = 3 , as X and b itself cannot fit in memory for higherdimensions with the GPUs we had available.
For details on how F 3 M scales across multiple GPUs,see Appendix G. [Table] [Table]
[Table] log( n )) 1. KMVM experiments We consider a wide variation of generated datasets to simulate different realworld scenarios to test F 3 M on.
For the k ( X , X ) -case we consider uniformly and normally distributeddata ( D = 1 , 2 , 3 , 4 , 5 , 6 , 7 ) together with data simulated from Brownian motion, fractional Brownianmotion, and Clustered data ( D = 1 , 2 , 3 ).
For the k ( X , Y ) -case we consider uniformly distributed xand normal distributed y ( D = 1 , 2 , 3 , 4 , 5 , 6 , 7 ).
See Appendix 9 for visualizations of data.
We haveto consider smaller n for the k ( X , Y ) -case when D ≥ 3 , as twice the amount of data needs to bestored.
For D = 3 / (4 , 5) / (6 , 7) we instead consider at most n = 5 · 10 8 / 2 . 5 · 10 8 / 10 8 .
It shouldbe noted that D = 7 is a hard limit for geometric partitioning-based methods, since for D = 8 , wewould have 2 8 · 2 · 2 8 · 2 ≈ 4 .
3 · 10 9 interactions after only 2 divisions.
This number of interactionscannot even be represented by a 32-bit integer.
We summarize the runs in Table 3 and plot the errorand time complexity in Figure 8 for each dataset when D = 3 .
We find that F 3 M maintains sub-linearempirical complexity up to D = 6 , where we have to set small field limit ρ to a larger number tonot run out of memory.
Further, the error increases in the higher dimensions since we use fewernodes per dimension when interpolating, owing to the sparse grid technique.
We note that D = 7 hasfaster run times than D = 6 which is explained by that for some values of EV, D = 7 doesn’t runwith acceptable errors which skew the run time to datasets where a larger portion of the data can beinterpolated.
We further replicate the data used in the first experiment in [3] and compare F 3 M against FFM (CPU)in Table 2.
Kernel Ridge Regression experiment We apply F 3 M to FALKON [24], where we replace their
KMVM operation with F 3 M and compare performance and speed in solving Kernel Ridge Regression(KRR).
The KMVM operation currently used for smaller dimensions is KeOps [11].
Given some data
X ∈ R N ×
d we want to find the solution
α
=
( k ( X , X ) + λI ) − 1 b where λ is the ridge parameter thatstabilizes the inverse.
FALKON is a Nyström approximation based solver that requires a subsample
[Table] X ′ ∈ R M × d of X to approximate the inverse computation.
We focus the experiments on tall andskinny data and take n = 10 9 , d ≤ 3 with M = 10 5 for all experiments.
We consider uniformly andnormally sampled data, the Open Street Map (OSM) dataset [1] and a classification task on the NYC
Taxi dataset [2], where we predict whether the customer will tip based on trip distance, trip time andfare cost.
To construct b on synthetic problems, we first take a subset D ∈ R 1000 × d of X and sampleα ∼ N (0 , I 1000 × 1000 ) .
We then calculate b
= k ( X , D ) · α + ε , where ε ∼ N (0 , 0 . 1) .
We run KRRfor EV = 0 .
1 , 1 , 10 on synthetic data, and report the average R 2 (AUC for NYC Taxi) and trainingtime in Table 4.
For the real world datasets OSM and NYC Taxi, we fix the lengthscale using themedian heuristic proposed in [16] averaged our results over the 3 runs.
Ablation study between FFM(GPU) and F 3 M As much of the improved performance can beattributed to our GPU implementation, we conduct an ablation study of FFM(GPU) against F 3 M and
KeOps in Table 5.
We first present KMVM run times averaged over D = 3 and real-world datasets
OSM and NYC Taxi.
For KeOps, we only computed the KMVM on uniform data.
Since KeOps isan exact method, the dataset distribution has no effect on computational time.
Here, the smoothnesscriteria and adaptive far-field technique improve computational time.
We find that F 3 M achieves aspeed-up between 2 .
− 33 .
3 × against FFM(GPU) and 8 .
− 8500 × speed-up against KeOps.
[Table]
Gaussian process regression experiments We further compare F 3 M as a drop-in KMVM operationapplied to Black-box Matrix Multiplication [15] for Gaussian Processes, compared to KISS-GP
[35],an approximate Gaussian process using cubic interpolation for kernel approximation.
We mimic thesetup in [32] and consider the datasets 3DRoad, Song, Buzz and House Electric, where we apply
PCA to the last three datasets and take the 3 first principal components for a fair comparison against
KISS-GP, which is limited by D ≤ 3 .
We demonstrate the results in Table 6.
As exact GP using F 3 Mdemonstrates competitive results even when compared to SVGP [18] and SGPR [30], we hypothesizethat many high-dimensional datasets conform to the manifold hypothesis [13], allowing F 3 M to bewidely applicable out-of-the-box even in high-dimensional settings.
[Table]
Limitations and Further Research
This work has introduced and implemented F 3 M on GPU, which enables fast KMVM for tall andskinny data up to n = 10 9 .
F 3 M has improved complexity which also is controllable throughη [Table] and retains linear memory.
Experiments in higher dimensions also exhibit linear complexity,however requiring more nodes for lower errors.
F 3 M can further be directly used as a drop-in KMVMoperation, as demonstrated with FALKON and Gaussian process regression, achieving significantspeedups and competitive performance on both tasks.
As an interpolation based approximation method, F 3 M is still limited by the exponential growth of interpolation nodes with respect to D ,although removing empty boxes, small field and sparse grids allow KMVM for D ≤ 7 .
A fruitfuldirection would be to extend ideas in F 3 M to accommodate higher-dimensional data by consideringrandomized partitioning [4], decoupling the dependency on D in geometry based partitioning.
Further,an exact characterization of how m far i , m smooth i , m small i , m empty i grows is left to future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B, D,
E.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
[Figure]
A Synthetic data
[Figure]
Note on synthetic datasets
We generated synthetic datasets of different types to measure theability of F 3 M to deal with dense or sparse data.
Dense datasets were generated as independentsamples with either uniform or normal distributions.
Clustered datasets were generated by samplingcluster centers from a normal distribution, and then recursively sampling sub-cluster centers froma normal distribution with reduced standard deviation and centered at each cluster center, until thedesired number of points is attained.
Fractional Brownian Motion and Brownian Motion sampleswere generated as samplings of Fractional Brownian Motion paths with respective Hurst index 0.75and 0.5.
Figure 9 shows samples of each dataset type in the 2D case.
B Note on maximal variance on an interval
Proposition 4.
Consider a random variable X ∈ R with finite variance with m = inf X and
M = sup X .
Then Var ( X ) ≤ ( M − m ) 2.
Proof.
Define a function g by g ( t ) =
E (cid:2) ( X − t ) 2 (cid:3) .
Computing the derivative g ′ , and solvingg ′ ( t ) =
− 2 E
[ X ] + 2 t = 0 yields that g achieves its minimum at t
= E [ X ] (note that g ′′ > 0 ).
Now, consider the value of the function g at the special point t =
M + m 2 .
It must be the case that
Var[ X ] =
g ( E [ X ]) ≤ g (cid:0) M + m 2 (cid:1) .
Evaluating yields the expression
[Equation]
Since X − m ≥ 0 and X − M ≤ 0 , we have
[Equation]
implying that
[Equation]
Hence
[Equation]
C Details on barycentric Lagrange interpolation
The barycentric lagrange interpolation is written as
[Equation]
where w i are known as the barycentric weights.
In case of singularities, i.e. when t = s j , we set
L i ( s j ) = δ ij .
In particular, [7] proposes Chebyshev nodes of the second kind s i
= cos θ i ,
[Equation]
i =
iπr , i = 0 , . . .
, r .
This choice of nodes combined with the scale invariance property of the barycentric form makes the calculation of w i particularly easy
[Equation]
and reduces the complexity of calculating w i from O ( r 2 ) to O ( r ) .
The Lagrange interpolationpolynomial can then be expressed as p r ( t ) =
(cid:80) ri
= wit
[Equation]
f i .
D Smooth field proof
Proposition.
Consider x , y ∈ X ⊂ R d such that d ( x , y ) :
= ∥ x − y ∥ 2 2 γ= 12 γ 2 (cid:80) di
( x ( i ) − y ( i ) ) 2 ≤ η <1 for all x , y .
When interpolating k ( x , y ) =
exp ( − d ( x , y )) using bivariate Lagrange interpolation
L r ( x , y ) :
= L TX · K · ( L Y · b ) with degree r = 2 p , for any p ∈ N > 0 there exist nodes s x , s y for
L r ( x , y ) such that the pointwise interpolation error is bounded by O ( η p +1 ) .
Proof.
Note that k
( x , y ) is analytic in d ( x , y ) with Taylor expansion given by T p
e − d ( x , y )
= 1 −d ( x , y )
+ ...
+ O (cid:0) d ( x
, y ) p +1 (cid:1) .
With d ( x , y ) ≤ η , it follows that | T p k ( x , y ) − k ( x , y ) | ≤ O ( η p +1 ) .
Using triangle inequality, we have |L r ( x , y ) − k
( x , y ) | ≤ | T p k ( x , y )
− k ( x , y ) | + | T p k
( x , y ) −
{+L r ( x , y ) | ≤ η p +
+}
+ | T p k
( x , y ) −L r ( x , y ) | .
We note that L r ( x , y ) contains all the terms of the Taylorexpansion, and we can thus choose the nodes s x , s y of L r ( x , y )
such that | T p k
( x , y ) −L r ( x , y )
| = 0as long as r = 2 p , meaning the polynomial orders are matched.
Note that the same proof strategy can be applied to any kernel k that admits a Taylor expansion.
E Complexity
Proposition.
A far-field interaction between two boxes containing n x and n y points respectively hastime complexity O ( n ) , where n = max( n x , n y ) .
Proof.
Far-field interactions are calculated as
[Equation]
[Equation]
As r X , r Y are independent of n x , n y , the complexity becomes O ( n ) .
Further see [3] for alternativeproof.
Proposition.
Given n data points in dimension D , the maximum number of divisions Tree max divisionsis given by
[Equation]
Proof.
To see this, simply solve for
[Equation]
Further see [3] for alternative proof.
Theorem.
Given a KMVM with edge E (dependent on data X , Y ), lengthscale γ , effective variancelimit η , n data points and data dimension D , F 3 M has time complexity O ( n · log 2 (cid:16) D ·E 2 γ 2 · 4 · η (cid:17) ) , whichcan be taken as O ( n · log 2 (cid:16) Cη (cid:17) ) where C ∝ D ·E 2 γ 2 .
Proof.
Recall that near-field interactions can be smoothly interpolated when D ·
[Figure]
[Table]
Then all interactions will be interpolated when tree_depth ≥ log 2 (cid:16) D
[Table]
(cid:17) , which implies we cantake Tree max divisions = log 2 (cid:16) D ·E 2 γ 2 · 4 · η (cid:17) .
Hence the complexity is O ( n · log 2 (cid:16) D ·E 2 γ 2 · 4 · η (cid:17) ) .
Theorem.
The number of interactions M i against tree depth i of FFM and F 3 M grows as
O (cid:16) M i − 1 2 2 · D − m far i (cid:17) and
[Equation]
respectively.
Here M − 1 = 12 2 D and m far= m smooth= m small= m empty= 0 andm far i , m smooth i , m small i , m empty i denotes the number of far-field, smooth field, small field interactions andthe number of empty boxes respectively at depth i > 0 .
Note that for i > 0 , these are dependent ondata.
Proof.
We prove through induction that the recursion holds for F 3 M. We start with the base case
M 0 = 1 , since at depth i = 0 , we only have one box and hence only one interaction.
M 0 =
[Equation]
m
[Equation]
= 1 .
Clearly at depth 0, there can not be anyempty boxes or possible approximations.
For the induction step, M i =
M i − 1 2 2 · D − ( m empty i )
2 −m far i − m smooth i − m small i .
To get M i +1 , each box at depth i is first divided into 2 D , hence thenumber of interactions grows by 2 2 · D .
At depth i + 1 , we can further remove ( m empty i +1 )
2 interactionsbetween empty boxes and further compute m far i +1 + m smooth i + + m small i +1 interactions.
{+Then M i +
+}
=
[Equation]
− m small i +
[Equation]
[Equation]
− m small i +1 .
Thus the base case and induction step holds whichcompletes our proof.
This proof also covers FFM, since FFM can be as a special case for F 3 Mwithout removing empty boxes, smooth field and small field computation.
F Algorithm summary
We present a summary of the algorithm presented in FFM in Algorithm 1 and the modifications F 3 Mdoes in boldface .
Input: Datasets X , Y , b , kernel k , average points threshold ζ
Result: v = k
( X , Y ) ·
b
Initialize treecodes τ x = T ( X ) , τ y =
T ( Y )
Initialize near-field interactions as I near = { 0 , 0 }
Initialize output v =while | I | > 0 and MaximumBoxSize ( τ x )
> ζ and MaximumBoxSize ( τ y ) > ζ do
Divide τ x , τ y
Calculate interactions left I :
= f
( I near )
Partition I to {I near , I far , I smooth }
Throw away interactions that are too far from eachother
Compute far-field interactions v += FarFieldCompute ( τ x , τ y , I far )
Compute smooth field interactions v += FarFieldCompute ( τ x , τ y , I smooth ) ; end
Compute remaining near-field interactions v += NearFieldCompute ( τ x , τ y , I near )
Instead of comparing the average box size to ζ we compare the maximum box size.
When points arenon-uniformly distributed, taking the maximum ensures that we don’t compute near-field interactionson boxes with many points, since it will be inefficient.
G Scalability Analysis
We conduct a scalability analysis over N GPU = 1 , 2 , 4 , 8 .
We parallelize the KMVM productby considering the k ( X , X ) -case and divide the work onto multiple GPUs by partitioning eachsubproduct of the KMVM (see for Figure 10 an example when N GPU = 8 ).
We take X to be Uniformand 3 dimensional.
We present results in Figure 11.
[Figure]
G P U :
G P U :
G P U :
G P U :...
[Figure]
We also use nvprof to analyze the % of peak throughput of the V100 cards F 3 M can utilize.
Werun nvprof for 3 dimensional uniform data for n = 10 6 , 10 7 , 10 8 , 10 9 .
We present our results in
Figure 12.
H Impact of η and r on performance
The performance of F 3 M is tuned by choosing η and r to trade speed against accuracy.
In Figure 13we plot how different choices of η and r impacts computation time for F 3 M on 3D data.
I Implementation overview
We provide a skiss of how data is stored and used for F 3 M in Figure 14.
[Figure]
We provide an illustration on near field computations are carried out for F 3 M in Figure 15.
[Figure]
J Additional ablation study between F 2 . 5 M and F 3 M
We provide and additional ablation study between F 2 .
5 M and F 3 M in Table 7.
The results arequite similar to the comparison bewteen F 3 M and FFM (GPU).
Here, we see that smooth field andadaptive far-field approximation both improves speed and also memory usage as smooth field helpsapproximate more interactions.
We can infer empirically that the m smooth i term in Theorem 2 issignificant.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
TheoreticalComplexity
[Figure]