Energy Transformer — part 1/5

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

# Energy Transformer

###### 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.

## 1 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 of 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 a taxonomy of different models. At present, however, the search for new transformer architectures is driven mostly by empirical evaluations, and the theoretical principles behind this growing list of architectural variations is missing.

Additionally, the computational role of the four elements remains the subject of discussions. Originally, [ 5 ] emphasized attention as 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 , 8 ] , have been gaining popularity in the machine learning community thanks to theoretical advancements pertaining to their memory storage capacity and novel architectural modifications [ 9 ] . Specifically, it has been shown that increasing the sharpness of the activation functions can lead to super-linear [ 10 ] and even exponential [ 11 ] 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 \mathrm{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 a 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 a form of Majorization-Minimization algorithms [ 16 ] to interpret the forward path in the transformer block as an 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 complementary 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 Associative 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 tokens describing a 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 mathematical 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 conventional transformer blocks is replaced with a single ET block, which iterates the token representations until they converge to a 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 a node given its neighbors, see Figure 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 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 is different from the attention mechanism commonly used in feed-forward transformers [ 5 ] . Lastly, our network may be viewed as an example of a broader class of Energy-Based Models [ 17 ] frequently discussed in the AI 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” as nodes on the graph.

The overall pipeline is similar to the Vision Transformer networks (ViTs) [ 18 ] and is shown in Figure 1 . 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 x i ​ A x_{iA} . In the following the indices i , j , k = 1 ​ … ​ D i,j,k=1...D are used to denote the token vector’s elements, indices A , B , C = 1 ​ … ​ N A,B,C=1...N are used to enumerate the patches and their corresponding tokens. It is helpful to think about each image patch as a physical particle, which has a complicated internal state described by a D D -dimensional vector 𝐱 A \mathbf{x}_{A} . This internal state describes the identity of the particle (representing the pixels of each patch), and the particle’s positional embedding (the patch’s location within the image). The ET block is described by a continuous time differential equation, which describes interactions between these particles. Initially, at t = 1 t=1 the network is given a set containing two groups of particles corresponding to 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 is designed so that it minimizes a global energy function, and is guaranteed to arrive at a 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, is represented by a vector 𝐱 ∈ ℝ D \mathbf{x}\in\mathbb{R}^{D} . At the same time, most of the operations inside the ET block are defined using a layer-normalized token representation g i = γ ​ x i − x ¯ 1 D ​ ∑ j ( x j − x ¯ ) 2 + ε + δ i , where x ¯ = 1 D ​ ∑ k = 1 D x k g_{i}=\gamma\frac{x_{i}-\bar{x}}{\sqrt{\frac{1}{D}\sum\limits_{j}\big(x_{j}-\bar{x}\big)^{2}+\varepsilon}}+\delta_{i},\ \ \ \ \ \text{where}\ \ \ \ \ \bar{x}=\frac{1}{D}\sum\limits_{k=1}^{D}x_{k} (1) The scalar γ \gamma and the vector elements δ i \delta_{i} are learnable parameters, ε \varepsilon is a small regularization constant. Importantly, this operation can be viewed as an activation function for the neurons and can be defined as a partial derivative of the Lagrangian function (see [ 19 , 20 , 21 ] for the discussion of this property) L = D ​ γ ​ 1 D ​ ∑ j ( x j − x ¯ ) 2 + ε + ∑ j δ j ​ x j , so that g i = ∂ L ∂ x i L=D\gamma\sqrt{\frac{1}{D}\sum\limits_{j}\big(x_{j}-\bar{x}\big)^{2}+\varepsilon}\ +\ \sum\limits_{j}\delta_{j}x_{j},\ \ \ \ \ \text{so that}\ \ \ \ \ g_{i}=\frac{\partial L}{\partial x_{i}} (2)

### Multi-Head Energy Attention

The first contribution to 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 not have a separate value matrix; instead the value matrix is a function of keys and queries). The goal of the energy based attention is to evolve the tokens in such a 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 α = 1 ​ … ​ Y \alpha=1...Y to denote elements of this internal space, and index h = 1 ​ … ​ H h=1...H to denote different heads of this operation. With these notations the energy-based attention operation is described by the following energy function: E = ATT − 1 β ∑ h = 1 H ∑ C = 1 N log ( ∑ B ≠ C exp ( β A h ​ B ​ C ) ) E{{}^{\text{ATT}}}=-\frac{1}{\beta}\sum\limits_{h=1}^{H}\sum\limits_{C=1}^{N}\textrm{log}\left(\sum\limits_{B\neq C}\textrm{exp}\left(\beta A_{hBC}\right)\right) (3) where the attention matrix A h ​ B ​ C A_{hBC} is computed from query and key tensors as follows: A h ​ B ​ C = ∑ α K α ​ h ​ B ​ Q α ​ h ​ C , 𝐀 ∈ ℝ H × N × N K α ​ h ​ B = ∑ j W α ​ h ​ j K ​ g j ​ B , 𝐊 ∈ ℝ Y × H × N Q α ​ h ​ C = ∑ j W α ​ h ​ j Q ​ g j ​ C , 𝐐 ∈ ℝ Y × H × N \begin{split}A_{hBC}&=\sum\limits_{\alpha}K_{\alpha hB}\;Q_{\alpha hC},\ \ \ \ \ \ \mathbf{A}\in\mathbb{R}^{H\times N\times N}\\ K_{\alpha hB}&=\sum\limits_{j}W^{K}_{\alpha hj}\;g_{jB},\ \ \ \ \ \ \ \ \ \ \ \mathbf{K}\in\mathbb{R}^{Y\times H\times N}\\ Q_{\alpha hC}&=\sum\limits_{j}W^{Q}_{\alpha hj}\;g_{jC},\ \ \ \ \ \ \ \ \ \ \ \mathbf{Q}\in\mathbb{R}^{Y\times H\times N}\end{split} (4) and the tensors 𝐖 K ∈ ℝ Y × H × D \mathbf{W}^{K}\in\mathbb{R}^{Y\times H\times D} and 𝐖 Q ∈ ℝ Y × H × D \mathbf{W}^{Q}\in\mathbb{R}^{Y\times H\times D} 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 are aligned with the keys of a small number of other patches connected by the attention map. Different heads (index h h ) contribute to the energy additively.

### Hopfield Network Module