11institutetext: Department of Bioengineering, University of Pennsylvania,
Philadelphia, PA, United States 22institutetext: Penn Image Computing and
Science Laboratory, Department of Radiology, University of Pennsylvania,
Philadelphia, PA, United States
22email<EMAIL_ADDRESS>
# Domain Generalizer: A Few-shot Meta Learning Framework for Domain
Generalization in Medical Imaging

Pulkit Khandelwal 1Department of Bioengineering, University of Pennsylvania,
Philadelphia, PA, United States 12Penn Image Computing and Science Laboratory,
Department of Radiology, University of Pennsylvania, Philadelphia, PA, United
States
<EMAIL_ADDRESS>Paul Yushkevich 2Penn Image Computing and Science
Laboratory, Department of Radiology, University of Pennsylvania, Philadelphia,
PA, United States
<EMAIL_ADDRESS>1Department of Bioengineering, University of
Pennsylvania, Philadelphia, PA, United States 12Penn Image Computing and
Science Laboratory, Department of Radiology, University of Pennsylvania,
Philadelphia, PA, United States
<EMAIL_ADDRESS>2Penn Image Computing and Science Laboratory, Department
of Radiology, University of Pennsylvania, Philadelphia, PA, United States
<EMAIL_ADDRESS>
###### Abstract

Deep learning models perform best when tested on target (test) data domains
whose distribution is similar to the set of source (train) domains. However,
model generalization can be hindered when there is significant difference in
the underlying statistics between the target and source domains. In this work,
we adapt a domain generalization method based on a model-agnostic meta-
learning framework [1] to biomedical imaging. The method learns a domain-
agnostic feature representation to improve generalization of models to the
unseen test distribution. The method can be used for any imaging task, as it
does not depend on the underlying model architecture. We validate the approach
through a computed tomography (CT) vertebrae segmentation task across healthy
and pathological cases on three datasets. Next, we employ few-shot learning,
i.e. training the generalized model using very few examples from the unseen
domain, to quickly adapt the model to new unseen data distribution. Our
results suggest that the method could help generalize models across different
medical centers, image acquisition protocols, anatomies, different regions in
a given scan, healthy and diseased populations across varied imaging
modalities.

###### Keywords:

domain adaptation domain generalization meta learning vertebrae segmentation
computed tomography

##  1 Introduction and Background

In biomedical imaging, deep learning models trained on one dataset are often
hard to generalize to other related datasets. Generally, biomedical images can
be represented as points on a high-dimensional non-linear manifold. Failure of
segmentation and classification algorithms to generalize across imaging
modalities, patients, image acquisition protocols, medical centers, healthy
and diseased populations, age, etc., can be explained by significant
differences in the statistical distributions of datasets on the image
manifolds, known as _covariate shift_ [2]. Addressing covariate shift by
retraining deep learning models on each new data domain is impractical in most
applications because of the scarcity of expert labeled data. Therefore, it is
important to develop deep learning methods that generalize well to new related
datasets not seen during training using few or no annotated examples from the
new dataset. Domain adaptation [3] and domain generalization [4] paradigms aim
at reducing the covariate shift between the training and test distributions by
learning domain invariant features. Domain adaptation learns a feature
representation that is invariant to the statistics of the source and target
domains, and is discriminative enough for the actual learning task. Domain
adaptation could either be unsupervised or semi-supervised. Domain
generalization, a relatively less studied and harder problem, trains models
using a variety of source domains to learn a generic feature representation
which should perform well on unseen target domains. This flavor of transfer
learning does not use any samples from the target distribution during
training. Relatedly, few-shot learning is a paradigm which adapts a trained
model to a completely new data distribution with very limited labeled training
examples [5].

The biomedical imaging community has witnessed several applications of domain
adaptation and few-shot learning. Adaptation across different medical centers
is a known challenge in image segmentation [6], and has been achieved through
both unsupervised [7], and supervised approaches [8]. Cross-modality domain
adaptation methods between magnetic resonance (MR) and CT images have been
proposed using variational autoencoders [9] for whole heart segmentation, and
CycleGANs for segmentation of the prostate [10], [11]. Decision forests have
been employed to adapt between in-vivo and in-vitro images [12] for
intravascular ultrasound tissue segmentation. A few-shot network [13] was
proposed to segment multiple organs in MR images. However, a priori knowledge
of the unseen test domain is not always available, which hinders model
generalizability as discussed above.

Very recently, some groups explored domain generalization for biomedical
imaging. (1) A series of nine data augmentation techniques were applied to the
training domains to mimic the test distribution [14] for heart ultrasound,
heart and prostate MR image segmentation. (2) An episodic training-based meta-
learning method [15] was applied to segment brain tissue in T1-weighted MRI
across four medical centers. (3) A variational auto-encoder [16] was used to
learn three latent subspaces to generalize across patients for a 2D cell
segmentation task via domain disentanglement. These methods have certain
limitations respectively: (1) In [14], the training data is very large and
heavily augmented, which might not be the general case for many problems in
medical imaging where the goal is to extract enough information from very
limited data for domain generalization; the method is not tested on diseased
populations which could vary significantly in anatomies and shapes from a
generic healthy population. (2) In [15], again, the training set contains
similar anatomies in both the train and test sets; the average performance is
quite marginal than the compared baseline, with an improvement of only around
0.8% in Dice score (Baseline: 90.6% vs proposed: 91.4%); does not evaluate on
cases with atrophied or irregular brain anatomies. (3) In [16], the method
uses domain labels as an additional cue, which we argue is difficult to define
precisely, due to its wide range of interpretation; the average performance is
not significant with an improvement of 0.4% in average Dice score (Baseline:
95.4% vs proposed: 95.8%); the method evaluates only one 2D dataset, not
extending to the 3D case; each patient is considered as different domain
acquired at the same medical center, and hence the test and train sets might
have similar statistical distributions.

Contributions. In the present work, we extend a gradient-based meta-learning
domain generalization method (MLDG) [1] that has shown promise for image
classification tasks to the context of biomedical image segmentation, termed
MLDG-Seg. To evaluate this approach, we focus on the problem of vertebrae
segmentation in CT images and utilize three publicly available databases. To
address the above-mentioned drawbacks of existing domain generalization
methods, we construct three domain generalization contexts: (a) generalization
to new anatomies: the vertebrae are divided into four domains: lumbar, lower,
middle, and upper thoracic regions. The model is trained on three domains, and
then tested on the unseen fourth domain. (b) generalization to a diseased
population with fractured vertebrae, dislocated discs: the model is trained on
a healthy population, then tested on unseen data of a diseased population from
another medical center. (c) generalization to unseen anatomies, surgical
implants, different acquisition protocols, arbitrary orientations, and field
of view (FoV): the model is tested on a very large dataset. Through these
three contexts, we show that MLDG-Seg is able to learn generalized
representations from very limited training examples. Finally, we show that the
learned generalized representation can be quickly adapted with a few examples
from the unseen target distribution in a k-shot learning setting to achieve
additional performance gains.

##  2 Methodology

Meta-learning, or learning to learn [17], aims at learning a variety of tasks,
and then quickly adapting to new tasks in different settings. We adopt the
optimization-based model-agnostic method proposed in [1], called Meta Learning
Domain Generalization (MLDG). Here, we briefly describe the method.

Description. Let there be two distributions: source $\mathcal{S}$, and target
$\mathcal{T}$. Both $\mathcal{S}$, and $\mathcal{T}$ share the same task, for
example, segmentation or classification with the same label space. The goal of
MLDG is to learn a single set of model parameters $\mathcal{\theta}$ via
gradient descent and two meta-learners: meta-train and meta-test procedures.
The model is trained on only the source $\mathcal{S}$ domains, and then tested
on target $\mathcal{T}$ domain. The source domains $\mathcal{S}$ are split
into two sets: meta-train domains $\hat{S}$, and meta-test domains $\bar{S}$ =
$\mathcal{S}$ \- $\hat{S}$. The goal of the two splits is to mimic the setting
of domain shifts, and thereby make it easier for the model to generalize on an
unseen target domain $\mathcal{T}$. We reproduce the learning procedure in
Algorithm 1, and show the extended version in Fig. 1.

Algorithm 1 MLDG

1:Input: Source domains $\mathcal{S}$

2:Model parameters $\mathcal{\theta}$ and Hyperparameters: $\alpha$, $\beta$,
$\gamma$.

3:for $iterations=1,2,\ldots$ do

4: Randomly Split source domains $\mathcal{S}$ into meta-train $\hat{S}$, and
meta-test $\bar{S}$

5: Meta-train: Gradients
$\nabla_{\theta}=F^{{}^{\prime}}_{\theta}(\hat{S};\theta)$

6: Updated parameters:
$\theta^{{}^{\prime}}{\leftarrow}\theta-\alpha\nabla_{\theta}$

7: Meta-test: Compute Loss with updated parameter $\theta^{{}^{\prime}}$ as
$G(\bar{S};\theta^{{}^{\prime}})$

8: Final Model parameters:
$\theta{\leftarrow}\theta-\gamma\frac{\partial(F(\hat{S};\theta)+\beta
G(\bar{S};\theta-\alpha\nabla_{\theta}))}{\partial\theta}$

9:end for

