具有 SNGP 的不确定性感知深度学习

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

在安全至关重要的 AI 应用程序（例如医疗决策和自动驾驶）中，或者在数据本身存在噪声的情况下（例如自然语言理解），深度分类器必须可靠地量化其不确定性。深度分类器应该能够意识到自己的局限性，以及何时应该将控制权交给人类专家。本教程展示了如何使用一种称为 **谱归一化神经高斯过程 (SNGP{.external})** 的技术来提高深度分类器量化不确定性的能力。

SNGP 的核心思想是通过对网络进行简单的修改来提高深度分类器的 **距离感知能力**。模型的 **距离感知能力** 是衡量其预测概率如何反映测试示例与训练数据之间距离的指标。这是一个理想的属性，对于黄金标准概率模型（例如，具有 RBF 内核的 高斯过程{.external}）来说很常见，但在具有深度神经网络的模型中却很缺乏。SNGP 提供了一种简单的方法，可以在保持其预测准确性的同时，将这种高斯过程行为注入深度分类器。

本教程在 scikit-learn 的双月数据集{.external} 上实现了基于深度残差网络 (ResNet) 的 SNGP 模型，并将其不确定性表面与其他两种流行的不确定性方法进行了比较：蒙特卡罗 dropout{.external} 和 深度集成{.external}。

本教程在二维玩具数据集上说明了 SNGP 模型。有关将 SNGP 应用于使用 BERT-base 的现实世界自然语言理解任务的示例，请查看 SNGP-BERT 教程。有关在各种基准数据集（如 CIFAR-100、ImageNet、Jigsaw 毒性检测 等）上实现 SNGP 模型（以及许多其他不确定性方法）的高质量实现，请参考 不确定性基线{.external} 基准。

关于 SNGP

SNGP 是一种简单的方法，可以提高深度分类器的不确定性质量，同时保持类似的准确性和延迟。给定一个深度残差网络，SNGP 对模型进行了两个简单的更改

• 它将谱归一化应用于隐藏的残差层。
• 它用高斯过程层替换了 Dense 输出层。

与其他不确定性方法（如蒙特卡罗 dropout 或深度集成）相比，SNGP 具有以下几个优点

• 它适用于各种最先进的基于残差的架构（例如，(Wide) ResNet、DenseNet 或 BERT）。
• 它是一种单模型方法——它不依赖于集成平均）。因此，SNGP 与单个确定性网络具有类似的延迟级别，并且可以轻松扩展到大型数据集，如 ImageNet{.external} 和 Jigsaw 毒性评论分类{.external}。
• 由于距离感知特性，它具有强大的域外检测性能。

该方法的缺点是

• SNGP 的预测不确定性是使用 拉普拉斯近似{.external} 计算的。因此，从理论上讲，SNGP 的后验不确定性与精确高斯过程不同。

• SNGP 训练需要在新纪元开始时进行协方差重置步骤。这会给训练管道增加少量的额外复杂性。本教程展示了一种使用 Keras 回调实现此功能的简单方法。

设置

pip install -U -q --use-deprecated=legacy-resolver tf-models-official tensorflow

# refresh pkg_resources so it takes the changes into account.
import pkg_resources
import importlib
importlib.reload(pkg_resources)

import matplotlib.pyplot as plt
import matplotlib.colors as colors

import sklearn.datasets

import numpy as np
import tensorflow as tf

import official.nlp.modeling.layers as nlp_layers

定义可视化宏

plt.rcParams['figure.dpi'] = 140

DEFAULT_X_RANGE = (-3.5, 3.5)
DEFAULT_Y_RANGE = (-2.5, 2.5)
DEFAULT_CMAP = colors.ListedColormap(["#377eb8", "#ff7f00"])
DEFAULT_NORM = colors.Normalize(vmin=0, vmax=1,)
DEFAULT_N_GRID = 100

双月数据集

从 scikit-learn 双月数据集{.external} 创建训练和评估数据集。

def make_training_data(sample_size=500):
  """Create two moon training dataset."""
  train_examples, train_labels = sklearn.datasets.make_moons(
      n_samples=2 * sample_size, noise=0.1)

  # Adjust data position slightly.
  train_examples[train_labels == 0] += [-0.1, 0.2]
  train_examples[train_labels == 1] += [0.1, -0.2]

  return train_examples, train_labels

