# Mitigating Transformer Overconfidence via Lipschitz Regularization

Wenqian Ye Department of Computer Science, University of Virginia,
Charlottesville, VA, USA  AI Lab, Shenzhen Children’s Hospital, Shenzhen,
China  Yunsheng Ma College of Engineering, Purdue University, West Lafayette,
IN, USA  AI Lab, Shenzhen Children’s Hospital, Shenzhen, China  Xu Cao
Department of Computer Science, University of Illinois Urbana-Champaign,
Urbana, IL, USA  AI Lab, Shenzhen Children’s Hospital, Shenzhen, China  Kun
Tang  T Lab, Tencent, Beijing, China

###### Abstract

Though Transformers have achieved promising results in many computer vision
tasks, they tend to be over-confident in predictions, as the standard Dot
Product Self-Attention (DPSA) can barely preserve distance for the unbounded
input domain. In this work, we fill this gap by proposing a novel Lipschitz
Regularized Transformer (LRFormer). Specifically, we present a new similarity
function with the distance within Banach Space to ensure the Lipschitzness and
also regularize the term by a contractive Lipschitz Bound. The proposed method
is analyzed with a theoretical guarantee, providing a rigorous basis for its
effectiveness and reliability. Extensive experiments conducted on standard
vision benchmarks demonstrate that our method outperforms state-of-the-art
single forward pass approaches in prediction, calibration, and uncertainty
estimation.

##  1 Introduction

(a) Deep Ensemble

(b) DUE

(c) SNGP

(d) LRFormer

Figure 1: Uncertainty heat map of LRFormer and baseline approaches on the two
moons 2D classification benchmark. Orange and blue points are positive and
negative training samples respectively. Background color visualizes the
predictive uncertainty of each model, where yellow stands for confidence and
blue indicates uncertainty. The proposed LRFormer (Figure 1(d)) achieves the
closest to ideal uncertainty quantification on this benchmark. Detail refer to
Section 4.4.

Deep learning (DL) has achieved remarkable performance, making it widely
employed in various inference and decision-making systems. However, DL models
still make mistakes, making trust and safety an increasingly important topic
[Amodei et al., 2016, Jiang et al., 2018], especially in critical applications
like self-driving cars [Huang and Chen, 2020] and medical diagnosis [Esteva et
al., 2017]. One solution to this problem is for models to not only achieve
high accuracy but also refrain from making overly confident predictions.

Transformer [Vaswani et al., 2017] and its variants, such as BERT [Devlin et
al., 2019], have made significant advances in Natural Language Processing
(NLP). Similarly, Vision Transformers (ViT) [Dosovitskiy et al., 2021] and its
variants [Liu et al., 2021, Li et al., 2022] have recently achieved state-of-
the-art performance on a variety of computer vision tasks [Yang et al., 2023,
Hu et al., 2023, Ma and Wang, 2023, Ma et al., 2023, Xu et al., 2022]. Despite
this, their propensity for overconfident predictions is cause for concern,
especially as they become one of the foundation architectures of deep
learning. To address this issue, we investigate the under-explored problem of
overconfidence issue in Transformers, which can aid subsequent tasks in the
construction of reliable models.

Overconfidence is a common problem in many machine learning models for both
in- and out-of-distribution inputs, including deep neural networks [Wei et
al., 2022, Kristiadi et al., 2020]. When a model is overconfident, it tends to
make highly confident predictions even when it is uncertain about the truth of
a given input. This can lead to poor performance and inaccurate results,
especially in real-world settings where uncertainty is prevalent. Uncertainty
estimation is a promising approach for addressing the issue of overconfidence
in machine learning models. By estimating uncertainty, a model can make more
informed predictions and provide a measure of confidence for each prediction.
This can help to improve the robustness and reliability of the model and
enable it to perform more effectively in a variety of applications, including
decision-making and risk assessment.

Previous techniques for estimating the model’s predictive uncertainty include
Bayesian deep learning [Wilson and Izmailov, 2020, Blundell et al., 2015] and
ensemble techniques[Lakshminarayanan et al., 2017, Gal and Ghahramani, 2016].
However, multiple forward passes at the test time are required by most of
these methods. In other words, these methods suffer heavy memory and
computation cost, which limits their adoption in real-world applications.

Recently, uncertainty quantification via single forward-pass neural networks,
which has similar latency as a single deterministic network, has received lots
of attention [Liu et al., 2020, van Amersfoort et al., 2021, Gillioz et al.,
2020]. SNGP [Liu et al., 2020] replaces the dense output layer with a Gaussian
Process (GP) layer and applies Spectral Normalization (SN) [Miyato et al.,
2018] to the hidden residual layers. DUE [van Amersfoort et al., 2021] builds
upon GPDNN [Bradshaw et al., 2017] and introduces additional constraints to
the feature extractor in the form of residual connections in combination with
SN [Miyato et al., 2018]. These methods perform well on uncertainty
estimation. However, they only focus on bounding the Lipschitz constants of
certain CNN modules i.e., convolution and batch normalization [Ioffe and
Szegedy, 2015] layers. Moreover, according to Lee et al. [2021], Transformer
blocks are very sensitive to the magnitude of the Lipschitz constant, and if
SN is employed in self-attention modules, training will progress exceptionally
slowly. Although some newest proposed Transformer architectures have been
proven to be Lipschitz continuous [Qi et al., 2023, Kim et al., 2021, Dasoulas
et al., 2021, Xu et al., 2023, Gupta and Verma, 2023], they still have not
solved the overconfidence problem of Transformer.

To address the issues above, we contribute as follows:

  * •

We propose a novel regularization technique, termed Lipschitz Regularized
Self-Attention (LRSA), that addresses distance awareness in both Lipschitzness
and Contraction. LRSA replaces the dot product similarity with the distance
within Banach Space and normalizes the term by a theoretical bound of the
Lipschitz constant. Furthermore, we provide a theoretical analysis of how our
method achieves these properties.

  * •

We develop the LRSA based Transformer called
LRFormer111https://github.com/SZCHAI/LRFormer, which integrates distance-
preserving hidden mappings in transformer blocks via LRSA and utilizes an
optional Gaussian Process (GP) distance-aware output layer for high-quality
uncertainty estimation.

  * •

We conduct extensive experiments on widely used OOD benchmarks, including
CIFAR-10/-100 versus SVHN and CIFAR-10/-100 versus CIFAR-100/-10. Compared to
state-of-the-art approaches, our experimental results demonstrate that the
proposed LRFormer model is superior in terms of prediction, calibration, and
uncertainty estimation, with minimal time complexity penalty.

##  2 Problem Statement

In the supervised multi-class classification setting, assume a data sample
$(\boldsymbol{x},y)\in\mathcal{X}\times\mathcal{Y}$ is sampled from an unknown
distribution, where $\mathcal{Y}=\\{1,\ldots,K\\}$ denote the label space with
$K$ classes and $\mathcal{X}=\mathbb{R}^{d}$ denote the feature space. A
learned classifier $f^{\theta}$: $\mathcal{X}\rightarrow\Delta^{K}$ can
produce a probability distribution for $\boldsymbol{x}$ on $K$ classes, where
$\Delta^{K}$ is the $K-1$ dimensional unit simplex. In this context, we
introduce a former definition of overconfidence for a general classifier.

######  Definition 2.1 (Overconfidence).

Assume $f^{\theta}$ as a composition of a non-probabilistic $K$-way classifier
$\boldsymbol{h}^{\theta}$ and a softmax function $\sigma$, i.e.
$\boldsymbol{f}^{\theta}=\boldsymbol{h}^{\theta}\circ\sigma$. Given a test
data sample $\boldsymbol{x},\boldsymbol{f}^{\theta}$ provides its probability
of assigning it to label $i$ as
$\frac{\exp\left(h_{i}^{\theta}(\boldsymbol{x})\right)}{\sum_{k=1}^{K}\exp\left(h_{k}^{\theta}(\boldsymbol{x})\right)}$,
where $h_{i}^{\theta}(\boldsymbol{x})$ denotes the $i$-th element of the logit
vector produced by $\boldsymbol{h}^{\theta}$. Then,
$\hat{y}:=\arg\max_{i}f_{i}^{\theta}(\boldsymbol{x})$ can be returned as the
predicted label and $\hat{p}:=\max_{j\neq y}f_{j}^{\theta}(\boldsymbol{x})$
can be treated as the confidence score. Overconfidence appears when the
prediction is wrong with high probability.

Figure 2: By breaking down a Transformer Layer into its fundamental
components, including GeLU Activation, MLP Layer, LayerNorm, and the Attention
Module, we gain a comprehensive understanding of their individual
contributions to the analysis of Lipschitzness in Transformers. Detailed
analysis can be found in Section 3.3.

Based on Definition 2.1, the logit vector can be decomposed into two
components:
$\boldsymbol{h}(\boldsymbol{x})=||\boldsymbol{h}(\boldsymbol{x})||\cdot\hat{\boldsymbol{h}(\boldsymbol{x})}$,
where $||\boldsymbol{h}(\boldsymbol{x})||$ is the L2-norm of the logit vector
and $\hat{\boldsymbol{h}(\boldsymbol{x})}$ is the unit vector in the same
direction as $\boldsymbol{h}(\boldsymbol{x})$. These two terms represent the
magnitude and direction of the logit vector, respectively. It is evident that
if $\arg\max_{k}(h_{k})=c$, then $\arg\max_{k}(\gamma\boldsymbol{h}_{k})=c$
always holds for any given constant value $\gamma>1$. This indicates that the
magnitude of the logit vector does not affect the predicted class $c$.
Additionally, for any given scalar $\gamma>1$, if $c=\arg\max_{k}(h_{k})$,
then $\sigma_{c}(\gamma\boldsymbol{h})\geq\sigma_{c}(\boldsymbol{h})$. From
the above claims, we observe that increasing the magnitude
$||\boldsymbol{h}(\boldsymbol{x})||$ will lead to a higher softmax confidence
score while leaving the final prediction unchanged.

During optimization, the cross-entropy loss is given as:

$\mathcal{L}_{\mathrm{CE}}(\boldsymbol{h}(\boldsymbol{x};\theta),y)=-\log
p(y\mid\boldsymbol{x})=-\log\frac{e^{\|\boldsymbol{h}\|\cdot\hat{h}_{y}}}{\sum_{i=1}^{k}e^{\|\boldsymbol{h}\|\cdot\hat{h}_{i}}}$

While the direction $\hat{\boldsymbol{h}(\boldsymbol{x})}$ remains constant,
increasing the magnitude will lead to a smaller $p(y\mid\boldsymbol{x})$. In
the standard Transformer, optimization on the training loss leads to an
increase in the magnitude of the network output to produce a higher softmax
confidence score, resulting in a smaller loss.

