[REVIEW] Experimental versions of GPU accelerated Kernel and Permutation SHAP

CHANGELOG.md

@@ -4,6 +4,7 @@
 - PR #3164: Expose silhouette score in Python
 - PR #2659: Add initial max inner product sparse knn
 - PR #2836: Refactor UMAP to accept sparse inputs
+- PR #3126: Experimental versions of GPU accelerated Kernel and Permutation SHAP
 
 ## Improvements
 - PR #3077: Improve runtime for test_kmeans

cpp/CMakeLists.txt

@@ -395,6 +395,8 @@ if(BUILD_CUML_CPP_LIBRARY)
     src/datasets/make_regression.cu
     src/dbscan/dbscan.cu
     src/decisiontree/decisiontree.cu
+    src/explainer/kernel_shap.cu
+    src/explainer/permutation_shap.cu
     src/fil/fil.cu
     src/fil/infer.cu
     src/glm/glm.cu

cpp/include/cuml/explainer/kernel_shap.hpp

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuml/cuml.hpp>
+
+namespace ML {
+namespace Explainer {
+
+/**
+ * Generates samples of dataset for kernel shap algorithm.
+ *
+ *
+ * @param[in]    handle             cuML handle
+ * @param[inout] X                  generated data [on device] 1-0 (row major)
+ * @param[in]    nrows_X            number of rows in X
+ * @param[in]    ncols              number of columns in X, background and dataset
+ * @param[in]    background         background data [on device]
+ * @param[in]    nrows_background   number of rows in background dataset
+ * @param[out]   dataset            generated data [on device] observation=background (row major)
+ * @param[in]    observation        row to scatter
+ * @param[in]    nsamples           vector with number of entries that are randomly sampled
+ * @param[in]    len_nsamples       number of entries to be sampled
+ * @param[in]    maxsample          size of the biggest sampled observation
+ * @param[in]    seed               Seed for the random number generator
+ *
+ * Kernel distrubutes exact part of the kernel shap dataset
+ * Each block scatters the data of a row of `observations` into the (number of rows of
+ * background) in `dataset`, based on the row of `X`.
+ * So, given:
+ * background = [[0, 1, 2],
+                 [3, 4, 5]]
+ * observation = [100, 101, 102]
+ * X = [[1, 0, 1],
+ *      [0, 1, 1]]
+ *
+ * dataset (output):
+ * [[100, 1, 102],
+ *  [100, 4, 102]
+ *  [0, 101, 102],
+ *  [3, 101, 102]]
+ * The first thread of each block calculates the sampling of `k` entries of `observation`
+ * to scatter into `dataset`. Afterwards each block scatters the data of a row of `X` into
+ * the (number of rows of background) in `dataset`.
+ * So, given:
+ * background = [[0, 1, 2, 3],
+ *               [5, 6, 7, 8]]
+ * observation = [100, 101, 102, 103]
+ * nsamples = [3, 2]
+ *
+ * X (output)
+ *      [[1, 0, 1, 1],
+ *       [0, 1, 1, 0]]
+ *
+ * dataset (output):
+ * [[100, 1, 102, 103],
+ *  [100, 6, 102, 103]
+ *  [0, 101, 102, 3],
+ *  [5, 101, 102, 8]]
+ */
+void kernel_dataset(const raft::handle_t& handle, float* X, int nrows_X,
+                    int ncols, float* background, int nrows_background,
+                    float* dataset, float* observation, int* nsamples,
+                    int len_nsamples, int maxsample, uint64_t seed = 0ULL);
+
+void kernel_dataset(const raft::handle_t& handle, float* X, int nrows_X,
+                    int ncols, double* background, int nrows_background,
+                    double* dataset, double* observation, int* nsamples,
+                    int len_nsamples, int maxsample, uint64_t seed = 0ULL);
+
+}  // namespace Explainer
+}  // namespace ML

cpp/include/cuml/explainer/permutation_shap.hpp

