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

AI 总结

• SD-KDE虽较经典KDE具更优渐近收敛率，但因经验score计算导致实际速度慢。
• 作者将SD-KDE计算重排为更适合矩阵乘法的形式，从而可充分利用GPU Tensor Cores。
• 在3.2万样本、16维任务上，相比强基线GPU版SD-KDE最高提速47倍。
• 相同设置下，相比scikit-learn KDE最高提速约3300倍。
• 在100万样本、16维且13.1万查询的大规模任务中，单GPU仅需2.3秒完成。

内容摘录

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.