# DPATD: Dual-Phase Audio Transformer for Denoising

Junhui Li1, Pu Wang2, Jialu Li3, Xinzhe Wang4, Youshan Zhang 5  12Department
of Mathematics, School of Science, University of Science and Technology,
Liaoning, Anshan, China
4School of business administration, University of Science and Technology,
Liaoning, Anshan, China
Email<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>3School of public policy, Cornell University, Ithaca, NY,
USA. Email<EMAIL_ADDRESS>5Department of Artificial Intelligence and
Computer Science, Yeshiva University, New York, NY, USA
Email<EMAIL_ADDRESS>
###### Abstract

Recent high-performance transformer-based speech enhancement models
demonstrate that time domain methods could achieve similar performance as
time-frequency domain methods. However, time-domain speech enhancement systems
typically receive input audio sequences consisting of a large number of time
steps, making it challenging to model extremely long sequences and train
models to perform adequately. In this paper, we utilize smaller audio chunks
as input to achieve efficient utilization of audio information to address the
above challenges. We propose a dual-phase audio transformer for denoising
(DPATD), a novel model to organize transformer layers in a deep structure to
learn clean audio sequences for denoising. DPATD splits the audio input into
smaller chunks, where the input length can be proportional to the square root
of the original sequence length. Our memory-compressed explainable attention
is efficient and converges faster compared to the frequently used self-
attention module. Extensive experiments demonstrate that our model outperforms
state-of-the-art methods.

###### Index Terms:

Audio denoising, transformer, audio chunks

##  I Introduction

Speech signals are inevitably accompanied by various types of background noise
in daily environments, such as automatic speech recognition systems, hearing
aids, vehicles and mobile phones, aircraft cockpits, and multi-party
conferencing devices [1]. In the field of speech communication, constructing
efficient models to eliminate background noise is still a difficult challenge.
The goal of speech enhancement learning is to find a transformation that makes
clean speech readily available from the original audio. Recent achievements in
transfer learning from extensive generative language models serve as a major
source of inspiration for our work. There are two main challenges: (1) audio
signals are continuous while textual representations are discrete; and (2) the
decoder is responsible for generating text that is very different from
traditional speech representation. Our efforts focus on directly creating an
audio-denoising model that can learn efficiently and applied to a variety of
datasets.

The transformer was originally proposed for natural language processing (NLP)
tasks [2] and has recently become popular in the fields of computer vision
(CV) [3] and audio processing (AP) [4]. The transformer model can effectively
solve the long-term dependency problem and can run well in parallel, showing
good performance on many natural language processing tasks. Transformer-based
methods also show promising performance in audio denoising. Yu et al. [5]
proposed a dual-branch federative magnitude and phase estimation framework,
named DBT-Net, for monaural speech enhancement, aiming at recovering the
coarse- and fine-grained regions of the overall spectrum in parallel. Wang et
al. [6] proposed a two-stage transformer neural network for end-to-end speech
denoising in the time domain. Yu et al. [7] proposed a cognitive computing-
based speech enhancement model termed SETransformer to improve speech quality
in unknown noisy environments, which takes advantage of the LSTM and multi-
head attention mechanisms. Dand et al. [8] proposed a dual-path transformer-
based full-band and sub-band fusion network (DPT-FSNet) for speech enhancement
in the frequency domain.

Speech typically requires long-range sequence modeling, convolutional neural
networks necessitate a greater number of convolutional layers to expand the
receptive field and thus continuously increase model complexity. In many
natural language processing and visual tasks, self-attention transformer-based
neural networks outperform deep learning models, which are constructed based
on convolutional neural networks (CNNs) [9]. Additionally, RNN models can be
commonly utilized for modeling long-term sequences with sequential
information, such as the Long Short-Term Memory (LSTM) and Gated Recurrent
Unit (GRU) [10]. However, RNN-based models are unable to be processed in
parallel due to the temporal structure of RNNs, leading to suboptimal
computational efficiency and increased complexity. Recent studies have
demonstrated that a self-attention methodology can be used for audio
denoising, which will relieve the challenges associated with modeling long-
sequence speech signals [4]. The ability to learn efficiently from raw audio
is crucial for constructing speech models for speech enhancement.

Inspired by the capability of the transformer in sequence modeling, we aim to
develop an efficient transformer approach that enhances audio quality by
learning general and meaningful speech features. We propose a dual-phase audio
transformer for denoising (DPATD) that incorporates explainable attention and
memory-compressed attention. The intuition is to divide the input audio into
shorter chunks and interleave dual-phase transformers, a local-chunk
transformer and a global-chunk transformer for local and global modeling,
respectively. Additionally, to solve low-occupancy or unneeded shared memory
reads and writes on the GPU, we employ FlashAttention-2, a unique attention
algorithm with superior work partitioning [11]. The contributions of this work
are threefold:

  1. 1.

We propose a dual-phase audio transformer that facilitates the modeling of
audio signals by organizing transformer layers. It first segments the input
audio into shorter chunks and employs dual-phase transformers in an
interleaved manner, namely, a local-chunk transformer for local modeling and a
global-chunk transformer for global modeling.

  2. 2.

We develop a novel transformer block that combines explainable attention and
memory-compressed attention to conserve memory resources and provide
interpretable attention maps that are noise-resistant and align with
informative input patterns naturally.

  3. 3.

Our extensive experiments on two benchmark datasets demonstrate that the DPATD
outperforms state-of-the-art methods across all evaluation criteria while
maintaining a relatively low model complexity.

Figure 1: System flowchart of the DPATD model. The segmentation stage splits
an audio input into chunks without overlaps and concatenates them to form a
2-D tensor. The transformer block, which consists of local-chunk and global-
chunk transformers, is applied to individual chunks in parallel to process
information. Multiple blocks are stacked to increase the total depth of the
network. A 2-D output of the last block is converted back to an audio output.

##  II Related work

Traditional speech denoising techniques. Traditional speech denoising
techniques primarily rely on statistical techniques that can be utilized to
build relevant denoising models and extract clean audio from noisy input
signals. The denoising performance can be improved by the Wiener filter [12].
Linden et al. [13] decomposed a spectral graph into a spectral basis matrix
and an encoding matrix. After the different sound sources are reconstructed
based on the clustering of the basis matrix and the corresponding encoding
information, the noise components are removed to facilitate more accurate
monitoring of biological sounds. Paliwal and Basu [14] proposed a Kalman
filtering method to improve speech enhancement performance. Ali et al. [15]
focused on the denoising of phonocardiogram (PCG) signals using different
families of discrete wavelet transforms, thresholding types and techniques,
and signal decomposition levels.

Deep learning for speech enhancement. Different speech enhancement models
based on deep neural networks have steadily taken center stage with the
advancement of deep learning technology. Based on different model inputs,
current voice augmentation techniques for DNNs can be broadly divided into two
groups: time-domain (T) techniques and time-frequency domain (TF) techniques.
Time-domain techniques employ an end-to-end model that directly estimates
clean waveforms using audio data in the time domain as raw waveform inputs.
The architecture foundation for time-domain approaches is WaveNet [16]. The
majority of speech enhancement techniques currently focus on the time-
frequency domain of speech. The TF techniques use the short-time Fourier
transform (STFT) and the inverse short-time Fourier transform (ISTFT). The
latest frequency-domain model, Band-Split RNN, explicitly splits the
spectrogram of the mixture into subbands and performs interleaved band-level
and sequence-level modeling for speech enhancement [17].

##  III Motivation

Most existing deep learning-based audio denoising methods study the magnitude
spectrum of images for audio denoising. However, these methods can be
constrained by computing power or limited filtering image regions, resulting
in low denoising performance. The transformer applications in audio
enhancement are still limited. Inspired by neural network approaches to text,
our model encodes audio information and is trained to understand what clean
audio should look like. Basically, we regard the audio signal as an ”audio
sequence” and further segment it into smaller chunks. The attention of each
audio chunk will be calculated based on other chunks in the given audio
sequence.

##  IV Model

In this section, we first review the audio denoising task, provide the
motivation for our model, then conduct an in-depth analysis of the
architecture of our dual-phase audio transformer for denoising (DPATD). Our
model first splits the input audio into several audio chunks and encodes the
audio sequence. Sequentially, generated sequence vectors are fed into a dual-
phase transformer to train to minimize the difference between denoised audio
and clean audio. Finally, we get the denoised audio as shown in Fig. 1.

We assume that the mixture speech signal $y(t)$ is a linear sum of the clean
speech signal $x(t)$ and noise $\varepsilon(t)$, and the noisy speech $y(t)$
can be typically expressed as Eq. (1):

$y(t)=x(t)+\varepsilon(t).$ (1)

###  IV-A Segmentation

Our DPATD framework first splits the input audio into several local
chunks,then calculates both representations and their relationship. The
segmentation stage encodes the audio sequence using patch embedding with the
maximum audio sequence size. The smaller utterances are zero-padded to match
the size of the maximum audio sequence. Given a sequence of acoustic input
vectors $Y(y_{1},y_{2},\cdots,y_{L})\in R^{1\times L}$ where $L$ is the audio
sequence length, the segmentation stage splits $Y$ into chunks of length $K$
and hop size $P$. Every sample in $Y$ appears and only appears in chunks,
generating $M$ equal size chunks $D_{m}\in R^{1\times K},m=1,\cdots,M$. All
chunks are then concatenated together to form a 2-D tensor
$T=[D_{1},\cdots,D_{M}]\in R^{K\times M}$. Positional embeddings are added to
patch embeddings to preserve positional information. We use standard learnable
1D position embeddings since the audio input is an ordered sequence. The
segmentation output tensor $T$ is then passed to the stack of $N$ transformer
blocks. Each block converts a 2-D tensor input to another tensor with the same
shape. We denote the input tensor for block $\textbf{Z}=1,\cdots,Z$ as
$T_{z}\in R^{K\times M}$, where $T_{1}=T.$ The segmentation phase of the model
is shown in Fig. 1.

###  IV-B Dual-Phase Audio Transformer

Basic Architecture. Since the input noisy audio and the output enhanced audio
have the same length, we introduce a simple but effective modification to the
transformer for audio sequences by removing the encoder module (almost
reducing model parameters by half for a given hyper-parameter set). Each
transformer block contains two sub-modules. As shown in Fig. 2, the first
module is a Memory-Compressed Explainable Multi-Heads Attention (MCE-MSA).
This attention consists of explainable multi-heads attention and memory-
compressed attention. The second module is a simple, position-wise, fully
connected feed-forward network. In addition, inspired by the effectiveness of
RNNs in tracking ordered sequential information, we replace the first fully
connected layer of the feed-forward network with a GRU layer to learn more
positional information. A residual connection was employed around each of the
two sub-modules, followed by layer normalization.

