∎

11institutetext:  Corresponding author: Li Zhang 22institutetext: 22email:
<EMAIL_ADDRESS>
1 School of Data Science, Fudan University, Shanghai, China
2 University of Surrey, Guildford, UK

# Softmax-free Linear Transformers

Jiachen Lu1 Junge Zhang1 Xiatian Zhu2 Jiafeng Feng1 Tao Xiang2 Li Zhang1

(15 March 2024)

###### Abstract

Vision transformers (ViTs) have pushed the state-of-the-art for visual
perception tasks. The self-attention mechanism underpinning the strength of
ViTs has a quadratic complexity in both computation and memory usage. This
motivates the development of approximating the self-attention at linear
complexity. However, an in-depth analysis in this work reveals that existing
methods are either theoretically flawed or empirically ineffective for visual
recognition. We identify that their limitations are rooted in the inheritance
of softmax based self-attention during approximations, that is, normalizing
the scaled dot-product between token feature vectors using the softmax
function. As preserving the softmax operation challenges any subsequent
linearization efforts. By this insight, a family of SOftmax-Free Transformers
(SOFT) are proposed. Specifically, a Gaussian kernel function is adopted to
replace the dot-product similarity, enabling a full self-attention matrix to
be approximated under low-rank matrix decomposition. For computational
robustness, we estimate the Moore-Penrose inverse using an iterative Newton-
Raphson method in the forward process only, while calculating its theoretical
gradients only once in the backward process. To further expand applicability
(e.g., dense prediction tasks), an efficient symmetric normalization technique
is introduced. Extensive experiments on ImageNet, COCO and ADE20K show that
our SOFT significantly improves the computational efficiency of existing ViT
variants. With linear complexity, much longer token sequences are permitted by
SOFT, resulting in superior trade-off between accuracy and complexity. Code
and models are available at https://github.com/fudan-zvg/SOFT.

###### Keywords:

Transformer linear complexity softmax normalization softmax-free Gaussian
attention

††journal: International Journal of Computer Vision

##  1 Introduction

(a)  (b)

Figure 1: Top-1 classification accuracy on ImageNet (Deng et al., 2009)
validation set with respect to parameters and the memory usage corresponding
to the token sequence length in practice compared to other methods. (a)
Comparison with CNN models: ResNet (He et al., 2016) and CoAtNet (Dai et al.,
2021) Transformer models: PVT (Wang et al., 2021), Swin (Liu et al., 2021),
DeiT (Touvron et al., 2021a), ViT (Dosovitskiy et al., 2021), T2T-ViT (Yuan et
al., 2021), Twins-SVT (Chu et al., 2021) and SAN10 (Zhao et al., 2020); (b)
Comparison with Transformer (Vaswani et al., 2017), Linformer (Wang et al.,
2020), Nyströformer (Xiong et al., 2021) and Performer (Choromanski et al.,
2021). The memory usage is measured with a batch size of 1 on a 16GB Tesla
V100.

Recently the step change brought by Transformers (Vaswani et al., 2017) in
natural language processing (NLP) (Devlin et al., 2019; Brown et al., 2020)
seems to have arrived in vision (Dosovitskiy et al., 2021; Yuan et al., 2021;
Zhu et al., 2021; Zheng et al., 2021). Indeed, with less inductive bias in its
architecture design than Convolution neural networks (CNNs), pure Vision
Transformer (ViT) (Dosovitskiy et al., 2021) and its variants have shown to be
able to outperform CNNs on various vision tasks (d’Ascoli et al., 2021; Jaegle
et al., 2021). However, there is a bottleneck in any Transformer based model,
namely its quadratic complexity in both computation and memory usage. This is
intrinsic to the self-attention mechanism: given a sequence of tokens (e.g.,
words or image patches) as input, the self-attention module iteratively learns
the feature representations by relating one token to all other tokens. This
results in a quadratic complexity $O(n^{2})$ with the token sequence length
$n$ in both computation (time) and memory (space) since an $n\times n$ sized
attention matrix needs to be computed and saved during inference. This problem
is particularly acute in vision: a 2D image after tokenization will produce a
far longer sequence than those in NLP even with a moderate spatial resolution.
This quadratic complexity thus prevents a ViT model from modeling images at
high spatial resolutions, which are often crucial for visual recognition
tasks.

A natural solution is to reduce the complexity of self-attention computation
via approximation. Indeed, there have been a number of attempts in NLP (Wang
et al., 2020; Choromanski et al., 2021; Kitaev et al., 2020; Xiong et al.,
2021). For example, (Wang et al., 2020) takes a naive approach by shortening
the length of Key and Value via learnable projections. Such a coarse
approximation would inevitably cause performance degradation. In contrast,
(Choromanski et al., 2021; Katharopoulos et al., 2020) both leverage the
kernel mechanism to approximate softmax normalization to linearize the
computation in self-attention. (Kitaev et al., 2020) instead adopts a hashing
strategy to selectively compute the most similar pairs. Recently, (Xiong et
al., 2021) uses Nyström matrix decomposition to reconstruct the full attention
matrix with polynomial iteration for approximating the pseudo-inverse of the
landmark matrix. Nonetheless, softmax normalization is simply duplicated
across the matrix decomposition process, which is theoretically unsound. We
empirically found that none of these methods are effective when applied to
vision (see Sec. 4.1).

In this work, we identify that the limitations of existing efficient
Transformers are caused by the use of softmax self-attention, and for the
first time propose a softmax-free Transformer. More specifically, in all
existing Transformers (with or without linearization), a softmax normalization
is needed on top of scaled dot-product between token feature vectors (Vaswani
et al., 2017). Keeping this softmax operation challenges any subsequent
linearization efforts. To overcome this obstacle, we introduce a novel
softmax-free self-attention mechanism, named as SOFT, with linear complexity
$O(n)$ in both space and time. Specifically, SOFT uses Gaussian kernel to
define the similarity (self-attention) function without the need for
subsequent softmax normalization. With this softmax-free attention matrix, we
further introduce a novel low-rank matrix decomposition algorithm for
approximation. The robustness of the approximation is theoretically guaranteed
by employing a Newton-Raphson method for reliably computing the Moore-Penrose
inverse of the matrix.

We make the following contributions. (I) We introduce a novel softmax-free
Transformer with linear space and time complexity. (II) Our attention matrix
approximation is achieved through a novel matrix decomposition algorithm with
theoretical guarantee. (III) To evaluate our method for visual recognition
tasks, we design a family of generic backbone architectures with varying
capacities using SOFT as the core self-attention component. Extensive
experiments show that with a linear complexity (Figure 1b), our SOFT models
can take in as input much longer image token sequences. As a result, with the
same model size, our SOFT outperforms the state-of-the-art CNNs and ViT
variants on ImageNet (Deng et al., 2009) classification in the
accuracy/complexity trade-off (Figure 1a).

A preliminary version of this work was presented in NeurIPS 2021 spotlight (Lu
et al., 2021). In this paper, we have further extended our conference version
as follows: (i) We improve the efficiency and robustness in computing Moore-
Penrose inverse by using an iterative method in the forward process only while
calculating its theoretical gradient only once in the backward propagation.
(ii) We analyze the limitations of our preliminary SOFT from a matrix spectral
norm perspective, revealing the importance of normalization for enhancing the
model’s task generalizability. (iii) We prove that the preliminary SOFT
experiences a second-order increase in matrix spectral norm relative to the
matrix size, making it fail in dense vision problems. (iv) To address these
limitations, we propose a normalized softmax-free self-attention, keeping
linear complexity while enhancing performance, supported by both theoretical
proof and extensive experiments. (v) The improved SOFT outperforms the state-
of-the-art CNNs and ViTs for classification on ImageNet (Deng et al., 2009),
object detection on COCO (Lin et al., 2014) and semantic segmentation on
ADE20K (Zhou et al., 2019).

##  2 Related work

###  2.1 Vision Transformers

There is a surge of research interests recently in exploiting Transformers for
visual recognition tasks (Wang et al., 2018, 2021; Guo et al., 2022; Yuan et
al., 2021; Touvron et al., 2021a; Zhang et al., 2020), inspired by their
remarkable success in NLP (Vaswani et al., 2017; Devlin et al., 2019; Brown et
al., 2020). Core to these NLP and vision transformers is the same self-
attention mechanism (Vaswani et al., 2017) that computes a self-attention
matrix by exhaustively comparing token pairs. This means a quadratic
complexity with the sequence length in both space and time, which thus limits
the scalability of Transformers in dealing with long sequences. This
limitation is more serious in vision than NLP: To process an image with at
least thousands of pixels, patch-wise tokenization is a must for Transformers
to control the computational cost. Given higher resolution images, the patch
size also needs to be enlarged proportionally sacrificing the spatial
resolution. This limits the capability of Transformers, e.g., learning fine-
grained feature representation as required in many visual recognition tasks.

###  2.2 Linear Transformers

Recently, there have been a number of linear/efficient variants (Choromanski
et al., 2021; Wang et al., 2020; Katharopoulos et al., 2020; Kitaev et al.,
2020; Tay et al., 2023; Peng et al., 2021; Kasai et al., 2021) of Transformers
in NLP. For example, (Wang et al., 2020) learns to shrink the length of Key
and Value based on a low-rank assumption. (Kitaev et al., 2020) adopts a
hashing strategy to selective the most similar pairs and only compute
attention among them. (Choromanski et al., 2021; Katharopoulos et al., 2020)
utilize different kernel functions for approximating softmax-based self-
attention matrix. (Peng et al., 2021) applies random feature mapping on the
sequences to approach the original softmax function. (Kasai et al., 2021)
decreases the time and memory consumption of the attention matrix by replacing
the softmax function with its linear-complexity recurrent alternative. When
applied to visual recognition tasks, however, we show that these models have
considerable performance degradation compared to the standard Transformers
(Vaswani et al., 2017) (see Sec. 4.1).

The most related work to SOFT is (Xiong et al., 2021) which uses the Nyström
matrix decomposition to avoid computing the full attention matrix. However,
this method suffers from several theoretical defects: (1) As the standard
self-attention needs to apply row-wise softmax normalization on the full
attention matrix, a direct application of matrix decomposition is infeasible.
As a workaround without solid theoretical support, softmax is simply applied
to all the ingredient matrices in (Xiong et al., 2021). Such an approximation
is not guaranteed theoretically. (2) With a polynomial iteration method, it is
not guaranteed that the generalized attention matrix inverse can be computed
when the matrix is a nearly singular one in practice. In contrast to all the
above methods, in this paper we propose a softmax-free self-attention
mechanism that facilitates matrix decomposition for complexity minimization
with theoretical guarantees.

##  3 Method

###  3.1 Softmax-free self-attention formulation

A schematic illustration of our model is given in Figure 2. Let’s first look
at our attention module design. Given a sequence of $n$ tokens
$X\in\mathbb{R}^{n\times d}$ with each token represented by a $d$-dimensional
feature vector, self-attention (Vaswani et al., 2017) aims to discover the
correlations of all token pairs exhaustively.

Formally, $X$ is first linearly projected into three $d_{e}$-dimensional
spaces (query, key, and values) as:

$Q=XW_{q}\in\mathbb{R}^{n\times d_{e}},K=XW_{k}\in\mathbb{R}^{n\times
d_{e}},V=XW_{v}\in\mathbb{R}^{n\times d_{e}},$ (1)

where $W_{q},W_{k},W_{v}\in\mathbb{R}^{d\times d_{e}}$ are learnable matrices.
For equation simplicity, we omit the multi-head notation in self-attention
operations. However, it should be noted that multi-head mechanisms are
employed throughout. Self-attention can be expressed in a generic formulation
as:

$y_{i,:}=\sum_{j=1}^{n}\alpha(Q_{i,:},K_{j,:})\odot V_{j,:},$ (2)

where $\odot$ is the Hadamard product, and $i,j\in\\{1,\cdots,n\\}$ index the
tokens. The key self-attention function
$\alpha:\mathbb{R}^{d_{e}}\times\mathbb{R}^{d_{e}}\rightarrow\mathbb{R}$ is
composed of a nonlinear function $\beta:\mathbb{R}\rightarrow\mathbb{R}$ and a
relation function
$\gamma:\mathbb{R}^{d_{e}}\times\mathbb{R}^{d_{e}}\rightarrow\mathbb{R}$. A
dominant instantiation of $\alpha$ is the scaled dot-product based softmax
self-attention (Vaswani et al., 2017), defined as

$\beta(\cdot)=\text{softmax}(\cdot),\quad\gamma(Q_{i,:},K_{j,:})=\frac{1}{\sqrt{d_{e}}}\cdot
Q_{i,:}^{\top}K_{j,:}.$ (3)

Whilst this softmax self-attention has been the de facto choice and seldomly
questioned, as discussed earlier it is not necessarily suited for
linearization. To facilitate the design of linear self-attention, we introduce
a softmax-free self-attention function with the dot-product replaced by a
Gaussian kernel as:

$\beta^{\prime}(\cdot)=\text{exp}(\cdot),\quad\gamma^{\prime}(Q_{i,:},K_{j,:})=-\frac{1}{2\sqrt{d_{e}}}\cdot\|Q_{i,:}-K_{j,:}\|_{2}^{2}.$
(4)

A dissection of Transformers by Tsai et al. (Tsai et al., 2019) reveals
negligible performance differences between asymmetric and symmetric kernels.
To maintain the symmetric properties of the attention matrix as defined in Eq.
(3), we opt for identical projection matrices $W_{q}$ and $W_{k}$ in Eq. (1),
effectively setting $Q=K$. To further investigate the impact of symmetric
kernels, additional experiments are conducted in our ablation studies of
Section 4.4. Our self-attention matrix is then written as:

$S_{i,j}=\text{exp}\left(-\frac{1}{2\sqrt{d_{e}}}\cdot\|Q_{i,:}-K_{j,:}\|_{2}^{2}\right).$
(5)

For notation simplicity, we define the matrix formulation as:
$S=\text{exp}\left(Q\ominus K\right)$.

