Lightweight transformer image feature extraction network

Linear attention basics

First, mathematically describe the general form of Transformer, where the given input was represented in the embedding space as X ∈ ℝ^(n×d), denoting the transformation that the input undergoes into the Transformer module as T, thereby the definition of transforming T is shown in Eq. (1): (1)${\displaystyle \mathrm{T}\left(\mathbf{X}\right)=\mathrm{F}\left(\mathrm{Att}\left(\mathbf{X}\right)+\mathbf{X}\right)}$

where F is a feedforward neural network containing residual connections; Att is a function used to compute the self-attention matrix.

The self-attention mechanism in the current Transformer model is scaling point product attention, and for the convenience of expression, the scaling factor of ATT is omitted, and the definition is shown in Eq. (2): (2)${\displaystyle \mathrm{Att}\left(\mathbf{Q},\mathbf{K},\mathbf{V}\right)=\text{Softmax}\left({\mathbf{QK}}^{\mathbf{T}}\right)\mathbf{V}}$

where Q ∈ℝ^(n×d), K ∈ ℝ^(n×d), V ∈ ℝ^(n×d), which can be calculated from the inputs, respectively.

An in-depth analysis of Eq. (2) shows that it is the Softmax operator defined in Eq. (2) that restricts the performance of the scaling point product attention mechanism and makes it quadratic, so we need to calculate QK^T first, and a matrix of n × n was obtained, so that the complexity is O (n²) level. If there is no Softmax operator, then the three matrices QK^TV will be multiplied, and using the combination law of matrix multiplication, the last two matrices K^TV can be calculated first to get a matrix of d_k ×d_v, and then let the matrix Q to do the left multiplication, due to the actual case d_k sum d_v is much smaller than n, so the overall complexity can be seen as O(n) level. The Softmax operator is defined as shown in Eq. (3): (3)${\displaystyle Softmax\left(z_i\right)=\frac{{\mathit{e}}^{z_i}}{\sum _{\mathit{c}=1}^{\mathit{c}}{\mathit{e}}^{z_c}}}$

where z_i represents the the i-th node’s output value, and C denotes the total number of output nodes.

The Softmax operator could convert the multiple clusters’ output values to a probability distribution at range between 0 and 1, with a sum of 1. With the introduction of Eqs. (3) and (2) can be rewritten as follows, as shown in Eq. (4): (4)${\displaystyle \mathrm{Att}{\left(\mathbf{Q},\mathbf{K},\mathbf{V}\right)}_i=\frac{\sum _{\mathit{j}=1}^{\mathit{n}}\mathrm{S}\left({\mathbf{q}}_i,{\mathbf{k}}_j\right){\mathbf{v}}_j}{\sum _{\mathit{j}=1}^{\mathit{n}}\mathrm{S}\left({\mathbf{q}}_i,{\mathbf{k}}_j\right)}}$

To preserve the similar properties of the attention mechanism, S(⋅) ≥0 should be constant. This general form of attention is also known as non-local networking (Wang et al., 2018).

For arriving the most ideal linear level on the complexity of attention, a decomposable method of calculation is needed to effectively approximate the similarity function S(⋅). The associative law of matrix multiplication can be used to calculate the product of the next two moment matrices first, reducing the complexity to the linear O(n) level (Zhuoran et al., 2021), as shown in Fig. 2.

Most existing linear attention mechanisms aim to approximate the Softmax operator estimation, but these methods are sensitive to the choice of sampling rate. Hence, it is worth considering whether a decomposable similarity function can be directly employed as a replacement for the Softmax operator under the premise of ensuring the experimental effect. To achieve this, it is necessary to identify the crucial characteristics of the current attention mechanism and design a decomposable similarity function that fulfills the prerequisites for linear attention.

Alternative function design