Figure 2: The architecture of the DPAT block. The Explainable Multi-Head
Attention (E-MHA) module is capable of providing interpretable attention maps
that are noise-resistant and align with informative input patterns naturally.
We utilize a strided convolution to limit the dot products between $Q$ and
$K$. We replace the first fully connected layer of the feed-forward network
with a GRU layer to learn more positional information.

###  IV-C Explainable Memory-Compressed Attention.

Memory-Compressed Attention To handle longer sequences, we modify the multi-
head attention to reduce memory usage by limiting the dot products between $Q$
and $K$ in Eq. (3). To achieve this goal, we take advantages of a strided
convolution with convolution kernels of size 3 with stride 3. Since the memory
cost of attention is constant for each block, this alteration allows us to
maintain the linear relationship between the number of activations and the
length of the sequence [18].

Explainable Multi-Heads Attention. In our work, attention blocks employ $h=8$
heads (the number of parallel attention layers) and map the input $h$ times to
get $Q$, $K$, and $V$ representations, respectively, as described in Eq. (2).
Given an input $Y$, each head $H_{h}$ holds an explainable attention weight
$A_{h}\in\mathcal{R}^{N\times d}$ that represents the relative importance of
input features. $A_{h}$ aims to learn explainable features for the output
through the MCE-MSA mechanism.

$Q_{i}=YW^{Q}_{i},\ K_{i}=YW^{K}_{i},\ V_{i}=YW^{V}_{i},$ (2)

where $Y\in R^{d\times k}$ is the input with sequence of length $L$ and
dimension $d$, $i=1,2,\cdots,h$ and $Q_{i},K_{i},V_{i}\in R^{l\times d/h}$ are
the mapped queries, keys and values respectively.
$W^{Q}_{i},W^{K}_{i},W^{V}_{i}\in R^{d\times d/h}$ denote the $i$-th linear
transformation matrix for queries, keys, and values, respectively.

The self-attention operation is constructed by Eq. (3). $W$ implies how much
attention is paid to each token.

$Att(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})V=softmax(W)V.$ (3)

The attention weight $A$ is defined as:

$A=\mathcal{L}(W+b)^{T},$ (4)

where $b$ is a trainable bias term, which is introduced as an initial
alignment for the input patterns. $\mathcal{L}$ is a non-linear function that
scales the L2 norm of its input.

In follows, the self-attention feature $P$ is formally expressed as:

$P=A^{T}V,$ (5)

According to Eq. (4), $\|A\|\leq 1$. There $P$ in Eq. (5) is upper-bounded as
follows:

$P=\|A\|\||V|\|cos(A,V)\leq\|V\|.$ (6)

When Eq. (6) is optimized, the attention weight $A$ is proportional to $V$. In
order to achieve maximal output, $A$ is driven to align with the
discriminative features in $V$, instead of the uninformative noise. Therefore,
$P$ can only achieve this upper bound if all possible solutions of $v\in V$
are encoded as eigenvectors in the weight $A$. This maximization suggests that
with the attention weight $A$, we will obtain an inherently explainable
decomposition of input patterns.

In our work, whole sequential explainable and memory-compressed transformer
blocks are computed as:

$\displaystyle S^{l}=MCE-MSA(Z^{l-1}),$ (7) $\displaystyle
Z^{l}=LayerNorm(Z+S^{l}),$ (8) $\displaystyle
FFN(Z^{l})=ReLU(GRU(Z^{l})W_{1}+b_{1})W_{2}+b_{2},$ (9) $\displaystyle
Output=LN(Z^{l}+FFN(Z^{l})),$ (10)

where $LayerNorm(\cdot)$ is the LayerNorm layer, $FFN(\cdot)$ denotes the
output of the position-wise feed-forward network, and $W_{1}\in
R^{d_{ff}\times d},b_{1}\in R^{d}$ and $d_{f}f=4\times d$.

###  IV-D Dual-phase Audio Transformer module

The dual-phase audio transformer module consists of four stacked dual-chunk
transformer blocks. Each block converts an input 2-D tensor into another
tensor with the same shape. We propose a dual-phase transformer block based on
explainable and memory-compressed attention. As shown in Fig.1, it has a
local-chunk transformer and a global-chunk transformer, which extract local
and global information, respectively. More specifically, the input is a 2-D
tensor ($[K,S]$), and the local-chunk transformer is first applied to
individual chunks to parallelly process inter information, which performs on
the last dimension $F$ of the input tensor. Then, the global-chunk transformer
is used to fuse the information of the output from the local-chunk transformer
to learn global dependency, which is implemented on the dimension of the
tensor. Besides, each transformer is followed by the group normalization
operation and utilizes residual connections.

Decoder. We use the patch-expanding layer in the decoder to upsample the
extracted deep features. The 2-D convolution with a filter size of (1, 1)
recovers the channel dimension of the enhanced speech feature into 1 and
produces the enhanced speech waveform by an overlap-add method.

Loss Function. The time-domain loss is based on the mean square error (MSE)
between the input clean audio $(x_{1},x_{2},\cdots,x_{N})$ and the predicted
audio$(\hat{x}_{1},\hat{x}_{2},\cdots,\hat{x}_{N})$. The model is optimized by
minimizing the MSE, which is defined as:

$Loss=\frac{1}{N}\sum_{i-1}^{N-1}(x_{i}-\hat{x}_{i})^{2},$ (11)

where $N$ denotes the number of samples.

The overall training algorithm is shown in Alg. 1.

Algorithm 1 DPATD: Dual-Phase Audio Transformer for Denoising. Batch of audio
input: $B(Y)=\\{Y^{1},...,Y^{n_{B}}\\}$, and their clean audio input
$B(X)=\\{X^{1},...,X^{n_{B}}\\}$, where ${n_{B}}$ is the total number of
batch. $I$ is the number of iterations

1: Input: Mixture audio signals $Y=\\{y_{i}\\}_{i=1}^{N}$ and clean audio
input $X=\\{x_{i}\\}_{i=1}^{N}$, where $N$ is the total sample number of
audios.

2: Output: Denoised audio signals: $\hat{X}$

3: Segmentation splits the input audio into a 2-D tensor
$T=[D_{1},\cdots,D_{M}]$ using 1-D CNNs.

4: for $iter=1$ to $I$ do

5: Derive batch-wise data: $Y^{k}$ and $X^{k}$ sampled from $B(Y)$ and $B(X)$

6: Optimize our audio generation model DPATD using Eq. (11)

7: end for

8: Output the denoised audio signals

##  V Experimental Procedures

###  V-A Dataset

VCTK+DEMAND dataset. We validate the effectiveness of our proposed model on a
standard speech dataset [19]. The clean speech datasets are selected from the
Voice Bank Corpus, including a training set of 11572 utterances from 28
speakers and a test set of 872 utterances from 2 speakers.

BirdSoundsDenoising. This dataset uses a variety of natural noises, such as
wind, rain, waterfalls, etc., in place of the normal intentionally generated
noise [20]. In particular, the dataset contains 14,120 audios from one second
to fifteen seconds and is a large-scale dataset of bird sounds collected,
containing 10,000/1,400/2,720 in training, validation, and testing sets,
respectively.

Figure 3: Denoising results comparisons. Raw audio is the original noisy
audio.

###  V-B Implementation details

We train the DPATD model for 100 epochs on mini-batches of 8 random samples,
where the segment chunk size is set to 1000. We used the Adam optimization
scheme with a maximum learning rate of 2.5e-4. Since the layernorm layer is
extensively used throughout the model, a simple weight initialization of
$N(0,0.02)$ was adequate. For the activation function, we used the Gaussian
Error Linear Unit (GELU). We used the learned position embeddings instead of
the sinusoidal version proposed in the original work. With a single NVIDIA GTX
3060 GPU and PyTorch to implement our model, it took around 180 hours of GPU
time to train our model. The smaller utterances in a batch are zero-padded to
the largest size of utterances. A dynamic strategy is used to adjust the
learning rate during the training stage [21].