Lipschitzness can help a classifier tackle the overconfidence issue by
limiting the amount of change in the classifier’s output when the input is
perturbed slightly. When a classifier is overconfident, it tends to assign
high confidence to incorrect predictions, which can result in poor performance
on unseen data. However, if the classifier is Lipschitz continuous, then the
amount of change in its output is limited when the input is perturbed, which
can prevent the classifier from making overly confident predictions on Out-of-
Distribution (OOD) samples.

######  Definition 2.2 (Lipschitz Continuity).

Lipschitz constant of a function $f$ is an upper bound on the ratio between
the output and the input variations of a function $f$. If $L\in[0,+\infty)$ is
such that, for every input $x\in\mathbb{R}^{d}$ and perturbation $\Delta
x\in\mathbb{R}^{d}$,

$\|f(x+\Delta x)-f(x)\|_{p}\leqslant L\|\Delta x\|_{p}$ (1)

then $L$ is a Lipschitz constant of $f$. $\|\cdot\|_{p}$ denotes the p-norm.
If $X$ is defined as the $\epsilon$-ball at point $x$, i.e.,
$\mathcal{X^{\prime}}=\left\\{x^{\prime}\mid\left\|x-x^{\prime}\right\|\leqslant\epsilon\right\\}$,
then $L$ is the local Lipschitz constant of $f$ at $x$.

Furthermore, the Lipschitz condition can be extended to the Bi-Lipschitz
condition. Given any two input samples $x_{1},x_{2}$ and a non-probabilistic
K-way classifier $\boldsymbol{h}^{\theta}$, the Bi-Lipschitz condition can be
defined as:

$L_{1}\|x_{1}-x_{2}\|\leq\|\boldsymbol{h}^{\theta}(x_{1})-\boldsymbol{h}^{\theta}(x_{2})\|\leq
L_{2}\|x_{1}-x_{2}\|$ (2)

where $\|.\|$ is a semantically meaningful distance for the data manifold for
positive and bounded Lipschitz constants $0<L_{1}<1<L_{2}$. These bounds
$L_{1}$, $L_{2}$ represent sensitivity and smoothness conditions which
prevents the hidden representations $\boldsymbol{h}^{\theta}(x)$ from being
unnecessarily invariant to the semantically meaningful changes in the input
manifold or being overly sensitive to the semantically meaningless
perturbations respectively.

Dot-Product Self-Attention is a fundamental building block of the Transformer
model. It enables the model to focus on the most relevant parts of the input
sequence by weighing the contribution of each input vector to the output, as
given by the equation $\operatorname{Attention}(X)=S(X)\cdot
V(X)=\operatorname{softmax}\left(\frac{Q\cdot
K^{\top}}{\sqrt{d_{k}}}\right)\cdot V(X)$. This mechanism can be further
generalized by incorporating a similarity function that measures the relevance
between input vectors [Katharopoulos et al., 2020]. While several similarity
functions, such as the cosine similarity [Qi et al., 2023] or the scaled dot
product, have been used in the original formulation, they may not be optimal
for all scenarios. Therefore, in this paper, we explore the use of a Lipschitz
similarity function to mitigate overconfidence issues in the Transformer. We
present a method for constructing a suitable Lipschitz similarity function and
demonstrate its effectiveness in improving the robustness and accuracy of the
model.

$\operatorname{Attention}(\mathbf{x}_{i},\mathbf{x}_{j})=\operatorname{softmax}\left(\frac{\operatorname{sim}(\mathbf{x}_{i},\mathbf{x}_{j})}{\sqrt{d_{k}}}\right)\cdot
V$ (3)

Our method aims to maintain a reasonable Lipschitz constant at the block level
to address the issue of overconfidence in neural networks. While other
methods, such as Bayesian approaches and label smoothing [Müller et al.,
2019], have been proposed to tackle overconfidence, our method incorporates
block-wise control. By constraining the Lipschitz constant at each block in
the Transformer, we can limit the growth of the magnitude of the network
output and reduce overconfidence. This block-wise design also allows our
method to be easily integrated into various Transformer-based architectures,
including those that have undergone large-scale pretraining. As a result, our
method can be scaled up to handle a wide range of tasks and datasets.

##  3 Our Method

###  3.1 Notations and Setup

  * •

$S^{(i)}:=\operatorname{diag}\left(S_{i:}\right)-S_{i:}^{\top},S_{i:}\in\mathbb{R}^{N\times
N}.$

  * •

Binary Matrix with one in the $(i,j)$ the entry and zeros elsewhere:
$E_{ij}\in\mathbb{R}^{N\times N}$

  * •

Kronecker delta: $\delta_{ij}\in\\{0,1\\}$

  * •

$(\infty,2)$-norm:
$\|M\|_{(\infty,2)}=\max_{i}\left(\sum_{j}M_{ij}^{2}\right)^{1/2}$

  * •

Frobenius norm: $\|M\|_{F}=\left(\sum_{i,j}M_{ij}^{2}\right)^{1/2}$

  * •

Lipschitz constant $L_{\mathbb{X},\mathbb{Y}}(f)$: for a function
$f:\mathbb{X}\rightarrow\mathbb{Y}$,
$L_{\mathbb{X},\mathbb{Y}}(f)=\sup_{X\in\mathbb{X}}\|\frac{\partial
f(X)}{\partial X}\|_{\mathbb{X},\mathbb{Y}}$

###  3.2 Lipschitz Regularization on Self Attention

Kim et al. [2021] proved that the Scaled Dot-Product Self-Attention does not
satisfy the bi-Lipschitz condition. To extend the generality of self-attention
with high-quality uncertainty estimation, we propose a new regularization
method Lipschitz Regularized on Self Attention (LRSA) by replacing the self-
attention function with a contractive Bi-Lipschitz expression without losing
the original ability of representation. We will explicitly discuss separate
aspects to see how to achieve both Lipschitzness and Contraction in our
method.

####  3.2.1 Lipschitzness

Given that Dot-Product Self-Attention is not Lipschitz, suppose there exists
such mapping $f(X)$, $X\in\mathbb{R}^{N}$:

$f(X)=S\cdot X=\operatorname{softmax}(aX\cdot X^{\top})\cdot X=\\\
\left[\begin{array}[]{c}f_{1}(X)\\\ \vdots\\\ f_{N}(X)\end{array}\right]$

Its Jacobian Matrix is $J_{f}=[J_{ij}]_{N\times N}$, each entry can be written
as:

$J_{ij}=aX^{\top}S^{(i)}\left[E_{ji}X+\delta_{ij}X\right]+S_{ij}I\in\mathbb{R}^{N\times
1}$

Thus for $i=j$:

$J_{ii}=aX^{\top}S^{(i)}E_{ii}X+aX^{\top}S^{(i)}X+S_{ii}$ (4)

$X^{\top}S^{(i)}X$ is in the form of a variance of a discrete distribution.
When $\mathbf{x}_{i}=\mathbf{0}$ for some $i$, some entries of the Jacobian of
$f$ grow proportionally to the sample variance of $\mathbf{x}_{\neq i}$.(The
softmax probabilities $S_{i:}$ are constant with respect to $\mathbf{x}_{\neq
i}$ when $\mathbf{x}_{i}=0$.) This will lead to an unbounded Jacobian matrix.

To avoid this pathology, we replace $Q\cdot K^{\top}$ by
$\operatorname{sim}(\mathbf{x}_{i},\mathbf{x}_{j})=-\|\mathbf{x}_{i}^{\top}Q-\mathbf{x}_{j}^{\top}K\|_{2}^{2}$
in $\operatorname{Attention}(X)$. Here, the new similarity measurement lies in
the Banach Space (complete vector space with norm $\|\cdot\|$), which is a
more generalized space over Hilbert Space (complete inner product space)
[Megginson, 2012]. This modification also gives a strong theoretical guarantee
on Lipschitzness with easy matrix multiplications during training.

####  3.2.2 Contraction

Contraction of the Scaled Dot-Product Self-Attention is another crucial issue
for achieving well-calibrated uncertainty. Deriving such contraction scalar
requires a theoretical lower bound of the Lipschitz constant on the Dot-
Product Self-Attention function. A desirable contraction scalar could be non-
strict but easy to compute during training.

######  Theorem 3.1 (Dasoulas et al. [2021]).

For $\alpha\geq 0$, if $\tilde{g}$ is Lipschitz and for all
$X\in\mathbb{R}_{d\times n}$, and $\tilde{g}$ satisfy the following
conditions:

  1. 1.

$\|\tilde{g}(X)\|_{\infty}\leqslant\alpha c(X)$,

  2. 2.

$\|X^{\top}\|_{(\infty,2)}\|\frac{\partial\tilde{g}(X)}{\partial
X}\|_{F,(2,\infty)}\leqslant\alpha c(X)$,

  3. 3.

$\|X^{\top}\|_{(\infty,2)}\|\frac{\partial c(X)}{\partial
X}\|_{F,1}\|\tilde{g}(X)\|_{(2,\infty)}\leqslant\alpha c(X)^{2}$,

where $c$ is a scalar function $c:\mathbb{R}^{d\times
n}\rightarrow\mathbb{R}_{+}$. Then $g(X)$ is Lipschitz:

$\displaystyle
g(X)=\frac{\alpha\tilde{g}(X)}{\max\left\\{\|\tilde{g}(X)\|_{(2,\infty)},\left\|X^{\top}\right\|_{(\infty,2)}L_{F,(2,\infty)}(\tilde{g})\right\\}}$
(5)

Inspired by 3.1, we introduce a proper regularization scalar function with a
Scalar Factor $\alpha$ by replacing $\tilde{g}(X)$ with $Q\cdot K^{\top}$:

$c(X)=\frac{\alpha}{\|Q\|_{F}\cdot\|X^{\top}\|_{(\infty,2)}}$ (6)

Here, we assign it as a hyperparameter in control of the corresponding
Lipschitz constants for proper contraction of the attention block. Small alpha
results in a loss of information while a large alpha causes the model tending
to be non-Lipschitz.

####  3.2.3 Summary

Here is the formal definition of the similarity function:

$S_{ij}:=b\cdot
c(X)=-\frac{\alpha\|\mathbf{x}_{i}^{\top}W_{Q}-\mathbf{x}_{j}^{\top}W_{K}\|_{2}^{2}}{\|Q\|_{F}\cdot\|X^{\top}\|_{(\infty,2)}}$
(7)

This pair-wise operation can alternatively be implemented as a matrix version
for improved computational efficiency:

$S(X)=\operatorname{softmax}(-\alpha\cdot\frac{\|Q\|_{\text{row
}}^{2}-2QK^{\top}+\|K\|_{\text{col
}}^{2^{\top}}}{\|Q\|_{F}\cdot\|X^{\top}\|_{(\infty,2)}})$ (8)

