Energy Transformer β€” part 4/5

Cleaned-text code points [44781, 60608). Bibliography omitted here; source.md includes it.

First, we compute the features that are passed to our energy-based transformer. Each node’s features 𝐲 A ∈ R F \mathbf{y}_{A}\in R^{F} are mapped into the token space 𝐱 A ∈ R D \mathbf{x}_{A}\in R^{D} , using a linear projection 𝐄 \mathbf{E} . Learnable positional embeddings Ξ» A \mathbf{\lambda}_{A} are added to this token at t = 1 t=1 , 𝐱 A t = 1 = 𝐄𝐲 A + Ξ» A \mathbf{x}_{A}^{t=1}=\mathbf{Ey}_{A}+\mathbf{\lambda}_{A} (14) At each time step the input to the ET-block is layer normalized: 𝐠 A t = LayerNorm ​ ( 𝐱 A t ) \mathbf{g}_{A}^{t}=\text{LayerNorm}(\mathbf{x}_{A}^{t}) (15) Let 𝐖 Q ∈ R Y Γ— H Γ— D \mathbf{W}^{Q}\in R^{Y\times H\times D} and 𝐖 K ∈ R Y Γ— H Γ— D \mathbf{W}^{K}\in R^{Y\times H\times D} be the query and key weight matrices, respectively. Here Y Y is the projection dimension in the attention operation, H H is the number of heads. We define K Ξ± ​ h ​ B = βˆ‘ j W Ξ± ​ h ​ j K ​ g j ​ B Q Ξ± ​ h ​ C = βˆ‘ j W Ξ± ​ h ​ j Q ​ g j ​ C \begin{split}K_{\alpha hB}&=\sum\limits_{j}W^{K}_{\alpha hj}\;g_{jB}\\ Q_{\alpha hC}&=\sum\limits_{j}W^{Q}_{\alpha hj}\;g_{jC}\end{split} (16) If we let h h indicate the index of the head, we have Ξ” ​ x i ​ A t = βˆ‘ C ∈ 𝒩 A βˆ‘ h , Ξ± [ W Ξ± ​ h ​ i Q ​ K Ξ± ​ h ​ C ​ Ο‰ C ​ A + W Ξ± ​ h ​ i K ​ Q Ξ± ​ h ​ C ​ Ο‰ A ​ C ] + βˆ‘ ΞΌ ΞΎ ΞΌ ​ i ​ r ​ ( βˆ‘ j ΞΎ ΞΌ ​ j ​ g j ​ A ) \Delta x_{iA}^{t}=\sum\limits_{C\in\mathcal{N}_{A}}\sum\limits_{h,\alpha}\Big[W^{Q}_{\alpha hi}\;K_{\alpha hC}\;\omega_{CA}+W^{K}_{\alpha hi}\;Q_{\alpha hC}\;\omega_{AC}\Big]+\sum\limits_{\mu}\xi_{\mu i}\;r\Big(\sum\limits_{j}\;\xi_{\mu j}g_{jA}\Big) (17) where Ο‰ C ​ A = softmax 𝐢 ​ ( Ξ² ​ βˆ‘ Ξ³ K Ξ³ ​ h ​ C ​ Q Ξ³ ​ h ​ A ) \omega_{CA}=\underset{C}{\text{softmax}}\Big(\beta\sum\limits_{\gamma}K_{\gamma hC}\;Q_{\gamma hA}\Big) (18)

Here Ξ² \beta controls the temperature of the softmax, 𝒩 A \mathcal{N}_{A} stands for the neighbors of node A A β€”a set of all the nodes connected to node A A , r r is the ReLU function. Restriction of the attention operation to the neighborhood of a given node is similar to that used in the Graph Attention Networks (GAT), see [ 54 ] . Finally, we have residual connection (which is a natural consequence of the discretized differential equation dynamics) 𝐱 A t + 1 = 𝐱 A t + 𝚫 ​ 𝐱 A t \mathbf{x}_{A}^{t+1}=\mathbf{x}_{A}^{t}+\mathbf{\Delta x}_{A}^{t} (19)

