#  CUE: An Uncertainty Interpretation Framework for Text Classifiers Built on
Pre-Trained Language Models

Jiazheng Li Department of Informatics, King’s College London, UK  Zhaoyue Sun
Department of Computer Science, University of Warwick, UK  Bin Liang  Joint
Lab of HITSZ-CMS, Harbin Institute of Technology, Shenzhen, China  Lin Gui
Department of Informatics, King’s College London, UK  Yulan He  Department of
Informatics, King’s College London, UK  Department of Computer Science,
University of Warwick, UK  The Alan Turing Institute, UK

###### Abstract

Text classifiers built on Pre-trained Language Models (PLMs) have achieved
remarkable progress in various tasks including sentiment analysis, natural
language inference, and question-answering. However, the occurrence of
uncertain predictions by these classifiers poses a challenge to their
reliability when deployed in practical applications. Much effort has been
devoted to designing various probes in order to understand what PLMs capture.
But few studies have delved into factors influencing PLM-based classifiers’
predictive uncertainty. In this paper, we propose a novel framework, called
CUE, which aims to interpret uncertainties inherent in the predictions of PLM-
based models. In particular, we first map PLM-encoded representations to a
latent space via a variational auto-encoder. We then generate text
representations by perturbing the latent space which causes fluctuation in
predictive uncertainty. By comparing the difference in predictive uncertainty
between the perturbed and the original text representations, we are able to
identify the latent dimensions responsible for uncertainty and subsequently
trace back to the input features that contribute to such uncertainty. Our
extensive experiments on four benchmark datasets encompassing linguistic
acceptability classification, emotion classification, and natural language
inference show the feasibility of our proposed framework. Our source code is
available at https://github.com/lijiazheng99/CUE.

##  1 Introduction

Text classifiers built on Pre-trained Language Models (PLMs) have made
remarkable progress on various Natural Language Processing (NLP) tasks [Devlin
et al., 2019, Liu et al., 2020b, Lan et al., 2020, Sanh et al., 2019].
However, their deployment in practical applications still faces significant
challenges. Of particular concern, these models tend to make over-confident
predictions in uncertain cases [Guo et al., 2017, He et al., 2020, Malinin and
Gales, 2018]. Since PLMs have been widely used in various applications, such
issues cause concerns about model trustworthiness and transparency, which
becomes a barrier to deploying PLMs in sensitive domains such as medicine and
finance.

Figure 1: The illustration of proposed framework: (a) The original
representations learned from a Pre-trained Language Model (PLM) and the
decision boundary separating two classes; (b) We can freeze the PLM-classifier
parameters and perturb the PLM-encoded representations in the latent space to
increase the aleatoric (data) uncertainty along some isotropic directions
while preserving the predictive labels. This would help determine the
uncertain areas; (c) Using the representations reconstructed from the
perturbed latent space, the uncertainty for out-of-distribution (OOD) areas
can be estimated.

Predictive uncertainty is generally believed to include two aspects -
_aleatoric uncertainty_ and _epistemic uncertainty_ , where the aleatoric
uncertainty measures the data uncertainty due to inherent random effects and
is irreducible, while the epistemic uncertainty measures the uncertainty
caused by the lack of knowledge from data and is reducible [Malinin and Gales,
2018]. Numerous approaches have been proposed to estimate the predictive
uncertainty of deep neural models, such as Deep Ensemble models
[Lakshminarayanan et al., 2017], Bayesian Neural Networks (BNN) [Blundell et
al., 2015] and Monte-Carlo (MC) Dropout [Gal and Ghahramani, 2016]. Similar
idea has been applied to PLMs in recent years, to study the uncertainty of
text classifiers. Particularly, quantifying uncertainty in PLM-based
classifiers can be done by incorporating weight uncertainty into the PLM
architecture. However, uncertainty can only be induced to a certain number of
layers (e.g., the last layer of the PLM feature extractor and/or the
classification layer) due to a large number of PLM layers and parameters.
Alternatively, one could use deep ensembles by aggregating classification
results generated from multiple PLM classifiers trained with different
initialisation [Lakshminarayanan et al., 2017], or apply MC dropout in the
inference stage to estimate the uncertainty of PLMs [Vazhentsev et al., 2022].

Previous studies [Vulić et al., 2020, Clark et al., 2019, Yang et al., 2021]
have also been devoted to designing various probes in order to understand what
PLMs capture. Nevertheless, they largely ignore the interpretation of the
source of the uncertainty, i.e., identifying the input features which cause
classification uncertainty, which can be crucial for understanding the model
and taking appropriate mitigating strategies. In text classification, recent
research tried to identify word tokens that lead to uncertainty via
perturbations on input sequences [Shen et al., 2020, Kim et al., 2020].
However, due to the discreteness of textual data, token replacement or removal
would require a large search space on the input sequence and incur expensive
computational costs.

In this paper, we aim to interpret the predictive uncertainty on PLM text
classifiers by identifying the input tokens that cause the uncertainty. We
propose a novel PLM Classifier Uncertainty Explanation (CUE) framework built
on Variational Auto-encoder (VAE) [Kingma and Welling, 2014, Card et al.,
2018] that generates perturbations on latent text representations to induce
uncertain predictions. As shown in Figure 1, we can perturb the PLM-encoded
representations in the latent space to increase the aleatoric uncertainty
(data uncertainty) along some isotropic directions while preserving the
predictive labels. As will be shown in §3.2, this is equivalent to decreasing
the predictive epistemic uncertainty. By examining the difference between the
original and the perturbed text representations, a subset of input features
(i.e., word tokens) can be identified as the interpretation of the original
model’s predictive uncertainty. We compared our framework with existing
approaches addressing the predictive uncertainty problem on three
classification tasks across four benchmark datasets. Extensive experimental
results show that our proposed method can identify the source of epistemic
uncertainty and calibrate text representations from four commonly used PLMs.

In summary, our contributions are: (1) We propose a novel framework CUE to
induce perturbations on PLM-encoded representations for uncertainty
interpretation of the PLM-based text classifiers. (2) We propose an
uncertainty feature identification algorithm to identify token-level features
which lead to model predictive uncertainty. (3) We validate the effectiveness
of our proposed framework by conducting extensive experiments using various
classifiers built on four commonly-used PLMs on three different tasks and four
datasets with class numbers ranging from 2 to 27. The results show that our
proposed framework achieves lower expected calibration errors compared to
existing approaches such as label smoothing, MC dropout, and BNN. To the best
of our knowledge, our framework is the first to study the token interpretation
of PLM-based classifiers’ predictive uncertainty from the representation
space, without editing the semantic meaning of the original input text.

##  2 Related work

Our work is related to two lines of research, interpretation of PLMs and
uncertainty estimation in ML.

#### Interpretation of PLMs

Transformer-based language models have achieved impressive performance across
various NLP tasks [Devlin et al., 2019, Liu et al., 2020b]. However, the
complex structure of these models has raised concerns about model transparency
and reliability. Thus, there has been growing interest in developing methods
to interpret PLMs. For example, Clark et al. [2019] proposed an attention-
based visualisation method to interpret the model parameters by probing the
feature space to determine the potential influence of the model output.
Brunner et al. [2020] studied the identifiability of attention weights in the
BERT model and found that the distribution of self-attentions cannot be
directly used as an interpretation. There has also been work focusing on
interpreting the representations from PLMs [Zhou and Srikumar, 2021] and
attention weights [Sun and Marasović, 2021, Mareček and Rosa, 2019].

#### Uncertainty Estimation

As the interpretation of PLMs cannot provide prediction confidence directly,
much effort has been devoted to developing approaches for uncertainty
estimation of neural models. A straightforward approach to uncertainty
estimation is using deep ensemble models [Lakshminarayanan et al., 2017].
Various Bayesian inference methods have also been developed to prevent
overfitting by attaching distributions to parameters in standard networks and
estimating parameters via posterior inference [Blundell et al., 2015].
Alternatively, uncertainty estimation can be performed using MC dropout [Gal
and Ghahramani, 2016], which performs multiple stochastic forwards passes with
dropout in a network during the inference stage to produce an ensemble of
predictions. Other approaches to uncertainty estimation include prior networks
[Malinin and Gales, 2018]. Taking advantage of the development of transformer
[Vaswani et al., 2017], there has been increasing interest in investigating
classification uncertainty of language models [Xiao and Wang, 2019, Desai and
Durrett, 2020]. Various methods have been developed, partly inspired by the
research in computer vision, from uncertainty quantification via input
marginalization [Shen et al., 2020, Kim et al., 2020] to MC dropout and
Bayesian inference methods such as SNGP [Shelmanov et al., 2021, Vazhentsev et
al., 2022, Liu et al., 2020a]. Nevertheless, the aforementioned approaches
cannot identify the cause of the uncertainty.

To overcome the limitation of existing methods, we propose an uncertainty
analysis framework CUE built on VAE [Kingma and Welling, 2014, Card et al.,
2018, Qiu et al., 2020], in which noise can be generated by perturbing the
latent representation space. This allows us to disentangle the source of
uncertainty via text representation dimensions and study PLM-based
classifiers’ predictive uncertainty at both the sequence- and the token-level.

##  3 Background

Figure 2: CUE requires two phases of training: Phase 1: Fine-tune a classifier
built on a PLM on a target dataset; Phase 2: Freeze the PLM and classification
layer parameters and plug-in the CUE module to train the model that perturbs
the latent representation $\bm{z}_{i}$ to generate $\bm{e}^{\prime}_{i}$ which
leads to uncertainty fluctuation.

###  3.1 Problem setup

We are given a labelled text classification dataset, where $X$ is input text
and $Y$ is the label set. $\forall\\{x_{i},y_{i}\\}\in X\times Y$, where
$(x_{i},y_{i}),i=1,2,...,N$, is an i.i.d. realisation of the random variables,
$P(Y|X)\sim\mathbb{D}$, where $\mathbb{D}$ is the unknown ground truth
conditional distribution of class labels. To train a text classifier built on
a PLM, we need to find an optimal feature extraction function $f$ and a
classification layer $g$ with trainable parameter $\bm{\vartheta}$ and
$\bm{\eta}$, respectively: $x_{i}\stackrel{{\scriptstyle
f_{\bm{\vartheta}}(\cdot)}}{{\longrightarrow}}\bm{e}_{i}\stackrel{{\scriptstyle
g_{\bm{\eta}(\cdot)}}}{{\longrightarrow}}\hat{y}_{i}$, which first encodes
text into a representation $\bm{e}_{i}$ and then outputs a probability
distribution over the label set with the predicted output close to the desired
true label $y_{i}$. In this work, we take one step further to analyse the
potential uncertainty in the two stages of the learning process: 1) in
$\bm{\eta}$: $\bm{e}_{i}\stackrel{{\scriptstyle
g_{\bm{\eta}(\cdot)}}}{{\longrightarrow}}\hat{y}_{i}$, which dimension(s) in
$\bm{e}_{i}$ is the source of uncertainty in prediction; and 2) in
$\bm{\vartheta}$: $x_{i}\stackrel{{\scriptstyle
f_{\bm{\vartheta}}(\cdot)}}{{\longrightarrow}}\bm{e}_{i}$, which input tokens
cause the uncertainty. Before we detail our proposed uncertainty estimation
approach, we give the formal definition of uncertainty first.