LRSA Attention can be represented by the expression
$\operatorname{LRSA}(X)=S(X)\cdot V(X)$, where $S(X)$ denotes the similarity
scores and $V(X)$ represents the value embeddings. In the following section,
we define the Lipschitz Constant of $S(X)$ as $L_{\operatorname{LRSA}}$. From
Supplementary Material, we can conclude that $L_{\operatorname{LRSA}}$ is
bounded by
$\frac{6\alpha}{\left\|X\right\|_{F}}\cdot\frac{(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|W_{Q}\right\|_{F}}^{2}$.

###  3.3 Bottom-Up Analysis on LRFormer

Analysis on Lipschitzness of GeLU Activation

Figure 3: GeLU $g(x)$ and the derivative of GeLU $g^{\prime}(x)$.

GeLU [Hendrycks and Gimpel, 2016] is the most commonly used activation
function in Transformers, especially the GPT series of models and Vision
Transformers. GeLU’s activation function form is
$\operatorname{GeLU}(x)=x\Phi(x)$, where $\Phi(x)$ is the standard Gaussian
cumulative distribution function. The derivative of GeLU is given as:

$\operatorname{GeLU}^{\prime}(x)=\frac{xe^{-\frac{x^{2}}{2}}}{\sqrt{2\pi}}+\frac{\operatorname{erf}(\frac{x}{\sqrt{2}})}{2}+\frac{1}{2}$
(9)

,where $\operatorname{erf}(x)=\frac{2}{\sqrt{\pi}}\int_{0}^{x}e^{-t^{2}dt}$.
The second derivative of GeLU is:

$\operatorname{GeLU}^{\prime\prime}(x)=\frac{1}{\sqrt{2\pi}}((1+\sqrt{2})\cdot
e^{-\frac{x^{2}}{2}}-x^{2}e^{-\frac{x^{2}}{2}})$ (10)

Compared with ReLU, GeLU is differentiable at zero and keeps the Lipschitz
continuous. The Lipschitz constant of GeLU is the value of $x$ from
$\max(\operatorname{GeLU}^{\prime}(x))$. Let
$\operatorname{GeLU}^{\prime\prime}(x)=0$, we can verify
$\max(\operatorname{GeLU}^{\prime}(x))\approx 1.129$.

Analysis on Lipschitzness of LayerNorm

The raw LayerNorm operation
($\text{LN}(\mathbf{x})=\frac{\mathbf{x}-\mu(\mathbf{x})}{\sqrt{\sigma^{2}(\mathbf{x})}}*\boldsymbol{\gamma}+\boldsymbol{\beta}$)
[Ba et al., 2016] is not Lipschitz continuous because the ill-defined input
with zero variance will lead to a Jacobian matrix filled with infinity.

However, the LayerNorm operation can be changed to a Lipschitz continuous
form, which is the LayerNorm used in our models. The form can be expressed as:

$\displaystyle\text{LayerNorm}(\mathbf{x})$
$\displaystyle=\frac{\mathbf{x}-\mu(\mathbf{x})}{\sqrt{\sigma^{2}(\mathbf{x})+\epsilon}}*\boldsymbol{\gamma}+\boldsymbol{\beta}$
(11)

where $\mathbf{x},\boldsymbol{\beta},\boldsymbol{\gamma}\in\mathbb{R}^{N}$,
$\mu(\mathbf{x})=\frac{1}{N}\sum_{i=1}^{N}x_{i}$,
$\sigma^{2}(\mathbf{x})=\frac{1}{N}\sum_{i=1}^{N}(x_{i}-\mu(\mathbf{x}))^{2}$.

From Supplementary Material, we can conclude that LayerNorm is Lipschitz with
the constant $\eta=\epsilon^{-\frac{1}{2}}\max_{i}|\gamma_{i}|N$.

Analysis on Lipschitzness of MLP Layer

In each Transformer block, the Attention Module is typically followed by the
Multi-Layer Perceptron (MLP) Layer. The MLP layer consists of two Fully
Connected (FC) Layers, a Dropout Layer, and a GeLU activation function. Since
the Dropout Layer has no impact on the Lipschitz constant, it can be
simplified as:

$\operatorname{MLP}(x)=\operatorname{FC}_{2}\circ\operatorname{GeLU}\circ\operatorname{FC}_{1}(x)$
(12)

The upper bound on the Lipschitz constant of a fully-connected layer can be
derived by analyzing the effect of each layer independently and considering a
product of the resulting spectral norms $\sigma(\cdot)$ [Miyato et al., 2018].

$L_{\operatorname{MLP}}=1.129\cdot\sigma(W_{1})\cdot\sigma(W_{2})$ (13)

Lipschitz Constant of LRFormer Layer

In the LRFormer, we analyze the LRFormer Layer to demonstrate our method
follows the Lipschitzness. Each LRFormer Layer can be expressed as:

$\displaystyle\operatorname{LRFormer}_{i}(x)$
$\displaystyle=\operatorname{LayerNorm}(\operatorname{LayerNorm}(x+\operatorname{LRSA}(x))$
(14)
$\displaystyle+\operatorname{MLP}(\operatorname{LayerNorm}(x+\operatorname{LRSA}(x))))$

Combined with the previous analysis, we conclude the following Lipschitz bound
of a single LRFormer layer:

$\displaystyle L_{\operatorname{Layer}}$
$\displaystyle=\eta_{1}\cdot\eta_{2}\cdot((1+L_{\operatorname{LRSA}})$ (15)
$\displaystyle+1.129\cdot\sigma(W_{1})\cdot\sigma(W_{2})\cdot(1+L_{\operatorname{LRSA}}))$

, where $\eta_{1}$, $\eta_{2}$ are the Lipschitz Constant of two LayerNorms.

##  4 Experiments

Method | Accuracy $(\uparrow)$ | ECE $(\downarrow)$ | NLL $(\downarrow)$ |  OOD AUPR $(\uparrow)$
---|---|---|---|---
SVHN | CIFAR-100
Deterministic∗ |  96.0 $\pm$ 0.01 |  0.023 $\pm$ 0.002 |  0.158 $\pm$ 0.01 |  0.781 $\pm$ 0.01 |  0.835 $\pm$ 0.01
MC Dropout∗ |  96.0 $\pm$ 0.01 |  0.021 $\pm$ 0.002 |  0.173 $\pm$ 0.01 |  0.971 $\pm$ 0.01 |  0.832 $\pm$ 0.01
MCD-GP∗ |  95.5 $\pm$ 0.02 |  0.024 $\pm$ 0.004 |  0.172 $\pm$ 0.01 |  0.960 $\pm$ 0.01 |  0.863 $\pm$ 0.01
DNN-SN∗ |  96.0 $\pm$ 0.01 |  0.025 $\pm$ 0.004 |  0.171 $\pm$ 0.01 |  0.974 $\pm$ 0.01 |  0.859 $\pm$ 0.01
DNN-GP∗ |  95.9 $\pm$ 0.02 |  0.029 $\pm$ 0.002 |  0.221 $\pm$ 0.02 |  0.976 $\pm$ 0.01 |  0.887 $\pm$ 0.01
DUQ∗ |  94.7 $\pm$ 0.02 |  0.034 $\pm$ 0.002 |  0.239 $\pm$ 0.02 |  0.973 $\pm$ 0.01 |  0.854 $\pm$ 0.01
DUE∗ |  95.6 $\pm$ 0.04 |  0.018 $\pm$ 0.002 |  0.187 $\pm$ 0.01 | - | -
SNGP∗ |  95.9 $\pm$ 0.01 |  0.018 $\pm$ 0.001 |  0.138 $\pm$ 0.01 |  0.990 $\pm$ 0.01 |  0.905 $\pm$ 0.01
Deep Ensemble∗† |  96.6 $\pm$ 0.01 |  0.010 $\pm$ 0.001 |  0.114 $\pm$ 0.01 |  0.964 $\pm$ 0.01 |  0.888 $\pm$ 0.01
LRFormer | 97.2 $\pm$ 0.01 | 0.012 $\pm$ 0.001 | 0.100 $\pm$ 0.01 | 0.993 $\pm$ 0.01 | 0.911 $\pm$ 0.01
Table 1: Comparison between proposed LRFormer and SOTA methods on CIFAR-10 vs SVHN/CIFAR-100 benchmarks, averaged over 10 seeds. The best method among single-network approaches is highlighted in bold. ∗Results from the original papers. † with 10 models. Method | Accuracy $(\uparrow)$ | ECE $(\downarrow)$ | NLL $(\downarrow)$ |  OOD AUPR $(\uparrow)$
---|---|---|---|---
SVHN | CIFAR-10
Deterministic∗ |  79.8 $\pm$ 0.02 |  0.085 $\pm$ 0.004 |  0.872 $\pm$ 0.01 |  0.882 $\pm$ 0.01 |  0.745 $\pm$ 0.01
MC Dropout∗ |  79.6 $\pm$ 0.02 |  0.050 $\pm$ 0.003 |  0.825 $\pm$ 0.01 |  0.832 $\pm$ 0.01 |  0.757 $\pm$ 0.01
MCD-GP∗ |  79.5 $\pm$ 0.04 |  0.085 $\pm$ 0.005 |  0.937 $\pm$ 0.01 |  0.873 $\pm$ 0.01 |  0.754 $\pm$ 0.01
DNN-SN∗ |  79.9 $\pm$ 0.02 |  0.098 $\pm$ 0.004 |  0.918 $\pm$ 0.01 |  0.879 $\pm$ 0.03 |  0.745 $\pm$ 0.01
DNN-GP∗ |  79.2 $\pm$ 0.03 |  0.064 $\pm$ 0.005 |  0.885 $\pm$ 0.01 |  0.876 $\pm$ 0.01 |  0.746 $\pm$ 0.02
DUQ∗ |  78.5 $\pm$ 0.03 |  0.119 $\pm$ 0.001 |  0.980 $\pm$ 0.02 |  0.878 $\pm$ 0.01 |  0.732 $\pm$ 0.01
SNGP∗ |  79.9 $\pm$ 0.03 |  0.025 $\pm$ 0.012 |  0.847 $\pm$ 0.01 |  0.923 $\pm$ 0.01 | 0.801 $\pm$ 0.01
Deep Ensemble∗† |  80.2 $\pm$ 0.01 |  0.021 $\pm$ 0.004 |  0.666 $\pm$ 0.02 |  0.888 $\pm$ 0.01 |  0.780 $\pm$ 0.01
LRFormer | 85.2 $\pm$ 0.03 | 0.018 $\pm$ 0.005 | 0.538 $\pm$ 0.01 | 0.955 $\pm$ 0.01 |  0.777 $\pm$ 0.01
Table 2: Comparison between proposed LRFormer and the SOTA methods on CIFAR-100 vs SVHN and CIFAR-100 vs CIFAR-10 benchmark, averaged over 10 seeds. The best method among single-network approaches is highlighted in bold. ∗Results from the original papers. † with 10 models. Method |  ECE $(\downarrow)$ |  NLL $(\downarrow)$
---|---|---
DP Attention [Vaswani et al., 2017] | 0.066 | 0.580
L2 Attention [Kim et al., 2021] | 0.048 | 0.582
SCSA [Qi et al., 2023] | 0.028 | 0.626
LRSA | 0.018 | 0.538
Table 3: Overconfident evaluation comparison among different attention
modules. The backbones are the same except for the kernel. We compare the
checkpoint when the accuracy achieves $0.85\pm 0.01$ with CIFAR-100. Note, we
select the best hyper-parameter for SCSA ($\nu=1.0,\tau=12,\epsilon=1e-8$)

In this section, we verify the effectiveness of LRFormer in OOD detection with
several benchmark datasets. We also design ablation experiments including
attention module comparison, searching for a proper scalar factor $\alpha$,
and validating the reliability of pretrained models.

$\alpha$ |  Accuracy $(\uparrow)$ |  NLL $(\downarrow)$
---|---|---
1 |  0.5637 $\pm$ 0.01 |  1.7063 $\pm$ 0.02
100 |  0.6601 $\pm$ 0.02 |  1.3712 $\pm$ 0.01
500 |  0.6404 $\pm$ 0.01 |  1.4041 $\pm$ 0.03
1000 |  0.6212 $\pm$ 0.02 |  1.4259 $\pm$ 0.01
Table 4: Effect of $\alpha$ in LRSA with CIFAR-100

###  4.1 Setup

####  4.1.1 Benchmarks

We evaluate the performance of the proposed LRFormer model on the OOD
benchmark [Miyato et al., 2018] using SVHN [Netzer et al., 2011] as the OOD
dataset for the model trained on CIFAR-10/-100 [Krizhevsky et al., 2009]. OOD
data is never seen during training, whereas ID samples are semantically
similar to training samples. We also show LRFormer’s performance on the Two
Moons dataset in Figure 1.

####  4.1.2 Baselines

Our baselines included the deterministic model and two ensemble models: MC
Dropout (with 10 dropout samples) and deep ensembles (with 10 models)
[Lakshminarayanan et al., 2017]. All models were trained with a dense output
layer and no spectral regularization. Besides, we also compared three single-
model approaches: MCD-GP (with 10 samples), DUQ [Van Amersfoort et al., 2020],
DUE [van Amersfoort et al., 2021], and SNGP series [Liu et al., 2020]
including DNN-SN and DNN-GP. For models that use a GP layer, we kept DL = 1024
and computed the predictive distribution using Monte Carlo averaging with 10
samples. For a fair comparison, we set the backbone with the same parameter
magnitude (19.9M parameters for LRFormer, and 36.5M for the SNGP series).

####  4.1.3 Evaluation Metrics

Expected Calibration Error (ECE) [Guo et al., 2017] quantifies the difference
between a model’s expected confidence (e.g., the maximum probability score)
and its actual accuracy. It achieves this by partitioning all the samples,
with $n$ representing the total number of samples, into $M$ equally sized bins
based on their confidence scores, then calculating the expected difference
between accuracy and the average confidence in each bin. In our task, ECE can
indicate the effectiveness of the model in dealing with overconfidence.

In addition to ECE, we employ Negative Log Likelihood (NLL), OOD Area Under
the Receiver Operating Characteristic Curve (AUROC), and OOD Area Under the
Precision-Recall Curve (AUPR) to evaluate the model’s performance in
overconfidence and uncertainty estimation ability.

####  4.1.4 Implementation Details

In the following experiments, we resize the input image to $224\times 224$
pixels and set the patch size of LRFormer to 16. We employ AdamW [Loshchilov
and Hutter, 2017] as the optimizer with a weight decay of 0.05. We use a
cosine learning rate scheduler [Loshchilov and Hutter, 2016] with the base
learning rate set to $5\times 10^{-5}$. All models are trained for 100 epochs
with 10 different random seeds on NVIDIA A100 GPUs.

Dataset | Pretained |  Accuracy $(\uparrow)$ |  NLL $(\downarrow)$ |  OOD AUROC $(\uparrow)$ |  OOD AUPR $(\uparrow)$
---|---|---|---|---|---
CIFAR-10 | W/O |  0.8528 $\pm$ 0.01 |  0.4447 $\pm$ 0.02 |  0.8500 $\pm$ 0.01 |  0.9078 $\pm$ 0.02
W/ | 0.8616 $\pm$ 0.01 | 0.4193 $\pm$ 0.01 | 0.9125 $\pm$ 0.02 | 0.9499 $\pm$ 0.01
CIFAR-100 | W/O |  0.6404 $\pm$ 0.01 |  1.4041 $\pm$ 0.03 |  0.8421 $\pm$ 0.01 |  0.9165 $\pm$ 0.01
W/ | 0.6679 $\pm$ 0.01 | 1.2122 $\pm$ 0.02 | 0.8689 $\pm$ 0.01 | 0.9319 $\pm$ 0.01
Table 5: Ablation study between with (W/) and without (W/O) pre-trained weights from ImageNet-1K dataset. The best method among single-network approaches is highlighted in bold. $\downarrow$ means lower is better. $\uparrow$ means higher is better. Dataset | Method |  Accuracy $(\uparrow)$ |  NLL $(\downarrow)$ |  OOD AUROC $(\uparrow)$ |  OOD AUPR $(\uparrow)$
---|---|---|---|---|---
CIFAR-10 | Transformer | 0.8592 $\pm$ 0.01 |  0.6972 $\pm$ 0.02 |  0.7552 $\pm$ 0.05 |  0.8521 $\pm$ 0.02
DUE + Transformer |  0.8556 $\pm$ 0.03 |  0.5337 $\pm$ 0.02 |  0.8348 $\pm$ 0.04 |  0.8921 $\pm$ 0.01
SNGP + Transformer |  0.8542 $\pm$ 0.02 |  0.4933 $\pm$ 0.01 |  0.8275 $\pm$ 0.03 |  0.8960 $\pm$ 0.01
LRFormer |  0.8528 $\pm$ 0.02 | 0.4447 $\pm$ 0.01 | 0.8500 $\pm$ 0.05 | 0.9078 $\pm$ 0.01
CIFAR-100 | Transformer |  0.6304 $\pm$ 0.02 |  1.7862 $\pm$ 0.01 |  0.7831 $\pm$ 0.02 |  0.8701 $\pm$ 0.02
SNGP + Transformer |  0.6298 $\pm$ 0.03 |  1.5413 $\pm$ 0.01 |  0.8134 $\pm$ 0.01 |  0.8929 $\pm$ 0.01
LRFormer | 0.6404 $\pm$ 0.02 | 1.4041 $\pm$ 0.02 | 0.8421 $\pm$ 0.01 | 0.9165 $\pm$ 0.01
Table 6: Ablation study between the proposed LRFormer and existing uncertainty
quantification methods with the same training backbone on the CIFAR-10/-100 vs
SVHN benchmark. The best method among single-network approaches is highlighted
in bold. $\downarrow$ means lower is better. $\uparrow$ means higher is
better.

###  4.2 Comparison with State-of-the-art models

Following Touvron et al. [2022], we adopt an existing training setup, namely
the A3 procedure of Wightman et al. [2021]. We adjust the learning rate of the
A3 procedure when training LRFormer. In our experiments, we set the learning
rate to 0.006 for LRFormer when pretraining and 0.004 while finetuning on
CIFAR-10/-100. Besides, Different from previous unfair comparison methods
[Fort et al., 2021, Xue et al., 2022, Cao and Zhang, 2022], pretrained models
from the extra datasets and few-shot outlier exposure settings are not used
during training.

To evaluate the model’s OOD detection performance, we adopt the two OOD tasks
suggested by SNGP: (1) using SVHN as the OOD dataset for a model trained on
CIFAR-10/-100; (2) using CIFAR-100/-10 as the OOD dataset for a model trained
on CIFAR-10/-100, respectively. Table 1 and Table 2 show the main comparison
results. LRFormer outperforms the other single forward pass approaches in all
the metrics of CIFAR-10 and most of the metrics of CIFAR-100. Moreover,
LRFormer also achieves similar results to Deep Ensemble, which contains 10
models and requires around $10\times$ as much time to execute as LRFormer and
other single forward pass approaches.

###  4.3 Ablation Study

####  4.3.1 Attention Blocks

To validate the ability to solve overconfidence issues of Transformers, we
compare the LRSA with the scale dot-product attention, L2 attention, and the
scaled cosine similarity attention (SCSA) using overconfidence evaluation
metrics. In Table3, we show the best Test ECE and NLL across training for each
of the Transformer models. The generalization performance of the best model
for each setting of self-attention is similar.

We find that L2 attention, SCSA, and LRSA work well under the Lipschitz
guarantee. Meanwhile, LRSA works better than SCSA and L2 attention.

####  4.3.2 Hyperparameter Analysis

The Scalar Factor $\alpha$ in Equation (6) controls the scale of the Lipschitz
constant of the Transformer blocks. In general, we propose running a grid
search for $\alpha\in\left\\{...,100,500,1000,...\right\\}$ to find the
highest possible value of $\alpha$ while retaining the predictive performance
of LRFormer. In our experiments (Table 5), we set scalar factor $\alpha=1000$
in CIFAR-10, and $\alpha=500$ in CIFAR-100. The model’s performance is not
very sensitive to the parameters in the GP output layer, we follow Liu et al.
[2020]’s suggestion and set the number of random features to $1024$, the
length-scale for the RBF kernel to $2$, and the $L_{2}$ regularization to 0. A
proper $\alpha$ value, i.e. 100, can preserve both the Lipschitzness and
Contraction properties in the model. Small alpha will cause loss of
information while large alpha will cause the model tending to be non-
Lipschitz, leading to degenerate performance.

####  4.3.3 Module Comparison

In this section, we compare LRFormer and other OOD detection methods (using
Transformer backbone) under the uncertainty estimation setting. We use a depth
6 shallow layer Transformer to conduct this experiment. For the Transformer
baseline model, we take the predictive entropy as uncertainty. For SNGP +
Transformer, the entropy of the average of the Monte Carlo softmax samples is
used as uncertainty. We do not compare with DUE for the CIFAR-100 dataset, as
its training does not converge. SGD is used as the optimizer with the initial
learning rate set to 0.01. All models are trained with batch size 128.

The accuracy, NLL, AUROC, and AUPR results are shown in Table 6. The AUROC
metric indicates the quality of uncertainty since it measures the probability
that in-distribution (ID) and OOD samples can be separated [Mukhoti et al.,
2021]. From the results, we have the following observations:

(1) For OOD detection, The proposed LRFormer model outperforms all other
methods with Transformer backbone on both CIFAR-10 vs SVHN and CIFAR-100 vs
SVHN benchmarks. This superior OOD detection performance benefits from the
proposed LRSA regularization method, which solves both Lipschitzness and
contraction problems in dot-product self-attention layers, and enables
distance-preserving mapping in Transformer blocks.

(2) Notably, superior performance in OOD is achieved without sacrificing
LRFormer’s predictive performance. On the contrary, LRFormer even outperforms
the standard Transformer baseline in terms of classification accuracy on the
CIFAR-10 dataset, making LRFormer achieve the best performance in terms of all
the metrics compared with all other single-network methods.

(3) Furthermore, the proposed LRSA self-attention can be computed efficiently
using matrix operations, with minimal overhead compared to the original dot-
product self-attention. This ensures LRFormer’s performance gains come without
compromising computation costs.

####  4.3.4 Pre-training

The recent work Plex [Tran et al., 2022] comprehensively validated the
reliability of the large pretrained models. The high performance of the
Transformer results from pre-training on large-size datasets such as
ImageNet-21K, and LAION-5B[Schuhmann et al., 2022]. It performs worse than
CNNs if trained from scratch on small-size datasets. We use a depth 6 shallow
layer Transformer to conduct pretrained weight experiments. The pretrained
weights are loaded from the standard Transformer, sharing the same weight
schemes of MLP layers and position embedding layer. This is because these
layers in LRFormer have the same structure as the standard Transformer, so
pre-trained weights can be directly applied to them. Our experiments in Table
5 show that LRFormer can also benefit from these pre-trained weights.

In summary, pre-trained weights of models can be directly transferred to
LRFormer, which is very convenient for real-world applications.

###  4.4 Visualization

In order to show the interpretability of our model, we visualize the
uncertainty heat map generated by LRFormer together with the baseline methods
on the two moons 2D classification benchmark, which consists of two moon-
shaped data clusters separable by a non-linear decision boundary. We employ a
tiny Transformer architecture for this task, in which the depth is set to 9,
the hidden dimension is set to 24 and the number of heads is set to 8.

The uncertainty heat map comparisons are shown in Figure 1. Background color
visualizes the predictive uncertainty of each model, where yellow stands for
confidence and blue indicates uncertainty. All methods achieve 100% test
accuracy. From the results, we can observe that Deep Ensemble (Figure 1(a))
estimates its uncertainty based on how far away test samples are from the
decision boundary, without considering the data distribution. In Figure 1(b),
we can see that DUE without restrictions in the feature extractor, produces
similar predictive uncertainty to Deep Ensemble which is heavily influenced by
the distance from the decision boundary. Although SNGP can make allowances for
data distribution, the decision boundary still has an impact on uncertainty
estimation. The proposed LRFormer model, on the other hand, achieves near-
ideal uncertainty quantification of this benchmark thanks to its bi-Lipschitz
constraint in the LRSA self-attention layers, which allows it to maintain
better distance awareness.

##  5 Conclusion

In this paper, we present LRSA, a regularization method designed to address
overconfidence issues in Transformer structure models. By enforcing Bi-
Lipschitz constraints and self-attention mapping contractions with theoretical
guarantees, LRSA encourages the model to generate conservative predictions for
out-of-distribution (OOD) inputs, thereby improving its ability to separate
in-distribution (ID) data.

Our approach primarily focuses on the attention mechanism of the Transformer
architecture, which is a powerful and widely-used component in various natural
language processing and vision tasks. Moving forward, it would be beneficial
to extend our approach to incorporate these other modules within the
Transformer architecture. Exploring how different combinations of modules can
be leveraged to enhance performance across various tasks represents a
promising avenue for future research. Additionally, investigating the
relationship between Lipschitz regularity and other regularization techniques,
including weight decay, dropout, and label smoothing, would provide valuable
insights. Although these techniques have demonstrated effectiveness in
preventing overfitting and improving generalization, their connection to
Lipschitz regularity is not yet well-understood. Gaining a deeper
understanding of this relationship could unlock insights into the inner
workings of deep learning models and potentially lead to further performance
improvements.

In conclusion, our proposed LRSA method addresses overconfidence issues in
Transformer structure models by encouraging conservative predictions for OOD
inputs. While our focus has been on the attention mechanism, future research
directions involve incorporating other modules, exploring the relationship
between Lipschitz regularity and other regularization techniques, and
expanding LRFormer’s applicability to diverse models and domains. These
efforts contribute to advancing the field of deep learning and improving the
robustness and performance of state-of-the-art models.

###### Acknowledgements.

We thank all the anonymous reviewers for their insightful and thoughtful
comments.

## References

  * Amodei et al. [2016] Dario Amodei, Christopher Olah, Jacob Steinhardt, Paul Francis Christiano, John Schulman, and Dandelion Mané.  Concrete problems in ai safety.  _ArXiv_ , abs/1606.06565, 2016.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Blundell et al. [2015] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In _International Conference on Machine Learning_ , pages 1613–1622. PMLR, 2015.
  * Bradshaw et al. [2017] J F Bradshaw, Alexander G. de G. Matthews, and Zoubin Ghahramani.  Adversarial examples, uncertainty, and transfer testing robustness in gaussian process hybrid deep networks.  _arXiv: Machine Learning_ , 2017.
  * Cao and Zhang [2022] Senqi Cao and Zhongfei Zhang.  Deep hybrid models for out-of-distribution detection.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 4733–4743, 2022.
  * Dasoulas et al. [2021] George Dasoulas, Kevin Scaman, and Aladin Virmaux.  Lipschitz normalization for self-attention layers with application to graph neural networks.  _arXiv preprint arXiv:2103.04886_ , 2021.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _ArXiv_ , abs/1810.04805, 2019.
  * Dosovitskiy et al. [2021] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  _ArXiv_ , abs/2010.11929, 2021.
  * Esteva et al. [2017] Andre Esteva, Brett Kuprel, Roberto A. Novoa, Justin M. Ko, Susan M. Swetter, Helen M. Blau, and Sebastian Thrun.  Dermatologist-level classification of skin cancer with deep neural networks.  _Nature_ , 542:115–118, 2017.
  * Fort et al. [2021] Stanislav Fort, Jie Ren, and Balaji Lakshminarayanan.  Exploring the limits of out-of-distribution detection.  _Advances in Neural Information Processing Systems_ , 34:7068–7081, 2021.
  * Gal and Ghahramani [2016] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _international conference on machine learning_ , pages 1050–1059. PMLR, 2016.
  * Gillioz et al. [2020] Anthony Gillioz, Jacky Casas, Elena Mugellini, and Omar Abou Khaled.  Overview of the transformer-based models for nlp tasks.  _2020 15th Conference on Computer Science and Information Systems (FedCSIS)_ , pages 179–183, 2020.
  * Guo et al. [2017] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In _International conference on machine learning_ , pages 1321–1330. PMLR, 2017.
  * Gupta and Verma [2023] Kavya Gupta and Sagar Verma.  Certvit: Certified robustness of pre-trained vision transformers.  _arXiv preprint arXiv:2302.10287_ , 2023.
  * Hendrycks and Gimpel [2016] Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (gelus).  _arXiv: Learning_ , 2016.
  * Hu et al. [2023] Yihan Hu, Jiazhi Yang, Li Chen, Keyu Li, Chonghao Sima, Xizhou Zhu, Siqi Chai, Senyao Du, Tianwei Lin, Wenhai Wang, Lewei Lu, Xiaosong Jia, Qiang Liu, Jifeng Dai, Yu Qiao, and Hongyang Li.  Planning-oriented Autonomous Driving.  March 2023.  10.48550/arXiv.2212.10156.  arXiv:2212.10156 [cs].
  * Huang and Chen [2020] Yu Huang and Yue Chen.  Autonomous driving with deep learning: A survey of state-of-art technologies.  _ArXiv_ , abs/2006.06091, 2020.
  * Ioffe and Szegedy [2015] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  _ArXiv_ , abs/1502.03167, 2015.
  * Jiang et al. [2018] Heinrich Jiang, Been Kim, and Maya R. Gupta.  To trust or not to trust a classifier.  In _NeurIPS_ , 2018.
  * Katharopoulos et al. [2020] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _International Conference on Machine Learning_ , pages 5156–5165. PMLR, 2020.
  * Kim et al. [2021] Hyunjik Kim, George Papamakarios, and Andriy Mnih.  The lipschitz constant of self-attention.  In _ICML_ , 2021.
  * Kristiadi et al. [2020] Agustinus Kristiadi, Matthias Hein, and Philipp Hennig.  Being bayesian, even just a bit, fixes overconfidence in relu networks.  In _International conference on machine learning_ , pages 5436–5446. PMLR, 2020.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Lakshminarayanan et al. [2017] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _NIPS_ , 2017.
  * Lee et al. [2021] Kwonjoon Lee, Huiwen Chang, Lu Jiang, Han Zhang, Zhuowen Tu, and Ce Liu.  Vitgan: Training gans with vision transformers.  In _International Conference on Learning Representations_ , 2021.
  * Li et al. [2022] Yanghao Li, Chao-Yuan Wu, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik, and Christoph Feichtenhofer.  MViTv2: Improved Multiscale Vision Transformers for Classification and Detection.  _Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 4804–4814, 2022.
  * Liu et al. [2020] Jeremiah Zhe Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax-Weiss, and Balaji Lakshminarayanan.  Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.  _ArXiv_ , abs/2006.10108, 2020.
  * Liu et al. [2021] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 10012–10022, 2021.
  * Loshchilov and Hutter [2016] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ , 2016.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Ma and Wang [2023] Yunsheng Ma and Ziran Wang.  ViT-DD: Multi-Task Vision Transformer for Semi-Supervised Driver Distraction Detection.  _IEEE Intelligent Vehicles Symposium_ , 2023.  10.48550/arXiv.2209.09178.
  * Ma et al. [2023] Yunsheng Ma, Liangqi Yuan, Amr Abdelraouf, Kyungtae Han, Rohit Gupta, Zihao Li, and Ziran Wang.  M2DAR: Multi-View Multi-Scale Driver Action Recognition with Vision Transformer.  _Conference on Computer Vision and Pattern Recognition Workshops_ , 2023.  10.48550/arXiv.2305.08877.
  * Megginson [2012] Robert E Megginson.  _An introduction to Banach space theory_ , volume 183.  Springer Science & Business Media, 2012.
  * Miyato et al. [2018] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for generative adversarial networks.  _ArXiv_ , abs/1802.05957, 2018.
  * Mukhoti et al. [2021] Jishnu Mukhoti, Andreas Kirsch, Joost R. van Amersfoort, Philip H. S. Torr, and Yarin Gal.  Deterministic neural networks with inductive biases capture epistemic and aleatoric uncertainty.  2021\.
  * Müller et al. [2019] Rafael Müller, Simon Kornblith, and Geoffrey E Hinton.  When does label smoothing help?  _Advances in neural information processing systems_ , 32, 2019.
  * Netzer et al. [2011] Yuval Netzer, Tiejie Wang, Adam Coates, A. Bissacco, Bo Wu, and A. Ng.  Reading digits in natural images with unsupervised feature learning.  2011\.
  * Qi et al. [2023] Xianbiao Qi, Jianan Wang, Yihao Chen, Yukai Shi, and Lei Zhang.  Lipsformer: Introducing lipschitz continuity to vision transformers.  In _International Conference on Learning Representations_ , 2023.
  * Schuhmann et al. [2022] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al.  Laion-5b: An open large-scale dataset for training next generation image-text models.  _arXiv preprint arXiv:2210.08402_ , 2022.
  * Touvron et al. [2022] Hugo Touvron, Matthieu Cord, Alaaeldin El-Nouby, Jakob Verbeek, and Herv’e J’egou.  Three things everyone should know about vision transformers.  _ArXiv_ , abs/2203.09795, 2022.
  * Tran et al. [2022] Dustin Tran, Jeremiah Liu, Michael W. Dusenberry, Du Phan, Mark Collier, Jie Ren, Kehang Han, Zi Wang, Zelda Mariet, Huiyi Hu, Neil Band, Tim G. J. Rudner, Karan Singhal, Zachary Nado, Joost van Amersfoort, Andreas Kirsch, Rodolphe Jenatton, Nithum Thain, Honglin Yuan, Kelly Buchanan, Kevin Murphy, D. Sculley, Yarin Gal, Zoubin Ghahramani, Jasper Snoek, and Balaji Lakshminarayanan.  Plex: Towards Reliability using Pretrained Large Model Extensions.  _ICML Workshop_ , 2022.  10.48550/arXiv.2207.07411.  pre-training: perspectives, pitfalls, and paths forward.
  * Van Amersfoort et al. [2020] Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal.  Uncertainty estimation using a single deep deterministic neural network.  In _International conference on machine learning_ , pages 9690–9700. PMLR, 2020.
  * van Amersfoort et al. [2021] Joost R. van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal.  On feature collapse and deep kernel learning for single forward pass uncertainty.  2021\.
  * Vaswani et al. [2017] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _ArXiv_ , abs/1706.03762, 2017.
  * Wei et al. [2022] Hongxin Wei, Renchunzi Xie, Hao Cheng, Lei Feng, Bo An, and Yixuan Li.  Mitigating neural network overconfidence with logit normalization.  In _International Conference on Machine Learning_ , pages 23631–23644. PMLR, 2022.
  * Wightman et al. [2021] Ross Wightman, Hugo Touvron, and Hervé Jégou.  Resnet strikes back: An improved training procedure in timm.  _arXiv e-prints_ , pages arXiv–2110, 2021.
  * Williams and Rasmussen [2006] Christopher K Williams and Carl Edward Rasmussen.  _Gaussian processes for machine learning_ , volume 2.  MIT press Cambridge, MA, 2006.
  * Wilson and Izmailov [2020] Andrew Gordon Wilson and Pavel Izmailov.  Bayesian deep learning and a probabilistic perspective of generalization.  _ArXiv_ , abs/2002.08791, 2020.
  * Xu et al. [2022] Runsheng Xu, Zhengzhong Tu, Hao Xiang, Wei Shao, Bolei Zhou, and Jiaqi Ma.  CoBEVT: Cooperative Bird’s Eye View Semantic Segmentation with Sparse Transformers.  _Conference on Robot Learning_ , July 2022.  10.48550/arXiv.2207.02202.
  * Xu et al. [2023] Xiaojun Xu, Linyi Li, Yu Cheng, Subhabrata Mukherjee, Ahmed Hassan Awadallah, and Bo Li.  Certifiably robust transformers with 1-lipschitz self-attention.  2023\.
  * Xue et al. [2022] Feng Xue, Zi He, Chuanlong Xie, Falong Tan, and Zhenguo Li.  Boosting out-of-distribution detection with multiple pre-trained models.  _arXiv preprint arXiv:2212.12720_ , 2022.
  * Yang et al. [2023] Chenyu Yang, Yuntao Chen, Hao Tian, Chenxin Tao, Xizhou Zhu, Zhaoxiang Zhang, Gao Huang, Hongyang Li, Yu Qiao, Lewei Lu, Jie Zhou, and Jifeng Dai.  BEVFormer v2: Adapting Modern Image Backbones to Bird’s-Eye-View Recognition via Perspective Supervision.  _Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 17830–17839, 2023.

Mitigating Transformer Overconfidence via Lipschitz Regularization
(Supplementary Material)

##  Appendix A Proof for the Lipschitz Constant of LayerNorm

The LayerNorm operation [Ba et al., 2016] used in LRFormer can be expressed
as:

$\displaystyle\text{LN}(\mathbf{x})$
$\displaystyle=\frac{\mathbf{x}-\mu(\mathbf{x})}{\sqrt{\sigma^{2}(\mathbf{x})+\epsilon}}*\boldsymbol{\gamma}+\boldsymbol{\beta}$

where $\mathbf{x},\boldsymbol{\beta},\boldsymbol{\gamma}\in\mathbb{R}^{N}$,
$\mu(\mathbf{x})=\frac{1}{N}\sum_{i=1}^{N}x_{i}$,
$\sigma^{2}(\mathbf{x})=\frac{1}{N}\sum_{i=1}^{N}(x_{i}-\mu(\mathbf{x}))^{2}$.

WLOG, assume $N>2$ and not all $x_{i}$ are equal.

The derivatives of $\mu$ and $\sigma^{2}$ w.r.t $x$:

$\frac{\partial\mu}{\partial\mathbf{x}}=\frac{1}{N}\mathds{1}^{\top}$
$\frac{\partial\sigma^{2}}{\partial\mathbf{x}}=\frac{2}{N}(\mathbf{x}-\mu)^{\top}$

Take the derivative of $\text{LN}(\mathbf{x})_{i}$, the $i$th element of
$\text{LN}(\mathbf{x})$, with respect to $\mathbf{x}$ is:

$\displaystyle\begin{split}\frac{\partial\text{LN}(\mathbf{x})_{i}}{\partial\mathbf{x}}&=\gamma_{i}(\sigma^{2}+\epsilon)^{-\frac{1}{2}}\bigg{[}(\mathbf{e}_{i}-\frac{1}{N}\mathds{1})^{\top}-\frac{1}{N}(\sigma^{2}+\epsilon)^{-1}(x_{i}-\mu)(\mathbf{x}-\mu)^{\top}\bigg{]}.\end{split}$
(16)

where $\mathbf{e}_{I}\in\mathbb{R}^{N}$ is a one-hot vector with $1$ at the
$i$th element. Therefore,

$\displaystyle\frac{\partial\text{LN}(\mathbf{x})}{\partial\mathbf{x}}$
$\displaystyle=(\sigma^{2}+\epsilon)^{-\frac{1}{2}}\bigg{[}\text{diag}(\boldsymbol{\gamma})-\frac{1}{N}\boldsymbol{\gamma}\mathds{1}^{\top}-\frac{1}{N}(\sigma^{2}+\epsilon)^{-1}\text{diag}(\boldsymbol{\gamma})(\mathbf{x}-\mu)(\mathbf{x}-\mu)^{\top}\bigg{]}.$

$\left\|\text{diag}(\boldsymbol{\gamma})-\frac{1}{N}\boldsymbol{\gamma}\mathds{1}^{\top}\right\|_{\infty}=\frac{2(N-1)}{N}\max_{i}|\gamma_{i}|,$
(17)

Take the infinity-norm on both sides, we have:

$\displaystyle\left\|\frac{\partial\text{LN}(\mathbf{x})}{\partial\mathbf{x}}\right\|_{\infty}$
$\displaystyle=(\sigma^{2}+\epsilon)^{-\frac{1}{2}}\left\|\text{diag}(\boldsymbol{\gamma})-\frac{1}{N}\boldsymbol{\gamma}\mathds{1}^{\top}-\frac{1}{N}(\sigma^{2}+\epsilon)^{-1}\text{diag}(\boldsymbol{\gamma})(\mathbf{x}-\mu)(\mathbf{x}-\mu)^{\top}\right\|_{\infty}$
$\displaystyle\leq\epsilon^{-\frac{1}{2}}\bigg{(}\frac{2(N-1)}{N}\max_{i}|\gamma_{i}|+\frac{1}{N}\max_{i}|\gamma_{i}|N(N-2)\bigg{)}$
$\displaystyle\leq\epsilon^{-\frac{1}{2}}\max_{i}|\gamma_{i}|N.$

##  Appendix B Proof for the Lipschitz Constant of LRSA

The pair-wise LRSA function is expressed as:

$\displaystyle
S_{ij}=-\frac{\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}^{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}}$
(18)

