Energy Transformer — part 2/5
Cleaned-text code points [14613, 29654). Bibliography omitted here; source.md includes it.
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 in realistic images. The energy of this sub-block is defined as: E HN = − ∑ B = 1 N ∑ μ = 1 K G ( ∑ j = 1 D ξ μ j g j B ) , ξ ∈ ℝ K × D E^{\text{HN}}=-\sum\limits_{B=1}^{N}\sum\limits_{\mu=1}^{K}G\Big(\sum\limits_{j=1}^{D}\xi_{\mu j}\;g_{jB}\Big),\ \ \ \ \ \ \ \ \ \ \mathbf{\xi}\in\mathbb{R}^{K\times D} (5) where ξ μ j \xi_{\mu j} is a set of learnable weights (memories in the Hopfield Network), and G ( ⋅ ) G(\cdot) is an integral of the activation function r ( ⋅ ) r(\cdot) , so that G ( ⋅ ) ′ = r ( ⋅ ) G(\cdot)^{\prime}=r(\cdot) . 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 ( ⋅ ) = r(\cdot)= ReLU), or as a modern continuous Hopfield Network [ 10 , 12 , 19 ] if the activation function is sharply peaked around the memories (e.g., r ( ⋅ ) = r(\cdot)= power, or softmax). The HN sub-block is analogous to 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) as the weights of the subsequent projection from the hidden space to 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 ξ \mathbf{\xi} , which represent memories, and high otherwise.
### 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 τ d x i A d t = − ∂ E ∂ g i A , where E = E ATT + E HN \tau\frac{dx_{iA}}{dt}=-\frac{\partial E}{\partial g_{iA}},\ \ \ \ \ \text{where}\ \ \ \ \ E=E^{\text{ATT}}+E^{\text{HN}} (6) Here x i A x_{iA} is the token representation (input and output from the ET block), and g i A g_{iA} 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 ξ \mathbf{\xi} ). The dynamical system ( 6 ) finds a trade-off between these two desirable properties of each token’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 d E d t = ∑ i , j , A ∂ E ∂ g i A ∂ g i A ∂ x j A d x j A d t = − 1 τ ∑ i , j , A ∂ E ∂ g i A M A i j ∂ E ∂ g j A ≤ 0 \frac{dE}{dt}=\sum\limits_{i,j,A}\frac{\partial E}{\partial g_{iA}}\ \frac{\partial g_{iA}}{\partial x_{jA}}\ \frac{dx_{jA}}{dt}=-\frac{1}{\tau}\sum\limits_{i,j,A}\frac{\partial E}{\partial g_{iA}}\ M^{A}_{ij}\ \frac{\partial E}{\partial g_{jA}}\leq 0 (7) The last inequality sign holds if the symmetric part of the matrix M i j A = ∂ g i A ∂ x j A = ∂ 2 L ∂ x i A ∂ x j A M^{A}_{ij}=\frac{\partial g_{iA}}{\partial x_{jA}}=\frac{\partial^{2}L}{\partial x_{iA}\partial x_{jA}} (8) is positive semi-definite (for each value of index A A ). The Lagrangian ( 2 ) 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 ( 3 ). 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 a continuous Hopfield Network with the softmax \mathrm{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 in 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): − ∂ E ATT ∂ g i A = ∑ C ≠ A ∑ α W α i Q K α C softmax 𝐶 ( β ∑ γ K γ C Q γ A ) + W α i K Q α C softmax 𝐴 ( β ∑ γ K γ A Q γ C ) -\frac{\partial E^{\text{ATT}}}{\partial g_{iA}}=\sum\limits_{C\neq A}\sum\limits_{\alpha}W^{Q}_{\alpha i}\;K_{\alpha C}\;\underset{C}{\text{softmax}}\Big(\beta\sum\limits_{\gamma}K_{\gamma C}\;Q_{\gamma A}\Big)+W^{K}_{\alpha i}\;Q_{\alpha C}\;\underset{A}{\text{softmax}}\Big(\beta\sum\limits_{\gamma}K_{\gamma A}\;Q_{\gamma C}\Big) In both terms the softmax \mathrm{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 [ 5 ] with the value matrix equal to 𝐕 = ( 𝐖 Q ) T 𝐊 = ∑ α W α i Q K α C \mathbf{V}=(\mathbf{W}^{Q})^{T}\mathbf{K}=\sum_{\alpha}W^{Q}_{\alpha i}K_{\alpha C} . 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 to 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 is 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 − ∂ E HN ∂ g i A = ∑ μ = 1 K ξ μ i G ′ ( ∑ j = 1 D ξ μ j g j A ) -\frac{\partial E^{\text{HN}}}{\partial g_{iA}}=\sum_{\mu=1}^{K}\xi_{\mu i}G^{\prime}\Big(\sum_{j=1}^{D}\xi_{\mu j}g_{jA}\Big) (9) which is applied to every token individually (there no mixing of different tokens).
## 3 Qualitative Inspection of the ET framework on ImageNet
We trained * * * The code is available: https://github.com/bhoov/energy-transformer-jax . the ET network on the masked image completion task using ImageNet-1k dataset [ 22 ] . Each image was broken into non-overlapping patches of 16x16 RGB pixels, which were projected with a single affine encoder into the token space. Half of these tokens were “masked”, by replacing them with a learnable MASK token. A distinct learnable position encoding vector was added to each token. Our ET block then processes all tokens recurrently for T T steps. The token representations after T T steps are passed to 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 B .
Examples of occluded/reconstructed images (unseen during training) are shown in Figure 3 . In general, our model learns to perform the task very well, capturing the texture in dog fur (column 3) and understanding meaningful boundaries of objects. However, we observe that our single ET block struggles to understand some global structure, e.g., failing to capture both eyes of the white dog (column 5) and completing irregular brick patterns in the name of extending the un-occluded borders (last column). We additionally inspect the positional encoding vectors associated with every token, Figure 2 , where the model learns a locality structure in the image plane that is very similar to the original ViT [ 18 ] . The positional embedding of each image patch has learned higher similarity values for neighboring tokens than for distant tokens.
Our network is unique compared to standard ViTs in that the iterative dynamics only move tokens around in the same space from which the final fixed point representation can be decoded back into the image plane. This functionality makes it possible to visualize essentially any token representation , weight , or gradient of the energy directly in the image plane. This feature is highly desirable from the perspective of interpretability, since it makes it possible to track the updates performed by the network directly in the image plane as the computation unfolds in time. In Figure 2 this functionality is used for inspecting the learned weights of the HN module directly in the image plane. According to our theory, these weights should represent basis vectors in the space of all possible image patches. These learned representations look qualitatively similar to the representations typically found in networks trained on image datasets, e.g., [ 23 ] .
We additionally visualize the gradients of the energy function (which are equal to the token updates, see Equation 6 ) of both ATTN block and the HN block, see Figure 4 . Early in time, almost all signal to the masked tokens comes from the ATTN block, which routes information from the open patches to the masked ones; no meaningful signal comes from the HN block to the masked patch dynamics. Later in time we observe a different phenomenon: almost all signal to masked tokens comes from the HN module while ATTN contributes a blurry and uninformative signal. Thus, the attention module is crucial early in the network dynamics, feeding signal to masked patches from the visible patches, whereas the HN is crucial later in the dynamics as the model approaches the final reconstruction, sharpening the masked patches. All the qualitative findings presented in this section are in accord with the core computational strategy of the ET block as it was designed theoretically in section 2 .
## 4 Graph Anomaly Detection
Having gained empirical insights in the image domain, we turn to quantitatively evaluating ET’s performance on the graph anomaly detection problem * * * The code is available: https://github.com/zhuergou/Energy-Transformer-for-Graph-Anomaly-Detection/ . , a task with plenty of strong and recently published baselines. Graph Convolutional Networks (GCN) [ 25 ] have been widely used for this task due to their capability of learning high level representations of graph structures and node attributes [ 26 , 27 ] . However, vanilla GCNs suffer from the over-smoothing problem [ 28 ] . In each layer of the forward pass, the outlier node aggregates information from its neighbors. This averaging makes the features of anomalies less distinguishable from the features of benign or normal nodes. Our approach does not suffer from this problem, since the routing of the information between the nodes is done through the energy based attention, and the information aggregation is based on the attention scores.
For anomaly detection on graphs in the ET framework, consider an undirected graph with N N nodes. Every node has a vector of raw attributes 𝐲 A ∈ ℝ F \mathbf{y}_{\!A}\in\mathbb{R}^{F} , where F F is the number of node’s features. Every node also has a binary label l A l_{A} indicating whether the node is benign or anomalous. We focus on node anomaly and assume that all edges are trusted. The task is to predict the label of a node given the graph structure and the node’s features. Since there are far more benign nodes in the graph than anomalous ones, anomaly detection can be regarded as an imbalanced node classification task.
First, similarly to images, the feature vectors for every node are converted to a token representation using a linear embedding 𝐄 \mathbf{E} and adding a learnable positional embedding λ A \mathbf{\lambda}_{A} 𝐱 A t = 1 = 𝐄 𝐲 A + λ A \mathbf{x}_{A}^{t=1}=\mathbf{E}\;\mathbf{y}_{\!A}+\mathbf{\lambda}_{A} (10) where the superscript t = 1 t=1 indicates the time of the update of the ET dynamics. This token representation is iterated through the ET block for T T iterations. When the retrieval dynamics becomes stable, we have the final representation for each node 𝐱 A t = T \mathbf{x}_{A}^{t=T} (or more precisely 𝐠 A t = T \mathbf{g}_{A}^{t=T} , since the outputs are additionally passed through a layer norm operation after the final ET update). This output is concatenated with the initial (layer normalized) token to form the final output 𝐠 A final = 𝐠 A t = 1 | | 𝐠 A t = T \mathbf{g}_{A}^{\text{final}}=\mathbf{g}_{A}^{t=1}\ ||\ \mathbf{g}_{A}^{t=T} (11) where | | || denotes concatenation. Following [ 24 ] , the node representation 𝐠 A final \mathbf{g}_{A}^{\text{final}} is fed into an MLP with a sigmoid activation function to compute the anomaly probabilities p A p_{A} . The weighted cross entropy Loss = ∑ A [ ω l A log ( p A ) + ( 1 − l A ) log ( 1 − p A ) ] \text{Loss}=\sum_{A}\Big[\omega\;l_{A}\log(p_{A})+(1-l_{A})\log(1-p_{A})\Big] (12) is used to train the network. Above, ω \omega is the ratio of the benign labels ( l A l_{A} = 0) to anomalous ones ( l A l_{A} = 1). Attention is restricted to 1-hop neighbors of the target node.
### 4.1 Experimental Evaluation
Four datasets are used for the experiments. YelpChi dataset [ 29 ] aims at opinion spam detection in Yelp reviews. Amazon dataset is used to detect anomalous users under the Musical Instrument Category on amazon.com [ 30 ] . T-Finance and T-Social datasets [ 24 ] are used for anomalous account detection in transactions and social networks, respectively. For these four datasets, the graph is treated as a homogeneous graph (i.e., all the edges are of the same type), and a feature vector is associated with each node. The task is to predict the label (anomaly status) of the nodes. For each dataset, either 1 % 1\% or 40 % 40\% of the nodes are used for training, and the remaining 99 % 99\% or 60 % 60\% are split 1 : 2 1:2 into validation and testing sets, see Appendix C for details.
We compare with state-of-the-art approaches for graph anomaly detection, which include GraphConsis [ 31 ] , CAREGNN [ 32 ] , PC-GNN [ 33 ] and BWGNN [ 24 ] . Additionally, multi-layer perceptrons (MLP) and Graph Transformer (GT) [ 34 ] are included in the baselines for completeness. Following previous work, macro-F1 score (unweighted mean of F1 score) and the Area Under the Curve (AUC) are used as the evaluation metrics on the test datasets [ 35 ] . See Appendix C for more details on training protocols and the hyperparameters choices. The results are reported in Table 1 . Our ET network demonstrates very strong results across all the datasets.