###  3.2 Uncertainty Estimation

According to established definitions found in prior literature, uncertainty
can be defined based on the probability of predictive error [Sullivan, 2015],
the mean squared error (MSE) [Cervera et al., 2021], or the conditional
entropy [Malinin and Gales, 2021]. We adopt the MSE-based definition as a
representative measure of uncertainty, which is chosen without compromising
the generality of our approach.

######  Definition 3.1.

$\forall P(y|x)\in\mathbb{D}$, the predictive epistemic uncertainty can be
defined by $\mathbb{E}\big{[}\big{(}\mathbb{E}[y]-\hat{y}\big{)}^{2}\big{]}$.

Here, $\hat{y}$ is the class label for input $x$ predicted by the trained
classifier. $\mathbb{E}[y]$ is the expectation of the ground truth label
distribution, which is however unknown to the learner, making it impossible to
calculate the epistemic uncertainty based on predictive variance directly.
Therefore, we propose to estimate the uncertainty by decomposing the variance
based on the observed training data, $\\{x_{i},y_{i}\\}_{i=1}^{N}$, which
yields:

$\small\mathbb{E}[(y_{i}-\hat{y}_{i})^{2}]=\underbrace{\mathbb{E}[(y_{i}-\mathbb{E}[y])^{2}]}_{\rm
aleatoric\,uncertainty}+\underbrace{\mathbb{E}[(\mathbb{E}[y]-\hat{y}_{i})^{2}]}_{\rm
epistemic\,uncertainty}$ (1)

Since the first term, $\mathbb{E}[(y_{i}-\mathbb{E}[y])^{2}]$, contains the
observed $y_{i}$, it can be defined as the aleatoric uncertainty. The detailed
derivation of Eq. (1) can be found in our Supplementary Material Section 1.1.
Similar to the setup in [Heiss et al., 2023], if we assume the conditional
distribution of class labels follows a Gaussian distribution:

######  Assumption 3.2.

$\forall P(y_{i}|x_{i})\in\mathbb{D}$, the true label distribution for a give
data $x_{i}$ follows a Gaussian noise based generating process:
$P(y_{i}|x_{i})=\bm{\mu}_{y}+\bm{\varepsilon}_{y_{i}}$, where
$\bm{\mu}_{y}=\mathbb{E}[y]$ and the noise $\bm{\varepsilon}_{y_{i}}$ follows
a Gaussian distribution of $\mathcal{N}(0,\bm{\sigma}_{y}^{2})$ and
$\bm{\sigma}_{y}^{2}=\mathbb{E}[(\mathbb{E}[y]-y_{i})^{2}]$.

Then, the epistemic uncertainty can be written as:

$\mathbb{E}\big{[}(\mathbb{E}[y]-\hat{y}_{i})^{2}\big{]}=\mathbb{E}[(y_{i}-\hat{y}_{i})^{2}]-\bm{\sigma}_{y}^{2}.$
(2)

Here, the term $\mathbb{E}[(y_{i}-\hat{y}_{i})^{2}]$ in the Eq. (2) is the
empirical MSE on the training data which can be optimised in the training
process. The term $\bm{\sigma}_{y}^{2}$ is based on the true label
distribution which is unseen to the learner. We need to clarify that
$P(y_{i}|x_{i})$ can be larger than $1$ under the assumption. Therefore, there
is necessary to stack a normalisation layer before the prediction to guarantee
the sum of predictive probabilities for different class labels is $1$.

We assume that the empirical MSE has been minimised by the trained PLM-based
classifier with parameters $\bm{\vartheta}$ and $\bm{\eta}$. To minimise the
epistemic uncertainty given by Eq. (2) for a given $x_{i}$ and its
corresponding representation $\bm{e}_{i}$, we have to increase
$\bm{\sigma}_{y}^{2}$, which however cannot be calculated directly. We propose
to use a VAE-based generative model parameterised by $\bm{\omega}$ to
reconstruct $\bm{e}_{i}$ by adding Gaussian noise while preserving the
predictive label, resulting in $\bm{e}^{\prime}_{i}$. The reconstructed
representation should be similar to the original input representation,
$\bm{e}^{\prime}_{i}\approx\bm{e}_{i}$, and the predictive class label
distribution from $\bm{e}^{\prime}_{i}$,
$\hat{y}^{\prime}_{i}=g_{\bm{\eta}}(\bm{e}^{\prime}_{i})$, should be close to
$\hat{y}_{i}=g_{\bm{\eta}}(\bm{e}_{i})$,
$\hat{y}^{\prime}_{i}\approx\hat{y}_{i}$. This allows us to manipulate the
latent code of the VAE to increase the variance of the Gaussian noise, which
leads to the resulting label distribution closer to a uniform distribution in
the out-of-distribution (OOD) area, thus achieving a lower epistemic
uncertainty. Accordingly, we define the learning objective function as:

Learning objective: $\forall P(y_{i}|x_{i})\in\mathbb{D}$ under the Assumption
3.2, the learning objective is to:

$\small\mathop{min}_{\bm{\omega}}\mathbb{E}[(y_{i}-\hat{y}_{i})^{2}]-\mathcal{H}_{e^{\prime}}(\hat{y}),\quad
s.t.\quad\bm{e}_{i}\approx\bm{e}^{\prime}_{i},\quad\hat{y}_{i}\approx\hat{y}^{\prime}_{i},$
(3)

where $\bm{\omega}$ denotes the parameters of VAE,
$\bm{e}_{i}=f_{\bm{\vartheta}}(x_{i})$,
$\bm{e}^{\prime}_{i}=f_{\bm{\omega}}(\bm{e}_{i})$,
$\hat{y}_{i}=g_{\bm{\eta}}(\bm{e}_{i})$,
$\hat{y}^{\prime}_{i}=g_{\bm{\eta}}(\bm{e}^{\prime}_{i})$,
$\mathcal{H}_{e^{\prime}}(\hat{y})$ is the estimated entropy by the predictive
label distribution from the reconstructed $\bm{e}^{\prime}_{i}$, which
approximates the variance of the true label distribution. The above learning
objective can be formulated using the method of Lagrange multipliers:

$\displaystyle\mathcal{L}(x_{i},y_{i})=$
$\displaystyle-\mathbb{E}[\mathcal{H}_{e^{\prime}}(\hat{y}_{i})]+\lambda_{1}\mathbb{E}[(f_{\bm{\vartheta}}(x_{i})-f_{\bm{\omega}}(f_{\bm{\vartheta}}(x_{i})))^{2}]$
$\displaystyle+\lambda_{2}\mbox{KL}\big{(}\hat{y}^{\prime}_{i}||\hat{y}_{i}\big{)}$
(4)

Therefore, by optimising Eq. (4), we can obtain an alternative representation
of $\bm{e}^{\prime}_{i}$ with the predictive distribution of
$\hat{y}^{\prime}_{i}$ using the parameters $\bm{\omega}$, where the lower
bound of the epistemic uncertainty can be obtained by increasing the aleatoric
uncertainty defined by $\mathcal{H}_{e^{\prime}}(\hat{y}_{i})$. In the next
section, we show how each term in Eq. (4) can be defined in our VAE-based
uncertainty interpretation framework CUE.

##  4 Uncertainty Interpretation

In this paper, we are interested in interpreting model uncertainty, that is,
what input features lead to the predictive uncertainty. To this end, we
propose a VAE-based uncertainty interpretation framework CUE, as shown in
Figure 2. Rather than directly perturbing the input features, perturbations
can be done in the latent space in CUE to generate the modified input
representation such that it still resides on the original data manifold while
the model’s predictive epistemic uncertainty on the modified input is reduced.
By examining the difference between the original and the perturbed text
representations, a subset of input features (i.e., word tokens) can be
identified as the interpretation of the original model’s predictive
uncertainty.

We will first present how to generate perturbations on latent space in order
to cause the prediction uncertainty change (§4.1). We will then describe how
to identify input features that lead to original prediction uncertainty to
facilitate the interpretation of model predictive uncertainty (§4.2).

###  4.1 Latent Space Perturbation for Epistemic Uncertainty Reduction

Once a classifier built on a PLM is fine-tuned on a target dataset, we freeze
the parameters of the PLM and the classification layer and then insert the CUE
between the PLM last layer and the task-specific classification layer. The
PLM-encoded representation $\bm{e}_{i}$ is mapped to a latent vector, denoted
by $\bm{z}_{i}$, via CUE which consists of two networks.

The _encoder_ network $\bm{\phi}$, learns the distribution of a lower
dimensional latent variable $\bm{z}_{i}$ given the PLM-encoded representation
by a random Gaussian noise $\epsilon$:
$\bm{z}_{i}=\bm{\mu_{\phi}}(\bm{e_{i}})+\epsilon\cdot\bm{\sigma_{\phi}}(\bm{e_{i}})$,
i.e., $\bm{z}_{i}\sim\mathcal{N}(\bm{\mu_{\phi}},\bm{\sigma_{\phi}}^{2})$.

The _decoder_ network $\bm{\theta}$, reconstructs the text representation
given the latent variable $\bm{z}_{i}$, defined as
$\bm{e}_{i}^{\prime}=p_{\bm{\theta}}(\bm{z}_{i})$. Although
$p_{\bm{\theta}}(\bm{z}_{i})$ can be any decoding network, our implementation
utilises a linear mapping $\bm{W_{\theta}}$ without a bias term. The benefit
is that $\bm{W_{\theta}}$ can be treated as a set of learnable vectors and the
reconstructed text representation $\bm{e}_{i}^{\prime}$ can be written as a
linear combination of the decoded output generated from each of the latent
dimensions of $\bm{z}_{i}$. As will be discussed in §4.2, such a decomposition
form of decoding as illustrated in Eq. (10) allows the identification of
latent dimensions of $\bm{z}_{i}$ which causes predictive uncertainty.

