[论文解读] Kernel Operations on the GPU, with Autodiff, without Memory Overflows

Benjamin Charlier, Jean Feydy+3 more|arXiv (Cornell University)|Mar 27, 2020

Parallel Computing and Optimization Techniques参考文献 8被引用 80

一句话总结

KeOps 提供一个快速、节省内存的 GPU 框架，用于核距离计算和自动微分，使用半符号化 LazyTensor 抽象来避免大内核缓冲区，并在大规模问题上超越传统 GPU 后端。

ABSTRACT

The KeOps library provides a fast and memory-efficient GPU support for tensors whose entries are given by a mathematical formula, such as kernel and distance matrices. KeOps alleviates the major bottleneck of tensor-centric libraries for kernel and geometric applications: memory consumption. It also supports automatic differentiation and outperforms standard GPU baselines, including PyTorch CUDA tensors or the Halide and TVM libraries. KeOps combines optimized C++/CUDA schemes with binders for high-level languages: Python (Numpy and PyTorch), Matlab and GNU R. As a result, high-level "quadratic" codes can now scale up to large data sets with millions of samples processed in seconds. KeOps brings graphics-like performances for kernel methods and is freely available on standard repositories (PyPi, CRAN). To showcase its versatility, we provide tutorials in a wide range of settings online at \url{www.kernel-operations.io}.

研究动机与目标

• 动机：在GPU上对大规模数据进行内存高效的核计算的需求。
• 介绍一种半符号化的规约框架，避免将完整的核矩阵物化。
• 描述 LazyTensor 抽象及其与 Python、NumPy、PyTorch 的集成方式。
• 介绍自动微分支持以及用于性能与可重复性的编译 C++ 后端。
• 评估与标准GPU库的性能对比，并讨论预期用途与局限性。

提出的方法

• 定义在符号矩阵条目 F(p, x_i, y_j) 上运行的通用规约框架。
• 解释构建符号计算并通过规约触发求值的 LazyTensor 包装。
• 描述内部 C++ 引擎为每个规约/公式编译优化的 CUDA 代码并存储二进制以便重复使用。
• 解释与 PyTorch autograd 的自动微分整合，可实现任意阶梯的梯度。
• 将 KeOps 的性能与 PyTorch、TensorFlow、Halide、TVM、JAX 以及领域特定库在核矩阵向量乘积上的表现进行比较。
• 突出预期的用例以及包括 Nyström 和 FFM 近似在内的未来改进。

实验结果

研究问题

• RQ1 如何在 GPU/CPU 上实现线性内存占用的核/距离矩阵计算？
• RQ2 半符号化的高层接口是否能同时提供易用性与高性能的适用于大规模核运算？
• RQ3 自动微分如何与专用的核运算引擎整合？
• RQ4 在不同数据规模下，KeOps 相对于主流 GPU 库的实际性能优势有哪些？
• RQ5 在该框架中近似和稀疏的局限性及未来方向有哪些？

主要发现

• KeOps 在数据数组上评估符号公式 F 时实现线性内存占用，避免存储完整的核矩阵。
• LazyTensor 抽象让用户编写类似 NumPy/PyTorch 的代码，惰性构建符号计算，只有在最终规约时才触发。
• 在高斯核向量乘积随数据规模增加的情况下，KeOps 的性能优于 PyTorch、Halide、TVM 和其他基线，并在显著更小的内存占用下提供具竞争力或更优的性能。
• 基准表显示 KeOps 在特定配置下达到诸如 0.4 ms（N=10k）和 1.38 s（N=1M）等结果，部分竞争对手在大 N 时出现内存错误。
• KeOps 支持与 torch.autograd 集成的自动微分，能够对 KeOps 计算进行反向传播。
• 本文认为 KeOps 最适用于维度在 100 以内、点数在 10k 以上的情况，并讨论未来在 Nyström/FFM 近似和 CUDA Tensor Core 支持方面的工作。