Through an in-depth reading of literature research, two critical properties that significantly impact the performance of current attention mechanisms are summarized. They are non-negative elements in the attention matrix (Katharopoulos et al., 2020; Tsai et al., 2019) and non-linear reweighting schemes (Zhen et al., 2022; Paoletti et al., 2023; Baldi & Vershynin, 2023).

First, only positive values are kept in the attention matrix, and features with negative correlation are ignored, thus effectively avoiding aggregation of irrelevant contextual information. Second, the nonlinear reweighting mechanism can make a focus on the distribution weights of attention, thereby stabilizing the training process. It can also help the model to amplify the locality, which means that a large part of the context dependence comes from neighboring markers.

Combining the above conclusions, this subsection proposes a combination function to replace Softmax. This combination function satisfies the above two properties and consists of two sub-functions. are used to implement non-negativity and non-linear reweighting, respectively, as shown in Eq. (5): (5)${\displaystyle \mathrm{Att}\left(\mathbf{Q},\mathbf{K},\mathbf{V}\right)=g\left(f\left(\mathbf{Q},\mathbf{K}\right)\right)\mathbf{V}}$ In the formula, f(⋅) and g(⋅) are two self-defined functions to ensure the above two properties, defined as shown in Eqs. (6) and (7):

(6)${\displaystyle f\left(x\right)=\left\{\begin{array}{cc}x+1\hfill & x\ge 0\hfill \\ e^x\hfill & x\text{¡}0\hfill \end{array}\right.}$ (7)${\displaystyle g\left({\mathbf{q}}_i,{\mathbf{k}}_j\right)={\mathbf{q}}_i{\mathbf{k}}_j^{\mathrm{T}}sin\left(\frac{\pi \left(i+n-j\right)}{2n}\right)}$

where i, j =1 , …, n.

Equation(6) is used to ensure the non-negativity of the attention matrix.

Equation (7) is formulated in this form for several reasons:

(1) Realize local deviation. When i and j are close, the corresponding trigonometric function value is close to 1. When i and j are far away, the difference between the two is close to n, and the corresponding trigonometric function value is close to 0, so that the corresponding The similarity function is negligible;

(2) The trigonometric function itself is nonlinear, so the design can focus on the distribution of weights, so as to achieve the purpose of stabilizing the training process;

(3) Eq. (7) can be decomposed by the sum-difference product formula, so that the associative law of matrix multiplication can be used to reduce the complexity of the attention mechanism to linear.

The Eq. (8) is the dismantling process of the sub-function g(⋅) for satisfying the non-linear reweighting property: (8)${\displaystyle \begin{array}{c}g\left({\mathbf{q}}_i^{\prime },{\mathbf{k}}_j^{\prime }\right)={\mathbf{q}}_i^{\prime }{\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}sin\left(\frac{\pi \left(i+n-j\right)}{2n}\right)\hfill \\ ={\mathbf{q}}_i^{\prime }{\mathbf{k}}_j^{\prime }\mathrm{T}sin\left(\frac{\pi \left(i-j\right)}{2n}+\frac{\pi }{2}\right)\hfill \\ ={\mathbf{q}}_i^{\prime }{\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}\left(cos\left(\frac{\pi i}{2n}\right)cos\left(\frac{\pi j}{2n}\right)+sin\left(\frac{\pi i}{2n}\right)sin\left(\frac{\pi j}{2n}\right)\right)\hfill \\ =\left({\mathbf{q}}_i^{\prime }cos\left(\frac{\pi i}{2n}\right)\right)\left({\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}cos\left(\frac{\pi j}{2n}\right)\right)+\left({\mathbf{q}}_i^{\prime }sin\left(\frac{\pi i}{2n}\right)\right)\left({\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}sin\left(\frac{\pi j}{2n}\right)\right)\hfill \end{array}}$

Among them, ${\mathbf{q}}_i^{\prime }=f\left({\mathbf{q}}_i^{}\right)$, ${\mathbf{k}}_j^{\prime }=f\left({\mathbf{k}}_j^{}\right)$, f(⋅) are shown in Eq. (2).

Let ${\mathbf{q}}_i^{cos}={\mathbf{q}}_i^{\prime }cos\left(\frac{\pi i}{2n}\right)$, ${\mathbf{q}}_i^{sin}={\mathbf{q}}_i^{\prime }sin\left(\frac{\pi i}{2n}\right)$, ${\mathbf{k}}_j^{cos}={\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}cos\left(\frac{\pi j}{2n}\right)$ and ${\mathbf{k}}_j^{sin}={\mathbf{k}}_j^{{}^{\prime }}\mathrm{T}sin\left(\frac{\pi j}{2n}\right)$.

Apparently, ${\mathbf{q}}_i^{sin}$ and ${\mathbf{q}}_i^{cos}$ are transformed from the i-th column vector ${\mathbf{q}}_i^{}$ of Q through self-defined non-negative function f(⋅) and sine or cosine function. Denote the Q after this transformation as Q^cos and Q^sin. Similarly, K after transformation is K^cos and K^sin. Thus, the final expression after linear decomposition is obtained, as shown in Eq. (9): (9)${\displaystyle \mathrm{Att}\left(\mathbf{Q},\mathbf{K},\mathbf{V}\right)=g\left(f\left(\mathbf{Q},\mathbf{K}\right)\right)\mathbf{V}={\mathbf{Q}}^{cos}\left({\mathbf{K}}^{cos}\mathbf{V}\right)+{\mathbf{Q}}^{sin}\left({\mathbf{K}}^{sin}\mathbf{V}\right)}$

The attention calculation Eq. (9) proposed in this section can make use of the associative law of matrix multiplication to minimize the computational complexity of the linear attention mechanism. Figure 3 shows the calculation process.

Score

After bypassing the path of training additional parameters to obtain the score of the input token, we can focus on the existing information of the image feature extraction network itself designed based on Transformer. All tokens input by the self-attention layer inside the model are divided into two parts. The first part only includes the first class token, which is additionally introduced to judge the category of the input image in the final stage, and it is placed in the first position of the token sequence. The second part is all the remaining tokens, which are obtained by dividing and transforming the input image.

To prune from the token dimension is to reduce the second part of tokens converted from images as much as possible and record the number of tokens in this part as m. That is to say, there are m + 1 tokens in the input, and there are r′ + 1 tokens after the output. Here r′ is the number of tokens in the second part retained after pruning, and the corresponding data range is r′ ≤ R ≤ m. Here R is a parameter given in advance, and its function is to control the maximum number of tokens retained after sampling.

In the standard self-attention layer, $\mathbf{Q}\in {\mathbb{R}}^{\left(\left(m+1\right)\times d_q\right)},\mathbf{K}\in {\mathbb{R}}^{\left(\left(m+1\right)\times d_k\right)},\mathbf{V}\in {\mathbb{R}}^{\left(\left(m+1\right)\times d_v\right)}$ is calculated from the input token, which is recorded as $\mathbf{X}\in {\mathbb{R}}^{\left(\left(m+1\right)\times d\right)}$. The self-attention matrix A is calculated from both Q and K^T. It should be noted that the attention calculation method at this time is still the standard scaling dot product self-attention. The attention matrix captures the association’s similarity or degree among all input tokens.

For example, in the self-attention matrix A, the elements of [i, j] or [j, i] represent the correlation degree between the i-th and the j-th tokens among all the currently input tokens. The larger the value, the more related the two are, indicating that each other is more important to another token.

It can be inspired from this, since the first class token of the currently input token sequence is used for the classification of the final stage. Then in the self-attention matrix, if other tokens are more correlated with the first class token used for classification, it means that these tokens are more important for classifying the class token of the input image category at least at the current stage. To measure the similarity between these tokens and the first-class token, all values except the 1st one in the 1st row of the self-attention matrix are utilized. Therefore, when designing the score formula for calculating the importance of each token except class token, the calculation method used is shown in Eq. (10): (10)${\displaystyle h_j=\frac{a_{1,j}}{\sum _{i=2}^m{a}_{1,i}}}$

In the Equation, a_1,j and a_1,i represent the elements of the 1st row, column j and column i respectively of the self-attention matrix A, h_j is the j-th token’s score.

In addition, for the multi-head attention layer, each head’s score can be calculated separately, and then all the heads can be added together.

Sampling

In the previous section, the importance scores of each token have been obtained in a parameterless way. Now some tokens can be removed from the self-attention matrix based on these scores. It is natural to think of first screening out some tokens with the highest scores, keeping them, and then directly removing the remaining tokens with lower scores.

But this method has two problems: First, the number of tokens that are screened out is not easy to determine, and it is difficult to achieve adaptive dynamic changes based on different input images. Secondly, if tokens with lower scores are directly screened out at a certain stage from the beginning, then these screened out tokens are not necessarily unimportant for the final classification. Just as the CNN-based image feature extraction network extracts only shallow features such as edges and textures at the beginning, but it extracts more advanced semantic features later. Different tokens may have different functions at different stages and represent different meanings. The current score is only obtained at the current stage. If it is simply deleted based on the low effect of an intermediate stage, the deleted token will not be able to enter the subsequent stage. But it may play an important role in a later stage, which will affect the final experimental results instead.

Therefore, according to these scores, some tokens can be randomly selected to be retained by sampling. If the score is lower, it has a lower probability of being selected at the current stage and thus cannot be retained, and vice versa. This method is indeed better than crudely deleting tokens with lower scores in terms of experimental results. This randomness can offset the deficiency of directly deleting token to a certain extent.

Then it needs to be sampled according to the scores of these tokens. Tokens with higher scores have a greater probability of being retained, and tokens with lower scores have a lower probability of being retained. That is to find a probability distribution that obeys these scores, and then sample according to this probability distribution. The cumulative distribution function corresponding to these scores can be calculated and inverse transformed to obtain the inverse function of this random distribution.

First, cumulative distribution function (CDF) can characterize the probability distribution of a random variable x, and it is denoted as X. For all real numbers x, the cumulative distribution function is defined (Han et al., 2021) as shown in Eq. (11): (11)${\displaystyle {\mathrm{F}}_{\mathrm{X}}\left(x\right)=\mathrm{P}\left(\mathrm{X}\le x\right).}$

If the cumulative distribution function F is a continuous strictly increasing function, then there is a corresponding inverse function F⁻¹(y), and the range of y is [0,1). The inverse function can be utilized to produce random variables that follow this random distribution. That is, if take F_X(x) as the CDF of X, and there is an inverse function $F_X^{-1}$. If a is a random variable with the range of [0, 1) and has a uniformly distribution, then $F_X^{-1}\left(a\right)$ obeys the X distribution. Therefore, the CDF can be calculated according to the following formula, as shown in Eq. (12): (12)${\displaystyle {\mathrm{CDF}}_i=\sum _{j=2}^i{h}_j.}$

Note here that the accumulation starts from the second token, because the first classification class token must be retained. After obtaining the CDF, the sampling function can be obtained according to its inverse form, as shown in Eq. (13): (13)${\displaystyle \eta \left(v\right)={\mathrm{CDF}}^{-1}\left(v\right)}$

where v is between 0 and 1.

Specifically, the sampling strategy in this section is to randomly select a number v from the uniform distribution between 0 and 1. Then calculate the corresponding η(v), and then select the nearest integer for sampling. This operation needs to be performed a fixed number of R times. In these R samplings, a token may be selected multiple times, so that the number of tokens actually sampled r′ is less than or equal to the fixed R.

After sampling the tokens to be retained, the initial attention matrix A changes from the initial m + 1 row to the current r′ + 1 row, and then adds it to the next calculation process.