The VAE parameters are denoted as $\bm{\omega}=\\{\bm{\phi},\bm{\theta}\\}$.
The classifier’s prediction on the reconstructed representation
$\bm{e}^{\prime}_{i}$ is denoted as
$\hat{y}^{\prime}_{i}=g_{\bm{\eta}}(\bm{e}^{\prime}_{i})$. Here, we choose to
use the Softmax based prediction layer to normalise the predictive
probability, but the representation $\bm{e}^{\prime}_{i}$ before the
normalisation should follow the Gaussian distribution since it is captured by
a linear combination of Gaussians. Besides, the latent representation
$\bm{z}_{i}$ can be perturbed which leads to uncertain predictions bounded by
a uniform distribution probability, $\log K$ ($K$ is the label set size)111The
proof is shown in Supplementary Material Section 1.2.. For the training of the
CUE model, we define various loss terms in Eq. (4) below:

Minimum change on both the perturbed representation and the model prediction.
The reconstructed $\bm{e}^{\prime}_{i}$ should be similar to the original
$\bm{e}_{i}$.

$\small\mathcal{L}_{r}=\mathinner{\\!\left\lVert\bm{e}^{\prime}_{i}-\bm{e}_{i}\right\rVert}^{2}$
(5)

The prediction, $p(\hat{y}^{\prime}_{i}|\bm{e}^{\prime}_{i})$, based on the
reconstructed $e^{\prime}_{i}$, should be close to the origin prediction
$p(\hat{y}_{i}|\bm{e}_{i})$.

$\small\mbox{KL}(\hat{y}^{\prime}_{i}||\hat{y}_{i})=\sum_{k=1}^{K}p(\hat{y}^{\prime}_{i_{k}}|\bm{e}^{\prime}_{i})\log\frac{p(\hat{y}^{\prime}_{i_{k}}|\bm{e}^{\prime}_{i})}{p(\hat{y}_{i_{k}}|\bm{e}_{i})},$
(6)

where $K$ denotes the size of the class label set.

Predictive Entropy Increment. We need to increase the predictive entropy
$\mathcal{H}_{e^{\prime}}(\hat{y}^{\prime}_{i})$ calculated based on the
reconstructed input representation $\bm{e}^{\prime}_{i}$, which approximates
the variance of the true label distribution, $\bm{\sigma}_{y}^{2}$, in order
to decrease the model epistemic uncertainty defined in Eq. (2).

$\small\mathcal{H}_{e^{\prime}}(\hat{y}^{\prime}_{i})=-\sum_{k=1}^{K}p(\hat{y}^{\prime}_{i_{k}}|\bm{e}^{\prime}_{i})\log
p(\hat{y}^{\prime}_{i_{k}}|\bm{e}^{\prime}_{i})$ (7)

In addition, we incorporate an orthogonality constraint within the decoder to
encourage independence among dimensions of the latent variable:

$\small\mathcal{L}_{o}=\|{\rm\bf
I}-\bm{W_{\theta}}\times\bm{W_{\theta}}^{\mathsf{T}}\|$ (8)

where $\rm\bf I$ is an identity matrix, $W_{\theta}$ is the weights in the
decoder. The final objective function is then defined as:

$\small\mathcal{L}=\gamma_{1}\mathcal{L}_{r}+\gamma_{2}\mbox{KL}(\hat{y}^{\prime}_{i}||\hat{y}_{i})-\gamma_{3}\mathcal{H}_{e^{\prime}}(\hat{y}^{\prime}_{i})+\gamma_{4}\mathcal{L}_{o}$
(9)

where the $\gamma$ coefficients are used to balance various loss terms.
Minimising the loss function defined in Eq. (9) is equivalent to introducing
perturbation in the latent space so as to increase the predictive entropy. We
can use the reconstruction error $||\bm{e}^{\prime}_{i}-\bm{e}_{i}||^{2}$ to
represent the perturbed noise that leads to predictive uncertainty difference
$\Delta\mathcal{H}$. As will be shown in Supplementary Material Section 1.2,
$\Delta\mathcal{H}$ is proportional to the reconstruction error
$||\bm{e}^{\prime}_{i}-\bm{e}_{i}||^{2}$. As such, the reconstruction error
can be used to interpret the predictive uncertainty. By retracing alterations
made in the input feature space, we can effectively identify features which
cause the uncertainty. To the best of our knowledge, we are the first to apply
perturbations in the latent representation space to interpret the predictive
uncertainty associated with PLM-based classifiers.

###  4.2 Input Feature Identification for Uncertainty Interpretation

| CoLA | GoEmotions
---|---|---
Model | Acc | F1 | $\mathcal{H}$ | ECE$\downarrow$ | Acc | F1 | $\mathcal{H}$ | ECE$\downarrow$
ALBERT (11M) | 0.7923±0.0192 | 0.8624±0.0106 | 0.4650±0.0695 | 0.0834±0.0179 | 0.6193±0.0051 | 0.4545±0.0106 | 0.3574±0.0042 | 0.0446±0.0080
ALBERT Label Smoothing | 0.7699±0.0452 | 0.8508±0.0195 | 0.5712±0.1834 | 0.0625±0.0345 | 0.6219±0.0025 | 0.4579±0.0133 | 0.3651±0.0163 | 0.0353±0.0146
ALBERT MC Dropout | 0.7893±0.0109 | 0.8583±0.0065 | 0.4575±0.0589 | 0.0854±0.0223 | 0.6152±0.0055 | 0.4448±0.0144 | 0.3697±0.0188 |  0.0345±0.0076
ALBERT w/ BNN | 0.7973±0.0011 | 0.8647±0.0006 | 0.4180±0.0015 | 0.0936±0.0016 | 0.6187±0.0009 | 0.4396±0.0015 | 0.3170±0.0003 | 0.0864±0.0011
ALBERT w/ CUE (Ours) | 0.8038±0.0005 | 0.8668±0.0004 |  0.5771±0.0004 |  0.0444±0.0031 | 0.6176±0.0021 | 0.4567±0.0046 |  0.3814±0.0294 | 0.0395±0.0098
DistilBERT (66M) | 0.7634±0.0032 | 0.8479±0.0019 | 0.5412±0.0151 | 0.0842±0.0060 | 0.6231±0.0018 | 0.4637±0.0047 | 0.3312±0.0024 | 0.0566±0.0039
DistilBERT Label Smoothing | 0.7632±0.0033 | 0.8477±0.0018 | 0.5620±0.0154 | 0.0765±0.0063 | 0.6233±0.0021 | 0.4643±0.0042 | 0.3412±0.0033 | 0.0520±0.0035
DistilBERT MC Dropout | 0.7787±0.0241 | 0.8559±0.0125 | 0.4773±0.0864 | 0.0897±0.0110 | 0.6224±0.0023 | 0.4670±0.0054 | 0.3246±0.0033 | 0.0623±0.0025
DistilBERT w/ BNN | 0.7659±0.0016 | 0.8491±0.0010 | 0.5133±0.0002 | 0.0966±0.0021 | 0.6237±0.0009 | 0.4550±0.0016 | 0.3080±0.0002 | 0.0802±0.0005
DistilBERT w/ CUE (Ours) | 0.7831±0.0012 | 0.8540±0.0010 |  0.8362±0.0008 |  0.0738±0.0029 | 0.6253±0.0017 | 0.4517±0.0022 |  0.4457±0.0004 |  0.0208±0.0031
BERT (110M) | 0.8000±0.0072 | 0.8696±0.0043 | 0.4026±0.0317 | 0.0995±0.0057 | 0.6266±0.0032 | 0.4829±0.0075 | 0.3342±0.0049 | 0.0537±0.0047
BERT Label Smoothing | 0.8036±0.0085 | 0.8717±0.0044 | 0.4099±0.0843 | 0.0967±0.0185 | 0.6268±0.0040 | 0.4819±0.0086 | 0.3444±0.0046 | 0.0492±0.0020
BERT MC Dropout | 0.8008±0.0054 | 0.8703±0.0037 | 0.4023±0.0321 | 0.0987±0.0100 | 0.6266±0.0026 | 0.4889±0.0085 | 0.3337±0.0056 | 0.0548±0.0053
BERT w/ BNN | 0.6104±0.1689 | 0.6545±0.3509 |  0.8837±0.0732 | 0.1095±0.1391 | 0.6296±0.0008 | 0.4855±0.0008 | 0.3102±0.0001 | 0.0775±0.0014
BERT w/ CUE (Ours) | 0.8123±0.0012 | 0.8762±0.0007 | 0.4991±0.0032 |  0.0677±0.0056 | 0.6282±0.0029 | 0.4712±0.0087 |  0.4433±0.0159 |  0.0326±0.0013
RoBERTa (125M) | 0.8050±0.0142 | 0.8721±0.0072 | 0.3310±0.0472 | 0.1100±0.0190 | 0.6226±0.0051 | 0.4877±0.0095 | 0.3310±0.0472 | 0.0602±0.0073
RoBERTa Label Smoothing | 0.8165±0.0128 | 0.8788±0.0058 | 0.3091±0.0415 | 0.1079±0.0120 | 0.6215±0.0021 | 0.4866±0.0098 | 0.3091±0.0415 | 0.0554±0.0084
RoBERTa MC Dropout | 0.8056±0.0056 | 0.8724±0.0023 | 0.3340±0.0452 | 0.1074±0.0165 | 0.6217±0.0034 | 0.4907±0.0082 | 0.3340±0.0452 | 0.0604±0.0060
RoBERTa w/ BNN | 0.7992±0.0022 | 0.8699±0.0014 | 0.3069±0.0010 | 0.1228±0.0021 | 0.6227±0.0002 | 0.4686±0.0056 | 0.3069±0.0010 | 0.0881±0.0002
RoBERTa w/ CUE (Ours) | 0.8075±0.0087 | 0.8744±0.0045 |  0.6077±0.0595 |  0.0465±0.0075 | 0.6255±0.0005 | 0.4540±0.0013 |  0.6077±0.0595 |  0.0316±0.0024
| Emotion | MultiNLI
Model | Acc | F1 | $\mathcal{H}$ | ECE$\downarrow$ | Acc | F1 | $\mathcal{H}$ | ECE$\downarrow$
ALBERT (11M) | 0.9284±0.0037 | 0.8862±0.0031 | 0.8862±0.0031 | 0.0348±0.0059 | 0.8362±0.0018 | 0.8358±0.0018 | 0.8358±0.0018 | 0.0465±0.0036
ALBERT Label Smoothing | 0.9310±0.0023 | 0.8897±0.0043 |  0.8897±0.0043 |  0.0231±0.0018 | 0.8327±0.0020 | 0.8317±0.0020 | 0.8317±0.0020 | 0.0364±0.0034
ALBERT MC Dropout | 0.9331±0.0033 | 0.8927±0.0036 | 0.8827±0.0036 | 0.0326±0.0032 | 0.8367±0.0013 | 0.8361±0.0019 |  0.8361±0.0019 | 0.0470±0.0039
ALBERT w/ BNN | 0.9265±0.0008 | 0.8862±0.0015 | 0.8862±0.0015 | 0.0411±0.0007 | 0.8339±0.0001 | 0.8338±0.0001 | 0.8338±0.0001 | 0.0527±0.0001
ALBERT w/ CUE (Ours) | 0.9269±0.0020 | 0.8897±0.0044 |  0.8897±0.0044 | 0.0282±0.0018 | 0.8331±0.0003 | 0.8329±0.0003 | 0.8329±0.0003 |  0.0338±0.0007
DistilBERT (66M) | 0.9287±0.0031 | 0.8886±0.0062 | 0.0441±0.0044 | 0.0393±0.0030 | 0.8067±0.0014 | 0.8059±0.0012 | 0.3737±0.0058 | 0.0376±0.0035
DistilBERT Label Smoothing | 0.9264±0.0031 | 0.8841±0.0056 | 0.0716±0.0051 | 0.0353±0.0020 | 0.8049±0.0012 | 0.8040±0.0011 | 0.3994±0.0070 | 0.0319±0.0044
DistilBERT MC Dropout | 0.9298±0.0018 | 0.8886±0.0025 | 0.0432±0.0051 | 0.0388±0.0020 | 0.8066±0.0020 | 0.8058±0.0019 | 0.3734±0.0121 | 0.0383±0.0066
DistilBERT w/ BNN | 0.9315±0.0008 | 0.8931±0.0009 | 0.0413±0.0001 | 0.0406±0.0000 | 0.8059±0.0002 | 0.8052±0.0002 | 0.3673±0.0000 | 0.0424±0.0002
DistilBERT w/ CUE (Ours) | 0.9295±0.0010 | 0.8911±0.0011 |  0.0900±0.0002 |  0.0265±0.0005 | 0.8058±0.0003 | 0.8051±0.0003 |  0.4600±0.0022 |  0.0229±0.0005
BERT (110M) | 0.9296±0.0030 | 0.8871±0.0057 | 0.0523±0.0011 | 0.0335±0.0020 | 0.8286±0.0029 | 0.8281±0.0027 | 0.3361±0.0071 | 0.0321±0.0033
BERT Label Smoothing | 0.9295±0.0042 | 0.8862±0.0074 | 0.0759±0.0024 |  0.0289±0.0061 | 0.8265±0.0016 | 0.8261±0.0017 | 0.3513±0.0116 | 0.0317±0.0043
BERT MC Dropout | 0.9285±0.0027 | 0.8872±0.0048 | 0.0527±0.0026 | 0.0325±0.0037 | 0.8287±0.0030 | 0.8281±0.0028 | 0.3363±0.0071 | 0.0315±0.0032
BERT w/ BNN | 0.9274±0.0008 | 0.8853±0.0011 | 0.0497±0.0001 | 0.0402±0.0006 | 0.3469±0.0194 | 0.1862±0.0179 |  0.9245±0.0835 | 0.1456±0.1111
BERT w/ CUE (Ours) | 0.9259±0.0009 | 0.8850±0.0015 |  0.1031±0.0082 |  0.0289±0.0043 | 0.8283±0.0005 | 0.8277±0.0005 | 0.3665±0.0030 |  0.0262±0.0021
RoBERTa (125M) | 0.9279±0.0033 | 0.8821±0.0062 | 0.0448±0.0043 | 0.0384±0.0050 | 0.8569±0.0043 | 0.8563±0.0044 | 0.2628±0.0127 | 0.0368±0.0074
RoBERTa Label Smoothing | 0.9301±0.0024 | 0.8896±0.0051 | 0.0675±0.0026 | 0.0341±0.0049 | 0.8551±0.0022 | 0.8546±0.0020 | 0.3029±0.0082 |  0.0255±0.0033
RoBERTa MC Dropout | 0.9305±0.0034 | 0.8919±0.0059 | 0.0514±0.0068 | 0.0315±0.0046 | 0.8586±0.0062 | 0.8581±0.0061 | 0.2516±0.0309 | 0.0403±0.0112
RoBERTa w/ BNN | 0.9290±0.0006 | 0.8923±0.0005 | 0.0478±0.0002 | 0.0412±0.0007 | 0.8528±0.0001 | 0.8527±0.0001 | 0.2620±0.0000 | 0.0416±0.0002
RoBERTa w/ CUE (Ours) | 0.9286±0.0005 | 0.8927±0.0008 |  0.0944±0.0004 |  0.0313±0.0033 | 0.8526±0.0003 | 0.8526±0.0003 |  0.3179±0.0003 | 0.0262±0.0004

