如何基于mlr3框架使用treeshap计算SHAP值？

问题场景 #

使用mlr3框架训练xgboost回归模型后，调用treeshap的unify()函数时触发错误：

unify.default(mm$learners[[1]], data):
Provided model is not of type supported by treeshap.

错误原因 #

mlr3的LearnerRegrXgboost是模型封装容器，mm$learners[[1]]返回的是mlr3封装后的学习器对象，而非treeshap支持的原始xgb.Booster模型。从模型信息中可以看到* Model: xgb.Booster，说明真实模型存储在学习器的$model属性中。

修正方案及完整代码 #

提取mlr3学习器中的原始模型，传入treeshap处理即可，同时确保输入数据仅包含特征列（不含目标变量）：

library(shapviz)
library(kernelshap)
library(mlr3)
library(mlr3verse)
library(mlr3learners)
library(treeshap)
library(xgboost)

# 数据预处理
data <- fifa20$data[colnames(fifa20$data) != 'work_rate']
target <- fifa20$target

# 创建mlr3回归任务
tsk <- mlr3::TaskRegr$new(id = "dd", backend = cbind(data, target = target), target = "target")

# 初始化xgboost学习器
learner <- lrn("regr.xgboost")

# 执行交叉验证并保存模型
mm <- resample(
  task = tsk,
  learner = learner,
  resampling = rsmp("cv", folds = 3),
  store_backends = TRUE,
  store_models = TRUE
)

# 关键：提取mlr3学习器中的原始xgb.Booster模型
xgb_raw_model <- mm$learners[[1]]$model

# 准备仅含特征的数据集（treeshap不需要目标列）
feature_data <- data

# 统一模型格式
unified_model <- unify(xgb_raw_model, feature_data)

# 计算SHAP值
shap_result <- treeshap(unified_model, feature_data)

# 可视化示例
sv <- shapviz(shap_result, X = feature_data)
sv_importance(sv)  # 绘制SHAP特征重要性图
sv_dependence(sv, feature = "overall")  # 绘制单个特征的SHAP依赖关系图

核心注意事项 #

• 所有treeshap支持的树模型（xgboost、lightgbm、randomForest、ranger等），在mlr3中都需要通过$model提取原始模型对象
• 传入treeshap的数据集必须仅包含特征列，不能包含目标变量或其他无关列
• 若使用其他树模型学习器（如regr.ranger），提取原始模型的方式一致：learner$model

内容的提问来源于stack exchange，提问作者gen linlin