Remarks: Our self-attention matrix $S$ has three important properties: (1) It
is symmetric; (2) All the elements lie in a unit range of $[0,1]$; (3) All
diagonal elements hold the largest value $1$ (self-reinforced), with the
bottom ones (corresponding to most dissimilar token pairs) being close to $0$.
As Gaussian kernel is a positive definite kernel (Fasshauer, 2011), $S$ is
deemed a Gram matrix. However, we find that when using our kernel-based self-
attention matrix $S$ without linearization, the training of a transformer
fails to converge.  More discussion can be found in Section 3.3.

Figure 2: Schematic illustration of the proposed softmax-free self-attention
(SOFT) method. P.E.: Position embedding. Dash lines: linear projection. dh:
the hidden dim of each attention head. $\circ$ denotes the matrix dot product.

###  3.2 Low-rank regularization via matrix decomposition with linear
complexity

To solve the convergence and quadratic complexity problems, we leverage matrix
decomposition as a unified solution with low-rank regularization. In
particular, we consider Nyström (Williams and Seeger, 2000), which is
originally a low-rank matrix approximation algorithm. This enables our model’s
complexity to be reduced significantly without computing the full self-
attention matrix $S$.

We make this choice because our $S$ is positive semi-definite (i.e., a Gram
matrix) without follow-up normalization which are all necessary conditions for
Nyström. In contrast, (Xiong et al., 2021) totally ignores these requirements,
leading to theoretical flaw in its approximation.

To define the Nyström method formally, let us express
$S=\text{exp}\left(Q\ominus K\right)$ as a block matrix:

$S=\left[\begin{array}[]{cc}A&B\\\
B^{\top}&C\end{array}\right]\in\mathbb{R}^{n\times n},$ (6)

where $A\in\mathbb{R}^{m\times m}$, $B\in\mathbb{R}^{m\times(n-m)}$,
$C\in\mathbb{R}^{(n-m)\times(n-m)}$ with $m\ll n$. Through Nyström
decomposition (see derivative details in Appendix A), an approximation can be
represented as:

$\hat{S}=\begin{bmatrix}A\\\
B^{\top}\end{bmatrix}A^{\dagger}\begin{bmatrix}A&B\end{bmatrix}=P^{\top}A^{\dagger}P,\quad\text{where}\quad
P=\begin{bmatrix}A&B\end{bmatrix},$ (7)

and $A^{\dagger}$ is the Moore-Penrose (a generalized) inverse of $A$.

Input: $Q\in\mathbb{R}^{n\times d_{e}}$, sampling function $f_{s}$

Sampling $\ \widetilde{Q}\leftarrow f_{s}(Q)$ ;

$A\leftarrow\text{exp}(\widetilde{Q}\ominus\widetilde{Q})$,
$P\leftarrow\text{exp}(\widetilde{Q}\ominus Q)$;

$\hat{S}\leftarrow P^{\top}\text{NR}(A)P$;

Output: $\hat{S}$

Algorithm 1 SOFT: Softmax-free attention

Input: $A\in\mathbb{R}^{m\times m}$, and $\mathcal{T}\in\mathbb{Z}^{+}$

$\alpha=2/\|A\|_{1}^{2}$.Initialize $A_{0}\leftarrow\alpha A$;

for  _$k$ from $1$ to $\mathcal{T}$ _ do

$A_{k}\leftarrow 2A_{k-1}-A_{k-1}AA_{k-1}$

end for

Output: $A_{\mathcal{T}}$

Algorithm 2 NR: Newton-Raphson iteration

Sampling:  In the standard Nyström formulation, $A$ and $B$ are sub-matrices
of $S$ obtained by randomly sampled $m$ tokens, denoted as $\widetilde{Q}$. We
call the sampled $\widetilde{Q}$ as bottleneck tokens. However, we find
empirically that random sampling is considerably sensitive to the choice of
$m$. We hence explore two additional options to leverage the structural prior
of visual data: (1) Using one convolutional layer with kernel size $k$ and
stride $k$ to learn $\widetilde{Q}$, and (2) Using average pooling with kernel
size $k$ and stride $k$ to generate $\widetilde{Q}$. For both, we need to
reshape $Q$ to the form of $\mathbb{R}^{H\times W\times d_{e}}$. Each slide of
convolution or pooling produces a token. We set $k$ according to the length of
$Q$ such that $m$ tokens can be obtained. Our experiments show that a
convolution layer performs better in accuracy. We therefore use a convolution
layer by default.

As $K$ is identical to $Q$, we have $\widetilde{K}=\widetilde{Q}$. Given these
$m$ tokens, we then compute bottleneck attention $A$ and $P$ as:

$A=\text{exp}(\widetilde{Q}\ominus\widetilde{K}),\quad
P=\text{exp}(\widetilde{Q}\ominus K).$ (8)

We finally obtain a regularized self-attention matrix $\hat{S}$ of SOFT as:

$\hat{S}=\text{exp}\left(Q\ominus\widetilde{K}\right)\left(\text{exp}\left(\widetilde{Q}\ominus\widetilde{K}\right)\right)^{\dagger}\text{exp}\left(\widetilde{Q}\ominus
K\right).$ (9)

The overall SOFT is summarized in Algorithm 1. The low-rank regularization is
conducted as follows. For computing the attention score between any two
tokens, we first correlate each of them with sampled tokens using our self-
attention function (Equation (5)); With this correlation representation we
then compute their similarity under the modulation of the generalized inverse
of $\widetilde{Q}$’s correlation matrix. Similar as the standard Nyström, our
design associates the input tokens w.r.t. a small space spanned by a set of
sampled tokens, giving a proper estimation of the original attention
relationships subject to a low-rank constraint. This method is proved in
Appendix A.

Moore-Penrose inverse: An accurate and commonly used way to calculate the
Moore-Penrose inverse is to use Singular Value Decomposition (SVD). Given
$A\in\mathbb{R}^{m\times m}$ and its SVD form $A=U\Sigma V^{\top}$ where $U,V$
are $m\times m$ unitary matrices and $\Sigma$ is a $m\times m$ diagonal
matrix, the Moore-Penrose inverse of $A$ is
$A^{\dagger}=V\Sigma^{\dagger}U^{\top}$. Nevertheless, SVD is not friendly to
the training process on GPU hence harming the model training efficiency. To
solve this issue, we adopt the Newton–Raphson method. It is an iterative
algorithm with the $(k+1)$-th iteration formulated given the previous
iteration as:

$A_{k+1}=2A_{k}-A_{k}AA_{k},\quad\text{and}\quad A_{0}=\alpha A.$ (10)

We now prove that $A_{k}$ finally converges to Moore-Penrose inverse of
$A_{m\times m}$, if $\alpha$ is sufficiently small (Ben-Israel and Cohen,
1966).

###### Proposition 1

When $\alpha$ is sufficiently small, $A_{k+1}=2A_{k}-A_{k}AA_{k}$, $A_{k}$
converges to $A^{\dagger}$.

The proposition is proved in Appendix B. Though $\alpha=2/\|A\|_{1}^{2}$ which
ensures good convergence behavior in Algorithm 2, in practice, we find that
using an alternative form gives more stable training and faster convergence.
Specifically, in $\|I-A\frac{2\beta^{n}}{\|A\|_{1}^{2}}\|_{1}\leq 1$ where
$\beta$ equals to $0.5$, we find the smallest $n_{i}$ that holds this
inequality. Then, we initialize $\alpha$ as
$\alpha=\frac{2\beta^{n_{i}}}{\|A\|_{1}^{2}}$.

The following proposition comes with the proof of Proposition 1:

###### Proposition 2

$\|AA_{k}A-A\|$ and $\|A_{k}-A^{\dagger}\|$ decreases to $0$ monotonously, if
$\alpha$ is sufficiently small.

###### Proof

Note that when we multiply $A$ on both sides of (33), the equation turns to
be:

$\displaystyle A-AA_{k+1}A$
$\displaystyle=A(A^{\dagger}-A_{k})A(A^{\dagger}-A_{k})A$ (11)
$\displaystyle=(AA^{\dagger}-AA_{k})(A-AA_{k}A).$

Similarly norm both sides of (11), considering that
$\|AA^{\dagger}-AA_{k}\|\rightarrow 0$ and $\|AA^{\dagger}-AA_{k}\|<1$ always
holds, $\|A-AA_{k}A\|$ monotonically decreases to $0$. The inequality (34)
implies that $\|A_{k}-A^{\dagger}\|$ decreases to $0$ monotonously .

Note although $\|A-AA_{k}A\|$ monotonically decreases to $0$,
$\|A_{k}AA_{k}-A_{k}\|$ cannot be proved to be so yet.

This ensures that our estimated inverse is sufficiently accurate for matrix
decomposition, subject to that our SOFT attention is regularized. In training
stage, we find that bottleneck matrix $A$ is always non-singular in practice
and its inverse $A^{-1}$ thus exists. Therefore, the iteration can be avoided
in back propagation because the differential of matrix inverse can be
explicitly expressed as

$\nabla_{x}\mathcal{L}=-Y^{\top}(\nabla_{Y}\mathcal{L})Y^{\top},$ (12)

where $X\in\mathbb{R}^{m\times m}$ is an non-singular matrix and $Y$ is the
inverse of $X$, i.e., $Y=X^{-1}$, $\nabla_{Y}\mathcal{L}$ is the gradient of
loss $\mathcal{L}$ on $Y$ and $\nabla_{X}\mathcal{L}$ is the gradient of loss
$\mathcal{L}$ on $X$. This can accelerate the training, as validated in Sec.
4.1. The theoretical proof is shown in Appendix B.

Complexity:  We summarize the complexity of SOFT in space and time. For time
complexity, it involves: (1) Sampling: $\mathcal{O}(nd_{e})$. (2) Calculating
three decomposed matrices:
$\mathcal{O}(nmd_{e}+mnd_{e}+m^{2}d_{e})=\mathcal{O}(2mnd_{e}+m^{2}d_{e})$;
(3) Moore-Penrose inverse: $\mathcal{O}(\mathcal{T}\times
m^{3})=\mathcal{O}(\mathcal{T}m^{3})$, where $\mathcal{T}$ is the iteration
steps. (4) All matrix multiplication:
$\mathcal{O}(nm^{2}+mnd_{e}+mnd_{e})=\mathcal{O}(nm^{2}+2mnd_{e})$. The total
time complexity is
$\mathcal{O}((d_{e}+4md_{e}+m^{2})n+\mathcal{T}m^{3}+d_{e}m^{2})$. The space
complexity is decided by four decomposed matrices with $\mathcal{O}(n\times
m)+\mathcal{O}(m\times m)+\mathcal{O}(m\times n)+\mathcal{O}(n\times
d_{e})=\mathcal{O}((2m+d_{e})n+m^{2})$. As we keep $m$ ($m\ll n$) a fixed
constant in our model, both time and space complexity are $\mathcal{O}(n)$,
making SOFT a linear self-attention.

###  3.3 Attention normalization

Whilst the above SOFT formulation is competitive for image classification,
transferring the pre-trained model to downstream tasks with different input
token sequence lengths is limited. We conduct analysis on the model’s
sensitivity against input perturbation. As suggested in (Yoshida and Miyato,
2017), all parts of a model should have small spectral norm (i.e., matrix
2-norm or the largest singular value of a matrix). Concretely, by the triangle
inequality $\|XY\|_{2}\leq\|X\|_{2}\|Y\|_{2}$ where $X$ and $Y$ are any real
matrix, any part with large spectral norm could lead to significant error
accumulation. This also applies for self-attention matrix, since it is often
symmetric and non-negative definite and its spectral norm corresponds to the
largest eigenvalue. We provide more theoretical analysis below. Note that both
scale dot-product kernel (linear kernel) and Gaussian kernel are positive
definite kernels, so they have non-negative eigenvalues.

###### Proposition 3

In scaled dot-product based softmax self-attention, assume
$\lambda_{1}\geq\lambda_{2}\geq\cdots\lambda_{n}\geq 0$ are eigenvalues of
self-attention matrix $S_{softmax}\in\mathbb{R}^{n\times n}$, then
$\lambda_{1}\leq 1$.

###### Proof

We rewrite the softmax self-attention as

$S_{softmax}=D^{-1}A,$ (13)

where $A\in\mathbb{R}^{n\times n}$ is a real symmetric matrix, and
$D=\text{diag}(A\mathbbm{1}_{n})$. We consider the graph Laplacian $L$ of
matrix $A$ defined by $L=D-A$, then the noramlized graph Laplacian can be
expressed as

$L_{rw}=D^{-1}L=D^{-1}(D-A)=I-D^{-1}A,$ (14)

where, $I$ is an identity matrix. According to (Von Luxburg, 2007), $L_{rw}$
is a real semi-definite matrix, so all the eigenvalue of $L_{rw}$ is non-
negative. Therefore, eigenvalues of $I-D^{-1}A$ are non-negative, which leads
to the fact that eigenvalues $\lambda_{1},\cdots,\lambda_{n}$ of $D^{-1}A$ is
less or equal to 1.

This proposition means that softmax normalization is useful in restricting the
range of self-attention matrix’s eigenvalues to $[0,1]$ and eventually the
effect of error accumulation. This is an important role softmax operation
plays in improving generalizability with standard self-attention. However,
this is not the case for our softmax-free self-attention as formulated above.

###### Proposition 4

In Gaussian kernel-based self-attention, if
$\lambda_{1}\geq\lambda_{2}\geq\cdots\lambda_{n}\geq 0$ are eigenvalues of
self-attention matrix $S_{gaussian}\in\mathbb{R}^{n\times n}$, then
$\lambda_{1}\leq n$.

###### Proof

$\sum_{i=1}^{n}\lambda_{i}=\text{Tr}(A)=n,$

since the diagonal elements of Gaussian kernel-based self-attention is always
1. Therefore, with the fact that all the eigenvalues are positive, we have
$\lambda_{1}\leq n$.