Intuitively, the first term in Equation 17 describes the influence (attention score) of the neighbor nodes with respect to the target node, the second term describes the influence of the target node with respect to each of its neighbor, and the third term is the contribution of the Hopfield Network module. It can be shown that the forward pass of our energy-based transformer layer minimizes the following energy function: E = βˆ’ 1 Ξ² βˆ‘ C βˆ‘ h log ( βˆ‘ B ∈ 𝒩 C exp ( Ξ² βˆ‘ Ξ± K Ξ± ​ h ​ B Q Ξ± ​ h ​ C ) ) βˆ’ 1 2 βˆ‘ C , ΞΌ G ( βˆ‘ j ΞΎ ΞΌ ​ j g j ​ C ) E=-\frac{1}{\beta}\sum\limits_{C}\sum\limits_{h}\textrm{log}\left(\sum\limits_{B\in\mathcal{N}_{C}}\textrm{exp}\left(\beta\sum\limits_{\alpha}K_{\alpha hB}\;Q_{\alpha hC}\right)\right)-\frac{1}{2}\sum\limits_{C,\mu}G\Big(\sum\limits_{j}\xi_{\mu j}\;g_{jC}\Big) (20) This energy function will decrease as the forward pass progresses until it reaches a local minimum.

After T T iterations when the retrieval is stable, we have the final representation for each node 𝐠 A final \mathbf{g}_{A}^{\text{final}} as 𝐠 A final = 𝐠 A t = 1 | | 𝐠 A t = T \mathbf{g}_{A}^{\text{final}}=\mathbf{g}_{A}^{t=1}\ ||\ \mathbf{g}_{A}^{t=T} (21) where | ⁣ | || is the concatenation sign. Following [ 24 ] , we treat anomaly detection as semi-supervised learning task in this work. The node representation 𝐠 A final \mathbf{g}_{A}^{\text{final}} is fed to another MLP with the sigmoid function to compute the abnormal probability p A p_{A} , weighted log-likelihood is then used to train the network. The loss function is as follow: Loss = βˆ‘ A [ Ο‰ ​ l A ​ log ⁑ ( p A ) + ( 1 βˆ’ l A ) ​ log ⁑ ( 1 βˆ’ p A ) ] \text{Loss}=\sum_{A}\Big[\omega\;l_{A}\log(p_{A})+(1-l_{A})\log(1-p_{A})\Big] (22) where Ο‰ \omega is the ratio of normal labels ( l A l_{A} = 0) to anomaly labels ( l A l_{A} = 1).

### C.2 Experimental Details

We train all models for 100 epochs using the Adam optimizer with a learning rate of 0.001, and use the model with the best Macro-F1 on the validation set for reporting the final results on the test set. Following [ 24 ] , we use training ratios 1% and 40% respectively (randomly select 1% and 40% nodes of the dataset to train the model, and use the remaining nodes for the validation and testing). These remaining nodes are split 1:2 for validation:testing. The statistics of the datasets are listed in Table 5 . For the four datasets used in the experiments, Amazon and Yelp datasets can be obtained from the DGL library, T-Finance and T-Social can be obtained from [ 24 ] .

We report the average performance of 5 runs on the test datasets. The hyperparameters of our model are tuned based on the validation set, selecting the best parameters within 100 epochs. To speedup the training process, for the large graph datasets T-Finance and T-Social, we sample a different subgraph to train for each epoch (subgraphs have 5 % 5\% of the nodes with respect to the whole training data). The hyperparameters include the number of hidden dimensions in ET-attention Y Y , the number of neurons K K in the hidden layer within the Hopfield Network Module, the number of time iterations T T , and the number of heads H H . The weights are learned via backpropagation, which includes embedding projection 𝐄 \mathbf{E} , positional embedding Ξ» A \mathbf{\lambda}_{A} , softmax inverse temperature parameter Ξ² \beta , ET-attention weight tensors 𝐖 Q \mathbf{W}^{Q} and 𝐖 K \mathbf{W}^{K} . The optimal hyperparameters used in Table 1 are reported in Table 6 . The last row in that table summarizes the range of the hyperparameter search that was performed in our experiments. In general, we have observed that for small datasets (Yelp, Amazon, T-Finance) a 1 or 2 applications of our network is sufficient for achieving strong results, for larger datasets (T-Social) more iterations (3) are necessary. For even bigger dataset (ImageNet) our network needs about 12 12 iterations.

## Appendix D Graph Classification with ET

Recently, there have been attempts of leveraging Transformer into the graph domain, but only certain key modules, such as feature aggregation, are replaced in GNN variants by the softmax attention [ 55 ] . However, there remains an interesting and open question about suitability of the Transformer architecture to model graphs and how to apply it in the graph classification domain. In this section we explain how ET can be used for graph classification.

### D.1 Details of Graph Classification ET Model