Explanation. Let’s consider a motivating example (carried out later in
Experiment 1) from the image manifold of lumbar ($\mathcal{D}$1), lower
thoracic ($\mathcal{D}$2) and middle thoracic regions ($\mathcal{D}$3),
comprising the set of source domains $\mathcal{S}$; as well as upper
($\mathcal{D}$4) thoracic region, which constitutes the unseen test
$\mathcal{T}$ domain. Refer to Algorithm 1. The MLDG method is supposed to
learn a single model parameter $\mathcal{\theta}$ with the help of two
optimization steps. At every iteration, the set of images in the source
domains, here ($\mathcal{D}$1, $\mathcal{D}$2, and $\mathcal{D}$3) are
randomly split into a meta-train (for example, consisting of images from
$\mathcal{D}$1, $\mathcal{D}$2 set), and meta-test with the set
$\mathcal{D}$3\. Now, two losses are computed. The first loss $\mathcal{F}$ is
computed using the training examples from meta-train set and the gradient is
computed with respect to the model parameter $\theta$. The second loss
$\mathcal{G}$ is computed on the meta-test set with the updated parameter
$\theta^{{}^{\prime}}$ = $\theta-\alpha\nabla_{\theta}$. The key idea by
introducing the second loss in the meta-test stage is that an improvement of
the model’s performance on the meta-train set should also improve the model’s
performance on the meta-test set. The final model parameter $\mathcal{\theta}$
is updated by taking the gradient of the weighted combination of the two
losses $\mathcal{F}$, and $\mathcal{G}$. By doing so, the model is tuned in
such a way that performance is improved in both meta-train, and meta-test
domains. In other words, the model is regularized and does not overfit to one
particular domain, by finding the best possible gradient direction due to the
joint optimization of the two losses. Compare this to a “vanilla” setup, where
a model is directly given images from the three domains $\mathcal{D}$1,
$\mathcal{D}$2, and $\mathcal{D}$3 without any meta-learning setup, might
overfit to a domain by minimizing its loss, and maximizing the loss for the
other domains.

Figure 1: MLDG-Seg: A schematic. Our method extends the MLDG algorithm to a
segmentation by using a 3D Unet-like architecture as the backbone. To explain
the segmentation procedure, consider the four domains: lumbar
($\mathcal{D}$1), lower ($\mathcal{D}$2), middle ($\mathcal{D}$3), upper
($\mathcal{D}$4) thoracic regions. As an example, the domains $\mathcal{D}$1,
and $\mathcal{D}$2 comprise the meta-train set, and $\mathcal{D}$3 the meta-
test. The domain $\mathcal{D}$4 is the held-out test domain.

##  3 Experiments

###  3.1 Databases

We validate our approach on three publicly available CT vertebrae segmentation
datasets. See supplementary material for sample images of the three datasets.

CSI challenge - healthy cases. The datasets of spine CT (MICCAI 2014 challenge
[18]) were acquired during daily clinical routine work in a trauma center from
10 adults (age: 16 to 35 years). In each subject, all 12 thoracic and 5 lumbar
vertebrae were manually segmented to obtain ground truth.

xVertSeg segmentation challenge - pathology cases. This database consists of
fractured and non-fractured CT lumbar spine images. We used the 15 subjects,
ranging in age from 40 to 90 years, made publicly available with their
corresponding lumbar ground truth segmentations [19].

VerSe MICCAI segmentation challenge 2020 - versatile dataset. This database
consists of labeled lumbar, thoracic, cervical vertebrae across 100 cases in
the released set [20], [21] as of June 2020. The data comes from several
medical centers, with a very wide range of acquisition protocols, certain
vertebrae with surgical implants, and a range of FoV with arbitrary image
orientations.

###  3.2 Experimental setup

We define the different implemented procedures, and then perform three
experiments. Baseline is the vanilla 3D Unet-like [23] architecture. The model
is trained on the source domains, and then tested on the unseen held-out test
domain. The procedure is then repeated for the other domains. MLDG-Seg is the
vanilla 3D Unet-like architecture trained under the scheme in Algorithm 1. The
model is trained on the source domains, but at every iteration the source
domains are divided into meta-train and meta-test splits. The model is then
tested on unseen held-out test domain. The procedure is then repeated for the
other domains. k-shot learning Once the MLDG-Seg model has been trained on
source domains, it can be quickly adapted via fine-tuning using very limited
labeled examples from the unseen target domain. The weights of the encoder,
and bottleneck layers of the 3D Unet-like architecture were frozen in the
k-shot learning experiments for fine-tuning. Our approach of k-shot learning
is different than that of test time adaptation methods [24, 25, 26, 27], where
usually the encoder weights are not frozen in a Y-shaped network architecture,
but are explicitly updated using a few examples at the time of testing. In
contrast, we fine-tune the decoder weights with k examples from the test
distribution. Oracle To establish a theoretical upper bound on segmentation
accuracy, we train the vanilla 3D Unet-like architecture using labeled
examples from the target domain, and then test on held-out examples from the
target domain. This is presumably the easiest task, since the training and
test domain distribution are similar. However, in our relatively small
datasets the amount of available data for the oracle experiment is smaller
than for the domain generalization experiments, so oracle performance reported
below should be read with caution.

Experiment 1. The images in all the subjects of the CSI database are divided
into different regions: lumbar (L1-L5), lower thoracic (T9-T12), middle
thoracic (T5-T8), and upper thoracic vertebrae (T1-T4). Each of these regions
comprise of a domain. Here, the aim is to let the model generalize across the
underlying vertebral anatomy, the intensity profile and the surrounding
intervertebral disc space which varies significantly along the vertebral
column.

Experiment 2. Here, the four domains i.e., lumbar, lower, middle, and upper
thoracic vertebrae from the CSI healthy database comprise the set of source
domains. The xVertSeg database with pathology cases comprise the unseen target
domain. Here, the aim is to let the model generalize from the vertebrae of
healthy subjects to the structure of fractured or dislocated vertebrae images
obtained at a different medical center.

Experiment 3. Here, the model is trained on CSI and xVertSeg datasets and
tested on VerSe dataset, with the aim to generalize to unseen anatomies,
different acquisition protocols, arbitrary orientations, FoVs, and pathology.

Train, Validation, and Test data split-up details. The supplementary material
details the training, validation, and test sets for all three databases.

Implementation details. All the images were converted to 1 mm3 isotropic
resolution using FreeSurfer [22]. The images were standardized, using mean
subtraction and division by standard deviation, and then normalized between 0
and 1. We implemented a 3D Unet-like [23] architecture in PyTorch [28]. This
3D Unet is the backbone architecture used for all the experiments. Each of the
three hyperparameters $\alpha$, $\beta$, and $\gamma$ in MLDG (Algorithm 1)
are set to 1. We use stochastic gradient descent as the optimizer with a
learning rate of 0.001, momentum of 0.9, and a weight decay of 5 x 10-5,
dropout with probability of 0.3, and groupnorm as the normalization technique.
See supplementary material for more details on the network architecture. The
loss function used is Generalized Dice Loss [29]. We randomly sample 50
patches of size 64x64x64 from each subject in a domain, and perform data
augmentation by randomly rotating and flipping 30 patches out of these 50
patches for every procedure. For the k-shot procedure, we randomly sampled 5
patches from each of the k-th subject from the unseen distribution. We trained
every procedure, as described above, in experiment 1 for 10 epochs, and
experiments 2 and 3 for 15 epochs. We use the model that gave the highest Dice
score on the validation set to evaluate the unseen test domain. The models
were trained on Nvidia P100 Tesla GPUs. A sliding window approach was used to
obtain the predicted segmentations, which were then post-processed by
retaining the largest connected component.

Evaluation details. We compute Dice coefficient (%), a volume-based
measurement, and Average Symmetric Surface Distance (ASSD) in mm, a surface-
based metric between the groundtruth image and a given model procedure
segmentation output. Desired: Higher Dice, and lower ASSD scores. Furthermore,
we perform pairwise Wilcoxon signed-rank test [30] between the baseline and
each of the other procedures in all the three experiments for both Dice score
and ASSD. In the three Tables 1, 2 and 3, we highlight the procedures which
reach significance at an $\alpha$=0.05 significance value using the following
notation to denote the level of significance: * ($p<$0.05), ** ($p<$0.005),
and *** ($p<$0.0005).

##  4 Results and Discussion

Experiment 1. Table 1 tabulates the results on different held-out test
distributions. Each model is trained on three domains and then tested on the
fourth unseen domain, except the oracle which is trained and tested on the
same domain. MLDG-Seg consistently outperforms the baseline on both Dice and
ASSD, and shows the desired low variance amongst the subjects. With a very few
labeled examples from the test distribution, we see a further boost in
performance. Fig. 2C shows that the MLDG-Seg is better able to segment the
region of interest (ROI) by having significantly reduced background error than
the baseline, in addition to the correct delineation of intervertebral discs
(IVDs) when the held-out distribution is the middle thoracic region. A further
improvement in performance is obtained using additional k-shot examples over
MLDG-Seg. Ideally, the oracle should have the best performance than the rest
of the procedures as the test domain distribution is similar to the training
domain distribution. Here in this particular experimental setup, it is not
surprising to see that the MLDG-Seg (and in some test domains, the baseline)
performs better than the oracle. This might be due to the limited number of
training (and validation) examples available to train the oracle. Furthermore,
since all the procedures were trained for a fixed number of epochs, the oracle
might have had a disadvantage of being trained for a lesser number of gradient
steps than the MLDG-Seg procedure. An alternative approach would be to train a
single oracle model which learns from all the domains simultaneously, and then
test on each domain separately.

Experiment 2. Table 2 and the Spaghetti plots in Fig. 3 shows that the MLDG-
Seg outperforms the baseline. Fig. 2B shows that MLDG-Seg is able to delineate
the vertebrae by not segmenting the undesired spinal cord, as incorrectly
segmented by the baseline. The k-shot setting further improves the
segmentation. Here again, it is not necessarily surprising to see that MLDG-
Seg performs better than the oracle. The oracle was trained with a smaller
number of subjects than the baseline and MDG-Seg (see supplement).