@@ -0,0 +1,140 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuml/cuml.hpp>
+
+namespace ML {
+namespace Explainer {
+
+/**
+ * Generates a dataset by tiling the `background` matrix into `out`, while
+ *  adding a forward and backward permutation pass of the observation `row`
+ * on the positions defined by `idx`. Example:
+ *
+ * background = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+ * idx = [2, 0, 1]
+ * row = [100, 101, 102]
+ * output:
+ * [[  0,   1,   2]
+ *  [  3,   4,   5]
+ *  [  6,   7,   8]
+ *  [  0,   1, 102]
+ *  [  3,   4, 102]
+ *  [  6,   7, 102]
+ *  [100,   1, 102]
+ *  [100,   4, 102]
+ *  [100,   7, 102]
+ *  [100, 101, 102]
+ *  [100, 101, 102]
+ *  [100, 101, 102]
+ *  [100, 101,   2]
+ *  [100, 101,   5]
+ *  [100, 101,   8]
+ *  [  0, 101,   2]
+ *  [  3, 101,   5]
+ *  [  6, 101,   8]
+ *  [  0,   1,   2]
+ *  [  3,   4,   5]
+ *  [  6,   7,   8]]
+ *
+ *
+ * @param[in]  handle          cuML handle
+ * @param[out] out             generated data [on device] [dim = (2 * ncols * nrows_bg + nrows_bg) * ncols]
+ * @param[in] background       background data [on device] [dim = ncols * nrows_bg]
+ * @param[in] nrows_bg         number of rows in background dataset
+ * @param[in] ncols            number of columns
+ * @param[in] row              row to scatter in a permutated fashion [dim = ncols]
+ * @param[in] idx              permutation indexes [dim = ncols]
+ * @param[in] row_major        boolean to generate either row or column major data
+ *
+ */
+void permutation_shap_dataset(const raft::handle_t& handle, float* out,
+                              const float* background, int nrows_bg, int ncols,
+                              const float* row, int* idx, bool row_major);
+
+void permutation_shap_dataset(const raft::handle_t& handle, double* out,
+                              const double* background, int nrows_bg, int ncols,
+                              const double* row, int* idx, bool row_major);
+
+/**
+ * Generates a dataset by tiling the `background` matrix into `out`, while
+ *  adding a forward and backward permutation pass of the observation `row`
+ * on the positions defined by `idx`. Example:
+ *
+ * background = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+ * idx = [2, 0, 1]
+ * row = [100, 101, 102]
+ * output:
+ * [[  0,   1,   2]
+ *  [  3,   4,   5]
+ *  [  6,   7,   8]
+ *  [  0,   1, 102]
+ *  [  3,   4, 102]
+ *  [  6,   7, 102]
+ *  [100,   1,   2]
+ *  [100,   4,   5]
+ *  [100,   7,   8]
+ *  [  0, 101,   2]
+ *  [  3, 101,   5]
+ *  [  6, 101,   8]]
+ *
+ *
+ * @param[in]  handle          cuML handle
+ * @param[out] out             generated data [on device] [dim = (2 * ncols * nrows_bg + nrows_bg) * ncols]
+ * @param[in] background       background data [on device] [dim = ncols * nrows_bg]
+ * @param[in] nrows_bg           number of rows in background dataset
+ * @param[in] ncols           number of columns
+ * @param[in] row              row to scatter in a permutated fashion [dim = ncols]
+ * @param[in] idx              permutation indexes [dim = ncols]
+ * @param[in] row_major        boolean to generate either row or column major data
+ *
+ */
+
+void shap_main_effect_dataset(const raft::handle_t& handle, float* out,
+                              const float* background, int nrows_bg, int ncols,
+                              const float* row, int* idx, bool row_major);
+
+void shap_main_effect_dataset(const raft::handle_t& handle, double* out,
+                              const double* background, int nrows_bg, int ncols,
+                              const double* row, int* idx, bool row_major);
+
+/**
+ * Function that aggregates averages of the averatge of results of the model
+ * called with the permutation dataset, to estimate the SHAP values.
+ * It is equivalent to the Python code:
+ *  for i,ind in enumerate(idx):
+ *     shap_values[ind] += y_hat[i + 1] - y_hat[i]
+ *  for i,ind in enumerate(idx):
+ *     shap_values[ind] += y_hat[i + ncols] - y_hat[i + ncols + 1]
+ *
+ * @param[in]  handle          cuML handle
+ * @param[out] shap_values     Array where the results are aggregated [dim = ncols]
+ * @param[in] y_hat            Results to use for the aggregation [dim = ncols + 1]
+ * @param[in] ncols            number of columns
+ * @param[in] idx              permutation indexes [dim = ncols]
+ */
+void update_perm_shap_values(const raft::handle_t& handle, float* shap_values,
+                             const float* y_hat, const int ncols,
+                             const int* idx);
+
+void update_perm_shap_values(const raft::handle_t& handle, double* shap_values,
+                             const double* y_hat, const int ncols,
+                             const int* idx);
+
+}  // namespace Explainer
+}  // namespace ML