A larger upper bound of eigenvalue with our proposed self-attention could thus
lead to less generalizability, due to a trend of higher error accumulation.
Specifically, for a softmax-free self-attention matrix,
$\hat{S}=P^{\top}A^{\dagger}P$, we have

$\|\hat{S}\|_{2}=\|P^{\top}A^{\dagger}P\|_{2}\leq\|P\|_{2}^{2}\|A^{\dagger}\|_{2}.$
(15)

For the bottleneck matrix of softmax-free attention $A\in\mathbb{R}^{m\times
m}$, we assume it is $k$-connected ($k<<m$), i.e., there are $k$ fields
disconnected with each other. This is because a field represents a semantic
part of an image and the number is often small.

###### Proposition 5

Assume the bottleneck matrix of softmax-free attention ,
$A\in\mathbb{R}^{m\times m}$ is $k$-connected. If
$\lambda_{1}\geq\lambda_{2}\geq\cdots\lambda_{m}\geq 0$ are eigenvalues of
$A^{\dagger}$, then $\lambda_{1}=\mathcal{O}(m^{2})$ and
$\|A^{\dagger}\|_{2}=\mathcal{O}(m^{2})$.

The proposition is proved in Appendix C. This proposition indicates that
$\|\hat{S}\|_{2}$ is proportional quadratically with the length $m$ of
bottleneck token sequence. It implies that the above SOFT formula would be
limited to the applications with short bottleneck token sequences (e.g., image
classification).

###### Proposition 6

For the bottleneck matrix of SOFT self-attention $A\in\mathbb{R}^{m\times m}$,
we have

$\|D^{-1/2}A^{\dagger}D^{-1/2}\|_{2}=\mathcal{O}(m),$ (16)

where $D=\text{diag}(A\mathbbm{1}_{m})$ and $\mathbbm{1}_{m}$ is an all one
$m$-D vector.

###### Proof

$\|D^{-1/2}A^{\dagger}D^{-1/2}\|_{2}\leq\|D^{-1/2}\|_{2}^{2}\|A^{\dagger}\|_{2}=\|D^{-1}\|_{2}\|A^{\dagger}\|_{2},$

also, $D$ is a diagonal matrix,

$\|D^{-1}\|_{2}\|A^{\dagger}\|_{2}=\|D^{-1}A^{\dagger}\|_{2}\
=\|A_{n}^{\dagger}\|^{2}=\mathcal{O}(m)$

This proposition suggests that symmetric normalization can lower the largest
eigenvalue of $A^{\dagger}$, reducing the spectral norm of $\hat{S}$.

Attention normalization:  In light of the above theorem, we further introduce
the normalization of SOFT as:

$\hat{S}=\text{exp}\left(Q\ominus\widetilde{K}\right)D^{-\frac{1}{2}}\left(\text{exp}\left(\widetilde{Q}\ominus\widetilde{K}\right)\right)^{\dagger}D^{-\frac{1}{2}}\text{exp}\left(\widetilde{Q}\ominus
K\right),$ (17)

where
$D=\text{diag}\left(\text{exp}\left(\widetilde{Q}\ominus\widetilde{K}\right)\mathbbm{1}_{m}\right)$.
This yields a tiny increase of $\mathcal{O}(m^{2})$ in the computational
complexity, which can be further reduced to $\mathcal{O}(\log{m})$ on GPU by
parallel reduction algorithm (Cheng et al., 2014). This discovery unleashes
new possibilities for SOFT and paves the way for its expanded usage,
particularly in applications that demand dense information reasoning (such as
object detection and semantic segmentation). This results in our SOFT++
formulation.

###  3.4 Instantiations

Figure 2 shows how our proposed softmax-free self-attention block (SOFT block)
can be implemented in a neural network. We replace the self-attention block
with our SOFT block in the traditional Transformer, that is, we stack a SOFT
block with a feed forward residual block (Dosovitskiy et al., 2021) to form a
softmax-free Transformer layer (SOFT layer).

Focusing on the general image recognition tasks, we integrate our SOFT layer
into the recent pyramidal Transformer architecture (Wang et al., 2021) to form
our final model SOFT. Further, several improvements are introduced in patch
embedding (i.e., tokenization). Specifically, unlike (Wang et al., 2021) that
uses a combination of non-overlapping convolution and layer normalization (Ba
et al., 2016), we adopt a stack of overlapping convolutions, batch
normalization (Ioffe and Szegedy, 2015) and ReLU non-linearity. Concretely,
the $\tt STEM$ is implemented by 3 units of $\tt 3x3\;Conv{\to}BN{\to}ReLU$,
with the stride of 2, 1, 2 respectively. Then, one such unit is applied to
each of three following down-sampling operations with stride of 2 in the
multi-stage architecture.

The architecture hyper-parameters of SOFT are: $d$: the input channel
dimension of SOFT layer. $d_{e}$: the embedding dimension of tokens in SOFT
block. In practice, we set $d_{e}=d$. $h$: the head number of SOFT block.
$d_{h}$: the channel dimension of each head and $d_{h}=d_{e}/h$. $n$: the
input token sequence length of a SOFT block. $m$: the bottleneck token
sequence length of SOFT block. $sp$: the sampling ratio of token sequence
length sampling, which is the ratio between input token sequence length and
the bottleneck token sequence length. $e$: the expansion ratio of the 2-layer
feed forward block. In SOFT, for all the stages we set $d_{h}=32$, $e=4$ and
$m=49$, $sp$ varies in each stage according to the input token sequence
length. Table 1 details the family of our SOFT configurations with varying
capacities (depth and width).

Table 1: Architecture specifications of SOFT variants. sp.: sampling ratio. -d: the hidden dimension. -h: the number of heads in the self-attention block. C33-BN-ReLU: three 3x3 Conv-BN-ReLU, with the stride of 2, 1, 2 respectively. C31-BN-ReLU: one 3x3 Conv-BN-ReLU, with a stride of 2.  | Tiny | Small | Medium | Large
---|---|---|---|---
Stage 1 | C33-BN-ReLU, 64-d
$\begin{bmatrix}\text{sp. 8x8,}\\\ \text{64-d, 2-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 8x8,}\\\ \text{96-d, 3-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 8x8,}\\\ \text{96-d, 3-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 8x8,}\\\ \text{128-d, 4-h}\end{bmatrix}$ x 2
Stage 2 | C31-BN-ReLU, 128-d
$\begin{bmatrix}\text{sp. 4x4,}\\\ \text{128-d, 4-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 4x4,}\\\ \text{192-d, 6-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 4x4,}\\\ \text{192-d, 6-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 4x4,}\\\ \text{256-d, 8-h}\end{bmatrix}$ x 2
Stage 3 | C31-BN-ReLU, 256-d
$\begin{bmatrix}\text{sp. 2x2,}\\\ \text{320-d, 10-h}\end{bmatrix}$ x 5 |  $\begin{bmatrix}\text{sp. 2x2,}\\\ \text{384-d, 12-h}\end{bmatrix}$ x 5 |  $\begin{bmatrix}\text{sp. 2x2,}\\\ \text{384-d, 12-h}\end{bmatrix}$ x 18 |  $\begin{bmatrix}\text{sp. 2x2,}\\\ \text{512-d, 16-h}\end{bmatrix}$ x 18
Stage 4 w. cls token | C31-BN-ReLU, 512-d
$\begin{bmatrix}\text{sp. 1x1,}\\\ \text{512-d, 16-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 1x1,}\\\ \text{768-d, 24-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 1x1,}\\\ \text{768-d, 24-h}\end{bmatrix}$ x 2 |  $\begin{bmatrix}\text{sp. 1x1,}\\\ \text{1024-d, 32-h}\end{bmatrix}$ x 2
Table 2: Comparison of different linear/efficient transformer variants on ImageNet (Deng et al., 2009), based on our multi-stage Tiny configuration (see Table 1). The memory usage is measured with the batch size of 1024 which is our standard training setting. Transformer is tested at a batch size of 256, which is the maximal number possible with the GPU resource at our disposal. Throughput is in format as $\text{Train throughput}\ /\ \text{inference throughput}$.  Methods | Memory | Params | FLOPs | Throughput (img/s) | Top-1 %
---|---|---|---|---|---
Transformer (Vaswani et al., 2017) | 19.0GB$\dagger$ | 13M | 3.9G | 1073 / 3240 | 80.0
Linformer (Wang et al., 2020) | 11.7GB | 13M | 1.9G | 2767 / 3779 | 78.2
Performer (Choromanski et al., 2021) | 15.0GB | 13M | 2.2G | 2037 / 3657 | 76.1
Nyströmformer (Xiong et al., 2021) | 17.2GB | 13M | 2.0G | 1891 / 3518 | 80.1
SOFT | 15.8GB | 13M | 1.9G | 1730 / 3436 | 80.9
SOFT++ | 15.8GB | 13M | 1.9G | 1730 / 3436 | 80.9
Table 3: Evaluation results on ILSVRC-2012 ImageNet-1K (Deng et al., 2009) $\tt validation$ set. We report the results using the input size of 224x224 pixels center cropped from resized images with 256x256 pixels. $\tt M.S.Out.$ stands for whether the model is designed for multi-scale output. $\dagger$: Corrected FLOPs by taking into account the cost of attention matrix multiplication overlooked in the origin paper.  Model | Style | Resolution | M.S. Out.? | Params | FLOPs | Top-1 %.
---|---|---|---|---|---|---
ResNet-18 (He et al., 2016) | Convolution | $224^{2}$ | ✓ | 11M | 1.9G | 69.8
PVT-Tiny (Wang et al., 2021) | Transformers | $224^{2}$ | ✓ | 13M | 1.9G† | 75.1
ConViT-Ti+ (d’Ascoli et al., 2021) | Transformers | $224^{2}$ | ✗ | 10M | 2.0G | 76.7
Coat-Lite Mini (Xu et al., 2021) | Transformers | $224^{2}$ | ✓ | 11M | 2.0G | 78.9
LambdaNets-50 (Bello, 2021) | Transformers | $224^{2}$ | ✓ | 16M | - | 78.9
ViP-Ti (Sun et al., 2021) | Transformers | $224^{2}$ | ✓ | 13M | 1.7G | 79.0
SOFT-Tiny | SOFT | $224^{2}$ | ✓ | 13M | 1.9G | 80.9
SOFT++-Tiny | SOFT++ | $224^{2}$ | ✓ | 13M | 1.9G | 80.9
ResNet-50 (He et al., 2016) | Convolution | $224^{2}$ | ✓ | 25M | 4.1G | 78.5
PVT-Small (Wang et al., 2021) | Transformer | $224^{2}$ | ✓ | 24M | 4.0G† | 79.8
Swin-T (Liu et al., 2021) | Transformer | $224^{2}$ | ✓ | 29M | 4.5G | 81.3
Twins-SVT-S (Chu et al., 2021) | Hybrid | $224^{2}$ | ✓ | 24M | 3.7G | 81.7
CoAtNet-0 (Dai et al., 2021) | Hybrid | $224^{2}$ | ✓ | 25M | 4.2G | 81.6
DeiT III-S (Touvron et al., 2022) | Transformer | $224^{2}$ | ✗ | 22M | 4.6G | 81.4
SwinV2-T (Liu et al., 2022a) | Transformer | $256^{2}$ | ✓ | 29M | 4.5G | 81.7
ConvNext-T (Liu et al., 2022b) | Hybrid | $224^{2}$ | ✓ | 29M | 4.5G | 82.1
SOFT-Small | SOFT | $224^{2}$ | ✓ | 27M | 4.5G | 82.5
SOFT++-Small | SOFT++ | $224^{2}$ | ✓ | 27M | 4.5G | 82.6
ResNet-101 (He et al., 2016) | Convolution | $224^{2}$ | ✓ | 44M | 7.9G | 79.8
PVT-Medium (Wang et al., 2021) | Transformer | $224^{2}$ | ✓ | 44M | 7.0G† | 81.2
ViT-Small/16 (Dosovitskiy et al., 2021) | Transformer | $224^{2}$ | ✗ | 48M | 9.9G | 80.8
Swin-S (Liu et al., 2021) | Transformer | $224^{2}$ | ✓ | 50M | 8.7G | 83.0
ConvNext-S (Liu et al., 2022b) | Hybrid | $224^{2}$ | ✓ | 50M | 8.7G | 83.1
CoAtNet-1 (Dai et al., 2021) | Transformer | $224^{2}$ | ✓ | 42M | 8.4G | 83.3
SwinV2-S (Liu et al., 2022a) | Transformer | $256^{2}$ | ✓ | 50M | 8.7G | 83.6
SOFT-Medium | SOFT | $224^{2}$ | ✓ | 48M | 8.7G | 83.2
SOFT++-Medium | SOFT++ | $224^{2}$ | ✓ | 48M | 8.7G | 83.7
CaiT-S36(Touvron et al., 2021b) | Transformer | $224^{2}$ | ✓ | 88M | 13.9G | 83.3
Swin-B(Liu et al., 2021) | Transformer | $224^{2}$ | ✓ | 88M | 15.4G | 83.3
Twins-SVT-L (Chu et al., 2021) | Hybrid | $224^{2}$ | ✓ | 99M | 14.8G | 83.3
DeiT III-B (Touvron et al., 2022) | Transformer | $224^{2}$ | ✗ | 87M | 15.5G | 83.8
ConvNext-S (Liu et al., 2022b) | Hybrid | $224^{2}$ | ✓ | 89M | 15.4G | 83.8
CoAtNet-2 (Dai et al., 2021) | Hybrid | $224^{2}$ | ✓ | 75M | 15.7G | 84.1
SwinV2-B (Liu et al., 2022a) | Transformer | $256^{2}$ | ✓ | 88M | 15.4G | 84.1
SOFT-Large | SOFT | $224^{2}$ | ✓ | 85M | 15.4G | 83.6
SOFT++-Large | SOFT++ | $224^{2}$ | ✓ | 85M | 15.4G | 84.1
Figure 3: A comparison of Top-1 classification accuracy on the ImageNet validation set (Deng et al., 2009) with respect to inference throughput for various models. Our comparison includes CNN models such as ConvNext (Liu et al., 2022b), as well as Transformer models like Swin (Liu et al., 2021). In this comparison, models positioned closer to the top-right indicate superior performance, balancing both accuracy and throughput effectively. Inference throughput is measured on a V100 GPU, following (Liu et al., 2021, 2022b). Table 4: Object detection (RetinaNet (Lin et al., 2017)) results on COCO (Lin et al., 2014) val2007. We report the results by training 12 epoch ($1\times$) schedule. #P stands for parameter size. APb represents bounding box AP. Note, SOFT cannot converge in training.  Backbone | RetinaNet 1$\times$
---|---
#P | AP | AP50 | AP75 | APS | APM | APL
ResNet18 (He et al., 2016) | 21M | 31.8 | 49.6 | 33.6 | 16.3 | 34.3 | 43.2
PVT-Tiny (Wang et al., 2021) | 23M | 36.7 | 56.9 | 38.9 | 22.6 | 38.8 | 50.0
PVTv2-B1 (Wang et al., 2022) | 23M | 41.2 | 61.9 | 43.9 | 25.4 | 44.5 | 54.3
SOFT++-Tiny | 23M | 41.9 | 62.7 | 44.7 | 27.8 | 45.4 | 55.6
ResNet50 (He et al., 2016) | 38M | 36.3 | 55.3 | 38.6 | 19.3 | 40.0 | 48.8
PVT-Small (Wang et al., 2021) | 34M | 40.4 | 61.3 | 43.0 | 25.0 | 42.9 | 55.7
ViL-S (Zhang et al., 2021) | 35M | 41.6 | 62.5 | 44.1 | 24.9 | 44.6 | 56.2
Swin-T(Liu et al., 2021) | 38M | 41.7 | 61.2 | 43.2 | 26.0 | 44.3 | 54.5
SOFT++-Small | 38M | 43.7 | 64.9 | 46.8 | 28.7 | 47.4 | 57.6
ResNet101 (He et al., 2016) | 57M | 38.5 | 57.8 | 41.2 | 21.4 | 42.6 | 51.1
PVT-Medium (Wang et al., 2021) | 54M | 41.9 | 63.1 | 44.3 | 25.0 | 44.9 | 57.6
ViL-M (Zhang et al., 2021) | 50M | 42.9 | 64.0 | 45.4 | 27.0 | 46.1 | 57.2
Swin-S(Liu et al., 2021) | 60M | 43.0 | 63.8 | 45.7 | 27.1 | 46.9 | 57.2
SOFT++-Medium | 59M | 44.3 | 64.7 | 47.4 | 29.0 | 48.2 | 59.9
ResNeXt101 (Xie et al., 2017) | 95M | 41.0 | 60.9 | 44.0 | 23.9 | 45.2 | 54.0
PVT-Large (Wang et al., 2021) | 71M | 42.6 | 63.7 | 45.4 | 25.8 | 46.0 | 58.4
Swin-B (Liu et al., 2021) | 98M | 44.7 | 65.9 | 49.2 | - | - | -
SOFT++-Large | 98M | 47.0 | 67.8 | 50.4 | 30.2 | 50.9 | 62.0
Table 5: Instance segmentation (Mask R-CNN (He et al., 2017)) results on COCO (Lin et al., 2014) val2007. We report the results by training 12 epoch ($1\times$) schedule. #P stands for parameter size. APm represents mask AP respectively. Note, SOFT cannot converge in training. Backbone | Mask R-CNN 1$\times$
---|---
#P | APb | AP${}^{b}_{50}$ | AP${}^{b}_{75}$ | APm | AP${}^{m}_{50}$ | AP${}^{m}_{75}$
ResNet18 (He et al., 2016) | 31M | 34.0 | 54.0 | 36.7 | 31.2 | 51.0 | 32.7
PVT-Tiny (Wang et al., 2021) | 33M | 36.7 | 59.2 | 39.3 | 35.1 | 56.7 | 37.3
PVTv2-B1 (Wang et al., 2022) | 33M | 41.8 | 64.3 | 45.9 | 38.8 | 61.2 | 41.6
SOFT++-Tiny | 32M | 41.2 | 63.7 | 44.7 | 38.2 | 61.0 | 41.0
ResNet50 (He et al., 2016) | 44M | 38.0 | 58.6 | 41.4 | 34.4 | 55.1 | 36.7
PVT-Small (Wang et al., 2021) | 44M | 40.4 | 62.9 | 43.8 | 37.8 | 60.1 | 40.3
ViL-S (Zhang et al., 2021) | 45M | 41.8 | 64.1 | 45.1 | 38.5 | 61.1 | 41.4
Swin-T(Liu et al., 2021) | 48M | 42.7 | 65.2 | 46.8 | 39.3 | 62.2 | 42.2
SOFT++-Small | 48M | 43.8 | 66.0 | 47.5 | 40.1 | 63.0 | 43.0
ResNet101 (He et al., 2016) | 63M | 40.4 | 61.1 | 44.2 | 36.4 | 57.7 | 38.8
PVT-Medium (Wang et al., 2021) | 64M | 42.0 | 64.4 | 45.6 | 39.0 | 61.6 | 42.1
ViL-M (Zhang et al., 2021) | 60M | 43.4 | 65.9 | 47.0 | 39.7 | 62.8 | 42.1
Swin-S(Liu et al., 2021) | 69M | 45.6 | 67.4 | 50.0 | 41.2 | 64.5 | 44.3
SOFT++-Medium | 69M | 46.6 | 67.8 | 51.2 | 42.0 | 64.8 | 45.2
ResNeXt101 (Xie et al., 2017) | 101M | 42.8 | 63.8 | 47.3 | 38.4 | 60.6 | 41.3
PVT-Large (Wang et al., 2021) | 81M | 42.9 | 65.0 | 46.6 | 39.5 | 61.9 | 42.5
Swin-B (Liu et al., 2021) | 107M | 45.5 | - | - | 41.3 | - | -
SOFT++-Large | 106M | 47.0 | 68.3 | 51.7 | 42.2 | 65.2 | 45.4