Experiment 3. Table 3 shows the improved performance of MLDG-Seg over the
baseline, and consistent performance boost in the few-shot learning regime.
The k-shot procedure is either on par with the oracle or outperforms the same
for k=4, 5 and 6. Fig. 2A depicts the superior performance of MLDG-Seg over
the baseline on a compression fraction subject from a different distribution
than the training set, where MLDG-Seg is able to segment more vertebrae
completely than the baseline. The k-shot setting further improves the
performance by segmenting the remaining vertebrae not segmented by baseline or
MLDG-Seg. Here, the oracle performs better than most of the procedures, which
is due to the fact that the training set for oracle is relatively larger than
in experiments 1 and 2.

Table 1: Dice score (%), and ASSD in mm (mean $\pm$ std. dev) for Experiment
1. Each row reports the result on held-out unseen test domain, where the model
was trained on the remaining three domains. Supplement contains results on
additional subjects.

Test domain |  Baseline |  MLDG-Seg |  k=1 |  k=2 |  Oracle
---|---|---|---|---|---
Lumbar |  81.67 $\pm$ 8.45 (1.85 $\pm$ 0.75) |  87.85 $\pm$ 2.73 (1.42 $\pm$ 0.31) |  88.57 $\pm$ 1.53 (1.46 $\pm$ 0.24) |  88.19 $\pm$ 2.47 (1.69 $\pm$ 0.67) |  83.60 $\pm$ 2.68 (2.74 $\pm$ 0.37)
Lower Th |  83.52 $\pm$ 4.12 (2.66 $\pm$ 0.98) |  86.17 $\pm$ 2.17 (1.44 $\pm$ 0.09) |  81.44 $\pm$ 2.46 (2.73 $\pm$ 0.49) |  82.28 $\pm$ 1.49 (2.74 $\pm$ 0.55) |  80.25 $\pm$ 4.16 (2.32 $\pm$ 0.41)
Middle Th |  55.72 $\pm$ 4.13 (10.41 $\pm$ 0.42) |  64.36 $\pm$ 11.45 (6.85 $\pm$ 3.18) |  75.57 $\pm$ 5.60 (3.56 $\pm$ 1.64) |  76.98 $\pm$ 8.66 (2.20 $\pm$ 1.45) |  83.60 $\pm$ 0.58 (2.21 $\pm$ 0.54)
Upper Th |  82.00 $\pm$ 1.45 (1.68 $\pm$ 0.18) |  83.70 $\pm$ 2.19 (1.46 $\pm$ 0.20) |  74.47 $\pm$ 6.21 (4.83 $\pm$ 2.56) |  81.50 $\pm$ 4.07 (1.75 $\pm$ 0.39) |  75.84 $\pm$ 4.00 (1.93 $\pm$ 0.36)

Figure 2: Qualitative illustrations for the three experiments. (A) Experiment
3: tested on VerSe dataset, here: k=7. (B) Experiment 2: tested on xVert
dataset, here: k=4. (C) Experiment 1: shown is the middle thoracic region as
the test set, here: k=2. Qualitative results for lumbar, lower, and upper
thoracic regions can be found in the supplement. A minor discrepancy at the
boundaries is noticed in groundtruth in C, perhaps due to registration errors.

Figure 3: Spaghetti plots are shown for Experiment 3. Left: Dice coefficient,
Right: ASSD score. Y-axis shows the Dice (%), and the ASSD (mm) scores for the
different procedures shown on the X-axis. Each of the plotted lines denote a
subject. Therefore, one can track the performance of a procedure for a given
subject. See supplementary material for the spaghetti plots for all the other
experiments. Table 2: Dice score (%), and ASSD in mm (mean $\pm$ std. dev) for
Experiment 2. The model was trained on CSI dataset and tested on xVertSeg
dataset, thereby generalizing to a pathology dataset when trained on a healthy
population.

Baseline |  MLDG-Seg |  k=1 |  k=2 |  k=3 |  k=4 |  Oracle
---|---|---|---|---|---|---
74.13 $\pm$ 13.69 (3.16 $\pm$ 1.05) |  75.34 $\pm$ 12.10 (2.42 $\pm$ 0.70∗∗) |  76.88 $\pm$ 7.35 (2.74 $\pm$ 0.94∗) |  76.51 $\pm$ 8.85 (2.31 $\pm$ 0.65∗∗) |  77.95 $\pm$ 8.79∗ (2.38 $\pm$ 0.74∗∗) |  79.20 $\pm$ 8.01∗ (2.29 $\pm$ 0.77∗∗) |  74.94 $\pm$ 7.84 (3.98 $\pm$ 1.62)

Table 3: Dice score (%), and ASSD in mm (mean $\pm$ std. dev) for Experiment
3. The model was trained on CSI, and xVertSeg dataset and tested on VerSe
dataset, thereby generalizing to fractured or dislocated vertebrae images
obtained at different medical centers, unseen anatomies, different acquisition
protocols, arbitrary orientations, FoVs, and various pathology.

Procedure |  Dice score |  ASSD
---|---|---
Baseline |  54.58 $\pm$ 18.16 |  4.23 $\pm$ 4.21
MLDG-Seg |  64.82 $\pm$ 13.13∗∗∗ |  2.99 $\pm$ 1.59∗
k=1 |  69.12 $\pm$ 10.48∗∗∗ |  2.57 $\pm$ 1.21∗∗
k=2 |  71.49 $\pm$ 9.04∗∗∗ |  2.31 $\pm$ 0.91∗∗∗
k=3 |  70.18 $\pm$ 9.53∗∗∗ |  2.42 $\pm$ 0.76∗∗
k=4 |  75.90 $\pm$ 4.99∗∗∗ |  2.29 $\pm$ 0.69∗
k=5 |  75.27 $\pm$ 5.11∗∗∗ |  2.31 $\pm$ 0.68∗
k=6 |  75.54 $\pm$ 5.71∗∗∗ |  2.13 $\pm$ 0.60∗∗
k=7 |  75.28 $\pm$ 5.67∗∗∗ |  2.17 $\pm$ 0.57∗∗
Oracle |  74.97 $\pm$ 5.81∗∗∗ |  2.31 $\pm$ 0.62∗

Conclusion In the present work, we benchmarked the performance of a gradient-
based meta-learning domain generalization segmentation method in the context
of biomedical image analysis, across a variety of training and test settings.
The method was not only able to generalize across multiple medical sites and
scanners, which is the most widely studied problem of generalization, but was
also able to generalize to newly introduced settings of unseen complex
vertebrae anatomies, surrounding inter-vertebral discs space, varying bone and
soft tissue intensities distribution, diseased populations, different
acquisition protocols, arbitrary orientations, and FoVs, thus resembling
actual clinical settings. In future, we will evaluate the method on other
modalities such as MRI, and US and compare with the recently proposed domain
generalization methods. Our source code, scripts, and dataset-split files are
available at:
https://github.com/Pulkit-Khandelwal/medical-mldg-seg.

Acknowledgments This work was supported by NIH grant R01 EB017255.

