Energy Transformer — part 3/5
Cleaned-text code points [29654, 44781). Bibliography omitted here; source.md includes it.
Dataset Method PROTEINS NCI1 NCI109 DD ENZYMES MUTAG MUTAGENICITY FRANKENSTEIN WKPI (kmeans) 78.5 ± 0.4 78.5_{\pm 0.4} ▼ \blacktriangledown ( 6.4 ) (6.4) 87.5 ± 0.5 87.5_{\pm 0.5} 85.9 ± 0.4 85.9_{\pm 0.4} ▼ \blacktriangledown ( 1.5 ) (1.5) 82.0 ± 0.5 82.0_{\pm 0.5} ▼ \blacktriangledown ( 13.7 ) (13.7) - 85.8 ± 2.5 85.8_{\pm 2.5} ▼ \blacktriangledown ( 14.2 ) (14.2) - - WKPI (kcenters) 75.2 ± 0.4 75.2_{\pm 0.4} ▼ \blacktriangledown ( 9.7 ) (9.7) 84.5 ± 0.5 84.5_{\pm 0.5} ▼ \blacktriangledown ( 3.0 ) (3.0) 87.4 ± 0.3 87.4_{\pm 0.3} 80.3 ± 0.4 80.3_{\pm 0.4} ▼ \blacktriangledown ( 15.4 ) (15.4) - 88.3 ± 2.6 88.3_{\pm 2.6} ▼ \blacktriangledown ( 11.7 ) (11.7) - - Spec-GN - 84.8 ± 1.6 84.8_{\pm 1.6} ▼ \blacktriangledown ( 2.7 ) (2.7) 83.6 ± 0.8 83.6_{\pm 0.8} ▼ \blacktriangledown ( 3.8 ) (3.8) - 72.5 ± 5.8 72.5_{\pm 5.8} ▼ \blacktriangledown ( 5.9 ) (5.9) - - - Norm-GN - 84.9 ± 1.7 84.9_{\pm 1.7} ▼ \blacktriangledown ( 2.6 ) (2.6) 83.5 ± 1.3 83.5_{\pm 1.3} ▼ \blacktriangledown ( 3.9 ) (3.9) - 73.3 ± 8.0 73.3_{\pm 8.0} ▼ \blacktriangledown ( 5.1 ) (5.1) - - - GWL-WL 75.8 ± 0.6 75.8_{\pm 0.6} ▼ \blacktriangledown ( 9.1 ) (9.1) - - - 71.3 ± 1.1 71.3_{\pm 1.1} ▼ \blacktriangledown ( 7.1 ) (7.1) - - 78.9 ± 0.3 {78.9_{\pm 0.3}} HGP-SL 84.9 ± 1.6 84.9_{\pm 1.6} 78.5 ± 0.8 78.5_{\pm 0.8} ▼ \blacktriangledown ( 9.1 ) (9.1) 80.7 ± 1.2 80.7_{\pm 1.2} ▼ \blacktriangledown ( 6.7 ) (6.7) 81.0 ± 1.3 81.0_{\pm 1.3} ▼ \blacktriangledown ( 14.7 ) (14.7) 68.8 ± 2.1 68.8_{\pm 2.1} ▼ \blacktriangledown ( 9.6 ) (9.6) - 82.2 ± 0.6 {82.2_{\pm 0.6}} - DSGCN 77.3 ± 0.4 {77.3_{\pm 0.4}} ▼ \blacktriangledown ( 7.6 ) (7.6) - - - 78.4 ± 0.6 {78.4_{\pm 0.6}} - - - U2GNN 80.0 ± 3.2 80.0_{\pm 3.2} ▼ \blacktriangledown ( 4.9 ) (4.9) - - 95.7 ± 1.9 95.7_{\pm 1.9} - 88.5 ± 7.1 88.5_{\pm 7.1} ▼ \blacktriangledown ( 11.5 ) (11.5) - - NDP 73.4 ± 3.1 73.4_{\pm 3.1} ▼ \blacktriangledown ( 11.5 ) (11.5) 74.2 ± 1.7 74.2_{\pm 1.7} ▼ \blacktriangledown ( 13.3 ) (13.3) - 72.8 ± 5.4 72.8_{\pm 5.4} ▼ \blacktriangledown ( 22.9 ) (22.9) 44.5 ± 7.4 44.5_{\pm 7.4} ▼ \blacktriangledown ( 34.9 ) (34.9) 87.9 ± 5.7 87.9_{\pm 5.7} ▼ \blacktriangledown ( 12.1 ) (12.1) 77.9 ± 1.4 77.9_{\pm 1.4} ▼ \blacktriangledown ( 4.3 ) (4.3) - ASAP 74.2 ± 0.8 74.2_{\pm 0.8} ▼ \blacktriangledown ( 10.7 ) (10.7) 71.5 ± 0.4 71.5_{\pm 0.4} ▼ \blacktriangledown ( 16.0 ) (16.0) 70.1 ± 0.6 70.1_{\pm 0.6} ▼ \blacktriangledown ( 17.3 ) (17.3) 76.9 ± 0.7 76.9_{\pm 0.7} ▼ \blacktriangledown ( 18.8 ) (18.8) - - - 66.3 ± 0.5 66.3_{\pm 0.5} ▼ \blacktriangledown ( 12.6 ) (12.6) EvoG - - - - 55.7 55.7 ▼ \blacktriangledown ( 22.7 ) (22.7) 100.0 \mathbf{100.0} - - ET (Ours) 90.3 ± 0.7 \mathbf{90.3_{\pm 0.7}} ▲ \blacktriangle (5.4) 90.1 ± 0.1 \mathbf{90.1_{\pm 0.1}} ▲ \blacktriangle (2.6) 90.5 ± 0.1 \mathbf{90.5_{\pm 0.1}} ▲ \blacktriangle (3.1) 95.9 ± 0.8 \mathbf{95.9_{\pm 0.8}} ▲ \blacktriangle (0.2) 99.8 \mathbf{99.8} ▲ \blacktriangle (21.4) 96.6 ± 0.2 96.6_{\pm 0.2} ▼ \blacktriangledown (3.4) 98.7 ± 0.1 \mathbf{98.7_{\pm 0.1}} ▲ \blacktriangle (16.5) 99.8 ± 0.1 \mathbf{99.8_{\pm 0.1}} ▲ \blacktriangle (20.9)
## 5 Graph Classification with ET
To fully explore the efficacy of ET, we further evaluate its performance on the graph classification problem * * * The code is available: https://github.com/Lemon-cmd/Energy-Transformer-For-Graph . . 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 G , where each node has a raw feature vector 𝐲 A ∈ ℝ F \mathbf{y}_{A}\in\mathbb{R}^{F} with F F feature-dimension. Each vector is projected to the token space yielding 𝐱 A ∈ ℝ D \mathbf{x}_{A}\in\mathbb{R}^{D} , where D D is the token dimension. A learnable CLS token 𝐱 CLS \mathbf{x}_{\text{CLS}} is concatenated to the set of tokens resulting in the token matrix 𝐗 ∈ ℝ ( N + 1 ) × D \mathbf{X}\in\mathbb{R}^{(N+1)\times D} , and a learnable linear projection of the top k 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 S ET blocks, where each block has the same number of temporal unfolding steps T T and its own LayerNorm. The final representation of the CLS token 𝐱 CLS ℓ = S , t = T \mathbf{x}^{\ell=S,\;t=T}_{\mathrm{CLS}} 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 ).
## Appendix A Notations Used in the Main Text and Appendices
Table 3 lists all the notations used in this paper.
## Appendix 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 × 224 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 D for a total of 196 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 100 tokens, 90 90 tokens are replaced with a learnable MASK token of dimension D D and 10 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 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 α = d t τ \alpha=\frac{dt}{\tau} which we treat as a non-learnable hyperparameter in our experiments. The token representations after T 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 100 occluded patches. We allow self attention as in the following formula for the energy of multiheaded attention E = ATT ∑ h − 1 β ∑ C log ( ∑ B exp ( β ∑ α K α h B Q α h C ) ) E{{}^{\text{ATT}}}=\sum\limits_{h}-\frac{1}{\beta}\sum\limits_{C}\textrm{log}\left(\sum\limits_{B}\textrm{exp}\left(\beta\sum\limits_{\alpha}K_{\alpha hB}\;Q_{\alpha hC}\right)\right) (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 768 , 12 12 heads each with Y = 64 Y=64 , softmax ’s β = 1 Y \beta=\frac{1}{\sqrt{Y}} , …), with several considerations learned from the qualitative image evaluations:
• The d t τ \frac{dt}{\tau} (step size) of 1 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 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.
### B.1 Exploring the Hopfield Memories
A distinctive aspect of our network is that any variable that has a vector index i i of tokens can be mapped into the data domain by applying the decoder network to this variable. This makes it possible to inspect all the weights in the model. For instance, the concept of “memories” is crucial to understanding how Hopfield networks function. The memories within the HN module represent the building blocks of all possible image patches in our data domain, where an encoded image patch is a superposition of a subset of memories. The complete set of memory vectors from the HN module is shown in Figure 5 . The same analysis can be applied to the weights of the ET-attention module. In Figure 6 , we show all the weights from this module mapped into the image plane.
### B.2 Positional Bias Correlations
The relationships between our position embeddings exhibit similar behavior to the position correlations of the original ViT in that they are highly susceptible to choices of the hyperparameters (Figure 10 of [ 18 ] ). In particular, we consider the effect of weight decay and the β \beta parameter that serves as the inverse temperature of the attention operation (see Equation 3 ). The lower the temperature (i.e., the higher the value of β \beta ), the spikier the softmax distribution. By using a lower β \beta , we encourage the attention energy to distribute its positional embeddings across a wider range of patches in the model.
## Appendix C Details of ET training on Anomaly Detection Task
Graph anomaly detection refers to the process of detecting outliers that deviate significantly from the majority of the samples. Neural network based methods are very popular due to their capability of learning sophisticated data representations. DOMINANT [ 26 ] utilizes an auto-encoder framework, using a GCN as an encoder and two decoders for structural reconstruction and attribute reconstruction. ALARM [ 27 ] aggregates the encoder information from multiple view of the node attributes. Another study [ 51 ] , propose a novel loss function to train graph neural networks for anomaly-detectable node representations. In [ 52 ] generative adversarial learning is used to detect anomaly nodes where a novel layer is designed to learn the anomaly-aware node representation. Recently, [ 24 ] pointed out that anomalies can lead to the “rightshift” of the spectral energy distribution – the spectral energy concentrates more on the high frequencies. They designed a filter that can better handle this phenomenon. We propose a new anomaly detection model from the perspective of Associative Memory (pattern matching), which does not have the over-smoothing problem often faced by GCNs, and has better model interpretability (outliers should be far from the common pattern). We also notice that Modern Hopfield Networks have been used before for node classification, link prediction, and graph coarsening tasks [ 53 ] .
### C.1 Detailed Model Structure for the Graph Anomaly Detection