Model configurations. The configuration of our model structure partially
follows the original transformer architecture. We trained a 12-layer decoder-
only transformer with self-attention heads (1000 dimensional states and 12
attention heads). Here, we focus on explainable and memory compression
attention to effectively compute and conserve memory. For the position-wise
feed-forward networks, we used 4000-dimensional inner states. Consider the sum
of the input audio length in a signal block denoted by $K\times S$ where the
hop size is set to be the same as the embedding size. It is simple to see that
$S=[L/K]$ where $[\cdot]$ is the ceiling function. To achieve the minimum
total input length $K+S=K+[L/K]$, $K$ is selected such that $K\approx
sqrt(5L)$. This gives us sublinear input length$(O(sqrt(L))$ rather than the
original linear input length $(O(L))$.

TABLE I: Results comparisons of different methods ($F1,IoU$, and $Dice$ scores are multiplied by 100. “$-$” means not applicable.  Networks | Validation | Test
---|---|---
$F1$ | $IoU$ | $Dice$ | $SDR$ | $F1$ | $IoU$ | $Dice$ | $SDR$
U2-Net [22] | 60.8 | 45.2 | 60.6 | 7.85 | 60.2 | 44.8 | 59.9 | 7.70
MTU-NeT [23] | 69.1 | 56.5 | 69.0 | 8.17 | 68.3 | 55.7 | 68.3 | 7.96
Segmenter [24] | 72.6 | 59.6 | 72.5 | 9.24 | 70.8 | 57.7 | 70.7 | 8.52
SegNet [25] | 77.5 | 66.9 | 77.5 | 9.55 | 76.1 | 65.3 | 76.2 | 9.43
DVAD [20] | 82.6 | 73.5 | 82.6 | 10.33 | 81.6 | 72.3 | 81.6 | 9.96
R-CED [26] | $-$ | $-$ | $-$ | 2.38 | $-$ | $-$ | $-$ | 1.93
Noise2Noise [27] | $-$ | $-$ | $-$ | 2.40 | $-$ | $-$ | $-$ | 1.96
TS-U-Net [28] | $-$ | $-$ | $-$ | 2.48 | $-$ | $-$ | $-$ | 1.98
DPATD | $-$ | $-$ | $-$ | 10.49 | $-$ | $-$ | $-$ | 10.43

###  V-C Result

Tab. II shows the comparison results of the VCTK+DEMAND dataset. Our model
surpasses most waveform-based approaches now in use in terms of the PESQ score
and achieves performance that is equivalent to other methods in other
evaluation metrics. For the BirdSoundsDenoising dataset, we report the
performance of eight state-of-the-art baselines. The results are shown in Tab.
I, where the bold text indicates the best outcomes for each statistic. The
results demonstrate that our model outperforms other state-of-the-art methods
in terms of the SDR. Results of F1, IoU, and Dice are not included because
these metrics are used for the audio image segmentation task [29, 30]. The
comparisons of raw bird audio, ground truth labeled denoised audio, and
denoised audio of other models are shown in Fig. 3. Additionally, our model
bears more resemblance to the labeled denoised signal. As a consequence, our
model enhances the audio-denoising capabilities of the BirdSoundDenoising
dataset.

TABLE II: Comparison results on the VoiceBank-DEMAND dataset. “$-$” means not applicable. Methods | Domain | PESQ | STOI | CSIG | CBAK | COVL
---|---|---|---|---|---|---
PGGAN [31] | T | 2.81 | 0.944 | 3.99 | 3.59 | 3.36
DCCRGAN [32] | TF | 2.82 | 0.949 | 4.01 | 3.48 | 3.40
S-DCCRN [33] | TF | 2.84 | 0.940 | 4.03 | 2.97 | 3.43
TSTNN [6] | T | 2.96 | 0.950 | 4.33 | 3.53 | 3.67
PHASEN [34] | TF | 2.99 | $-$ | 4.18 | 3.45 | 3.50
DEMUCS [35] | T | 3.07 | 0.95 | 4.31 | 3.40 | 3.63
SE-Conformer [36] | T | 3.13 | 0.95 | 4.45 | 3.55 | 3.82
MetricGAN+ [37] | TF | 3.15 | 0.927 | 4.14 | 3.12 | 3.52
MANNER [38] | T | 3.21 | 0.950 | 4.53 | 3.65 | 3.91
CMGAN [39] | T | 3.41 | 0.96 | 4.63 | 3.94 | 4.12
DPATD | T | 3.55 | 0.97 | 4.78 | 3.96 | 4.22

Ablation study. First, we examine the performance of our method by comparing
it with different transformer layers and the effect of chunks in segmentation.
In order to further demonstrate the effectiveness of our proposed transformer
block, we also designed another architecture for comparison. In this
architecture, we use different transformer blocks rather than the 12 blocks in
the DPATD, and we increase or reduce the number of heads. In addition, we also
set chunk sizes at 500 and 2000, while only 1000 in our DPATD. From
Tab.LABEL:tab:layers, a 12-transformer-block DPATD has better scores than
other models. As shown in Tab.LABEL:tab:chunks, the chunk has only a slight
influence on the performance, and we use chunk=1000 by default for its
efficiency.

TABLE III: Effect of layers and chunks in transformer block in the DPATD.

Heads | 6 | 8 | 12 | 16
---|---|---|---|---
PESQ | 2.99 | 3.23 | 3.45 | 3.15

Chunks | 500 | 1000 | 2000
---|---|---|---
PESQ | 3.26 | 3.45 | 3.38

##  VI Conclusion

In this study, we present a framework for using a dual-phase audio transformer
for denoising (DPATD) to provide robust speech enhancement. The DPATD splits
the audio input into non-overlapping chunks in the segmentation stage, which
are then passed as input to the transformer model. In a DPAT block, the local-
chunk transformer and global-chunk transformer process the local chunks and
all the chunks, respectively. We modified the transformer model using
explainable multi-head attention and memory-compressed attention. Extensive
experiments on datasets have demonstrated the effectiveness and superiority of
the proposed DPATD architecture. Finally, our method is still computationally
demanding, and future directions of the work could improve on these
limitations.

## References

  * [1] Wenbin Jiang, Zhijun Liu, Kai Yu, and Fei Wen.  Speech enhancement with neural homomorphic synthesis.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 376–380. IEEE, 2022.
  * [2] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [3] Zhengsu Chen, Lingxi Xie, Jianwei Niu, Xuefeng Liu, Longhui Wei, and Qi Tian.  Visformer: The vision-friendly transformer.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 589–598, 2021.
  * [4] Cem Subakan, Mirco Ravanelli, Samuele Cornell, Mirko Bronzi, and Jianyuan Zhong.  Attention is all you need in speech separation.  In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 21–25. IEEE, 2021.
  * [5] Guochen Yu, Andong Li, Hui Wang, Yutian Wang, Yuxuan Ke, and Chengshi Zheng.  Dbt-net: Dual-branch federative magnitude and phase estimation with attention-in-attention transformer for monaural speech enhancement.  IEEE/ACM Transactions on Audio, Speech, and Language Processing, 30:2629–2644, 2022.
  * [6] Kai Wang, Bengbeng He, and Wei-Ping Zhu.  Tstnn: Two-stage transformer based neural network for speech enhancement in the time domain.  In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 7098–7102. IEEE, 2021.
  * [7] Weiwei Yu, Jian Zhou, HuaBin Wang, and Liang Tao.  Setransformer: Speech enhancement transformer.  Cognitive Computation, pages 1–7, 2022.
  * [8] Feng Dang, Hangting Chen, and Pengyuan Zhang.  Dpt-fsnet: Dual-path transformer based full-band and sub-band fusion network for speech enhancement.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 6857–6861. IEEE, 2022.
  * [9] Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang.  Swin-unet: Unet-like pure transformer for medical image segmentation.  In Computer Vision–ECCV 2022 Workshops: Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part III, pages 205–218. Springer, 2023.
  * [10] Yi Luo, Zhuo Chen, and Takuya Yoshioka.  Dual-path rnn: efficient long sequence modeling for time-domain single-channel speech separation.  In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 46–50. IEEE, 2020.
  * [11] Tri Dao.  Flashattention-2: Faster attention with better parallelism and work partitioning.  arXiv preprint arXiv:2307.08691, 2023.
  * [12] Marwa A Abd El-Fattah, Moawad I Dessouky, Alaa M Abbas, Salaheldin M Diab, El-Sayed M El-Rabaie, Waleed Al-Nuaimy, Saleh A Alshebeili, and Fathi E Abd El-samie.  Speech enhancement with an adaptive wiener filter.  International Journal of Speech Technology, 17:53–64, 2014.
  * [13] Tzu-Hao Lin, Shih-Hua Fang, and Yu Tsao.  Improving biodiversity assessment via unsupervised separation of biological sounds from long-duration recordings.  Scientific reports, 7(1):4547, 2017.
  * [14] K Paliwal and Anjan Basu.  A speech enhancement method based on kalman filtering.  In ICASSP’87. IEEE International Conference on Acoustics, Speech, and Signal Processing, volume 12, pages 177–180. IEEE, 1987.
  * [15] Mohammed Nabih Ali, EL-Sayed A El-Dahshan, and Ashraf H Yahia.  Denoising of heart sound signals using discrete wavelet transform.  Circuits, Systems, and Signal Processing, 36:4482–4497, 2017.
  * [16] Daniel Stoller, Sebastian Ewert, and Simon Dixon.  Wave-u-net: A multi-scale neural network for end-to-end audio source separation.  arXiv preprint arXiv:1806.03185, 2018.
  * [17] Jianwei Yu and Yi Luo.  Efficient monaural speech enhancement with universal sample rate band-split rnn.  In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1–5. IEEE, 2023.
  * [18] Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer.  Generating wikipedia by summarizing long sequences.  arXiv preprint arXiv:1801.10198, 2018.
  * [19] Cassia Valentini-Botinhao, Xin Wang, Shinji Takaki, and Junichi Yamagishi.  Speech enhancement for a noise-robust text-to-speech synthesis system using deep recurrent neural networks.  In Interspeech, volume 8, pages 352–356, 2016.
  * [20] Youshan Zhang and Jialu Li.  Birdsoundsdenoising: Deep visual audio denoising for bird sounds.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 2248–2257, 2023.
  * [21] Jingjing Chen, Qirong Mao, and Dong Liu.  Dual-path transformer network: Direct context-aware modeling for end-to-end monaural speech separation.  arXiv preprint arXiv:2007.13975, 2020.
  * [22] Xuebin Qin, Zichen Zhang, Chenyang Huang, Masood Dehghan, Osmar R Zaiane, and Martin Jagersand.  U2-net: Going deeper with nested u-structure for salient object detection.  Pattern recognition, 106:107404, 2020.
  * [23] Hongyi Wang, Shiao Xie, Lanfen Lin, Yutaro Iwamoto, Xian-Hua Han, Yen-Wei Chen, and Ruofeng Tong.  Mixed transformer u-net for medical image segmentation.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 2390–2394. IEEE, 2022.
  * [24] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid.  Segmenter: Transformer for semantic segmentation.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 7262–7272, 2021.
  * [25] Vijay Badrinarayanan, Alex Kendall, and Roberto Cipolla.  Segnet: A deep convolutional encoder-decoder architecture for image segmentation.  IEEE transactions on pattern analysis and machine intelligence, 39(12):2481–2495, 2017.
  * [26] Se Rim Park and Jinwon Lee.  A fully convolutional neural network for speech enhancement.  arXiv preprint arXiv:1609.07132, 2016.
  * [27] Madhav Mahesh Kashyap, Anuj Tambwekar, Krishnamoorthy Manohara, and S Natarajan.  Speech denoising without clean training data: A noise2noise approach.  arXiv preprint arXiv:2104.03838, 2021.
  * [28] Eloi Moliner and Vesa Välimäki.  A two-stage u-net for high-fidelity denoising of historical recordings.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 841–845. IEEE, 2022.
  * [29] Junhui Li, Pu Wang, and Youshan Zhang.  Deeplabv3+ vision transformer for visual bird sound denoising.  IEEE Access, 2023.
  * [30] Youshan Zhang and Jialu Li.  Complex Image Generation SwinTransformer Network for Audio Denoising.  In Proc. INTERSPEECH 2023, pages 186–190, 2023.
  * [31] Yihao Li, Meng Sun, and Xiongwei Zhang.  Perception-guided generative adversarial network for end-to-end speech enhancement.  Applied Soft Computing, 128:109446, 2022.
  * [32] Huixiang Huang, Renjie Wu, Jingbiao Huang, Jucai Lin, and Jun Yin.  Dccrgan: Deep complex convolution recurrent generator adversarial network for speech enhancement.  In 2022 International Symposium on Electrical, Electronics and Information Engineering (ISEEIE), pages 30–35. IEEE, 2022.
  * [33] Shubo Lv, Yihui Fu, Mengtao Xing, Jiayao Sun, Lei Xie, Jun Huang, Yannan Wang, and Tao Yu.  S-dccrn: Super wide band dccrn with learnable complex feature for speech enhancement.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 7767–7771. IEEE, 2022.
  * [34] Dacheng Yin, Chong Luo, Zhiwei Xiong, and Wenjun Zeng.  Phasen: A phase-and-harmonics-aware speech enhancement network.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 9458–9465, 2020.
  * [35] Alexandre Défossez, Gabriel Synnaeve, and Yossi Adi.  Real time speech enhancement in the waveform domain.  2020\.
  * [36] Eesung Kim and Hyeji Seo.  Se-conformer: Time-domain speech enhancement using conformer.  In Interspeech, pages 2736–2740, 2021.
  * [37] Szu-Wei Fu, Cheng Yu, Tsun-An Hsieh, Peter Plantinga, Mirco Ravanelli, Xugang Lu, and Yu Tsao.  Metricgan+: An improved version of metricgan for speech enhancement.  arXiv preprint arXiv:2104.03538, 2021.
  * [38] Hyun Joon Park, Byung Ha Kang, Wooseok Shin, Jin Sob Kim, and Sung Won Han.  Manner: Multi-view attention network for noise erasure.  In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 7842–7846. IEEE, 2022.
  * [39] Ruizhe Cao, Sherif Abdulatif, and Bin Yang.  Cmgan: Conformer-based metric gan for speech enhancement.  arXiv preprint arXiv:2203.15149, 2022.


# Error-Driven Uncertainty Aware Training

Pedro Mendes Corresponding Author. Email<EMAIL_ADDRESS>Paolo Romano
David Garlan  Software and Societal Systems Department, Carnegie Mellon
University  INESC-ID and Instituto Superior Tecnico, Universidade de Lisboa

###### Abstract

Neural networks are often overconfident about their predictions, which
undermines their reliability and trustworthiness. In this work, we present a
novel technique, named Error-Driven Uncertainty Aware Training (EUAT), which
aims to enhance the ability of neural models to estimate their uncertainty
correctly, namely to be highly uncertain when they output inaccurate
predictions and low uncertain when their output is accurate. The EUAT approach
operates during the model’s training phase by selectively employing two loss
functions depending on whether the training examples are correctly or
incorrectly predicted by the model. This allows for pursuing the twofold goal
of i) minimizing model uncertainty for correctly predicted inputs and ii)
maximizing uncertainty for mispredicted inputs, while preserving the model’s
misprediction rate. We evaluate EUAT using diverse neural models and datasets
in the image recognition domains considering both non-adversarial and
adversarial settings. The results show that EUAT outperforms existing
approaches for uncertainty estimation (including other uncertainty-aware
training techniques, calibration, ensembles, and DEUP) by providing
uncertainty estimates that not only have higher quality when evaluated via
statistical metrics (e.g., correlation with residuals) but also when employed
to build binary classifiers that decide whether the model’s output can be
trusted or not and under distributional data shifts.