## References

  * [1] Li, D., Yang, Y., Song, Y.Z. and Hospedales, T.M., 2018, April. Learning to generalize: Meta-learning for domain generalization. In Thirty-Second AAAI Conference on Artificial Intelligence.
  * [2] Meng, Q., Rueckert, D., and Kainz, B. (2020). Learning Cross-domain Generalizable Features by Representation Disentanglement. arXiv preprint arXiv:2003.00321.
  * [3] Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M. and Lempitsky, V., 2016. Domain-adversarial training of neural networks. The Journal of Machine Learning Research, 17(1), pp.2096-2030.
  * [4] Li, D., Yang, Y., Song, Y.Z. and Hospedales, T.M., 2017. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision (pp. 5542-5550).
  * [5] Snell, J., Swersky, K. and Zemel, R., 2017. Prototypical networks for few-shot learning. In Advances in neural information processing systems (pp. 4077-4087).
  * [6] Glocker, B., Robinson, R., Castro, D.C., Dou, Q. and Konukoglu, E., 2019. Machine learning with multi-site imaging data: An empirical study on the impact of scanner effects. arXiv preprint arXiv:1910.04597.
  * [7] Kamnitsas, K., Baumgartner, C., Ledig, C., Newcombe, V., Simpson, J., Kane, A., Menon, D., Nori, A., Criminisi, A., Rueckert, D. and Glocker, B., 2017, June. Unsupervised domain adaptation in brain lesion segmentation with adversarial networks. In International conference on information processing in medical imaging (pp. 597-609). Springer, Cham.
  * [8] Valindria, V.V., Lavdas, I., Bai, W., Kamnitsas, K., Aboagye, E.O., Rockall, A.G., Rueckert, D. and Glocker, B., 2018. Domain adaptation for MRI organ segmentation using reverse classification accuracy. arXiv preprint arXiv:1806.00363.
  * [9] Ouyang, C., Kamnitsas, K., Biffi, C., Duan, J. and Rueckert, D., 2019, October. Data efficient unsupervised domain adaptation for Cross-Modality image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention (pp. 669-677). Springer, Cham.
  * [10] Liu, Y., Khosravan, N., Liu, Y., Stember, J., Shoag, J., Bagci, U. and Jambawalikar, S., 2019. Cross-Modality Knowledge Transfer for Prostate Segmentation from CT Scans. In Domain Adaptation and Representation Transfer and Medical Image Learning with Less Labels and Imperfect Data (pp. 63-71). Springer, Cham.
  * [11] Dou, Q., Ouyang, C., Chen, C., Chen, H., Glocker, B., Zhuang, X. and Heng, P.A., 2019. Pnp-adanet: Plug-and-play adversarial domain adaptation network at unpaired cross-modality cardiac segmentation. IEEE Access, 7, pp.99065-99076.
  * [12] Conjeti, S., Katouzian, A., Roy, A.G., Peter, L., Sheet, D., Carlier, S., Laine, A. and Navab, N., 2016. Supervised domain adaptation of decision forests: Transfer of models trained in vitro for in vivo intravascular ultrasound tissue characterization. Medical image analysis, 32, pp.1-17.
  * [13] Roy, A.G., Siddiqui, S., Pölsterl, S., Navab, N. and Wachinger, C., 2020. ‘Squeeze and excite’guided few-shot segmentation of volumetric images. Medical image analysis, 59, p.101587.
  * [14] Zhang, L., Wang, X., Yang, D., Sanford, T., Harmon, S., Turkbey, B., Wood, B.J., Roth, H., Myronenko, A., Xu, D. and Xu, Z., 2020. Generalizing deep learning for medical image segmentation to unseen domains via deep stacked transformation. IEEE Transactions on Medical Imaging.
  * [15] Dou, Q., de Castro, D.C., Kamnitsas, K. and Glocker, B., 2019. Domain generalization via model-agnostic learning of semantic features. In Advances in Neural Information Processing Systems (pp. 6450-6461).
  * [16] Ilse, M., Tomczak, J.M., Louizos, C. and Welling, M., 2019. Diva: Domain invariant variational autoencoders. arXiv preprint arXiv:1905.10427.
  * [17] Finn, C., Abbeel, P. and Levine, S., 2017, August. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 (pp. 1126-1135). JMLR. org.
  * [18] Yao, J., Burns, J.E., Forsberg, D., Seitel, A., Rasoulian, A., Abolmaesumi, P., Hammernik, K., Urschler, M., Ibragimov, B., Korez, R. and Vrtovec, T., 2016. A multi-center milestone study of clinical vertebral CT segmentation. Computerized Medical Imaging and Graphics, 49, pp.16-28.
  * [19] R. Korez, B. Ibragimov, B. Likar, F. Pernuš and T. Vrtovec, ”A Framework for Automated Spine and Vertebrae Interpolation-Based Detection and Model-Based Segmentation,” in IEEE Transactions on Medical Imaging, vol. 34, no. 8, pp. 1649-1662, Aug. 2015.
  * [20] Sekuboyina, A., Bayat, A., Husseini, M.E., Löffler, M., Rempfler, M., Kukačka, J., Tetteh, G., Valentinitsch, A., Payer, C., Urschler, M. and Chen, M., 2020. VerSe: A Vertebrae Labelling and Segmentation Benchmark. arXiv preprint arXiv:2001.09193.
  * [21] Löffler, M.T., Sekuboyina, A., Jacob, A., Grau, A.L., Scharr, A., El Husseini, M., Kallweit, M., Zimmer, C., Baum, T. and Kirschke, J.S., 2020. A Vertebral Segmentation Dataset with Fracture Grading. Radiology: Artificial Intelligence, 2(4), p.e190138.
  * [22] Fischl, B., 2012. FreeSurfer. Neuroimage, 62(2), pp.774-781.
  * [23] Ronneberger, O., Fischer, P. and Brox, T., 2015, October. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention (pp. 234-241). Springer, Cham.
  * [24] Wang, G., Li, W., Zuluaga, M.A., Pratt, R., Patel, P.A., Aertsen, M., Doel, T., David, A.L., Deprest, J., Ourselin, S. and Vercauteren, T., 2018. Interactive medical image segmentation using deep learning with image-specific fine tuning. IEEE transactions on medical imaging, 37(7), pp.1562-1573.
  * [25] Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A.A. and Hardt, M., 2019. Test-time training for out-of-distribution generalization. arXiv preprint arXiv:1909.13231.
  * [26] Karani, N., Chaitanya, K. and Konukoglu, E., 2020. Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation. arXiv preprint arXiv:2004.04668.
  * [27] Zhang, J., Liu, Z., Zhang, S., Zhang, H., Spincemaille, P., Nguyen, T.D., Sabuncu, M.R. and Wang, Y., 2020. Fidelity imposed network edit (FINE) for solving ill-posed image reconstruction. NeuroImage, 211, p.116579.
  * [28] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L. and Desmaison, A., 2019. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (pp. 8024-8035).
  * [29] Crum, W.R., Camara, O. and Hill, D.L., 2006. Generalized overlap measures for evaluation and validation in medical image analysis. IEEE transactions on medical imaging, 25(11), pp.1451-1461.
  * [30] Wilcoxon, F., 1992. Individual comparisons by ranking methods. In Breakthroughs in statistics (pp. 196-202). Springer, New York, NY.

Domain Generalizer: A Few-shot Meta Learning Framework for Domain
Generalization in Medical Imaging [Supplementary Material] Pulkit Khandelwal
✉Paul Yushkevich

##  5 Dataset split-up for train, validation, and test sets

Figure 4: The CSI database consists of 10 subjects each comprising the four
domains: lumbar, lower, middle, and upper thoracic vertebrae. The numerals in
the above figure denotes the subject number. The data split-up on the left
side of the black vertical bar gives the split up for the results in Table 1
of the main paper. The data split for the oracle (upper-bound) procedure in
Table 1 in the main paper is depicted on the right side of the vertical bar.
The results in Table 1 in the main paper is for the first four subjects. The
results are then reported for all the four domains, leaving one out as test.
The results on all the 10 subjects are given in Table 4 in Section 9 of this
supplementary material. Note: Dx denotes any of the domains. Figure 5: The
xVertSeg database consists of 15 subjects. Depicted above is the split for the
results in Table 2 in the main paper. Note that oracle is the upper-bound in
the figure. Figure 6: The VerSe segmentation challenge dataset consists of 100
cases at the time of writing this manuscript in June 2020. We randomly select
55 cases, and assign them to the four categories as listed in the table. Shown
are the subject IDs.

##  6 3D Unet-like architecture

Figure 7: 3D Unet-like architecture implemented. IN and OUT are the input and
the output respectively each with dimensions 1x64x64x64, and 2X64x64x64 (for
the two classes). There are five encoder blocks (with the last block acting as
the bottleneck layer). There are four decoder blocks. The dashed lines
represent the skip connections. We have used dropout with probability of 0.3,
and _groupnorm_ as the normalization method. Note: the first encoder block
does not consist of the MaxPool operation, and the first decoder block does
not consist of the first set of conv3D, groupnorm, and ReLU units. The encoder
block has the following output dimensions at each of the encoder block: E1:
16x64x64x64, E2: 32x32x32x32, E3: 64x16x16x16, E4: 128x8x8x8, E5: 256x4x4x4.
The decoder block has the following output dimensions at each of the decoder
block: D1: 256x8x8x8, D2: 128x16x16x16, D3: 64x32x32x32, D4: 32x64x64x64.

##  7 Sample image slices for the three datasets

Figure 8: CSI dataset Figure 9: xVertSeg dataset Figure 10: VerSe dataset

##  8 t-SNE plots

Figure 11: Shown are the t-SNE [1] plots for Baseline (left), and MLDG-Seg
(right). The trained bottleneck layer’s feature vector with dimension of 16384
(E5: 256x4x4x4) are plotted for the experiment when the model was trained on
CSI and xVertSeg dataset (Experiment 3). The four colors represent the four
domains: lumbar, lower thoracic, middle thoracic, and upper thoracic regions
in the CSI dataset. It can be seen that the data is clustered more separately
in the trained baseline than the MLDG-Seg model. Whereas, the trained MLDG-Seg
method learns domain agnostic features to better generalize the model. The
following hyperparameters were used to produce the t-sne plot: perplexity: 45,
learning rate: 10, supervision: 30, iterations: 700.

##  9 Experiment 1 additional results

Table 4: Results for Experiment 1 on all the 10 subjects. See Figure 1 for the data split-up. Dice score (%), and Average Symmetric Surface Distance (ASSD) in mm (mean $\pm$ std. dev) for Experiment 1 within brackets. Each row reports the result on the held-out unseen test domain, where the model was trained on the remaining three domains. We performed the pairwise Wilcoxon signed-rank test between the baseline and MLDG-Seg for both Dice score and ASSD. We highlight the procedures which reach significance at an $\alpha$=0.05 significance value using the following notation to denote the level of significance: * ($p$$<$0.05), ** ($p$$<$0.005), and *** ($p$$<$0.0005). Test domain | Baseline | MLDG-Seg
---|---|---
Lumbar | 81.22 $\pm$ 9.48 (1.90 $\pm$ 0.80) | 88.31 $\pm$ 2.55∗ (1.39 $\pm$ 0.32∗)
Lower Thoracic | 83.58 $\pm$ 6.71 (2.39 $\pm$ 1.21) | 86.77 $\pm$ 3.35 (1.41 $\pm$ 0.57∗)
Middle Thoracic | 65.5 $\pm$ 15.14 (7.49 $\pm$ 4.17) | 71.22 $\pm$ 14.4 (4.93 $\pm$ 3.80∗∗)
Upper Thoracic | 81.55 $\pm$ 3.99 (1.59 $\pm$ 0.33) | 81.75 $\pm$ 5.87 (1.70 $\pm$ 0.88)

##  10 Experiment 1 additional qualitative results

Figure 12: Qualitative illustrations for the three held-out domains in
experiment 1. (A) Held-out domain: Lumbar, here: k=2. (B) Held-out domain:
Lower Thoracic, here: k=2. (C) Held-out domain: Middle Thoracic, here: k=2. A
minor discrepancy at the boundaries is noticed in the groundtruth
segmentations, perhaps due to registration errors when converting to an
isotropic resolution.

##  11 Spaghetti Plots

In the following set of figures, Spaghetti plots are shown for all the
experiments. Y-axis shows the Dice (%), and the ASSD (mm) scores for the
different procedures, shown on the X-axis. Each of the plotted lines denote a
subject. Therefore, one can track the performance of a procedure for a given
subject.

###  11.1 Experiment 1 (4 subjects) [Refer Table 1 in the main paper]

Figure 13: Left: Dice coefficient, Right: ASSD score. Held-out domain: Lumbar.

Figure 14: Left: Dice coefficient, Right: ASSD score. Held-out domain: Lower
thoracic.

Figure 15: Left: Dice coefficient, Right: ASSD score. Held-out domain: Middle
thoracic.

Figure 16: Left: Dice coefficient, Right: ASSD score. Held-out domain: Upper
thoracic.

###  11.2 Experiment 1 (10 subjects) [Refer Table 1 in this supplement in
Section 9 above]

Figure 17: Left: Dice coefficient, Right: ASSD score. Held-out domain: Lumbar.

Figure 18: Left: Dice coefficient, Right: ASSD score. Held-out domain: Lower
thoracic.

Figure 19: Left: Dice coefficient, Right: ASSD score. Held-out domain: Middle
thoracic.

Figure 20: Left: Dice coefficient, Right: ASSD score. Held-out domain: Upper
thoracic.

###  11.3 Experiment 2 [Refer Table 2 in the main paper]

Figure 21: Left: Dice coefficient, Right: ASSD score. Experiment 2.

## References

  * [1] Maaten, L.V.D. and Hinton, G., 2008. Visualizing data using t-SNE. Journal of machine learning research, 9(Nov), pp.2579-2605.


# Uncertainty Estimation of Transformers’ Predictions via Topological Analysis
of the Attention Matrices

Elizaveta Kostenok
Moscow Institute of
Physics and Technology
<EMAIL_ADDRESS>
&Daniil Cherniavskii
AIRI

<EMAIL_ADDRESS>
&Alexey Zaytsev
Skolkovo Institute of
Science and Technology
<EMAIL_ADDRESS>

###### Abstract

Determining the degree of confidence of deep learning model in its prediction
is an open problem in the field of natural language processing. Most of the
classical methods for uncertainty estimation are quite weak for text
classification models. We set the task of obtaining an uncertainty estimate
for neural networks based on the Transformer architecture. A key feature of
such mo-dels is the attention mechanism, which supports the information flow
between the hidden representations of tokens in the neural network. We explore
the formed relationships between internal representations using Topological
Data Analysis methods and utilize them to predict model’s confidence. In this
paper, we propose a method for uncertainty estimation based on the topological
properties of the attention mechanism and compare it with classical methods.
As a result, the proposed algorithm surpasses the existing methods in quality
and opens up a new area of application of the attention mechanism, but
requires the selection of topological features.

##  1 Introduction

In deep learning theory, Transformers are a class of neural networks for
processing text data. During the training stage model recognizes grammatical
and semantic relationships between words in a sentence through the interaction
between their internal representations. The permanent exchange of information
between token representations allows the model to determine which of the words
are most important for understanding each other. This interaction is called
the Attention mechanism and ensures high accuracy of Transformers predictions
on the main types of NLP problems, including text classification.

However, for the practical use of neural networks, one should consider not
only the percentage of correct predictions, but also the degree of confidence
in each prediction. Transformers solving a text classification problem can
produce unreliable answers. First, as typical classification models, they are
poorly calibrated in probability Guo et al. (2017). We expect that the
incorrect predictions of the neural network correspond to a low (close to
$0.5$) probability of the selected class. But classification models tend to
overestimate their confidence in the prediction, resulting in a shift in
probabilities closer to one. Secondly, Transformers suffer from adversarial
attacks Guo et al. (2021): a small perturbation in the input of the attacked
model causes a significant change in its output. Therefore, it is crucial to
have a mechanism to identify untrustworthy predictions in situations where
model errors are unacceptable.

This issue is addressed to uncertainty estimation methods. The simplest way to
get an uncertainty estimate for classification task is the output of the
Softmax layer of the model, interpreted as a probability. More advanced
approaches are Bayesian and ensemble methods, as well as MC dropout, which
will be discussed in detail in the 2 Section. When applied to Transformers,
these approaches have low computational efficiency or require significant
modifications of the model setting.

The use of hidden representations of data has a potential to improve the
quality of uncertainty estimates for Transformers. An effective method for
detecting objects that such models regard as outliers uses the distances
between the embeddings of the last layer Vazhentsev et al. (2022). It is based
on the fact that the neural network maps from the space of objects to the
space of internal representations, in which the concept of distance can be
defined. Objects from the same distribution correspond to closely spaced
embeddings, while internal outlier representations are far away from them.
This method is resource efficient, shows consistently high results in metrics,
and points to the promise of using Transformers’ internal mechanism to
estimate uncertainty.

In this paper, we develop the idea of determining the degree of model
confidence in the prediction by its internal response to the input tokens. To
characterize the structure of the Attention mechanism, we use its graph
representation and apply methods of topological data analysis to it. Our main
contributions are listed below:

  * •

We obtain a complete description of the attention mechanism as a set of
topological statistics. Topological features can characterize each attention
matrix independently (feature type $1$) or pairs of attention matrices located
in the different parts of the network (feature type $2$). While in earlier
works Cherniavskii et al. (2022), Kushnareva et al. (2021) only features of
type $1$ are considered, we are using features of type $2$ to analyze
Transformers for the first time. We show that topological statistics help
improve the estimate of model uncertainty.

  * •

We propose an approximation of the model confidence by a trainable analog of
the probability for each prediction. We introduce suitable training pipeline
and loss function for connecting topological statistics with uncertainty of
predictions.

  * •

Our confidence prediction algorithm outperforms baseline methods on three
models determining grammatical correctness of sentences in English, Russian
and Italian, respectively.

##  2 Related work

There are several ways to estimate uncertainty for deep neural networks. For
classification models one can utilize the output of the Softmax layer,
interpreted as a class probability. The authors of Geifman and El-Yaniv (2017)
propose a Softmax Response method based on the maximum probability $p(y|x)$ of
belonging to the class $y=c\in C$ at the model output. The higher this
probability, the more confident the model is in the prediction:

$u_{SR}(x)=1-\max_{c\in C}p(y=c|x)$ (1)

The main advantage of this approach is that it does not require additional
calculations: the required estimate is obtained automatically at the output of
the last layer of the model. However, Softmax probabilities are not really
credible, as models tend to be overconfident Guo et al. (2017).

The classical method for uncertainty estimation, which has a fundamental
theoretical proof, is the Bayesian approach Blundell et al. (2015). It
initially sets a prior distribution on the model parameters, and then, based
on the training samples, calculates the posterior distribution on the weights.
In order to apply the Bayesian approach in practice, various approximations
are used Van Amersfoort et al. (2020) and the choice of approximation greatly
affects the quality of the estimate. In general, Bayesian neural networks
capable of obtaining reliable estimates of uncertainty require significant
changes in the training scheme and learn more slowly than classical neural
networks.

For classical neural networks, there is another way to introduce variability
into the model parameters without changing the architecture and learning
process Lakshminarayanan et al. (2017). It is based on training deep
ensembles: several copies of the original model are trained independently and
after that predictions of the resulting uncorrelated models are averaged.
However, the quality of the uncertainty estimate increases with the number of
models in the ensemble, and, accordingly, with the computational cost of their
training.

The ideas of the previous two methods are developed by the more resource-
efficient Monte Carlo (MC) Dropout method Gal and Ghahramani (2016). It
interprets the dropout layer, widely known regularization technique in deep
neural networks, as a Bayesian approximation of a probabilistic model: a
Gaussian process. Neural networks with different sets of deactivated neurons
can be considered as Monte Carlo samples from the space of all possible neural
networks. So, the same uncertainty estimation methods work for them as for
ensembles. It is important that dropout must be enabled at the testing stage
for the diversity of the predictions, as opposed to the classical use of the
layer at the training stage. The main advantage of this approach is that
single trained model is enough to estimate the uncertainty, while the
disadvantage is the amount of computation for neural network runs with
different sets of disabled neurons. Other approaches that can provide
uncertainty estimation for deep neural networks train a separate head for this
goal Kendall and Gal (2017); Kail et al. (2022). While these approaches often
focuses on computer vision, they can be adopted as well in the natural
language processing.

There are also methods for uncertainty estimation which use the Mahalanobis
distance - a generalization of the Euclidean distance for vectors of random
variables - between the internal representations of the training and test
objects. The article Lee et al. (2018) explores the use of the Mahalanobis
distance between the test samples and the nearest class conditional Gaussian
distribution as an uncertainty estimate:

$u_{md}=\min_{c\in C}(h_{i}-\mu_{c})^{T}\Sigma^{-1}(h_{i}-\mu_{c})$ (2)

where $h_{i}$ is the hidden representation of the test sample, $\mu_{c}$ is
the centroid of the $C$ class, $\Sigma$ is the covariance matrix of the hidden
representations of the training samples. An advantage of this method is that
it does not require significant changes in the architecture of the model and
the memory and time costs of training and storing multiple copies of the
model. Obtaining an uncertainty estimate only adds a small amount of
computation to the model testing phase. A disadvantage of this approach is the
need to obtain internal representations for the whole training examples, which
may take more time than several forward passes of the neural network on the
test dataset for MC dropout evaluation. This is due to the fact that the
training dataset is much larger than the test dataset.

As a result, each of the existing methods for uncertainty estimation of neural
network predictions has its own benefits and limitations, which, in turn,
strongly depend on the task, the amount of available resources and the
required quality of the estimate. In terms of the ratio of computational
efficiency and reliability of the obtained estimate, the MC dropout methods
and the Mahalanobis distance between internal representations are leading.
However, most comparisons of methods for uncertainty estimation in the
literature are carried out for computer vision problems and their
effectiveness for problems of natural analysis is still unclear. Also, in
works where uncertainty is evaluated for Transformers Shelmanov et al. (2021),
Vazhentsev et al. (2022), the approaches and modifications proposed to them
only occasionally consider the key feature of the architecture, namely the
Attention mechanism. The proposed algorithm for calibrating neural model
predictions by the degree of confidence not only produces high-quality
estimates specifically for the NLP task, but also opens up another potential
area of application of the attention mechanism - uncertainty estimation.

