Energy Transformer — part 5/5

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

In the graph domain, it is common to concatenate all of the feature vectors of all graphs in a batch together. However, in order for ET to work, we form the batch dimension by separating the feature vectors of all graphs in a given batch and utilize the largest node count to pad all graphs such that they all share the same number of nodes. Additionally, we set a limit on the number of nodes equal to 500, to prevent out-of-memory error. Specifically, if a graph has a node count exceeding the limit, the number of utilized nodes is equal to the limit. Hence, a portion of the graph structural information is ignored as a result. However, it is worth mentioning such a graph is rare in the experimental datasets. Additionally, instead of ignoring the padded entries, they are altered to be sink nodes in a graph, such that they are connected to the actual nodes of the graph but not vice versa. We found this approach to be helpful when it comes to computing the positional embedding and training the model. Additionally, following [ 37 , 56 ] , the sign of the top k k eigen-vectors is flipped randomly during training and fixed during evaluation. The whole experiment is implemented using JAX [ 49 ] , Flax [ 50 ] , Optax [ 60 ] , and PyTorch Geometric [ 61 ] frameworks.

For the eight datasets of TUDataset, we train all models for 300 epochs using AdamW [ 62 ] and Gradient Centralization [ 63 ] . The best model is selected based on its performance obtained from the 10-fold cross validation process delineated in [ 36 ] . Since the task is classification, all models are trained with the cross-entropy loss function and label smoothing [ 64 ] , where the smoothing parameter is set to 0.05. Additionally, the cosine-annealing with warm-up learning rate scheduler [ 65 ] is utilized, where the initial and end learning rates are both set as 5 ​ e − 6 5e-6 while the peak learning rate is 0.001 0.001 . The number of warm-up steps is set to 50 epochs while the batch size is 32 for all datasets. We report the average performance of 100 runs on the 10-fold cross validation process with random seeding. The hyperparameters of our model are tuned based on the performance of the cross validation, selecting within 100 epochs. The optimal hyper-parameters are reported in Table 9 and the statistics of the used datasets are reported in Table 8 . We also include the number of gpu-devices used for training ET.

Similarly, for the five datasets of GNNBenchmark, we train all models using using AdamW [ 62 ] and Gradient Centralization [ 63 ] . The best model is selected based on its performance obtained from the validation set across 10 runs. For the CIFAR10 and MNIST datasets, the additional node positional features are concatenated to the node features (i.e., pixel values coordinates). All models, with the exception of ZINC, CIFAR10, and MNIST models, are trained using the cross-entropy loss function with label smoothing, where the smoothing parameter is set to 0.05. The ZINC model is trained with the mean absolute error function. Meanwhile, as a means to reduce overfitting, for the CIFAR10 and MNIST models, the mean squared error loss function is utilized to reconstruct the edge feature matrix ℰ \mathcal{E} in addition to the cross entropy loss function. Specifically, the parameterized adjacency tensor 𝐀 ^ \hat{\mathbf{A}} is projected to match P ′ P^{\prime} raw feature dimension of ℰ \mathcal{E} , where both tensors are as part of the mean squared error function. In addition, all models are trained with the cosine-annealing with warm-up learning rate scheduler, where the number of warm-up steps is 50. We report the average performance of 100 evaluation runs on the testing set with random seeding. The hyperparameters of each model are tuned based on the performance of cross validation, selecting within 50 epochs. The optimal hyper-parameters are reported in Tables 10 and 11 and the statistics of the used datasets are reported in [ 37 ] .

Dataset Parameter CIFAR10 MNIST PATTERN CLUSTER ZINC batch_size 128 128 128 128 128 epochs 150 150 100 300 500 initial lr 1e-3 1e-3 1e-3 1e-3 5e-7 ending lr 5e-5 5e-5 5e-7 5e-7 5e-7 peak lr 1e-3 1e-3 1e-3 1e-3 1e-3 warmup_epochs 50 50 50 50 50 b1, b2 (ADAM) (0.9, 0.99) (0.9, 0.99) (0.9, 0.99) (0.9, 0.99) (0.9, 0.99) weight_decay 0.05 0.05 0.05 0.05 0.05 grad_clipping None None None None None num. of gpu devices 4 4 1 1 8

Dataset Parameter CIFAR10 MNIST PATTERN CLUSTER ZINC token_dim 128 128 128 128 128 num_heads 12 12 12 12 12 head_dim 64 64 64 64 64 β \beta 1 64 \frac{1}{\sqrt{64}} 1 64 \frac{1}{\sqrt{64}} 1 64 \frac{1}{\sqrt{64}} 1 64 \frac{1}{\sqrt{64}} 1 64 \frac{1}{\sqrt{64}} train_betas Yes Yes Yes Yes Yes step size α \alpha 0.25 0.25 0.1 0.1 0.1 k eigenvalues 15 15 15 15 15 noise σ ϵ \sigma_{\epsilon} 0.02 0.02 0.005 0.005 0.005 depth 4 4 4 4 1 block_size 2 2 2 2 4 kernel_size [3, 3] [3, 3] [3, 3] [3, 3] [3, 3] dilation_size [1, 1] [1, 1] [1, 1] [1, 1] [1, 1] hidden_dim HN 512 512 512 512 512 bias in HN None None None None None bias in ATT None None None None None bias in LNORM Yes Yes Yes Yes Yes total num. of params 529224 528968 527936 528964 1052409

