标题:Out-of-Distribution Detection with Deep Nearest Neighbors

作者:Yiyou Sun, Yifei Ming, Xiaojin Zhu, Yixuan Li (University of Wisconsin - Madison)

原文链接:https://arxiv.org/pdf/2204.06507.pdf

代码链接:https://github.com/deeplearning-wisc/knn-ood

1. 介绍

由于经典的机器学习方法通常假设模型训练和测试的数据是独立同分布的(IID, Independent Identical Distribution),这里训练和测试的数据都可以说是 In Distribution(ID) 。

在实际应用当中,模型部署上线后得到的数据往往不能被完全控制的,会出先一些样本外的数据,也就是说模型接收的数据有可能是 Out-of-Distribution (OOD) 样本,也可以叫异常样本(outlier) 指的是模型能够检测出 OOD 样本,而 OOD 样本是相对于 ID 样本来说的。

如今的深度模型常常会对一个 OOD 样本认为是ID样本中的某一个类,并给出高的置信度,这显然是不合理的。举个通俗的例子,我们利用一个包含“猫”和“狗”两类的数据集训练一个二分类器。

如果测试的时候,出现了“人”的样本,如果利用传统的softmax对输出的logit进行处理,那么我们会认为该样本属于“猫”或“狗”的一类,这显然对模型的部署是有影响的,所以我们希望我们的模型能够不把它判定为“猫”或者“狗”,而是判定它为OOD。如何让模型识别出 OOD 样本对 AI 的发展有很重要的意义,特别是 AI 安全。

最近出现了一系列丰富的 OOD 检测算法,其中基于距离的方法 (distance-based methods) 显示出非常卓越的性能。基于距离的方法(Lee et al., 2018 ; Tack et al., 2020; Sehwag et al., 2021) 利用从模型中提取的特征嵌入(feature embedding),并假设:在测试中OOD的样本相对训练中的ID样本相对来讲更远。

但这些方法都把feature embedding的空间空间分布建模为多变量的混合高斯分布 (mixture of multivariate Gaussian distributions GMM),这显然是作用有限的,因为很显然视觉信息非常复杂,不是用一个很简单的GMM就能描述的,所以可能这种方式叶比较受阻。其实利用参数相关的分布假设来描述真实世界中的各种分布显然是非常有挑战的。马毅老师也在表达过高维空间特征的复杂性。

所以本文提出了一个问题:

Can we leverage the non-parametric nearest neighbor approach for OOD detection?

我们可以利用非参数最近邻方法进行 OOD 检测吗?既然基于参数估计的分布假设可能不work,那我们能不能基于一些很简单的非参数方式来解决OOD问题?来自University of Wisconsin - Madison的团队设计了一种非常简洁的方式,给出了肯定的回答。

为了检测 OOD 样本,我们计算第 k 个嵌入之间的最近邻 (KNN) 距离测试输入和训练集的嵌入和使用基于阈值的标准,用于确定输入是否为 OOD与否。简而言之,提出的方法进行非参数的估计,基于深度k-NN distance将数据分成两组(ID vs. OOD)。

方法具有以下几个优点:

  • 无分布假设:非参数最近邻方法不对底层特征空间施加分布假设。因此,KNN 提供了更强的灵活性和通用性,即使在特征空间不符合高斯混合。-
  • 不依赖OOD数据:测试过程不依赖未知数据的信息。距离阈值仅根据 ID 数据进行估计。
  • 易于使用:近似最近邻搜索的现代实现允许我们在几毫秒内完成此操作,即使数据库包含数十亿张图像(Johnson 等人,2019 年)。相比之下,马氏距离需要计算协方差矩阵的逆矩阵,这在数值上可能是不稳定的。
  • 模型无关:测试过程适用于各种模型架构,包括 CNN 和最近的基于 Transformer 的 ViT 模型。

此外,该工作说明了 KNN 也与训练过程无关,并且与在不同损失函数下训练的模型兼容(cross-entropy loss和contrastive loss)。但可以看到基于对比学习的方式,由于特征聚合的更加紧凑,显然依照k-NN的性能也更优。