蒙亦有道：用共形预测为机器学习赋予置信度

写在前面

这是一篇科普向 Notebook，其文字材料大量参考了文章 A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification 及其作者 Stephen Bates 在 MIT 2023 年秋季学期讲授的 6.7900 Machine Learning 课程。

本文所需知识背景：

• 机器学习基础
• 概率论与数理统计

从二分类任务的概率诠释谈起

从线性 SVM 到最为复杂的神经网络，无论架构如何，很多二分类模型都会将最后一层的输出 $z\in \mathbb{R}$ 通过 Logistic 函数得到最终结果 $\hat{y}=\frac{1}{1+{\text{e}}^z}\in [0,1]$ ，并在 $\hat{y}>0.5$ 时分类为 1，在 $\hat{y}<0.5$ 时分类为 0。直觉告诉我们，同样是分类为 1，$\hat{y}=0.999$ 和 $\hat{y}=0.501$ 的“含金量”是不同的：如果这一数据点的真实分类 $y=1$，那么二分类任务常用的交叉熵损失函数 $-L=y\log \hat{y}+(1-y)\log (1-\hat{y})$ 的值也在 $\hat{y}\to 1$ 时更小。似乎前者表明模型更加确信这一分类结果，而后者则是一个字面意思上“模棱两可”的输出。因此，我们可以将 $\hat{y}$ 解读为模型预测这一数据点分类为 1 的概率。

$\hat{y}$ 确实常常作为主动学习（active learning）中不确定性采样（uncertainty sampling）标准：在训练集中初训模型后，我们要从外部选择一批新的数据点进行标注并加入训练集继续训练。此时模型在外部数据上预测的 $\hat{y}$ 是否接近 0.5 可以作为其是否加入训练集的准绳——我们要挑选那些模型预测具有最小置信度（Least Confidence）的点来加强模型的认识。

在多分类任务中，Logistic 函数扩展为 Softmax 函数。如果一共有 $d$ 个类别，则最后一层输出 $\mathbf{z}\in {\mathbb{R}}^d$ 对应的最终结果为 ${\hat{y}}_i=\frac{{\text{e}}^{z_i}}{{\sum }_i{\text{e}}^{z_i}}$，并将其分类为最大的 ${\hat{y}}_i$ 对应的类别 $i$。由于 ${\sum }_i{\hat{y}}_i=1$ 的良好性质，我们依然可以将 ${\hat{y}}_i$ 解读为模型所预测的数据点属于类别 $i$ 的概率。在主动学习中，模型在外部数据上预测的最大和第二大的 ${\hat{y}}_i$ 之差是否接近 0 也可以作为是否加入训练集的准绳——我们要挑选那些模型预测区分度低的点来加强模型的认识。这称之为边缘（Margin）采样，可见最小置信度采样就是边缘采样在二分类问题中的特殊情况。

这些解读和应用看起来直观而美好，在二分类任务的理想情况下，我们希望对于任何数据点的预测结果 $\hat{y}$ 和真实标签 $y$ 满足 $P(y=1|\hat{y}=\alpha )=\alpha$，但事实上仅通过常规的模型训练，以上概率恒等式没有任何理论保证。为了衡量我们预想的 $\hat{y}$ 的概率意义和实际的概率的差距，我们可以将区间 $[0,1]$ 分割为互不相交的区间 $I_1,\cdots ,I_K$，并在 ${\hat{y}}^i\in I_k$ 时将数据点 $i$ 投入集合 $B_k$ 中，然后计算校准误差（calibration error） $$\sum _{k=1}^K\frac{|B_k|}{n}\left|\frac{{\sum }_{i\in B_k}{1}_{y^i=1}}{|B_k|}-\frac{{\sum }_{i\in B_k}{\hat{y}}^i}{|B_k|}\right|$$ 如果数据点足够多，其中的每一项都反映了 $|B_k|$ 中从数据点估计的经验概率 $P(y=1|\hat{y})$ 和从模型中解读的预测概率 $\hat{y}$ 之间的平均差异。

借助这一指标，我们可以尝试通过调整模型输出来解决这一问题：我们将数据集的一部分划分为校准集（calibration set）。模型训练完成后，在推理所用的 Logistic 函数中加入超参数 $\beta$，使其变为 $$\hat{y}=\frac{1}{1+{\text{e}}^{\beta z}}$$ $\beta$ 称为温度参数，较高的温度参数将使得 $\hat{y}$ 更接近 0 或 1，使模型更“自信”。于是，我们可以调节 $\beta$ 使校准集上的校准误差尽可能地小，从而使模型预测的 $\hat{y}$ 更接近真实的概率意义。

借助可靠的概率输出，我们不仅能更有效地开展主动学习中的不确定度采样，还能够给出预测结果的置信度；当单一结果的置信度不高时，高置信度结果的集合也能为决策提供参考。然而，即使我们有充分大的校准集，也不能保证找到合适的 $\beta$ 使校准误差小于我们所期待的值。换言之，即使通过校准和温度参数的手段，逻辑回归或 softmax 输出的依然是“伪概率”。

