基于 BERT-SNGP 的不确定性感知深度语言学习

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

查看 TF Hub 模型

在 SNGP 教程中，您学习了如何基于深度残差网络构建 SNGP 模型，以提高其量化不确定性的能力。在本教程中，您将通过在深度 BERT 编码器之上构建 SNGP，将其应用于自然语言理解 (NLU) 任务，从而提升深度 NLU 模型检测范围外 (out-of-scope) 查询的能力。

具体来说，您将完成以下内容：

构建 BERT-SNGP，一个经过 SNGP 增强的 BERT 模型。
加载 CLINC 范围外 (OOS) 意图检测数据集。
训练 BERT-SNGP 模型。
评估 BERT-SNGP 模型在不确定性校准和领域外检测方面的性能。

除了 CLINC OOS 之外，SNGP 模型还被应用于大规模数据集（如 Jigsaw 毒性检测）以及图像数据集（如 CIFAR-100 和 ImageNet）。有关 SNGP 及其他不确定性方法的基准测试结果，以及包含端到端训练/评估脚本的高质量实现，您可以查看 Uncertainty Baselines 基准测试库。

设置

pip uninstall -y tensorflow tf-text

pip install "tensorflow-text==2.11.*"

pip install -U tf-models-official==2.11.0

import matplotlib.pyplot as plt

import sklearn.metrics
import sklearn.calibration

import tensorflow_hub as hub
import tensorflow_datasets as tfds

import numpy as np
import tensorflow as tf

import official.nlp.modeling.layers as layers
import official.nlp.optimization as optimization

本教程需要 GPU 才能高效运行。请检查 GPU 是否可用。

tf.__version__

gpus = tf.config.list_physical_devices('GPU')
gpus

assert gpus, """
  No GPU(s) found! This tutorial will take many hours to run without a GPU.

  You may hit this error if the installed tensorflow package is not
  compatible with the CUDA and CUDNN versions."""

首先，按照使用 BERT 进行文本分类教程实现一个标准的 BERT 分类器。我们将使用 BERT-base 编码器，并将内置的 ClassificationHead 作为分类器。

标准 BERT 模型

PREPROCESS_HANDLE = 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3'
MODEL_HANDLE = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3'

class BertClassifier(tf.keras.Model):
  def __init__(self, 
               num_classes=150, inner_dim=768, dropout_rate=0.1,
               **classifier_kwargs):

    super().__init__()
    self.classifier_kwargs = classifier_kwargs

    # Initiate the BERT encoder components.
    self.bert_preprocessor = hub.KerasLayer(PREPROCESS_HANDLE, name='preprocessing')
    self.bert_hidden_layer = hub.KerasLayer(MODEL_HANDLE, trainable=True, name='bert_encoder')

    # Defines the encoder and classification layers.
    self.bert_encoder = self.make_bert_encoder()
    self.classifier = self.make_classification_head(num_classes, inner_dim, dropout_rate)

  def make_bert_encoder(self):
    text_inputs = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')
    encoder_inputs = self.bert_preprocessor(text_inputs)
    encoder_outputs = self.bert_hidden_layer(encoder_inputs)
    return tf.keras.Model(text_inputs, encoder_outputs)

  def make_classification_head(self, num_classes, inner_dim, dropout_rate):
    return layers.ClassificationHead(
        num_classes=num_classes, 
        inner_dim=inner_dim,
        dropout_rate=dropout_rate,
        **self.classifier_kwargs)

  def call(self, inputs, **kwargs):
    encoder_outputs = self.bert_encoder(inputs)
    classifier_inputs = encoder_outputs['sequence_output']
    return self.classifier(classifier_inputs, **kwargs)

构建 SNGP 模型

要实现 BERT-SNGP 模型，您只需将 ClassificationHead 替换为内置的 GaussianProcessClassificationHead。谱归一化 (Spectral normalization) 已经预封装在此分类头中。与 SNGP 教程中一样，向模型添加一个协方差重置回调函数，以便模型在每个新 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()

class SNGPBertClassifier(BertClassifier):

  def make_classification_head(self, num_classes, inner_dim, dropout_rate):
    return layers.GaussianProcessClassificationHead(
        num_classes=num_classes, 
        inner_dim=inner_dim,
        dropout_rate=dropout_rate,
        gp_cov_momentum=-1,
        temperature=30.,
        **self.classifier_kwargs)

  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)

加载 CLINC OOS 数据集

现在加载 CLINC OOS 意图检测数据集。该数据集包含 15000 条用户语音查询，涵盖 150 个意图类别，还包含 1000 条不属于任何已知类别的领域外 (OOD) 句子。

(clinc_train, clinc_test, clinc_test_oos), ds_info = tfds.load(
    'clinc_oos', split=['train', 'test', 'test_oos'], with_info=True, batch_size=-1)

准备训练和测试数据。

train_examples = clinc_train['text']
train_labels = clinc_train['intent']

# Makes the in-domain (IND) evaluation data.
ind_eval_data = (clinc_test['text'], clinc_test['intent'])

