Flash-SD-KDE: Accelerating SD-KDE with Tensor Cores

TL;DR

This paper introduces Flash-SD-KDE, a GPU-accelerated method for score-debiased kernel density estimation that significantly improves speed, enabling practical large-scale applications.

Contribution

The paper re-structures SD-KDE computations to leverage Tensor Cores, achieving substantial speedups over existing implementations and making large-scale SD-KDE feasible.

Findings

Up to 47x faster than baseline SD-KDE GPU implementation

3,300x faster than scikit-learn's KDE

Completes large-scale estimation in 2.3 seconds on a single GPU

Abstract

Score-debiased kernel density estimation (SD-KDE) achieves improved asymptotic convergence rates over classical KDE, but its use of an empirical score has made it significantly slower in practice. We show that by re-ordering the SD-KDE computation to expose matrix-multiplication structure, Tensor Cores can be used to accelerate the GPU implementation. On a 32k-sample 16-dimensional problem, our approach runs up to $47\times$ faster than a strong SD-KDE GPU baseline and $3{,}300\times$ faster than scikit-learn's KDE. On a larger 1M-sample 16-dimensional task evaluated on 131k queries, Flash-SD-KDE completes in $2.3$ s on a single GPU, making score-debiased density estimation practical at previously infeasible scales.