Fix shap with vector intercept. (#11764)

Author: trivialfis

src/predictor/cpu_predictor.cc

@@ -257,14 +257,14 @@ struct DataToFeatVec {
 
 template <typename EncAccessor>
 class SparsePageView : public DataToFeatVec<SparsePageView<EncAccessor>> {
-  EncAccessor const &acc_;
+  EncAccessor acc_;
   HostSparsePageView const view_;
 
  public:
   bst_idx_t const base_rowid;
 
-  SparsePageView(HostSparsePageView const p, bst_idx_t base_rowid, EncAccessor const &acc)
-      : acc_{acc}, view_{p}, base_rowid{base_rowid} {}
+  SparsePageView(HostSparsePageView const p, bst_idx_t base_rowid, EncAccessor acc)
+      : acc_{std::move(acc)}, view_{p}, base_rowid{base_rowid} {}
   [[nodiscard]] std::size_t Size() const { return view_.Size(); }
 
   [[nodiscard]] bst_idx_t DoFill(bst_idx_t ridx, float *out) const {
@@ -283,7 +283,7 @@ template <typename EncAccessor>
 class GHistIndexMatrixView : public DataToFeatVec<GHistIndexMatrixView<EncAccessor>> {
  private:
   GHistIndexMatrix const &page_;
-  EncAccessor const &acc_;
+  EncAccessor acc_;
   common::Span<FeatureType const> ft_;
 
   std::vector<std::uint32_t> const &ptrs_;
@@ -295,10 +295,10 @@ class GHistIndexMatrixView : public DataToFeatVec<GHistIndexMatrixView<EncAccess
   bst_idx_t const base_rowid;
 
  public:
-  GHistIndexMatrixView(GHistIndexMatrix const &_page, EncAccessor const &acc,
+  GHistIndexMatrixView(GHistIndexMatrix const &_page, EncAccessor acc,
                        common::Span<FeatureType const> ft)
       : page_{_page},
-        acc_{acc},
+        acc_{std::move(acc)},
         ft_{ft},
         ptrs_{_page.cut.Ptrs()},
         mins_{_page.cut.MinValues()},
@@ -365,11 +365,11 @@ template <typename Adapter, typename EncAccessor>
 class AdapterView : public DataToFeatVec<AdapterView<Adapter, EncAccessor>> {
   Adapter const *adapter_;
   float missing_;
-  EncAccessor const &acc_;
+  EncAccessor acc_;
 
  public:
-  explicit AdapterView(Adapter const *adapter, float missing, EncAccessor const &acc)
-      : adapter_{adapter}, missing_{missing}, acc_{acc} {}
+  explicit AdapterView(Adapter const *adapter, float missing, EncAccessor acc)
+      : adapter_{adapter}, missing_{missing}, acc_{std::move(acc)} {}
 
   [[nodiscard]] bst_idx_t DoFill(bst_idx_t ridx, float *out) const {
     auto const &batch = adapter_->Value();
@@ -408,7 +408,7 @@ struct EncAccessorPolicy {
   [[nodiscard]] auto MakeAccessor(Context const *ctx, enc::HostColumnsView new_enc,
                                   gbm::GBTreeModel const &model) {
     auto [acc, mapping] = MakeCatAccessor(ctx, new_enc, model.Cats());
-    this->mapping_ = std::move(mapping);
+    std::swap(mapping, this->mapping_);
     return acc;
   }
 };
@@ -923,7 +923,7 @@ class CPUPredictor : public Predictor {
     CHECK_NE(ncolumns, 0);
     auto device = ctx_->Device().IsSycl() ? DeviceOrd::CPU() : ctx_->Device();
     auto base_margin = info.base_margin_.View(device);
-    auto base_score = model.learner_model_param->BaseScore(device)(0);
+    auto base_score = model.learner_model_param->BaseScore(device);
 
     // parallel over local batch
     common::ParallelFor(batch.Size(), this->ctx_->Threads(), [&](auto i) {
@@ -962,7 +962,7 @@ class CPUPredictor : public Predictor {
           CHECK_EQ(base_margin.Shape(1), ngroup);
           p_contribs[ncolumns - 1] += base_margin(row_idx, gid);
         } else {
-          p_contribs[ncolumns - 1] += base_score;
+          p_contribs[ncolumns - 1] += base_score(gid);
         }
       }
     });

src/predictor/gpu_predictor.cu

@@ -1129,6 +1129,7 @@ class GPUPredictor : public xgboost::Predictor {
     // allocate space for (number of features + bias) times the number of rows
     size_t contributions_columns = model.learner_model_param->num_feature + 1;  // +1 for bias
     auto dim_size = contributions_columns * model.learner_model_param->num_output_group;
+    // Output shape: [n_samples, n_classes, n_features + 1]
     out_contribs->Resize(p_fmat->Info().num_row_ * dim_size);
     out_contribs->Fill(0.0f);
     auto phis = out_contribs->DeviceSpan();
@@ -1159,11 +1160,11 @@ class GPUPredictor : public xgboost::Predictor {
     const auto margin = p_fmat->Info().base_margin_.Data()->ConstDeviceSpan();
 
     auto base_score = model.learner_model_param->BaseScore(ctx_);
-    dh::LaunchN(p_fmat->Info().num_row_ * model.learner_model_param->num_output_group,
-                ctx_->CUDACtx()->Stream(), [=] __device__(size_t idx) {
-                  phis[(idx + 1) * contributions_columns - 1] +=
-                      margin.empty() ? base_score(0) : margin[idx];
-                });
+    bst_idx_t n_samples = p_fmat->Info().num_row_;
+    dh::LaunchN(n_samples * ngroup, ctx_->CUDACtx()->Stream(), [=] __device__(std::size_t idx) {
+      auto [_, gid] = linalg::UnravelIndex(idx, n_samples, ngroup);
+      phis[(idx + 1) * contributions_columns - 1] += margin.empty() ? base_score(gid) : margin[idx];
+    });
   }
 
   void PredictInteractionContributions(DMatrix* p_fmat, HostDeviceVector<float>* out_contribs,
@@ -1219,14 +1220,13 @@ class GPUPredictor : public xgboost::Predictor {
 
     auto base_score = model.learner_model_param->BaseScore(ctx_);
     size_t n_features = model.learner_model_param->num_feature;
-    dh::LaunchN(p_fmat->Info().num_row_ * model.learner_model_param->num_output_group,
-                ctx_->CUDACtx()->Stream(), [=] __device__(size_t idx) {
-                  size_t group = idx % ngroup;
-                  size_t row_idx = idx / ngroup;
-                  phis[gpu_treeshap::IndexPhiInteractions(row_idx, ngroup, group, n_features,
-                                                          n_features, n_features)] +=
-                      margin.empty() ? base_score(0) : margin[idx];
-                });
+    bst_idx_t n_samples = p_fmat->Info().num_row_;
+    dh::LaunchN(n_samples * ngroup, ctx_->CUDACtx()->Stream(), [=] __device__(size_t idx) {
+      auto [ridx, gidx] = linalg::UnravelIndex(idx, n_samples, ngroup);
+      phis[gpu_treeshap::IndexPhiInteractions(ridx, ngroup, gidx, n_features, n_features,
+                                              n_features)] +=
+          margin.empty() ? base_score(gidx) : margin[idx];
+    });
   }
 
   void PredictLeaf(DMatrix* p_fmat, HostDeviceVector<float>* predictions,

tests/cpp/predictor/test_predictor.cc

@@ -860,7 +860,7 @@ void ShapExternalMemoryTest::Run(Context const *ctx, bool is_qdm, bool is_intera
                                  .Classes(n_classes));
   std::unique_ptr<Learner> learner{Learner::Create({p_fmat})};
   learner->SetParam("device", ctx->DeviceName());
-  learner->SetParam("base_score", "0.5");
+  learner->SetParam("base_score", "[0.5, 0.5, 0.5]");
   learner->SetParam("num_parallel_tree", "3");
   learner->SetParam("max_bin", std::to_string(max_bin));
   for (std::int32_t i = 0; i < 4; ++i) {
@@ -869,8 +869,10 @@ void ShapExternalMemoryTest::Run(Context const *ctx, bool is_qdm, bool is_intera
   Json model{Object{}};
   learner->SaveModel(&model);
   auto j_booster = model["learner"]["gradient_booster"]["model"];
-  auto model_param = MakeMP(n_features, 0.0, n_classes, ctx->Device());
 
+  auto base_score = linalg::Tensor<float, 1>{{0.0, 0.0, 0.0}, {3}, ctx->Device()};
+  LearnerModelParam model_param(n_features, std::move(base_score), n_classes, 1,
+                                MultiStrategy::kOneOutputPerTree);
   gbm::GBTreeModel gbtree{&model_param, ctx};
   gbtree.LoadModel(j_booster);
 

tests/python-gpu/test_gpu_prediction.py

@@ -628,3 +628,44 @@ def test_dtypes(self):
 
 def test_base_margin_vs_base_score() -> None:
     run_base_margin_vs_base_score("cuda")
+
+
+@pytest.mark.skipif(**tm.no_sklearn())
+def test_shap_multiclass() -> None:
+    from sklearn.datasets import make_classification
+
+    X, y = make_classification(n_classes=3, random_state=2025, n_informative=16)
+    param = {
+        "tree_method": "hist",
+        "device": "cuda",
+        "num_class": 3,
+        "base_score": [1.0, 2.0, 3.0],
+    }
+    Xy = xgb.DMatrix(X, y)
+    bst = xgb.train(param, Xy, 8)
+
+    d_shap = bst.predict(Xy, pred_contribs=True)
+    d_margin = bst.predict(Xy, output_margin=True)
+
+    bst.set_param({"device": "cpu"})
+
+    h_shap = bst.predict(Xy, pred_contribs=True)
+    h_margin = bst.predict(Xy, output_margin=True)
+
+    np.testing.assert_allclose(d_shap, h_shap, atol=1e-6)
+    np.testing.assert_allclose(d_margin, h_margin, atol=1e-6)
+
+    # Compare base margin and base score
+    margin = np.stack(
+        [
+            np.ones(X.shape[0]),
+            np.full(X.shape[0], fill_value=2.0),
+            np.full(X.shape[0], fill_value=3.0),
+        ],
+        axis=1,
+    )
+    Xy = xgb.DMatrix(X, y, base_margin=margin)
+
+    bst.set_param({"device": "cuda"})
+    d_shap = bst.predict(Xy, pred_contribs=True)
+    np.testing.assert_allclose(d_shap, h_shap, atol=1e-6)