## Appendix E ET for heterogeneous Graph

In this section, we show the performance of our ET model in the heterogeneous graph case. For the heterogeneous graph case (where more than one type of edges exist in the graph), we first run our ET model for different subgraphs (corresponding to different types of edges), and then aggregate the final representations using two methods. We have tried max pooling and concatenation for the aggregation step. Max pooling means to pick the largest value across all the subgraph representations, and concatenation means to concatenate the representations obtained from different subgraphs. Table 12 shows the comparison between these two variants of our model and BWGNN (heterogeous case). ET performs better than heterogeneous BWGNN on Amazon, but loses to heterogeneous BWGNN on Yelp. Interestingly, BWGNN in the heterogeneous setting performs worse than BWGNN in the homogeneous setting on Amazon.

## Appendix F Ablation Study for Attention and Hopfield Network Modules

As we described in the main text the the ET network consists of two modules processing the tokens in parallel: the attention module (ATT) and the Hopfield Network module (HN). The ATT module is responsible for routing the information between the tokens, while the HN module is responsible for reinforcing the token representation to be consistent with the general expectation about the particular data domain. It is interesting to explore the contribution that these two subnetworks produce on the task performed by the network. In this section we ablate the ET architecture by dropping each of the two subnetworks and measuring the impact of the ablation on the performance.

### F.1 On Graphs

The results on graphs are reported in Table 13 . From this table it is clear that most of the computation is performed by the ATT block on this task, which pools the information about other tokens to the token of interest. When ATT block is kept, but HN block is removed the network looses 1 % 1\% or less relative to the full ET (occasional improvements of the ablated model compared to the full ET are within the statistical error bars). In contrast, removing ATT module and keeping only the HN, the ET network effectively turns into an MLP with shared weights that is recurrently applied. In this regime the network can only use the features of a given node for that node’s anomalous status prediction. This results in a more significant drop in performance, which is about 5 % 5\% on average.

### F.2 On Images

The ablation results for image reconstruction are shown in Table 14 . Each experiment was trained using the same hyperparameter settings as shown in Table 4 . After training the model on IN1K, we calculate the average MSE on the reconstructed masked tokens for the validation set (using the same 50 50 % masking ratio used for training) across 10 10 different random seeds for the masking.

We make several conclusions from these ablation studies.

• We gain several insights regarding the use of “self-attention” in our ET (when a token patch query is allowed to consider itself as a key in the attention weights). When both self-attention and HN are present (ET-Full+Self), there is no noticeable benefit over ET-Full for a token to attend to itself. In fact, preventing the ATTN energy module from attending to itself slightly improves the performance. However, when the HN is removed (ET-NoHN*), we notice that allowing self-attention (ET-NoHN+Self) outperforms the version that prevents self-attention (ET-NoHN).

• On its own, allowing self-attention (ET-NoHN+Self) in the ATTN module performs nearly as well as the full ET at a fraction of the total parameters. However, MSE is a forgiving metric for blurry reconstructions. While ATTN can capture the global structure of the image quite well, it does so at the expense of image sharpness ( Figure 4 ).

• As expected, removal of the ATTN energy module performs the worst, because the HN operates tokenwise and has no way to aggregate token information across the global image without ATTN.

Figure 8 shows our best performing model (ET-Full) on the qualitative image reconstructions corresponding to the largest errors across IN1K validation images, averaged across all masking seeds. Likewise, Figure 9 shows the lowest errors across IN1K validation images and masking seeds. In general, image reconstructions that require ET to produce sharp, high frequency, and high contrast lines negatively impact MSE performance.

## Appendix G Parameter comparison

The energy function enforces symmetries in our model, which means ET has fewer parameters than its ViT counterparts. In particular, ET has no “Value Matrix” 𝐖 V \mathbf{W}^{V} in the attention mechanism, and the HN module has only one of the two matrices in the standard MLP of the traditional Transformer block. We report these differences in Table 15 . We take the ET configuration used in this paper, which has an architecture fully comparable to the original ViT-base [ 18 ] with patch_size=16 , and report the number of parameters against ViT-base and an “ALBERT” version of ViT [ 14 ] where a single ViT block is shared across layers. We saw no benefit when including biases in ET, so we also exclude the biases from the total parameter count in the configuration of ViT and ALBERT-ViT. We report both the total number of parameters and the number of parameters per Transformer block. Furthermore, as a means for a fairer comparison, we reduce the parameters of ViT-Base such that the parameter count is similar to that of ET. We contrast ET and the reduced ViT on the image reconstruction task using PSNR (Peak-Signal-to-Noise-Ratio) and SSIM (Structural-SIMilarity) [ 66 ] as the evaluation metrics of the reconstructions. The results are demonstrated in Table 16 , where the performance of ET is very close to that of the reduced ViT, which has a slight advantage in parameter count.

## Appendix H Formal Algorithm for Training and Inference of ET

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .