使用 SNGP 的不確定性感知深度學習

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視 下載筆記本

在安全攸關的 AI 應用程式中 (例如醫療決策和自動駕駛)，或是資料本質上雜訊過多的情況下 (例如自然語言理解)，深度分類器務必能可靠地量化其不確定性。深度分類器應能意識到自身的限制，並在應將控制權交給人類專家時告知。本教學課程說明如何使用名為頻譜正規化神經高斯過程 (SNGP{.external}) 的技術，來提升深度分類器量化不確定性的能力。

SNGP 的核心概念是透過對網路進行簡單的修改，來提升深度分類器的距離感知能力。模型的距離感知能力是一種衡量指標，可衡量模型的預測機率如何反映測試範例與訓練資料之間的距離。這是黃金標準機率模型 (例如具有 RBF 核心的高斯過程{.external}) 的常見理想屬性，但深度神經網路模型卻缺乏此屬性。SNGP 提供一種簡單的方式，可將此高斯過程行為注入深度分類器，同時維持其預測準確性。

本教學課程在 scikit-learn 的雙月形{.external}資料集上實作了以深度殘差網路 (ResNet) 為基礎的 SNGP 模型，並將其不確定性表面與另外兩種常見的不確定性方法進行比較：蒙地卡羅 dropout{.external} 和 深度集成{.external}。

本教學課程在玩具 2D 資料集上說明 SNGP 模型。如要瞭解如何將 SNGP 應用於使用 BERT-base 的真實世界自然語言理解工作，請參閱 SNGP-BERT 教學課程。如需 SNGP 模型 (和許多其他不確定性方法) 在各種基準資料集 (例如 CIFAR-100、ImageNet、Jigsaw 毒性偵測等等) 上的高品質實作方式，請參閱 Uncertainty Baselines{.external} 基準。

關於 SNGP

SNGP 是一種簡單的方法，可在維持相似準確度和延遲時間的情況下，提升深度分類器的不確定性品質。在給定深度殘差網路的情況下，SNGP 對模型進行了兩項簡單的變更

  • 它將頻譜正規化套用至隱藏的殘差層。
  • 它將「密集」輸出層取代為高斯過程層。

SNGP

相較於其他不確定性方法 (例如蒙地卡羅 dropout 或深度集成)，SNGP 有幾個優點

  • 它適用於各種最先進的以殘差為基礎的架構 (例如 (Wide) ResNet、DenseNet 或 BERT)。
  • 它是一種單一模型方法，不仰賴集成平均)。因此，SNGP 的延遲時間與單一確定性網路相似，並且可以輕鬆擴充至大型資料集，例如 ImageNet{.external} 和 Jigsaw 毒性評論分類{.external}。
  • 由於距離感知能力屬性，它具有強大的領域外偵測效能。

此方法的缺點如下

  • SNGP 的預測不確定性是使用 Laplace 近似{.external} 計算而得。因此，從理論上來說，SNGP 的後驗不確定性與精確高斯過程的不確定性不同。

  • SNGP 訓練在新 epoch 開始時需要一個共變異數重設步驟。這可能會為訓練管道增加少許額外的複雜性。本教學課程示範如何使用 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

評估模型在整個 2D 輸入空間中的預測行為。

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)。

本教學課程使用具有 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 optimizer。

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 個 epoch。

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

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

視覺化不確定性

現在視覺化確定性模型的預測。首先繪製類別機率

\[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 模型章節，瞭解 SNGP 的實作方式。)

SpectralNormalization 包裝函式

SpectralNormalization{.external} 是 Keras 層包裝函式。它可以像這樣套用至現有的「密集」層

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\) 是可學習的核心權重，與「密集」層的權重類似。

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 回呼，以在新 epoch 開始時重設共變異數矩陣。

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 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 和深度集成模型之間的不確定性行為。

資源和延伸閱讀