SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tam Nguyen
Department of Electrical & Computer Engineering
Houston, USA
mn72@rice.edu
Tan M. Nguyen
tanmn@nus.edu.sg
Richard G. Baraniuk
richb@rice.edu
Transformers have achieved remarkable success in a wide range of natural language processing and computer vision applications. However, the representation capacity of a deep transformer model is degraded due to the over-smoothing issue in which the token representations become identical when the model's depth grows. In this work, we show that self-attention layers in transformers minimize a functional which promotes smoothness, thereby causing token uniformity. We then propose a novel regularizer that penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens to preserve the fidelity of the tokens. Minimizing the resulting regularized energy functional, we derive the Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO), a novel class of transformer models that can mitigate the over-smoothing issue. We empirically demonstrate the advantages of NeuTRENO over the baseline transformers and state-of-the-art methods in reducing the over-smoothing of token representations on various practical tasks, including object classification, image segmentation, and language modeling.
Transformer models [62] have achieved substantial success in natural language processing [16, 2, 13, 10, 47, 4, 6, 14], reinforcement learning [9, 32], computer vision [19, 40, 59, 49, 44, 3, 41, 71, 27], and other practical applications [50, 33, 70, 26, 66]. Transformers also excel at transferring knowledge from pre-trained models to new tasks, even when limited supervision is available [45, 46, 16, 69, 39]. At the heart of transformers lies the self-attention mechanism, which computes a weighted average of token representations within a sequence. These weights are determined based on the similarity scores between pairs of tokens, determining their relative importance in the sequence [11, 43, 38]. This flexibility in capturing diverse syntactic and semantic relationships has been identified as a crucial factor contributing to the success of transformers [57, 63, 12, 64, 31].
1.1 Background: Self-Attention
For a given input sequence  $\mathbf{X} \coloneqq [\pmb{x}(1), \dots, \pmb{x}(N)]^{\top} \in \mathbb{R}^{N \times D_x}$  of  $N$  feature vectors, self-attention transforms  $\mathbf{X}$  into the output sequence  $\mathbf{H}$  in the following two steps:
ImageNet Classification
Figure 1: The cosine similarity between tokens representations across layers of NeuTRENO DeiT vs. the baseline DeiT models on the Imagenet classification and ADE20K image segmentation tasks. In both tasks, the DeiT baseline suffers from over-smoothing as tokens become similar to identical when the model gets deeper. In contrast, tokens in NeuTRENO models are significantly more diverse, suggesting a reduction in over-smoothing. Further details regarding this analysis can be found in Appendix E.
ADE20K Segmentation
Step 1. The input sequence  $\mathbf{X}$  is projected into the query matrix  $\mathbf{Q}$ , the key matrix  $\mathbf{K}$ , and the value matrix  $\mathbf{V}$  via three linear transformations
$$
\mathbf {Q} = \mathbf {X} \mathbf {W} _ {Q} ^ {\top}; \mathbf {K} = \mathbf {X} \mathbf {W} _ {K} ^ {\top}; \mathbf {V} = \mathbf {X} \mathbf {W} _ {V} ^ {\top}, \tag {1}
$$
where  $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{D_{q_k} \times D_x}$ , and  $\mathbf{W}_V \in \mathbb{R}^{D \times D_x}$  are the weight matrices. We denote  $Q \coloneqq [q(1), \ldots, q(N)]^\top$ ,  $\mathbf{K} \coloneqq [k(1), \ldots, k(N)]^\top$ , and  $\mathbf{V} \coloneqq [v(1), \ldots, v(N)]^\top$ , where the vectors  $q(i), k(i)$ , and  $v(i)$ , for  $i = 1, \ldots, N$  are the query, key, and value vectors, respectively.
Step 2. The output sequence  $\mathbf{U} \coloneqq [\pmb{u}(1), \dots, \pmb{u}(N)]^{\top} \in \mathbb{R}^{N \times D_{qk}}$  is then computed as follows
$$
\mathbf {U} = \operatorname {s o f t m a x} \left(\mathbf {Q} \mathbf {K} ^ {\top} / \sqrt {D _ {q k}}\right) \mathbf {V} := \mathbf {A V}, \tag {2}
$$
where the softmax function is applied to each row of the matrix  $\mathbf{Q}\mathbf{K}^{\top} / \sqrt{D_{qk}}$ . The matrix  $\mathbf{A} := \operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D_{qk}}}\right) \in \mathbb{R}^{N \times N}$  and its component  $a_{ij}$  for  $i, j = 1, \dots, N$  are called the attention matrix and attention scores, respectively. For each query vector  $\pmb{q}(i)$  for  $i = 1, \dots, N$ , an equivalent form of Eqn. (2) to compute the output vector  $\pmb{u}(i)$  is given by
$$
\boldsymbol {u} (i) = \sum_ {j = 1} ^ {N} \operatorname {s o f t m a x} \left(\boldsymbol {q} (i) ^ {\top} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right) \boldsymbol {v} (j). \tag {3}
$$
The self-attention computed by Eqn. (2) and (3) is referred as softmax attention. In our work, we refer to a transformer that uses softmax attention as a softmax transformer.
1.2 Over-smoothing in Transformers
Despite their remarkable success, deep transformer-based models have been observed to suffer from the over-smoothing issue, in which all token representations become identical when more layers are added to the models [55, 65, 18]. This over-smoothing phenomenon, also known as the "token uniformity" problem, significantly limits the representation capacity of transformers. To illustrate this phenomenon, we examine the average cosine similarity between pairs of token representations across different layers in a softmax transformer trained for theImagenet object classification and ADK20 image segmentation tasks [73]. As depicted in Fig. 1, in both tasks, this cosine similarity between tokens increases as the models become deeper. Particularly, in the last two layers, the cosine similarity scores are approximately 0.9, indicating a high degree of similarity among tokens.
1.3 Contribution
We develop a nonlocal variational denoising framework for self-attention, providing insights into the over-smoothing phenomenon in transformers. In particular, by viewing self-attention as a gradient descent step toward minimizing a nonlocal functional that penalizes high-frequency noise in the signal, we uncover the diffusive nature of self-attention, which explains the over-smoothing issue of transformers. Motivated by this understanding, we propose the Neural Transformer with a Regularized
DeiT NeuTRENO DeiT
Nonlocal Functional (NeuTRENO), a novel class of transformers designed to mitigate over-smoothing. NeuTRENO is derived by optimizing a regularized nonlocal functional, which includes an additional convex fidelity term. This fidelity term penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens, thereby reducing the over-smoothing effect. Our contribution is three-fold.
1. We develop a nonlocal variational denoising framework for self-attention and shed light on the over-smoothing issue that hampers the representation capacity of transformers.
2. We develop NeuTRENO, a novel class of transformers that are capable of alleviating the over-smoothing issue.
3. We theoretically prove that transformers with softmax self-attention are prone to oversmoothing while NeuTRENO can avoid this issue.
We empirically demonstrate the benefits of NeuTRENO on various large-scale applications, including the ImageNet object classification, ADE20K image segmentation, and WikiText-103 language modeling tasks.
Organization: We organize our paper as follows: in Section 2, we develop a nonlocal variational denoising framework for self-attention and provide an explanation for the over-smoothing issue in transformer-based models. In section 3, we propose NeuTRENO, and present a theoretical result that guarantees NeuTRENO's capability of mitigating over-smoothing. In Section 4, we empirically validate the benefits of NeuTRENO. We discuss the related work in Section 6. Finally, we conclude our main contributions and remarks. Further results, details, and proofs are provided in the Appendix.
2 A Nonlocal Variational Denoising Framework for Self-attention
We first consider the output matrix  $\mathbf{U} \coloneqq [\pmb{u}(1), \dots, \pmb{u}(N)]^\top \in \mathbb{R}^{N \times D}$  in self-attention as given by Eqn. 2 in Section 1.1. Let  $\Omega \subset \mathbb{R}$ ,  $x \in \Omega$ , and  $\pmb{u}(x) \coloneqq [u_1(x), \dots, u_D(x)]^T$  be a real vector-valued function,  $\pmb{u} : \Omega \to \mathbb{R}^D$ ,  $\pmb{u} \in L^2(\Omega)$ . The output matrix  $\mathbf{U}$  in self-attention discretizes the function  $\pmb{u}(x)$  on a 1-D grid. In the context of signal/image denoising,  $\mathbf{U}$  can be considered as the desired clean signal, and  $\pmb{u}(x)$  is its corresponding intensity function denoting the signal values at the position  $x \in \Omega$ . We further let the observed intensity function  $\pmb{f}(x)$  denote the values of the observed noisy signal at  $x \in \Omega$ ,  $\pmb{f} : \Omega \to \mathbb{R}^D$ ,  $\pmb{f} \in L^2(\Omega)$ . For example,  $\pmb{f}(x)$  can be given as
$$
\boldsymbol {f} (x) = \boldsymbol {u} (x) + \boldsymbol {n} (x), \tag {4}
$$
where  $\pmb{n}$  is the additive noise. We wish to reconstruct  $\pmb{u}(x)$  from  $\pmb{f}(x)$ . Following the variational denoising method proposed in [23] and [24], the denoised image  $\pmb{u}(x)$  can be obtained by minimizing the following regularized functional with respect to  $\pmb{u}$ :
$$
\begin{array}{l} E (\boldsymbol {u}, \boldsymbol {f}) = J (\boldsymbol {u}) + G (\boldsymbol {u}, \boldsymbol {f}) \tag {5} \\ = \frac {1}{2} \int_ {\Omega \times \Omega} \| \boldsymbol {u} (x) - \boldsymbol {u} (y) \| _ {2} ^ {2} k (x, y) d x d y + \frac {\lambda}{2} \int_ {\Omega} \| \boldsymbol {u} (x) - \boldsymbol {f} (x) \| _ {2} ^ {2} d x. \\ \end{array}
$$
Here,  $J(\pmb{u}) = \frac{1}{2}\int_{\Omega \times \Omega} \| \pmb{u}(x) - \pmb{u}(y) \|_2^2 k(x,y)dxdy$  is a nonlocal functional of weighted differences. The weights  $k(x,y)$  represent the affinity between signal values at positions  $x$  and  $y$ . For example, for images,  $k(x,y)$  captures the proximity between pixels  $x$  and  $y$  in the image.  $J(\pmb{u})$  works as a regularizer. Minimizing  $J(\pmb{u})$  promotes the smoothness of  $\pmb{u}$  and penalizes high-frequency noise in the signal. Adding the convex fidelity term  $G(\pmb{u},\pmb{f}) = \frac{\lambda}{2}\int_{\Omega} \| \pmb{u}(x) - \pmb{f}(x) \|_2^2 dx$  to the functional  $J(\pmb{u})$  allows the denoised signal  $\pmb{u}(x)$  to preserve relevant information in the observed noisy signal  $\pmb{f}(x)$ . The regularized functional  $E(\pmb{u},\pmb{f})$  can be considered as an energy functional.
2.1 Self-attention as a Gradient Descent Step to Minimize the Nonlocal Functional  $J$
We show that self-attention is equivalent to taking a gradient descent step toward minimizing the functional  $J(\pmb{u})$  in the energy functional  $E(\pmb{u},\pmb{f})$ . We expand  $J(\pmb{u})$  as follows
$$
J (\boldsymbol {u}) = \frac {1}{2} \int_ {\Omega \times \Omega} \sum_ {j = 1} ^ {D} \left(u _ {j} (x) - u _ {j} (y)\right) ^ {2} k (x, y) d x d y \tag {6}
$$
The gradient of  $J$  with respect to  $\pmb{u}$  is then given by
$$
\nabla_ {\boldsymbol {u}} J (\boldsymbol {u}) = \left[ \frac {\partial J}{\partial u _ {1}}, \frac {\partial J}{\partial u _ {2}}, \dots , \frac {\partial J}{\partial u _ {D}} \right] ^ {T}. \tag {7}
$$
The partial derivative  $\partial J / \partial u_j$ ,  $j = 1,2,\ldots ,D$ , is defined through its dot product with an arbitrary function  $h_j\in L^2 (\Omega)$  as follows
$$
\begin{array}{l} \frac {\partial J}{\partial u _ {j}} \cdot h _ {j} (x) = \left. \frac {d}{d \tau} J (u _ {j} + \tau h _ {j}) \right| _ {\tau = 0} \\ = \frac {1}{2} \left(\frac {d}{d \tau} \int_ {\Omega \times \Omega} (u _ {j} (x) - u _ {j} (y) + \tau h _ {j} (x) - \tau h _ {j} (y)) ^ {2} k (x, y) d x d y\right) \Bigg | _ {\tau = 0} \\ = \left(\int_ {\Omega \times \Omega} (u _ {j} (x) - u _ {j} (y) + \tau h _ {j} (x) - \tau h _ {j} (y)) (h _ {j} (x) - h _ {j} (y)) k (x, y) d x d y\right) \Bigg | _ {\tau = 0} \\ = \int_ {\Omega \times \Omega} \left(u _ {j} (x) - u _ {j} (y)\right) \left(h _ {j} (x) - h _ {j} (y)\right) k (x, y) d x d y \\ = \int_ {\Omega \times \Omega} \left(u _ {j} (x) - u _ {j} (y)\right) h _ {j} (x) k (x, y) d x d y - \int_ {\Omega \times \Omega} \left(u _ {j} (x) - u _ {j} (y)\right) h _ {j} (y) k (x, y) d x d y \\ \end{array}
$$
Applying a change of variables  $(x,y)\to (y,x)$  to the second term of the above integral, we have
$$
\begin{array}{l} \frac {\partial J}{\partial u _ {j}} \cdot h _ {j} (x) = \int_ {\Omega \times \Omega} \left(u _ {j} (x) - u _ {j} (y)\right) h _ {j} (x) k (x, y) d x d y - \int_ {\Omega \times \Omega} \left(u _ {j} (y) - u _ {j} (x)\right) h _ {j} (x) k (y, x) d x d y \\ = \int_ {\Omega \times \Omega} (u _ {j} (x) - u _ {j} (y) (k (x, y) + k (y, x)) d y h _ {j} (x) d x \\ \end{array}
$$
Thus, the Frechet derivative of  $J$  with respect to  $u_{j}$  is given by
$$
\frac {\partial J}{\partial u _ {j}} = \int_ {\Omega} \left(u _ {j} (x) - u _ {j} (y) (k (x, y) + k (y, x)) d y. \right. \tag {8}
$$
Substituting the formula for  $\partial J / \partial u_{j}$  in Eqn. 8 into Eqn. 7 for  $\nabla_{\pmb{u}}J(\pmb{u})(x)$ , we obtain the following gradient flow
$$
\frac {d \boldsymbol {u} (x , t)}{d t} = - \nabla_ {\boldsymbol {u}} J (\boldsymbol {u}) = \int_ {\Omega} \left(\boldsymbol {u} (y, t) - \boldsymbol {u} (x, t)\right) \left(k (x, y) + k (y, x)\right) d y, \tag {9}
$$
where  $t$  is the time variable we introduce to capture the dynamics of  $\pmb{u}$  when gradient descent is applied to minimize  $J(\pmb{u})$ . Let  $\pmb{v}(x) \coloneqq [v_1(x), \dots, v_D(x)]^T$  be a real vector-valued function,  $\pmb{v} : \Omega \to \mathbb{R}^D$ ,  $\pmb{v} \in L^2(\Omega)$ . We discretize  $\pmb{v}(x)$  on a 1-D grid to attain the value vectors  $\pmb{v}(1), \dots, \pmb{v}(N) \in \mathbb{R}^D$ , which form the value matrix  $\mathbf{V} \coloneqq [\pmb{v}(1), \dots, \pmb{v}(N)]^\top \in \mathbb{R}^{N \times D}$  in self-attention as defined in Eqn. 2. We initialize  $\pmb{u}$  at  $t = 0$  with  $\pmb{v}(x)$ , i.e.,  $\pmb{u}(x, 0) = \pmb{v}(x)$ .
Self-attention is an Euler Discretization of the Gradient Flow Given in 9. We discretize the gradient flow in Eqn. 9 using the Euler method [21] with step size  $\Delta t(x) = 1 / \int_{\Omega} (k(x, y) + k(y, x)) dy$  and obtain the following update
$$
\begin{array}{l} \boldsymbol {u} (x, \Delta t (x)) = \boldsymbol {u} (x, 0) + \Delta t (x) \int_ {\Omega} (\boldsymbol {u} (y, 0) - \boldsymbol {u} (x, 0)) (k (x, y) + k (y, x)) d y \\ = \int_ {\Omega} \frac {\left(k (x , y) + k (y , x)\right) \boldsymbol {u} (y , 0)}{\int_ {\Omega} \left(k (x , y ^ {\prime}) + k \left(y ^ {\prime} , x\right)\right) d y ^ {\prime}} d y = \int_ {\Omega} \frac {K (x , y) \boldsymbol {v} (y)}{\int_ {\Omega} K (x , y ^ {\prime}) d y ^ {\prime}} d y. \tag {10} \\ \end{array}
$$
Here,  $K(x,y) \coloneqq k(x,y) + k(y,x)$  is a symmetric kernel and  $\pmb{u}(y,0) = \pmb{v}(y)$  since  $\pmb{u}$  is initialized at  $t = 0$  with  $\pmb{v}$  as aforementioned. Let  $\pmb{k}(x) \coloneqq [k_1(x),\dots,k_{D_{q_k}}(x)]^T$  be a real vector-valued function,  $\pmb{k}:\Omega \to \mathbb{R}^{D_{q_k}}$ ,  $\pmb{k} \in L^2(\Omega)$ . Similar to  $\pmb{u}(x)$  and  $\pmb{v}(x)$ , we can discretize  $\pmb{k}(x)$  on a 1-D grid to attain the key vectors  $\pmb{k}(1),\dots,\pmb{k}(N) \in \mathbb{R}^{D_{q_k}}$ , which form the key matrix  $\mathbf{K} \coloneqq [\pmb{k}(1),\dots,\pmb{k}(N)]^\top \in \mathbb{R}^{N\times D_{q_k}}$  in self-attention as defined in Eqn. 2. We choose  $K(x,y) = \exp\bigl(\pmb{k}(x)^T\pmb{k}(y)/\sqrt{D_{q_k}}\bigr)$  and rewrite Eqn. 10 as follows
$$
\boldsymbol {u} (x, \Delta t (x)) = \int_ {\Omega} \frac {\exp (\boldsymbol {k} (x) ^ {T} \boldsymbol {k} (y) / \sqrt {D _ {q k}})}{\int_ {\Omega} \exp (\boldsymbol {k} (x) ^ {T} \boldsymbol {k} (y ^ {\prime}) / \sqrt {D _ {q k}}) d y ^ {\prime}} \boldsymbol {v} (y) d y. \tag {11}
$$
Estimating the integrals in Eqn. 11 via Monte-Carlo approximation using the key vectors  $\pmb{k}(1),\dots,\pmb{k}(N)\in \mathbb{R}^{D_{qk}}$  and value vectors  $\pmb{v}(1),\dots,\pmb{v}(N)\in \mathbb{R}^D$ , we obtain
$$
\boldsymbol {u} (x, \Delta t (x)) \approx \sum_ {j = 1} ^ {N} \frac {\exp \left(\boldsymbol {k} (x) ^ {T} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right)}{\sum_ {j ^ {\prime} = 1} ^ {N} \exp \left(\boldsymbol {k} (x) ^ {T} \boldsymbol {k} \left(j ^ {\prime}\right) / \sqrt {D _ {q k}}\right)} \boldsymbol {v} (j). \tag {12}
$$
Discretizing  $\pmb{u}(x, \Delta t(x))$  on another 1-D grid, we attain
$$
\begin{array}{l} \boldsymbol {u} (i) \approx \sum_ {j = 1} ^ {N} \frac {\exp \big (\boldsymbol {k} (i) ^ {T} \boldsymbol {k} (j) / \sqrt {D _ {q k}} \big)}{\sum_ {j ^ {\prime} = 1} ^ {N} \exp \big (\boldsymbol {k} (i) ^ {T} \boldsymbol {k} (j ^ {\prime}) / \sqrt {D _ {q k}} \big)} \boldsymbol {v} (j) \\ = \sum_ {j = 1} ^ {N} \operatorname {s o f t m a x} \left(\boldsymbol {k} (i) ^ {\top} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right) \boldsymbol {v} (j), i = 1, \dots , N. \tag {13} \\ \end{array}
$$
Comparing Eqn. 13 and Eqn. 3, we observe that Eqn. 13 implement a symmetric self-attention, in which the query matrix  $\mathbf{Q}$  and the key matrix  $\mathbf{K}$  are the same, i.e.  $\mathbf{W}_Q = \mathbf{W}_K$  where  $\mathbf{W}_Q$  and  $\mathbf{W}_K$  are the linear projections that map the input sequence  $\mathbf{X}$  into  $\mathbf{Q}$  and  $\mathbf{K}$  as given in Eqn. 1. This symmetry of the attention scores is desirable in some image processing tasks due to the symmetric similarities between pixels, but can be relaxed for other tasks. To break the symmetry of attention scores in Eqn. 13, we replace the key vectors  $k(i)$  by the query vectors  $\pmb{q}(i), i = 1,\dots ,N$ , to obtain the exact formula of self-attention given by Eqn. 3. The following theorem summarizes our results:
Theorem 1 (Self-attention as a Gradient Descent Step to Minimize a Nonlocal Functional). Given the nonlocal functional  $J(\pmb{u}) = \frac{1}{2} \int_{\Omega \times \Omega} \| \pmb{u}(x) - \pmb{u}(y) \|_2^2 k(x, y) \, dx \, dy$  of a vector-valued function  $\pmb{u}: \Omega \to \mathbb{R}^D$ ,  $\pmb{u} \in L^2(\Omega)$ , and let  $K(x, y) := k(x, y) + k(y, x) = \exp \big( k(x)^T \pmb{k}(y) / \sqrt{D_{qk}} \big)$ , where  $\pmb{k}: \Omega \to \mathbb{R}^{D_{qk}}$ ,  $\pmb{k} \in L^2(\Omega)$ . Then, taking a gradient descent step on  $\pmb{u}$  at time  $t = 0$ , where  $\pmb{u}(x, 0) = \pmb{v}(x)$ , with an adaptive step size  $\Delta t(x) := \frac{1}{\int_{\Omega} (k(x, y) + k(y, x)) \, dy}$  to minimize  $J$  is equivalent to updating  $\pmb{u}$  via a symmetric self-attention
$$
\boldsymbol {u} (x, \Delta t (x)) = \sum_ {j = 1} ^ {N} \operatorname {s o f t m a x} \left(\boldsymbol {k} (x) ^ {\top} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right) \boldsymbol {v} (j),
$$
which results in
$$
\boldsymbol {u} (i) = \sum_ {j = 1} ^ {N} \operatorname {s o f t m a x} \left(\boldsymbol {k} (i) ^ {\top} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right) \boldsymbol {v} (j), i = 1, \dots , N. \tag {14}
$$
Here,  $\pmb{u}(n)$ ,  $\pmb{v}(n)$ , and  $\pmb{u}(n)$ ,  $n = 1, \dots, N$ , are the key, value, and output vectors in self-attention, respectively. Breaking the symmetry of the attention scores by replacing  $\pmb{k}(i)$  with  $\pmb{q}(i)$ ,  $i = 1, \dots, N$ , in Eqn. 14, we obtain the exact formula of self-attention
$$
\boldsymbol {u} (i) = \sum_ {j = 1} ^ {N} \operatorname {s o f t m a x} \left(\boldsymbol {q} (i) ^ {\top} \boldsymbol {k} (j) / \sqrt {D _ {q k}}\right) \boldsymbol {v} (j), i = 1, \dots , N.
$$
Remark 1. In Eqn. 9, the change in  $\mathbf{u}$  at position  $x$  is proportional to the sum of differences between  $\mathbf{u}(x)$  and  $\mathbf{u}$  at other positions in the domain  $\Omega$ . In particular, when  $\mathbf{u}(x)$  is smaller or larger than the values at other positions, it will increase or decrease, respectively. This is analogous to a diffusion process in which particles or substances move from high-concentration to low-concentration regions. It has been proved that a diffusion process converges to a saturating state in which the concentrations at all positions are the same. This suggests that  $\mathbf{u}(x)$  tends to suffer from the over-smoothing issue.
2.2 Random Walk Analysis of Over-smoothing
The diffusion process and random walk are closely related concepts, as diffusion can be seen as a collective behavior of numerous random walks performed by individual particles or molecules. Inspired by the analogy between the dynamics of  $\pmb{u}$  in Eqn 9 and a diffusion process, as well as the relationship between diffusion process and random walk, in this section, we show the connection
between the evolution of  $\pmb{u}$  and a random walk. By adopting a random walk perspective on graph neural network [58], we demonstrate that  $\pmb{u}(\pmb{x})$  under the dynamics given in Eqn 9 suffers from over-smoothing.
Recall from the gradient flow in Eqn 9, by using Euler method discretization, after  $k$  update steps starting from the initial  $\pmb{u}(x,0) = \pmb{v}(x)$ , with adaptive stepsize  $\Delta t = 1 / \int_{\Omega} (k(x,y) + k(y,x)) dy$ , we obtain the following
$$
\boldsymbol {u} (x, k \Delta t (x)) = \int_ {\Omega} \frac {K (x , y) \boldsymbol {u} (y , (k - 1) \Delta t (x))}{\int_ {\Omega} K (x , y ^ {\prime}) d y ^ {\prime}} d y. \tag {15}
$$
Discretizing  $\pmb{u}(x, k\Delta t(x))$  and using Monte-Carlo approximation for the integrals in 15, we attain
$$
\boldsymbol {u} ^ {(k)} (i) = \sum_ {j = 1} ^ {N} \mathbf {A} _ {i j} \boldsymbol {u} ^ {(k - 1)} (j) \tag {16}
$$
where  $\mathbf{A}_{ij}$  is computed using the keys and queries as either  $\mathrm{softmax}\Bigl (k(i)^{\top}k(j) / \sqrt{D_{qk}}\Bigr)$  or  $\mathrm{softmax}\Bigl (q(i)^{\top}k(j) / \sqrt{D_{qk}}\Bigr)$ . Let  $\{\mathbf{B}^{(k)}(i)\}_{k\in K}$  be a random walk on  $\{\pmb {v}(i)\}_{i = 1}^{N}$  as defined:
$$
\mathbf {B} ^ {(0)} (i) = \boldsymbol {v} (i)
$$
$$
\mathbb {P} \left(\mathbf {B} ^ {(k + 1)} (l) = \boldsymbol {v} (j) \mid \mathbf {B} ^ {(k)} (l) = \boldsymbol {v} (i)\right) = \mathbf {A} _ {i j} \tag {17}
$$
where  $\mathbf{B}^{(k)}(n)$  is the random value of a  $k$ -step walk, starts at node  $n$ , and  $\pmb{v}(n)$  is the initial value at node  $n$ , respectively, for  $n = 1,2,\dots,N$ . The transition probability  $\mathbf{A}$  is defined as above. To investigate the connection between the update process of  $\pmb{u}$  and the random walk defined in 17, we show that, for  $i = 1,2,\dots,N$ , after  $k$  update steps as in 16, with initial value  $\pmb{u}^{(0)}(i) = \pmb{v}(i)$ ,  $\pmb{u}(i)^{(k)}$  equals to the expected value of the  $k$ -step walk, starting at node  $i$ :
Lemma 1. Let  $\mathbf{u}^{(k)}(i)$  defined in 16 and  $\{\mathbf{B}^{(k)}(i)\}_{k\in K}$  is the random walk defined by 17. Then
$$
\boldsymbol {u} ^ {(k)} (i) = \mathbb {E} \left[ \mathbf {B} ^ {(k)} (i) \right]. \tag {18}
$$
We next present the Lemma 2 which is necessary to show the convergence of  $\pmb{u}^{(k)}(i)$ .
Lemma 2. The random walk  $\mathbf{B}^{(k)}(i)$  in 17 with the transition matrix  $\mathbf{A}$  either be  $\mathbf{A}_{ij} = \mathrm{softmax}\Bigl (\pmb {k}(i)^{\top}\pmb {k}(j) / \sqrt{D_{qk}}\Bigr)$  or  $\mathbf{A}_{ij} = \mathrm{softmax}\Bigl (\pmb {q}(i)^{\top}\pmb {k}(j) / \sqrt{D_{qk}}\Bigr)$ , has a unique stationary distribution  $\pi = [\pi_1,\pi_2,\dots ,\pi_N]$  such that  $\pi_i\coloneqq P(\mathbf{B}^{(k)}(j) = \pmb {v}(i))$ , for  $i,j = 1,2,\ldots ,N$ ,  $\sum_{i = 1}^{N}\pi_{i} = 1$ , and  $\pmb{\pi}^T = \pmb{\pi}^T\mathbf{A}$ .
If  $\mathbf{A}_{ij} = \mathrm{softmax}\Bigl (\pmb {k}(i)^{\top}\pmb {k}(j) / \sqrt{D_{qk}}\Bigr)$ , the stationary distribution is:
$$
\boldsymbol {\pi} = \left(\frac {d _ {1}}{\sum_ {j = 1} ^ {N} d _ {j}}, \frac {d _ {2}}{\sum_ {j = 1} ^ {N} d _ {j}}, \dots , \frac {d _ {n}}{\sum_ {j = 1} ^ {N} d _ {j}}\right), \tag {19}
$$
where  $d_{i} = \sum_{j = 1}^{N}\exp \Bigl {(}\pmb {k}(i)^{\top}\pmb {k}(j) / \sqrt{D_{qk}}\Bigr),$ $\pmb {k}(1),\pmb {k}(2),\dots ,\pmb {k}(N)$  are the key vectors.