##  4 Experiments

###  4.1 Image classification

Dataset: We evaluate the proposed SOFT and SOFT++ on the ILSVRC-2012
ImageNet-1K dataset (Deng et al., 2009) with 1.28M training images and 50K
validation images from 1,000 classes. Following the common practice, we train
a model on the training set and evaluate on the validation set.

Metrics: For model performance, the top-1 accuracy on a single crop is
reported. To assess the cost-effectiveness, we also report the model size and
floating point operations (i.e., FLOPs).

Implementation details: We use the code base (Wightman, 2019) with the default
setting to train and test all the models. Specifically, we use weight decay of
0.05 and 10 epochs of linear warm-up. We conduct 300 epochs training with an
$\tt AdamW$ optimizer and decreasing learning rate with the cosine annealing
schedule. During training, random flipping, mixup (Zhang et al., 2018) and
cutmix (Yun et al., 2019) are adopted for data augmentation. Label smoothing
(Szegedy et al., 2016) is used for loss calculation. All our variants are
trained with a batch size of 1024 on 32G NVIDIA V100 GPUs. We also implement
our method using the Mindspore (Mindspore, 2020).

Comparison with existing linear Transformers: We compare our method with three
existing linear Transformer models: Linformer (Wang et al., 2020), Performer
(Choromanski et al., 2021), Nyströmformer (Xiong et al., 2021) in terms of
model complexity and accuracy.

Two experimental settings are adopted. Under the first setting, for all
methods we use the same $\tt Tiny$ (Table 1) architecture for a fair
comparison. That is, we replace the core self-attention block in SOFT/SOFT++
with each baseline’s own attention block with the rest of the architecture
unchanged. Note that the spatial reduction module of (Wang et al., 2021) is a
special case of Linformer (Wang et al., 2020). We set the reduction ratio to
be identical to ours. With the same uniform sampling idea, we replace the 1D
window averaging of Nyströmformer (Xiong et al., 2021) (for NLP tasks) with 2D
average pooling (for images). The downsampling ratio remains identical to
ours. It is also worth mentioning that there is no official code released for
Reformer (Kitaev et al., 2020) and the local Sensitive Hash (LSH) module has
strict requirements on the length of input tokens. We thus do not include this
method in our comparison.

From Table 2 we can make the following observations: (i) Linear Transformer
methods substantially reduce the memory and FLOPs while maintain similar
parameter size comparing to the Transformer on the $\tt Tiny$ architecture;
(ii) Our approach SOFT/SOFT++ achieve the best classification accuracy among
all the linearization methods. (iii) Our inference speed is on-par with other
compared linear Transformers and our training speed is slightly slower than
Nystromformer and both are slower than Performer and Linformer. Note that the
slow training speed of our model is mostly due to the Newton-Raphson iteration
which can only be applied sequentially for ensuring the accuracy of Moore-
Penrose inverse. In summary, due to the on-par inference speed we consider the
training cost increase is a price worth paying for our superior accuracy.

Under the second setting, we focus on the memory efficiency of SOFT against
the baselines. Here we follow the ViT (Dosovitskiy et al., 2021) network
structure, stacking 12 attention layers with hidden dimension $d=384$, heads
$h=12$, bottleneck token sequence length $m=49$. Different attention blocks
from the three linearized Transformer variants, Linformer (Wang et al., 2020),
Performer (Choromanski et al., 2021), and Nyströmformer (Xiong et al., 2021)
are studied. For each Transformer variant, we adjust its token sequence length
$n$ in a linear increment. Specifically, we use a token sequence length of
$784\times p$ where $p=1,2,3,4,5,6,7,8$ and set batch size 1 to verify whether
the memory consumption increases “quadratically” or “linearly”. Figure 1b
shows all compared transformer variants including our SOFT indeed have a
linear memory usage complexity. This is in contrast with the standard
Transformer which cannot cope with long token sequences with a quadratic
complexity.

Comparison with state-of-the-art CNNs and ViTs:  We compare with state-of-the-
art alternatives and report the top-1 accuracy on the ImageNet-1K validation
set. FLOPs are calculated at batch size 1. From Figure 1a and Table 3, the
following observations are made: (i) Overall, ViT and its variants yield
better classification accuracy over CNNs. (ii) We achieve the best performance
among the recent pure vision Transformer based methods including ViT
(Dosovitskiy et al., 2021) and DeiT (Touvron et al., 2021a), as well as the
state-of-the-art CNN RegNet (Radosavovic et al., 2020). (iii) Our SOFT and
SOFT++ outperform the most similar (in architecture configuration) Transformer
counterparts PVT (Wang et al., 2021) at all variants. Since the attention
module is the main difference, this validates directly the effectiveness of
our model. (iv) We can also beat the latest ViT variants Twins (Chu et al.,
2021) which is designed to address the efficiency limitation of ViT. We have
done so with less parameters and fewer float point computation.

Figure 4: Comparison of attention heatmaps for a selected query patch
(indicated by a cross ”+”) against all patches in an image. Heatmaps are
derived from the first head’s corresponding row in the attention maps, as
calculated by Equation 17. These heatmaps are normalized to a 0-1 scale, with
warmer colors indicating higher relevance. The model variants compared are:
(a) Transformer (Vaswani et al., 2017), (b) Performer (Choromanski et al.,
2021), (c) Nystromformer (Xiong et al., 2021), and (d) Our SOFT approach. For
additional examples, refer to Appendix E.

To gain insights into how attention is learned using our SOFT and the
alternatives, Figure 4 shows the attention masks. For each model, we visualize
the first two attention heads. It is evident that SOFT exhibits robustness and
versatility in capturing local and long distance relations among pixels. We
note that, although SOFT is trained for object categorization (ImageNet (Deng
et al., 2009)), it seems to be able to learn both semantic concepts shared
across instances in the same category and instance specific features. For
instance, in the bottom-right example of a bird class, one attention head
focuses on the black bird only, while the other attends to both birds in the
image. More examples are given in Appendix E.

Throughputs comparison:  Figure 3 illustrates that our model achieves a
superior balance between speed and performance compared to the Swin
Transformer (Liu et al., 2021), and maintains a comparable balance with the
current CNN state-of-the-art, ConvNext (Liu et al., 2022b). As detailed in
Table 2, we employ a series-computed Newton iteration method to ensure
numerical accuracy. While this approach slightly reduces speed, it preserves
the accuracy of our model.

Comparison between SOFT and SOFT++ As shown in Table 3, SOFT++ performs
equally well as SOFT for small models, and outperforms the preliminary version
for larger models. Further, Table 6 validates the superiority of SOFT++ over
SOFT in processing long token sequences in model deployment.

Table 6:  The evaluation results on the validation set of ImageNet-1K (Deng et al., 2009) at various input sizes. The models were trained on the size of $224\times 224$. The down-sampling ratio of the patch embedding remains at 4.  Test input size | $224\times 224$ | $384\times 384$ | $512\times 512$
---|---|---|---
Token sequence length | 3136 | 9216 | 16384
SOFT | 82.5 | 33.7 | 0.14
SOFT++ | 82.6 | 82.6 | 80.2

###  4.2 Object detection on COCO

Dataset:  We evaluate the object detection performance of our SOFT++ on the
COCO benchmark (Lin et al., 2014) including the train2017 (118k images) and
val2017 (5k images) sets.

Implementation details:  We consider two representative detectors: RetinaNet
(Lin et al., 2017) and Mask R-CNN (He et al., 2017). We use AdamW optimizer
with base learning rate of $1\times 10^{-4}$ and weight decay of $0.01$. We
train all the model with batch size 16 on 8 V100 GPUs. Following the practices
of MMDetection (Chen et al., 2019), we adopt the 1$\times$ and 3$\times$
training schedules. In the training stage without multi-scale, images are
resized to make the shorter sides at 800 pixels and the longer sides no
exceeding 1333 pixels. In the training stage with multi-scale, the shorter
sides of images are randomly resized to between 480 to 800.