##  3 Problem Statement

###  3.1 Confidence Score

We formalize the problem of estimating uncertainty for a pretrained
Transformer model with fixed weights. For each sample $x$, we need to get the
Confidence Score $c(x)\in[0,1]$ expressing the degree of confidence of the
neural network in the prediction $f_{\theta}(\bf{x})$. The task of uncertainty
estimation in our case is to find the function $c(Attention(x))$.

###  3.2 Testing Method

The search for a suitable metric for our problem is complicated due to the
fact that uncertainty estimation belongs to the class of unsupervised
problems. There are no targets for a measure of the model’s confidence in its
prediction in the training dataset.

In order to correctly evaluate the estimate of model uncertainty, we focus on
the correlation between the confidence of the model and the accuracy of its
predictions. We use the following test pipeline:

  * •

For each of the $N$ test samples, we get the prediction of the neural network
and the Confidence Score.

  * •

We rank objects by Confidence Score in ascending order: from less confident
objects to more confident ones. We calculate and save the accuracy score on
the whole test dataset.

  * •

We iteratively remove the $r$ least confident samples and calculate the metric
on $N-ir$ remaining samples, where $i$ is the iteration number.

  * •

On the graph, we display the dependence of the accuracy score for the
remaining subset on the portion of removed objects $\frac{ir}{N}$

The graph of the obtained dependence is called Accuracy Rejection Curve Nadeem
et al. (2009) and the area under the graph is used as the numerical value of
the metric. It is common to calculate the area from the accuracy level for the
whole test dataset, not taking into account the rectangular area below this
level.

The plot of the Accuracy Rejection curve can be interpreted as follows: a
model correctly calibrated for uncertainty makes errors mainly on examples
with a low Confidence Score. The graph shows that after removing a small
number of such uncertain examples, the accuracy of predictions on the
remaining subset increases significantly. In practice, one can ask experts to
make predictions on these examples instead of a neural network, and thus the
probability of error for a hybrid system from a model and experts will be
minimized. The most effective method for assessing uncertainty in this case
minimizes the number of examples that must be given to experts for analysis.

##  4 Method Description

In this section, we consider in more detail topological statistics of the
Attention mechanism and propose a training scheme for our method.

###  4.1 Topological analysis of the Attention mechanism

Attention matrix has an equivalent representation in the form of a weighted
directed graph $G$. The vertices of such a graph are the tokens of the input
sequence, and the weights of the edges are the attention weights for each pair
of tokens. The direction of the edge is chosen from the request token to the
key token. Using this method, for each input sentence, we get $N_{l}*N_{h}$
attention graphs, where $N_{l}$ is the number of layers, $N_{h}$ is the number
of Transformer heads.

Our task is to extract from graph representations a set of numerical
statistics for each sentence Kushnareva et al. (2021). First, we calculate
graph features: the number of vertices, edges, connected components, simple
cycles, and Betti numbers. This set of statistics is not exhaustive, as it
does not involve edge weights. The way to take into account the presence of
weights is to construct a filtration - a family of graphs $\\{G^{t_{i}}\\}$
obtained from the original one by removing edges with a weight less than the
threshold $t_{i}$ barcode. Successive reduction of edges leads to a change in
the structure of the graph and its main properties listed above. The methods
of topological data analysis make it possible to quantitatively describe the
evolution of graph properties by determining for each property the time of its
appearance and disappearance in the filtration (the corresponding thresholds
are denoted by $t_{birth}$ and $t_{death}$). The set of intervals between
$t_{birth}$ and $t_{death}$ is called a barcode and is a measure of the
stability of the attention structure of Cherniavskii et al. (2022) graphs. The
most stable and pronounced graph properties will have the largest barcode
length. We extract the following numerical statistics from barcodes: sum,
mean, variance and entropy of barcode lengths, number of barcodes with
appearance/disappearance times less/greater than the threshold value.

An analysis of attention graphs Kovaleva et al. (2019) taken from different
heads and layers of the network shows that in the process of filtering graphs
have similar structural elements/patterns. In the article Clark et al. (2019)
these patterns are divided into several main types: attention to the current,
previous and next tokens, attention to service [SEP] and [CLS] tokens. The
authors Kushnareva et al. (2021) propose a graph representation method for
attention patterns and introduce template features based on it. Numerically,
it is equal to the Frobenius norm of the difference between the attention
matrix and the incidence matrix of the attention pattern.

The topological features listed above are calculated independently for
attention graphs from different layers and heads of the network, however, the
attention mechanisms in different parts of the model are trained jointly and
it is better to take into account the connection between them. The cross-
barcode method Barannikov et al. (2021) develops the idea of barcodes and
allows to compare the distributions of attention weights on different layers
and heads of the neural network. Consider two attention graphs $G^{w}$ and
$G^{w^{\prime}}$, where $w$ and $w^{\prime}$ are attention weights, as well as
a matrix composed of pairwise minimum of the weights $M=min(w,w^{\prime})$ and
the corresponding graph $G^{min(w,w^{\prime})}$. The main difference from
barcodes is that the filtering is performed for the graph $G^{w,w^{\prime}}$
constructed from $G^{w}$ and $G^{min(w,w^{\prime})}$. After filtering, we get
a set of intervals $(t_{birth},t_{death})$ and calculate the value of the
topological feature as the total length of the cross-barcode segments.
Intuitively, the cross-barcode expresses simpler graph properties (the number
of simple cycles, connected components, etc.), which, at a fixed $\alpha$
threshold, have already appeared/disappeared in one of the graphs, but have
not yet appeared/disappeared in another graph .

Using the methods above, from the graph representation of the attention
matrices for each input sentence, we obtain a set of topological features of 4
types:

  * •

graph statistics (graph features)

  * •

features obtained from barcodes
(barcode/ripser features)

  * •

features obtained from attention patterns
(template features)

  * •

features obtained from cross-barcodes
(cross-barcode features)

The first 3 types of features are calculated for attention matrices
independently and belong to the first type, the last type of features is
calculated for pairs of matrices and belongs to the second type.

The linguistic interpretation of features is that graph structures are used in
lexicology to describe the laws of semantic changes in a language Hamilton et
al. (2016). The evolution of the meaning of a word over time can be
represented as a graph in which the edges are proportional to the semantic
shift relative to different meanings of the word. A similar evolutionary
process occurs within a neural model as the hidden representation of a word
changes from the initial layers of the network to deeper layers. Groups of
tokens with strong mutual influence form separate clusters in a graph
representation, just as words related in meaning and syntactically in a
sentence form phrases. Barcodes highlight the dominant, most stable semantic
relationships between words, and attention patterns highlight the main
directions of interaction between context tokens.

###  4.2 Confidence predictor model and its training

Formally, given the input set of features $(z_{1},z_{2},\cdots z_{n})$, the
algorithm should produce the number Confidence Score $c\in[0,1]$. This range
of values allows us to interpret the Confidence Score as a probability, thus
the problem of estimating the confidence of the Transformer in the prediction
is reduced to a standard binary classification problem. We use a separate
Score Predictor model for it, consisting of several fully connected layers and
activation functions between them. The last activation function should be
sigmoid to limit the output of the model to $[0,1]$. The simplicity of the
architecture is explained by the high risk of overfitting. We get an increase
in the predictive ability of the model by more careful selection and
aggregation of topological features instead of making the network deeper.

The key point is to build the loss function for our model. Cross-entropy,
which is widely used by models for binary classification, does not suit us,
since the target distribution of the Confidence Score is not available. Also,
optimizing the loss function should reward the neural network for obtaining
uncertainty estimates that truly reflect Transformer’s ability to predict the
correct class for any input. Similar to the method proposed in DeVries and
Taylor (2018), we choose the cross-entropy modification as the loss function,
in which the probability at the output of the model is calibrated using the
Confidence Score.

Let embeddings of words in a sentence $(x_{1},x_{2},\cdots x_{n})$, the
corresponding set of topological features $(z_{1},z_{2},\cdots z_{m})$,
Transformer output
$\mathbf{p}=\mathrm{Softmax}(f(\mathbf{x},\mathbf{\theta}))$ and the
uncertainty estimate for this prediction obtained from the Score Predictor
model $c=\mathrm{Sigmoid}(g(\mathbf{z},\mathbf{w}))$, while
$p_{1},p_{2},c\in[0,1],p_{1}+p_{2}=1$. Then the calibrated probability is
obtained by interpolating between the Transformer’s prediction and its one-hot
target distribution $\mathbf{y}$ with the degree of interpolation determined
by the Confidence Score: $p^{\prime}_{i}=cp_{i}+(1-c)y_{i}$. The final loss
function

$L=-\sum_{i=1}^{2}\log(p^{\prime}_{i})y_{i}-\lambda\log(c)$ (3)

is the sum of cross-entropy with calibrated probability and regularization,
which penalizes the model for uncertain predictions. The ratio of the main
loss and regularization is determined by the $\lambda$ hyperparameter.

Let’s analyze how Confidence Score affects the dynamics of the loss function.
For $c\rightarrow 1$ (case of sure prediction) $p\rightarrow p^{{}^{\prime}}$
and the value of the loss is determined by the original Transformer prediction
without any uncertainty estimate. On the contrary, for $c\rightarrow 0$,
instead of an uncertain prediction, the loss function receives the target
distribution, preventing the error from growing. If $c$ is between 0 and 1,
then the calibrated probability becomes closer to the target penalty cost in
the regularization term. In general, a decrease in the loss function
corresponds to an increase in the Confidence Score on correct predictions and
a decrease in it on erroneous predictions, which is what we need. It is also
important that all transformations in the error functional are differentiable,
but only the parameters of the Score Predictor model are trained. Any change
in the weights of the Transformer will lead to a change in the topology of the
attention mechanism and will require recalculation of topological features, so
the weights of the Transformer are fixed in the process of estimating the
uncertainty.

