Giga-scale Kernel Matrix-Vector Multiplication on GPU
Abstract
Kernel matrix-vector multiplication (KMVM) is a foundational operation in machine learning and scientific computing.
However, as KMVM tends to scale quadratically in both memory and time, applications are often limited by these computational constraints.
In this paper, we propose a novel approximation procedure coined Faster-Fast and Free Memory Method (F 3 M) to address these scaling issues of KMVM for tall ( 10 8 ∼ 10 9 ) and skinny ( D ≤ 7 ) data.
Extensive experiments demonstrate that F 3 M has empirical linear time and memory complexity with a relative error of order 10 − 3 and can compute a full KMVM for a billion points in under a minute on a high-end GPU, leading to a significant speed-up in comparison to existing CPU methods.
We demonstrate the utility of our procedure by applying it as a drop-in for the state-of-the-art GPU-based linear solver FALKON, improving speed 1.5-5.5 times at the cost of < 1% drop in accuracy.
We further demonstrate competitive results on Gaussian Process regression coupled with significant speedups on a variety of real-world datasets.
Introduction
Kernel matrix-vector multiplication (KMVM) is one of the most important operations needed in scientific 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 size n , KMVM using direct computation has complexity and memory footprint O ( n 2 ) , both unfeasible for modern large scale applications where n ≈ 10 9 is becoming increasingly common.
Pioneering contributions presented in the Fast Multipole Method (FMM)
[10] amend the complexity of these problems to O ( n log ( ϵ − 1 )) , where ϵ is the chosen error tolerance, with varying reductions in memory footprint for data restricted to dimension D
= 2 .
Subsequent developments in [8, 17] mainly focused on extending approximations for a broader set of kernels for a fixed dimensionality D ≤ 3 , tailored for problems in physics with narrow data such as electrostatics, stellar dynamics, Stokes flow, and acoustic problems, amongst others.
In this paper, we introduce Faster-Fast and Free Memory Method (F 3 M), a novel algorithm built upon the FFM [3] framework to perform KMVM on a GPU for tall and skinny ( D ≤ 7 ) data of order
[Figure]
n ∼ 10 9 in under a minute with user-specified error tolerance, providing between 2 − 8500 times speed-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 a kernel 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.
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 developments have been made to improve both of these constraints, ranging from hardware acceleration using GPUs in KeOps [11], to various approximation techniques proposed in [36, 33, 34, 3, 9].
In this work, we focus 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 product on-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 , as the GPU hardware acceleration allows the KMVM product to be computed in less than a second on a conventional GPU.
Moreover, the method places no constraint on the number of features D it can be applied to, making it favourable for KMVM on medium size datasets.
In application contexts, KeOps is 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 as the O ( n 2 ) complexity would imply a computational time of 10 6 seconds, or roughly 11 days.
To overcome this billion points barrier, Aussal et al.
[3] deploys a geometric space partitioning scheme, and proposed the Fast and Free Memory Method (FFM), a KMVM approach that extends the FMM [10] family of algorithms.
In contrast to 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 any conventional kernel and further enables the user to trade off accuracy with computational efficiencies by 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 many machine learning techniques infeasible.
Further, as recursive partitioning of the data space scales poorly with D [5], both FMM and FFM can only be applied to D ≤ 3 data, a price to pay for the speed-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 (bottom row in Appendix 10).
Our contribution.
To surpass the billion point barrier while maintaining high-speed and stable computation, we propose F 2 .
5 M and our main algorithm F 3 M, the first pair of KMVM algorithms that 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 stabilize and improve the original geometric partitioning scheme in FFM to significantly reduce memory constraints, leading to a relaxation of dimensionality constraints from 3 to 7 .
At last, we introduce an adaptive far-field and smooth field approximation scheme for kernel interpolation, resulting in our main algorithm F 3 M, which runs 2 .
− 33 .
3 times quicker and more stable than a direct port of FFM on GPU.
See Fig.
1 and Table 1 for an overview and comparisons of the methods.
We summarise our contribution as follows:
We propose Faster-Fast and Free Memory Method ( F 3 M ) , a KMVM algorithm building on top of FFM by applying multiple low-level enhancements, GPU parallelisation, and algorithmic computational and memory enhancements, allowing for KMVM operations on n ≤ 10 9 data in under a minute.
Codebase is released here [14].
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– times 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 gradient solver 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 competitive results on Gaussian process regression against KISS-GP [35], SVGP [18] and SVGR [30] with significant speed-ups.
Background
The FFM method considers KMVM for a kernel k evaluated on two data matrices X , Y and b 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 x and 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
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 X and Y such that certain calculations can be approximated in a fast manner, based on the pairwise distances between partitions.
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 and similarly 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 corresponds to a subset of rows in the data matrix.
Let us also denote b q as the partition of b j ’s grouped with the same indices as B Yq .
To calculate the
KMVM between two boxes B Xp , B Yq with the grouped vector b q , for each x i ∈ B X q , we compute
[Equation]
with v p,q =
[ v p,q 1 . . .
v p,qn x ] .
Now the target v can be computed as v = σ ([ v p =1 , . .
. , v p = P ] ⊤ ) , where v p =
(cid:80) Qq
=1 v p,q , P, Q denote the total number of boxes and σ ( · ) a permutation such that v pi appear in 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 partitioning procedure in FFM is recursive, where the recursion depth tree_depth controls the size of the edge l =
E 2 tree_depth of the box.
An interaction is defined to be in the far-field if the distance between the two center points of the boxes exceeds 2 l , i.e. IsFarField :
=
∥ x center − y center ∥ ≥ 2 l .
While l for each box will decrease with the number of divisions, this rule ensures a fixed minimal distance for a given depth for far-field interactions.
Figure 3 illustrates how far(green) and near(orange)-field interactions arise 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 , there exists a unique polynomial p r ( x ) of degree ≤ r that interpolates f at p r
( s i ) =
f ( s i ) .
The Lagrange polynomial 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 is especially 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 Lagrange interpolation 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 xi , s yj (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 the grid of interpolation nodes for B Xp and B Yq .
Note that since x ∈ R D , we take L i ( x ) :
= (cid:81)
[Equation]
i = 0 , . .
r. These operations can be vectorized and computed sequentially on-the-fly with linear memory footprint v ≈ L TX ·
[Equation]
(cid:124) (cid:123)(cid:122) (cid:125)
[Equation]
) , which is done by first computing v 1 , then v 2 and lastly v .
A far-field KMVM between two boxes v p,q = k ( B Xp , B Yq ) ·
b q is 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 for improved 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 parallelized all major computations.
These parallelizations are non-trivial and require low-level algorithmic optimizations, with challenges such as:
Box-to-threadblock alignment – A major challenge in the implementation of both the parallel far-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.
To minimize memory usage of box and block indicators for our implementation, we represented the box belonging of each point as index intervals (i.e. box 1 consists of points with i ∈ [1 , . .
. , 500] and box 2 with i ∈ [501 , . .
. , 1337] , etc.)
and modulo arithmetic to infer the block belonging.
This clearly requires that the points are sorted or grouped according to their box belonging.
However, as we detail in the next paragraph, arranging the points could not be done straightforwardly with native sorting methods.
We further illustrate how 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 implies
[Figure]
allocating 8GB of memory, 25% of the 32GB card used, making it a necessity to avoid native sorting methods
In-place grouping data on boxes – Due to infeasible LibTorch sorting methods, we additionally had to design an algorithm that finds a permutation that would group X into its corresponding boxes in linear time and memory.
We used a count and increment-based strategy that would:
Count the number of points in each box during the assignment operation ( O ( n ) ) and store the count in a vector ξ .
Then run a cumulative sum over ξ , starting from 0.
Initialize a n long permutation vector π .
Us-
ing the counting vector ξ , we would re-run the assignment operation and arrange a point with index i as following π [ atomicAdd ( ξ ′ [ box_index ] , 1) + ξ [ box_index ]] =
i , where ξ ′ is a running count of points in each box.
We specifically have to use the function atomicAdd to increment the count for each box in parallelized GPU environments to avoid thread locks.
We refer to the box_division_cum_hash and box_division_assign_hash function in n_tree.cu for exact details.
Ensuring interactions are sorted – To avoid any unnecessary sorting, we ensure that the matrix containing interactions is always sorted by recursively dividing old interactions.
We illustrate the procedure in Figure 5.
We refer to the get_new_interactions function in n_tree.cu for the exact implementation.
However, we found that these optimizations and porting alone were not enough to scale to n = 10 9 on 3D datasets, as Figure 6 demonstrates.
FFM doesn’t remove empty boxes or handle boxes with few points in them and keeps exponentially creating new empty boxes and interactions, thus leading to out-of-memory (OOM) errors on non-uniform data (see Appendix 10).
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 =
Removing empty boxes with hash list indexing To ensure linear memory on GPU, we only keep a reindexing vector σ of size n x (resp.
n y )
in memory during the computation of the algorithm in addition to a list of interactions and box centers.
This reindexing vector rearranges the data points so they appear in the order of the box they belong to.
We optimize both the computation and the memory 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.
As the number of centers grows exponentially with depth tree_depth , this method quickly becomes pathological.
To amend this, we propose a linear complexity formula to retrieve the box index β i a point x ∈ X ⊂ R D belongs to β i =
[Figure]
, where α d denotes the minimum value of X in dimension D and x d is the value of x in dimension D .
To prevent the number of boxes from growing exponentially, we remove empty boxes with each division.
To assign points to the corresponding boxes, we use a hash list to store β
i and the order i .
We can then group 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 box can 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 X q if both have a small number of points, i.e. if
| B Xp | + | B Xq | ≤ ρ , for some threshold number ρ .
To minimize the computations needed, ρ can be set to ρ =
r X + r Y .
This intuitively allows F 3 M to directly compute 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 stopping partitions from dividing further than necessary.
In higher dimensions where the division rate is faster, ρ can be set to a higher value to limit memory usage at the expense of more direct computations which are slower.
Sparse grids As the number of Lagrange polynomials increases exponentially with dimension, we implement 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 sparse grid versus a full grid in 2D in Figure 7.
4.3 Speeding up
[Equation]
Smoothness criteria FFM speeds up its computations with minimal loss in accuracy by selectively interpolating interactions that are far apart.
To improve speed, we introduce the smoothness criterion to 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
[Equation]
[Figure]
[Equation]
(cid:17) , with lengthscale γ , the smoothness criteria is defined as is_smooth :=
[Figure]
between an 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 the smoothness criteria with the following proposition.
Proposition 1.
Consider x , y ∈ X
⊂
R d such that d ( x , y ) :=
∥ x
[Equation]
y
[Equation] [Equation]
1 2 γ 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 y for 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 small for well specified L r ( x , y ) .
To avoid calculating the sample variance during computations which costs O ( n ) , we exploit that data is partitioned into hypercubes with a known edge E and take the upper bound of the variance in each cube as E 2 4 along a dimension.
A proof for this bound is provided in Appendix B.
Adjacent interactions are then classified as smooth when (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.
Adaptive far-field approximation
To further improve speed we introduce an adaptive rule to select the number of interpolation nodes used when calculating far-field interactions.
Error bounds for multidimensional Lagrange interpolation have been proposed in [22], however, these bounds cannot be directly used to create an adaptive interpolation rule.
We thus simulate 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 Y and vary the squared of this distance between boxes against nodes in Figure 8.
We use a Gaussian Kernel with γ = 1 √ 2 .
Based on Figure 8, we use the following rule for selecting the number of interpolation nodes for far-field interactions r far ( r ) =
[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 implementing barycentric Lagrange interpolation [7] evaluated at the chebyshev nodes of the second kind.
As this is a well-known technique, we refer to the appendix for more details.
It should be noted that the above methods can straightforwardly be extended to any translationinvariant kernel by recalculating the Taylor expansion for smoothness criteria and rerunning the simulation for adaptive far-field approximation .
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 to obtain a complexity that is dependent on the effective variance limit η (chosen parameter) and the box 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 respectively has 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 divisions multiplied 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 time complexity, hence the results hold.
Theorem 1.
Given a KMVM with edge E (dependent on data X , Y ), lengthscale γ , effective variance limit η , 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) ) , which can 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 not accurately reflect the memory footprint of the actual implementations, whose memory mostly depends on 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
[Equation]
1 2 2 · D − m far
i
[Equation]
and
[Equation]
respectively.
Here M − 1 = 12 2 D and m far= m smooth= m small= m empty= 0 and m far i , m smooth i , m small i , m empty i denotes the number of far-field, smooth field, small field interactions and the 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 the additional ( 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 with points very close to each other would significantly benefit m smooth i more, as the closeness of points would imply more smooth interactions.
If points are sparsely spread out, ( m empty i ) 2 , m small i would provide the most benefit as they remove empty boxes and stops boxes with few points to divide unnecessarily.
All proofs 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 ) .
3 We generate data such that the EV (see section 4.3) varies between 0 . 1 , 1 , 10 for data of sizes n = 10 6 , 10 7 , 10 8 , 10 9 .
The parameters used for F 3 M are η = 0 .
5 and r = 2 D , 3 D , 4 D with a cap at r = 2048 .
The error for the approximated KMVM product ˆ v is calculated as Relative error :
=
∥ ˆ v − v ∥ 2 ∥ v ∥ 2 , where the true KMVM product v is obtained by calculating the full KMVM on a subset X ′ consisting of the first 5000 points in X against the entire dataset in double 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 chosen since the extra graphic memory is necessary to fit the data on one card when n = 10 9 .
It should be noted that n = 10 9 can only be run up to D = 3 , as X and b itself cannot fit in memory for higher dimensions with the GPUs we had available.
For details on how F 3 M scales across multiple GPUs, see Appendix G.
[Table] [Table]
[Table] [Table]
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 distributed 3 F 3 M is kernel agnostic, however we choose the Gaussian kernel for simplicity.
[Table]
data ( D = 1 , 2 , 3 , 4 , 5 , 6 , 7 ) together with data simulated from Brownian motion, fractional Brownian motion, and Clustered data ( D = 1 , 2 , 3 ).
For the k ( X , Y ) -case we consider uniformly distributed x and normal distributed y ( D = 1 , 2 , 3 , 4 , 5 , 6 , 7 ).
See Appendix 10 for visualizations of data.
We have to consider smaller n for the k ( X , Y ) -case when D ≥ 3 , as twice the amount of data needs to be stored.
For D = 3 / (4 , 5) / (6 , 7) we instead consider at most n = 5 · 10 8 / 2 . 5 · 10 8 / 10 8 .
It should be noted that D = 7 is a hard limit for geometric partitioning-based methods, since for D = 8 , we would have 2 8 · 2 · 2 8 · 2 ≈ 4 .
3 · 10 9 interactions after only 2 divisions.
This number of interactions cannot even be represented by a 32-bit integer.
We summarize the runs in Table 3 and plot the error and time complexity in Figure 9 for each dataset when D = 3 .
We find that F 3 M maintains sub-linear empirical complexity up to D = 6 , where we have to set small field limit ρ to a larger number to not run out of memory.
Further, the error increases in the higher dimensions since we use fewer nodes per dimension when interpolating, owing to the sparse grid technique.
We note that D = 7 has faster run times than D = 6 which is explained by that for some values of EV, D = 7 doesn’t run with acceptable errors which skew the run time to datasets where a larger portion of the data can be interpolated.
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 that stabilizes the inverse.
FALKON is a Nyström approximation based solver that requires a subsample
X ′ ∈ R M × d of X to approximate the inverse computation.
We focus the experiments on tall and skinny data and take n = 10 9 , d ≤ 3 with M = 10 5 for all experiments.
We consider uniformly and normally 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 and fare 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 KRR for EV = 0 .
1 , 1 , 10 on synthetic data, and report the average R 2 (AUC for NYC Taxi) and training time in Table 4.
For the real world datasets OSM and NYC Taxi, we fix the lengthscale using the median 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 be attributed 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 is an exact method, the dataset distribution has no effect on computational time.
Here, the smoothness criteria and adaptive far-field technique improve computational time.
We find that F 3 M achieves a speed-up between 2 .
− 33 .
3 × against FFM(GPU) and 8 .
− 8500 × speed-up against KeOps.
[Table]
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 6.
The results are quite similar to the comparison between F 3 M and FFM(GPU).
Here
[Table]
we see that smooth field and adaptive far-field approximation (F 3 M) both improve speed and also memory usage as smooth field helps approximate more interactions.
We can thus infer empirically that the m smooth i term in Theorem 2 has a significant impact on reducing the memory footprint of interactions
[Table] [Table]
Gaussian process regression experiments We further compare F 3 M as a drop-in KMVM operation applied 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 the setup 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 7.
As exact GP using F 3 M demonstrates competitive results even when compared to SVGP [18] and SGPR [30], we hypothesize that many high-dimensional datasets conform to the manifold hypothesis [13], allowing F 3 M to be widely 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 and skinny 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 KMVM operation, as demonstrated with FALKON and Gaussian process regression, achieving significant speedups and competitive performance on both tasks.
As an interpolation based approximation method
[Table]
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 fruitful direction would be to extend ideas in F 3 M to accommodate higher-dimensional data by considering randomized 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.
Acknowledgments
The authors sincerely thank Lood van Niekerk and Jean-François Ton for their helpful comments.
A Synthetic data Note on synthetic datasets
We generated synthetic datasets of different types to measure the ability of F 3 M to deal with dense or sparse data.
Dense datasets were generated as independent samples with either uniform or normal distributions.
Clustered datasets were generated by sampling cluster centers from a normal distribution, and then recursively sampling sub-cluster centers from a normal distribution with reduced standard deviation and centered at each cluster center, until the desired number of points is attained.
Fractional Brownian Motion and Brownian Motion samples were generated as samplings of Fractional Brownian Motion paths with respective Hurst index 0. and 0.5.
Figure 10 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 solving g ′ ( 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 interpolation polynomial can then be expressed as p r ( t ) =
(cid:80)
[Equation]
D Smooth field proof
Proposition.
Consider x , y ∈ X ⊂ R d such that d ( x , y ) :
= ∥ x − y ∥ 2 2 γ
[Equation]
≤ η < 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 Taylor expansion, 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 )
|
=as 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 has time complexity O ( n ) , where n = max( n x , n y ) .
Proof.
Far-field interactions are calculated as
[Equation]
As r X , r Y are independent of n x , n y , the complexity becomes O ( n ) .
Further see [3] for alternative proof.
Proposition.
Given n data points in dimension D , the maximum number of divisions Tree max divisions is 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 variance limit η , 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) ) , which can 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 ·
[Equation]
Then all interactions will be interpolated when tree_depth ≥ log 2 (cid:16) D
[Equation] [Equation]
4 · η (cid:17) , which implies we can take Tree max divisions = log 2 (cid:16)
[Equation]
(cid:17) .
Hence the complexity is O ( n · log
[Equation]
(cid:16)
[Equation]
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
[Equation]
= m
[Equation]
= 0 and m far i , m smooth i , m small i , m empty i denotes the number of far-field, smooth field, small field interactions and the number of empty boxes respectively at depth i > 0 .
Note that for i > 0 , these are dependent on data.
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 =
M
[Equation]
m
[Equation] [Equation]
− m
[Equation]
= 1 .
Clearly at depth 0, there can not be any empty 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 the number of interactions grows by 2 2 · D .
At depth i + 1 , we can further remove ( m empty i +1 ) 2 interactions between empty boxes and further compute m far i +
+ m smooth i + + m small i +1 interactions.
Then M i +
=
[Equation]
m
[Equation]
− m small i +1 .
Thus the base case and induction step holds which completes our proof.
This proof also covers FFM, since FFM can be as a special case for F 3 M without 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 M does 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 are non-uniformly distributed, taking the maximum ensures that we don’t compute near-field interactions on 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 product by considering the k ( X , X ) -case and divide the work onto multiple GPUs by partitioning each subproduct of the KMVM (see for Figure 11 an example when N GPU = 8 ).
We take X to be Uniform and 3 dimensional.
We present results in Figure 12.
[Figure]
P
[Figure]
G P U
[Figure] [Figure]
P U
[Figure] [Figure]
We also use nvprof to analyze the % of peak throughput of the V100 cards F 3 M can utilize.
We run nvprof for 3 dimensional uniform data for n = 10 6 , 10 7 , 10 8 , 10 9 .
We present our results in
Figure 13.
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 Figurewe 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 15.
[Figure]
We provide an illustration on near field computations are carried out for F 3 M in Figure 16.
[Figure]