Comparison with state-of-the-art CNNs and ViTs:  We compare with the state-of-
the-art alternatives on RetinaNet (Lin et al., 2017) in Table 4 and Mask R-CNN
(He et al., 2017) in Table 5. Our SOFT++ outperforms the CNN ResNet (He et
al., 2016) and Transformer counterparts PVT (Wang et al., 2021) across all
complexity groups in both object detection and instance segmentation,
validating the superior trade-off between model complexity and performance by
our design.

###  4.3 Semantic segmentation

Dataset:  We evaluate the semantic segmentation performance of our SOFT++ on
ADE20K (Zhou et al., 2019) and Cityscapes (Cordts et al., 2016).

Implementation details:  UperNet (Xiao et al., 2018) is used as the framework.
AdamW optimizer with $6\times 10^{-4}$ learning rate is applied to train
ADE20K for 160k iterations and Cityscapes for 40k iterations. Images are
cropped as $512\times 512$ for ADE20K and $768\times 768$ for Cityscapes
during training. Multi-scale training and test time augmentation are not used.

Comparison with state-of-the-art CNNs and ViTs:  Table 7 shows that SOFT++
surpasses ResNet clearly, and the newest designed vision transformer Swin (Liu
et al., 2021) and CovNext (Liu et al., 2022b) slightly under the similar
parameters and float point computation cost.

Table 7: Semantic segmentation performance using our HLG Transformer with UperNet (Xiao et al., 2018) and SETR-PUP on Cityscapes validation set. Single-scale inference and 40k training schedules are used. Note, SOFT cannot converge in training.  Method | Backbone | #P | Cityscapes | ADE20K
---|---|---|---|---
FLOPs | mIoU | FLOPs | mIoU
FCN (Long et al., 2015) | ResNet-101 | 68M | 619G | 76.6 | 276G | 39.9
PSPNet (Zhao et al., 2017) | ResNet-101 | 68M | 576G | 78.5 | 256G | 44.4
DLabV3+ (Chen et al., 2018) | ResNet-101 | 62M | 571G | 79.3 | 262G | 46.9
CCNet (Huang et al., 2019) | ResNet-101 | 68M | 625G | 80.2 | 278G | 43.7
SETR (Zheng et al., 2021) | ViT-Large | 318M | 1340G | 79.2 | 602G | 48.5
OCRNet (Yuan et al., 2020) | HRNet-W48 | 70M | 972G | 81.1 | 164G | 43.2
UperNet (Xiao et al., 2018) | Swin-T | 60M | - | - | 236G | 44.5
UperNet (Xiao et al., 2018) | Swin-S | 81M | - | - | 259G | 47.6
UperNet (Xiao et al., 2018) | Swin-B | 121M | - | - | 297G | 48.1
UperNet (Xiao et al., 2018) | ConvNext | 60M | - | - | 235G | 46.0
UperNet (Xiao et al., 2018) | ConvNext | 82M | - | - | 256G | 48.7
UperNet (Xiao et al., 2018) | ConvNext | 122M | - | - | 293G | 49.1
UperNet (Xiao et al., 2018) | SOFT++-Small | 60M | 545G | 81.2 | 237G | 46.5
UperNet (Xiao et al., 2018) | SOFT++-Med | 81M | 607G | 82.0 | 260G | 48.9
UperNet (Xiao et al., 2018) | SOFT++-Large | 121M | 899G | 82.6 | 301G | 49.2

###  4.4 Ablation studies

Pyramidal architecture:  Unlike the earlier non-pyramidal vision Transformers
(e.g., ViT (Dosovitskiy et al., 2021)), most recent pyramidal (multi-scale)
Transformers (e.g., PVT (Wang et al., 2021)) use convolution layers to reduce
the spatial resolution (i.e., token sequence length) between stages. In this
study, we ablate SOFT++ with a pyramidal architecture (our default SOFT++-$\tt
Small$), SOFT++ w/o a pyramidal architecture and DeiT-S (Touvron et al.,
2021a) (no pyramidal architecture either). We replace the Transformer layer
with a SOFT++ layer to get SOFT++ w/o a pyramidal architecture. Note all three
variants have similar parameters and FLOPs. Table 11a shows that the conv-
based pyramidal architecture is clearly superior to a non-pyramidal design,
and our non-pyramidal counterpart is even slightly better than DeiT-S (Touvron
et al., 2021a) whilst enjoying linear complexity.

Table 8: Ablations on assessing the Impact of Symmetric Kernels on ImageNet-1K (Deng et al., 2009) image classification. As baselines for our study, we employ two of the most common backbones: ViT-small/16 (Dosovitskiy et al., 2021) and Swin-T (Liu et al., 2021).  Model | Asymmetric ($Q\neq K$) | Symmetric ($Q=K$)
---|---|---
ViT-Small/16 (Dosovitskiy et al., 2021) | 80.8 | 80.5
Swin-T (Liu et al., 2021) | 81.3 | 80.9
Table 9: Ablations on bottleneck token sequence length. Bottleneck | Memory | FLOPs | Top-1 %
---|---|---|---
36 | 15.1GB | 1.9G | 80.5
49 | 15.8GB | 1.9G | 80.9
64 | 16.9GB | 2.0G | 80.9
81 | 18.5GB | 2.1G | 80.4
Table 10: Ablations on sampling methods. Sampling methods | Params | FLOPs | Top-1 %
---|---|---|---
Convolution | 13.07M | 2.0G | 80.9
Random sampling | 12.96M | 1.9G | 80.8
Biased sampling | 12.96M | 1.9G | 79.9
Average pooling | 12.96M | 1.9G | 80.8

Symmetric kernels:  Tsai et al. (Tsai et al., 2019) conducted a comprehensive
analysis of Transformers, revealing minimal performance disparities between
asymmetric and symmetric kernels in NLP tasks, such as neural machine
translation (NMT) and sequence prediction (SP). In contrast, our research,
detailed in Table 8, extends this investigation to the domain of computer
vision, specifically focusing on image classification tasks. Our findings
indicate that while symmetric kernels do have a marginal detrimental effect in
computer vision tasks, this impact is relatively limited. Therefore,
maintaining a symmetric kernel represents a balanced approach for enabling
effective decomposition in our model.

Bottleneck token sequence length:

(a)  (b)  (c)

Figure 5: Convergence analysis for the approximation of Moore-Penrose inverse
on ImageNet, COCO and ADE20K separately. SOFT-Tiny is used. We measure
$\|AA_{k}A-A\|_{p}/\|A\|_{2}$ for 100 input images on each dataset. The solid
line shows the average convergence metric, while the shallow area indicates
the upper bound and lower bound.

In this study, we examine how the bottleneck token sequence length $m$,
sampled from $n$ tokens, influences the model’s performance. We change the
bottleneck token sequence length in all stages to $36,49,64,81$. Table 9 shows
that longer bottleneck token would increase the memory cost and the
computational overhead. $m=49$ seems to give the best trade-off between the
performance and computational overhead. The memory usage is measured with the
batch size of 128.

Table 11: (a) Ablations on pyramidal architecture. (b) Ablations on overlapped convolution. Methods | Pyramidal? | Top-1 % |  | Methods | Overlapped? | Top-1 %
---|---|---|---|---|---|---
DeiT-S (Touvron et al., 2021a) | ✗ | 79.8 |  | PVT (Wang et al., 2021) | ✗ | 75.1
SOFT++ | ✗ | 80.1 |  | SOFT++ | ✗ | 78.4
SOFT++ | ✓ | 82.4 |  | SOFT++ | ✓ | 80.9

Token sampling:  The sampling function in SOFT++ can assume different forms.
Convolution: The sequence $Q\in\mathbb{R}^{n\times d_{e}}$ is first reshaped
to a feature map $\mathbb{R}^{H\times W\times d_{e}}$. $r\times r$ convolution
kernel with stride of $r$ is applied for downsampling, where $r=\sqrt{sp}$.
The output channel size is also kept and no bias is used. At last, the feature
map is reshaped back to the sequence. Average pooling: using a $r\times r$
kernel and $r$ stride, where $r=\sqrt{sp}$. Random sampling: $m$ tokens are
randomly picked from $n$ tokens. Biased sampling:  We pick $m$ tokens with a
biased policy. Here, the first $m$ tokens are picked. Table 10 shows that both
convolution yields the best performance. Biased sampling can miss the most
salient samples, and there is no guarantee that random sampling can keep the
uniformity of the chosen samples. This result thus justifies the choice of
using convolution in SOFT++.

Overlapped convolution:  We ablate SOFT++ with overlapped convolution (our
default choice, same as many recent works) and with non-overlapped convolution
in our $\tt Tiny$ configuration. Table 11b shows that overlapped convolution
is a better choice. Our non-overlapped convolution variant still outperforms
PVT (Wang et al., 2021) with the same non-overlapped convolution by a clear
margin.

Newton-Raphson’s convergence: We study how many iterations the Newton-Raphson
method needs to converge when computing the Moore-Penrose inverse
$A^{\dagger}$. We use $\|AA_{k}A-A\|_{p}/\|A\|_{p}$ with $p=2$ (see
Proposition 2) as the convergence metric to quantify the difference between
$A_{k}$ and $A^{\dagger}$. Figure 5 shows that our approximation converges
within 20 iterations across all stages.

Table 12: Comparison of different linear/efficient Transformer variants on Long Range Arena (Tay et al., 2020), based on its official configuration. Our SOFT surpasses previous efficient methods on three tasks.  Methods | Listops(2K) | Text(4K) | Retrieval(4K) | Image(1K) | Avg. %
---|---|---|---|---|---
Transformer (Vaswani et al., 2017) | 37.10 | 65.02 | 79.35 | 38.20 | 54.92
Reformer (Kitaev et al., 2020) | 19.05 | 64.88 | 78.64 | 43.29 | 51.47
Linformer (Wang et al., 2020) | 37.25 | 55.91 | 79.37 | 37.84 | 52.59
Performer (Choromanski et al., 2021) | 18.80 | 63.81 | 78.62 | 37.07 | 49.58
Nyströmformer (Xiong et al., 2021) | 37.15 | 65.52 | 79.56 | 41.58 | 55.95
SOFT | 37.40 | 63.49 | 81.77 | 46.91 | 57.39

Effect of attention normalization: Our attention normalization enables the
model to perform more challenging vision tasks such as object detection and
segmentation. This is because the leading eigenvalue of Gaussian kernel self-
attention is conditioned on the token sequence length (Proposition 4), as
shown in Figure 6a. And this effect can be clearly mitigated by our
normalization scheme, as indicated in Figure 6b.

(a)  (b)

Figure 6: (a) Eigenvalues of a specific bottleneck matrix
$A\in\mathbb{R}^{1024\times 1024}$. The auxiliary (red dash) line connects all
the eigenvalues. (b) Comparing the spectral norm of self-attention matrix
$A^{\dagger}$ (without normalization, blue) and $D^{-1/2}A^{\dagger}D^{-1/2}$
(with normalization, orange) under a variety of bottleneck token sequence
lengths. “+” represents individual samples; The lines connect the average norm
at each token sequence length.

###  4.5 Additional experiments on NLP tasks

We compare our method with previous linear counterparts on four tasks of the
Long Range Arena (LRA) (Tay et al., 2020) benchmark, including Listops (Nangia
and Bowman, 2018), byte-level IMDb reviews text classification (Maas et al.,
2011), byte-level document retrieval (Radev et al., 2013), and image
classification on sequences of pixels (Krizhevsky, 2009).

Implementations. We use the Pytorch version of LRA (Tay et al., 2020)
benchmark, provided by (Xiong et al., 2021). For the evaluation protocol, we
strictly follow (Tay et al., 2020; Xiong et al., 2021). We omit the
Pathfinder(1K) task as we cannot replicate the result of Nyströmformer (Xiong
et al., 2021). For our SOFT, we simply use the average pooling with window
size 4, stride 4 to sample the bottlenecks. We follow the configurations of
(Xiong et al., 2021) with 2 layers, 64 and 128 hidden dimension respectively,
and 2 attention heads. Table 12 shows that our SOFT outperforms both the
standard and alternative efficient Transformers on three out of four tasks, as
well as the average performance.

##  5 Conclusions

In this work, we have introduced a novel softmax-free self-attention (SOFT)
mechanism for linearizing Transformer’s complexity in space and time. Unlike
existing linear Transformers that aim to approximate the conventional softmax
based self-attention, SOFT employs a Gaussian kernel based attention which
eliminates the need for softmax normalization. This design enables a full
self-attention matrix to be approximated via low-rank matrix decomposition.
The robustness of this proposed approximation is achieved by calculating its
Moore-Penrose inverse using a Newton-Raphson method and an efficient symmetric
attention normalization. Extensive experiments show that SOFT yields superior
trade-off in accuracy and complexity on a variety of vision and language
tasks.

Acknowledgements This work was supported in part by STI2030-Major Projects
(Grant No. 2021ZD0200204), National Natural Science Foundation of China (Grant
No. 62106050 and 62376060), Natural Science Foundation of Shanghai (Grant No.
22ZR1407500) and Lingang Laboratory (Grant No. LG-QS-202202-07).

