Pyspark中Naive Bayes分类器threshold阈值参数应如何设置？

核心参数逻辑 #

Spark ML 库中所有分类器（含朴素贝叶斯）的thresholds参数传入的是每个类别的预测权重阈值，而非直接的类别概率阈值，内部预测规则为：

1. 模型输出每个类别的预测概率
2. 对每个类别i，计算值为 概率[i] / thresholds[i]
3. 最终预测结果为计算值最大的对应类别

二分类场景自定义阈值配置 #

如果你需要实现「样本属于类别1的概率超过0.7才预测为1」的需求，推导过程如下：

• 目标条件：P(类别1) > 0.7
• 由概率和为1可得：P(类别1) > 0.7 * (P(类别0) + P(类别1))
• 移项整理：0.3 * P(类别1) > 0.7 * P(类别0)
• 变形匹配Spark预测规则：P(类别1)/0.7 > P(类别0)/0.3

因此你需要传入的thresholds值为[0.3, 0.7]，注意列表顺序要和模型输出的概率向量的类别顺序完全匹配，顺序错配会直接导致阈值逻辑完全相反，也是你之前ROC曲线结果不符合预期的常见原因。

官方示例参数逻辑说明 #

官方文档示例nb = nb.setThresholds([0.01, 10.00])的作用是极度偏向预测类别0：

• 类别0的阈值仅为0.01，对应计算值会被放大100倍
• 类别1的阈值高达10，对应计算值会被缩小到1/10
  除非类别1的概率是类别0的1000倍以上，否则都会预测为类别0，适合类别极度不平衡、需要严格抬高少数类预测门槛的场景。

内容的提问来源于stack exchange，提问作者Juan Cruz Carrau