评估模型在整个二维输入空间上的预测行为。

def make_testing_data(x_range=DEFAULT_X_RANGE, y_range=DEFAULT_Y_RANGE, n_grid=DEFAULT_N_GRID):
  """Create a mesh grid in 2D space."""
  # testing data (mesh grid over data space)
  x = np.linspace(x_range[0], x_range[1], n_grid)
  y = np.linspace(y_range[0], y_range[1], n_grid)
  xv, yv = np.meshgrid(x, y)
  return np.stack([xv.flatten(), yv.flatten()], axis=-1)

为了评估模型不确定性，添加一个属于第三类的域外 (OOD) 数据集。模型在训练期间从未观察过这些 OOD 示例。

def make_ood_data(sample_size=500, means=(2.5, -1.75), vars=(0.01, 0.01)):
  return np.random.multivariate_normal(
      means, cov=np.diag(vars), size=sample_size)

# Load the train, test and OOD datasets.
train_examples, train_labels = make_training_data(
    sample_size=500)
test_examples = make_testing_data()
ood_examples = make_ood_data(sample_size=500)

# Visualize
pos_examples = train_examples[train_labels == 0]
neg_examples = train_examples[train_labels == 1]

plt.figure(figsize=(7, 5.5))

plt.scatter(pos_examples[:, 0], pos_examples[:, 1], c="#377eb8", alpha=0.5)
plt.scatter(neg_examples[:, 0], neg_examples[:, 1], c="#ff7f00", alpha=0.5)
plt.scatter(ood_examples[:, 0], ood_examples[:, 1], c="red", alpha=0.1)

plt.legend(["Positive", "Negative", "Out-of-Domain"])

plt.ylim(DEFAULT_Y_RANGE)
plt.xlim(DEFAULT_X_RANGE)

plt.show()

这里，蓝色和橙色代表正类和负类，红色代表 OOD 数据。一个能够很好地量化不确定性的模型，预计在接近训练数据时会很有信心（即，\(p(x_{test})\) 接近 0 或 1），而在远离训练数据区域时会不确定（即，\(p(x_{test})\) 接近 0.5）。

确定性模型

定义模型

从（基线）确定性模型开始：一个具有 dropout 正则化的多层残差网络 (ResNet)。

class DeepResNet(tf.keras.Model):
  """Defines a multi-layer residual network."""
  def __init__(self, num_classes, num_layers=3, num_hidden=128,
               dropout_rate=0.1, **classifier_kwargs):
    super().__init__()
    # Defines class meta data.
    self.num_hidden = num_hidden
    self.num_layers = num_layers
    self.dropout_rate = dropout_rate
    self.classifier_kwargs = classifier_kwargs

    # Defines the hidden layers.
    self.input_layer = tf.keras.layers.Dense(self.num_hidden, trainable=False)
    self.dense_layers = [self.make_dense_layer() for _ in range(num_layers)]

    # Defines the output layer.
    self.classifier = self.make_output_layer(num_classes)

  def call(self, inputs):
    # Projects the 2d input data to high dimension.
    hidden = self.input_layer(inputs)

    # Computes the ResNet hidden representations.
    for i in range(self.num_layers):
      resid = self.dense_layers[i](hidden)
      resid = tf.keras.layers.Dropout(self.dropout_rate)(resid)
      hidden += resid

    return self.classifier(hidden)

  def make_dense_layer(self):
    """Uses the Dense layer as the hidden layer."""
    return tf.keras.layers.Dense(self.num_hidden, activation="relu")

  def make_output_layer(self, num_classes):
    """Uses the Dense layer as the output layer."""
    return tf.keras.layers.Dense(
        num_classes, **self.classifier_kwargs)

本教程使用一个具有 128 个隐藏单元的六层 ResNet。

resnet_config = dict(num_classes=2, num_layers=6, num_hidden=128)

resnet_model = DeepResNet(**resnet_config)

resnet_model.build((None, 2))
resnet_model.summary()

训练模型

配置训练参数以使用 SparseCategoricalCrossentropy 作为损失函数和 Adam 优化器。

loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metrics = tf.keras.metrics.SparseCategoricalAccuracy(),
optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=1e-4)

train_config = dict(loss=loss, metrics=metrics, optimizer=optimizer)

使用批次大小 128 训练模型 100 个纪元。

fit_config = dict(batch_size=128, epochs=100)

resnet_model.compile(**train_config)
resnet_model.fit(train_examples, train_labels, **fit_config)

可视化不确定性

def plot_uncertainty_surface(test_uncertainty, ax, cmap=None):
  """Visualizes the 2D uncertainty surface.

  For simplicity, assume these objects already exist in the memory:

    test_examples: Array of test examples, shape (num_test, 2).
    train_labels: Array of train labels, shape (num_train, ).
    train_examples: Array of train examples, shape (num_train, 2).

  Arguments:
    test_uncertainty: Array of uncertainty scores, shape (num_test,).
    ax: A matplotlib Axes object that specifies a matplotlib figure.
    cmap: A matplotlib colormap object specifying the palette of the
      predictive surface.

  Returns:
    pcm: A matplotlib PathCollection object that contains the palette
      information of the uncertainty plot.
  """
  # Normalize uncertainty for better visualization.
  test_uncertainty = test_uncertainty / np.max(test_uncertainty)

  # Set view limits.
  ax.set_ylim(DEFAULT_Y_RANGE)
  ax.set_xlim(DEFAULT_X_RANGE)

  # Plot normalized uncertainty surface.
  pcm = ax.imshow(
      np.reshape(test_uncertainty, [DEFAULT_N_GRID, DEFAULT_N_GRID]),
      cmap=cmap,
      origin="lower",
      extent=DEFAULT_X_RANGE + DEFAULT_Y_RANGE,
      vmin=DEFAULT_NORM.vmin,
      vmax=DEFAULT_NORM.vmax,
      interpolation='bicubic',
      aspect='auto')

  # Plot training data.
  ax.scatter(train_examples[:, 0], train_examples[:, 1],
             c=train_labels, cmap=DEFAULT_CMAP, alpha=0.5)
  ax.scatter(ood_examples[:, 0], ood_examples[:, 1], c="red", alpha=0.1)

  return pcm

现在可视化确定性模型的预测。首先绘制类概率

\[p(x) = softmax(logit(x))\]

resnet_logits = resnet_model(test_examples)
resnet_probs = tf.nn.softmax(resnet_logits, axis=-1)[:, 0]  # Take the probability for class 0.

_, ax = plt.subplots(figsize=(7, 5.5))

pcm = plot_uncertainty_surface(resnet_probs, ax=ax)

plt.colorbar(pcm, ax=ax)
plt.title("Class Probability, Deterministic Model")

plt.show()

在此图中，黄色和紫色是两个类的预测概率。确定性模型在使用非线性决策边界对两个已知类（蓝色和橙色）进行分类方面做得很好。但是，它不是距离感知的，并且自信地将从未观察到的红色域外 (OOD) 示例分类为橙色类。

通过计算 预测方差 来可视化模型不确定性

\[var(x) = p(x) * (1 - p(x))\]

resnet_uncertainty = resnet_probs * (1 - resnet_probs)

_, ax = plt.subplots(figsize=(7, 5.5))

pcm = plot_uncertainty_surface(resnet_uncertainty, ax=ax)

plt.colorbar(pcm, ax=ax)
plt.title("Predictive Uncertainty, Deterministic Model")

plt.show()

在此图中，黄色表示高不确定性，紫色表示低不确定性。确定性 ResNet 的不确定性仅取决于测试示例与决策边界的距离。这会导致模型在超出训练域时过度自信。下一节将展示 SNGP 在此数据集上的不同行为。

SNGP 模型

定义 SNGP 模型

现在让我们实现 SNGP 模型。SNGP 的两个组件 SpectralNormalization 和 RandomFeatureGaussianProcess 都可以在 tensorflow_model 的 内置层 中找到。

让我们更详细地检查这两个组件。（您也可以跳到 完整的 SNGP 模型 部分，了解 SNGP 的实现方式。）

SpectralNormalization 包装器