Data availability statement The datasets generated during and/or analysed
during the current study are available in the Imagenet (Deng et al., 2009)
(https://www.image-net.org/), COCO (Lin et al., 2014)
(https://cocodataset.org), ADE20K (Zhou et al., 2019)
(https://groups.csail.mit.edu/vision/datasets/ADE20K/), Cityscapes (Cordts et
al., 2016) (https://www.cityscapes-dataset.com), Long Range Arena (Tay et al.,
2020) (https://github.com/google-research/long-range-arena) repositories.

##  Appendix A Nyström method

Nyström method Williams and Seeger (2000) aims to calculate a low-rank
approximation for a Gram matrix. For Transformers, the self-attention matrix
can be viewed as a Gram matrix $S$ with a Gaussian kernel $k$ applied to the
query $Q$, with each element $S_{ij}$ expressed as:

$S_{ij}=k\big{(}Q_{i,:},Q_{j,:}\big{)}=\text{exp}(-\frac{\|Q_{i,:}-Q_{j,:}\|_{2}^{2}}{2\sqrt{d}}),$
(18)

$k(x,y)$ means operating Gaussian kernel $k$ to $(x,y)$, which can be written
in the feature space as:

$k(x,y)=\sum\limits_{i=1}^{n}\lambda_{i}\phi_{i}(x)\phi_{i}(y),$ (19)

$n$ is the dimension of a feature space, $\lambda_{i}$ denotes the eigenvalue
and $\phi_{i}$ denotes the eigenfunction of kernel $k$. According to the
eigenfunction’s definition, we can get:

$\int k(y,x)\phi_{i}(x)p(x)dx=\lambda_{i}\phi_{i}(y),$ (20)

where $p(x)$ is the probability distribution of $x$. And {$\phi_{i}(x)$} are
$p$-orthogonal:

$\int\phi_{i}(x)\phi_{j}(x)p(x)dx=\delta_{ij}.$ (21)

$\delta_{ij}$ is $0$ when $i\neq j$, $1$ when $i=j$. To get an approximation
of the eigenfunctions, we sample $\\{x_{1},x_{2},\cdots,x_{q}\\}$ from $p(x)$,
then:

$\frac{1}{q}\sum\limits_{t=1}^{q}k(y,x_{t})\phi_{i}(x_{t})\approx\lambda_{i}\phi_{i}(y),$
(22)
$\frac{1}{q}\sum\limits_{t=1}^{q}\phi_{i}(x_{t})\phi_{j}(x_{t})\approx\delta_{ij}.$
(23)

This inspires us to approximate the Gram matrix $S$. Let $S^{(m)}$ be a
submatrix of $S$, consisting of $m\times m$ elements from $S$. Gram matrix is
a symmetric positive semi-definite matrix, so it has a spectral decomposition:

$S^{(m)}U^{(m)}=U^{(m)}\Lambda^{(m)},$ (24)

where $U^{(m)}$ is column orthogonal and $\Lambda^{(m)}$ is a diagonal matrix
with the diagonal elements as the eigenvalues of $S^{(m)}$. Substituting the
$y$ to $x_{j}$ and applying the approximation above to $S$, we can get:

$\phi_{i}(x_{j})\approx\sqrt{m}U_{j,i}^{(m)},\quad\lambda_{i}\approx\frac{\lambda_{i}^{(m)}}{m},$
(25)
$\phi_{i}(y)\approx\frac{\sqrt{m}}{\lambda_{i}^{(m)}}\sum\limits_{t=1}^{m}k(y,x_{t})\phi_{i}(x_{t}),$
(26)

$\lambda_{i}$ is eigenvalue of $S$ and $\lambda_{i}^{(m)}$ is the eigenvalue
of $S^{(m)}$. Denote $\tilde{S}$ as the rank-$m$ approximation of $S$ and
$\tilde{U},\tilde{\Lambda}$ as the approximation for spectral decomposition of
$S$. Now we can get an approximation of $S$ with rank $m$:

$\tilde{S}=\tilde{U}\tilde{\Lambda}\tilde{U}^{T}=\sum\limits_{t=1}^{m}\tilde{\lambda_{t}}^{(n)}\tilde{u}_{t}^{(n)}(\tilde{u}_{t}^{(n)})^{T}.$
(27)

Similarly, we have:

$\phi_{i}(x_{j})\approx\sqrt{n}U_{j,i}(n),\quad\lambda_{i}\approx\frac{\tilde{\lambda_{i}}^{(n)}}{n}.$
(28)

Thus

$\tilde{\lambda_{i}}^{(n)}\approx\frac{n\lambda_{i}^{(m)}}{m},$ (29)
$\tilde{u}_{t}^{(n)}\approx\sqrt{\frac{m}{n}}\frac{1}{\lambda_{t}^{(m)}}S_{n,m}u_{t}^{(m)}.$
(30)

Then we get an approximation of $S$: $\tilde{S}\approx
S_{n,m}S_{m,m}^{\dagger}S_{m,n}$. $S$ has a block representation below:

$S=\begin{bmatrix}S_{m,m}&S_{m,n-m}\\\ S_{n-m,m}&S_{n-m,n-m}\end{bmatrix}.$
(31)

##  Appendix B Newton method

Proof of Proposition 1 When $\alpha$ is sufficiently small,
$A_{k+1}=2A_{k}-A_{k}AA_{k}$, $A_{k}$ converges to $A^{\dagger}$.

###### Proof

$A$ is a symmetric positive semi-definite matrix and $A_{ij}\leq 1$, $\forall
1\leq i,j\leq n$, $A_{ii}=1,\quad 1\leq i\leq n$ in our case. $A_{0}$ is
chosen to be $\alpha A$, so the $A_{k}$ can be written as
$A_{k}=C_{k}A=AD_{k}$ for some matrix $C_{k},D_{k}$, leading to the fact that

$A^{\dagger}AA_{k}=A_{k},\quad A_{k}AA^{\dagger}=A_{k}.$ (32)

This is because $A_{k+1}=A_{k}(2I_{n}-AA_{k})=(2I_{n}-A_{k}A)A_{k}$ and
$A_{0}=\alpha A$. We make a difference between $A^{\dagger}$ and $A_{k+1}$:

$\displaystyle A^{\dagger}-A_{k+1}$
$\displaystyle=A^{\dagger}-2A_{k}+A_{k}AA_{k}$
$\displaystyle=A^{\dagger}-A_{k}AA^{\dagger}-A^{\dagger}AA_{k}+A_{k}AA_{k}$
$\displaystyle=(A^{\dagger}-A_{k})A(A^{\dagger}-A_{k}).$ (33)

We norm both sides of the equation above:

$\displaystyle\|A^{\dagger}-A_{k+1}\|$
$\displaystyle=\|(A^{\dagger}-A_{k})A(A^{\dagger}-A_{k})\|$
$\displaystyle\leq\|A^{\dagger}-A_{k}\|\|A(A^{\dagger}-A_{k})\|.$ (34)

And we left multiply $A$ on the both sides of (Eq 33), then norm the equation:

$\displaystyle\|AA^{\dagger}-AA_{k+1}\|$
$\displaystyle=\|A(A^{\dagger}-A_{k})A(A^{\dagger}-A_{k})\|$
$\displaystyle\leq\|AA^{\dagger}-AA_{k}\|^{2}.$ (35)

We choose $\alpha$ sufficiently small so that the initial value satisfy
$\|AA^{\dagger}-AA_{0}\|<1$. We set $\alpha=\frac{2}{\|A\|_{1}^{2}}$ to ensure
it is small enough Ben-Israel and Cohen (1966). Then the
$\|AA^{\dagger}-AA_{k}\|\rightarrow 0$, when $k\rightarrow\infty$. The
inequality (34) implies that $A_{k}\rightarrow A^{\dagger},\
k\rightarrow\infty$.

Proof of Equation 12

$\nabla_{x}\mathcal{L}=-Y^{\top}(\nabla_{Y}\mathcal{L})Y^{\top},$

###### Proof

Since $XY=I$, we take derivatives on the both sides

$\displaystyle YdX+XdY$ $\displaystyle=0$ $\displaystyle dY$
$\displaystyle=-YdXY$ (36)

For the loss function $\mathcal{L}$, we have

$d\mathcal{L}=\langle\nabla_{X}\mathcal{L},dX\rangle=\langle\nabla_{Y}\mathcal{L},dY\rangle$

Recall that the inner product $\langle A,B\rangle=\text{Tr}(A^{\top}B)$ and
$\text{Tr}(ABC)=\text{Tr}(CAB)=\text{Tr}(BCA)$, we have

$\displaystyle\langle\nabla_{X}\mathcal{L},dX\rangle$
$\displaystyle=\langle\nabla_{Y}\mathcal{L},dY\rangle$
$\displaystyle=\langle\nabla_{Y}\mathcal{L},-YdXY\rangle$
$\displaystyle=-\text{Tr}(\nabla_{Y}\mathcal{L}^{\top}YdXY)$
$\displaystyle=\langle-Y^{\top}\nabla_{Y}\mathcal{L}Y^{\top},dX\rangle$

Therefore, Eq (12) is proved.

##  Appendix C Attention normalization

Proof of Porpostion 5 Assume the bottleneck matrix of softmax-free attention ,
$A\in\mathbb{R}^{m\times m}$ is $k$-connected. If
$\lambda_{1}\geq\lambda_{2}\geq\cdots\lambda_{m}\geq 0$ are eigenvalues of
$A^{\dagger}$, then $\lambda_{1}=\mathcal{O}(m^{2})$ and
$\|A^{\dagger}\|_{2}=\mathcal{O}(m^{2})$.

To prove Proposition 5, we first introduce the following lemmas.

###### Lemma 1

If $F$ is a non-singular and symmetric matrix, then $\|F\|_{2}\leq\|F\|_{1}$.

###### Proof

According to the definition of 2-Norm, we have $F^{T}Fz=\mu^{2}z$ with
$\mu=\|F\|_{2}$. Therefore, $\mu^{2}\|z\|_{1}=\|F^{T}Fz\|_{2}\leq$
$\|F^{T}\|_{1}\|F\|_{1}\|z\|_{1}=\|F\|_{\infty}\|\|F_{1}\|z\|_{1}$. Since $A$
is non-singular and symmetric,

$\|F\|_{2}\leq\sqrt{\|F\|_{\infty}\|F\|_{1}}=\|F\|_{1}$ (37)

###### Lemma 2

If $F$ is a real matrix and $\|F\|_{p}<1$, then

$(I-F)^{-1}=\sum_{k=0}^{\infty}F^{k}$ (38)

with

$\|(I-F)^{-1}\|_{p}\leq\frac{1}{1-\|F\|_{p}}$ (39)

###### Proof

$\left(\sum_{k=0}^{N}F^{k}\right)(I-F)=I-F^{N+1}$

Since $\|F\|_{p}<1$, $\lim_{k\rightarrow\infty}F^{k}=0$. Thus,

$\left(\lim_{N\rightarrow\infty}\sum_{k=0}^{N}F^{k}\right)(I-F)=I$

Multiply the equation by $(I-F)^{-1}$ we obtain Eq (38). From that, we can
easily get

$\|(I-F)^{-1}\|_{p}\leq\sum_{k=0}^{\infty}\|F\|_{p}^{k}=\frac{1}{1-\|F\|_{p}}$

Now we begin to prove Proposition 5.

###### Proof

Since the bottleneck matrix softmax-free attention $\|A\|_{1}$ is not less
than 1, we normalize the matrix as $A_{n}=D^{-1}A$, where
$D=\text{diag}(A\mathbbm{1}_{m})$. From Lemma 2, it follows that

$\|A_{n}^{-1}\|_{1}\leq\frac{1}{1-\|1-A_{n}\|_{1}}$

Since we assume the bottleneck matrix is k-connected, and $k<<m$,

$\|I-A_{n}\|_{1}\leq\frac{(m-k)}{m}$

Then,

$\|A_{n}^{-1}\|_{1}\leq\frac{1}{1-\frac{(m-k)}{m}}=\frac{m}{k}=\mathcal{O}(m)$

Note that we $D=\mathcal{O}(m)$ by assumption that the bottleneck matrix is
k-connected, it follows that

$\|A^{-1}\|_{1}=\|A_{n}^{-1}D\|_{1}\leq\|A_{n}^{-1}\|_{1}\|D\|_{1}=\mathcal{O}(m^{2})$

##  Appendix D Non-linearized Gaussian kernel attention

Instead of directly calculating the Gaussian kernel weights, in our
formulation they are approximated. More specifically, the relation between any
two tokens is reconstructed via sampled bottleneck tokens. As the number $m$
(e.g., 49) of bottleneck tokens is much smaller than the entire token sequence
length, our attention matrix is of low-rank. This brings about two favorable
consequences: (I) The model now focuses the attentive learning on latent
salient information captured by the $m$ bottleneck tokens.  (II) The model
becomes more robust against the underlying token noise due to the auto-encoder
style reconstruction Geng et al. (2021).  This explains why a model with an
approximated gram matrix performs better than one with a directly computed
matrix. Further, we find that exact Gaussian kernel attention leads to
training difficulties. As Proposition 4 reveals, this is due to lacking
normalization that leads to explosion of the spectral norm especially in long
token sequence cases. Big spectral norm could jeopardize the training and tend
to collapse the model.

##  Appendix E Attention visualization

Figure 7 shows more visualization of the attention masks by various
Transformers Vaswani et al. (2017); Choromanski et al. (2021); Xiong et al.
(2021) and our SOFT. For each model, we show the output from the first two
attention heads (up and down row). It is noteworthy that SOFT exhibits better
semantic diversity of the multi-head mechanism than other methods. Moreover,
when we sample the patch at the boundary of multiple objects, SOFT is able to
more precisely capture all these objects.

Figure 7:  Comparison of attention heatmaps for a selected query patch
(indicated by a cross ”+”) against all patches in an image. Heatmaps are
derived from the first head’s corresponding row in the attention maps, as
calculated by Equation 17. These heatmaps are normalized to a 0-1 scale, with
warmer colors indicating higher relevance. The model variants compared are:
(a) Transformer (Vaswani et al., 2017), (b) Performer (Choromanski et al.,
2021), (c) Nystromformer (Xiong et al., 2021), and (d) Our SOFT approach.

## References

  * Ba et al. (2016) Ba JL, Kiros JR, Hinton GE (2016) Layer normalization. arXiv preprint arXiv:160706450
  * Bello (2021) Bello I (2021) Lambdanetworks: Modeling long-range interactions without attention. In: International Conference on Learning Representations
  * Ben-Israel and Cohen (1966) Ben-Israel A, Cohen D (1966) On iterative computation of generalized inverses and associated projections. SIAM Journal on Numerical Analysis 3(3):410–419
  * Brown et al. (2020) Brown T, Mann B, Ryder N, Subbiah M, Kaplan JD, Dhariwal P, Neelakantan A, Shyam P, Sastry G, Askell A, et al. (2020) Language models are few-shot learners. Advances in Neural Information Processing Systems 33:1877–1901
  * Chen et al. (2019) Chen K, Wang J, Pang J, Cao Y, Xiong Y, Li X, Sun S, Feng W, Liu Z, Xu J, et al. (2019) Mmdetection: Open mmlab detection toolbox and benchmark. arXiv preprint arXiv:190607155
  * Chen et al. (2018) Chen LC, Zhu Y, Papandreou G, Schroff F, Adam H (2018) Encoder-decoder with atrous separable convolution for semantic image segmentation. In: European Conference on Computer Vision, pp 801–818
  * Cheng et al. (2014) Cheng J, Grossman M, McKercher T (2014) Professional CUDA c programming
  * Choromanski et al. (2021) Choromanski K, Likhosherstov V, Dohan D, Song X, Gane A, Sarlos T, Hawkins P, Davis J, Mohiuddin A, Kaiser L, et al. (2021) Rethinking attention with performers. In: International Conference on Learning Representations
  * Chu et al. (2021) Chu X, Tian Z, Wang Y, Zhang B, Ren H, Wei X, Xia H, Shen C (2021) Twins: Revisiting the design of spatial attention in vision transformers. Advances in Neural Information Processing Systems 34:9355–9366
  * Cordts et al. (2016) Cordts M, Omran M, Ramos S, Rehfeld T, Enzweiler M, Benenson R, Franke U, Roth S, Schiele B (2016) The cityscapes dataset for semantic urban scene understanding. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 3213–3223
  * Dai et al. (2021) Dai Z, Liu H, Le QV, Tan M (2021) Coatnet: Marrying convolution and attention for all data sizes. In: Advances in Neural Information Processing Systems, vol 34, pp 3965–3977
  * Deng et al. (2009) Deng J, Dong W, Socher R, Li LJ, Li K, Fei-Fei L (2009) Imagenet: A large-scale hierarchical image database. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 248–255
  * Devlin et al. (2019) Devlin J, Chang MW, Lee K, Toutanova K (2019) Bert: Pre-training of deep bidirectional transformers for language understanding. In: Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT, pp 4171–4186
  * Dosovitskiy et al. (2021) Dosovitskiy A, Beyer L, Kolesnikov A, Weissenborn D, Zhai X, Unterthiner T, Dehghani M, Minderer M, Heigold G, Gelly S, et al. (2021) An image is worth 16x16 words: Transformers for image recognition at scale. In: International Conference on Learning Representations
  * d’Ascoli et al. (2021) d’Ascoli S, Touvron H, Leavitt ML, Morcos AS, Biroli G, Sagun L (2021) Convit: Improving vision transformers with soft convolutional inductive biases. In: International Conference on Machine Learning, PMLR, pp 2286–2296
  * Fasshauer (2011) Fasshauer GE (2011) Positive definite kernels: past, present and future. Dolomites Research Notes on Approximation 4:21–63
  * Geng et al. (2021) Geng Z, Guo MH, Chen H, Li X, Wei K, Lin Z (2021) Is attention better than matrix decomposition? In: International Conference on Learning Representations
  * Guo et al. (2022) Guo MH, Liu ZN, Mu TJ, Hu SM (2022) Beyond self-attention: External attention using two linear layers for visual tasks. IEEE Transactions on Pattern Analysis and Machine Intelligence 45(5):5436–5447
  * He et al. (2016) He K, Zhang X, Ren S, Sun J (2016) Deep residual learning for image recognition. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 770–778
  * He et al. (2017) He K, Gkioxari G, Dollár P, Girshick R (2017) Mask r-cnn. In: IEEE International Conference on Computer Vision, pp 2961–2969
  * Huang et al. (2019) Huang Z, Wang X, Huang L, Huang C, Wei Y, Liu W (2019) Ccnet: Criss-cross attention for semantic segmentation. In: IEEE International Conference on Computer Vision, pp 603–612
  * Ioffe and Szegedy (2015) Ioffe S, Szegedy C (2015) Batch normalization: Accelerating deep network training by reducing internal covariate shift. In: International Conference on Machine Learning, pmlr, pp 448–456
  * Jaegle et al. (2021) Jaegle A, Gimeno F, Brock A, Vinyals O, Zisserman A, Carreira J (2021) Perceiver: General perception with iterative attention. In: International Conference on Machine Learning, PMLR, pp 4651–4664
  * Kasai et al. (2021) Kasai J, Peng H, Zhang Y, Yogatama D, Ilharco G, Pappas N, Mao Y, Chen W, Smith NA (2021) Finetuning pretrained transformers into rnns. In: Conference on Empirical Methods in Natural Language Processing, pp 10630–10643
  * Katharopoulos et al. (2020) Katharopoulos A, Vyas A, Pappas N, Fleuret F (2020) Transformers are rnns: Fast autoregressive transformers with linear attention. In: International Conference on Machine Learning, PMLR, pp 5156–5165
  * Kitaev et al. (2020) Kitaev N, Kaiser Ł, Levskaya A (2020) Reformer: The efficient transformer. In: International Conference on Learning Representations
  * Krizhevsky (2009) Krizhevsky A (2009) Learning multiple layers of features from tiny images. URL https://api.semanticscholar.org/CorpusID:18268744
  * Lin et al. (2014) Lin TY, Maire M, Belongie S, Hays J, Perona P, Ramanan D, Dollár P, Zitnick CL (2014) Microsoft coco: Common objects in context. In: European Conference on Computer Vision, Springer, pp 740–755
  * Lin et al. (2017) Lin TY, Goyal P, Girshick R, He K, Dollár P (2017) Focal loss for dense object detection. In: IEEE International Conference on Computer Vision, pp 2980–2988
  * Liu et al. (2021) Liu Z, Lin Y, Cao Y, Hu H, Wei Y, Zhang Z, Lin S, Guo B (2021) Swin transformer: Hierarchical vision transformer using shifted windows. In: Proceedings of the IEEE/CVF international conference on computer vision, pp 10012–10022
  * Liu et al. (2022a) Liu Z, Hu H, Lin Y, Yao Z, Xie Z, Wei Y, Ning J, Cao Y, Zhang Z, Dong L, et al. (2022a) Swin transformer v2: Scaling up capacity and resolution. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 12009–12019
  * Liu et al. (2022b) Liu Z, Mao H, Wu CY, Feichtenhofer C, Darrell T, Xie S (2022b) A convnet for the 2020s. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 11976–11986
  * Long et al. (2015) Long J, Shelhamer E, Darrell T (2015) Fully convolutional networks for semantic segmentation. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 3431–3440
  * Lu et al. (2021) Lu J, Yao J, Zhang J, Zhu X, Xu H, Gao W, Xu C, Xiang T, Zhang L (2021) Soft: softmax-free transformer with linear complexity. Advances in Neural Information Processing Systems 34:21297–21309
  * Maas et al. (2011) Maas A, Daly RE, Pham PT, Huang D, Ng AY, Potts C (2011) Learning word vectors for sentiment analysis. In: Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pp 142–150
  * Mindspore (2020) Mindspore (2020) https://www.mindspore.cn/
  * Nangia and Bowman (2018) Nangia N, Bowman SR (2018) Listops: A diagnostic dataset for latent tree learning. In: Conference of the North American Chapter of the Association for Computational Linguistics, Student Research Workshop, pp 92–99
  * Peng et al. (2021) Peng H, Pappas N, Yogatama D, Schwartz R, Smith NA, Kong L (2021) Random feature attention. In: International Conference on Learning Representations
  * Radev et al. (2013) Radev DR, Muthukrishnan P, Qazvinian V, Abu-Jbara A (2013) The acl anthology network corpus. Language Resources and Evaluation 47:919–944
  * Radosavovic et al. (2020) Radosavovic I, Kosaraju RP, Girshick R, He K, Dollár P (2020) Designing network design spaces. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 10428–10436
  * Sun et al. (2021) Sun S, Yue X, Bai S, Torr P (2021) Visual parser: Representing part-whole hierarchies with transformers. arXiv preprint arXiv:210705790
  * Szegedy et al. (2016) Szegedy C, Vanhoucke V, Ioffe S, Shlens J, Wojna Z (2016) Rethinking the inception architecture for computer vision. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 2818–2826
  * Tay et al. (2020) Tay Y, Dehghani M, Abnar S, Shen Y, Bahri D, Pham P, Rao J, Yang L, Ruder S, Metzler D (2020) Long range arena: A benchmark for efficient transformers. In: International Conference on Learning Representations
  * Tay et al. (2023) Tay Y, Dehghani M, Bahri D, Metzler D (2023) Efficient transformers: A survey. ACM Computing Surveys 55(6):109:1–109:28
  * Touvron et al. (2021a) Touvron H, Cord M, Douze M, Massa F, Sablayrolles A, Jégou H (2021a) Training data-efficient image transformers & distillation through attention. In: International Conference on Machine Learning, PMLR, pp 10347–10357
  * Touvron et al. (2021b) Touvron H, Cord M, Sablayrolles A, Synnaeve G, Jégou H (2021b) Going deeper with image transformers. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 32–42
  * Touvron et al. (2022) Touvron H, Cord M, Jégou H (2022) Deit iii: Revenge of the vit. In: European Conference on Computer Vision, pp 516–533
  * Tsai et al. (2019) Tsai YHH, Bai S, Yamada M, Morency LP, Salakhutdinov R (2019) Transformer dissection: A unified understanding of transformer’s attention via the lens of kernel. In: Conference on Empirical Methods in Natural Language Processing and International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp 4344–4353
  * Vaswani et al. (2017) Vaswani A, Shazeer N, Parmar N, Uszkoreit J, Jones L, Gomez AN, Kaiser Ł, Polosukhin I (2017) Attention is all you need. Advances in Neural Information Processing Systems 30
  * Von Luxburg (2007) Von Luxburg U (2007) A tutorial on spectral clustering. Statistics and computing 17:395–416
  * Wang et al. (2020) Wang S, Li BZ, Khabsa M, Fang H, Ma H (2020) Linformer: Self-attention with linear complexity. arXiv preprint arXiv:200604768
  * Wang et al. (2021) Wang W, Xie E, Li X, Fan DP, Song K, Liang D, Lu T, Luo P, Shao L (2021) Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In: IEEE International Conference on Computer Vision, pp 568–578
  * Wang et al. (2022) Wang W, Xie E, Li X, Fan DP, Song K, Liang D, Lu T, Luo P, Shao L (2022) Pvt v2: Improved baselines with pyramid vision transformer. Computational Visual Media 8(3):415–424
  * Wang et al. (2018) Wang X, Girshick R, Gupta A, He K (2018) Non-local neural networks. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 7794–7803
  * Wightman (2019) Wightman R (2019) Pytorch image models. https://github.com/rwightman/pytorch-image-models
  * Williams and Seeger (2000) Williams C, Seeger M (2000) Using the nyström method to speed up kernel machines. Advances in Neural Information Processing Systems 13
  * Xiao et al. (2018) Xiao T, Liu Y, Zhou B, Jiang Y, Sun J (2018) Unified perceptual parsing for scene understanding. In: European Conference on Computer Vision, pp 418–434
  * Xie et al. (2017) Xie S, Girshick R, Dollár P, Tu Z, He K (2017) Aggregated residual transformations for deep neural networks. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 1492–1500
  * Xiong et al. (2021) Xiong Y, Zeng Z, Chakraborty R, Tan M, Fung G, Li Y, Singh V (2021) Nyströmformer: A nyström-based algorithm for approximating self-attention. In: AAAI Conference on Artificial Intelligence, vol 35, pp 14138–14148
  * Xu et al. (2021) Xu W, Xu Y, Chang T, Tu Z (2021) Co-scale conv-attentional image transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision, pp 9981–9990
  * Yoshida and Miyato (2017) Yoshida Y, Miyato T (2017) Spectral norm regularization for improving the generalizability of deep learning. arXiv preprint arXiv:170510941
  * Yuan et al. (2021) Yuan L, Chen Y, Wang T, Yu W, Shi Y, Jiang ZH, Tay FE, Feng J, Yan S (2021) Tokens-to-token vit: Training vision transformers from scratch on imagenet. In: IEEE International Conference on Computer Vision, pp 558–567
  * Yuan et al. (2020) Yuan Y, Chen X, Wang J (2020) Object-contextual representations for semantic segmentation. In: European Conference on Computer Vision, pp 173–190
  * Yun et al. (2019) Yun S, Han D, Oh SJ, Chun S, Choe J, Yoo Y (2019) Cutmix: Regularization strategy to train strong classifiers with localizable features. In: IEEE International Conference on Computer Vision, pp 6023–6032
  * Zhang et al. (2018) Zhang H, Cissé M, Dauphin YN, Lopez-Paz D (2018) mixup: Beyond empirical risk minimization. In: International Conference on Learning Representations
  * Zhang et al. (2020) Zhang L, Xu D, Arnab A, Torr PH (2020) Dynamic graph message passing networks. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 3726–3735
  * Zhang et al. (2021) Zhang P, Dai X, Yang J, Xiao B, Yuan L, Zhang L, Gao J (2021) Multi-scale vision longformer: A new vision transformer for high-resolution image encoding. In: IEEE International Conference on Computer Vision, pp 2998–3008
  * Zhao et al. (2017) Zhao H, Shi J, Qi X, Wang X, Jia J (2017) Pyramid scene parsing network. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 2881–2890
  * Zhao et al. (2020) Zhao H, Jia J, Koltun V (2020) Exploring self-attention for image recognition. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 10076–10085
  * Zheng et al. (2021) Zheng S, Lu J, Zhao H, Zhu X, Luo Z, Wang Y, Fu Y, Feng J, Xiang T, Torr PH, et al. (2021) Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers. In: IEEE Conference on Computer Vision and Pattern Recognition, pp 6881–6890
  * Zhou et al. (2019) Zhou B, Zhao H, Puig X, Xiao T, Fidler S, Barriuso A, Torralba A (2019) Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision 127:302–321
  * Zhu et al. (2021) Zhu X, Su W, Lu L, Li B, Wang X, Dai J (2021) Deformable detr: Deformable transformers for end-to-end object detection. In: International Conference on Learning Representations


# gym-saturation: an OpenAI Gym environment for saturation provers

Boris Shminke

††margin:  DOI:
https://doi.org/10.21105/joss.03849https://doi.org10.21105/joss.03849 Software
• https://github.com/openjournals/joss-
reviews/issues/3849https://doi.orgReview • https://github.com/inpefess/gym-
saturationhttps://doi.orgRepository •
https://doi.org/10.5281/zenodo.6324282https://doi.orgArchive Editor:
https://github.com/VivianePonshttps://doi.org@VivianePons
Reviewers: • https://github.com/lutzhamelhttps://doi.org@lutzhamel •
https://github.com/quickbeam123https://doi.org@quickbeam123 Submitted: 01
October 2021
Published: 03 March 2022 License
Authors of papers retain copyright and release the work under a Creative
Commons Attribution 4.0 International License
(http://creativecommons.org/licenses/by/4.0/https://doi.orgCC BY 4.0).

gym-saturation is an OpenAI Gym (Brockman et al., 2016) environment for
reinforcement learning (RL) agents capable of proving theorems. Currently,
only theorems written in a formal language of the Thousands of Problems for
Theorem Provers (TPTP) library (Sutcliffe, 2017) in clausal normal form (CNF)
are supported. gym-saturation implements the ‘given clause’ algorithm (similar
to the one used in Vampire (Kovács & Voronkov, 2013) and E Prover (Schulz et
al., 2019)). Being written in Python, gym-saturation was inspired by PyRes
(Schulz & Pease, 2020). In contrast to the monolithic architecture of a
typical Automated Theorem Prover (ATP), gym-saturation gives different agents
opportunities to select clauses themselves and train from their experience.
Combined with a particular agent, gym-saturation can work as an ATP. Even with
a non trained agent based on heuristics, gym-saturation can find refutations
for 688 (of 8257) CNF problems from TPTP v7.5.0.

## Statement of need

Current applications of RL to saturation-based ATPs like Enigma (Jakubuv et
al., 2020) or Deepire (Suda, 2021) are similar in that the environment and the
agent are not separate pieces of software but parts of larger systems that are
hard to disentangle. The same is true for non saturation-based RL-friendly
provers too (e.g. lazyCoP, Rawson & Reger (2021)). This monolithic approach
hinders free experimentation with novel machine learning (ML) models and RL
algorithms and creates unnecessary complications for ML and RL experts willing
to contribute to the field. In contrast, for interactive theorem provers,
projects like HOList (Bansal, Loos, Rabe, Szegedy, & Wilcox, 2019) or GamePad
(Huang et al., 2019) separate the concepts of environment and agent. Such
modular architecture may lead to the development of easily comparable agents
based on diverse approaches (see, e.g. Paliwal et al. (2020) or Bansal, Loos,
Rabe, & Szegedy (2019)). gym-saturation is an attempt to implement a modular
environment-agent architecture of an RL-based ATP. In addition, some RL
empowered saturation ATPs are not accompanied with their source code
(Abdelaziz et al., 2022), while gym-saturation is open-source software.

## Usage example

Suppose we want to prove an extremely simple theorem with a very basic agent.
We can do that in the following way:

# first we create and reset a OpenAI Gym environmentfrom importlib.resources
import filesimport gymenv = gym.make( "gym_saturation:saturation-v0", # we
will try to find a proof shorter than 10 steps step_limit=10, # for a
classical syllogism about Socrates problem_list=[
files("gym_saturation").joinpath( "resources/TPTP-
mock/Problems/TST/TST003-1.p" ) ],)env.reset()# we can render the environment
(that will become the beginning of the proof)print("starting
hypotheses:")print(env.render("human"))# our 'age' agent will always select
clauses for inference# in the order they appeared in current proof
attemptaction = 0done = Falsewhile not done: observation, reward, done, info =
env.step(action) action += 1# SaturationEnv has an additional method# for
extracting only clauses which became parts of the proof# (some steps were
unnecessary to find the proof)print("refutation
proof:")print(env.tstp_proof)print(f"number of attempted steps: {action}")

The output of this script includes a refutation proof found:

    starting hypotheses:
    cnf(p_imp_q, hypothesis, ~man(X0) | mortal(X0)).
    cnf(p, hypothesis, man(socrates)).
    cnf(q, hypothesis, ~mortal(socrates)).
    refutation proof:
    cnf(_0, hypothesis, mortal(socrates), inference(resolution, [], [p_imp_q, p])).
    cnf(_2, hypothesis, $false, inference(resolution, [], [q, _0])).
    number of attempted steps: 6

## Architecture

gym-saturation includes several sub-packages:

  * •

parsing (happens during env.reset() in example code snippet)

  * •

logic operations (happen during env.step(action) in the example)

  * •

AI Gym environment implementation

  * •

agent testing (a bit more elaborated version of the while loop from the
examle)

gym-saturation relies on a deduction system of four rules which is known to be
refutationally complete (Brand, 1975):

$\displaystyle{\frac{C_{1}\vee A_{1},C_{2}\vee\neg
A_{2}}{\sigma\left(C_{1}\vee
C_{2}\right)}},\sigma=mgu\left(A_{1},A_{2}\right)\quad\text{(resolution)}$
$\displaystyle{\frac{C_{1}\vee s\approx t,C_{2}\vee
L\left[r\right]}{\sigma\left(L\left[t\right]\vee C_{1}\vee
C_{2}\right)}},\sigma=mgu\left(s,r\right)\quad\text{(paramodulation)}$
$\displaystyle{\frac{C\vee A_{1}\vee A_{2}}{\sigma\left(C\vee
L_{1}\right)}},\sigma=mgu\left(A_{1},A_{2}\right)\quad\text{(factoring)}$
$\displaystyle\frac{C\vee s\not\approx
t}{\sigma\left(C\right)},\sigma=mgu\left(s,t\right)\quad\text{(reflexivity
resolution)}$

where $C,C_{1},C_{2}$ are clauses, $A_{1},A_{2}$ are atomic formulae, $L$ is a
literal, $r,s,t$ are terms, and $\sigma$ is a substitution (most general
unifier). $L\left[t\right]$ is a result of substituting the term $t$ in
$L\left[r\right]$ for the term $r$ at only one chosen position.

For parsing, we use the LARK parser (Shinan, 2021). We represent the clauses
as Python classes forming tree-like structures. gym-saturation also includes a
JSON serializer/deserializer for those trees. For example, a TPTP clause

    cnf(a2,hypothesis,
        ( ~ q(a) | f(X) = X )).

becomes

Clause( literals=[ Literal( negated=True, atom=Predicate( name="q",
arguments=[Function(name="a", arguments=[])] ), ), Literal( negated=False,
atom=Predicate( name="=", arguments=[ Function(name="f",
arguments=[Variable(name="X")]), Variable(name="X"), ], ), ), ], label="a2",)

This grammar serves as the glue for gym-saturation sub-packages, which are, in
principle, independent of each other. After switching to another parser or
another deduction system, the agent testing script won’t break, and RL
developers won’t need to modify their agents for compatibility (for them, the
environment will have the same standard OpenAI Gym API).

Figure 1: A diagram showing interactions between four main subpackages of gym-
saturation: 1) parsing; 2) logic operations (including the given clause
algorithm); 3) OpenAI Gym Env implementation; 4) the agent testing script.

Agent testing is a simple episode pipeline (see Figure 1). It is supposed to
be run in parallel (e.g. using GNU Parallel, Tange (2021)) for a testing
subset of problems. See the following table for the testing results of two
popular heuristic-based agents on TPTP v7.5.0 (trained RL agents should strive
to be more successful than those primitive baselines):

| size agent | age agent | size&age agent
---|---|---|---
proof found | 509 | 206 | 688
step limit | 1385 | 35 | 223
out of memory | 148 | 149 | 148
5 min time out | 6215 | 7867 | 7198
total | 8257 | 8257 | 8257

size agent is an agent which always selects the shortest clause.

age agent is an agent which always selects the clause which arrived first to
the set of unprocessed clauses (‘the oldest one’).

size&age agent is an agent which selects the shortest clause five times in a
row and then one time — the oldest one.

‘Step limit’ means an agent didn’t find proof after 1000 steps (the longest
proof found consists of 287 steps). This can work as a ‘soft timeout.’

## Mentions

At the moment of writing this paper, gym-saturation was used by its author
during their PhD studies for creating experimental RL-based ATPs.

## Acknowledgements

This work has been supported by the French government, through the 3IA Côte
d’Azur Investments in the Future project managed by the National Research
Agency (ANR) with the reference number ANR-19-P3IA-0002. This work was
performed using HPC resources from GENCI-IDRIS (Grant 2021-AD011013125).

## References

Abdelaziz, I., Crouse, M., Makni, B., Austel, V., Cornelio, C., Ikbal, S.,
Kapanipathi, P., Makondo, N., Srinivas, K., Witbrock, M., & Fokoue, A. (2022).
Learning to guide a saturation-based theorem prover. _IEEE Transactions on
Pattern Analysis and Machine Intelligence_ , 1–1.
https://doi.org/10.1109/TPAMI.2022.3140382

Bansal, K., Loos, S. M., Rabe, M. N., & Szegedy, C. (2019). Learning to reason
in large theories without imitation. _CoRR_ , _abs/1905.10501_.
http://arxiv.org/abs/1905.10501

Bansal, K., Loos, S. M., Rabe, M. N., Szegedy, C., & Wilcox, S. (2019).
HOList: An environment for machine learning of higher order logic theorem
proving. In K. Chaudhuri & R. Salakhutdinov (Eds.), _Proceedings of the 36th
international conference on machine learning, ICML 2019, 9-15 june 2019, long
beach, california, USA_ (Vol. 97, pp. 454–463). PMLR.
http://proceedings.mlr.press/v97/bansal19a.html

Brand, D. (1975). Proving theorems with the modification method. _SIAM Journal
on Computing_ , _4_(4), 412–430. https://doi.org/10.1137/0204036

Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang,
J., & Zaremba, W. (2016). OpenAI gym. _CoRR_ , _abs/1606.01540_.
http://arxiv.org/abs/1606.01540

Huang, D., Dhariwal, P., Song, D., & Sutskever, I. (2019). GamePad: A learning
environment for theorem proving. _7th International Conference on Learning
Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_.
https://openreview.net/forum?id=r1xwKoR9Y7

Jakubuv, J., Chvalovský, K., Olsák, M., Piotrowski, B., Suda, M., & Urban, J.
(2020). ENIGMA anonymous: Symbol-independent inference guiding machine (system
description). In N. Peltier & V. Sofronie-Stokkermans (Eds.), _Automated
reasoning - 10th international joint conference, IJCAR 2020, paris, france,
july 1-4, 2020, proceedings, part II_ (Vol. 12167, pp. 448–463). Springer.
https://doi.org/10.1007/978-3-030-51054-1_29

Kovács, L., & Voronkov, A. (2013). First-order theorem proving and vampire. In
N. Sharygina & H. Veith (Eds.), _Computer aided verification - 25th
international conference, CAV 2013, saint petersburg, russia, july 13-19,
2013. proceedings_ (Vol. 8044, pp. 1–35). Springer.
https://doi.org/10.1007/978-3-642-39799-8_1

Paliwal, A., Loos, S. M., Rabe, M. N., Bansal, K., & Szegedy, C. (2020). Graph
representations for higher-order logic and theorem proving. _The Thirty-Fourth
AAAI Conference on Artificial Intelligence, AAAI 2020, the Thirty-Second
Innovative Applications of Artificial Intelligence Conference, IAAI 2020, the
Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI
2020, New York, NY, USA, February 7-12, 2020_ , 2967–2974.
https://aaai.org/ojs/index.php/AAAI/article/view/5689

Rawson, M., & Reger, G. (2021). lazyCoP: Lazy paramodulation meets neurally
guided search. In A. Das & S. Negri (Eds.), _Automated reasoning with analytic
tableaux and related methods - 30th international conference, TABLEAUX 2021,
birmingham, UK, september 6-9, 2021, proceedings_ (Vol. 12842, pp. 187–199).
Springer. https://doi.org/10.1007/978-3-030-86059-2_11

Schulz, S., Cruanes, S., & Vukmirovic, P. (2019). Faster, higher, stronger: E
2.3. In P. Fontaine (Ed.), _Automated deduction - CADE 27 - 27th international
conference on automated deduction, natal, brazil, august 27-30, 2019,
proceedings_ (Vol. 11716, pp. 495–507). Springer.
https://doi.org/10.1007/978-3-030-29436-6_29

Schulz, S., & Pease, A. (2020). Teaching automated theorem proving by example:
PyRes 1.2 - (system description). In N. Peltier & V. Sofronie-Stokkermans
(Eds.), _Automated reasoning - 10th international joint conference, IJCAR
2020, paris, france, july 1-4, 2020, proceedings, part II_ (Vol. 12167, pp.
158–166). Springer. https://doi.org/10.1007/978-3-030-51054-1_9

Shinan, E. (2021). _Lark-parser_ (Version 0.12.0).
https://pypi.org/project/lark-parser/

Suda, M. (2021). Improving ENIGMA-style clause selection while learning from
history. In A. Platzer & G. Sutcliffe (Eds.), _Automated deduction - CADE 28 -
28th international conference on automated deduction, virtual event, july
12-15, 2021, proceedings_ (Vol. 12699, pp. 543–561). Springer.
https://doi.org/10.1007/978-3-030-79876-5_31

Sutcliffe, G. (2017). The TPTP Problem Library and Associated Infrastructure.
From CNF to TH0, TPTP v6.4.0. _Journal of Automated Reasoning_ , _59_(4),
483–502.

Tange, O. (2021). _GNU parallel 20210822 (’kabul’)_. Zenodo.
https://doi.org/10.5281/zenodo.5233953