Table 1: Results for CUE compare against baseline methods on four language
models trained on four datasets. Values shown in parentheses indicate the
model size. The highest entropy are underlined and the lowest ECE values are
in bold.

In this subsection, we discuss how to quantify the prediction uncertainty that
caused by the input features based on the latent space perturbation in §4.1.
The discussion is built on an inner product space defined by our noise
generation methods. During the inference stage, we identify the possible
feature that caused predictive uncertainty by our proposed Uncertain Feature
Identification (UFI) algorithm222We provide the UFI algorithm implementation
in Supplementary Material Section 2.. For a given input, we can retrieve three
different representations from the CUE framework, the original PLM-encoded
representation $\bm{e}_{i}$, the reconstructed representation
$\bm{e}^{\prime}_{i}$, and the difference between two representation, $\Delta
e_{i}=\bm{e}^{\prime}_{i}-\bm{e}_{i}$. The reconstructed representation
$\bm{e}^{\prime}_{i}$ can be rewritten as the weighted sum of each latent
dimension from $\bm{z}_{i}$, where the weight is given by the decoder:

$\small\bm{e}^{\prime}_{i}=\Sigma_{d=1}^{\rm
dim}{q_{\phi}({\bm{z}_{i}}_{d}|\bm{e}_{i})}\cdot{\bm{r}_{z_{i}}}_{d},\quad{\bm{r}_{z_{i}}}_{d}=\mu_{\theta}({\bm{z}_{i}}_{d}),$
(10)

where $\rm dim$ is the size of the latent space and ${\bm{r}_{z_{i}}}_{d}$
denotes the representation generated via the $d$-th dimension’s code
corresponding to the latent vector $\bm{z}_{i}$ from the decoder. As mentioned
in §4.1, $\Delta\mathcal{H}$ is proportional to the reconstruction error
$||\bm{e}^{\prime}_{i}-\bm{e}_{i}||^{2}$. We thus use the norm (calculated as
the inner product) of the reconstruction error,
$||\Delta\bm{e}_{i}||^{2}=\langle\Delta\bm{e}_{i},\Delta\bm{e}_{i}\rangle$, to
measure the entropy change as:

$\displaystyle\langle\Delta\bm{e}_{i},\Delta\bm{e}_{i}\rangle=$
$\displaystyle\langle\Sigma_{d=1}^{\rm
dim}{q_{\phi}({\bm{z}_{i}}_{d}|\bm{e}_{i})}\cdot{\bm{r}_{z_{i}}}_{d}-\bm{e}_{i},\Delta\bm{e}_{i}\rangle$
$\displaystyle=$ $\displaystyle\langle\Sigma_{d=1}^{\rm
dim}{q_{\phi}({\bm{z}_{i}}_{d}|\bm{e}_{i})}\cdot{\bm{r}_{z_{i}}}_{d},\Delta\bm{e}_{i}\rangle-\langle\bm{e}_{i},\Delta\bm{e}_{i}\rangle,$
(11)

where $\langle\cdot,\cdot\rangle$ denotes the inner product. In the first line
of Eq. (4.2), we substitute the first $\Delta\bm{e}_{i}$ with
$\bm{e}^{\prime}_{i}-\bm{e}_{i}$, and further substitute $\bm{e}^{\prime}_{i}$
with Eq. (10). When determining the relative importance of each latent
dimension with respect to the predictive entropy change,
$\langle\bm{e}_{i},\Delta\bm{e}_{i}\rangle$ can be ignored as it is the same
for all latent dimensions. Therefore, the inner product of
$\langle{\bm{r}_{z_{i}}}_{d},\Delta\bm{e}_{i}\rangle$, which dominates the
norm value of $||\Delta\bm{e}_{i}||$ in the $d$-th dimension can be used to
measure predictive uncertainty caused by each dimension from the latent space
$\bm{r}_{z_{i}}$, and thus determine each dimension’s importance.

On the other hand, the input text representation $\bm{e}_{i}$ output by the
PLM at layer-$L$ can be written as a Softmax-based weighted sum of each
token’s representation from the previous layer $L-1$ by333Note that all
representations in the RHS are from Layer $L-1$. We drop the superscript $L-1$
to simplify the notations.:

$\displaystyle\bm{e}_{i}^{L}$ $\displaystyle=\Sigma_{j=1}^{n}{\rm
Softmax}(\left\langle\bm{e}_{i},{\bm{e}_{i}}_{j}\rangle\right)\cdot{\bm{e}_{i}}_{j}$
$\displaystyle\propto\Sigma_{j=1}^{n}{\rm
exp}(\langle\bm{e}_{i},{\bm{e}_{i}}_{j}\rangle)\cdot{\bm{e}_{i}}_{j},$ (12)

where ${\bm{e}_{i}}_{j}$ denotes the representation of the $j$-th input token.
We assume that the adjacent layers in the transformer share similar
representations. Then, $\bm{e}_{i}$ at layer $L$ is:

$\displaystyle\bm{e}_{i}^{L}\propto$ $\displaystyle\Sigma_{j=1}^{n}{\rm
exp}(\langle\Sigma_{d=1}^{\rm
dim}{q_{\phi}({\bm{z}_{i}}_{d}|\bm{e}_{i})}\cdot{\bm{r}_{z_{i}}}_{d},{\bm{e}_{i}}_{j}\rangle)\cdot{\bm{e}_{i}}_{j}$
$\displaystyle\propto$ $\displaystyle\Sigma_{j=1}^{n}(\Sigma_{d=1}^{\rm
dim}{q_{\phi}({\bm{z}_{i}}_{d}|\bm{e}_{i})}\cdot\langle{\bm{r}_{z_{i}}}_{d},{\bm{e}_{i}}_{j}\rangle)\cdot{\bm{e}_{i}}_{j}$
(13)

Therefore, the influence on prediction uncertainty changes $\Delta\mathcal{H}$
of the $j$-th token is decided by the generative probability of the encoder
and the inner product $\langle{\bm{r}_{z_{i}}}_{d},{\bm{e}_{i}}_{j}\rangle$.
However, seeking the optimal ${\bm{z}_{i}}_{d}$, by minimizing the
reconstruction loss Eq. (5), is a typical Knapsack problem, which is an NP-
complete problem. Hence, intuitively, we use greedy search to find a locally
optimal solution by identifying the most influential latent dimensions of
$\bm{z}_{i}$ first and then estimating the influential score for each token.