SpectralNormalization{.external} 是一个 Keras 层包装器。它可以像这样应用于现有的 Dense 层

dense = tf.keras.layers.Dense(units=10)
dense = nlp_layers.SpectralNormalization(dense, norm_multiplier=0.9)

谱归一化通过逐渐引导其谱范数（即 \(W\) 的最大特征值）朝向目标值 norm_multiplier 进行正则化隐藏权重 \(W\)。

高斯过程 (GP) 层

RandomFeatureGaussianProcess{.external} 实现了一个 基于随机特征的近似{.external} 到高斯过程模型，该模型可以与深度神经网络进行端到端训练。在幕后，高斯过程层实现了一个两层网络

\[logits(x) = \Phi(x) \beta, \quad \Phi(x)=\sqrt{\frac{2}{M} } * cos(Wx + b)\]

这里，\(x\) 是输入，\(W\) 和 \(b\) 分别是从高斯分布和均匀分布中随机初始化的冻结权重。（因此，\(\Phi(x)\) 被称为“随机特征”。）\(\beta\) 是类似于 Dense 层的可学习内核权重。

batch_size = 32
input_dim = 1024
num_classes = 10

gp_layer = nlp_layers.RandomFeatureGaussianProcess(units=num_classes,
                                               num_inducing=1024,
                                               normalize_input=False,
                                               scale_random_features=True,
                                               gp_cov_momentum=-1)

GP 层的主要参数是

• units：输出 logits 的维度。
• num_inducing：隐藏权重 \(W\) 的维度 \(M\)。默认为 1024。
• normalize_input：是否对输入 \(x\) 应用层归一化。
• scale_random_features：是否对隐藏输出应用比例 \(\sqrt{2/M}\)。

• gp_cov_momentum 控制模型协方差的计算方式。如果设置为正值（例如，0.999），则协方差矩阵是使用基于动量的移动平均更新（类似于批次归一化）计算的。如果设置为 -1，则协方差矩阵在没有动量的情况下更新。

给定一个形状为 (batch_size, input_dim) 的批次输入，GP 层返回一个 logits 张量（形状为 (batch_size, num_classes)）用于预测，以及 covmat 张量（形状为 (batch_size, batch_size)），它是批次 logits 的后验协方差矩阵。

embedding = tf.random.normal(shape=(batch_size, input_dim))

logits, covmat = gp_layer(embedding)

从理论上讲，可以扩展算法以计算不同类的不同方差值（如 原始 SNGP 论文{.external} 中介绍的那样）。但是，这很难扩展到具有大型输出空间的问题（例如，使用 ImageNet 或语言建模进行分类）。

完整的 SNGP 模型

给定基类 DeepResNet，可以通过修改残差网络的隐藏层和输出层来轻松实现 SNGP 模型。为了与 Keras model.fit() API 兼容，还需要修改模型的 call() 方法，使其在训练期间仅输出 logits。

class DeepResNetSNGP(DeepResNet):
  def __init__(self, spec_norm_bound=0.9, **kwargs):
    self.spec_norm_bound = spec_norm_bound
    super().__init__(**kwargs)

  def make_dense_layer(self):
    """Applies spectral normalization to the hidden layer."""
    dense_layer = super().make_dense_layer()
    return nlp_layers.SpectralNormalization(
        dense_layer, norm_multiplier=self.spec_norm_bound)

  def make_output_layer(self, num_classes):
    """Uses Gaussian process as the output layer."""
    return nlp_layers.RandomFeatureGaussianProcess(
        num_classes,
        gp_cov_momentum=-1,
        **self.classifier_kwargs)

  def call(self, inputs, training=False, return_covmat=False):
    # Gets logits and a covariance matrix from the GP layer.
    logits, covmat = super().call(inputs)

    # Returns only logits during training.
    if not training and return_covmat:
      return logits, covmat

    return logits

使用与确定性模型相同的架构。

resnet_config

sngp_model = DeepResNetSNGP(**resnet_config)

sngp_model.build((None, 2))
sngp_model.summary()

实现一个 Keras 回调，在新纪元开始时重置协方差矩阵。