$\displaystyle P_{i}=S_{i}(X)$

$\displaystyle P_{ij}=\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\leq 1$

To take the derivative $P_{ij}$, there are two cases.

When $t=j$:

$\displaystyle\frac{\partial P_{ij}}{\partial S_{it}}$
$\displaystyle=\frac{\partial P_{ij}}{\partial
S_{ij}}=\frac{\partial}{\partial
S_{ij}}\bigg{(}\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\bigg{)}=\frac{e^{S_{ij}}(\sum_{t=1}^{n}e^{S_{it}})-(e^{S_{ij}})^{2}}{(\sum_{t=1}^{n}e^{S_{it}})^{2}}$
(19)
$\displaystyle=\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\bigg{(}1-\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\bigg{)}=P_{ij}(1-P_{ij})$

When $t\neq j$:

$\displaystyle\frac{\partial P_{ij}}{\partial S_{it}}=\frac{\partial}{\partial
S_{it}}\bigg{(}\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\bigg{)}=-\frac{e^{S_{ij}}}{\sum_{t=1}^{n}e^{S_{it}}}\frac{e^{S_{it}}}{\sum_{t=1}^{n}e^{S_{it}}}=-P_{ij}P_{it}$

$\displaystyle\frac{\partial P_{ij}}{\partial
x_{k}}=\sum_{t=1}^{n}\frac{\partial P_{ij}}{\partial S_{it}}\frac{\partial
S_{it}}{\partial x_{k}}=P_{ij}(1-P_{ij})\frac{\partial S_{ij}}{\partial
x_{k}}-\sum_{t=1,t\neq j}^{n}P_{ij}P_{it}\frac{\partial S_{it}}{\partial
x_{k}}=P_{ij}\frac{\partial S_{ij}}{\partial
x_{k}}-P_{ij}\sum_{t=1}^{n}P_{it}\frac{\partial S_{it}}{\partial x_{k}}$ (20)