创建 OOD 评估数据集。为此，将领域内测试数据 clinc_test 和领域外数据 clinc_test_oos 合并。我们还将为领域内样本分配标签 0，为领域外样本分配标签 1。

test_data_size = ds_info.splits['test'].num_examples
oos_data_size = ds_info.splits['test_oos'].num_examples

# Combines the in-domain and out-of-domain test examples.
oos_texts = tf.concat([clinc_test['text'], clinc_test_oos['text']], axis=0)
oos_labels = tf.constant([0] * test_data_size + [1] * oos_data_size)

# Converts into a TF dataset.
ood_eval_dataset = tf.data.Dataset.from_tensor_slices(
    {"text": oos_texts, "label": oos_labels})

训练和评估

首先设置基本的训练配置。

TRAIN_EPOCHS = 3
TRAIN_BATCH_SIZE = 32
EVAL_BATCH_SIZE = 256

def bert_optimizer(learning_rate, 
                   batch_size=TRAIN_BATCH_SIZE, epochs=TRAIN_EPOCHS, 
                   warmup_rate=0.1):
  """Creates an AdamWeightDecay optimizer with learning rate schedule."""
  train_data_size = ds_info.splits['train'].num_examples

  steps_per_epoch = int(train_data_size / batch_size)
  num_train_steps = steps_per_epoch * epochs
  num_warmup_steps = int(warmup_rate * num_train_steps)  

  # Creates learning schedule.
  lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
      initial_learning_rate=learning_rate,
      decay_steps=num_train_steps,
      end_learning_rate=0.0)  

  return optimization.AdamWeightDecay(
      learning_rate=lr_schedule,
      weight_decay_rate=0.01,
      epsilon=1e-6,
      exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias'])

optimizer = bert_optimizer(learning_rate=1e-4)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metrics = tf.metrics.SparseCategoricalAccuracy()

fit_configs = dict(batch_size=TRAIN_BATCH_SIZE,
                   epochs=TRAIN_EPOCHS,
                   validation_batch_size=EVAL_BATCH_SIZE, 
                   validation_data=ind_eval_data)

sngp_model = SNGPBertClassifier()
sngp_model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
sngp_model.fit(train_examples, train_labels, **fit_configs)

评估 OOD 性能

评估模型检测陌生领域外查询的能力。为了进行严谨的评估，请使用前面构建的 OOD 评估数据集 ood_eval_dataset。

def oos_predict(model, ood_eval_dataset, **model_kwargs):
  oos_labels = []
  oos_probs = []

  ood_eval_dataset = ood_eval_dataset.batch(EVAL_BATCH_SIZE)
  for oos_batch in ood_eval_dataset:
    oos_text_batch = oos_batch["text"]
    oos_label_batch = oos_batch["label"] 

    pred_logits = model(oos_text_batch, **model_kwargs)
    pred_probs_all = tf.nn.softmax(pred_logits, axis=-1)
    pred_probs = tf.reduce_max(pred_probs_all, axis=-1)

    oos_labels.append(oos_label_batch)
    oos_probs.append(pred_probs)

  oos_probs = tf.concat(oos_probs, axis=0)
  oos_labels = tf.concat(oos_labels, axis=0) 

  return oos_probs, oos_labels

计算 OOD 概率为 \(1 - p(x)\)，其中 \(p(x)=softmax(logit(x))\) 为预测概率。

sngp_probs, ood_labels = oos_predict(sngp_model, ood_eval_dataset)

ood_probs = 1 - sngp_probs

现在评估模型的置信度分数 ood_probs 预测领域外标签的效果。首先计算 OOD 概率与 OOD 检测准确率之间的精确率-召回率曲线下面积 (AUPRC)。

precision, recall, _ = sklearn.metrics.precision_recall_curve(ood_labels, ood_probs)

auprc = sklearn.metrics.auc(recall, precision)
print(f'SNGP AUPRC: {auprc:.4f}')

这与 CLINC OOS 基准测试中在 Uncertainty Baselines 下报告的 SNGP 性能相符。

接下来，检查模型在不确定性校准方面的质量，即模型的预测概率是否与其预测准确率相符。一个经过良好校准的模型被认为是值得信赖的，因为例如其预测概率 \(p(x)=0.8\) 意味着模型有 80% 的时间是正确的。

prob_true, prob_pred = sklearn.calibration.calibration_curve(
    ood_labels, ood_probs, n_bins=10, strategy='quantile')

plt.plot(prob_pred, prob_true)

plt.plot([0., 1.], [0., 1.], c='k', linestyle="--")
plt.xlabel('Predictive Probability')
plt.ylabel('Predictive Accuracy')
plt.title('Calibration Plots, SNGP')

plt.show()

资源与进一步阅读

有关从零开始实现 SNGP 的详细演练，请参阅 SNGP 教程。
有关在各种基准数据集（例如 CIFAR、ImageNet、Jigsaw 毒性检测等）上实现 SNGP 模型（以及许多其他不确定性方法）的信息，请参阅 Uncertainty Baselines。
要更深入地了解 SNGP 方法，请参阅论文《Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness》（通过距离感知实现确定性深度学习的简单且原则性的不确定性估计）。