class ResetCovarianceCallback(tf.keras.callbacks.Callback):

  def on_epoch_begin(self, epoch, logs=None):
    """Resets covariance matrix at the beginning of the epoch."""
    if epoch > 0:
      self.model.classifier.reset_covariance_matrix()

将此回调添加到 DeepResNetSNGP 模型类中。

class DeepResNetSNGPWithCovReset(DeepResNetSNGP):
  def fit(self, *args, **kwargs):
    """Adds ResetCovarianceCallback to model callbacks."""
    kwargs["callbacks"] = list(kwargs.get("callbacks", []))
    kwargs["callbacks"].append(ResetCovarianceCallback())

    return super().fit(*args, **kwargs)

训练模型

使用 tf.keras.model.fit 训练模型。

sngp_model = DeepResNetSNGPWithCovReset(**resnet_config)
sngp_model.compile(**train_config)
sngp_model.fit(train_examples, train_labels, **fit_config)

可视化不确定性

首先计算预测 logits 和方差。

sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)

sngp_variance = tf.linalg.diag_part(sngp_covmat)[:, None]

现在计算后验预测概率。计算概率模型的预测概率的经典方法是使用蒙特卡罗采样，即

\[E(p(x)) = \frac{1}{M} \sum_{m=1}^M logit_m(x), \]

其中 \(M\) 是样本大小，\(logit_m(x)\) 是从 SNGP 后验 \(MultivariateNormal\)(sngp_logits,sngp_covmat) 中随机抽取的样本。但是，对于延迟敏感的应用程序（如自动驾驶或实时竞价），这种方法可能很慢。相反，您可以使用 平均场方法{.external} 来近似 \(E(p(x))\)

\[E(p(x)) \approx softmax(\frac{logit(x)}{\sqrt{1+ \lambda * \sigma^2(x)} })\]

其中 \(\sigma^2(x)\) 是 SNGP 方差，\(\lambda\) 通常选择为 \(\pi/8\) 或 \(3/\pi^2\)。

sngp_logits_adjusted = sngp_logits / tf.sqrt(1. + (np.pi / 8.) * sngp_variance)
sngp_probs = tf.nn.softmax(sngp_logits_adjusted, axis=-1)[:, 0]

这种平均场方法被实现为一个内置函数 layers.gaussian_process.mean_field_logits

def compute_posterior_mean_probability(logits, covmat, lambda_param=np.pi / 8.):
  # Computes uncertainty-adjusted logits using the built-in method.
  logits_adjusted = nlp_layers.gaussian_process.mean_field_logits(
      logits, covmat, mean_field_factor=lambda_param)

  return tf.nn.softmax(logits_adjusted, axis=-1)[:, 0]

sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)
sngp_probs = compute_posterior_mean_probability(sngp_logits, sngp_covmat)

SNGP 总结

def plot_predictions(pred_probs, model_name=""):
  """Plot normalized class probabilities and predictive uncertainties."""
  # Compute predictive uncertainty.
  uncertainty = pred_probs * (1. - pred_probs)

  # Initialize the plot axes.
  fig, axs = plt.subplots(1, 2, figsize=(14, 5))

  # Plots the class probability.
  pcm_0 = plot_uncertainty_surface(pred_probs, ax=axs[0])
  # Plots the predictive uncertainty.
  pcm_1 = plot_uncertainty_surface(uncertainty, ax=axs[1])

  # Adds color bars and titles.
  fig.colorbar(pcm_0, ax=axs[0])
  fig.colorbar(pcm_1, ax=axs[1])

  axs[0].set_title(f"Class Probability, {model_name}")
  axs[1].set_title(f"(Normalized) Predictive Uncertainty, {model_name}")

  plt.show()

现在您可以将所有内容整合在一起。整个过程——训练、评估和不确定性计算——只需五行代码即可完成

def train_and_test_sngp(train_examples, test_examples):
  sngp_model = DeepResNetSNGPWithCovReset(**resnet_config)

  sngp_model.compile(**train_config)
  sngp_model.fit(train_examples, train_labels, verbose=0, **fit_config)

  sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)
  sngp_probs = compute_posterior_mean_probability(sngp_logits, sngp_covmat)

  return sngp_probs

sngp_probs = train_and_test_sngp(train_examples, test_examples)