Take the infinity-norm on $S_{it}$, we get:

$\displaystyle\left\|\frac{\partial S_{it}}{\partial x_{k}}\right\|_{\infty}$
$\displaystyle=\left\|\frac{\partial}{\partial
x_{k}}\bigg{(}-\frac{\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}^{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}}\bigg{)}\right\|_{\infty}$
$\displaystyle=\left\|-\frac{2\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}}\frac{\partial\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}}{\partial
x_{k}}+\frac{\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}^{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}^{2}}\frac{\partial\left\|X^{\top}\right\|_{(\infty,2)}}{\partial
x_{k}}\right\|_{\infty}$
$\displaystyle\leq\left\|\frac{2\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}}\frac{\partial\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}}{\partial
x_{k}}\right\|_{\infty}+\left\|\frac{\alpha\left\|x_{i}^{\top}W_{Q}-x_{j}^{\top}W_{K}\right\|_{2}^{2}}{\left\|Q\right\|_{F}\left\|X^{\top}\right\|_{(\infty,2)}^{2}}\frac{\partial\left\|X^{\top}\right\|_{(\infty,2)}}{\partial
x_{k}}\right\|_{\infty}$
$\displaystyle\leq\frac{2\alpha}{\left\|Q\right\|_{F}}\frac{\left\|x_{i}^{\top}W_{Q}\right\|_{2}+\left\|x_{j}^{\top}W_{K}\right\|_{2}}{\left\|X^{\top}\right\|_{(\infty,2)}}\bigg{(}\frac{\partial\left\|x_{j}^{\top}W_{Q}\right\|_{2}}{\partial
x_{k}}+\frac{\partial\left\|x_{j}^{\top}W_{K}\right\|_{2}}{\partial
x_{k}}\bigg{)}+\frac{\alpha}{\left\|Q\right\|_{F}}\bigg{(}\frac{\left\|x_{i}^{\top}W_{Q}\right\|_{2}+\left\|x_{j}^{\top}W_{K}\right\|_{2}}{\left\|X^{\top}\right\|_{(\infty,2)}}\bigg{)}^{2}$
$\displaystyle\leq\frac{2\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}+\frac{\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}$
$\displaystyle=\frac{3\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}$