Given a graph G = ( V , 𝐀 , β„° ) G=(V,\mathbf{A},\mathcal{E}) , we have the set of nodes V = { v 1 , v 2 , … , v N } V=\{v_{1},v_{2},\dots,v_{N}\} , the adjacency matrix 𝐀 ∈ { 0 , 1 } N Γ— N \mathbf{A}\in\{0,1\}^{N\times N} , and if it exists, the edge feature matrix β„° ∈ ℝ N Γ— N Γ— P β€² \mathcal{E}\in\mathbb{R}^{N\times N\times P^{\prime}} where each edge has P β€² P^{\prime} raw features. Each feature vector 𝐲 A ∈ ℝ F \mathbf{y}_{A}\in\mathbb{R}^{F} corresponding to a node v A v_{A} is first projected to the token space 𝐱 𝐀 ∈ ℝ D \mathbf{x}_{\mathbf{A}}\in\mathbb{R}^{D} via a linear embedding. Then, the CLS token 𝐱 CLS \mathbf{x}_{\text{CLS}} is concatenated to the set of tokens resulting in 𝐗 ∈ ℝ ( N + 1 ) Γ— D \mathbf{X}\in\mathbb{R}^{(N+1)\times D} and the positional embedding Ξ» ∈ ℝ ( N + 1 ) Γ— D \lambda\in\mathbb{R}^{(N+1)\times D} is added to it afterwards.

To obtain the positional embedding Ξ» ∈ ℝ ( N + 1 ) Γ— D \lambda\in\mathbb{R}^{(N+1)\times D} , the adjacency matrix 𝐀 \mathbf{A} is first padded in the upper left corner with ones resulting in 𝐀 ~ ∈ { 0 , 1 } ( N + 1 ) Γ— ( N + 1 ) \tilde{\mathbf{A}}\in\{0,1\}^{(N+1)\times(N+1)} as a means to provide the CLS token full connectivity with all of the nodes in a given graph. Following [ 37 ] , the top k k smallest eigen-vectors Ξ» ~ ∈ ℝ ( N + 1 ) Γ— k \tilde{\lambda}\in\mathbb{R}^{(N+1)\times k} are extracted from the normalized Laplacian matrix 𝐋 ~ \tilde{\mathbf{L}} obtained from 𝐀 ~ \tilde{\mathbf{A}} 𝐋 ~ = 𝐈 βˆ’ 𝐃 ~ βˆ’ 1 2 ​ 𝐀 ~ ​ 𝐃 ~ βˆ’ 1 2 , \tilde{\mathbf{L}}=\mathbf{I}-\tilde{\mathbf{D}}^{-\frac{1}{2}}\tilde{\mathbf{A}}\tilde{\mathbf{D}}^{-\frac{1}{2}}, (23) and projected to the token space via a linear embedding to form Ξ» \lambda . Note, 𝐈 \mathbf{I} and 𝐃 ~ \tilde{\mathbf{D}} are the identity matrix and the degree matrix of 𝐀 ~ \tilde{\mathbf{A}} .

Meanwhile, the attention in ET is modified to take in 𝐀 ^ ∈ ℝ ( N + 1 ) Γ— ( N + 1 ) Γ— H \mathbf{\hat{A}}\in\mathbb{R}^{(N+1)\times(N+1)\times H} the parameterized adjacency tensor, which acts as the weighted β€˜attention mask’ that enables the model to consider the graph structural information. To obtain 𝐀 ^ \mathbf{\hat{A}} , a 2D-convolutional layer with H H filters equals to the number of heads in the attention block, β€˜SAME’ padding, and a stride of one is performed on the outer product of 𝐗 \mathbf{X} to itself. The result is then multiplied with the tensor 𝐀 β€² ∈ ℝ ( N + 1 ) Γ— ( N + 1 ) Γ— H \mathbf{A}^{\prime}\in\mathbb{R}^{(N+1)\times(N+1)\times H} element-wise (denoted by βŠ™ \odot ) via broadcasting

𝐀 ^ = Conv2D ​ ( 𝐗 βŠ— 𝐗 ) βŠ™ 𝐀 β€² . \mathbf{\hat{A}}=\text{Conv2D}(\mathbf{X}\otimes\mathbf{X})\odot\mathbf{A}^{\prime}. (24)