##  5 Experiments

This section provides implementation details of our uncertainty estimation
algorithm, data and model details, and experimental results.

###  5.1 Data

We are experimenting on 3 Corpus of Linguistic Acceptability (CoLA) Warstadt
et al. (2019) datasets, consisting of sentences in English, Italian and
Russian. Each sentence has a target value: 1 for grammatically correct
sentences, 0 otherwise. Basic information about datasets is given in the
Appendix A.

###  5.2 Models

The standard approach to using BERT-like models is to finetune the pre-trained
models Devlin et al. (2019) that are publicly available. Pre-trained neural
networks have sufficient generalization ability to condition on a specific
task in just a few epochs of retraining with a low learning rate.

In this work, we use pre-trained BERT-base-cased models from the Transformers
library. Hyperparameters of finetuning and final model metrics are given in
the Appendix A.

###  5.3 Basic Methods

We compare the results of our algorithm with three classical methods of
uncertainty estimation: Softmax Response, MC Dropout, Mahalanobis estimator.
The implementation of the last two methods was based on the code for the
article Vazhentsev et al. (2022). Also, we checked the strategy of training
the Score Predictor on the BERT embeddings extracted from its classification
part instead of topological statistics. This approach is successfully used in
the article DeVries and Taylor (2018) for computer vision tasks, however, the
decrease in the performance of this method for NLP tasks provides motivation
for applying TDA methods instead.

###  5.4 Extraction of topological features

Each of the BERT-like neural networks considered in this paper contains 12
layers and 12 heads on each layer, respectively, for each input object, we
extract 144 attention matrices. Then, for each matrix, we extract 3 types of
type 1 topological features, discussed in detail in the 4.1 section: graph
features, barcode features, template features. We calculate features of type 2
between pairs of attention matrices $A_{ik}$ and $A_{kj}$, where the first
index corresponds to the layer number, the second index is the head number,
$i,j,k\in\\{0,12\\}$. An example of a barcode for a test sample object is
shown in the Appendix D. Barcodes are calculated using the Ripser++ Zhang et
al. (2020) library, and cross-barcodes are calculated using the MTopDiv
Barannikov et al. (2021) library. The algorithms implemented in these
libraries use advanced optimizations, which can significantly reduce the time
for computing topological features. As a result, for each input feature we get
vectors of dimension (12, 12, 7) graph features, (12, 12, 14) barcode
features, (12, 12, 5) template features, (12, 12, 144) cross-barcode features.
The first dimension corresponds to the layer number, the second - to the head
number, the third - to the subtype of the topological feature. All subtypes of
topological features are listed in the Appendix C. To train the Score
Predictor model, the selected feature vectors are concatenated along the third
dimension.

###  5.5 Training setup of the Score Predictor

The model that predicts the Confidence Score takes a set of topological
statistics as input and converts it into a scalar value of the uncertainty
score as an output. It is training by minimizing the Confidence Loss
introduced in the 4.2 section, so outputs of the Transformer Softmax layer and
the distribution of target values for the original classification task are
also given to the model as input. High quality estimates were obtained by
Score Predictor model containing 2 fully connected layers with sigmoids as
activation functions. Our optimal setup also includes the Adam optimization
method and the training hyperparameters listed in the Appendix E.

###  5.6 Analysis of topological features of the first type

Let us consider separately the topological features of the first type. The
total number of components in the vectors of topological features of 3 types
for one object is $12*12*(7+14+5)=$3744\. Considering that the number of
objects in the training samples does not exceed 10000, the use of 3744
embeddings would guarantee overfitting. So it was necessary to reduce the
number of components in the feature vector, and we considered two methods:

  * •

averaging over all layers and all heads of the Transformer

  * •

selection of the components that contribute most to the prediction using
Shapley values Lundberg and Lee (2017)

Feature aggregation method | En-CoLA | Ita-CoLA | Ru-CoLA
---|---|---|---
averaging over all heads and layers | 0.079 | 0.088 | 0.075
selection via Shapley values | 0.087 | 0.093 | 0.080
Table 1: Area under Accuracy Rejection curves for different methods of
aggregation of topological features

The results of training the Score Predictor model only on the first type of
topological features are given in the Table 1.

For each model, selection of the most significant components improved the
uncertainty estimate in relation to simple averaging. Therefore, in further
experiments, we feed the Score Predictor model with two components of
topological features of the first type that most affect the result. Feature
selection process via Shapley values is considered in detail in Appendix B.

###  5.7 Analysis of topological features of the second type

Since we found out earlier that the last layer attentions are the most
informative for us, we also calculate cross-barcodes along this layer. We fix
the index $k$ equal to 12, vary the indices $i$, $j$ and observe how the
quality of the uncertainty estimation changes, when the Score Predictor
receives the pairwise attention statistics of the matrices $Aik$ and $Akj$.

Table 2 shows the change in the metric when adding cross-barcodes to the
English-language setup. According to the experimental results, there is a
small portion of the optimal pairs of attention matrices. Cross-barcode
calculated for such pair improves the uncertainty estimate for our topological
method, for example, the pair $A_{12,12}$ and $A_{12,9}$. These pairs are
located in the lower right corner of the table, which corresponds to the 7-12
layer and head indices in the Transformer. The remaining pairs of attention
matrices barely improve the performance of Score Predictor.

| (12, 1) | (12, 3) | (12, 5) | (12, 7) | (12, 9) | (12, 11)
---|---|---|---|---|---|---
(2, 12) | 0.085 | 0.085 | 0.085 | 0.086 | 0.086 | 0.086
(4, 12) | 0.086 | 0.089 | 0.085 | 0.087 | 0.088 | 0.085
(6, 12) | 0.086 | 0.087 | 0.087 | 0.086 | 0.087 | 0.088
(8, 12) | 0.086 | 0.087 | 0.085 | 0.086 | 0.088 | 0.091
(10, 12) | 0.085 | 0.086 | 0.087 | 0.086 | 0.087 | 0.086
(12, 12) | 0.088 | 0.086 | 0.085 | 0.090 | 0.096 | 0.094
Table 2: Area under Accuracy Rejection curves for the topological method with
the addition of cross-barcodes between the $A_{ik}$ and $A_{kj}$ matrices. The
indices $(i,k)$ attention of the matrix $A_{ik}$ are placed vertically, the
indices $(k,j)$ attention of the matrix $A_{kj}$ are placed horizontally

As a result, we supplement the set of topological features of the first type
with the two most informative cross-barcodes. Further, our final results are
given for the Score Predictor, trained only on best-performing topological
features of each type.

##  6 Results

The area under the Accuracy Rejection curve for each method is presented in
the Table 3. This metric has a theoretical upper bound, which depends on the
accuracy of Transformer’s predictions on the full dataset. It can be
interpreted as confidence of an Oracle, for which the condition $\forall
x,x^{\prime}\rightarrow C(x)>C(x^{\prime})$ is satisfied, where $x$ is a
correctly recognized object, $x^{\prime}$ is an incorrectly recognized object
, $C(x)$ and $C(x^{\prime})$ are the corresponding Confidence Score values.
Consequently, the plot of the Accuracy Rejection curve for an Oracle increases
linearly while only incorrectly classified objects are disappearing from the
subset, and then reaches a constant value of $1$.

Method type | Method name | En-CoLA | Ita-CoLA | Ru-CoLA
---|---|---|---|---
Basic Methods | Softmax Response | 0.068 | 0.085 | 0.073
MC Dropout 1 | 0.071 | 0.084 | 0.076
Mahalanobis estimator 2 | 0.083 | 0.091 | 0.080
Embedding estimator 3 | 0.075 | 0.090 | 0.074
Our methods |  Topological estimator without cross-barcodes | 0.087 | 0.092 | 0.080
Topological estimator with cross-barcodes | 0.098 | 0.099 | 0.085
Oracle Upper Bound | 0.124 | 0.121 | 0.118
Table 3: Area under Accuracy Rejection curves for the considered methods

For all benchmarks we have considered, topological methods outperforms basic
methods, and the addition of cross-barcode statistics leads to a significant
increase in metric. This effect is most pronounced for the English-language
BERT and leads to increase of 12 percent relative to the method using only
topological statistics of first type, and the least pronounced for the
Russian-language BERT, where the increase is 2 times less. Among the
baselines, Softmax Response and MC Dropout show poor quality of uncertainty
estimates, while Mahalanobis estimator gives a consistently reliable estimate
and is practically comparable to our topological method without cross-
barcodes. The results of the Embedding estimator method are not stable: for
the Transformer working with Italian texts, they are close to the topological
method without cross-barcodes, but for other models the Embedding estimator is
inferior to the topological methods.

We highlight the key features of compared methods on the example of English-
language benchmark. The graph 1 shows the Accuracy Rejection curves for the
two basic methods, the Oracle upper estimate and our topological method using
cross-barcodes. Firstly, the Accuracy Rejection curve for our method lies
above the corresponding curves of the base methods and reaches a constant
value earlier. Secondly, the main interest in practice is the initial part of
the curve with rejection rate in $[0,0.2]$. For the topological method, this
section is convex, while for the basic methods, the character of the initial
sections is closer to linear. In practical application, convexity is
preferable, since it means a noticeable increase in metric when only a small
number of objects are removed.

Figure 1: Accuracy rejection curves of the selected methods for the BERT-base
model on the En-CoLA test set

##  7 Conclusion

Existing studies in the intersection of TDA and NLP domains prove that use of
topological features of attention maps can boost classification performance.
We show that topological statistics also allow to obtain an uncertainty
estimates of Transformer predictions, and the quality of the estimates is
superior to that of baselines. We prove this by experiments on three BERT
models to determine the linguistic acceptability of sentences in English,
Russian and Italian, respectively. The confidence prediction by topological
method demonstrates an increase compared to the best of the basic methods up
to 16 percent in metric.