我们可以使用最简单的逻辑回归算法和经典的乳腺癌二分类数据集作为示例：首先将数据集分割为训练集和校准集

定义带正则化的逻辑回归的损失函数，其中 $\mathbf{w}$ 是模型权重

训练模型

以下函数把 $[0,1]$ 分为 num 等份来计算校准误差，并引入温度参数 $\beta$

作出不同温度参数下，模型预测的概率和校准集上估计的经验概率的折线图及其校准误差

可见适当调节 $\beta$ 可以使得预测概率和经验概率相接近，但很难实现完美吻合。

通过共形预测得到给定置信度的结果范围

共形预测（conformal prediction）能够将上述这种并不严格但具有启发性的不确定度估计转化为统计上严格的不确定度估计，并适用于任何已经训练好的机器学习模型。在 AI4S 领域，它已经应用于毒性预测 (Zhang et. al., 2021) 和生物分子设计 (Fannjiang et. al., 2022) 等任务上。其基本步骤如下：

1. 将数据分为训练集、校准集和测试集。在训练集上训练好一个模型，它根据数据特征 $x$ 给出标签 $y$ 的预测值 $\hat{y}$。
2. 根据所选择的模型的意义，定义一个打分函数 $s(x,y)$：这个函数的值越大，表明模型对特征 $x$ 的预测结果 $\hat{y}$ 和真实标签 $y$ 的一致性越差。
3. 对于给定的置信度 $1-\alpha$，在整个校准集上计算 $s(x_1,y_1),\cdots ,s(x_n,y_n)$ 的 $\frac{\lceil (n+1)(1-\alpha )\rceil }{n}$ 分位数 $\hat{q}$。
4. 对于测试集中的新数据点 $(x_{\text{test}},y_{\text{text}})$，共形预测结果是一个集合 $\mathcal{C}(x_{\text{test}})=\{y:s(x_{\text{test}},y)\le \hat{q}\}$。

这个集合就是共形预测给出的结果范围，能够证明：

如果校准集和测试集中的所有数据 $(x,y)$ 独立同分布，那么只要按照以上步骤 3 和 4 计算 $\hat{q}$ 并输出 $\mathcal{C}(x_{\text{test}})$，就有 $$\mathbb{P}(y_{\text{text}}\in \mathcal{C}(x_{\text{test}}))\ge 1-\alpha$$ 且在相当宽松的条件下，还有 $$\mathbb{P}(y_{\text{text}}\in \mathcal{C}(x_{\text{test}}))\le 1-\alpha +\frac{1}{n+1}$$

这说明真实标签以近似为 $1-\alpha$ 的置信度落在集合 $\mathcal{C}(x_{\text{test}})$ 中。

我们会发现上述数学结论几乎完全不依赖于 $s(x,y)$ 的定义，千万不要被这样美好的结果所迷惑！事实上，共形预测方法保证了输出的结果范围具有“置信区间”的统计意义，但这个集合对决策是否有帮助仍要依赖于我们对 $s(x,y)$ 的定义——$s(x,y)$ 包含了我们对模型和任务的理解，是所得的结果范围提供有效信息的保证。我们来看一个没有任何信息量的共形预测的极端例子：定义 $s(x,y)$ 为一个随机数：

将之前的校准集重新分割为校准集和测试集

在不同的置信度下计算共形预测结果

将测试集上估计的经验概率 $P(y_{\text{text}}\in \mathcal{C}(x_{\text{test}}))$ 对置信度 $1-\alpha$ 作图

可见吻合良好。但实际上这一共形预测的输出结果只是按照给定的置信度随机猜测集合 $\{\},\{0\},\{1\},\{0,1\}$，与模型输出结果没有任何关系。我们可以统计一下 90% 置信度条件下输出结果中 $\{0,1\}$ 的比例

如此之高——这样的结果没有任何价值，因为只要输出足够多的 $\{0,1\}$ 就能保证 $y$ 大概率落入 $C(x)$。它确保了置信度，却未能对决策提供任何帮助。因此，我们仍需设计好的打分函数 $s(x,y)$，在借助共形预测的统计严格性的同时，给出有意义的结果。

设计有意义的打分函数

为了利用我们训练的逻辑回归模型提供的信息，可以这样设计打分函数：设模型对于输入特征 $x$ 的预测为 $\hat{y}(x)$，则 $s(x,y)=|y-\hat{y}(x)|$。这一打分函数非常直观地反映了 $\hat{y}$ 和 $y$ 的偏差，我们来做以下实验：

可见测试集上经验概率和置信度的吻合依然良好。接下来观察高置信度 90% 时共形预测结果有多少包含 $\{0,1\}$。

远小于前面的随机打分函数！这说明共形预测确实从模型输出中受益：它能够更加精准地推断高置信度集合，而不是通过大量猜测全集 $\{0,1\}$ 来提高对真实标签的覆盖率。

后记

关于共形预测，更多的代码案例可以参考文章 A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification 的配套 Github 仓库中的示例 Notebook。