[論文レビュー] 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ラッパーを説明する。
• 各リダクション/式ごとに最適化されたCUDAコードをコンパイルし、再利用のためにバイナリを格納する内部C++エンジンを説明する。
• 任意次数の勾配に対するPyTorch autogradとの自動微分統合を説明する。
• カーネル行列-ベクトル積を用いてKeOpsの性能をPyTorch、TensorFlow、Halide、TVM、JAX、およびドメイン固有ライブラリと比較する。
• Nyström法およびFFM近似を含む将来の強化を含む、想定される使用例と今後の拡張を強調する。

実験結果

リサーチクエスチョン

• RQ1 GPU/CPU上で線形メモリフットプリントでカーネル/距離行列計算を実行するにはどうすればよいか？
• RQ2 半象徴的で高レベルのインターフェースは、大規模なカーネル計算に対して使いやすさと高性能の両立を提供できるか？
• RQ3 自動微分は特殊なカーネル演算エンジンとどのように統合されるか？
• RQ4 データサイズの異なる範囲で、KeOpsと主流のGPUライブラリとの実用的な性能利得はどの程度か？
• RQ5 このフレームワークにおける近似とスパース性の制限と今後の方向性は何か？

主な発見

PyTorch	PyTorch-TPU	TF-XLA	Halide	TVM	PyKeOps	KeOps++
9 ms	10 ms	13 ms	1.0 ms	3.80 ms	0.7 ms	0.4 ms
out of mem	out of mem	89 ms	34.1 ms	36.8 ms	15.0 ms	14.6 ms
out of mem	out of mem	3.8 s	2.79 s	1.39 s	1.38 s

• 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サポートに関する将来の作業を議論します。