可视化 SNGP 模型的类概率（左）和预测不确定性（右）。

plot_predictions(sngp_probs, model_name="SNGP")

请记住，在类概率图（左侧）中，黄色和紫色代表类概率。当接近训练数据域时，SNGP 正确地对示例进行分类，并具有高度置信度（即，分配接近 0 或 1 的概率）。当远离训练数据时，SNGP 的置信度逐渐降低，其预测概率接近 0.5，而（归一化）模型不确定性上升到 1。

将此与确定性模型的不确定性表面进行比较

plot_predictions(resnet_probs, model_name="Deterministic")

如前所述，确定性模型不是距离感知的。其不确定性由测试示例与决策边界的距离定义。这会导致模型对域外示例（红色）产生过度自信的预测。

与其他不确定性方法的比较

本节将 SNGP 的不确定性与蒙特卡罗 dropout{.external} 和 深度集成{.external} 进行比较。

这两种方法都基于对确定性模型的多个前向传递进行蒙特卡罗平均。首先，设置集成大小 \(M\)。

num_ensemble = 10

蒙特卡罗 dropout

给定一个训练好的具有 Dropout 层的神经网络，蒙特卡罗 dropout 计算平均预测概率

\[E(p(x)) = \frac{1}{M}\sum_{m=1}^M softmax(logit_m(x))\]

通过对多个启用 Dropout 的前向传递 \(\{logit_m(x)\}_{m=1}^M\) 进行平均。

def mc_dropout_sampling(test_examples):
  # Enable dropout during inference.
  return resnet_model(test_examples, training=True)

# Monte Carlo dropout inference.
dropout_logit_samples = [mc_dropout_sampling(test_examples) for _ in range(num_ensemble)]
dropout_prob_samples = [tf.nn.softmax(dropout_logits, axis=-1)[:, 0] for dropout_logits in dropout_logit_samples]
dropout_probs = tf.reduce_mean(dropout_prob_samples, axis=0)

dropout_probs = tf.reduce_mean(dropout_prob_samples, axis=0)

plot_predictions(dropout_probs, model_name="MC Dropout")

深度集成

深度集成是深度学习不确定性的一种最先进（但昂贵）方法。要训练深度集成，首先训练 \(M\) 个集成成员。

# Deep ensemble training
resnet_ensemble = []
for _ in range(num_ensemble):
  resnet_model = DeepResNet(**resnet_config)
  resnet_model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
  resnet_model.fit(train_examples, train_labels, verbose=0, **fit_config)

  resnet_ensemble.append(resnet_model)

收集 logits 并计算平均预测概率 \(E(p(x)) = \frac{1}{M}\sum_{m=1}^M softmax(logit_m(x))\)。

# Deep ensemble inference
ensemble_logit_samples = [model(test_examples) for model in resnet_ensemble]
ensemble_prob_samples = [tf.nn.softmax(logits, axis=-1)[:, 0] for logits in ensemble_logit_samples]
ensemble_probs = tf.reduce_mean(ensemble_prob_samples, axis=0)

plot_predictions(ensemble_probs, model_name="Deep ensemble")

蒙特卡罗 Dropout 和深度集成方法都通过使决策边界不太确定来提高模型的不确定性能力。但是，它们都继承了确定性深度网络在缺乏距离感知方面的局限性。

总结

在本教程中，您已

• 在深度分类器上实现了 SNGP 模型，以提高其距离感知能力。
• 使用 Keras Model.fit API 对 SNGP 模型进行端到端训练。
• 可视化了 SNGP 的不确定性行为。
• 比较了 SNGP、蒙特卡罗 dropout 和深度集成模型之间的不确定性行为。

资源和进一步阅读

• 查看 SNGP-BERT 教程，了解将 SNGP 应用于 BERT 模型以进行不确定性感知自然语言理解的示例。
• 访问 Uncertainty Baselines GitHub 仓库{.external}，了解在各种基准数据集（例如，CIFAR、ImageNet、Jigsaw 毒性检测 等）上实现 SNGP 模型（以及许多其他不确定性方法）。
• 要更深入地了解 SNGP 方法，请查看题为 通过距离感知实现确定性深度学习的简单且原则性的不确定性估计{.external} 的论文。