(a) BERT CUE on CoLA.

(b) BERT CUE on GoEmotions.

(c) BERT CUE on Emotion.

(d) BERT CUE on MultiNLI.

Figure 3: Evaluation results by removing latent dimensions. The $x$-axis
represents the index of removed dimensions ranked by their relevance to
$\Delta\bm{e}_{i}$, smaller index number indicates higher relevance.
Histograms show the ECE scores after removing the corresponding latent
dimensions. The blue curve shows the predictive entropy. The green and red
curves show classification accuracy and F1, respectively.

##  5 Experiments

We first present the experimental setup followed by evaluation results.

#### Datasets

We evaluate our proposed framework on four datasets for _linguistic
acceptability classification_ , _natural language inference_ , and _emotion
classification_.

#### Baselines

We compare our method with three baselines: Label Smoothing [Gupta et al.,
2021], MC Dropout [Gal and Ghahramani, 2016] and Bayesian Neural Network
(BNN). Label Smoothing and MC Dropout are implemented in PLMs and directly
fine-tuned on the target datasets. The BNN works as a plug-in component, same
as CUE, for which the base PLM encoding and the classification layer are
firstly fine-tuned and then parameters are frozen for the plug-in layer
training.

Evaluation Metrics Accuracy (Acc), macro-averaged F1 (F1), average entropy
($\mathcal{H}$), and Expected Calibration Error (ECE) are used as metrics for
classification performance, uncertainty and model calibration measurement.

More details on dataset statistics, baseline setup, evaluation metrics and
hyperparameter settings are in Supplementary Material Section 3.

Examples |  Predicted |  True
---|---|---
GoEmotions
Despite having lived here for 10 years, I’ve never been to portillos, and given this, it’s _somewhat unlikely_ I start going now… |  Disapproval 0.42$\rightarrow$ 0.35 |  Neutral
Somehow I got banned for replying to a troll. The mods over there have _itchy trigger fingers_. |  Disappointment
0.28 $\rightarrow$ 0.19 |  Disapproval
Boundaries. Have some boundaries. _Say no_. _Don’t go_. This is frustrating to read, honestly. Don’t do things that you hate doing. |  Disgust
0.34 $\rightarrow$ 0.25 |  Fear
Emotion
I were to _go overseas or cross the border_ then I become a foreigner and will feel that way but never in my _beloved land_. |  Joy
0.51 $\rightarrow$ 0.43 |  Love
I started feeling a little stressed about leaving on time and making sure we _got the getting ready pictures_ I wanted but everything seemed to work out perfectly. |  Sadness
0.59 $\rightarrow$ 0.40 |  Anger
I _wont_ lie this week has been _abit of a difficult_ week for me ive been feeling very stressed and anxious this week plus i think im coming down with the flu but it _has definately_ helped me to appreciate the little things. |  Sadness
0.64 $\rightarrow$ 0.44 |  Anger
MultiNLI
P: There are no shares of a stock that might someday come back, just piles of options _as worthless as_ those shares of cook’s american business alliance. |  Neutral
0.43 $\rightarrow$ 0.40 |  Contradiction
H: Cook’s american business alliance caused shares of stock to come back. |  |
P: Until all _members_ of our _society_ are afforded that access, this promise of our government will continue to be _unfulfilled_. |  Entailment
0.48 $\rightarrow$ 0.43 |  Neutral
H: the government is flawed and _unfulfilled_. |  |
P: Here you‘ll find the finest leather goods and of - the - moment fashions from all the predictable high - _priests_ (valentino, armani, versace, gucci, missoni, etc.). A number of _classic_ men’s clothing meccas such as cucci (with a c), brioni, and battistoni are still _going strong_. |  Contradiction
0.35 $\rightarrow$ 0.33 |  Entailment
H: You will find only the highest quality goods, be they _high - fashion icons_ or top - _notch designer clothing_ here. |  |

Table 2: Visualisation of token-level uncertainty interpretation. The
‘Predicted’ column denotes the incorrect predictions made by the original
fine-tuned model. Values below each predicted label denotes the predictive
probability change after applying our framework. The ‘True’ column denotes the
gold-standard class labels. Italic text highlighted in blue are word tokens
identified by UFI Algorithm that cause predictive uncertainty.

###  5.1 Overall Comparison

Table 1 presents the performance of methods with four state-of-the-art PLMs,
namely, BERT [Devlin et al., 2019], ALBERT [Lan et al., 2020], DistilBERT
[Sanh et al., 2019] and RoBERTa [Liu et al., 2020b], as backbones. Our
framework with a plug-in CUE module obtains the lowest ECE scores and highest
average predictive entropy on all tasks and with different base model choices
while maintaining a comparable level of Acc/F1 scores as the original model.
Although the BNN model achieves the highest entropy with BERT on MultiNLI and
CoLA dataset, we can observe a significant drop in its Acc/F1 scores. This
indicates the BNN encoder hardly generates reliable perturbations that
maintain predicted labels unchanged. Interestingly, while the classification
performance of all compared models shows slight degradation with the injection
of uncertainty into the PLMs, our framework achieves steady accuracy gains on
the CoLA dataset.

###  5.2 Effectiveness of the Uncertainty Feature Identification Algorithm

#### Results with Latent Dimension Removal

As presented in §4.2, we can use the CUE’s reconstruction difference
$\Delta{\bm{e}_{i}}$, to disentangle the most influential latent dimensions
${\bm{z}_{i}}_{d}$s which cause predictive uncertainty. Since each latent
dimension is associated with an influential score, we can sort the latent
dimensions accordingly. We speculate that by removing latent dimensions with
higher influence scores, we should be able to observe a reduction in
predictive uncertainty. As shown in Figure 3, we visualise the evaluation
results by removing latent dimensions from $\bm{z}_{i}$ according to their
relevance to $\Delta\bm{e}_{i}$ (the rank is shown on the $x$-axis) on BERT
models. In our experiments, the latent vector $\bm{z}_{i}$ has 100 dimensions,
we thus sort them into 10 bins in descending order based on their influential
scores. In practice, the latent dimension removal is achieved by assigning $0$
as the value of the dimension on $\bm{z}_{i}$ to create a modified latent
variable $\bm{z}_{i}^{\prime}$, new prediction is made with
$\hat{y}_{i}^{\prime}=g_{\bm{\eta}}(p_{\bm{\theta}}(\bm{z}_{i}^{\prime}))$.

We can observe a remarkable increasing trend of ECE (the histograms) and
average entropy (the blue curve) when removing the most influential latent
dimensions of $\bm{z}_{i}$ on GoEmotions, Emotion and MultiNLI datasets. This
indicates the top-ranked dimensions (smaller index number) contribute more to
increase the predictive uncertainty and reducing the overconfidence
prediction, while lower-ranked dimensions have less effect. Therefore, we can
select the appropriate threshold for each dataset during token-level
uncertainty identification by observing the trend of ECE and entropy. We also
notice that across all datasets, removing any latent dimensions does not
affect much the classification accuracy (the green curve) and macro-F1 (the
red curve). However, we did not observe a similar trend of ECE and entropy on
the CoLA dataset. We suspect this is due to a relatively simple setup in CoLA
as it is only a binary classification task. For datasets with more classes,
such as GoEmotion with 27 classes, the trend of ECE with latent dimension
removal becomes more obvious. We also performed the same analysis and observed
similar phenomena on other PLMs, DistilBERT, ALBERT and RoBERTa, in
Supplementary Material Section 4.1.

#### Case Study of Token-Level Uncertainty Identification

In this subsection, we demonstrate the effectiveness of our uncertainty
identification algorithm by visualising the tokens that our framework finds
contributing to predictive uncertainty. We present several examples in which
our framework reduces overconfident predictions in Table 2. We only show the
results with BERT as the base model due to page limits. Tokens coloured in
blue are the influential tokens444For words split into subword tokens, we take
the average importance score of the constituent subword tokens. identified by
the UFI Algorithm.

For emotion classification, we found classifiers tend to be confused by idioms
or phrases carrying emotions different from the true emotion labels. For
example, the second sentence in GoEmotion contains a metaphorical phrase, ‘
_itchy trigger finger_ ’, making it a tricky case for emotion classification.
We conducted additional experiments in which we substituted the phrase "itchy
trigger finger" with either the [MASK] token or commonly used words to express
the same meaning. In both cases, the model uncertainty is reduced by replacing
the original phrase with the mask tokens leading to label switching. Replacing
the identified phrase with more commonly-used words increases the predictive
probability and leads to a more confident prediction. These results verify the
validity of our approach for identifying words/phrases causing predictive
uncertainties. The first and last sentences in GoEmotion and also the last
sentence in Emotion contain phrases which are somewhat more closely related to
the incorrectly predicted labels than the true labels, confusing the
classifier to generate wrong predictions. Tokenisation may also cause a
problem. For example, the word ‘ _beloved_ ’ in the first sentence in Emotion
is split into three parts after tokenisation, making it difficult for the
classifier to recognise the ‘Love’ emotion. For the natural language inference
task, we found classifiers tend to make overconfident predictions when the
same words are found in both premise and hypothesis. For examples, the second
instance in MultiNLI has the word ‘ _unfulfilled_ ’ in both its premise and
hypothesis. This leads to the wrong prediction of ‘Entailment’. In the last
instance, the classifier misunderstood that the ‘ _classic men’s clothing_ ’
contradicts with ‘ _high - fashion icons_ ’ and thus failed to recognise the
‘Entailment’ relation. Nevertheless, in all these cases, our proposed
framework managed to increase the predictive entropy by reducing the
confidence of predictions, alleviating the overconfidence problem.

We provide further experimental results and the ablation study, including
stability of various additional training loss terms and latent space
orthogonality in Supplementary Material Section 4.2.

##  6 Conclusion

In this paper, we have proposed a new framework CUE for uncertainty
interpretation of PLM classifiers. By comparing our method with previous
solutions, we show that CUE can achieve lower expected calibration errors
across four datasets. In some cases, it can also mitigate the confidence of
previously wrong predictions. Further experiments and case studies demonstrate
CUE is effective in identifying tokens/latent dimensions that could
potentially cause predictive uncertainty. Our work sheds light on a new
direction of uncertainty interpretation for PLMs in various NLP tasks.

###### Acknowledgements.

This work was supported in part by the UK Engineering and Physical Sciences
Research Council (grant no. EP/T017112/2, EP/V048597/1, EP/X019063/1). YH is
supported by a Turing AI Fellowship funded by the UK Research and Innovation
(grant no. EP/V020579/2). The authors would like to thank Yuxiang Zhou, Hanqi
Yan and Xingwei Tan for their invaluable feedback on this paper.