547

##  1 Introduction

Deep Neural Networks (DNNs) have achieved remarkable performance across
various domains and are increasingly utilized to automate intricate decision-
making processes. However, a critical limitation of current neural models is
their tendency to display overconfidence in their predictions [11, 2]. This
overconfidence persists even when erroneous predictions are made, ultimately
compromising the reliability and trustworthiness of the models.

Recent research efforts [10, 5, 9, 12, 34] have been dedicated to enhancing
the trustworthiness of DNNs by estimating the model’s predictive uncertainty
through various approaches. Bayesian neural networks (BNNs) [34, 43, 44] offer
an elegant framework for modeling uncertainty [38]. However, while BNNs
provide theoretically sound uncertainty estimates, they incur prohibitive
costs, being impractical for large datasets and complex models. To mitigate
these challenges, various approximations have been introduced. For instance,
Monte Carlo (MC) dropout [10], which leverages dropout regularization during
both training and inference stages to approximate the behavior of BNNs.

Furthermore, numerous studies [27, 23, 45, 41, 33] have focused on calibrating
the models’ predicted uncertainty in different ways. These works can be
categorized into two primary groups based on whether they: 1) account for the
uncertainty during training by introducing an additional term in the loss
function to quantify the model’s predictive uncertainty, or ii) implement a
post-processing stage to calibrate the model’s predicted probabilities using a
validation set. Although post-processing methods have empirically proven to be
effective and cost-efficient [27], they present additional calibration
parameters that are sensitive to the method and data used. On the other hand,
despite being more expensive, learning-based methods have achieved better
performance for uncertainty estimation [26, 32].

In this work, we mix both approaches by introducing Error-Driven Uncertainty
Aware Training (EUAT), a specialized training procedure for classification
tasks that aims at improving the model’s uncertainty estimation by imposing
high uncertainty for erroneous outputs and low uncertainty for accurate
predictions. To achieve this two-fold goal, during training, EUAT iterates
between two loss functions depending on whether the training examples are
correctly or incorrectly predicted by the model. More in detail, our approach
extends a base loss function, which aims to minimize the classification error
rate (e.g., cross-entropy (CE)), with an additional term whose objective is to
maximize the model’s uncertainty for misclassifications and minimize
uncertainty for correct classified inputs. However, to separate the correctly
and incorrectly classified inputs and speed up the training procedure, we
first consider pre-train the model and then we apply EUAT to conduct a post-
learning-based phase to improve its uncertainty.

We conducted an extensive evaluation of EUAT on classification tasks using
popular image recognition models and benchmarks, where we compared our
approach against several state-of-the-art methods for uncertainty estimation
using six different evaluation metrics. Further, we extended our assessment to
a binary classification problem, which presents a particularly interesting
case involving the class inversion of the high uncertainty outputs that are
likely to be wrong classified. We also evaluate our technique in an out-of-
distribution detection task, where corrupted inputs are used to evaluate the
model, and at last, we integrate our function into adversarial training
settings in order to identify possible misclassifications based on
uncertainty. We detail the challenges encountered in each domain/task. In
general, EUAT presents the best performance in more than 60% of the metrics
considered, and in the majority of the cases where the baselines are more
competitive, EUAT is still able to achieve similar performance metrics.
Further, in all the scenarios considered, we show that our strategy can better
separate wrong and accurate predictions based on uncertainty, increasing the
reliability and trustworthiness of the models.

##  2 Related Work

In this section, we first review different formulations of the problem of
estimating models’ uncertainty, and the corresponding metrics, proposed in the
literature. Subsequently, we analyze existing methods to estimate the
uncertainty of DNNs. Finally, we discuss methods aimed at improving
uncertainty estimation by adjusting the model’s outputs via post-processing or
uncertainty-aware training techniques.

##### Problem definition and metrics.

The problem of accurately estimating model uncertainty has been formalized
using two main theoretical frameworks. One such formalization is based on the
notion of calibration, which aims at aligning the probabilities output by the
model with the true likelihood of the predicted outcomes [14]. An alternative
formalization is based on the Uncertainty Confusion Matrix (UCM) [3, 20], as
defined in Table 1. The UCM specializes the concept of confusion matrix to
evaluate the ability to leverage the model’s uncertainty to discern correct
predictions. For example, different metrics such as the expected calibration
error (ECE) [33], adaptive calibration error [36], or test-based calibration
error [31], have been proposed to measure the calibration error of a model. On
the other hand, leveraging the UCM, several works [20, 3] have exploited
additional metrics such as uncertainty accuracy (uA) and the uncertainty area
under the curve (uAUC) to enhance the reliability of uncertainty estimates.

Table 1: Uncertainty Confusion Matrix Correctness |  | Uncertainty
---|---|---
| Certain | Uncertain
Correct |  True Certainty TC |  False Uncertainty FU
Wrong |  False Certainty FC |  True Uncertainty TU

##### Uncertainty Estimation.

One important foundation of these works lies in the computation of
uncertainty. Uncertainty in DNNs plays a pivotal role in quantifying the
reliability and robustness of their predictions. There are two main types of
uncertainty: epistemic uncertainty, associated with the model’s lack of
knowledge or data, and aleatoric uncertainty, linked to the inherent
randomness and unpredictability within the training data. Various metrics are
employed to measure these types of uncertainty such as predictive entropy (PE)
and mutual information (MI) [1]. However, quantifying uncertainty with DNNs is
a challenging task. Bayesian methods [34, 28] can directly provide an estimate
of the uncertainty by parameterizing the parameters of the network with
distributions. However, training BNNs usually comes with a prohibitive cost.
Thus, several approximations have been developed. Monte Carlo (MC) dropout
[10], which is one of the most popular techniques for uncertainty
quantification on DNNs [1, 4], proposed as a Bayesian approximation to
estimate the uncertainty by sampling multiple dropout masks during training
and inference phases and aggregating the predictions, providing a
probabilistic estimate that allows the quantification of uncertainty. One can
also approximate and estimate model uncertainty by computing the variance of
the output predictions [5, 9, 25]. Additionally, Markov Chain Monte Carlo
(MCMC) methods [12] offer another avenue for drawing the posterior
distribution (albeit, those require a long time to converge to the final
distribution [1, 34]). Further, Variational Inference (VI) is a technique used
to estimate the uncertainty of BNNs by approximating the posterior
distribution over the model’s weights, which are treated as random variables
with associated probability distributions. Training the network to approximate
these distributions enables the capture of inherent uncertainty in the model’s
predictions. Moreover, we can also resort to models that output an uncertainty
estimation, such as Deep Gaussian Processes [6], or use Laplace approximations
for uncertainty [29], or deep ensemble methods that offer yet another approach
to estimate model uncertainty [24, 1], by aggregating the outputs of each
learner in the ensemble and computing relevant metrics such as the entropy or
MI. At last, to account for the computation overhead of the previous methods
(resulting, for example, from training multiple models to the ensemble or
performing several forward passes using MC dropout), DEUP [23] trains a new
simple estimator to predict the uncertainty of the base model exploiting its
error/loss, and DUN [2] leverages the outputs of different layers of a DNN to
measure the uncertainty.

In this work, we resort to MC dropout to estimate the uncertainty of our
models and compute the PE using the outputted distributions. Distinctly, we
explicitly aim to increase the uncertainty of misclassifications by maximizing
the PE of the wrong prediction, while minimizing the global error rate.

##### Post-processing Calibration Methods.

Model calibration becomes especially critical in tasks where estimated
uncertainties play a fundamental role in risk assessment. Thus, a first set of
works aims at calibrating a fully trained model by applying a post-processing
phase to align the output probabilities with the true likelihood of the
predicted outcomes of events occurring [14]. DNNs are not inherently
calibrated, and several techniques, such as Platt Scaling [39], Isotonic
Regression [45], Temperature Scaling [14], or Beta Calibration [22], can be
employed to fine-tune the probabilities outputted by the model ensuring a more
accurate alignment with true outcome frequencies. Moreover, Krishnan et al.
[20] introduced an accuracy versus uncertainty calibration (AvUC) loss
function devised to obtain well-calibrated uncertainties while simultaneously
preserving or enhancing model accuracy, and they extended their approach by
proposing a post-hoc calibration phase that extends Temperature Scaling using
AvUC. In addition, Karandikar et al. [18] proposed an extension of ECE and
AvUC by developing a soft version of the binning operation underlying these
calibration-error estimators, and also extended their approach for post-
processing calibration by optimizing the temperature parameter in the
temperature scaling method based on the soft calibration error. Complementary,
Gupta et al. [15] presented a binning-free calibration approach. However,
these calibration methods can be very sensitive to both the model and the
validation set [27] and perform suboptimally when faced with shifts in data
distribution [37].