Our algorithm for determining the confidence of the Transformer in its
prediction uses two main types of topological statistics: features calculated
for each attention matrix independently and pairwise statistics. The use of
features of the second type for attention mechanism is our innovation and
cross-barcodes significantly boost the quality of uncertainty estimation by
proposed methods.

We found out that the location of the attention matrix in the Transformer
affects the contribution of topological statistics. On average, the last layer
features are the most informative, however, to achieve the highest possible
quality of the estimate, careful selection of features is required. For
topological features calculated independently, this process can be automated
by selecting components with highest Shapley values.

## Limitations

Our method also has some limitations: 1) Feature extraction slows down the
inference (on average 20 sec per sample is required for feature computation
with GPU acceleration). 2) Selection of the most informative attention
matrices should be fully automatized. Calculation of topological statistics of
both types only on a limited small set of attention matrices will
significantly speed up our method and, quite likely, not reduce the quality of
the uncertainty estimation. Optimization of mentioned calculations will be our
main direction for further work in this area.

## References

  * Barannikov et al. (2021) Serguei Barannikov, Ilya Trofimov, Grigorii Sotnikov, Ekaterina Trimbach, Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. 2021.  Manifold topology divergence: a framework for comparing data manifolds. In _Advances in Neural Information Processing Systems_ , volume 34, pages 7294–7305. Curran Associates, Inc.
  * Blundell et al. (2015) Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. 2015.  Weight uncertainty in neural network.  In _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37, pages 1613–1622. PMLR.
  * Cherniavskii et al. (2022) Daniil Cherniavskii, Eduard Tulchinskii, Vladislav Mikhailov, Irina Proskurina, Laida Kushnareva, Ekaterina Artemova, Serguei Barannikov, Irina Piontkovskaya, Dmitri Piontkovski, and Evgeny Burnaev. 2022.  Acceptability judgements via examining the topology of attention maps.  In _Findings of the Association for Computational Linguistics: EMNLP 2022_ , pages 88–107, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. 2019.  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, Florence, Italy. Association for Computational Linguistics.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota. Association for Computational Linguistics.
  * DeVries and Taylor (2018) Terrance DeVries and Graham W. Taylor. 2018.  Learning confidence for out-of-distribution detection in neural networks.  _ArXiv_ , abs/1802.04865.
  * Gal and Ghahramani (2016) Yarin Gal and Zoubin Ghahramani. 2016.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _Proceedings of The 33rd International Conference on Machine Learning_ , volume 48 of _Proceedings of Machine Learning Research_ , pages 1050–1059. PMLR.
  * Geifman and El-Yaniv (2017) Yonatan Geifman and Ran El-Yaniv. 2017.  Selective classification for deep neural networks.  In _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc.
  * Guo et al. (2017) Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. 2017.  On calibration of modern neural network.  In _Proceedings of the 34th International Conference on Machine Learning_ , volume 70, pages 1321–1330.
  * Guo et al. (2021) Chuan Guo, Alexandre Sablayrolles, Hervé Jégou, and Douwe Kiela. 2021.  Gradient-based adversarial attacks against text transformers.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 5747–5757, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Hamilton et al. (2016) William L. Hamilton, Jure Leskovec, and Dan Jurafsky. 2016.  Diachronic word embeddings reveal statistical laws of semantic change.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1489–1501, Berlin, Germany. Association for Computational Linguistics.
  * Kail et al. (2022) Roman Kail, Kirill Fedyanin, Nikita Muravev, Alexey Zaytsev, and Maxim Panov. 2022\.  Scaleface: Uncertainty-aware deep metric learning.  _arXiv preprint arXiv:2209.01880_.
  * Kendall and Gal (2017) Alex Kendall and Yarin Gal. 2017.  What uncertainties do we need in bayesian deep learning for computer vision?  _Advances in neural information processing systems_ , 30.
  * Kovaleva et al. (2019) Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. 2019.  Revealing the dark secrets of BERT.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 4365–4374, Hong Kong, China. Association for Computational Linguistics.
  * Kushnareva et al. (2021) Laida Kushnareva, Daniil Cherniavskii, Vladislav Mikhailov, Ekaterina Artemova, Serguei Barannikov, Alexander Bernstein, Irina Piontkovskaya, Dmitri Piontkovski, and Evgeny Burnaev. 2021.  Artificial text detection via examining the topology of attention maps.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 635–649, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Lakshminarayanan et al. (2017) Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. 2017.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc.
  * Lee et al. (2018) Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. 2018.  A simple unified framework for detecting out-of-distribution samples and adversarial attacks.  In _Advances in Neural Information Processing Systems_ , volume 31. Curran Associates, Inc.
  * Lundberg and Lee (2017) Scott M Lundberg and Su-In Lee. 2017.  A unified approach to interpreting model predictions.  In _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc.
  * Nadeem et al. (2009) Malik Sajjad Ahmed Nadeem, Jean-Daniel Zucker, and Blaise Hanczar. 2009.  Accuracy-rejection curves (arcs) for comparing classification methods with a reject option.  In _International Workshop on Machine Learning in Systems Biology_.
  * Shelmanov et al. (2021) Artem Shelmanov, Evgenii Tsymbalov, Dmitri Puzyrev, Kirill Fedyanin, Alexander Panchenko, and Maxim Panov. 2021.  How certain is your Transformer? In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , pages 1833–1840, Online. Association for Computational Linguistics.
  * Van Amersfoort et al. (2020) Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. 2020.  Uncertainty estimation using a single deep deterministic neural network.  In _Proceedings of the 37th International Conference on Machine Learning_ , volume 119, pages 9690–9700. PMLR.
  * Vazhentsev et al. (2022) Artem Vazhentsev, Gleb Kuzmin, Artem Shelmanov, Akim Tsvigun, Evgenii Tsymbalov, Kirill Fedyanin, Maxim Panov, Alexander Panchenko, Gleb Gusev, Mikhail Burtsev, Manvel Avetisian, and Leonid Zhukov. 2022.  Uncertainty estimation of transformer predictions for misclassification detection.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 8237–8252, Dublin, Ireland. Association for Computational Linguistics.
  * Warstadt et al. (2019) Alex Warstadt, Amanpreet Singh, and Samuel Bowman. 2019.  Neural network acceptability judgments.  _Transactions of the Association for Computational Linguistics_ , 7\.
  * Zhang et al. (2020) Simon Zhang, Mengbai Xiao, and Hao Wang. 2020.  Gpu-accelerated computation of vietoris-rips persistence barcodes.  In _36th International Symposium on Computational Geometry (SoCG 2020)_. Schloss Dagstuhl-Leibniz-Zentrum für Informatik.

##  Appendix A Datasets and pretrained models info

Size of training and test sets for each benchmark is given in Table 4.
Finetuning setups for BERT-base-cased models and final metrics are presented
in the Table 5.

##  Appendix B Shapley values

Shapley values were introduced in game theory to distribute the payoff fairly
among the players in a team according to the contribution of each of them to
the result. In our case, the individual components of the vectors act as
players, and the Shapley values express the influence of each component on the
prediction. To calculate these values, we use the SHAP library. An example of
the analysis of one subtypes of a graph feature (number of vertices) is shown
in the graph 2 The largest variance of Shapley values corresponds to the
greatest influence of the component on the prediction. According to the graph,
the components from the last layer of the Transformer turned out to be the
most important, since their indices are in the range of 130-144. Probably, the
reason for the significance of the components from the last layer are caused
by BERT finetuning. This process only affects the weights of the last layer,
so they probably catches more specific properties of data than initial layers.

Figure 2: Analysis of graph feature components using Shapley values. The number of the component is plotted along the vertical axis, and the Shapley values along the horizontal axis. The color gradient from blue to pink corresponds to the increase in the absolute value of the graph feature Dataset | Training set size | Test set size
---|---|---
En-CoLA | 8551 | 527
Ita-CoLA | 7801 | 946
Ru-CoLA | 7869 | 984
Table 4: Number of sentences in datasets Model | Dataset | Epochs | Batch size | Learning rate | Accuracy
---|---|---|---|---|---
BERT-base | En-CoLA | 3 | 32 | 3e-5 | 0.850
Ita-CoLA | 3 | 64 | 3e-5 | 0.866
RuBERT | Ru-CoLA | 3 | 32 | 3e-5 | 0.802
Table 5: Model configuration and final metrics on the test set

##  Appendix C Feature sybtypes

  1. 1.

Graph statistics:

     * •

Number of vertices

     * •

Number of simple loops

     * •

Number of connectivity components

     * •

Number of edges

     * •

Average vertex degree

     * •

Betti numbers

  2. 2.

Features received from barcodes:

     * •

Sum of barcode lengths

     * •

Variance of barcode lengths

     * •

Entropy of barcode lengths

     * •

Birth time of the longest barcode

     * •

Number of barcodes along the homology dimension

     * •

Number of barcodes with birth/death times greater than/less than a fixed
threshold

  3. 3.

Features obtained from attention patterns:

     * •

Distance to previous token

     * •

Distance to current token

     * •

Distance to next token

     * •

Distance to classification token

     * •

Distance to punctuation marks

  4. 4.

Features obtained from cross-barcodes:

     * •

Sum of lengths of cross-barcode segments

##  Appendix D Cross-barcode sample

Figure 3: An example of a cross-barcode between pairs of Attention matrices,
given by the numbers of the layer and the head of the Transformer (6, 11) and
(11, 6). H0 and H1 correspond to 0- and 1-dimensional homology

##  Appendix E Training configuration of Score Predictor

By grid search we selected the optimal set of training hyperparameters:

N epochs = 250, learning rate = $1e^{-3}$, $\lambda=0.01$