## References

  * Blundell et al. [2015] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra.  Weight uncertainty in neural network.  In _Proceedings of the 32nd ICML_ , pages 1613–1622. PMLR, 2015.
  * Brunner et al. [2020] G. Brunner, Y. Liu, D. Pascual, O. Richter, M. Ciaramita, and R. Wattenhofer.  On identifiability in transformers.  In _ICLR_ , 2020.
  * Card et al. [2018] D. Card, C. Tan, and N. A. Smith.  Neural models for documents with metadata.  In _Proceedings of the 56th Annual Meeting of the ACL_ , pages 2031–2040, July 2018.  10.18653/v1/P18-1189.
  * Cervera et al. [2021] M. R. Cervera, R. Dätwyler, F. D’Angelo, H. Keurti, B. F. Grewe, and C. Henning.  Uncertainty estimation under model misspecification in neural network regression.  _CoRR_ , abs/2111.11763, 2021.
  * Clark et al. [2019] K. Clark, U. Khandelwal, O. Levy, and C. D. Manning.  What does BERT look at? An Analysis of BERT’s Attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 276–286, Aug. 2019.  10.18653/v1/W19-4828.
  * Desai and Durrett [2020] S. Desai and G. Durrett.  Calibration of pre-trained transformers.  In _Proceedings of the 2020 Conference on EMNLP_ , pages 295–302, Nov. 2020.  10.18653/v1/2020.emnlp-main.21.
  * Devlin et al. [2019] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the NACCL_ , pages 4171–4186, June 2019.  10.18653/v1/N19-1423.
  * Gal and Ghahramani [2016] Y. Gal and Z. Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _Proceedings of The 33rd ICML_ , pages 1050–1059. PMLR, 2016.
  * Guo et al. [2017] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger.  On calibration of modern neural networks.  In _Proceedings of the 34th ICML_ , pages 1321–1330. PMLR, 2017.
  * Gupta et al. [2021] A. Gupta, G. Kvernadze, and V. Srikumar.  Bert & family eat word salad: Experiments with text understanding.  _AAAI_ , 35(14):12946–12954, May 2021.  10.1609/aaai.v35i14.17531.
  * He et al. [2020] J. He, X. Zhang, S. Lei, Z. Chen, F. Chen, A. Alhamadani, B. Xiao, and C. Lu.  Towards more accurate uncertainty estimation in text classification.  In _Proceedings of the 2020 Conference on EMNLP_ , pages 8362–8372, Nov. 2020.  10.18653/v1/2020.emnlp-main.671.
  * Heiss et al. [2023] J. Heiss, J. Weissteiner, H. Wutte, S. Seuken, and J. Teichmann.  Nomu: Neural optimization-based model uncertainty, 2023.
  * Kim et al. [2020] S. Kim, J. Yi, E. Kim, and S. Yoon.  Interpretation of NLP models through input marginalization.  In _Proceedings of the 2020 Conference on EMNLP_ , pages 3154–3167, Nov. 2020.  10.18653/v1/2020.emnlp-main.255.
  * Kingma and Welling [2014] D. P. Kingma and M. Welling.  Auto-encoding variational bayes.  In Y. Bengio and Y. LeCun, editors, _2nd ICLR, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.
  * Lakshminarayanan et al. [2017] B. Lakshminarayanan, A. Pritzel, and C. Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _NeurIPS_ , volume 30. Curran Associates, Inc., 2017.
  * Lan et al. [2020] Z. Lan, M. Chen, S. Goodman, K. Gimpel, P. Sharma, and R. Soricut.  Albert: A lite bert for self-supervised learning of language representations.  In _ICLR_ , 2020.
  * Liu et al. [2020a] J. Liu, Z. Lin, S. Padhy, D. Tran, T. Bedrax Weiss, and B. Lakshminarayanan.  Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.  In _NeurIPS_ , volume 33, pages 7498–7512, 2020a.
  * Liu et al. [2020b] Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov.  RoBERTa: A robustly optimized BERT pretraining approach, 2020b.
  * Malinin and Gales [2021] A. Malinin and M. Gales.  Uncertainty estimation in autoregressive structured prediction.  In _ICLR_ , 2021.
  * Malinin and Gales [2018] A. Malinin and M. J. F. Gales.  Predictive uncertainty estimation via prior networks.  In _NeurIPS_ , 2018.
  * Mareček and Rosa [2019] D. Mareček and R. Rosa.  From balustrades to pierre vinken: Looking for syntax in transformer self-attentions.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 263–275, Aug. 2019.  10.18653/v1/W19-4827.
  * Qiu et al. [2020] X. Qiu, E. Meyerson, and R. Miikkulainen.  Quantifying point-prediction uncertainty in neural networks via residual estimation with an i/o kernel.  In _ICLR_ , 2020.
  * Sanh et al. [2019] V. Sanh, L. Debut, J. Chaumond, and T. Wolf.  Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter.  _CoRR_ , abs/1910.01108, 2019.
  * Shelmanov et al. [2021] A. Shelmanov, E. Tsymbalov, D. Puzyrev, K. Fedyanin, A. Panchenko, and M. Panov.  How certain is your Transformer?  In _Proceedings of the 16th Conference of the EACL_ , pages 1833–1840, Apr. 2021.  10.18653/v1/2021.eacl-main.157.
  * Shen et al. [2020] Y. Shen, W. Chen, and H. Jin.  Modeling token-level uncertainty to learn unknown concepts in SLU via calibrated dirichlet prior RNN.  _CoRR_ , abs/2010.08101, 2020.
  * Sullivan [2015] T. J. Sullivan.  _Introduction to uncertainty quantification_ , volume 63.  Springer, 2015.
  * Sun and Marasović [2021] K. Sun and A. Marasović.  Effective attention sheds light on interpretability.  In _Findings of the ACL-IJCNLP 2021_ , pages 4126–4135, Aug. 2021\.  10.18653/v1/2021.findings-acl.361.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin.  Attention is all you need.  In _NeurIPS_ , volume 30. Curran Associates, Inc., 2017.
  * Vazhentsev et al. [2022] A. Vazhentsev, G. Kuzmin, A. Shelmanov, A. Tsvigun, E. Tsymbalov, K. Fedyanin, M. Panov, A. Panchenko, G. Gusev, M. Burtsev, M. Avetisian, and L. Zhukov.  Uncertainty estimation of transformer predictions for misclassification detection.  In _Proceedings of the 60th Annual Meeting of the ACL_ , pages 8237–8252, May 2022.
  * Vulić et al. [2020] I. Vulić, E. M. Ponti, R. Litschko, G. Glavaš, and A. Korhonen.  Probing pretrained language models for lexical semantics.  In _Proceedings of the 2020 Conference on EMNLP_ , pages 7222–7240, Nov. 2020.  10.18653/v1/2020.emnlp-main.586.
  * Xiao and Wang [2019] Y. Xiao and W. Y. Wang.  Quantifying uncertainties in natural language processing tasks.  In _Proceedings of the 33rd AAAI_. AAAI Press, 2019.  ISBN 978-1-57735-809-1.  10.1609/aaai.v33i01.33017322.
  * Yang et al. [2021] L. Yang, J. Li, P. Cunningham, Y. Zhang, B. Smyth, and R. Dong.  Exploring the efficacy of automatically generated counterfactuals for sentiment analysis.  In _ACL 2021_ , pages 306–316, Online, Aug. 2021. Association for Computational Linguistics.  10.18653/v1/2021.acl-long.26.
  * Zhou and Srikumar [2021] Y. Zhou and V. Srikumar.  DirectProbe: Studying representations without classifiers.  In _Proceedings of the 2021 Conference of NAACL_ , pages 5070–5083, June 2021.  10.18653/v1/2021.naacl-main.401.


$\left|\left|\left[\widehat{\Delta}(x)\right]_{S_{\ell}(x)^{c}}\right|\right|_{1}\leq\frac{r_{\ell}^{2}\left|\left|x\right|\right|_{2}^{2}}{(a_{0}-1)\kappa_{R}^{4}}\cdot\frac{\nu}{\gamma}+\left(\frac{a_{0}+1}{a_{0}-1}\right)\left|\left|\left[\widehat{\Delta}(x)\right]_{S_{\ell}(x)}\right|\right|_{1}.$

On the one hand, if
$\frac{r_{\ell}^{2}\left|\left|x\right|\right|_{2}^{2}}{(a_{0}-1)\kappa_{R}^{4}}\cdot\frac{\nu}{\gamma}\leq\left(\frac{a_{0}+1}{a_{0}-1}\right)\sum_{k\in
S_{\ell}(x)}\left|\widehat{\ell}_{k}(x)-\widetilde{\ell}_{k}(x)\right|$, then

$\left|\left|\left[\widehat{\Delta}(x)\right]_{S_{\ell}(x)^{c}}\right|\right|_{1}\leq
2\left(\frac{a_{0}+1}{a_{0}-1}\right)\left|\left|\left[\widehat{\Delta}(x)\right]_{S_{\ell}(x)}\right|\right|_{1}\quad\text{
so that
}\quad\widehat{\ell}(x)-\widetilde{\ell}(x)\in\mathcal{C}_{1}\left(S_{\ell}(x),2\left(\frac{a_{0}+1}{a_{0}-1}\right)\right).$

On the other hand, if
$\frac{r_{\ell}^{2}\left|\left|x\right|\right|_{2}^{2}}{(a_{0}-1)\kappa_{R}^{4}}\cdot\frac{\nu}{\gamma}>\left(\frac{a_{0}+1}{a_{0}-1}\right)\sum_{k\in
S_{\ell}(x)}\left|\widehat{\ell}_{k}(x)-\widetilde{\ell}_{k}(x)\right|$, then

$\left|\left|\widehat{\ell}(x)-\widetilde{\ell}(x)\right|\right|_{1}\leq\frac{(a_{0}^{2}+1)r_{\ell}^{2}\left|\left|x\right|\right|_{2}^{2}}{(a_{0}-1)\kappa_{R}^{4}}\cdot\frac{\nu}{\gamma}.$

Thus, we conclude that

$\widehat{\ell}(x)-\widetilde{\ell}(x)\in\mathcal{C}_{1}\left(S_{\ell}(x),2\left(\frac{a_{0}+1}{a_{0}-1}\right)\right)\bigcup
B_{d}^{(1)}\left(0,\,\frac{(a_{0}^{2}+1)r_{\ell}^{2}\left|\left|x\right|\right|_{2}^{2}}{(a_{0}-1)\kappa_{R}^{4}}\cdot\frac{\nu}{\gamma}\right).$

The proof is completed. ∎