##### Uncertainty Aware Training Methods.

One fundamental aspect of training DNNs is the choice of a loss function.
Although CE stands out as a common choice for addressing classification
problems, it tends to increase the over-confidence of the resulting model [14,
32]. Thus, since accounting for predictive uncertainty during training
improves model calibration [20], several loss functions have been developed to
extend existing ones, aiming to incorporate the model’s uncertainty. These
learning-based methods such as focal loss [26] or label smoothing [26]
introduce an additional term addressing model uncertainty. These terms are
typically balanced using corresponding weights introducing as a consequence
additional hyper-parameters. Similarly, Shamsi et al. [41] proposed two loss
functions that extend the CE by adding a new term to address the model’s
uncertainty that can be determined through the PE or the ECE. Additionally,
Einbinder et al [8] introduced an uncertainty-aware conformal loss function CE
loss by adding a new term that quantifies uncertainty via conformal
prediction. CALS [27] exploits the Augmented Lagrangian Multiplier method to
adaptively learn the weights of the penalties to balance each term in the new
loss function. Separately, class uncertainty-aware (CUA) loss [19] tailored
object detection introduces the uncertainty of each class to augment the loss
value when prediction results are uncertain. Further, Ding et al. [7]
developed an uncertainty-aware loss for selective medical image segmentation
that considers uncertainty in the training process to directly maximize the
accuracy on the confident segmentation subset, rather than the accuracy on the
whole dataset. Differently from the aforementioned works, our approach takes a
distinct path by focusing on leveraging a pre-trained model to deliberately
increase the uncertainty associated with incorrectly classified inputs without
degrading the overall error rate.

Further, through our novel method, the objective is to empower Machine
Learning (ML) systems to recognize possible model misclassification in
production and take customized actions accordingly. This idea can be further
extended to adversarial training scenarios, where the deliberate increase in
uncertainty for misclassified adversarial inputs enhances the system’s ability
to detect and respond to potential attacks in production.

##  3 Error-Driven Uncertainty Aware Training

DNNs often exhibit overconfidence in their predictions [11, 2]. Thus, in this
work, we address this problem by proposing a loss function called Error-Driven
Uncertainty Aware Training that shifts the focus toward refining the
uncertainty guarantees of a pre-trained model by increasing the uncertainty
associated with misclassifications while reducing the error rate and
uncertainty of correct predictions. We assert that if a model reaches a
stagnation phase in terms of quality and the pursuit of further training
offers negligible benefit, it should, at the very least, produce high
uncertainty values for incorrect predictions.

In order to achieve our design goals, we start the process by querying a pre-
trained model to determine which inputs of the training dataset are wrong and
correctly classified. Subsequently, we create two sets, one with incorrect
classified inputs $\mathcal{W}$ and the other with the correct ones
$\mathcal{C}$. We resize the correct classified inputs’ set by randomly
selecting samples until the two sets have the same size. Then, in order to
reduce overfitting, we mix wrong and correct classified inputs. Our approach
employs distinct loss functions for each set. Since our objective is to
deliberately increase the uncertainty of misclassifications, we minimize the
CE and maximize the uncertainty (determined via MC dropout and PE) for the
wrong-classified inputs while, for the correct-classified inputs, we minimize
the CE and the uncertainty, i.e.,

$L_{\text{EUAT}}(f_{\theta}(\textbf{x}),\textbf{y})\text{=}\begin{cases}L_{\text{CE}}(f_{\theta}(\textbf{x}),\textbf{y})\text{-}L_{U}(f_{\theta}(\textbf{x}),\textbf{y})&\forall(x,y)\in\mathcal{W}\\\
L_{\text{CE}}(f_{\theta}(\textbf{x}),\textbf{y})\text{+}L_{U}(f_{\theta}(\textbf{x}),\textbf{y})\hskip
5.69054pt&\forall(x,y)\in\mathcal{C}\end{cases}$ (1)

where the cross-entropy loss is given by

$L_{\text{CE}}(f_{\theta}(x),y)=-\frac{1}{K}\sum_{i=1}^{K}t(x_{i})\log(f_{\theta}(x_{i})),$
(2)

($t(x)$ denotes the true label given the input $x$), and the uncertainty loss
is measured by resorting to predictive entropy $H$

$L_{U}(f_{\theta}(x),y)=H[P(y|x)]=-\sum_{y\in\mathcal{Y}}P(y|x)\log P(y|x),$
(3)

where $P(y|x)$ is the model’s output distribution over the set of possible
outcomes $\mathcal{Y}$ obtained via MC dropout by approximating the model’s
output predictions using the average across parameters $\theta_{i}$ sampled
from a dropout distribution

$p(y|\mathcal{D},x)\approx\frac{1}{N}\sum_{i=1}^{N}p(y|\theta_{i},x).$ (4)

We decided to also minimize uncertainty for correct-classified inputs in order
to avoid a peak in the global model’s uncertainty. By using this approach, the
model is able to better separate the erroneous and correct predictions using
the uncertainty of each individual forecast. We also implement an early
stopping policy based on the different evaluation metrics considered. It
should also be noted that this proposed loss function is differentiable and
continuous. The pseudo-code of the method proposed is described in Algorithm
1.

Algorithm 1 Pseudo-code to train a model with EUAT loss function.

1: Input: model $f$, training set $\mathcal{S}$, optimizer opt

2: Output: model $f$

3: no_inputs = 0

4: while True do $\triangleright$ spot condition

5: $\mathcal{W}$, $\mathcal{C}$ $\leftarrow$ $f$($\mathcal{S}$)
$\triangleright$ Wrong $\mathcal{W}$ & correct $\mathcal{C}$ classified sets

6: $\mathcal{C}$ $\leftarrow$ Random($\mathcal{C}$, $|\mathcal{W}|$)

7:
$L_{\mathcal{W}}(f(\textbf{x}),\textbf{y})$=$L_{CE}(f(\textbf{x}),\textbf{y})$-$L_{U}(f(\textbf{x}),\textbf{y})$
$\forall(\textbf{x},\textbf{y})$$\in$$\mathcal{W}$ … $\triangleright$ Minimize
CE-PE for wrong outputs

8:
$L_{\mathcal{C}}(f(\textbf{x}),\textbf{y})$=$L_{CE}(f(\textbf{x}),\textbf{y})$+$L_{U}(f(\textbf{x}),\textbf{y})$
$\forall(\textbf{x},\textbf{y})$$\in$$\mathcal{C}$ … $\triangleright$ Minimize
CE+PE for correct outputs

9:
$L(f(\textbf{x}),\textbf{y})$=$L_{\mathcal{W}}(f(\textbf{x}),\textbf{y})$+$L_{\mathcal{C}}(f(\textbf{x}),\textbf{y})$
$\triangleright$ Add losses

10: $L$.backward() $\triangleright$ Gradient estimation

11: opt.step() $\triangleright$ Back-propagation computation

12: no_inputs += ($|\mathcal{W}|$ \+ $|\mathcal{C}|$)

13: if no_inputs $>$ epoch_size then

14: $f$.test() $\triangleright$ Test $f$ if trained with the same amount of
data as the dataset

15: end if

16: end while

17: return $f$

By resorting to EUAT, which effectively distinguishes between the uncertainty
of accurate and erroneous predictions, we enhance the quality of uncertainty
estimates and ultimately improve the reliability and trustworthiness of the
model. Operationalizing the model in production involves assessing the
uncertainty associated with each prediction. When the prediction falls below
an uncertainty threshold, the model outcome can be trusted. Otherwise, when
the uncertainty is above the threshold, the prediction is untrustworthy,
warranting further scrutiny or review by human evaluators. Consequently,
tuning the uncertainty threshold becomes fundamental. Akin to the
implementation of the early stopping policy, we also explore various
evaluation metrics to select the threshold.

In this work, we have considered classification tasks using different image
recognition models. Thus, given the model’s prediction, we can promptly
categorize it as either accurate or erroneous, and place it in the respective
set. The early stopping policy and threshold selection are individually
optimized for each metric.

We also consider a binary classification problem where the high-uncertainty
output predictions that are likely to be wrongly classified can be inverted to
the opposite class. Given the particularities of this problem, we tune the
uncertainty threshold for EUAT differently by maximizing the number of correct
predictions (i.e., true positives and true negatives) that are below the
threshold and the number of incorrect predictions that are above the threshold
(false positives and false negatives). Then, in production, when the outputted
uncertainty is larger than the threshold, we can directly flip the prediction
for the opposite class, and this way improve the model’s quality.

At last, we can also use EUAT to perform adversarial training. It should be
noted that our approach does not aim to identify adversarial attacks using
uncertainty; rather, when applied to AT focuses on identifying
misclassifications based on the output uncertainty, regardless of whether they
result from an attack or not. Equivalently, we can separate the wrong and
correct predictions (clean and adversarial) in the two sets and train directly
with our loss function. Since the function is differentiable, we can solve the
optimization problem to find the perturbation using adversarial training
methods like Fast Gradient Descent Method (FGSM) [13] or Project Gradient
Descent (PGD) [30]. Similar to Smith and Gal [42], we acknowledge that neither
dropout nor our method alone can be considered a reliable adversarial defense.
Rather, we advocate for the joint use of all these methods when training a
model adversarially in order to achieve better guarantees and trade-offs of
accuracy, robustness, and uncertainty, yielding resilient models harder to
attack.

##  4 Evaluation

In this section, we report the evaluation of the EUAT on a variety of domains
and tasks.

###  4.1 Experimental Setup, Benchmarks, and Baselines

To evaluate EUAT, we employed four models and datasets widely used in the
image recognition domain namely, ResNet50 [16] with ImageNet [40], Wide-
ResNet-28x10 [46] with Cifar100 [21], ResNet18 with Cifar10, and ResNet18 with
SVHN [35]. We also considered a binary classification model (using ResNet18
with Cifar10 to verify if there is a cat in an image), and an out-of-
distribution (OOD) detection task, where corrupted inputs using distributional
data shifts are used to evaluate the model. Finally, we evaluate our approach
in adversarial training settings using three models/datasets mentioned above
in the image recognition domain.