The tensor 𝐀 β€² \mathbf{A}^{\prime} is a linear projection, via a linear embedding layer, of the edge feature matrix β„° \mathcal{E} if it exists, where P β€² P^{\prime} is projected to H H dimension; or the padded adjacency matrix 𝐀 ~ \tilde{\mathbf{A}} where P β€² = 1 P^{\prime}=1 . Altogether, the resulting energy equation is E = ATT βˆ’ 1 Ξ² βˆ‘ h βˆ‘ C log ( βˆ‘ B β‰  C exp ( Ξ² βˆ‘ Ξ± K Ξ± ​ h ​ B Q Ξ± ​ h ​ C βŠ™ A ^ h ​ C ) ) . E{{}^{\text{ATT}}}=-\frac{1}{\beta}\sum\limits_{h}\sum\limits_{C}\textrm{log}\left(\sum\limits_{B\neq C}\textrm{exp}\left(\beta\sum\limits_{\alpha}K_{\alpha hB}\;Q_{\alpha hC}\odot\hat{A}_{hC}\right)\right). (25) Moreover, to only consider the energy of edges in G G , we set non-connected entries in 𝐀 ^ \hat{\mathbf{A}} to be βˆ’ ∞ -\infty , such that the gradient of such entries would be zero following the dynamics. Note, if non-connected entries are not ignored, the gradient of such entries would be non-zero and hence, the resultant attention energy might violate some of the graph structural information.

Meanwhile, in this implementation, the overall model consists of S S vertically stacked ET blocks, where each block shares the same number of T T depth and has its own LayerNorm. Similarly, the token representation 𝐗 β„“ , t \mathbf{X}^{\ell,\;t} at dynamic step t t corresponding to a block β„“ \ell is layer-normalized, 𝐠 β„“ , t = LayerNorm ​ ( 𝐗 β„“ , t ) . \mathbf{g}^{\ell,\;t}=\text{LayerNorm}(\mathbf{X}^{\ell,\;t}). (26) Keep in mind, 𝐗 = 𝐗 β„“ = 1 , t = 1 \mathbf{X}=\mathbf{X}^{\ell=1,\;t=1} is the initial token representation.

Following the dynamic equations 6 , we inject a small amount of noise Ο΅ β„“ , t ∈ ( 0 , 1 ) \epsilon^{\ell,\;t}\in(0,1) , generated from a normal distribution using a standard deviation Οƒ Ο΅ \sigma_{\epsilon} and zero mean, into the gradient of energy function to produce 𝐗 β„“ , t + 1 \mathbf{X}^{\ell,\;t+1} , the new token representation of block β„“ \ell . The premise of this noise injection is to β€˜robustify’ the model and help it escape saddle points of the energy function. 𝐗 β„“ , t + 1 = 𝐗 β„“ , t βˆ’ Ξ± ​ βˆ‡ 𝐠 E β„“ , t + Ξ± ​ Ο΅ β„“ , t , Ο΅ β„“ , t ∼ 𝒩 ⁑ ( 0 , Οƒ Ο΅ 2 ) . \mathbf{X}^{\ell,\;t+1}=\mathbf{X}^{\ell,\;t}-\alpha\nabla_{\mathbf{g}}E^{\ell,\;t}+\sqrt{\alpha}\,\epsilon^{\ell,\;t},\qquad\epsilon^{\ell,\;t}\sim\mathcal{N}(0,\sigma_{\epsilon}^{2}). (27) Once stability is reached in the retrieval dynamics of a block β„“ \ell , the final representation 𝐗 β„“ , t = T \mathbf{X}^{\ell,\;t=T} is then passed on to the next block β„“ + 1 \ell+1 and the whole process is repeated again. When the final token representation 𝐗 β„“ = S , t = T \mathbf{X}^{\ell=S,\;t=T} is computed by the last block S S , the resultant CLS token 𝐱 CLS β„“ = S , t = T ∈ ℝ D β€² \mathbf{x}_{\text{CLS}}^{\ell=S,\;t=T}\in\mathbb{R}^{D^{\prime}} , extracted from 𝐗 β„“ = S , t = T \mathbf{X}^{\ell=S,\;t=T} , is utilized as the predictor of the current graph G G .

### D.2 Experimental Evaluation

As mentioned prior, eight datasets from TUDataset [ 36 ] are used for experimentation. Based on the collected results in Table 2 , we observed that the modified ET demonstrates the best performance across all datasets with the exception of MUTAG. Additionally, as a means to contrast ET with other graph transformers, ET is trained and evaluated with five datasets from GNNBenchmark [ 37 ] . Specifically, PATTERN and CLUSTER for node classification; CIFAR10 and MNIST for graph classification; and lastly, ZINC for graph-regression. We compare ET to the recently introduced graph transformers – namely, EGT [ 56 ] , Graphormer [ 57 ] , Exphormer [ 58 ] , GT [ 59 ] , and GAT [ 54 ] . The results are depicted in table 7 .

