Energy Transformer
Authors: Benjamin Hoover
Yuchen Liang
Bao Pham
Rameswar Panda
Hendrik Strobelt
Duen Horng Chau
Mohammed J Zaki
Dmitry Krotov
Venue: NeurIPS 2023 poster (2023)
Keywords: Hopfield Network
Dense Associative Memory
Energy-based models
Attention Mechanism
TL;DR: We propose a network, which describes the forward pass in a transformer as a gradient decent on an energy function.

Abstract: Our work combines aspects of three promising paradigms in machine learning, namely, attention mechanism, energy-based models, and associative memory. Attention is the power-house driving modern deep learning successes, but it lacks clear theoretical foundations. Energy-based models allow a principled approach to discriminative and generative tasks, but the design of the energy functional is not straightforward. At the same time, Dense Associative Memory models or Modern Hopfield Networks have a well-established theoretical foundation, and allow an intuitive design of the energy function. We propose a novel architecture, called the Energy Transformer (or ET for short), that uses a sequence of attention layers that are purposely designed to minimize a specifically engineered energy function, which is responsible for representing the relationships between the tokens. In this work, we introduce the theoretical foundations of ET, explore its empirical capabilities using the image completion task, and obtain strong quantitative results on the graph anomaly detection and graph classification tasks.

Energy Transformer
Benjamin Hoover"
Yuchen Liang' Department of CS RPI liangy7@rpi. edu
Bao Pham" Department of CS RPI phamborpi.cm .ddu edu
IBM Research
Georgia Tech
benj jmmn. aaaaaminoooo www.aaamoe...ee.e.eom www.ooooorriie. wwweerr@bbm.c.. com
Rameswar Panda
Hendrik Strobelt
Duen Horng Chau College of Computing Georgia Tech
MIT-IBM Watson AI Lab IBM Research
MIT-IBM Watson Al Lab
IBM Research
rpanda@ibig bm. com
polo@gatech. edu
Mohammed J. Zaki
Dmitry Krotov
Department of CS RPI zaki@cs edu
MIT-IBM Watson AI Lab
IBM Research
krotov@ibm.cm com
Abstract
Our work combines aspects of three promising paradigms in machine learning, namely, attention mechanism, energy-based models, and associative memory. At- tention is the power-house driving modern deep learning successes, but it lacks clear theoretical foundations. Energy-based models allow a principled approach to discriminative and generative tasks, but the design of the energy functional is not straightforward. At the same time, Dense Associative Memory models or Modern Hopfield Networks have well-established theoretical foundation, and allow intuitive design of the energy function. We propose a novel architecture, called the Energy Transformer (or ET for short), that uses sequence of attention layers that are purposely designed to minimize specifically engineered energy function, which is responsible for representing the relationships between the tokens. In this work, we introduce the theoretical foundations of ET, explore its empirical capabilities using the image completion task, and obtain strong quantitative results on the graph anomaly detection and graph classification tasks.
Introduction
Transformers have become pervasive models in various domains of machine learning, including language, vision, and audio processing. Every transformer block uses four fundamental operations: attention, feed-forward multi-layer perceptron (MLP), residual connection, and layer normalization. Different variations transformers result from combining these four operations in various ways. For instance, [1] proposes to frontload additional attention operations and backload additional MLP layers in a sandwich-like manner instead of interleaving them, [2] prepends an MLP layer before the attention in each transformer block, [3] uses neural architecture search methods to evolve even more sophisticated transformer blocks, and so on. Various methods exist to approximate the attention operation, multiple modifications of the norm operation, and connectivity of the block; see, for example, 4 for taxonomy of different models. At present, however, the search for new transformer
Hoover, Y.Liang, and B. Pham equally contributed to this work.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