Thus,

$\displaystyle\left\|\frac{\partial P_{ij}}{\partial x_{k}}\right\|_{\infty}$
$\displaystyle=\left\|P_{ij}\frac{\partial S_{ij}}{\partial
x_{k}}-P_{ij}\sum_{t=1}^{n}P_{it}\frac{\partial S_{it}}{\partial
x_{k}}\right\|_{\infty}\leq
P_{ij}\frac{3\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}+P_{ij}\sum_{t=1}^{n}P_{it}\frac{3\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}$
$\displaystyle\leq\frac{6\alpha(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|Q\right\|_{F}}^{2}\leq\frac{6\alpha}{\left\|X\right\|_{F}}\cdot\frac{(\left\|W_{Q}\right\|_{2}+\left\|W_{K}\right\|_{2})}{\left\|W_{Q}\right\|_{F}}^{2}$

##  Appendix C Gaussian Process Layer

As an optional module in LRFormer, Gaussian Process (GP) with an RBF kernel
following SNGP [Liu et al., 2020] is capable of perserving the distance
awareness between input test sample and previously seen training data. This
approach makes sure the model returns a uniform distribution over output
labels when the input sample is OOD.

To make it end-to-end trainable, the Gaussian Process layer can be implemented
a two-layer network:

$\operatorname{logits}(x)=\Phi(x)\beta,\quad\Phi(x)=\sqrt{\frac{2}{M}}*\cos(Wx+b)$
(21)

Here, $x$ is the input, and $W$ and $b$ are frozen weights initialized
randomly from Gaussian and uniform distributions, respectively. $\Phi(x)$ is
Random Fourier Features (RFF) [Williams and Rasmussen, 2006]. $\beta$ is the
learnable kernel weight similar to that of a Dense layer. The layer outputs
the class prediction
$\operatorname{logits}(x)\in\mathbb{R}_{\operatorname{NumClasses}}$ .

##  Appendix D Experimental Details

In Table 7, we provide the training details used for reproducing the main
results in Tables above. The $Depth=12$ (pretraining) is the experimental
setup of the ImageNet1K dataset pretraining. The other hyperparameters follows
the same setting from DeiT III [Touvron et al., 2022].

Table 7: Hyperparameters for LRFormer Training. Hyperparameters | $Depth=6$ | $Depth=12$ |  $Depth=12$ (pretraining)
---|---|---|---
Layer depth | 6 | 12 | 12
Input size | $224\times 224$ | $224\times 224$ | $224\times 224$
Batch size | 128 | 32 | 32
Warm-up steps | 5 | 5 | 5
Optimizer | SGD | AdamW | AdamW
Learning rate | 0.01 | 0.006 | 0.004
Weight decay | 0.05 | 0.05 | 0.05
Learning rate scheduler | cosine | cosine | cosine
Training epochs | 100 | 100 | 100


# Permutation Invariant Recurrent Neural Networks for Sound Source Tracking
Applications

###### Abstract

Many multi-source localization and tracking models based on neural networks
use one or several recurrent layers at their final stages to track the
movement of the sources. Conventional recurrent neural networks (RNNs), such
as the long short-term memories (LSTMs) or the gated recurrent units (GRUs),
take a vector as their input and use another vector to store their state.
However, this approach results in the information from all the sources being
contained in a single ordered vector, which is not optimal for permutation-
invariant problems such as multi-source tracking. In this paper, we present a
new recurrent architecture that uses unordered sets to represent both its
input and its state and that is invariant to the permutations of the input set
and equivariant to the permutations of the state set. Hence, the information
of every sound source is represented in an individual embedding and the new
estimates are assigned to the tracked trajectories regardless of their order.

Keywords: Sound source tracking (SST), permutation-invariant recurrent neural
networks (PI-RNN)

##  1 Introduction

In recent years, the state-of-the-art of sound source localization established
by classic signal processing techniques has been surpassed by new systems
using deep-learning models [1]. These models use different input features and
network architectures, but most of them track the temporal evolution of the
signals using convolutional layers followed by recurrent layers [2, 3, 4].
Using these architectures, the latent representations at every hidden layer
are difficult to interpret and we cannot exploit the permutation invariance of
the tracking problem where, if we cannot apply any criteria to order or
classify the sources, any permutation of the sources should be considered
equally correct.

In [5], we proposed an icosahedral convolutional neural network (icoCNN) for
single source localization where the output of the last convolutional layer
can be interpreted as the probability distribution of the direction of arrival
(DOA) and we can obtain the estimated DOA as its expected value. Extending
this model to multi-source scenarios is straightforward and we just need to
increase the number of channels of the last convolutional layers to the
maximum number of concurrent sources $M$ that the model should be able to
localize. Following this approach, after computing the expected value of every
one of the $M$ probability distributions generated by the icoCNN, we obtain a
set of $M$ DOAs that should be considered invariant to the permutations of its
elements. In order to incorporate a recurrent neural network (RNN) after the
localization model to increase its temporal perceptive field and improve its
tracking capabilities, we could concatenate every element of the DOA set into
a single vector and use it as the input of a gated recurrent unit (GRU) [6] or
a long short-term memory (LSTM) layer [7]. However, we should expect the
output of a tracking system to not be affected by the order of the new
estimates at every time frame (i.e., to be invariant to the permutations of
the input set), and a conventional RNN operating over the concatenation of the
estimates would need to learn this property during the training instead of
being part of its architecture. In addition, in a tracking system, we can also
expect the association of a new estimate to the tracked trajectories be done
regardless of their order (i.e., be equivariant to the permutations of the
state set) but the state vector generated by a conventional RNN would contain
the information of every tracked trajectory in an unstructured way so we would
not be able to exploit this property either.

In this paper we present a permutation-invariant recurrent neural network (PI-
RNN) that takes an unordered set of embeddings as input (each one with the
information of one of the sources detected by the localization network) and
generates a recursive output, or state, that is also an unordered set of
embeddings with the information of every tracked trajectory. As we could
expect from a tracking system, the proposed architecture associates the
embeddings in the input set to the embeddings of the state set in a way that
is invariant to the permutations of the input set and equivariant to the
permutations of the state set.

To the best of our knowledge, this is the first recurrent layer that works
with sets instead of with vectors. The closest proposal in the literature is
probably the TrackFormer [8], a model for multiple object tracking on video
signals that is based on the DETR transformer [9, 10], a model for object
detection on images. The recursivity of the TrackFormer model is built around
the decoder of the DETR transformer by using the output obtained for a video
frame as the input for the following frame. Compared with the TrackFormer, the
PI-RNN is not a model but a layer that can be integrated easily into many
different models. In addition, it is based on an architecture, the
conventional GRU, that, unlike the transformer, was designed to be used in
recurrent loops.

Thanks to be taking into account the symmetries of the problem, the proposed
PI-RNN, compared with the conventional RNN, scales better with the number of
tracked sources and the amount of information stored in each one. Furthermore,
we present experiments proving that they can obtain better tracking results
than the conventional GRUs.

##  2 Network architecture

Figure 1: Architecture of the proposed permutation invariant recurrent layer.

Conventional RNNs use a $\textbf{h}(t)\in\mathbb{R}^{d_{h}}$ vector to store
the tracking state, which is updated at every time frame based on an input
vector $\textbf{x}(t)\in\mathbb{R}^{d_{x}}$ using fully connected perceptrons
whose computational complexity and number of trainable parameters grow
linearly with $d_{x}$ and quadratically with $d_{h}$. When applied to track up
to $M$ sources, the information of all the sources and tracked trajectories
are stored in these vectors without any structure, so there is a trade-off
between the number of sources $M$ that we can track, the amount of information
that we store about each one, and the model size and complexity.

In contrast to conventional RNNs, we propose to replace the input and state
vectors $\textbf{x}(t)$ and $\textbf{h}(t)$ with the sets of embedding
$\textbf{X}(t)=\\{\textbf{x}_{1}(t),\textbf{x}_{2}(t),...,\textbf{x}_{M_{X}}(t)\\}$
and
$\textbf{H}(t)=\\{\textbf{h}_{1}(t),\textbf{h}_{2}(t),...,\textbf{h}_{M_{H}}(t)\\}$
where every element $\textbf{x}_{i}(t)\in\mathbb{R}^{d_{x}}$ and
$\textbf{h}_{i}(t)\in\mathbb{R}^{d_{h}}$ contains information about a single
input detection or tracked trajectory respectively. For the sake of
simplicity, we will keep $M_{X}=M_{H}=M$ and $d_{x}=d_{h}=d$ during the rest
of the paper, however the proposed architecture can work with $M_{X}\neq
M_{H}$ or even with dynamic values that change during time, and it can be
easily extended to configurations with $d_{x}\neq d_{h}$.

In order to match every new embedding of the input set with the embeddings of
the state set, we can use a multi-head attention module [11], which is well
known for its use in transformer models and is invariant to the permutation of
the elements of its input sets:

$\begin{split}\mathbf{C}(t)=\operatorname{MultiHead}(&\mathbf{H}(t-1),\ \\\
&\mathbf{X}(t)\cup\mathbf{H}(t-1),\\\
&\mathbf{X}(t)\cup\mathbf{H}(t-1))\end{split}$ (1)
$\operatorname{MultiHead}(\mathbf{Q},\mathbf{K},\mathbf{V})=\operatorname{Concat}(\mathbf{head}_{1},...,\mathbf{head}_{N_{\textrm{heads}}})$
(2)
$\mathbf{head}_{i}=\operatorname{Attention}(\mathbf{Q}\mathbf{W}^{Q}_{i},\mathbf{K}\mathbf{W}^{K}_{i},\mathbf{V}\mathbf{W}^{V}_{i})$
(3)
$\operatorname{Attention}(\mathbf{Q}_{i},\mathbf{K}_{i},\mathbf{V}_{i})=\operatorname*{soft-
max}\left(\frac{\mathbf{Q}_{i}\mathbf{K}_{i}^{T}}{\sqrt{d_{k}}}\right)\mathbf{V}_{i},$
(4)

with the $\operatorname*{soft-max}(\cdot)$ operating across rows. With this
configuration, the generated set $\mathbf{C}(t)$ is invariant to the
permutations of $\mathbf{X}(t)$ and equivariant to the permutations of
$\mathbf{H}(t-1)$ as we would expect from a tracking system.

Figure 2: Architecture of the icoCNN used for evaluation. B is the batch size,
T is the number of temporal frames of the acoustic scenes, $H=2^{r}=8$ and
$W=2^{r+1}=16$ are the height and the width of the projections of the
icosahedral grid.

Finally, as shown in Fig. 1, once we have assigned the input embeddings to
their corresponding state embedding, we can just update every element of the
state set according to this assignation:

$\mathbf{h}_{i}(t)=\left[1-\mathbf{z}_{i}(t)\right]\odot\mathbf{h}_{i}(t-1)+\tilde{\mathbf{h}}_{i}(t)$
(5) $\mathbf{z}_{i}(t)=\sigma(\mathbf{c}_{i}(t)\mathbf{W}^{z})$ (6)
$\tilde{\mathbf{h}}_{i}(t)=\tanh(\mathbf{c}_{i}(t)\mathbf{W}^{h}),$ (7)

where $\odot$ denotes element-wise vector multiplication and
$\sigma(\cdot),\tanh(\cdot)$ denote sigmoid and hyperbolic tangent functions
respectively, applied to each element of their vector arguments. This gated
architecture is based on a simplified version of the minimal gated recurrent
unit [12], but we could design different architectures based on different
conventional recurrent architectures. As in conventional RNNs, the number of
trainable parameters grows quadratically with $d$, but in the case of the PI-
GRUs we have $M$ embeddings of size $d$ containing the information of every
tracked trajectory. Hence, we can expect our model to scale better when we
increase the number of sources we want to track or the amount of information
that we want to be able to represent for each one of them.

##  3 Evaluation

###  3.1 Experiment design

As a preliminary study of the performance of this new architecture, we decided
to add a PI-RNN after the icoCNN presented in [5] for single source
localization. As shown in Fig. 2, in order to extend the icoCNN to multi-
source localization, we just increased the number of output channels from $1$
to $M$. Fig. 3 represents the PI-RNN we used after the icoCNN: we first used a
multi-layer perceptron to project every ACCDOA [13] generated by the icoCNN
into an embedding of size $d$ and then we used those embeddings as the input
set of our PI-RNN. After the PI-RNN had associated every new estimate from the
icoCNN to one of the tracked trajectories, we added a conventional GRU
(operating independently over the embedding of every tracked trajectory so it
did not break the permutation invariance of the model) and, finally, we used a
linear layer to project the $d$-size embedding into a 3D ACCDOA. The initial
state of every embedding of the state set of the PI-RNN was learned during the
training of the model while, at every time frame, the embeddings of all the
inactive trajectories were reset (i.e., those who had lead to ACCDOAs with a
norm lower than 0.5).

Figure 3: Architecture of the PI-RNN used after the icoCNN in the evaluated
model. Figure 4: Architecture of the conventional RNN used after the icoCNN in
the baseline model.

(a)  (b)  (c)  (d)

Figure 5: Evaluation metrics for proposed PI-RNN and the baseline models.
Figure 6: Example of one of the test acoustic scenes. The solid line
represents the ground truth trajectories of the sources, the crosses the
defections estimated at the icoCNN (i.e., the input of the PI-RNN), and the
dashed line the trajectories estimated by the whole model (i.e., the output of
the PI-RNN). The color indicates to which of the outputs they correspond, so
the IDS are visible.

The method was compared to two baselines, a) the icoCNN without any kind of
recurrent layers, and b) the icoCNN with two conventional GRUs designed to
have a similar number of trainable parameters as the evaluated model (see Fig.
4). In order to avoid identity switches (IDSs) in the tracked trajectories, we
trained all the models using sliding permutation invariant training (sPIT)
[14]. To facilitate the training of the icoCNN, we added an auxiliary frame-
level permutation invariant training (fPIT) at its output in the models that
included recurrent layers after it.