Dataset Method CIFAR10 MNIST PATTERN CLUSTER ZINC (Accuracy) ↑ \uparrow (Accuracy) ↑ \uparrow (Accuracy) ↑ \uparrow (Accuracy) ↑ \uparrow (MAE) ↓ \downarrow GT - - 84.808 Β± 0.068 84.808_{\pm 0.068} β–Ό \blacktriangledown ( 2.013 ) (2.013) 73.169 Β± 0.622 73.169_{\pm 0.622} β–Ό \blacktriangledown ( 6.063 ) (6.063) 0.226 Β± 0.014 0.226_{\pm 0.014} β–Ό \blacktriangledown ( 0.112 ) (0.112) GAT 64.223 Β± 0.455 64.223_{\pm 0.455} β–Ό \blacktriangledown ( 10.531 ) (10.531) 95.535 Β± 0.205 95.535_{\pm 0.205} β–Ό \blacktriangledown ( 2.879 ) (2.879) 78.271 Β± 0.186 78.271_{\pm 0.186} β–Ό \blacktriangledown ( 8.550 ) (8.550) 70.587 Β± 0.447 70.587_{\pm 0.447} β–Ό \blacktriangledown ( 8.645 ) (8.645) 0.384 Β± 0.007 0.384_{\pm 0.007} β–Ό \blacktriangledown ( 0.276 ) (0.276) EGT 68.702 Β± 0.409 68.702_{\pm 0.409} β–Ό \blacktriangledown ( 6.052 ) (6.052) 98.173 Β± 0.087 98.173_{\pm 0.087} β–Ό \blacktriangledown ( 0.241 ) (0.241) 86.821 Β± 0.020 86.821_{\pm 0.020} 79.232 Β± 0.348 \mathbf{79.232_{\pm 0.348}} 0.108 Β± 0.009 0.108_{\pm 0.009} Graphormer 65.978 Β± 0.579 65.978_{\pm 0.579} β–Ό \blacktriangledown ( 8.776 ) (8.776) 97.905 Β± 0.176 97.905_{\pm 0.176} β–Ό \blacktriangledown ( 0.509 ) (0.509) 86.650 Β± 0.033 86.650_{\pm 0.033} β–Ό \blacktriangledown ( 0.171 ) (0.171) 74.660 Β± 0.236 74.660_{\pm 0.236} β–Ό \blacktriangledown ( 4.572 ) (4.572) 0.122 Β± 0.006 0.122_{\pm 0.006} β–Ό \blacktriangledown ( 0.014 ) (0.014) Exphormer 74.754 Β± 0.194 \mathbf{74.754_{\pm 0.194}} 98.414 Β± 0.038 \mathbf{98.414_{\pm 0.038}} 86.734 Β± 0.008 86.734_{\pm 0.008} β–Ό \blacktriangledown ( 0.087 ) (0.087) 78.220 Β± 0.045 78.220_{\pm 0.045} β–Ό \blacktriangledown ( 1.012 ) (1.012) - ET (Ours) 63.920 Β± 0.173 63.920_{\pm 0.173} β–Ό \blacktriangledown ( 10.834 ) (10.834) 97.011 Β± 0.045 97.011_{\pm 0.045} β–Ό \blacktriangledown ( 1.403 ) (1.403) 90.690 Β± 0.013 \mathbf{90.690_{\pm 0.013}} β–² \blacktriangle ( 3.869 ) (3.869) 77.169 Β± 0.044 77.169_{\pm 0.044} β–Ό \blacktriangledown ( 2.063 ) (2.063) 0.096 Β± 0.000 \mathbf{0.096_{\pm 0.000}} β–² \blacktriangle ( 0.012 ) (0.012)

Based on the collected results, ET outperforms its baselines on PATTERN and ZINC datasets, and has competitive performance on MNIST and CLUSTER datasets. Meanwhile, ET performs the worst on the CIFAR10 dataset, in which its performance is relatively close to GAT [ 54 ] and Graphformer [ 57 ] . The reason for such a bad performance can be attributed to the overfitting problem seen during the training of the model, which is also experienced by [ 56 , 59 ] . Hence, while ET performs very well on TUDataset, further explorations are required to enable ET to work better on CIFAR10, MNIST, and CLUSTER datasets. However, ET has overall demonstrated great performance on node binary-classification and graph-regression tasks.

### D.3 Experimental Details