SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Examples of occluded/reconstructed images (unseen during training) are shown in Figure 3. In general, our model learns to perform the task very well, capturing the texture in dog fur (column 3) and understanding meaningful boundaries of objects. However, we observe that our single ET block struggles to understand some global structure, e.g., failing to capture both eyes of the white dog (column 5) and completing irregular brick patterns in the name of extending the un-occluded borders (last column). We additionally inspect the positional encoding vectors associated with every token, Figure 2, where the model learns a locality structure in the image plane that is very similar to the original ViT [18]. The positional embedding of each image patch has learned higher similarity values for neighboring tokens than for distant tokens.
Our network is unique compared to standard ViTs in that the iterative dynamics only move tokens around in the same space from which the final fixed point representation can be decoded back into the image plane. This functionality makes it possible to visualize essentially any token representation, weight, or gradient of the energy directly in the image plane. This feature is highly desirable from the perspective of interpretability, since it makes it possible to track the updates performed by the network directly in the image plane as the computation unfolds in time. In Figure 2 this functionality is used for inspecting the learned weights of the HN module directly in the image plane. According to our theory, these weights should represent basis vectors in the space of all possible image patches. These learned representations look qualitatively similar to the representations typically found in networks trained on image datasets, e.g., [23].
Figure 4: Token representations and gradients are visualized using the decoder at different times during the dynamics. The Energy Attention (ATTN) block contributes general structure information to the masked patches at earlier time steps, whereas the Hopfield Network (HN) significantly sharpens the quality of the masked patches at later time steps.
Table 1: Performance on Yelp, Amazon, T-Finance, and T-Social datasets with different training ratios. Following [24], mean and standard deviation over 5 runs with different train/dev/test split are reported (standard deviations are only included if they are available in the prior work). Best results are in bold. Our model is state of the art or near state of the art on every category.
Datasets	Split	GraphConsis	CAREGNN	PC-GNN	BWGNN	MLP	GT	ET (Ours)
Macro-F1	Yelp	1%	56.8±2.8	62.1±1.3	59.8±1.4	61.1±0.4	53.9±0.2	61.7±0.4	63.0±0.6
40%	58.7±2.0	63.3±0.9	63.0±2.3	71.0±0.9	57.5±0.8	68.7±0.4	71.5±0.1
Amazon	1%	68.5±3.4	68.7±1.6	79.8±5.6	90.9±0.7	74.6±1.2	88.6±0.5	89.3±0.7
40%	75.1±3.2	86.3±1.7	89.5±0.7	92.2±0.4	79.1±1.2	91.7±0.8	92.8±0.3
T-Finance	1%	71.7	73.3	62.0	84.8	61.0	81.5	85.1±1.0
40%	73.4	77.5	63.1	86.8	70.5	83.6	88.2±1.0
T-Social	1%	52.4	55.8	51.1	75.9	50.0	64.3	79.1±0.7
40%	56.5	56.2	52.1	83.9	50.3	68.2	83.5±0.4
AUC	Yelp	1%	66.4±3.4	75.0±3.8	75.4±0.9	72.0±0.5	59.8±0.4	72.5±0.6	73.2±0.8
40%	69.8±3.0	76.1±2.9	79.8±0.1	84.0±0.9	66.5±1.0	81.9±0.5	84.9±0.3
Amazon	1%	74.1±3.5	88.6±3.5	90.4±2.0	89.4±0.3	83.6±1.7	89.0±1.2	91.9±1.0
40%	87.4±3.3	90.5±1.6	95.8±0.1	98.0±0.4	89.8±1.0	95.4±0.6	97.3±0.4
T-Finance	1%	90.2	90.5	90.7	91.1	82.9	90.0	92.8±1.1
40%	91.4	92.1	91.2	94.3	87.1	88.2	95.0±3.0
T-Social	1%	65.2	71.2	59.8	88.0	56.3	81.4	91.9±0.6
40%	71.2	71.8	68.4	95.2	56.9	82.5	93.9±0.2
We additionally visualize the gradients of the energy function (which are equal to the token updates, see Equation 6) of both ATTN block and the HN block, see Figure 4. Early in time, almost all signal to the masked tokens comes from the ATTN block, which routes information from the open patches to the masked ones; no meaningful signal comes from the HN block to the masked patch dynamics. Later in time we observe a different phenomenon: almost all signal to masked tokens comes from the HN module while ATTN contributes a blurry and uninformative signal. Thus, the attention module is crucial early in the network dynamics, feeding signal to masked patches from the visible patches, whereas the HN is crucial later in the dynamics as the model approaches the final reconstruction, sharpening the masked patches. All the qualitative findings presented in this section are in accord with the core computational strategy of the ET block as it was designed theoretically in section 2.
4 Graph Anomaly Detection
Having gained empirical insights in the image domain, we turn to quantitatively evaluating ET's performance on the graph anomaly detection problem*, a task with plenty of strong and recently published baselines. Graph Convolutional Networks (GCN) [25] have been widely used for this task due to their capability of learning high level representations of graph structures and node attributes [26, 27]. However, vanilla GCNs suffer from the over-smoothing problem [28]. In each layer of the forward pass, the outlier node aggregates information from its neighbors. This averaging makes the features of anomalies less distinguishable from the features of benign or normal nodes. Our approach does not suffer from this problem, since the routing of the information between the nodes is done through the energy based attention, and the information aggregation is based on the attention scores.
For anomaly detection on graphs in the ET framework, consider an undirected graph with  $N$  nodes. Every node has a vector of raw attributes  $\mathbf{y}_A \in \mathbb{R}^F$ , where  $F$  is the number of node's features. Every node also has a binary label  $l_A$  indicating whether the node is benign or anomalous. We focus on node anomaly and assume that all edges are trusted. The task is to predict the label of a node given the graph structure and the node's features. Since there are far more benign nodes in the graph than anomalous ones, anomaly detection can be regarded as an imbalanced node classification task.
First, similarly to images, the feature vectors for every node are converted to a token representation using a linear embedding  $\mathbf{E}$  and adding a learnable positional embedding  $\lambda_{A}$
$$
\mathbf {x} _ {A} ^ {t = 1} = \mathbf {E} \mathbf {y} _ {A} + \lambda_ {A} \tag {10}
$$
where the superscript  $t = 1$  indicates the time of the update of the ET dynamics. This token representation is iterated through the ET block for  $T$  iterations. When the retrieval dynamics becomes
*The code is available: https://github.com/zhuergou/Energy-Transformer-for-Graph-Anomaly-Detection/.
Table 2: Graph classification performance on eight datasets of TUDataset. Following [36], mean and standard deviation obtained from 100 runs of 10-fold cross validation are reported. For baselines standard deviations are only included if they are available in prior work. If the entry is unavailable in prior literature it is denoted by  $-$ ; best results are in bold. The performance difference between non-baseline approaches (including ours) and the baselines (specified by their gray cell) is indicated by  $\nabla$  (decrease),  $\triangle$  (increase), and  $\nabla$  (no change within the error bars) along with the value.
Method	Dataset
PROTEINS	NCI1	NCI109	DD	ENZYMES	MUTAG	MUTAGENICITY	FRANKENSTEIN
WKPI (kmeans)	78.5±0.4▼(6.4)	87.5±0.5	85.9±0.4▼(1.5)	82.0±0.5▼(13.7)	-	85.8±2.5▼(14.2)	-	-
WKPI (kcenters)	75.2±0.4▼(9.7)	84.5±0.5▼(3.0)	87.4±0.3	80.3±0.4▼(15.4)	-	88.3±2.6▼(11.7)	-	-
Spec-GN	-	84.8±1.6▼(2.7)	83.6±0.8▼(3.8)	-	72.5±5.8▼(5.9)	-	-	-
Norm-GN	-	84.9±1.7▼(2.6)	83.5±1.3▼(3.9)	-	73.3±8.0▼(5.1)	-	-	-
GWL-WL	75.8±0.6▼(9.1)	-	-	-	71.3±1.1▼(7.1)	-	-	78.9±0.3
HGP-SL	84.9±1.6	78.5±0.8▼(9.1)	80.7±1.2▼(6.7)	81.0±1.3▼(14.7)	68.8±2.1▼(9.6)	-	82.2±0.6	-
DSGCN	77.3±0.4▼(7.6)	-	-	-	78.4±0.6	-	-	-
U2GNN	80.0±3.2▼(4.9)	-	-	95.7±1.9	-	88.5±7.1▼(11.5)	-	-
NDP	73.4±3.1▼(11.5)	74.2±1.7▼(13.3)	-	72.8±5.4▼(22.9)	44.5±7.4▼(34.9)	87.9±5.7▼(12.1)	77.9±1.4▼(4.3)	-
ASAP	74.2±0.8▼(10.7)	71.5±0.4▼(16.0)	70.1±0.6▼(17.3)	76.9±0.7▼(18.8)	-	-	-	66.3±0.5▼(12.6)
EvoG	-	-	-	-	55.7▼(22.7)	100.0	-	-
ET (Ours)	90.3±0.7▲(5.4)	90.1±0.1▲(2.6)	90.5±0.1▲(3.1)	95.9±0.8▲(0.2)	99.8▲(21.4)	96.6±0.2▼(3.4)	98.7±0.1▲(16.5)	99.8±0.1▲(20.9)
stable, we have the final representation for each node  $\mathbf{x}_A^{t = T}$  (or more precisely  $\mathbf{g}_A^{t = T}$ , since the outputs are additionally passed through a layer norm operation after the final ET update). This output is concatenated with the initial (layer normalized) token to form the final output
$$
\mathbf {g} _ {A} ^ {\text {f i n a l}} = \mathbf {g} _ {A} ^ {t = 1} \mid \mid \mathbf {g} _ {A} ^ {t = T} \tag {11}
$$
where  $||$  denotes concatenation. Following [24], the node representation  $\mathbf{g}_A^{\mathrm{final}}$  is fed into an MLP with a sigmoid activation function to compute the anomaly probabilities  $p_A$ . The weighted cross entropy
$$
\operatorname {L o s s} = \sum_ {A} \left[ \omega l _ {A} \log \left(p _ {A}\right) + \left(1 - l _ {A}\right) \log \left(1 - p _ {A}\right) \right] \tag {12}
$$
is used to train the network. Above,  $\omega$  is the ratio of the benign labels  $(l_{A} = 0)$  to anomalous ones  $(l_{A} = 1)$ . Attention is restricted to 1-hop neighbors of the target node.
4.1 Experimental Evaluation
Four datasets are used for the experiments. YelpChi dataset [29] aims at opinion spam detection in Yelp reviews. Amazon dataset is used to detect anomalous users under the Musical Instrument Category on amazon.com [30]. T-Finance and T-Social datasets [24] are used for anomalous account detection in transactions and social networks, respectively. For these four datasets, the graph is treated as a homogeneous graph (i.e., all the edges are of the same type), and a feature vector is associated with each node. The task is to predict the label (anomaly status) of the nodes. For each dataset, either  $1\%$  or  $40\%$  of the nodes are used for training, and the remaining  $99\%$  or  $60\%$  are split 1:2 into validation and testing sets, see Appendix C for details.
We compare with state-of-the-art approaches for graph anomaly detection, which include GraphConsis [31], CAREGNN [32], PC-GNN [33] and BWGNN [24]. Additionally, multi-layer perceptrons (MLP) and Graph Transformer (GT) [34] are included in the baselines for completeness. Following previous work, macro-F1 score (unweighted mean of F1 score) and the Area Under the Curve (AUC) are used as the evaluation metrics on the test datasets [35]. See Appendix C for more details on training protocols and the hyperparameters choices. The results are reported in Table 1. Our ET network demonstrates very strong results across all the datasets.
5 Graph Classification with ET
To fully explore the efficacy of ET, we further evaluate its performance on the graph classification problem*. Unlike the anomaly detection task, where we predict a label for every node, in the graph classification task, a single label is predicted for the entire graph.
Consider a graph  $G$ , where each node has a raw feature vector  $\mathbf{y}_A \in \mathbb{R}^F$  with  $F$  feature-dimension. Each vector is projected to the token space yielding  $\mathbf{x}_A \in \mathbb{R}^D$ , where  $D$  is the token dimension. A learnable CLS token  $\mathbf{x}_{\mathrm{CLS}}$  is concatenated to the set of tokens resulting in the token matrix  $\mathbf{X} \in \mathbb{R}^{(N + 1) \times D}$ , and a learnable linear projection of the top  $k$  smallest eigen-vectors of the normalized Laplacian matrix, following [37], is added to  $\mathbf{X}$ . Graph structural information is provided to ET within the attention operation. We also use a stacked version of the Energy Transformer consisting of  $S$  ET blocks, where each block has the same number of temporal unfolding steps  $T$  and its own LayerNorm. The final representation of the CLS token  $\mathbf{x}_{\mathrm{CLS}}^{\ell = S, t = T}$  is projected via a linear embedding into the predictor space. A softmax over the number of classes is applied to this predictor representation, and the cross-entropy loss is used to train the network. See Appendix D for full details.
5.1 Experimental Evaluation
Eight graph datasets from the TUDataset [36] collection are used for experimentation. NCI1, NCI109, MUTAG, MUTAGENICITY, and FRANKENSTEIN are a common class of graph datasets consisting of small molecules with class labels representing toxicity or biological activity determined in drug discovery projects [36]. Meanwhile, DD, ENZYMES, and PROTEINS represent macromolecules. The task for both DD and PROTEINS is to classify whether a protein is an enzyme. Lastly, for ENZYMES, the task is to assign enzymes to one of the six classes, which reflect the catalyzed chemical reaction [36]. See Table 8 for more details of the datasets.
We compare ET with the current state-of-the-art approaches for the mentioned datasets, which include WKPI-kmeans [38], WKPI-kcenters [38], DSGCN [39], HGP-SL [40], U2GNN [41], and EvoG [42]. Additionally, approaches [39, 43, 44, 45, 46], which are close to the baselines, are included to further contrast the performance of our model. Following the 10-fold cross validation process delineated in [36], accuracy score is used as the evaluation metric and reported in Table 2. In general, we have observed that the modified ET demonstrates strong performance across the eight datasets. With the exception of MUTAG, ET beats other methods on all the baselines by a substantial margin.
6 Discussion and Conclusions
A lot of recent research has been dedicated to understanding the striking analogy between Hopfield Networks and the attention mechanism in transformers. At a high level, the main message of our work is that the entire transformer block (including feed-forward MLP, layer normalization, and residual connections) can be viewed as a single large Hopfield Network, not just attention alone. At a deeper level, we use recent advances in the field of Hopfield Networks to design a novel energy function that is tailored for dynamical information routing between the tokens; and representation of a large number of relationships between those tokens. We have tested the ET network qualitatively on the image completion task, and quantitatively on node anomaly detection and graphs classification. The qualitative investigation reveals the perfect alignment between the theoretical design principles of our network and its empirical computation. The quantitative evaluation demonstrates strong results, which stand in line or exceed the methods recently developed specifically for these tasks. We believe that the proposed network will be useful for other tasks and domains (e.g., NLP, audio, and video), which serve as future directions for a comprehensive investigation.
There are two metrics describing computational costs of our architecture: memory footprint of the model and the number of flops. In terms of the memory footprint our model wins (is smaller) compared to feedforward transformers with independent weights and even ALBERT-style shared-weight transformers (see Appendix G and Table 15). Regarding the number of flops, our model has the same parametric scaling as the conventional transformers (quadratic in the number of tokens), but the energy attention has a constant factor of two more flops (due to the second term in the attention-induced updates of the tokens).
Finally, we have developed an auto-grad framework called HAMUX [47] that makes building flexible energy-based architectures, including ET, convenient and efficient (see HAMUX).
*The code is available: https://github.com/Lemon-cmd/Energy-Transformer-For-Graph.
[53] Yuchen Liang, Dmitry Krotov, and Mohammed J Zaki. Modern hopfield networks for graph embedding. Frontiers in big Data, 5, 2022.
[54] Petar Velicković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017.
[55] Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform bad for graph representation? CoRR, abs/2106.05234, 2021.
[56] Md Shamim Hussain, Mohammed J. Zaki, and Dharmashankar Subramanian. Global self-attention as a replacement for graph convolution. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '22, page 655–665, New York, NY, USA, 2022. Association for Computing Machinery.
[57] Junhan Yang, Zheng Liu, Shitao Xiao, Chaozhuo Li, Defu Lian, Sanjay Agrawal, Amit Singh, Guangzhong Sun, and Xing Xie. Graphformers: Gnn-nested transformers for representation learning on textual graph. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 28798-28810. Curran Associates, Inc., 2021.
[58] Hamed Shirzad, Ameya Velingker, Balaji Venkatachalam, Danica J Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. arXiv preprint arXiv:2303.06147, 2023.
[59] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699, 2020.
[60] Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020.
[61] Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
[62] Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. CoRR, abs/1711.05101, 2017.
[63] Hongwei Yong, Jianqiang Huang, Xiansheng Hua, and Lei Zhang. Gradient-centralization: A new optimization technique for deep neural networks. 2020.
[64] Rafael Müller, Simon Kornblith, and Geoffrey E Hinton. When does label smoothing help? Advances in neural information processing systems, 32, 2019.
[65] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016.
[66] Alain Hore and Djemel Ziou. Image quality metrics: Psnr vs. sSIM. In 2010 20th International Conference on Pattern Recognition, pages 2366-2369, 2010.
A Notations Used in the Main Text and Appendices
Table 3 lists all the notations used in this paper.
Table 3: Notations used in the paper.
Notation	Description
F	dimension of node's feature space
D	dimension of token space
N	number of tokens
Y	number of hidden dimensions in the attention
K	number of hidden dimensions in the Hopfield Network module
T	number of recurrent time steps
H	number of heads
kh	height of each image patch
kw	width of each image patch
P	number of pixels per image patch (3 × kh × kw)
yA	input feature vector of node A
xA	vector representation of token A
xiA	each element of vector representation of token A
gA	vector representation of token A after layernorm
giA	each element of vector representation of token A after layernorm
K	key tensor
Q	query tensor
KαhB	each element of the key tensor K
QαhC	each element of the query tensor Q
AhBC	each element of the attention tensor
lA	label of node A on graph
r(·)	activation function of the Hopfield Network module
G(·)	energy contribution of the Hopfield Network module
B Details of Training on ImageNet
We trained the ET network on a masked-image completion task on the ImageNet-1k (IN1K) dataset. We treat all images in IN1K as images of shape  $224 \times 224$  that are normalized according to standard IN1K practices (mean 0, variance 1 on the channel dimension) and use data augmentations provided by the popular timm library [48] (See Table 4). Following the conventional ViT pipeline [18], we split these images into non-overlapping patches of 16x16 RGB pixels which are then projected with a single affine encoder into the token dimension  $D$  for a total of 196 encoded tokens per image. We proceed to randomly and uniformly assign 100 of these tokens as "occluded" which are the only tokens considered by the loss function. "Occluded" tokens are designated as follows: of the 100 tokens, 90 tokens are replaced with a learnable MASK token of dimension  $D$  and 10 we leave untouched (which we find important for the HN to learn meaningful patch representations). To all tokens we then add a distinct learnable position bias.
These tokens are then passed to our Energy Transformer block which we recur for  $T$  steps (the "depth" of the model in conventional Transformers). At each step, the feedback signal (the sum of the energy gradients from our attention block and HN block) is subtracted from our original token representation with a scalar step size  $\alpha = \frac{dt}{T}$  which we treat as a non-learnable hyperparameter in our experiments. The token representations after  $T$  steps are passed to a simple linear decoder (consisting of a layer norm and an affine transformation) to project our representations back into the image plane. We then use the standard MSE Loss between the original pixels and reconstructed pixels for only the 100 occluded patches. We allow self attention as in the following formula for the energy of multiheaded attention
$$
E ^ {\mathrm {A T T}} = \sum_ {h} - \frac {1}{\beta} \sum_ {C} \log \left(\sum_ {B} \exp \left(\beta \sum_ {\alpha} K _ {\alpha h B} Q _ {\alpha h C}\right)\right) \tag {13}
$$
We give details of our architectural choices in Table 4. In the main paper we present our Energy Transformer with a configuration similar to the standard base Transformer configuration (e.g., token dimension 768, 12 heads each with  $Y = 64$ , softmax's  $\beta = \frac{1}{\sqrt{Y}}$ , ...) with several considerations learned from the qualitative image evaluations:
- The  $\frac{dt}{\tau}$  (step size) of 1 implicitly used in the traditional transformer noticeably degrades our ability to smoothly descend the energy function. We find that a step size of 0.1 provides a smoother descent down the energy function and benefits the image reconstruction quality.
- We observe that our MSE loss must include some subset of un-occluded patches in order for the HN to learn meaningful filters.
- Values of  $\beta$  in the energy attention that are too high prevent our model from training. This is possibly due to vanishing gradients in our attention operation from a softmax operation that is too spiky.
- Without gradient clipping, our model fails to train at the learning rates we tried higher than 1e-4. We observe that gradient clipping not only helps our model train faster at the trainable learning rates, it also allows us to train at higher learning rates.
Our architecture and experiments for the image reconstruction task were written in JAX [49] using Flax [50]. This engineering choice means that our architecture definitions are quite lightweight, as we can define the desired energy function of the ET and use JAX's autograd to automatically calculate the desired update.
Table 4: Hyperparameter, architecture, and data augmentation choices for ET model during ImageNet-1k masked training experiments. Data augmentations are listed as parameters passed to the equivalent timm datalloader functionality.
Training	Architecture	Data Augmentation
batch_size	768	token_dim	768	random_lease	None
epochs	100	num_heads	12	horizontalFlip	0.5
lr	5e-4	head_dim	64	verticalFlip	0
warmup_epochs	2	β	1/8	color_jitter	0.4
start & end lr	5e-7	train_betas	No	scale	(0.08, 1)
b1, b2 (ADAM)	0.9, 0.99	step size α	0.1	ratio	(3/4, 4/3)
weight_Decay	0.05	depth	12	auto_augment	None
grad_clipping	1.	hidden_dim HN	3072
bias in HN	None
bias in ATT	None
bias in LNORM	Yes
B.1 Exploring the Hopfield Memories