We know from (b) of Theorem 2 that

$\displaystyle\frac{\gamma}{a_{0}n}$ $\displaystyle\asymp
r_{\gamma}(n,\delta)$
$\displaystyle=\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}\cdot
r_{\pi}$

will satisfy the bound in (58) with probability at least $1-\delta$. The
following lemma suggests a (probabilistic) lower bound for
$\frac{\nu}{a_{0}n}$ in (58) using the similar arguments. It is worth noting
that the lower bound for $\frac{\gamma}{a_{0}n}$ can converge to 0 as
$n\to\infty$ under certain conditions, while the lower bound for
$\frac{\nu}{a_{0}n}$ is asymptotically non-vanishing.

######  Lemma 10.

Let $\delta\in(0,1)$ and $x\in\mathbb{R}^{d}$ be a fixed covariate vector.
Suppose that Assumption 2 holds. Then, with probability at least $1-\delta$,

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\pi_{i}\left[X_{i}^{T}v\right]^{2}\right|$
$\displaystyle\lesssim(2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}\left[1+\sqrt{\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}}+\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}\right].$

If, in addition, Assumption 5 holds, then with probability at least
$1-\delta$,

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}\left[X_{i}^{T}v\right]^{2}\right|$
$\displaystyle\lesssim\left[1+r_{\pi}(n,\delta)\right](2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}\left[1+\sqrt{\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}}+\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}\right].$

###### Proof of Lemma 10.

We only prove the second statement, as the proof of the first statement is
already contained in the following proof. By Assumption 5, we know that with
probability $1-\delta$,

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}\left[X_{i}^{T}v\right]^{2}\right|$
$\displaystyle\leq\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\pi_{i}\left[X_{i}^{T}v\right]^{2}\right|+\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\left(\widehat{\pi}_{i}-\pi_{i}\right)\left[X_{i}^{T}v\right]^{2}\right|$
$\displaystyle\leq\left[1+r_{\pi}(n,\delta)\right]\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\left(X_{i}^{T}v\right)^{2}\right|$
$\displaystyle\leq\left[1+r_{\pi}(n,\delta)\right]\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left[\left|\frac{1}{n}\sum_{i=1}^{n}\left\\{\left(X_{i}^{T}v\right)^{2}-\mathrm{E}\left[\left(X_{i}^{T}v\right)^{2}\right]\right\\}\right|+\left|v^{T}\mathrm{E}\left[X_{1}X_{1}^{T}\right]v\right|\right].$

By Lemma 7, we know that

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\left\\{\left(X_{i}^{T}v\right)^{2}-\mathrm{E}\left[\left(X_{i}^{T}v\right)^{2}\right]\right\\}\right|$
$\displaystyle\lesssim(2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}\left[\sqrt{\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}}+\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}\right]$

with probability at least $1-\delta$. In addition, we know from Lemma 4 that
there exists a discrete set $M_{d,\ell}\subset B_{d}(0,1)$ such that
$\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)\subset 2(2+\vartheta_{\ell})\cdot\mathrm{conv}(M_{d,\ell})$ with
$\left|\left|v\right|\right|_{0}\leq s_{\ell}(x)$ for all $v\in M_{d,\ell}$.
By Lemma 5, we obtain that

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|v^{T}\mathrm{E}\left[X_{1}X_{1}^{T}\right]v\right|$
$\displaystyle\leq\sup_{v\in
2(2+\vartheta_{\ell})\mathrm{conv}(M_{d,\ell})}\left|v^{T}\mathrm{E}\left[X_{1}X_{1}^{T}\right]v\right|$
$\displaystyle=\sup_{v\in 2(2+\vartheta_{\ell})\cdot
M_{d,\ell}}\left|v^{T}\mathrm{E}\left[X_{1}X_{1}^{T}\right]v\right|$
$\displaystyle\leq 4(2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}.$

Therefore, we conclude that

$\displaystyle\sup_{v\in\mathcal{C}_{2}\left(S_{\ell}(x),\vartheta_{\ell}\right)\cap
B_{d}(0,1)}\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}\left[X_{i}^{T}v\right]^{2}\right|$
$\displaystyle\lesssim\left[1+r_{\pi}(n,\delta)\right](2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}\left[\sqrt{\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}}+\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}\right]$
$\displaystyle\quad+4(2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}$
$\displaystyle\lesssim\left[1+r_{\pi}(n,\delta)\right](2+\vartheta_{\ell})^{2}\phi_{s_{\ell}(x)}\left[1+\sqrt{\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}}+\frac{\log(\frac{1}{\delta})+s_{\ell}(x)\log\left(\frac{ed}{s_{\ell}(x)}\right)}{n}\right]$

with probability at least $1-\delta$. The result follows. ∎

######  Lemma 11.

Let $\delta\in(0,1)$ and $x$ be a fixed covariate vector. Suppose that
Assumptions 1, 2 and 4 hold. Then, with probability at least $1-\delta$,

$\left|\left|\frac{1}{n}\sum_{i=1}^{n}\widetilde{Z}_{i}X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}\lesssim\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}\left[\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right],$

where the above statement holds when $\widetilde{Z}_{i}=\pi(X_{i}),i=1,...,n$
or $\widetilde{Z}_{i}=R_{i},i=1,...,n$. If, in addition, Assumption 5 holds,
then with probability at least $1-\delta$,

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}\left[\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]$
$\displaystyle\quad+\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}\left[1+\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]r_{\pi}(n,\delta).$

###### Proof of Lemma 11.

We only prove the second statement, as the proof of the first statement is
already contained in the following proof. Notice that

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}$
$\displaystyle\leq\underbrace{\left|\left|\frac{1}{n}\sum_{i=1}^{n}\pi_{i}X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}}_{\textbf{Term
I}}+\underbrace{\left|\left|\frac{1}{n}\sum_{i=1}^{n}\left(\widehat{\pi}_{i}-\pi_{i}\right)X_{i}X_{i}^{T}\ell_{0}(x)\right|\right|_{\infty}}_{\textbf{Term
II}}.$

$\bullet$ Term I: Recall from (18) that
$\ell_{0}(x):=-2\left[\mathrm{E}\left(RXX^{T}\right)\right]^{-1}x=-2\left[\mathrm{E}\left(\pi(X)XX^{T}\right)\right]^{-1}x$
and thus,

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\pi(X_{i})X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}$
$\displaystyle=\max_{1\leq k\leq
d}\left|-\frac{2}{n}\sum_{i=1}^{n}\pi(X_{i})X_{ik}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x+2x_{k}\right|,$

and all the random variables

$2x_{1}-\frac{2}{n}\sum_{i=1}^{n}\pi(X_{i})\cdot
X_{i1}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x,\dots,2x_{d}-\frac{2}{n}\sum_{i=1}^{n}\pi(X_{i})\cdot
X_{id}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x$

have zero means and sub-exponential tails (but neither identically distributed
nor independent). By Assumption 2 and Lemma 12, we know that

$\displaystyle\max_{1\leq k\leq d}\left|\left|x_{k}-\pi(X_{i})\cdot
X_{ik}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x\right|\right|_{\psi_{1}}$
$\displaystyle\lesssim\max_{1\leq k\leq d}\left|\left|\pi(X_{i})\cdot
X_{ik}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x\right|\right|_{\psi_{1}}$
$\displaystyle\lesssim\max_{1\leq k\leq d}\max_{1\leq i\leq
n}\left|\left|X_{ik}\right|\right|_{\psi_{2}}\left|\left|\pi(X_{i})X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x\right|\right|_{\psi_{2}}$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}$

for all $1\leq i\leq n$. Therefore, by the union bound and Bernstein’s
inequality (Corollary 2.8.3 in Vershynin 2018), we know that for all $t>0$,
there exists an absolute constant $A_{1}>0$ such that

$\displaystyle\mathrm{P}\left(\left|\left|\frac{1}{n}\sum_{i=1}^{n}\pi(X_{i})X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}>t\right)$
$\displaystyle\leq d\cdot\max_{1\leq k\leq
d}\mathrm{P}\left(\left|\frac{2}{n}\sum_{i=1}^{n}\left[\pi(X_{i})\cdot
X_{ik}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x-x_{k}\right]\right|>t\right)$
$\displaystyle\leq
2d\cdot\exp\left(-nA_{1}\min\left\\{\frac{t\kappa_{R}}{\sqrt{\phi_{1}}\left|\left|x\right|\right|_{2}},\,\frac{t^{2}\kappa_{R}^{2}}{\phi_{1}\left|\left|x\right|\right|_{2}^{2}}\right\\}\right).$

Hence, we conclude that with probability at least $1-\delta$,

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\pi(X_{i})X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}\left[\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right].$

$\bullet$ Term II: Note that under Assumption 5,

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\left(\widehat{\pi}_{i}-\pi_{i}\right)X_{i}X_{i}^{T}\ell_{0}(x)\right|\right|_{\infty}$
$\displaystyle=\max_{1\leq k\leq
d}\left|\frac{1}{n}\sum_{i=1}^{n}\left(\widehat{\pi}_{i}-\pi_{i}\right)X_{ik}X_{i}^{T}\ell_{0}(x)\right|$
$\displaystyle\lesssim\max_{1\leq i\leq
n}\left|\widehat{\pi}_{i}-\pi_{i}\right|\cdot\max_{1\leq k\leq
d}\left(\frac{1}{n}\sum_{i=1}^{n}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\right)$
$\displaystyle\lesssim r_{\pi}(n,\delta)\left[\max_{1\leq k\leq
d}\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|+\max_{1\leq k\leq
d}\left|\frac{1}{n}\sum_{i=1}^{n}\Big{(}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\Big{)}\right|\right]$

with probability at least $1-\delta$. By Lemma 8, we know that
$\ell_{0}(x)\in\mathcal{C}_{2}(S_{\ell}(x),\vartheta_{\ell})$ with
$\vartheta_{\ell}=\frac{r_{\ell}}{\sqrt{1-r_{\ell}^{2}}}$. Hence, by Cauchy-
Schwarz inequality,