We used the same synthetic dataset as in [14], where acoustic sources randomly
appeared and disappeared along 20-second-length scenes. As source signals, we
used speech utterances from the LibriSpeech corpus and we simulated them
following random trajectories in rooms with reverberation times from $T_{60}=$
0.2 to $1.3\text{\,}\mathrm{s}$ with the image source method. The maximum
number of concurrent active sources in a time frame was 3.

We used $M=10$ as the number of ACCDOA outputs of all our models since we
observed that it was beneficial to use a higher number than the maximum
possible number of active sources in the dataset (i.e., 3) and we used $d=128$
as embedding size for the input and state sets of the PI-RNN. This is a
preliminary study of this new architecture and further experiments should be
conducted for a better optimization of these hyperparameters.

###  3.2 Results

As we can see in Fig. 5, the proposed PI-RNN clearly outperforms the baselines
in terms of localization error and the frequency of the identity switches
while, as we can see in the detection error tradeoff (DET) curve, the trade-
off between false positives and misses remains is for all the evaluated
models. It is worth saying that both the conventional and the permutation-
invariant RNNs are receiving only spatial information about the estimated
sources. By modifying the model to include spectral information in their input
we could expect both models to improve their performance, with the PI-RNNs
scaling better to the amount of spectral information of each source and
therefore being able to better exploit it.