All the models and training procedures were implemented in Python3 via the
Pytorch framework and trained using a single Nvidia RTX A4000.

To train the models, we used a dropout rate of 0.3 and resorted to stochastic
gradient descent to minimize the loss function using a momentum of 0.9 and a
batch size of 64 for all the models, a learning rate of 0.01 and weight decay
of $10^{-5}$ for ResNet50/ImageNet and 0.1 and 0 for the remaining ones,
respectively. Before training the model using EUAT, we pre-trained the models
using CE loss and then decreased the learning rate by $10^{3}\times$ when
applying EUAT. Additionally, we exploited automatic mixed precision to train
the ResNet50/ImageNet and Wide-ResNet/Cifar100. The models were trained during
60 epochs (except in the binary classification problem where it was trained
for 200 epochs). More in detail, we pre-trained the models for 30 epochs (100
epochs in the binary classification problem) before starting the second phase
of training where we applied the EUAT. The implementation of the training
pipeline and additional information to ensure the reproducibility of results
are provided in the supplemental material.

We compared EUAT against the CE loss, model calibration, DEUP [23], an
ensemble of five learners [24], CALS [27], and a loss function incorporating
both CE and PE (CE+PE) [41]. To calibrate the model and train DEUP’s
additional error predictor, we created a validation set comprising 10%
randomly selected samples from the test set. Further, we resort to Isotonic
regression [45] to calibrate the model, achieving superior results compared to
other methods like Platt scaling [39], temperature scaling [14], and beta
calibration [22]. Although we experimented DEUP with different validation set
sizes, to ensure fairness, we maintained consistency by employing the same
validation set size in both cases.

We evaluate the different baselines using six different metrics. First, we
report the uncertainty accuracy (uA) (Eq. 5) and the uncertainty area under
the curve (uAUC), which are computed based on the Uncertainty Confusion Matrix
[3] defined in Table 1.

$uA=\frac{TC+TU}{TC+TU+FC+FU}$ (5)

We also evaluate the models using the correlation between the residuals of the
model and predicted uncertainties (Corr. w/ res.) [23], and the Wasserstein
distance of the uncertainty between the sets of correct and wrong predictions
(Wasser. dist.). At last, we report the ECE and the model’s error. The
uncertainty/confidence of the models is always computed via MC dropout using
the normalized PE, except for DEUP, which resorts to the loss values of the
base model to estimate the quality of its predictions (for a fair comparison,
after testing the model trained with DEUP, we had to normalize the loss
values). Moreover, we optimize each metric considered independently by
employing an early stopping policy and tuning the uncertainty threshold, and
we report the best values obtained for each metric.

Table 2: Comparison of EUAT against the baselines using different evaluation metrics and considering four benchmarks. Benchmark | Baseline | uA | uAUC | Corr. w/ res. | Wasser. dist. | ECE | Error
---|---|---|---|---|---|---|---
ResNet50/ImageNet | EUAT | 0.804 | 0.878 | 0.655 | 0.301 | 0.223 | 0.439
CE | 0.749 | 0.812 | 0.546 | 0.220 | 0.224 | 0.513
Calibration | 0.743 | 0.807 | 0.538 | 0.193 | 0.274 | 0.535
DEUP | 0.580 | 0.591 | 0.291 | 0.030 | 0.429 | 0.522
Ensemble | 0.745 | 0.809 | 0.536 | 0.213 | 0.245 | 0.506
CALS | 0.746 | 0.809 | 0.542 | 0.214 | 0.235 | 0.525
CE+PE | 0.755 | 0.786 | 0.530 | 0.141 | 0.382 | 0.524
Wide-ResNet/Cifar100 | EUAT | 0.858 | 0.891 | 0.711 | 0.216 | 0.162 | 0.273
CE | 0.794 | 0.774 | 0.546 | 0.128 | 0.235 | 0.296
Calibration | 0.787 | 0.836 | 0.580 | 0.226 | 0.146 | 0.312
DEUP | 0.699 | 0.601 | 0.300 | 0.052 | 0.252 | 0.328
Ensemble | 0.742 | 0.696 | 0.466 | 0.089 | 0.294 | 0.332
CALS | 0.813 | 0.831 | 0.589 | 0.233 | 0.114 | 0.264
CE+PE | 0.779 | 0.740 | 0.518 | 0.099 | 0.252 | 0.300
ResNet18/Cifar10 | EUAT | 0.914 | 0.921 | 0.626 | 0.410 | 0.018 | 0.099
CE | 0.905 | 0.866 | 0.576 | 0.273 | 0.025 | 0.103
Calibration | 0.898 | 0.893 | 0.533 | 0.329 | 0.030 | 0.108
DEUP | 0.917 | 0.563 | 0.297 | 0.032 | 0.039 | 0.094
Ensemble | 0.919 | 0.840 | 0.545 | 0.224 | 0.039 | 0.087
CALS | 0.907 | 0.875 | 0.563 | 0.283 | 0.012 | 0.101
CE+PE | 0.907 | 0.837 | 0.566 | 0.213 | 0.052 | 0.102
ResNet18/SVHN | EUAT | 0.960 | 0.927 | 0.638 | 0.479 | 0.011 | 0.047
CE | 0.956 | 0.841 | 0.572 | 0.232 | 0.021 | 0.047
Calibration | 0.953 | 0.902 | 0.537 | 0.346 | 0.026 | 0.051
DEUP | 0.960 | 0.564 | 0.312 | 0.040 | 0.024 | 0.044
Ensemble | 0.960 | 0.756 | 0.515 | 0.162 | 0.029 | 0.043
CALS | 0.963 | 0.867 | 0.569 | 0.264 | 0.011 | 0.041
CE+PE | 0.959 | 0.799 | 0.547 | 0.184 | 0.029 | 0.045

(a) ResNet50/ImageNet

(b) Wide-ResNet/Cifar100

(c) ResNet18/Cifar10

(d) ResNet18/SVHN

Figure 1: Normalized uncertainty distribution of correct and incorrect predictions for the different baselines (the average value of each distribution is marked with a black square, and the median with a pink triangle). Table 3: Comparison of EUAT against the baselines using different evaluation metrics and considering a binary classification problem. Baseline | uA | uAUC | Corr. | Wasser. | ECE | Error | Error | F1 | Precision | TPR | TNR
---|---|---|---|---|---|---|---|---|---|---|---
w/ res. | dist. | w/o flip | w/ flip
EUAT | 0.861 | 0.816 | 0.446 | 0.405 | 0.112 | 0.152 | 0.139 | 0.860 | 0.863 | 0.858 | 0.864
CE | 0.845 | 0.759 | 0.435 | 0.350 | 0.113 | 0.152 | 0.221 | 0.763 | 0.823 | 0.712 | 0.847
Calib. | 0.783 | 0.779 | 0.366 | 0.356 | 0.313 | 0.159 | 0.290 | 0.733 | 0.670 | 0.809 | 0.603
DEUP | 0.844 | 0.591 | 0.323 | 0.049 | 0.124 | 0.172 | 0.178 | 0.812 | 0.859 | 0.771 | 0.874
Ensemble | 0.861 | 0.683 | 0.358 | 0.216 | 0.112 | 0.147 | 0.163 | 0.834 | 0.847 | 0.822 | 0.852
CALS | 0.831 | 0.748 | 0.407 | 0.326 | 0.119 | 0.164 | 0.163 | 0.834 | 0.847 | 0.822 | 0.852
CE+PE | 0.811 | 0.690 | 0.360 | 0.274 | 0.153 | 0.184 | 0.225 | 0.759 | 0.817 | 0.709 | 0.842
Table 4: Comparison of EUAT against the baselines using different evaluation metrics and tested with out-of-distribution samples with Cifar10. Baseline | uA | uAUC | Corr. | Wasser. | ECE | Error
---|---|---|---|---|---|---
w/ res. | dist.
EUAT | 0.754 | 0.796 | 0.529 | 0.255 | 0.143 | 0.489
CE | 0.691 | 0.676 | 0.311 | 0.126 | 0.292 | 0.539
Calib. | 0.619 | 0.553 | 0.110 | 0.036 | 0.497 | 0.619
DEUP | 0.617 | 0.509 | 0.072 | 0.004 | 0.551 | 0.617
Ensemble | 0.753 | 0.734 | 0.426 | 0.177 | 0.237 | 0.555
CALS | 0.686 | 0.716 | 0.383 | 0.147 | 0.216 | 0.464
CE+PE | 0.681 | 0.663 | 0.299 | 0.103 | 0.418 | 0.568
Table 5: Comparison of EUAT against the baselines considered using different evaluation metrics in the adversarial training scenario using three benchmarks. |  | Standard Settings | Adversarial Settings
---|---|---|---
Benchmark | Baseline | uA | uAUC | Corr. w/ | Wasser. | ECE | Error | uA | uAUC | Corr. w/ | Wasser. | ECE | Error
res. | dist. | res | dist.
Cifar100 | EUAT | 0.850 | 0.902 | 0.709 | 0.254 | 0.154 | 0.303 | 0.791 | 0.853 | 0.643 | 0.224 | 0.268 | 0.436
CE | 0.786 | 0.793 | 0.565 | 0.154 | 0.241 | 0.329 | 0.694 | 0.716 | 0.430 | 0.104 | 0.377 | 0.460
Calibration | 0.771 | 0.839 | 0.576 | 0.271 | 0.087 | 0.335 | 0.734 | 0.801 | 0.529 | 0.222 | 0.202 | 0.456
DEUP | 0.677 | 0.585 | 0.222 | 0.036 | 0.235 | 0.337 | 0.577 | 0.584 | 0.214 | 0.033 | 0.391 | 0.493
Ensemble | 0.800 | 0.817 | 0.597 | 0.161 | 0.183 | 0.288 | 0.865 | 0.852 | 0.750 | 0.194 | 0.334 | 0.436
CALS | 0.788 | 0.787 | 0.567 | 0.159 | 0.245 | 0.332 | 0.689 | 0.714 | 0.430 | 0.105 | 0.383 | 0.466
CE+PE | 0.773 | 0.750 | 0.535 | 0.113 | 0.286 | 0.354 | 0.625 | 0.633 | 0.346 | 0.059 | 0.422 | 0.482
Cifar10 | EUAT | 0.905 | 0.921 | 0.635 | 0.414 | 0.031 | 0.117 | 0.845 | 0.883 | 0.601 | 0.353 | 0.031 | 0.196
CE | 0.895 | 0.898 | 0.571 | 0.306 | 0.011 | 0.116 | 0.826 | 0.847 | 0.542 | 0.259 | 0.028 | 0.198
Calibration | 0.890 | 0.893 | 0.528 | 0.339 | 0.051 | 0.119 | 0.819 | 0.845 | 0.513 | 0.278 | 0.019 | 0.201
DEUP | 0.904 | 0.591 | 0.235 | 0.044 | 0.018 | 0.104 | 0.804 | 0.603 | 0.237 | 0.034 | 0.092 | 0.205
Ensemble | 0.923 | 0.894 | 0.535 | 0.310 | 0.007 | 0.081 | 0.901 | 0.881 | 0.780 | 0.348 | 0.083 | 0.179
CALS | 0.895 | 0.898 | 0.579 | 0.317 | 0.009 | 0.117 | 0.833 | 0.842 | 0.542 | 0.264 | 0.015 | 0.196
CE+PE | 0.895 | 0.850 | 0.564 | 0.232 | 0.046 | 0.119 | 0.819 | 0.789 | 0.505 | 0.177 | 0.104 | 0.197
SVHN | EUAT | 0.944 | 0.947 | 0.649 | 0.447 | 0.030 | 0.072 | 0.794 | 0.848 | 0.588 | 0.307 | 0.020 | 0.281
CE | 0.931 | 0.893 | 0.590 | 0.335 | 0.009 | 0.080 | 0.749 | 0.773 | 0.473 | 0.196 | 0.134 | 0.307
Calibration | 0.929 | 0.916 | 0.521 | 0.428 | 0.138 | 0.082 | 0.758 | 0.815 | 0.516 | 0.261 | 0.020 | 0.299
DEUP | 0.933 | 0.560 | 0.258 | 0.035 | 0.029 | 0.071 | 0.667 | 0.535 | 0.302 | 0.022 | 0.313 | 0.354
Ensemble | 0.951 | 0.893 | 0.545 | 0.333 | 0.009 | 0.053 | 0.881 | 0.827 | 0.681 | 0.221 | 0.234 | 0.308
CALS | 0.940 | 0.894 | 0.592 | 0.340 | 0.009 | 0.069 | 0.820 | 0.757 | 0.481 | 0.206 | 0.117 | 0.200
CE+PE | 0.932 | 0.874 | 0.596 | 0.274 | 0.021 | 0.078 | 0.823 | 0.736 | 0.449 | 0.145 | 0.127 | 0.187

###  4.2 Experimental results

Next, we report the results obtained using EUAT in the different domains
evaluated.

####  4.2.1 Image Recognition Models

We start by reporting in Table 2 the results obtained using four
models/datasets for image recognition: ResNet50/ImageNet, Wide-
ResNet/Cifar100, ResNet18/Cifar10, and ResNet18/SVHN. Across all baselines and
metrics, it is evident that EUAT consistently demonstrates superior
performance, outperforming others in 16 out of 24 cases. Notably, in the cases
where the baselines are more competitive, the differences are marginal (e.g.,
the error using EUAT increases by 3.4%, 13.8%, and 14.6% compared to the best
baseline training a Wide-ResNet/Cifar100, ResNet18/Cifar10, and ResNet18/SVHN,
respectively).

EUAT presents the best uA when training a ResNet50 with ImageNet and a Wide-
ResNet with Cifar100, and yields minimal reduction of 0.5% and 0.3% in the uA
when training a ResNet18 with Cifar10 and SVHN compared with the best
baselines (namely, ensemble and CALS, respectively). Additionally, EUAT
consistently outperforms existing approaches in terms of uAUC and the
correlation between the residuals of the model and predicted uncertainties.
EUAT achieves an improvement on the uAUC by up to 15.1%, 8.9%, 64.3%, 28.1%,
8.5%, and 20.4% compared to CE, calibration, DEUP, deep ensemble, CALS, and
CE+PE, respectively. Moreover, the correlation between the model’s residuals
and predicted uncertainties improves by 20%, 20.7%, 8.6%, and 11.5% when using
EUAT for training ResNet50 with ImageNet, Wide-ResNet with Cifar100, and
ResNet18 with Cifar10 and SVHN compared to the best-performing baseline in
each scenario.

Next, the effectiveness of EUAT in distinguishing correct predictions from
misclassifications based on the predicted uncertainty is assessed using the
Wasserstein distance of the uncertainty between the sets of correct and wrong
predictions. More in detail, on average across all models/datasets, the
Wasserstein distance increases by up to 1.7$\times$, 1.3$\times$, 9.7$\times$,
2.3$\times$, 1.4$\times$, and 2.2$\times$ using EUAT compared to CE,
calibration, DEUP, ensemble, CALS, and CE+PE, respectively. At last, we
evaluate the impact of using EUAT on the ECE and models’ misclassification
rate. EUAT results in lower ECE in two cases (ResNet50/ImageNet and
ResNet18/SVHN), while a slight increase is observed in the other two (Wide-
ResNet/Cifar100 and ResNet18/Cifar10) Importantly, the error rate remains
consistent across all benchmarks for all baselines, with a notable improvement
observed in the ResNet50/ImageNet benchmark, showcasing a reduction of 13.2%
(compared to the best baseline, namely, a deep ensemble). Furthermore, the
large gains of EUAT were obtained using larger models and datasets, where the
model is less accurate, and thus, by prioritizing the training of incorrect
predictions, the EUAT is able to yield significant improvements.

Further, in Figure 1, we plot the distribution of the normalized uncertainty
of the correctly and incorrectly predicted sets using the different baselines.
By visualizing these distributions, we verify an improvement in the separation
of the uncertainty of these two sets using EUAT (which is confirmed by the
computation of the Wasserstein distance in Table 2). These results confirm the
trustworthiness of the EUAT to improve the model’s uncertainty to separate
accurate and incorrect predictions.

####  4.2.2 Binary Classification Problem

Next, we proceed by assessing the effectiveness of the EUAT through a binary
classification scenario, wherein high-uncertainty predictions can be inverted
to the opposite class. We compared the models obtained using the different
baselines and tuned the uncertainty threshold as described in Section 3.
Additionally, we conducted an extensive evaluation utilizing supplementary
metrics such as F1-score, precision, True Positive Rate (TPR), True Negative
Rate (TNR), and error rates concerning the inversion or retention of high
uncertainty predictions.

Table 3 presents a comprehensive overview of these performance metrics.
Remarkably, the EUAT outperforms other baselines across 9 out of 11 evaluated
metrics. More in detail, employing EUAT leads to significant enhancements in
the F1-score, precision, and TPR, on average, by up to 9.2%, 11.3%, and 11.3%,
respectively, with a corresponding 29.9% reduction in error rates through high
uncertainty prediction inversion. Furthermore, our method improves the uA,
uAUC, and the correlation with model residuals by up to 9.9%, 52.5%, and
113.3%, while the Wasserstein distance of the uncertainty between the correct
and wrong predicted sets is enhanced by up to 17.5$\times$ (averaging at
4.1$\times$). Lastly, it is noteworthy that the ECE achieved using the EUAT
aligns closely with other baselines, namely CE, DEUP, Ensemble, and CALS, and
reduces on average across baselines by 18.1%, while the misclassification rate
lowers on average by 6%. These findings underscore the importance of using
EUAT to improve the model quality in a binary classification task.

####  4.2.3 Out-Of-Distribution Detection Task

Predicted uncertainty can be leveraged to reject difficult examples with high
uncertainty, even when these samples present huge distribution shifts from the
dataset used to train. Further, the OOD detection task varies significantly
when using different methods and datasets [18]. Thus, in this section, we
evaluate the effectiveness of EUAT to detect and reject OOD examples based on
the predicted uncertainty. For each baseline, we trained a ResNet18 using the
Cifar10 dataset and then tested it using a corrupted version with Gaussian
noise of this dataset (called Cifar10-C [17]). For a fair comparison, when
calibrating the model and training using DEUP, we use an additional validation
set containing 10% of the clean inputs of the original test set.

In Table 4, we compared EUAT against the other baselines considering the
previous metrics. Notably, in 83% of the metrics considered, EUAT performs
better than the baselines. More in detail, EUAT improves the uA by 9.2%,
21.9%, 22.2%, 0.25%, 9.9%, and 10.7% compared to CE, calibration, DEUP,
ensembles, CALS, and CE+PE, respectively, while the uAUC increases by 17.7%,
43.9%, 56.7%, 8.4%, 11.2%, and 20.1%. We also verify an enhancement in the
correlation between model residuals and the predicted uncertainty, Wasserstein
distance of the uncertainty between the correct and wrong predicted sets, and
ECE of 3.8$\times$, 16.9$\times$, 2.8$\times$ on average compared to the other
baselines. Finally, all the baselines, except CALS, yielded a model with a
larger error rate than EUAT.

####  4.2.4 Adversarial Training

Finally, we evaluate the EUAT in adversarial training settings. We opted to
exclusively train our models with adversarial examples, utilizing the FGSM
[13] to generate perturbations, with a predefined perturbation bound
$\epsilon$ set to 4/255. Due to resource constraints and the overhead
introduced by adversarial training, we did not deploy the ResNet50/ImageNet
benchmark in adversarial settings.

We report the results of the different baselines on standard and adversarial
settings in Table 5. Despite the challenges posed by adversarial scenarios,
overall, we see similar trends compared to the standard training. As expected
the adversarial error increases in all baselines and benchmarks considered. In
half of the scenarios/metrics assessed, EUAT outperforms the baselines.
Quantitatively, employing EUAT yields an average increase in uAUC of 20.2%,
10.9%, and 17.0%, when training a Wide-ResNet on Cifar100, a ResNet18 on
Cifar10, and a ResNet18 on SVHN datasets, respectively. Moreover, while the
gains in uA are slightly smaller, they still present significant improvements
of 15.2%, 1.4%, and 2.1% across the same models/datasets. Additionally, the
error rates across baselines exhibit negligible variance, and we verify a
larger separation of the uncertainty of the incorrect and correct predictions
when performing AT with EUAT, which highlights the robustness of our method
across different tasks and domains, reaffirming its efficacy in the
challenging context of adversarial attacks.

##  5 Conclusion and Future Work

This paper introduces Error-Driven Uncertainty Aware Training, a novel
approach designed to refine the estimation of model uncertainty. EUAT is
engineered to achieve two primary objectives: first, to heighten uncertainty
when models generate inaccurate predictions, and second, to output low
uncertainty when predictions are correct. This dual-purpose strategy is
achieved through the usage of two loss functions, which adapt based on whether
training examples are correctly or incorrectly predicted by the model. By
minimizing uncertainty for accurate predictions and maximizing it for
mispredictions while retaining error rates, EUAT aims to enhance model
reliability.