$\displaystyle\max_{1\leq k\leq
d}\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|$
$\displaystyle\leq\max_{1\leq k\leq
d}\sqrt{\mathrm{E}X_{ik}^{2}}\cdot\sqrt{\mathrm{E}\left[(X_{i}^{T}\ell_{0}(x))^{2}\right]}$
$\displaystyle\lesssim\sqrt{\phi_{1}}\cdot\sqrt{\sup_{v\in\mathcal{C}_{2}(S_{\ell}(x),\vartheta_{\ell})\cap\mathbb{S}^{d-1}}\mathrm{E}\left[(X_{i}^{T}v)^{2}\right]}\cdot\left|\left|\ell_{0}(x)\right|\right|_{2}$
$\displaystyle\stackrel{{\scriptstyle\text{(i)}}}{{\lesssim}}\sqrt{\phi_{1}}\cdot\sqrt{\sup_{v\in
2(2+\vartheta_{\ell})\cdot\mathrm{conv}(M_{d,\ell})}\mathrm{E}\left[(X_{i}^{T}v)^{2}\right]}\cdot\frac{\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}$
$\displaystyle\stackrel{{\scriptstyle\text{(ii)}}}{{\asymp}}\sqrt{\phi_{1}}\cdot\sqrt{\sup_{v\in
M_{d,\ell}}\mathrm{E}\left[(X_{i}^{T}v)^{2}\right]}\cdot\frac{\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}$
$\displaystyle\lesssim\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\cdot\frac{\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}},$

where we argue by Lemma 4 in the (asymptotic) inequality (i) that there exists
a discrete set $M_{d,\ell}\subset B_{d}(0,1)$ with cardinality
$|M_{d,\ell}|\leq\frac{3}{2}\left(\frac{5ed}{s_{\ell}(x)}\right)^{s_{\ell}(x)}$
such that $\mathcal{C}_{1}(S_{\ell}(x),3)\cap\mathbb{S}^{d-1}\subset
10\cdot\mathrm{conv}(M_{d,\ell})$ and $\left|\left|u\right|\right|_{0}\leq
s_{\ell}(x)$ for all $u\in M_{d,\ell}$. Also, in the asymptotic equality (ii),
we leverage Lemma 5 to show that the supremum of a convex function on
$\mathrm{conv}(M_{d,\ell})$ is attained at its vertices, which are contained
in $M_{d,\ell}$.

Now, we also know that all the random variables

$\frac{1}{n}\sum_{i=1}^{n}\Big{[}\left|X_{i1}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{i1}X_{i}^{T}\ell_{0}(x)\right|\Big{]},\,...,\,\frac{1}{n}\sum_{i=1}^{n}\Big{[}\left|X_{id}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{id}X_{i}^{T}\ell_{0}(x)\right|\Big{]}$

have zero means and sub-exponential tails (but neither identically distributed
nor independent). By Assumption 2 and Lemma 12, we know that for all $1\leq
i\leq n$,

$\displaystyle\max_{1\leq k\leq
d}\left|\left|\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\right|\right|_{\psi_{1}}$
$\displaystyle\lesssim\max_{1\leq k\leq
d}\left|\left|X_{ik}X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x\right|\right|_{\psi_{1}}$
$\displaystyle\lesssim\max_{1\leq k\leq d}\max_{1\leq i\leq
n}\left|\left|X_{ik}\right|\right|_{\psi_{2}}\left|\left|X_{i}^{T}\mathrm{E}\left[\pi(X_{i})X_{i}X_{i}^{T}\right]^{-1}x\right|\right|_{\psi_{2}}$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}}{\kappa_{R}^{2}}\left|\left|x\right|\right|_{2}.$

Therefore, by the union bound and Bernstein’s inequality (Corollary 2.8.3 in
Vershynin 2018), we know that there exists an absolute constant $A_{2}>0$ such
that for all $t>0$,

$\displaystyle\mathrm{P}\left(\max_{1\leq k\leq
d}\left|\frac{1}{n}\sum_{i=1}^{n}\Big{(}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\Big{)}\right|>t\right)$
$\displaystyle\leq d\cdot\max_{1\leq k\leq
d}\mathrm{P}\left(\left|\frac{1}{n}\sum_{i=1}^{n}\Big{(}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\Big{)}\right|>t\right)$
$\displaystyle\leq
2d\cdot\exp\left(-nA_{2}\min\left\\{\frac{t\kappa_{R}^{2}}{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\left|\left|x\right|\right|_{2}},\,\frac{t^{2}\kappa_{R}^{4}}{\phi_{1}\phi_{s_{\ell}(x)}\left|\left|x\right|\right|_{2}^{2}}\right\\}\right).$

Hence, we conclude that with probability at least $1-\delta$,

$\displaystyle\max_{1\leq k\leq
d}\left|\frac{1}{n}\sum_{i=1}^{n}\Big{(}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|-\mathrm{E}\left|X_{ik}X_{i}^{T}\ell_{0}(x)\right|\Big{)}\right|$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}}{\kappa_{R}^{2}}\left|\left|x\right|\right|_{2}\left[\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]$

Therefore, we conclude that with probability at least $1-\delta$,

$\displaystyle\begin{split}\textbf{Term
II}&=\left|\left|\frac{1}{n}\sum_{i=1}^{n}\left(\widehat{\pi}_{i}-\pi_{i}\right)X_{i}X_{i}^{T}\ell_{0}(x)\right|\right|_{\infty}\\\
&\lesssim\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}\left[1+\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]r_{\pi}(n,\delta).\end{split}$
(59)

Combining the bounds for Term I and Term II yields that with probability at
least $1-\delta$,

$\displaystyle\left|\left|\frac{1}{n}\sum_{i=1}^{n}\widehat{\pi}_{i}X_{i}X_{i}^{T}\ell_{0}(x)+2x\right|\right|_{\infty}$
$\displaystyle\lesssim\frac{\sqrt{\phi_{1}}}{\kappa_{R}}\left|\left|x\right|\right|_{2}\left[\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]$
$\displaystyle\quad+\frac{\sqrt{\phi_{1}\phi_{s_{\ell}(x)}}\left|\left|x\right|\right|_{2}}{\kappa_{R}^{2}}\left[1+\sqrt{\frac{\log(d/\delta)}{n}}+\frac{\log(d/\delta)}{n}\right]r_{\pi}(n,\delta).$

The proof is completed. ∎

######  Lemma 12 (Orlicz-norm of products of random variables).

Let $X_{1},...,X_{N}$ be random variables with finite $\psi_{\alpha}$-Orlicz
norm (14). Then, for any $N\geq 1$,
$\left|\left|\prod_{i=1}^{N}X_{i}\right|\right|_{\psi_{\frac{\alpha}{N}}}\leq\prod_{i=1}^{N}\left|\left|X_{i}\right|\right|_{\psi_{\alpha}}$.

###### Proof of Lemma 12.

The result generalizes Lemma 2.7.7 in Vershynin (2018), which follows from the
Young’s inequality as:
$\prod_{i=1}^{N}x_{i}^{\lambda_{i}}\leq\sum_{i=1}^{N}\lambda_{i}x_{i}$ for all
$x_{i},\lambda_{i}\geq 0,i=1,...,N$ with $\sum_{i=1}^{N}\lambda_{i}=1$.
Without the loss of generality, we assume that
$\left|\left|X_{i}\right|\right|_{\psi_{\alpha}}=1$ for all $i=1,...,N$. Then,

$\displaystyle\mathrm{E}\left[\exp\left(\prod_{i=1}^{N}|X_{i}|^{\frac{\alpha}{N}}\right)\right]\leq\mathrm{E}\left[\exp\left(\frac{1}{N}\sum_{i=1}^{N}|X_{i}|^{\alpha}\right)\right]$
$\displaystyle=\mathrm{E}\left[\prod_{i=1}^{N}\exp\left(\frac{1}{N}|X_{i}|^{\alpha}\right)\right]$
$\displaystyle\leq\frac{1}{N}\sum_{i=1}^{N}\mathrm{E}\left[\exp\left(|X_{i}|^{\alpha}\right)\right]\leq
2,$

where we apply the Young’s inequality to obtain the first and last
inequalities. ∎

######  Lemma 13 (Giessing 2022).

Let $\mathcal{F}\subset L_{1}(\mathcal{X},\mathcal{A},\mathrm{P})$ and
$\varrho$ be a pseudo-metric on $\mathcal{F}\times\mathcal{F}$. For $\eta>0$,
we define
$\mathcal{F}_{\eta}=\left\\{f-g:f,g\in\mathcal{F},\varrho(f,g)\leq\eta\right\\}$.
Suppose that for $\alpha\in(0,2]$, there exists a constant $K>0$ such that for
all $f,g\in\mathcal{F}$,

$\left|\left|\left(f-\mathrm{P}f\right)-\left(g-\mathrm{P}g\right)\right|\right|_{\mathrm{P},\psi_{\alpha}}\leq
K\cdot\varrho(f,g).$

Then, the following statements hold true:

  1. (a)

If $\alpha\in(0,1]$, then for all $t>0$, with probability at least $1-e^{-t}$,

$\displaystyle\left|\left|\mathbb{G}_{n}\right|\right|_{\mathcal{F}_{\eta}}\leq
C_{\alpha}K\left[\int_{0}^{\eta}\sqrt{\log
N(\epsilon,\mathcal{F},\varrho)}\,d\epsilon+n^{-\frac{1}{2}}\int_{0}^{\eta}\big{[}\log
N(\epsilon,\mathcal{F},\varrho)\big{]}^{\frac{1}{\alpha}}d\epsilon+\eta\left(\sqrt{t}+n^{-\frac{1}{2}}t^{\frac{1}{\alpha}}\right)\right],$

where $C_{\alpha}>0$ is a constant depending only on $\alpha$.

  2. (b)

If $\alpha\in(1,2]$, then for all $t>0$, with probability at least $1-e^{-t}$,

$\displaystyle\left|\left|\mathbb{G}_{n}\right|\right|_{\mathcal{F}_{\eta}}\leq
C_{\alpha}K\left[\int_{0}^{\eta}\sqrt{\log
N(\epsilon,\mathcal{F},\varrho)}\,d\epsilon+n^{-\frac{1}{2}+\frac{1}{\beta}}\int_{0}^{\eta}\big{[}\log
N(\epsilon,\mathcal{F},\varrho)\big{]}^{\frac{1}{\alpha}}d\epsilon+\eta\left(\sqrt{t}+n^{-\frac{1}{2}+\frac{1}{\beta}}t^{\frac{1}{\alpha}}\right)\right],$

where $C_{\alpha}>0$ is a constant depending only on $\alpha$ and
$\frac{1}{\alpha}+\frac{1}{\beta}=1$.

######  Remark 3 (Giessing and Wang 2021).

If $\mathcal{F}=\\{f_{t}:t\in T\\}$ and
$\left|\left|\left(f_{s}-\mathrm{P}f_{s}\right)-\left(f_{t}-\mathrm{P}f_{t}\right)\right|\right|_{\mathrm{P},\psi_{\alpha}}\leq
K\cdot\varrho(s,t)$ for all $f_{s},f_{t}\in\mathcal{F}$ and a pseudo-metric
$\varrho$ on $T\times T$, where $K>0$ is some constant, then we can replace
$N(\epsilon,\mathcal{F},\varrho)$ by $N(\epsilon,T,\varrho)$.