As an example, in Fig. 6 we can see one of the test acoustic scenes. We can
see how the output of the icoCNN had a high number of identity switches even
when only one source was active but the PI-RNN was able to fix these switches
and also reduce the localization error.

###  3.3 Attention matrices

We can interpret the attention matrix of the multi-head attention module of
the PI-RNN as an assignment matrix where each row indicates which elements of
the input and state set were employed to compute each element of the output
set.

(a)

(b)

(c)

(d)

Figure 7: Examples of attention matrices from the multi-head attention module
of the PI-RNN.

The attention matrix shown in Fig. 7a corresponds to the first frame where a
source appeared and we can see how it was detected at the 8 output of the
icoCNN (i.e., the 8 input of the PI-RNN) and the PI-RNN assigned it to its 9
output. In the attention matrix of the next time frame (Fig. 7b) we can see
that the 9 output of the PI-RNN was computed combining the information of the
new estimate at that frame with the corresponding recurrent state. A new
source was detected by the icoCNN at its 4 output in the time frame
corresponding to Fig. 7c and it was assigned to the 10 output of the PI-RNN.
Finally, in Fig. 7d we can see how, after an identity switch at the output of
the icoCNN, the PI-RNN was able to assign every new estimate to the correct
tracked trajectory fixing the identity switch.

##  4 Conclusions

We have presented a new RNN architecture whose input and state are presented
with sets instead of vectors and that is invariant to the permutation of the
elements of the input and equivariant to the permutations of the elements of
the state set. This new architecture is able to exploit the permutation
symmetries of the tracking problem and to outperform the conventional RNN in
the preliminary experiments presented in this paper. We expect the difference
between the performance of the PI-RNNs and the conventional RNNs to become
even greater when including more information of every source at their input.

## References

  * [1] P.-A. Grumiaux, S. Kitić, L. Girin, and A. Guérin, “A survey of sound source localization with deep learning methods,” J. Ac. Soc. America, vol. 152, pp. 107–151, 2022.
  * [2] S. Adavanne, A. Politis, J. Nikunen, and T. Virtanen, “Sound event localization & detection of overlapping sources using convolutional recurrent neural networks,” IEEE J. Sel. Topics Sig. Proc., vol. 13, pp. 34–48, 2019.
  * [3] L. Perotin, R. Serizel, E. Vincent, and A. Guérin, “CRNN-Based Multiple DoA Estimation Using Acoustic Intensity Features for Ambisonics Recordings,” IEEE J. Sel. Topics Sig. Proc., vol. 13, pp. 22–33, 2019\.
  * [4] Y. Cao, T. Iqbal, Q. Kong, M. B. Galindo, W. Wang, and M. D. Plumbley, “Two-Stage Sound Event Localization and Detection Using Intensity Vector and Generalized Cross-Correlation,” in Proc. of the Detection and Classification of Acoustic Scenes and Events 2019 Workshop (DCASE2019), New York University, 2019.
  * [5] D. Diaz-Guerra, A. Miguel, and J. R. Beltran, “Direction of Arrival Estimation of Sound Sources Using Icosahedral CNNs,” IEEE/ACM Trans. Audio, Speech, and Lang. Proc., vol. 31, pp. 313–321, 2023.
  * [6] K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio, “Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation,” in Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), (Doha, Qatar), pp. 1724–1734, Association for Computational Linguistics, Oct. 2014.
  * [7] S. Hochreiter and J. Schmidhuber, “Long Short-Term Memory,” Neural Computation, vol. 9, pp. 1735–1780, Nov. 1997.
  * [8] T. Meinhardt, A. Kirillov, L. Leal-Taixe, and C. Feichtenhofer, “TrackFormer: Multi-Object Tracking with Transformers,” in 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), (New Orleans, LA, USA), pp. 8834–8844, IEEE, June 2022.
  * [9] N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko, “End-to-End Object Detection with Transformers,” in Computer Vision – ECCV 2020 (A. Vedaldi, H. Bischof, T. Brox, and J.-M. Frahm, eds.), Lecture Notes in Computer Science, (Cham), pp. 213–229, Springer International Publishing, 2020.
  * [10] X. Zhu, W. Su, L. Lu, B. Li, X. Wang, and J. Dai, “Deformable DETR: Deformable Transformers for End-to-End Object Detection,” in International Conference on Learning Representations, Feb. 2022.
  * [11] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is All you Need,” in Advances in Neural Information Processing Systems, vol. 30, Curran Associates, Inc., 2017.
  * [12] G.-B. Zhou, J. Wu, C.-L. Zhang, and Z.-H. Zhou, “Minimal gated unit for recurrent neural networks,” International Journal of Automation and Computing, vol. 13, pp. 226–234, June 2016.
  * [13] K. Shimada, Y. Koyama, N. Takahashi, S. Takahashi, and Y. Mitsufuji, “ACCDOA: Activity-coupled cartesian direction of arrival representation for sound event localization and detection,” in IEEE Int. Conf. on Acoustics, Speech and Sig. Proc. (ICASSP), pp. 915–919, 2021.
  * [14] D. Diaz-Guerra, A. Politis, and T. Virtanen, “Position tracking of a varying number of sound sources with sliding permutation invariant training,” in 2023 31th European Signal Processing Conference (EUSIPCO), Sept. 2023.