architectures is driven mostly by empirical evaluations, and the theoretical principles behind this growing list of architectural variations missing.
Additionally, the computational role of the four elements remains the subject of discussions. Originally, [5] emphasized attention the most important part of the transformer block, arguing that the learnable long-range dependencies are more powerful than the local inductive biases of convolutional networks. On the other hand more recent investigations [6] argue that the entire transformer block is important. The "correct" way to combine the four basic operations inside the block remains unclear, as does an understanding of the core computational function of the entire block and each of its four elements.
On the other hand, Associative Memory models, also known as Hopfield Networks 7 1,, have been gaining popularity in the machine learning community thanks to theoretical advancements pertaining to their memory storage capacity and novel architectural modifications ( Specifically, it has been shown that increasing the sharpness of the activation functions can lead to super-linear [1]] and even exponential [I]] memory storage capacity for these models, which is important for machine learning applications. This new class of Hopfield Networks is called Dense Associative Memories or Modern Hopfield Networks. Study 12] additionally describes how the attention mechanism in transformers is closely related to a special model of this family with the softmax activation function.
There are high-level conceptual similarities between transformers and Dense Associative Memories, since both architectures are designed for some form of denoising of the input. Transformers are typically pre-trained on a masked-token task, e.g., in the domain of Natural Language Processing (NLP), certain tokens in the sentence are masked and the model predicts the masked tokens. Dense Associative Memory models are designed for completing the incomplete patterns. They can be trained in self-supervised way by predicting the masked parts of the pattern, or denoising the pattern.
There are also high-level differences between the two approaches. Associative Memories are recurrent networks with a global energy function so that the network dynamics converges to a fixed point attractor state corresponding to a local minimum of the energy function. Transformers are typically not described as dynamical systems at all. Rather, they are thought of as feed forward networks built from the four computational elements discussed above. Even if one thinks about them as dynamical systems with tied weights, e.g., 13]. there is no reason to expect that their dynamics converge to a fixed point attractor (see the discussion in [14
Additionally, a recent study [15] uses form of Majorization- Majorization-MMiiinzztooo algorithms [6] to interpret the forward path in the transformer block as optimization process. This interpretation requires imposing certain constraints on the operations inside the block, and attempting to find an energy function that describes the constrained block. We take a complementrry approach by using the intuition developed in Associative Memory models to start with an energy function that is engineered for the problem of interest. The optimization process and the resulting architecture of the transformer block in our approach is a consequence of this specifically chosen energy function.
Concretely, we use the recent theoretical advancements and architectural developments in Dense As- sociative Memories to design an energy function tailored to route the information between the tokens. The goal of this energy function is to represent the relationships between the semantic contents of to- kens describing given data point (e.g., the relationships between the contents of the image patches in the vision domain, or relationships between the nodes' attributes in the graph domain). The core math- ematical idea of our approach is that the sequence of these unusual transformer blocks, which we call the Energy Transformer (ET), minimizes this global energy function. Thus, the sequence of conven- tional transformer blocks is replaced with a single ET block, which iterates the token representations until they converge to fixed point attractor state. In the image domain, this fixed point corresponds to the completed image with masked tokens replaced by plausible auto-completions of the occluded image patches. In the graph domain, the fixed point reveals the anomaly status of node given its neighbors, see Figurre 1] or the graph label. The energy function in our ET block is designed with the goal to describe the relationships between the tokens. Examples of relationships in the image domain are: straight lines tend to continue through multiple patches, given a face with one eye being masked the network should inpaint the missing eye, etc. In the graph domain, these are the relationships between the features of nodes; or features of nodes and graph label in graph classification.
The core mathematical principle of the ET block the existence of the global energy function
dictates strong constraints on the possible operations inside the block, the order in which these
2

---

ET for graphs
g(x')
x
ET
energy
Energy
Transformer
Block
ET
time
E
pos
ET for images
Figure I: Overview of the Energy Transformer (ET). Instead of a sequence of conventional transformer blocks, single recurrent ET block is used. The operation of this block dictated by the global energy function. The token representations are updated according to a continuous time differential equation with the dine-discretized update step a dt/T. On the image domain, images are split into non-overlapping patches that are linearly encoded into tokens with added learnable positional embeddings (POS). Some patches are randomly masked. These tokens are recurrently passed through ET, and each iteration reduces the energy of the set of tokens. The token representations at or near the fixed point are then decoded using the decoder network obtain the reconstructed image. The network is trained by minimizing the mean squared error loss between the reconstructed image and the original image. On the graph domain, the same general pipeline used. Each token represents node and each node has its own positional encoding. The token representations at or near the fixed point are used for the prediction of the anomaly status of each node, or the graph label.
operations are executed in the inference pass, and the symmetries of the weights in the network. As a corollary of this theoretical principle, the attention mechanism of ET different from attention mechanism commonly used in feed-forward transformers [5). Lastly, our network may be
viewed an example a broader class mmergy-Basees Models [7 frequently discussed the
community. The proposed model is defined through the specific choice of the energy function, which, on the one hand, is suitable for the computational task of interest, and, on the other hand, results in optimization equations closely related to the forward pass in feed-forward transformers.
2 Energy Transformer Block
We now introduce the theoretical framework of the ET network. For clarity of presentation, we use language associated with the image domain. For the graph domain, one should think about "image patches" nodes on the graph.
The overall pipeline similar to the Vision Transformer networks (ViTs) and is shown in Figure I] An input image is split into non-overlapping patches. After passing these patches through the encoder and adding the positional information, the semantic content of each patch and its position is encoded in the token F1A- In the following the indices i, k = 1...D are used to denote the token 1...N are used to enumerate the patches and their corresponding physical particle, which has complicated
vector's elements, indices A, B,
tokens. It is helpful to think about each image patch internal state described by a D-dimensional vector XAA This internal state describes the identity the particle (representing the pixels of each patch), and the particle's positional embedding (the patch's s

---

Ap
Position Embedding Correlations
Weights in HN
1.0
Addition
0.5
Patch Row
Mutti-Hoad
Hopfilld
Energy Attention
Netwerk
0.0
10
Layer Norm
18
x
Patch Column
Figure 2: Left: Inside the ET block. The input token x()) passes through a sequence of operations and gets updated to produce the output token x (++1), The operations inside the ET block are carefully engineered so that the entire network has global energy function, which decreases with time and is bounded from below. In contrast to conventional transformers, the ET-based analogs of the attention module and the feed-forward MLP module are applied in parallel as opposed to consecutively. Center: The cosine similarity between the learned position embedding of each patch and every other patch. In each cell, the brightest patch indicates the cell of consideration. Right: 100 selected memories stored in the HN memory matrix, visualized by the decoder as 16x16 RGB image patches. This visualization unique to our model, as traditional Transformers cannot guarantee image representations in the learned weights.
location within the image). The ET block is described by continuous time differential equation, which describes interactions between these particles. Initially, at f = 1 the network is given a set containing two groups of particles corresponding open and masked patches. The "open" particles know their identity and location in the image. The "masked" particles only know where in the image they are located, but are not provided the information about what image patch they represent. The goal of ET's non-linear dynamics is to allow the masked particles to find an identity consistent with their locations and the identities of open particles. This dynamical evolution designed that minimizes global energy function, and is guaranteed to arrive at fixed point attractor state. The identities of the masked particles are considered to be revealed when the dynamical trajectory reaches the fixed point. Thus, the central question is: how can we design the energy function that accurately captures the task that the Energy Transformer needs to solve?
The masked particles' search for identity is guided by two pieces of information: identities of the open particles, and the general knowledge about what patches are in principle possible in the space of all possible images. These two pieces of information are described by two contributions to the ET's energy function: the energy based attention and the Hopfield Network, respectively, for reasons that will become clear in the next sections. Below we define each element of the ET block in the order they appear in Figure 2
Layer Norm
Each token, or a particle, represented by a vector x E R. At the same time, most of the operations inside the ET block are defined using layer-normalized token representation
D
gi =
T
8,, where
T = D
Σ T
(I)
k=
The scalar y and the vector elements 5; are learnable parameters, is small regularization constant. Importantly, this operation can be viewed as an activation function for the neurons and can be defined partial derivative of the Lagrangian function (see [1 20 21
for the discussion of this property)
aL
== DT D
Σ x
x)2 + + Σ 85
so that
9:
3z

---

Multi-Heate Energy Attention
The first contribution the ET's energy function is responsible for exchanging information between the particles (tokens). Similarly to the conventional attention mechanism, each token generates a pair of queries and keys (ET does have a separate value matrix; instead the value matrix function of keys and queries). The goal of the energy based attention is to evolve the tokens in such way that the keys of the open patches are aligned with the queries of the masked patches in the internal space of the attention operation. Below we use index a 1...Y ...Y to denote elements of this internal I...H denote different heads of this operation. With these notations the space, and index energy-based attention operation is described by the following energy function:
N
1
E
log
Σ exp
(3AABC)
(3)
3
B&C
where the attention matrix BBBC is computed from query and key tensors as follows:
ABBC =
"KahD Qakcs
A E HxNNN
KahB =
Σ Wakg 9jB,
K E RYKHNN
(4)
Qarc = Σ
W.O. adj IJC,
Q QERYXHXN
and the tensors W E RYXHxD and W E RXXHxD are learnable parameters.
From the computational perspective each patch generates two representations: query (given the position of the patch and its current content, where in the image should it look for the prompts on how to evolve in time?), and key (given the current content of the patch and its position, what should be the contents of the patches that attend to it?). The log-sum energy function (3) is minimal when
for every patch in the image its queries aligned with the keys of a small number of other patches connected by the attention map. Different heads (index h) contribute to the energy additively.
Hopfield Network Module
The next step of the ET block, which we call the Hopfield Network (HN), is responsible for ensuring that the token representations are consistent with what one expects to see realistic images. The energy of this sub-block is defined as:
N
K
D
EHN
Σ B=1p=1
Σ j=1
Eps 9ja).
E RERKX
(5)
where Euj is set learnable weights (memories in the Hopfield Network), and G(.) is an integral of the activation function r(-), that G(.)' = r(-). Depending on the choice of the activation function this step can be viewed either as a classical continuous Hopfield Network [8] if the activation function grows slowly (e.g., r(-) =ReUU), activation function is sharply peaked around the memories (e.g., r(.) = power, or softmax). The HN sub-block analogous the feed-forward MLP step in the conventional transformer block but requires that the weights of the projection from the token space to the hidden neurons' space to be the same (transposed matrix) the weights of the subsequent projection from the hidden space the token space. Thus, the HN module here is an MLP with shared weights that is applied recurrently. The energy contribution of this block is low when the token representations are aligned with some rows of the matrix &, which represent memories, and high otherwise.
modern continuous Hopfield Network [TO] 12 [9]
if the
Dynamics of Token Updates
The inference pass of the ET network is described by the continuous time differential equation, which minimizes the sum of the two energies described above
dxiA
dE
where
E' = EAT + HN
(6)
dl
5

---

Here 24,4 is the token representation (input and output from the ET block), and 914 is its layer- normalized version. The first energy is low when each patch's queries are aligned with the keys of its neighbors. The second energy is low when each patch has content consistent with the general expectations about what an image patch should look like (memory slots of the matrix Q The dynamical system 6 finds trade-off between these two desirable properties of each token's s representation. For numerical evaluations, equation (6) is discretized in time.
To demonstrate that the dynamical system (6 minimizes the energy, consider the temporal derivative
dE
OE dy:a dajA - Σ
0E
0E
M
S 0
dt
θτja dt ogia
Egia
JJ.4
The last inequality sign holds if the symmetric part of the matrix
M Ogia =
0LL
(8)
Đzja a
positive semi-definite (for each value of index
A). The Lagrangian (
satisfies this condition.
Relationship to Modern Hopfield Networks and Conventional Attention
One of the theoretical contributions of our work is the design of the energy attention mechanism and the corresponding energy function 33) Although heavily inspired by prior work on Modern Hopfield Networks, our approach is fundamentally different from it. Our energy function (3) may look somewhat similar to the energy function of continuous Hopfield Network with the softmax activation function. The main difference, however, is that in order to use Modern Hopfield Networks recurrently (as opposed to applying their update rule only once) the keys must be constant parameters (called memories in the Hopfield language). In contrast, in our energy attention network the keys are dynamical variables that evolve time with the queries.
To emphasize this further, it is instructive to write explicitly the ET attention contribution to the update dynamics
(6). It is given by (for clarity, assume only one head of attention):
ØEATT
ΣΣ CAΛ
w Kao softmax B EKnC KgO Qua) +Wos Qae softmax (3 (
KGA Qnc)
OgLA
In both terms the softmax normalization is done over the token index of the keys, which is indicated by the subscript in the equation. The first term in this formula is the conventional attention mechanism [S] with the value matrix equal to V = (WQ)"K = Ea W The second term is the brand new contribution that is missing in the original attention mechanism. The presence of this second term is crucial to make sure that the dynamical system (6) minimizes the energy function if applied recurrently. This second term is the main difference of our approach compared the Modern Hopfield Networks. The same difference applies compared to the other recent proposals [15).
Lastly, we want to emphasize that our ET block contains two different kinds of Hopfield Networks acting in parallel, see Figure 2| The first one is the energy attention module, which inspired by, but not identical to, Modern Hopfield Networks. The second one is the "Hopfield Network" module, which can be either a classical or Modern Hopfield Network. These two should not be confused.
For completeness, the contribution of the "Hopfield Network" module to the update equation 6 can be written as
HE
K EMn
D (
(9)
0g.A
#=1
f=1
which is applied to every token individually (there mixing of different tokens).
3 Qualitative Inspection of the ET framework on ImageNet
We trained the ET network on the masked image completion task using ImageNet- Ik dataset [22] Each image was broken into non-overlapping patches of 16x 16 RGB pixels, which were projected
The code is available: https//githbb.cco//bhoovvvergrrr haa.://oeengggtttaaaaannrreaaaaa

---

with single affine encoder into the token space. Half of these tokens were "masked". by replacing them with learnable MASK token. distinct learnable position encoding vector was added each token. Our ET block then processes all tokens recurrently for 7 steps. The token representations after steps are passed a simple linear decoder (consisting of a layer norm and an affine transformation). The loss function is the standard MSE loss on the occluded patches. See more details on the implementation and the hyperparameters in Appendix
Figure 3: Reconstruction examples of our Energy Transformer using images from the Image/Ntt- 1 validation set. Top row: input images where 50% the patches are masked with the learned MASK token. Middle row: output reconstructions after 12 time steps. Bottom row: original images.
Examples of occluded/recoostruuted images (unseen during training) are shown in Figure 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 extending the un-occluded borders (last column). We additionally inspect the positional encoding vectors associated with every token, Figure 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 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, c.g.,
Early time
Late time
HN Update
HN Update
Tokens
Tokens (time s+1)
Tokens (time i)
Tokens
(time s)
(time [+1]
Time
Update
ATTN Update
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 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 they are available in the prior work). Best results are in bold. Our model is state of the art near state of the art every category.
Datasets
Split
GraphConsis CAREGNN
PC-GNN
BWGNN
MLP
GT
(Ours)
1%
56.8+2.8
+2.
59.8+1.4
61.1
53.9
61.7+0.4
63.0+a.6
Yelp
40%
58.7120
63.310.9
63.0
71.020.9
57.510.8
68. 68.710.4
71.51011
1%
68.5ga.
68.7
79.8
90.9
74.6
88.6 68.610.5
89.32a.7
Amazon
Macro-FI
40%
75.1
86.3
89.5
92.2
79.1
91.710.
92.8103
T- Finance
1%
71.7
T3.3
62.0
84.8
61.0
81.5
85.1110
40%
73.4
77.5
63.1
86.8
70.5
83.6
88.2110
52.4
55.8
51.1
75.9
50.0
64.3
79.1
T-Social
40%
56.5
56.2
52.1
83.9
50.3
68.2
83.5
1%
66.423.4
75.043.6
75.420.9
72.020.5
59.820.4
72.510.6
73.220.8
Yelp
40%
69.813.0
76.112.9
79.810.1
84.010.9
66.5110
81.910.5
84.910.3
1%
74.1
88.613.5
90.4120
80.410.3
83.611.7
89.011.2
91 01.9110
Amazon
40%
87.4+3.3
90.5+1.6
93.8+0.1
98.0+0.4
89.8+1.0
95.4+0.6
97 32.44
AUC
1%
90.2
90.5
90.7
91.1
82.9
90.0
92.8+1.
T-Finance
40%
91.4
92.1
91.2
94.3
87.1
88.2
95 05.0
1%
65.2
71.2
59.8
88.0
56.3
81.4
91 9..6
T-Social
40%
71.2
71.8
68.4
95.2
56.9
82.5
93.9
We additionally visualize the gradients of the energy function (which are equal the token updates, see Equation 6) of both ATTN block and the HN block, see beelFigure 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; meaningful signal comes from the HN block to the masked patch dynamics. Later in time we observe different phenomenon: almost all signal to masked tokens comes from the HN module while ATTN contributes 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 sharpening the masked patches. All the qualitative findings presented this section are accord with the core computational strategy of the ET block was designed theoretically in section 2
the model approaches the final reconstruction.
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 F task with plenty of strong and recently published baselines. Graph Convolutional Networks (GCN) 25 have been widely used for this task due 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 vector of raw attributes YA E R where F is the number of node's features. Every node also has a binary label LA indicating whether the node 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 images, the feature vectors for every node converted to token representation using a linear embedding E and adding a learnable positional embedding XA
YA + XA
(10)
where the superscript f indicates the time of the update of the ET dynamics. This token representation is iterated through the ET block for 7' iterations. When the retrieval dynamics becomes
The code is available:
hrnn Sntgppph-nomay-eeeeteetioo Detection/

---

Table 2: Graph classification performance on eight datasets of TUDaaset. 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 prior work. the entry 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 (decrease), (increase), and V (no change within the ertor bars) along with the value.
Method
PROTEINS
NCII
NC1108
DD
ENZYMES
MUTAG
MUTAGENICITY FRANKENSTEIN
WKPI
78.5gn4 v(6.4)
$7.5ga5
85.9 T(L.)
62.0gas ((8..7)
85 85.8 V(142)
(kmeans)
(koenters)
WKPI
75.2 ((9.7)
$4.5.a. ((UU)
87.4.aa
80.3a.4 (15.4)
88.laz ((11.7)
Spre-GN
84.8 8LLL v(2.7)
63.6 v(3.8)
72.5 v15.)
Norm-GN-
64.9g1 ((2.6)
83.5 ((3.9)
73.3 v[5.1)
GWL-WW
75.8 (00.)
71.3a11 v(7.1)
78 78.9
HGF-SL
84.916
78.5as v(U1)
8072.1 v(6.7)
$100ga ((4.T)
68.8ga ((9.6)
82.2.
DSGCN
77.3gnc
v(7.6)
78.4gu
U2GNN
80.0gas
(4.9)
96.711
88.5gg ((11.))
NDP
73.4ga
T(1.5)
742gur 7(15.))
728 ((2.0)
46.5mm v(S40) 87.0aar 7(12.1) 77.92 v(d.3)
ASAP
74:2uu
(10.7) 7L5224 ((18.) 70.1,a4 v(17.3) 78.9p (188))
08.3263 (12.6)
EvoG
66.7 (227)
100.0
(Ours) 90.3g07 4(54) 00:1.01 A(288) 90.5.aa 40.1) 95.0 2 90.8 4(21.))
90 M.0 TO3.)) 08 7:aa A(16.5) 99 8:81 A(2.9)
stable, we have the final representation for each node x T (or more precisely g since the outputs additionally passed through layer norm operation after the final ET update). This output is concatenated with the initial (layer normalized) token to form the final output
final gA = g
gA
(11)
where || denotes concatenation. Following 24), the node representation g fed into an MLP with sigmoid activation function to compute the anomaly probabilities PA- The weighted cross entropy
Loss = Σ w log(pA) + (1 - log(1 -PA)
(12)
is used to train the network. Above, w is the ratio of the benign labels (l. = 0) to anomalous ones (lA 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 amazon.com 330 T T-FFnaacce and T-Social datasets [4 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 into validation and testing sets, see (Appendix for details.
We compare with satte-of-the-ann approaches for graph anomaly detection, which include GraphConsis [I]] CAREGNN (2),, PC-GNN 3 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 Fl score) and the Area Under the Curve (AUC) are used as the evaluation metrics on the test datasets 35 See Aeeennddiixx C for more details on training protocols and the hyperparameters choices. The results are reported Table 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 problen Unlike the anomaly detection task, where we predict a label for every node, in the graph classification task, single label is predicted for the entire graph.
Consider graph G, where each node has a raw feature vector Each vector is projected to the token space yielding X,4 E Ik D, where D is the token dimension. learnable CLS token XCLS is concatenated the set of tokens resulting in the token matrix X e R(N+1)xD and learnable linear projection of the top k smallest ligenvecctor of the normalized Laplacian matrix, following [37], is added to X. Graph structural information is provided to ET within the attention operation. We also use stacked version of the Energy Transformer consisting of S ET blocks, where each block has the same number of temporal unfolding steps 7 and its own LayerNorm. <=S, t=f is projected via 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 SeepAddddxx for full details.
RF with F feature-dimension.
The final representation of the CLS token XCLS
5.1 Experimental Evaluation
Eight graph datasets from the TUDataset 33 collection are used for experimentation. NCII, NCT109, MUTAG, MUTAGENICITY, and FRANKENSTEIN are 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 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 [6). See Table 8 for more details of the datasets.
We compare ET with the current state-of-the-are 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 33 43
44
45
46 which are close 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
lot of recent research has been dedicated to understanding the striking analogy between Hopfield Networks and the attention mechanism in transformers. At 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 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-sty weight transformers (see Appendix G [nd 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).
shared-
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: http://gttii...cooooooooooooooon
10

---

References
[1] Ofir Press, Noah A Smith, and Omer Levy. Improving transformer models by reordering their sublayers. arXiv preprint arXiv:1911.0386. 2019.
[2] Yiping Lu, Zhuohan Li, Di He. Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie- Yan Liu. Understanding and improving transformer from a multi multi-particle dynamic system point of view. arXiv preprint arXiv: 027600 2019.
[3] David So, Quoc Le, and Chen Liang. The evolved transformer. In International Conference Machine Leamiin, pages 5877-5836. PMLR, 2019.
[4] Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. arXiv preprint arXiv:2106.045550 2021.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[6] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. arXiv preprint vrXiv:2/11.11415: 2021.
[7] John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7554):2554-2558 1982.
[8] John J Hopfield. Neurons with graded response have collective computational properties like those of two-state neurons. Proceedings of the national academy of sciences, ((10))3088-3092 1984.
[9] Dmitry Krotov. A new frontier for hopfield networks. Nature Reviews Physics, pages 1-2, 2023.
[10] Dmitry Krotov and John J. Hopfield. Dense associative memory for pattern recognition. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
[11] Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On a model of associative memory with huge storage capacity. Journal of Statistical Physics, 168(2):288-299, May 2017.
[12] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.022 02211, 2020.
[13] Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
[14] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020.
[15] Yongyi Yang, Zengfeng Huang, and David Wipf. Transformers from an optimization perspective. arXiv preprint arXiv:2205. 13991,
[16] Ying Sun, Prabhu Babu, and Daniel P Palomar. processing, communications, and machine learning. IEEE Transactions on Signal Processing. 65(3):794 816, 2016.
algorithms in signal
[17] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy- based learning. Predicting structured data, 1(0), 2006.
[18] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlshy. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.

---

[19] Dmitry Krotov and John J. Hopfield. Large associative memory problem in neurobiology and machine learning. In International Conference on Learning Representations, 2021.
[20] Fei Tang and Michael Kopp. A remark on a paper of krotov and hopfield [arxiv: 2008.06996]. arXiv preprint arXiv:2105. 15034, 2021.
[21] Dmitry Krotov. Hierarchical associative memory. arXiv preprint arXiv:2107.06446, 2021.
[22] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. Icce, 2009.
[23] Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In European conference on computer vision, pages 818-833. Springer, 2014.
[24] Jianheng Tang, Jiajin Li, Ziqi Gao, and Jia Li. Rethinking graph neural networks for anomaly detection. arXiv preprint arXiv:2205. 2022.
[25] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional 2016.
networks, arXiv preprint
[26] Kaize Ding, Jundong Li, Rohit Bhanushali, and Huan Liu. Deep anomaly detection on attributed networks. In Proceedings of the 2019 SIAM International Conference on Data Mining, pages 594 602. SIAM, 2019.
[27] Zhen Peng, Minnan Luo, Jundong Li, Luguo Xue, and Qinghua Zheng. A deep multi-view framework for anomaly detection on attributed networks. IEEE Transactions on Knowledge and Data Engineering. 2020.
[28] Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. Simplifying graph convolutional networks. In International conference on machine learning, pages 6861-6871. PMLR, 2019.
[29] Shebuti Rayana and Leman Akoglu. Collective opinion spam detection: Bridging review networks and metadata. Proceedings of the 21th aem sigkdd international conference on knowledge discovery and data mining, pages 985-994, 2015.
[30] Julian John McAuley and Jure Leskovec. From amateurs connoisseurs: modeling the evolution of user expertise through online reviews. In Proceedings of the 22nd international conference on World Wide Web, pages 897-908, 2013.
[31] Zhiwei Liu, Yingtong Dou, Philip S Yu, Yutong Deng, and Hao Peng. Alleviating the inconsis- tency problem of applying graph neural network to fraud detection. In Proceedings of the 43rd international ACM SIGIR conference on research and development in information retrieval, pages 1569-1572, 2020.
[32] Yingtong Dou, Zhiwei Liu, Li Sun, Yutong Deng, Hao Peng, and Philip S Yu. Enhancing graph neural network- fraud detectors against camouflaged fraudsters. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pages 315-331 3244, 2020.
[33] Yang Liu, Xiang Ao, Zidi Qin, Jianfeng Chi, Jinghua Feng, Hao Yang, and Qing He. Pick and choose: a gnn-based imbalanced learning approach for fraud detection. In Proceedings of the Web Conference 2021, pages 3168-3177, 2021.
[34] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09669, 2020.
[35] Jesse Davis and Mark Goadrich. The relationship between precision-recall and roc curves. In Proceedings of the 23rd international conference on Machine learning, pages 233-240, 2006.
[36] Christopher Morris, Nils M. Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. CoRR, abs/2007.08660 2020.
12

---

[37] Vijay Prakash Dwivedi, Chaitanya K. Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks, 2022.
[38] Qi Zhao and Yusu Wang. Learning metrics for persistence-based summaries and applications for graph classification. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[39] Muhammet Balcilar, Guillaume Renton, Pierre Héroux, Benoit Gaüzère, Sébastien Adam, and Paul Honeine. Bridging the gap between spectral and spatial domains in graph neural networks. CoRR, abs/2003.1 1702, 2020.
[40] Zhen Zhang. Jiajun Bu, Martin Ester, Jianfeng Zhang, Zhao Li, Chengwei Yao, Dai Huifen, Zhi Yu, and Can Wang. Hierarchical multi-view graph pooling with structure learning. IEEE Transactions on Knowledge and Data Engineering, pages 1-1, 2021.
[41] Dai Quoe Nguyen, Tu Dinh Nguyen, and Dinh Phung. Universal graph transformer self-attention networks. In Companion Proceedings of the Web Conference 2022, WWW '22, page 193-196, New York, NY, USA, 2022. Association for Computing Machinery.
[42] Miguel Domingue, Rohan Dhamdhere, Naga Durga Harish Kanamarlapudi, Sunand Raghupathi, and Raymond Ptucha. Evolution of graph classifiers. In 20/9 JEEE Western New York Image and Signal Processing Workshop (WNYISPW), pages 1-5, 2019.
[43] Mingqi Yang. Yanming Shen, Rui Li, Heng Qi, Qiang Zhang, and Baocai Yin. A new perspec- tive on the effects of spectrum in graph neural networks. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 25261-25279, PMLR, 17-23 Jul 2022.
[44] Francesco Orsini, Paolo Frasconi, and Luc De Raedt. Graph invariant kernels. Proceedings of the 24th International Conference on Arlificial Intelligence, IJCAI 15, page 3756-3762. AAAI Press, 2015.
[45] Filippo Maria Bianchi, Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Hierarchical repre- sentation learning in graph neural networks with node decimation pooling. IEEE Transactions on Neural Networks and Learning Systems, 2020.
[46] Ekagra Ranjan, Soumya Sanyal, and Partha Pratim Talukdar. ASAP: Adaptive structure aware pooling for learning hierarchical graph representations. arXiv preprint urXiv:191.00797) 2019.
[47] Benjamin Hoover, Duen Horng Chau, Hendrik Strobelt, and Dmitry Krotov. A universal ab- straction for hierarchical hopfield networks. In The Symbiosis of Deep Learning and Differential Equations II.
[48] Ross Wightman. Pytorch image models.
httpp:://github.ccm/rrrgggg
pytorch- iethh--mmaeemmoo models 2019.
[49] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin -Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
[50] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020.
[51] Tong Zhao, Chuchen Deng, Kaifeng Yu, Tianwen Jiang, Daheng Wang, and Meng Jiang. Error-bounded graph anomaly loss for gnns. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pages 1873-1882, 2020.
[52] Kaize Ding, Jundong Li, Nitin Agarwal, and Huan Liu. Inductive anomaly detection on at- tributed networks. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence, pages 1288-120 2294,129,,2 2021.
13

---

[53] Yuchen Liang, Dmitry Krotov, and Mohammed J Zaki. Modern hopfield networks for graph embedding. Frontiers big Data, 5, 2022.
[54] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 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 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 Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. arXiv preprint arXiv:23033061147 2023.
(59] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09697 2020.
[60]
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, 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, lurii 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 Stokowiee, 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-cenntralzzatinn 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: /608. 03983, 2016.
[66] Alain Horé and Djemel Ziou. Image quality metrics: Psnr vs. ssim. In 2010 20th International Conference on Pattern Recognition, pages 2366-2369, 2010.
14

---

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
dimension of node's feature space
D
dimension of token space number of tokens number of hidden dimensions in the attention
Y
number of hidden dimensions the Hopfield Network module number of recurrent time steps
T
number of heads
kn
height of each image patch
km
width of each image patch number of pixels per image patch ki x kw) input feature vector of node A vector representation of token each element of vector representation of token A vector representation of token A after layernorm
P
YA
XA
ZAA
gA
9:4
each element of vector representation of token A after layernorm key tensor query tensor
K
Q
KOAB
each element of the key tensor each element of the query tensor Q each element of the attention tensor label of node A graph
Qatc
AABC
LA
r(.)
activation function of the Hopfield Network module energy contribution of the Hopfield Network module
G(-)
B Details of Training on ImageNet
We trained the ET network on masked-image completion task on the ImageNet-1k (INIK) dataset. We treat all images in INIK as images of shape 224 x 224 that are normalized according to standard INIK practices (mean variance 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. "Oecluded" tokens are designated as follows: of the 100 tokens, 90 tokens are replaced with learnable MASK token of dimension 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 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 scalar step size a dt which we treat as a non-learnable hyperparameter The token representations after T steps are passed to simple linear decoder (consisting of layer norm and an affine transformation) to project our representations back into the image plane. We then 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
our experiments.
KahB Qatc) ))
EATT_ Σ
log Σ exp
(13)
3
15

---

We give details of our architectural choices in Table 4 In the main paper we present our Energy Transformer with a configuration similar the standard base Transformer configuration (e.g. token dimension 768, 12 heads each with Y = 64, softmax's is B Vy....) with several considerations leamned from the qualitative image evaluations:
The de (step size) of implicitly used in the traditional transformer noticeably degrades 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 learn meaningful filters.
Values of 3 in the energy attention that are too high prevent our model from training. This is possibly due to vanishing gradients our attention operation from a smat operation that is too spiky.
Without gradient clipping, our model fails to train at the learning rates we tried higher than le-4. We observe that gradient clipping not only helps our model train faster 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- Ik masked training experiments. Data augmentations are listed as parameters passed to the equivalent timm dataloader functionality.
Training
Architecture
Data Augmentation
batch_size
768
token_ding
768
random_erane
None
epochs
100
num nea_heads
12
horizontal flip
0.5
Ir
5e-4
head dim
64
vertical flip
0
warmup @ermuu_epochhs
2
B
1/8
color jitter
0.4
start & end Ir
5e-7
train betas
No
scale
(0.08, 1)
bl, b2 (ADAM)
0.9, 0.99
step size a
0.1
ratio auto auto_augment
(3/4, 4/3)
weight_ decay grad_cllippinn
0.05
depth
12
None
hidden dim HN
3072
bias in HN
None
bias in ATT
None
bias in LNORM
Yes
B.1 Exploring the Hopfield Memories
distinctive aspect of our network is that any variable that has vector index 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 subset of memories. The complete set of memory vectors from the HN module is shown in Figure 3] 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 correla- tions of the original ViT in that they are highly susceptible to choices of the hyperparameters (Figure 10 of [88). particular, we consider the effect of weight decay and the parameter that serves as the inverse temperature of the attention operation (see Eqquation 3). The lower the temperature (i.e., the higher the value of 3), the spikier the softmax distribution. By using a lower 3, we encourage the attention energy to distribute its positional embeddings across a wider range of patches in the model.
16

---

Figure 5: Visualizing a randomly selected 3025 patch memories of the 3072 learned by weight matrix in the Hopfield Network module (HN) of our model. These memories are vectors of the same dimensions D as the patch tokens, stored as rows the weight matrix &. Each image patch visualized using the model's 's trained decoder.
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 very popular due to their capability of learning sophisticated data representations. DOMINANT (66 utilizes an auto-encoder framework, using 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 [TI] propose a novel loss function train graph neural networks for anomaly-detectabl node representations. In [52] generative adversarial learning is used to detect anomaly nodes where a novel layer designed learn the anomaly-awar 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 filter that can better handle this phenomenon. We propose 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
17

---

W keys
W queries
Figure Visualizing the token dimension the "key" and "query" matrices of the attention as image patches. Each head represented as cell the 4 3 grid above. We use the trained decoder of our model to visualize each weight.
WD=0.001; =0.02
WD=0.001; 800.125
WD=0.05; 8=0.125
0.0
00.5
11.0
Figure 7: The cosine similarity between position biases of patches when the ET model is trained under different hyperparameter choices for 3 (inverse temperature of the attention energy) and weight decay. Our ET sees trend where smoother correlations are observed with smaller / and weight decay.
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
First, compute the features that passed to energy-based transformer. Each node's features Y E R are mapped into the token space XA E R& using a linear projection E. Learnable positional embeddings da are added this token at 1,
Ey 4 XA
(14)
At each time step the input to the ET-block layer normalized:
gAA LayerNorm(xx)
(15)
18

---

Let W E RXXNX and w E RXXHxD be the query and key weight matrices, respectively. Here Y is the projection dimension in the attention operation. H is the number of heads. We define
KahB =
Σ
Wahg abs 9jB
(16)
Qaic
Σ
W ohj 9jC
If we let h indicate the index of the head, we have
A = Σ
Σ
Woo ahi Kanc woa + W.K Qaic waa
+
E Ens rE ENslsA)
(17)
CENA h.a
where
WCA = softmax B Σ K Qrnas)
(18)
Here & controls the temperature of the softmax, .A stands for the neighbors of node A - set of all the nodes connected to node A, r 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 a natural consequence of the discretized differential equation dynamics)
= x + Ax'a
(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 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
-ー
ΣΣ
log
Σ exp
[<
B &
Kahn Qanc)
Σ C2Ps
G Σ Eus gjc)
(20)
BeNc
This energy function will decrease as the forward pass progresses until it reaches a local minimum.
After T iterations when the retrieval stable, we have the final representation for each node final as
8A
BA
gA
(21)
where| || is the concatenation sign. Following (24). we treat anomaly detection as semi-supervised learning task this work. The node representation final is fed to another MLP with the sigmoid function to compute the abnormal probability pa, weighted logeliiel is then used train the network. The loss function is as follow:
Loss = Σ
w la log(ps) + (1 - log(1 PA)
(22)
where w is the ratio of normal labels (A = 0) to anomaly labels (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-FI 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 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 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 224. We report the average performance of 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 different subgraph to train for each epoch (subgraphs have 5% of the nodes with respect to the whole training data). The hyperparameters
19

---

Dataset
V
[
Anomaly( (6)
Features
Amazon
11944
4398392
6.87%
25
Yelp
45954
3846979
14.53%
32
T-Finance
39357
21222543
4.58%
10
T-Social
5781065
73105508
3.01%
10
Table Summary of all the datasets.
include the number of hidden dimensions in ET-attention Y, the number of neurons in the hidden layer within the Hopfield Network Module, the number of time iterations T, and the number of heads H. The weights are learned via backpropagation, which includes embedding projection E, positional embedding A,, softmax inverse temperature parameter ETT-attention weight tensors W and W The optimal hyperparameters used in Table are reported in Table 6 The last row in that table summarizes the range of the hyperparameter search that was performed our experiments. In general, we have observed that for small datasets (Yelp, Amazon, T-Finance) 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 iterations.
Dataset
K
T
H
Amazon (40%)
128
Amazon (1%
64
128
Yelp (40%)
128
256
1
Yelp (1%
128
256
T-Finance (40%)
128
256
T-Finance (1 (1%)
T-Social (40%)
128
256
3
3
T-Social (1%)
128
256
3
3
Range of hyperparameters
(64, 128, 256]
(2Y,3Y,4Y,
[1,233
(1,2,3)
Table 6: Hyperparameters choice of our method on all the datasets.
D Graph Classification with ET
Recently, there have been attempts 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 the Transformer architecture to model graphs and how 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 graph (V, E)), we have the of nodes V {U1, 02, .2,..... uN }, the adjacency matrix A e {0, NNN, and it exists, the edge feature matrix E N where each edge has p raw features. Each feature vector Y RH'" corresponding token space XA E D via linear embedding. Then, the CLS token XCLS is concatenated to the set of tokens resulting in E R(N-1)xD and positional embedding R (N+1)xD added afterwards.
node VA is first projected to the
To obtain the positional embedding À (N+1)xD the adjacency matrix A first padded in the means to provide the CLS
upper left corner with ones resulting À 10,1)NN+ token full connectivity with all of the nodes in a given graph. Following [37], the top A smallest eigen-vectors À E IR +X+1))* are extracted from the normalized Laplacian matrix L obtained from Á
i-I-D -
(23)
and projected to the token space via linear embedding to form X. Note, I and D are the identity matrix and the degree matrix of A.
20

---

Meanwhile, the attention in ET is modified to take in Å R(N+1)(N+1)XX the parameterized adjacency tensor, which acts the weighted "attention mask' that enables the model to consider the graph structural information. To obtain A, 20-coonolttional layer with H filters equals the number of heads in the attention block, 'SAME" padding, and stride of one is performed on the outer product of X to itself. The result is then multiplied with the tensor A' (NNN+)) element-wee (denoted by )) via broadcasting
Á Conv2D X) A
(24)
The tensor A' is a linear projection, via linear embedding layer, of the edge feature matrix & if it exists, where P' is projected to // dimension; or the padded adjacency matrix A where p 1. Altogether, the resulting energy equation is
Σ log Σ exp E KanB Qaro O Anc) Anc)) )
E -
(25)
B4C
Moreover, to only consider the energy of edges in G, we set non-connected entries in Á to be -oc, 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 vertically stacked ET blocks, where each block shares the same number of T depth and has its own LayerNorm. Similarly, the token representation X at dynamic step corresponding to a block & is layer-normalized,
g LayerNorm (X ').
(26)
Keep in mind, X X f=1, is the initial token representation.
Following the dynamic equations 6 we inject small amount of noise e f (0, 1). generated from normal distribution using a standard deviation o, and zero mean, into the gradient of energy function to produce X 2+1, the new token representation of block f. The premise of this noise injection is to "robustify" the model and help it escape saddle points of the energy function.
t = x aVgEL Vael
el ((,,0(0,02))
(27)
Once stability is reached in the retrieval dynamics block 4, the final representation X t-T is then passed on to the next block L + and the whole process is repeated again. When the final token 8= t=T RD
representation 8=S, t=T' is computed by the last block S, the resultant CLS token XCLS extracted from X t= t=T is utilized as the predictor of the current graph 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 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; CIFARIO 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 588 GT [59], and GAT [54].
The results are depicted in table 7
Based on the collected results, ET outperforms its baselines 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 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 CIFAR 10, MNIST, and CLUSTER datasets. However, ET has overall demonstrated great performance on node binary-clasifficttion and graph-regression tasks.
21

---

Table 7: Graph classification performance on five datasets of GNNBenchmark [37]. The reported mean and standard deviation are obtained from 100 runs of evaluation on the test set of each dataset. For baselines standard deviations are only included if they are available in prior work. the entry unavailable in prior literature it is denoted by "-"; best results are in bold. The performance difference between non-baseline approaches (including ours) and the baseline (specified by the gray cell in each column) is indicated by V (drop in performance with respect to the baseline), A (increase in performance with respect to the baseline), and V (no change). Additionally, + indicates lower is better and indicates higher is better.
Dataset
Method
CIFARIO (Accuracy)
MNIST (Accurcc))
PATTERN (Accurcc))
CLUSTER (Accurcc))
ZINC (MAF)
GT
84.808 ose (2.013)
73.16920.622 (6.063) 0.226 (0.1122)
GAT
61.23 10.455 V (10.531)
93.535
(2.879)
78.271
(8.550)
70.587 (8.613)
0.361gu.or (0.276)
EGT
68.712 7027002 10.00 V (6.162)
98.173 0.087 (0.211)
86.821 86.821 1o0000
79. 19.232
0.108
Graphormer 65.97810.57 (8.776)
97.905_ V (0.509)
36.650 asa (0.171))
14.600_n..0m V (4.372)
0.
0.122
((000144
Exphormer
74.75420.194
98.4 4.000
96.73420.00e
(0.0877;
78. 78.220=0.045 (1.012)
(Ours)
1rs (10.834) 97.01120e4s ((L403)
90.690.0.01a á (3.869)
27.169_0.04 V (2.063)
0.09620.000 (0.012)
D.3
Experimental Details
In the graph domain, it is common to concatenate all of the feature vectors of all graphs in batch together. However, in order for ET to work, we form the batch dimension by separating the feature vectors of all graphs in given batch and utilize the largest node count 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 graph has 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 a result. However, is worth mentioning such 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 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/5)] the sign of the top eigen- vectors flipped randomly during training and fixed during evaluation. The whole experiment is implemented using JAX 49 Flax (50, Optax [60]. and PyTorch Geometric 66 frameworks.
Table 8: The statistics and properties of the eight datasets of TUDataset (additional node attributes are indicated by '+' if exist).
Dataset MUTAG ENZYMES
Graphs 188
Avg. Nodes
Avg. Edges
Node Attr
Classes
17.93
19,79
600
32.63
62.14
18 + 3
PROTEINS
1113
39.06
72.82
0 044
DD
1178
284.32
715.66
89
2
NCTI
4110
29.87
32.30
37
2
NCT109
4127
29.68
32.13
38
MUTAGENICITY
4337
30.32
30.77
FRANKENSTEIN
4337
16.90
17.88
780
For the eight datasets of TUDataset, we train all models for 300 epochs using AdamW 62 and Gradient Centralization 633. 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 smoothing parameter is set to 0.05. Additionally, the cosine-annealie
with warm-up learning rate scheduler (65) utilized, where the initial and end learning rates are both set as 5e while peak learning
rate is 0.001. The number of warm-up steps is set to 50 epochs while the batch size is 32 for all
22

---

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 Sible/andd the statistics of the used datasets are reported in Table g We also include the number of gpu-devices used for training ET.
Table 9: Hyperparameter and architecture choices for ET during TUDataset experiments.
Training
Architecture
batch size
32
token dim
128
epochs 300
num heads
num_heads
12
peak
le-3
head_ding
64
warmup_epodee
50
B
initial and ending
5e-6
tetas_
Yes
vei
b1, b2 (ADAM)
0.9, 0.99
step size a eigenvalues
0.01
weight decay
0.05
15
grad_ cel__lip num. of gpu devices
None
noise de depth
0.02
2
1
block_size
4
kernel_size
[3, 3]
dilation_size
[1,
hidden_ding dim HN
512
bias in HN
None
bias in ATT None
bias in LNORM num. of params per ET block
Yes
262,929
avg. total num. of params
1071294
Similarly, for the five datasets of GNNBenchmark, train all models using using Adam W 62 and Gradient Centralization [63. The best model is selected based on its performance obtained from the validation set across 10 runs. For the CIFARIO 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 means reduce overfitting, for the CIFAR I0
the mean absolute error function. Meanwhile, and MNIST models, the mean squared error loss function is utilized to reconstruct the edge feature matrix & in addition the cross entropy loss function. Specifically, the parameterized adjacency tensor Á is projected to match p' raw feature dimension of &, where both tensors are as part of the mean squared error function. In addition, all models are trained with the cosine-annealin 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 paraeetas are reported in Tables 10 and I1 and the statistics of the used datasets are reported in (37].
E ET for heterogeneous Graph
In this section, we show the performance 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 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 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 the heterogeneous setting performs worse than BWGNN in the homogeneous setting on Amazon.
23

---

Table 10: Hyperparameter choice for ET training on GNNBenchmark datasets.
Dataset
Parameter
CIFARIO
MNIST
PATTERN
CLUSTER
ZINC
batch_sizz
128
128
128
128
128
epochs
150
150
100
300
500
initial
le-3
le-3
le-3
le-3
5e-7
ending Ir
5e-5
5e-5
5e-7
5e-7
5e-7
peak Ir warmup_epont
le-3
le-3
le-3
le-3
le-3
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 _dccttddecyyee cld__ciiping
0.05
0.05
0.05
0.05
0.05
None
None
None
None
None
num. of gpu devices
1
8
Table 11: Architecture choice for ET on GNNBenchmark datasets.
Dataset
Parameter
CIFAR10
MNIST
PATTERN
CLUSTER
ZINC
token_din
128
128
128
128
128
num Lum_heads
12
12
12
12
12
head_ddm
64
64
64
64
64
B
1
1
1
1
664
V64
V64
/64
V64
vain_betar
Yes
Yes
Yes
Yes
Yes
step size
0.25
0.25
0.1
0.1
0.1
eigenvalues
15
15
15
15
15
noise ae
0.02
0.02
0.005
0.005
0.005
depth
4
4
1
block_sio size
2
2
2
2
4
kernel_sizz size
[3, 3]
[3
[3, 3]
[3, 3]
[3, 3]
dilation_sio hidden, dim HN
[1, 1]
[1,
[1, 1]
[1, 1]
[1, 1]
512
512
512
512
512
bias in HN bias in ATT bias in LNORM
None
None
None
None
None
None
None
None
None
None
Yes
Yes
Yes
Yes
Yes
total num. of params
529224
528968
527936
528964
1052409
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 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 consistent with the general expectation about particular data domain. It interesting explore the contribution that these two subnetworks produce on the task performed by the network. this section we ablate the ET architecture by dropping each the two subnetworks and measuring the impact of the ablation on the performance.

---

Table 12: ET for anomaly detection in heterogeneous graph setting, Best results are in bold.
Datasets
Split
MaxPool
Concatenation BWGNN (Heterogenous)
Macro-Fi
Yelp
1%
61.540.4
61.7
67.02
40%
70.7+0.6
71.140.1
76.96
1%
88.3
87.4+1.4
83.83
Amazon
40%
92.1
91.8+0.2
91.72
1%
72.240.5
72.8+0.2
76.95
Yelp
AUC
40%
84.3
85.1 80.1
90.54
1%
90.841.4
90.641.0
86.59
Amazon
40%
97.510.1
97.240.6
97.42
F1
On Graphs
The results on graphs are reported in Table 13 From this table it is clear that most of the computation performed by the ATT block on this task, which pools the information about other tokens the token of interest. When ATT block is kept, but HN block is removed the network looses 1% or less relative to the full ET (occasional improvements of the ablated model compared 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 given node for that node's s anomalous status prediction. This results a more significant drop in performance, which is about 5% on average.
Table 13: Ablation study with respect to ATT block and HN Block. Best results are bold.
Datasets
Split
ATTV HNX
HT
full model (Ours)
1%
62.540.3
57.4+0.5
(-5.6)
63.0+0.6
Yelp 40%
70.640.5
71.240.7 V
71.540.1
1%
89.5+0.9
87.4+1.
89.3+0.7
Amazon
Macro-FI
40%
91.7+0.66
(-1.1)
88.7+0.3
(-4.1)
92.840.3
1%
84.7+1.0
80.3
(-4.8)
85.1+1.0
T-Finance
40%
87.4+0.7
82.3+0.8
(-5.99)
88.2+1.0
1%
79.8
72.7+1.0
V--..4)
79.1
T-Social
40%
82.9+1.0
78.6+1.2
(-4.9)
83.5+0.4
1%
72.9
67.4
(-5.8)
73.2
Yelp
10%
83.5+0.4
(-1.4)
83.1
(-1.8)
84.9+0.3
1%
90.7
89.8+1.2
91.9+1.0
Amazon
40%
96.8+0.6
95.7+0.5
(-1.6)
97 97.3+0.4
AUC
1%
91.7+11.2
90.2+0.8
(-2.6)
92.8+1.1
T-Finance
40%
94.342.6
90.2+2.1
95.0+3.0
1%
92.2+0.8
86.4+0.7
(-5.5)
91.9+0.6
T-Social
40%
93. 83.1400
88.341.3
(-5.6)
93.9
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 INIK, we calculate the average MSE the reconstructed masked tokens for the validation set (using the same 50% masking ratio used for training) across 10 different random seeds for the masking.
We make several conclusions from these ablation studies.
25

---

We gain several insights regarding the use of "self-attention" in our ET (when a token patch query is allowed consider itself as key in the attention weights). When both self-attention and HN are present (ET-Full+Selff) there is no noticeable benefit over ET-Fuul for 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-NoHNN*), we notice that allowing self-attention (ET-NoHN+Self) outperforms the version that prevents self-attention (ET-NoHNN.
On its own, allowing self-attention (ET-NoHN+Self) in the ATTN module performs nearly as well the full ET fraction of the total parameters. However, MSE is forgiving metric for blurry reconstructions. While ATTN can capture the global structure of the image quite well, does so 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 best performing model (ET-Full) on the qualitative image reconstructions responding to the largest errors across IN1K validation images, averaged across all masking seeds. Likewise, Figure shows the lowest errors across INIK 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.
Table 14: Module ablation tests for image reconstruction task, reporting average INIK validation MSE masked tokens after 100 epochs. Reported number of parameters excludes the constant number parameters in the affine encoder and decoder.
Allow self-attn?
NParams (in ET block)
Model
Has ATTN?
Has HN?
MSE
ET-Full (Ours)
3.7M
0.306
ET-Full+Self
/
3.7M
0.312+0.10
ET-NoHN+Self
V
X
1.3M
0.343+0..10
ET-NoHN
X
X
1.3M
0.103
ET-NoATT
X
X
2.5M
0. 0.8254...0
Avg MSE:
MSE:
MSE:
Avg MSE:
MSE: 0.734
0.762
0.744
0.731
Figure 8: Reconstruction examples of images with the worst MSE from the INIk validation set. Top row: input images where 50% of the patches are masked with the learned MASK token. Middle row: all tokens reconstructed after 12 time steps. Bottom row: original images.
G Parameter comparison
The energy function enforces symmetries in our model, which means ET has fewer parameters than its ViT counterparts. particular, ET has no "Value Matrix" W in the attention mechanism, and the HN module has only one of the two matrices in the standard MLP of the traditional Transformer
26

---

Avg MSE: 0,0099
Ávg MSE:
Avg MSE: 0.0269
Avg MSE: 0.032
Avg MSE:
Avg MSE:
Avg MSE: 0.0389
Avg MSE:
Avg MSE: 0.0421
Avg MSE: 0.0423
0.0246
0.0324
0.0341
0.0396
Figure 9: Reconstruction examples of images with the best (lowest) MSE from the IN1k validation set. Top row: input images where 50% of the patches are masked with the learned MASK token. Middle row: all tokens reconstructed after 12 time steps. Bottom row: original images.
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_ssiii size=16, and report the number of parameters against ViT-base and an "ALBERT" version of ViT [14] where single ViT block is shared across layers. saw no benefit when including biases ET, so we also exclude 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 reduce the parameters of ViT-Base such that the parameter count
means for fairer comparison, similar to that ET. We contrast ET and the reduced ViT on the image reconstruction task using PSNR (Peak-Siz (Peak-Signal-toe -ooenal-o-NNoie-RRRe Ratio)) and SSIM Prructural-SMMilarte of the reconstructions. The results are demonstrated in Table 16 where the performance of ET is very close that of the reduced ViT, which has slight advantage in parameter count.
[66] the evaluation metrics
Table 15:
Comparison between the number of parameters in a standard ViT, an ALBERT version of ViT where standard Transformer blocks are shared across layers, and our ET. Comparison is done assuming no biases in any operation.
NParams (per block)
Model
NParams
ViT-Base
86.28M
00.00%
7,08M
00.0%%
ALBERT VIT-Base
8.41M
90.25%
7.08M
70.00%
ET
4.87M
94.36%
3.54M
50.02% %
Table 16: Comparison between and 'comparable-size' VIT on image reconstruction task. Given
VIT- VaseBase, reduce its parameter count down to number similar to that of ET for the image domain. The metrics, PSNR (Peak-Signal Signal-to-Noise -oisee-----)))) and SSIM (Structural- -MMuuurrt-SMillaittt)) are recorded for the image reconstruction evaluations.
Model
NParams
PSNR
SSIM
VIT
5.52M
+0.00%
22.11
V0.00
0.715
W0.00
ET
4.28M
22.46%%
21.51
V0.59
0.681
70.03
27

---

H Formal Algorithm for Training and Inference of ET
Algorithm I: Training and inference pseudocode of ET for image reconstruction task
HyperParameters
a: Energy descent stepsize
Learning
p: Token mask probability
b: batch size
Parameters
W E RXXXX W E ERXXXX; Key. Query kernels of the Energy Attention
CERMAD Kernel Hopfield Network
Thorm E R. ònana PR Scale, bias of LayerNorm
MASK RD Mask token
6pos R N Position bias, added to each token
Was E RRPxD dane E D: Kernel, bias of affine Encoder
12
Wi E DxP. Ose E PR Kernel, bias of affine Decoder
13
Infer
15
Inputs
L Corrupted image tokens X RND
17
Add position biases: X <- X + Opusi for timesteps f 1,...: do Normalize each token:
19
I F Layeerorrm(M X; Thorm Snom); Calculate Energy of tokens:
9 RXD
20
21
<- Energy Transformer(sr W W :));
22
e
X X aV
23
return X
25
Train
Inputs
Dataset Swain with elements gechnnele
27
Initialize
Randomly initialize from (O00 0.02):
W.K W E, MASK, W enc, W doc, Opus ~ (0,00 Set other biases zero: denc, ddec, Quorm +-
31
32
Set LayerNorm scale to one: "Ynorm f-
33
for epoch = =,,,,,..... 1,...,Nopoeh do
Sepoch Swrain
for batch B Sepoch
B R xexcctameexxxae xhw.caambbbhiiiiiixxxooo
36
37
Convert image into non non-ovene conoolellappingggg... patches:
Bpach +- Patchify(B):
Bpach E RRxNNxP
39
Embed image patches into tokens:
49
+- Encode Bpacni ene, 6mn));
X E R
41
Replace image tokens randomly by MASK:
42
X, Laaak f- Mask(XX (X; MASK, p)
X E RbxNDD Lmac E {0, 1
Reconstruct tokens with ET:
X + Infer(X))
44
Decode tokens:
Bpuch + Decode(X(]Tal Ware: ddec);
Bpach E R
Calculate MSE loss corrupted tokens:
L +- 4+MMean(1 (lean(|Bpuchh [mask] Bpusch mmssk] )2a)
R
49
params +- params eV params
50
Sepoch + Sepoch B
S1
return params
28


Reviews:
[{"review_id": "tJP0Aj9hXb", "reviewer": "Reviewer_q5cg", "summary": "The paper presents a new architecture for transformer (ET). The architecture is designed by first defining an energy function that forces the alignment of the query of a patch to the keys of a few other patches and to make each patch to look like a natural image patch. A dynamics (differential equation) is induced to minimize the energy function, giving rise to a transformer layer containing the novel attention mechanism and a Hopfield network component.  Experiments are performed using this new transformer to fill partially masked images and to detect graph anomalies and classify graphs. Effectiveness is demonstrated. ", "questions": "In the introduction (beginning of page 2), you mention that large storage capacity of associative memories is important for learning applications. I feel this is a correct statement intuitively however have not seen a rigorous argument connecting the storage capacity to learning. Could you explain the reasoning behind the statement or give a reference?", "limitations": "NONE", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The design is well motivated and clearly rationalized, nicely exploiting the recently observed connection between Hopfield networks/associative memories and transformers,. The proposed architecture is novel and interesting to this reviewer. \n\nThe experimental validation looks convincing. The interpretability aspect of the model relative to ViT is particularly appealing to this reviewer. The paper is also very well written, pleasant to read. The proposed structure is well positioned in relation to the literature of attention and associative memories. ", "weaknesses": "Despite its well-motivated design, the proposed structure still lacks rigorous theoretical justification for its effectiveness. \n\nI have a few minor editorial comments:\n\n1. In Eq (4), the notations g_{jB} are not explained until later in the paper.\n2. You perhaps need a notation for the range of \\alpha, otherwise in the unlabelled equation in page 6, sum over \\gamma is unclear. — If I am not mistaken,  \\gamma and \\alpha range in the same set. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Despite its well-motivated design, the proposed structure still lacks rigorous theoretical justification for its effectiveness.\n\nThank you for your insightful comments! We are glad to hear that you find ET's design well-motivated. We discuss the general theory of ET in Section 2. To summarize, ET is a unification of the empirically successful Transformer architecture with the theoretically rich literature around Associative Memories. This unification imbues ET with the following desired properties that existing transformers cannot claim:\n    \n- An ET block is a function $E: \\mathbf{x} \\mapsto \\mathbb{R}$ that computes a scalar energy value on a collection of $N$ input tokens each of dimension $D$, i.e., $\\mathbf{x} \\in \\mathbb{R}^{D \\times N}$. This energy is a Lyapunov function with guaranteed fixed points convergence.\n- The forward pass through ET is a continuous dynamical system on the tokens, defined through the update equation $\\tau \\frac{d x_{jB}}{dt} = - \\frac{\\partial E}{\\partial g_{jB}}$ for each element $j$ of token $B$. Thus, each token seeks to minimize the ET block's total energy. These dynamics are guaranteed to converge to a fixed point or, equivalently, a \"memory\" of the system.\n\n\n> In Eq (4), the notations $g_{jB}$ are not explained until later in the paper.\n\nThanks for pointing this out. We defined $g_{jB}$ in Eq (1) as element $j$ of token $B$'s `layernorm` representation. We will clarify and reemphasize this double-index notation right around Eq(4).\n        \n> You perhaps need a notation for the range of $\\alpha$, otherwise in the unlabelled equation in page 6, sum over $\\gamma$ is unclear. If I am not mistaken, $\\gamma$ and $\\alpha$ range in the same set.\n\nCorrect! The range of $\\alpha$ is defined in line 129, and the range of $\\gamma$ is the same. We will remind the readers about this notation around the unnumbered equation on page 6. In general, we use lowercase latin indices $i,j,k\\in 1...D$ for individual elements of token vectors, capital latin indices $A,B,C\\in 1...N$ for labeling the tokens vectors, greek letters from the beginning of the alphabet $\\alpha, \\gamma \\in 1...Y$ for the internal space of the attention operation. The parameter $\\beta$ is reserved for the inverse temperature of the softmax attention and is not an index. Finally, greek letters from the middle of the alphabet $\\mu,\\nu\\in 1...K$ are reserved for enumerating memory patterns in the Hopfield module. We will remind the readers about these notations after line 154 in the revised manuscript. Please also see Table 3 of Appendix A where we summarized all the notations that we used in the paper.  \n\n> Could you explain the statement: “large storage capacity of associative memories is important for learning applications”\n\nWhen training associative memories with the backpropagation algorithm on real-world data (so that there is a high degree of correlation between individual data points) an important desideratum is that similar data points land into the same basin of attraction, and dissimilar data points land into different basins of attraction. When training [old Hopfield networks](https://www.pnas.org/doi/10.1073/pnas.81.10.3088) it often happens that the entire training set falls into a single basin of attraction, so that it becomes impossible to differentiate between distinct patterns within the dataset. Replacing old Hopfield networks with modern Hopfield networks generally solves this problem if the network is given enough capacity to distribute the training samples among multiple basins of attraction. In our model data points are collections of tokens (e.g. a collection of patches in an image) and the network needs to memorize approximate rules of how patches fit together in an image (e.g. if there is a head of an animal - there must be a tail of that animal properly positioned in the image plane). The number of such \"rules\" is large for natural datasets. Therefore we expect that the network needs large memory capacity to properly represent these rules. \n\nSince the learning phenomenon by definition requires training the network on some natural dataset (even if just MNIST)--- it is difficult to theoretically quantify these statements, since theoretical tools almost always assume random uncorrelated patterns that are stored in associative memory. From an empirical perspective this phenomenon has been studied in a recent [work](https://arxiv.org/abs/2306.03209). Similarly to our model, the memory storage capacity of their model increases if the inverse temperature parameter $\\beta$ gets bigger. As evidenced by the results reported in Table 6 of that paper, the parameter $\\beta$ (which plays similar role to the parameter $\\beta$ in our Energy Attention module, Eq 3) needs to be sifficiently large in order for the network to cluster the data successfully (compare to the range of $\\beta$ in Table 5 of their work). \n"}, {"review_id": "nT0EwxuQX7", "reviewer": "Reviewer_bQx9", "summary": "This paper proposes a new transformer architecture namely Energy Transformer, consists of layer norm, multi-head energy attention and Hopfield network, within a single recurrent ET block. The ET block models the dynamics of token updates by minimising the sum of energy from both the multi-head attention and the Hopfield network.\n\nThe main contribution is the new designed ET architecture.", "questions": " - Is there any theoretical justifications that show the ET has either faster/slower convergence in training compared to a vanilla transformer/energy-based models, or ET obtains a better optimum than EBM after optimisation?\n\n - Why there are no quantitative experiments on e.g. image generation? How ET performs when the input image is 100% masked(generate from noise)? Can ET be applied to image anomaly detection as in [1]?", "limitations": "I don't see a detailed discussion on limitations of the work.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a novel transformer architecture, which uses the Hopfield network, and models the dynamics between tokens with a time differential equation. The proposed architecture seems promising in the applications for CV, NLP and other domains. ", "weaknesses": "The experiments can not fully support the claim. \n\n - Empirical experiments on image reconstruction is hard to quantify and are also not compared to any of the SOTA algorithms. The value of the reconstruction error shown in the Appendix is hard to understand/evaluate without comparing to any other algorithms. \n\n - For instance, FID scores in most of the image generation tasks, or AUC in [1] can also be quantitative measures for the ET. \n\n - Transformers are widely used in NLP tasks, but this is almost completely ignored in the discussion. Since there is no theoretical justification for the proposed architecture, showing empirical results on various applications (including but not limited to NLP) may help to convince the reader. \n\n - In my humble opinion, quantitative experiments on graph anomaly detection and classification are not sufficient to show the full power of such generative models. \n\n[1] Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. In NeurIPS, 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Empirical experiments on image reconstruction is hard to quantify and are also not compared to any of the SOTA algorithms. The value of the reconstruction error shown in the Appendix is hard to understand/evaluate without comparing to any other algorithms. \n\nThank you for your valuable comment. Our goal in presenting images is to foster an in-depth understanding of the inner workings of our model. Furthermore, the results on MSE loss are rarely reported in Computer Vision papers. Typically, studies report results on downstream tasks (e.g., image segmentation, object detection, image classification, etc.) rather than for model inspection. We want to emphasize that our work is not a full scale Computer Vision paper and the results on images are only included for intuitive understanding of our model.\n\n\n> For instance, FID scores in most of the image generation tasks, or AUC in [[1]](https://papers.nips.cc/paper_files/paper/2019/hash/378a063b8fdb1db941e34f4bde584c7d-Abstract.html) can also be quantitative measures for the ET. \n\nOur setting is different from that studied in [1] in several ways. First, our image model was trained on a fixed masking ratio (50%), and would not work for fully occluded images (as a generative model in the conventional sense). In other words, the open patches are required to meaningfully guide the temporal evolution of the masked patches. Second, only unlabeled images were used in our setting and conditional generation requires model modifications and retraining. Third, we used 224x224 images, while the largest one available in [1] is 128x128. Having these caveats in mind, we have run FID evaluation of our model and obtained FID score = 40.22 on the validation set of ImageNet 224x224. Images are masked out by 50% (patch-wise) and ET performs image restoration through temporal dynamics. The restored images are then compared to the original images using FID. This number can be somewhat compared with the results in the table of Fig. 4 of Ref [1], e.g., CIFAR10 unconditional ranging 21.7-65.93, and ImageNet 128x128 conditional ranging 27.62-43.7. Generally, unconditional generation is harder than conditional (typically leads to higher FID score), and generation of larger images (224x224) is harder than generation of smaller images (128x128). For these reasons we believe that the FID score of 40.22 on unconditionally generated images of size 224x224 is a very decent number. \n\n\nThe AUC setting in [1] is fundamentally different from ours. They use anomaly detection on entire images. In our setting, node anomaly detection (studied in section 4 of our paper) is analogous to patch anomaly detection where each patch in the image receives the binary label indicating its anomaly status. We are not aware of any datasets that would use this setting. \n\n> Transformers are widely used in NLP tasks, but this is almost completely ignored in the discussion. Since there is no theoretical justification for the proposed architecture, showing empirical results on various applications (including but not limited to NLP) may help to convince the reader. \n\nOur paper is not a full scale NLP study, but in order to partly address your question we have trained our model on the Wikitext dataset. We have included examples of input sentences with masked tokens together with top-5 predictions for the masked token and the corresponding probabilities of those words (please see the pdf file attached to the general response). We conclude that our model qualitatively outputs meaningful predictions for the masked tokens.\n\n> Is there any theoretical justifications that show the ET has either faster/slower convergence in training compared to a vanilla transformer/energy-based models, or ET obtains a better optimum than EBM after optimisation?\n\nWe are not aware of such results and do not see an easy path towards such a claim. We hope that our work inspires exploration of these exciting questions in subsequent publications. \n\n> Why there are no quantitative experiments on e.g. image generation? How ET performs when the input image is 100% masked(generate from noise)? Can ET be applied to image anomaly detection as in [[1]]((https://papers.nips.cc/paper_files/paper/2019/hash/378a063b8fdb1db941e34f4bde584c7d-Abstract.html))?\n\nET requires an appreciable number of open tokens to guide the temporal evolution of masked tokens. Our energy attention mechanism can route the information from open tokens to masked tokens, but cannot generate information from nowhere. Please notice, that there is no noise injection in our pipeline, unlike Ref [1]. For this reason starting with a fully masked image will deterministically generate a homogeneous gray image.\n\nAs we explained above, ET cannot be applied to image anomaly detection in an easy way. The results of section 4 (node anomaly detection) would be translated to patch anomaly detection in the image domain, not image anomaly detection. Additionally, there are no baselines which we could compare with for patch anomaly detection on images. "}, {"review_id": "YkAnk28r0X", "reviewer": "Reviewer_zsPk", "summary": "This paper proposes a new Transformer-like architecture which is is inspired by modern Hopfield networks, and falls into the category of energy-based models. The authors argue that the existence of an energy function leads to several advantages such as more principled design and theoretical interpretation. The energy transformer (ET) they design a recurrent network similar to a deep transformer architecture with shared weights, however with the additional property that the recurrent dynamics is shown to minimize an energy function, and can be seen as an instance of a Hopfield network. The authors evaluate the performance of the architecture on image completion (qualitative), graph anomaly detection and graph classification tasks, demonstrating strong and competitive performance.\n\n\nPlease note that due to receiving 6 papers to review from NeurIPS alone, I have allocated a time budget of 4h per paper, and my review is based on that. In particular I did not study the appendix in great depth. I regret this situation and apologize for possible inaccuracies.\n", "questions": "- In lines 148-152, the functional similarity between the HN sub-block and a MLP seems plausible but was not fully clear to me. Could this discussion be made more explicit in terms of equations?\n- The two descriptions of the difference to previous modern Hopfield networks in lines 172-175 and lines 183-184 did not seem consistent to me: The first states that the difference is that the keys and queries are dynamical variables, while the second states the main difference is the additional symmetrizing term in the attention mechanism.\n- in line 192, the text states \"e.g. by replacing them with a learnable MASK token\". The \"e.g.\" makes it unclear in how far this is a generic example, and what was actually done in the experiment.\n- In the architecture used in sect. 4, the main text does not explain how the adjacency structure of the graph enters the model. From the appendix I understand that the attention operation is constrained to the 1-step environment of the node. Maybe this would be good to discuss in the main text.\n- In general, the ET block can be understood as a deep transformer model, however with shared weights instead of independent weights in each layer. What limitations does this imply on the type of tasks which are well suitable for the ET architecture? Is the stacked model of used in sect. 5 a result of the need to overcome the shared weight limitation?\n", "limitations": "While some weaknesses in the empirical performance are pointed out, e.g. lines 200-203, there is no broader discussion of technical and theoretical limitations.\n", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Clear presentation\n- The combination of attention mechanism and Hopfield networks makes this architecture interesting for theoreticians.\n- Strong evaluation results on the graph anomaly detection and graph classification tasks.\n", "weaknesses": "- Insufficient discussion of limitations.\n- Evaluation on the masked image completion task is not quantitative, and not compared to related architectures/baselines such as ViT.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your enthusiastic evaluation of our work! \n> Insufficient discussion of limitations.\n\nThank you for pointing this out, we will add the following discussion to the revised manuscript. So far, the proposed model has only been empirically tested on the graph domain. It remains unclear whether or not it would provide strong results on other domains such as images or NLP (although we do see promising qualitative behaviors on images reconstructions and NLP masked token reconstructions in the pdf attached to this rebuttal). We hope to deeply look into other data domains in subsequent publications. \n\nDespite ET's smaller memory footprint (in parameters) compared to feed-forward transformers, the energy-ATTN operation requires twice more flops than conventional attention. Although this is just a constant factor that doesn't grow with the number of tokens, in certain flop-limiting applications this factor could become important. \n\nThe nature of representations learned by ET (global vs. local) requires more investigation. \n\nDespite these limitations, our work is the first one to propose a rigorous mathematical formulation of the ET network, that unites transformers, associative memory, and some of the EBMs ideas into a single well-defined neural architecture. We believe that this result on its own justifies acceptance of our work.\n\n> Evaluation on the masked image completion task is not quantitative, and not compared to related architectures/baselines such as ViT.\n\nIndeed, our evaluation on masked image reconstructions is intended to visually explain the inner workings of our architecture and not to show superiority to existing transformers. This said, we have taken a ViT with a comparable number of parameters to ET and trained it on the same image reconstruction task. We report the PSNR (Peak-Signal-to-Noise-Ratio) and SSIM (Structural-SIMilarity) of ViT vs. ET in Table 1 of the attached pdf. We see that the numbers are in the same ballpark, though the ViT slightly outperforms ET (please also notice that ViT still has approximately 25% more parameters than ET). \n\n\n> In lines 148-152, the functional similarity between the HN sub-block and a MLP seems plausible but was not fully clear to me. Could this discussion be made more explicit in terms of equations?\n\nLet us show this for one token (index $B=1$). Since the Hopfield module does not mix tokens, the following equations will hold for any number of tokens. The contribution of $E^{\\text{HN}}$ (Eq 5) to dynamics (Eq 6) is \n\n$$\n\\tau\\frac{dx_i}{dt} = - \\frac{\\partial E^{\\text{HN}}}{\\partial g_i} = \\sum\\limits_{\\mu} \\xi_{\\mu i} G^\\prime (\\sum\\limits_j \\xi_{\\mu j} g_j)\n$$\n\nOn the other hand, the MLP in a Transformer block would give \n\n$$\n\\tau\\frac{dx_i}{dt} = \\sum\\_{\\mu} \\tilde{W}\\_{i \\mu} r ( \\sum\\_j W\\_{\\mu j} g_j)\n$$\n\nThese two equations are identical if $G^\\prime(\\cdot) = r(\\cdot)$ (i.e., $r =\\mathrm{relu}$ as in conventional transformers) and $\\mathbf{W} = \\mathbf{\\xi} = \\mathbf{\\tilde{W}^T}$. Thus, the Hopfield Network module can be viewed as an MLP with shared weights.\n        \n> The two descriptions of the difference to previous modern Hopfield networks in lines 172-175 and lines 183-184 did not seem consistent to me: The first states that the difference is that the keys and queries are dynamical variables, while the second states the main difference is the additional symmetrizing term in the attention mechanism.\n\nThese two statements are equivalent. The first term in line 177 appears when the derivative $\\frac{\\partial}{\\partial g_{iA}}$ (in Eq 6) acts on the queries $\\mathbf{Q}$, while the second symmetrizing term appears when the same derivative acts on keys $\\mathbf{K}$. Thus, the second symmetrizing term appears only because the keys are promoted to dynamical variables. If the keys were constant (reference to line 172-175) the second symmetrizing terms would disappear (reference to lines 183-184). \n\n\n> in line 192, the text states \"e.g. by replacing them with a learnable MASK token\". The \"e.g.\" makes it unclear in how far this is a generic example, and what was actually done in the experiment.\n\nApologies for the confusion, this statement is clarified in Appendix B (lines 489-493). We followed the standard training procedure of ViTs -- given a MASKing ratio of 50%, we set 40% of all tokens equal to a learnable MASK token and the remaining 10% we leave untouched but still consider them in the loss function. We will explain this in the main text too.\n\n> In the architecture used in sect. 4, the main text does not explain how the adjacency structure of the graph enters the model. From the appendix I understand that the attention operation is constrained to the 1-step environment of the node. Maybe this would be good to discuss in the main text.\n\nYes, agreed, thank you! We will add this discussion to the main text.\n\n> In general, the ET block can be understood as a deep transformer model, however with shared weights instead of independent weights in each layer. What limitations does this imply on the type of tasks which are well suitable for the ET architecture? Is the stacked model of used in sect. 5 a result of the need to overcome the shared weight limitation?\n\nWe stacked ET blocks in part because we saw small performance gains using a stacked model on graph classification but also to show that we can modularly embed this block into deeper architectures. Additionally, in classification tasks it might be helpful to have several layers with different sets of weights. This may help transform input data into labels by passing the data through several layers with independent weights."}, {"review_id": "nGVEWqmH7B", "reviewer": "Reviewer_qJXt", "summary": "This paper presents a new transformer architecture that is based on minimizing energy with two networks viewed as Hopfield networks (one for the attention within the data point, and another for global features). The result energy-based model resembles (and perhaps motivates) empirical choices in the vanilla transformer, and unlocks interpretations of the learned features, as well as better performance on graph-based benchmarks. The authors present the idea that their energy transformer could be useful for other domains, such as computer vision, NLP, audio, etc. ", "questions": "Q1: How does your visualization from Figure 2 relate to patterns learned by CNNs?\n\nQ2: How does the MLP weights relate to the weights $\\xi_{\\mu j}$ in Eqation 5? Can you explain in more detail?\n\nQ3: Have you tested the crucial importance (for minimizing the energy) of the second term of the RHS in the equation of Line 177? I think you have done that in the Appendix (in the SM). Could you comment on the findings of your ablations a bit?\n\nQ4: Have you tried downstream tasks of the representations for ImageNet? It would be interesting to compare with MAE. I have a suspicion that the features are not going to be useful for transfer learning because it seems like they only learn local information (the example with the missing eye for the dog that you give in Fig. 3 is insightful). Then, my question becomes: Why is MAE (speculatively) more useful than ET in terms of obtaining features for transfer learning? How can we let ET learn more global information?\n\nQ5: Can you show reported results from previous works in Table 1?", "limitations": "the authors adequately addressed the limitations", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1: The problem is highly relevant. Finding better transformer architectures is important in modern machine learning. The ET is a reasonable and well-motivated contribution. \n\nS2: The improvements on graph benchmarks seem to be really strong. \n\nS3: I like the visual interpretation of the intermediate states of the tokens. \n\nS4: The paper is very well written.", "weaknesses": "W1: Limited evaluation (only on graphs). I suspect the authors didn't evaluate large-image representation learning because of the issue of learning global features (see question Q4 below). I am not sure why the authors didn't try NLP. I think trying NLP is a very crucial step for convincing the community that ET works.\n\nW2: Missing discussion about the computational cost of the architecture (at least I do not see it in the main text). How easy it is to implement your method and scale it? The answers to these questions should be a part of the main text. They are also integral to convincing the community about the scalability of ET.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Limited evaluation (only on graphs)... I am not sure why the authors didn't try NLP. I think trying NLP is a very crucial step for convincing the community that ET works.\n\nThank you for your valuable comment. We have included a more thorough discussion on this in our general response and in our response to your Q4 below. In summary, we believe the graph domain is a perfectly valid domain on which to conduct quantitative evaluation of our ideas. Additionally, our initial qualitative results on NLP (attached pdf) suggest that NLP is a promising domain to further test this architecture.\n\n> Missing discussion about the computational cost of the architecture... How easy it is to implement your method and scale it?\n\nThere are two metrics describing computational costs: memory footprint of the model and the number of flops. Regarding the number of parameters (memory footprint) our model wins (is smaller) compared to feedforward transformers with independent weights and 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 (energy attention is quadratic in the number of tokens, like conventional attention), but the energy attention has a constant factor of two more flops. The reason for this overhead can be understood from the unnumbered equation on page 6. The second term in the attention equation has the same complexity as the first, but doubles the overall number of flops required to compute energy-based attention.  We will add a few sentences about this in the revised manuscript.\n\n> How does your visualization from Figure 2 relate to patterns learned by CNNs?\n\nIndeed, they look very similar. We tried to emphasize this visual similarity in lines 212-216 and comparing them with ref [22] for CNNs.\n\n> How does the MLP weights relate to the weights in Eqation 5? Can you explain in more detail?\n\nLet us show this for one token (index $B=1$). Since the Hopfield module does not mix tokens, the following equations will hold for any number of tokens. The contribution of $E^{\\text{HN}}$ (Eq 5) to dynamics (Eq 6) is $$\\tau\\frac{dx_i}{dt} = - \\frac{\\partial E^{\\text{HN}}}{\\partial g_i} = \\sum\\limits_{\\mu} \\xi_{\\mu i} G^\\prime (\\sum\\limits_j \\xi_{\\mu j} g_j)$$On the other hand, the conventional MLP would give $$\\tau\\frac{dx_i}{dt} =  \\sum\\limits_{\\mu} \\tilde{W}\\_{i \\mu} r (\\sum\\limits_j W_{\\mu j} g_j)$$ These two equations are identical if $G^\\prime(\\cdot) = r(\\cdot)$ (i.e., $r =\\mathrm{relu}$ as in conventional transformers) and $\\mathbf{W} = \\mathbf{\\xi} = \\mathbf{\\tilde{W}^T}$. Thus, the Hopfield Network module can be viewed as an MLP with shared weights. These are the relationships between the weights in Eq 5 and the MLP weights.\n\n> Have you tested the crucial importance (for minimizing the energy) of the second term of the RHS in the equation of Line 177?... Could you comment on the findings of your ablations a bit?\n\n\nThe ablation studies in Tables 13,14 (Appendix F2) explore how the *energy function* of ET should be designed (e.g., Do we include self-attention in Eq. (3)? Do we benefit from adding the energy of an HN as described in Eq. (5)?) rather than how the *update rule* (attention step) should be designed. The update rule is the exact gradient of the designed energy function. It is impossible to define an energy function for attention where the second term of the equation in line 177 is not present. The top-two performing configurations (`ET-Full` and `ET-Full+Self` of Table 14) have both an ATTN and HN module that are further broken down and meaningfully interpreted (see Appendix B3 lines 539-548 and Fig. 7).\n \n> Have you tried downstream tasks of the representations for ImageNet? It would be interesting to compare with MAE. I have a suspicion that the features are not going to be useful for transfer learning because it seems like they only learn local information... Why is MAE (speculatively) more useful than ET in terms of obtaining features for transfer learning? How can we let ET learn more global information?\n\nOn images our model design indeed seems to favor aligning with local patches over global patches which could cause issues on downstream tasks (on graphs this behavior is enforced - attention is constrained to the 1-hop environment of each node). Like the ViT we represent all image patches as a bag of tokens where position is encoded as a learnable vector added to each patch representation. The correlations between these positions (see Fig. 2 middle and Fig. 6) provide a strong local signal and favor the expected structure of image patches connectivity. In the graph classification setting, however, our model attends to all the nodes. Therefore we expect it to benefit from global information.\n\nMAE uses an encoding module that creates meaningful representations from only visible patches in the image. This approach has two benefits: it is more computationally efficient (using fewer tokens), and it accumulates global image information into visible patches *before* introducing masked tokens that evolve from them. ET can be viewed as the decoder step of MAE where mask and visible tokens evolve together from an initial representation that captures only local pixel information. Indeed, the MAE approach of first encoding global information into the visible patches (either by performing ET dynamics on visible patches first, or by using a separate ET block to first encode the global relations to visible patches) seems highly beneficial to ET applied on downstream tasks.\n\n> Can you show reported results from previous works in Table 1?\n\nIn fact, this is what Table 1 shows. It compares ET to published results from relevant previous works (GraphConsis, CAREGNN, PC-GNN, BWGNN, MLP). Consider Table 4 from Ref [27]: the only difference is that we rounded up their numbers to 1 decimal in our table. GT was implemented ourselves, since that baseline is missing from the literature. "}]

Author rebuttal:
We sincerely thank all reviewers for their thoughtful feedback. We are excited that they highlight the strengths of the Energy Transformer (ET):

1. A **novel architecture** (q5cg, bQx9, zsPk, qJXt) 
2. with a **well-motivated design** (q5cg, qJXt) built on top of Associative Memory fundamentals
3. that is **highly interpretable** (q5cg, qJXt) throughout the entire dynamics
4. with **strong evaluation results** on graph classification tasks (zsPk, qJXt)
5. that is **interesting for theoreticians** (zsPk) 

In this general response, we would like to address two comments that have been expressed by more than one reviewer.

1. **What is the theory supporting ET?** (q5cg, bQx9)

    Thank you for providing this opportunity to clarify the theoretical contributions of our work. We discuss the general theory of ET in Section 2. To summarize, ET is a unification of the empirically successful Transformer architecture with the theoretically rich literature around Associative Memories. This unification imbues ET with the following desired properties that existing transformers cannot claim:
    
    - An ET block is a function $E: \mathbf{x} \mapsto \mathbb{R}$ that computes a scalar energy value on a collection of $N$ input tokens each of dimension $D$, i.e., $\mathbf{x} \in \mathbb{R}^{D \times N}$. This energy is a Lyapunov function with guaranteed fixed points convergence.
    - The forward pass through ET is a continuous dynamical system on the tokens, defined through the update equation $\tau \frac{d x_{jB}}{dt} = - \frac{\partial E}{\partial g_{jB}}$ for each element $j$ of token $B$. Thus, each token seeks to minimize the ET block's total energy. These dynamics are guaranteed to converge to a fixed point or, equivalently, a "memory" of the system.


2. **Where are quantitative results on NLP?** (bQx9, qJXt)

    Our work is the first to combine attention, energy, and associative memory into a single novel ET architecture that offers unparalleled interpretability while operating as an attractor network under continuous time dynamics. We have chosen the graph domain for quantitative evaluation of this idea because it is natural to perceive ET as an efficient compute engine to learn relationships between the nodes of large and possibly sparse input graphs. In addition, we illustrate qualitative behavior on images to help readers benefit from more intuitive and visual understanding of our proposed architecture (as we have benefitted ourselves). Lastly, we want to emphasize that ML on graphs is a burgeoning area of research with a large interest behind it; we do not believe that graphs are any less valuable than images or NLP.

    We do agree with the reviewers that NLP is an important domain, and we believe our first results presented in this work will inspire the community to extend our ideas to more domains including NLP. This said, we quickly trained our model on the Wikitext-103 dataset (which is tiny relative to other NLP datasets). In Figure 1 of the attached pdf, we have included examples of input sentences with masked tokens together with top-5 predictions (and their corresponding probabilities) for each masked token. We conclude that our model outputs qualitatively meaningful predictions on masked tokens. More thorough hyperparameter tuning and evaluations are outside the scope of this submission. 
    

We also want to emphasize that our work is the first to propose a rigorous mathematical formulation of the ET network, that unites transformers, associative memory, and some of the ideas of EBMs into a single, well-defined neural architecture. In our humble opinion this result on its own justifies acceptance of our work. 


We respond individually to reviewers' additional questions and comments.

Meta-review: Four reviewers evaluated the submission and the scores are 8764. The strengths are

* a novel architecture with a well-motivated design
* strong results in graph-related experiments
* paper is very well-written

The only negative score is based on the lack of SoTA results in image reconstruction and experiments in NLP tasks. However, after examining the paper and reviews more closely, the area chair determined these additional experiments were not necessary to support the paper's core claims. Given these considerations, I recommend acceptance.
Decision: Accept (poster)