We evaluate EUAT against six different baselines and using six metrics, and
the results consistently demonstrate EUAT’s superior performance across the
majority of the considered cases. Even when faced with competitive baselines,
EUAT is still able to achieve comparable performance. Furthermore, we extend
our evaluation to encompass diverse problems, including binary classification,
out-of-distribution detection, and adversarial training settings. Across all
evaluated domains, EUAT demonstrates an enhanced ability to differentiate
between erroneous and accurate predictions based on uncertainty levels,
thereby increasing model trustworthiness.

Given our current focus on evaluating EUAT solely within classification tasks
using image recognition models, we envision broadening our scope to encompass
other domains such as regression models or language models for machine
translation or summarization, which present new and diverse challenges.

## References

  * Abdar et al. [2021] M. Abdar, F. Pourpanah, S. Hussain, D. Rezazadegan, L. Liu, M. Ghavamzadeh, P. Fieguth, X. Cao, A. Khosravi, U. R. Acharya, V. Makarenkov, and S. Nahavandi.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  _Information Fusion_ , 76:243–297, 2021.
  * Antorán et al. [2020] J. Antorán, J. U. Allingham, and J. M. Hernández-Lobato.  Depth uncertainty in neural networks.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.
  * Asgharnezhad et al. [2022] H. Asgharnezhad, A. Shamsi, R. Alizadehsani, S. Khosravi, Abbas andNahavandi, Z. A. Sani, D. Srinivasan, and S. M. S. Islam.  Objective evaluation of deep uncertainty predictions for covid-19 detection.  _Scientific Reports_ , 12, 2022.
  * Brach et al. [2020] K. Brach, B. Sick, and O. Dürr.  Single shot mc dropout approximation.  _ArXiv_ , abs/2007.03293, 2020.  URL https://api.semanticscholar.org/CorpusID:220381176.
  * Carlini and Wagner [2017] N. Carlini and D. Wagner.  Adversarial examples are not easily detected: Bypassing ten detection methods.  In _Proceedings of the 10th ACM Workshop on Artificial Intelligence and Security_ , AISec ’17, page 3–14, New York, NY, USA, 2017. Association for Computing Machinery.  ISBN 9781450352024.  URL https://doi.org/10.1145/3128572.3140444.
  * Damianou and Lawrence [2013] A. Damianou and N. D. Lawrence.  Deep Gaussian processes.  In _Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics_ , volume 31 of _Proceedings of Machine Learning Research_ , pages 207–215, Scottsdale, Arizona, USA, 2013. PMLR.
  * Ding et al. [2020] Y. Ding, J. Liu, X. Xu, M. Huang, J. Zhuang, J. Xiong, and Y. Shi.  Uncertainty-aware training of neural networks for selective medical image segmentation.  In _Proceedings of the Third Conference on Medical Imaging with Deep Learning_ , volume 121 of _Proceedings of Machine Learning Research_ , pages 156–173. PMLR, 2020.
  * Einbinder et al. [2022] B.-S. Einbinder, Y. Romano, M. Sesia, and Y. Zhou.  Training uncertainty-aware classifiers with conformalized deep learning.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, _Advances in Neural Information Processing Systems_ , volume 35, 2022.
  * Feinman et al. [2017] R. Feinman, R. R. Curtin, S. Shintre, and A. B. Gardner.  Detecting adversarial samples from artifacts.  In _International Conference on Machine Learning_ , 2017.
  * 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 International Conference on Machine Learning_ , volume 48. PMLR, 2016.
  * Gawlikowski et al. [2023] J. Gawlikowski, C. Tassi, M. Ali, J. Lee, M. Humt, J. Feng, A. Kruspe, R. Triebel, P. Jung, R. Roscher, M. Shahzad, W. Yang, R. Bamler, and X. Zhu.  A survey of uncertainty in deep neural networks.  _Artificial Intelligence Review_ , 56, 2023.
  * Godsill [2001] S. J. Godsill.  On the relationship between markov chain monte carlo methods for model uncertainty.  _Journal of Computational and Graphical Statistics_ , 10(2):230–248, 2001.
  * Goodfellow et al. [2015] I. Goodfellow, J. Shlens, and C. Szegedy.  Explaining and harnessing adversarial examples.  In _ICLR_ , 2015.
  * 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 International Conference on Machine Learning_ , volume 70. PMLR, 2017.
  * Gupta et al. [2021] K. Gupta, A. Rahimi, T. Ajanthan, T. Mensink, C. Sminchisescu, and R. Hartley.  Calibration of neural networks using splines.  In _International Conference on Learning Representations_ , 2021.
  * He et al. [2016] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016.
  * Hendrycks and Dietterich [2019] D. Hendrycks and T. Dietterich.  Benchmarking neural network robustness to common corruptions and perturbations.  In _International Conference on Learning Representations_ , 2019.
  * Karandikar et al. [2021] A. Karandikar, N. Cain, D. Tran, B. Lakshminarayanan, J. Shlens, M. C. Mozer, and B. Roelofs.  Soft calibration objectives for neural networks.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , volume 34, 2021.
  * Kim et al. [2021] J. U. Kim, S. T. Kim, H. J. Lee, S. Lee, and Y. M. Ro.  Cua loss: Class uncertainty-aware gradient modulation for robust object detection.  _IEEE Transactions on Circuits and Systems for Video Technology_ , 31(9), 2021.
  * Krishnan and Tickoo [2020] R. Krishnan and O. Tickoo.  Improving model calibration with accuracy versus uncertainty optimization.  In _Advances in Neural Information Processing Systems_ , volume 33, 2020.
  * Krizhevsky and Hinton [2009] A. Krizhevsky and G. Hinton.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * Kull et al. [2017] M. Kull, T. S. Filho, and P. Flach.  Beta calibration: a well-founded and easily implemented improvement on logistic calibration for binary classifiers.  In _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_ , volume 54, 2017.
  * Lahlou et al. [2023] S. Lahlou, M. Jain, H. Nekoei, V. I. Butoi, P. Bertin, J. Rector-Brooks, M. Korablyov, and Y. Bengio.  DEUP: Direct epistemic uncertainty prediction.  _Transactions on Machine Learning Research_ , 2023.  ISSN 2835-8856.  URL https://openreview.net/forum?id=eGLdVRvvfQ.
  * Lakshminarayanan et al. [2017] B. Lakshminarayanan, A. Pritzel, and C. Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , 2017.
  * Leibig et al. [2017] C. Leibig, V. Allken, M. S. Ayhan, P. Berens, and S. Wahl.  Leveraging uncertainty information from deep neural networks for disease detection.  _Scientific reports_ , 7(1), 2017.
  * Liu et al. [2022] B. Liu, I. B. Ayed, A. Galdran, and J. Dolz.  The devil is in the margin: Margin-based label smoothing for network calibration.  In _Computer Vision and Pattern Recognition Conference_ , 2022.
  * Liu et al. [2023] B. Liu, J. Rony, A. Galdran, J. Dolz, and I. Ben Ayed.  Class adaptive network calibration.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 16070–16079, June 2023.
  * MacKay [1992] D. J. C. MacKay.  A Practical Bayesian Framework for Backpropagation Networks.  _Neural Computation_ , 4(3):448–472, 1992.
  * MacKay [2003] D. J. C. MacKay.  _Information Theory, Inference, and Learning Algorithms_.  Cambridge University Press, 2003.
  * Madry et al. [2018] A. Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu.  Towards deep learning models resistant to adversarial attacks.  In _ICLR_ , 2018.
  * Matsubara et al. [2023] T. Matsubara, N. Tax, R. Mudd, and I. Guy.  TCE: A test-based approach to measuring calibration error.  In _Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence_ , 2023.
  * Mukhoti et al. [2020] J. Mukhoti, V. Kulharia, A. Sanyal, S. Golodetz, P. Torr, and P. Dokania.  Calibrating deep neural networks using focal loss.  In _Advances in Neural Information Processing Systems_ , volume 33, 2020.
  * Naeini et al. [2015] M. P. Naeini, G. F. Cooper, and M. Hauskrecht.  Obtaining well calibrated probabilities using bayesian binning.  In _Proceedings of the Twenty-Ninth AAAI Conference on Artificial Intelligence_ , 2015.
  * Neal [1996] R. M. Neal.  _Bayesian Learning for Neural Networks_.  Springer-Verlag, 1996.
  * Netzer et al. [2011] Y. Netzer, T. Wang, A. Coates, A. Bissacco, B. Wu, and A. Y. Ng.  Reading digits in natural images with unsupervised feature learning.  In _NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011_ , 2011.
  * Nixon et al. [2019] J. Nixon, M. W. Dusenberry, L. Zhang, G. Jerfel, and D. Tran.  Measuring calibration in deep learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops_ , June 2019.
  * Ovadia et al. [2019] Y. Ovadia, E. Fertig, J. Ren, Z. Nado, D. Sculley, S. Nowozin, J. Dillon, B. Lakshminarayanan, and J. Snoek.  Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Pearce et al. [2020] T. Pearce, F. Leibfried, and A. Brintrup.  Uncertainty in neural networks: Approximately bayesian ensembling.  In _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , 2020.
  * Platt [1999] J. C. Platt.  Probabilistic outputs for support vector machines and comparisons to regularized likelihood methods.  In _Advances in Large Margin Classifiers_ , 1999.
  * Russakovsky et al. [2015] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  _International Journal of Computer Vision (IJCV)_ , 115(3):211–252, 2015.  10.1007/s11263-015-0816-y.
  * Shamsi et al. [2023] A. Shamsi, H. Asgharnezhad, A. Tajally, S. Nahavandi, and H. Leung.  An uncertainty-aware loss function for training neural networks with calibrated predictions.  _ArXiv_ , abs/2110.03260, 2023.
  * Smith and Gal [2018] L. Smith and Y. Gal.  Understanding measures of uncertainty for adversarial example detection.  In _Proceedings of the Thirty-Fourth Conference on Uncertainty in Artificial Intelligence_ , 2018.
  * Wang and Yeung [2020] H. Wang and D.-Y. Yeung.  A survey on bayesian deep learning.  _ACM Comput. Surv._ , 53(5), sep 2020.
  * Wang et al. [2018] K.-C. Wang, P. Vicol, J. Lucas, L. Gu, R. Grosse, and R. Zemel.  Adversarial distillation of Bayesian neural network posteriors.  In _Proceedings of the 35th International Conference on Machine Learning_ , volume 80, 2018.
  * Zadrozny and Elkan [2001] B. Zadrozny and C. Elkan.  Obtaining calibrated probability estimates from decision trees and naive bayesian classifiers.  In _Proceedings of the Eighteenth International Conference on Machine Learning_ , ICML ’01, 2001.
  * Zagoruyko and Komodakis [2016] S. Zagoruyko and N. Komodakis.  Wide residual networks.  In _Proceedings of the British Machine Vision Conference 2016_ , 2016.