SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Benjamin Hoover*
benjamin.hoover@ibm.com
Yuchen Liang*
Department of CS
liangy7@rpi.edu
Bao Pham*
phamb@rpi.edu
rpanda@ibm.com
Hendrik Strobelt
hendrik.strobelt@ibm.com
Duen Horng Chau
College of Computing
polo@gatech.edu
Mohammed J. Zaki
zaki@cs.rpi.edu
Dmitry Krotov
krotov@ibm.com
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.
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
$^{*}$  B.Hoover, Y.Liang, and B.Pham equally contributed to this work.
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 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
Figure 1: Overview of the Energy Transformer (ET). Instead of a sequence of conventional transformer blocks, a single recurrent ET block is used. The operation of this block is dictated by the global energy function. The token representations are updated according to a continuous time differential equation with the time-discretized update step  $\alpha = dt / \tau$ . 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 to 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 is used. Each token represents a 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 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_{iA}$ . In the following the indices  $i,j,k = 1\dots D$  are used to denote the token vector's elements, indices  $A,B,C = 1\dots 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$ -dimensional vector  $\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
Figure 2: Left: Inside the ET block. The input token  $\mathbf{x}^{(t)}$  passes through a sequence of operations and gets updated to produce the output token  $\mathbf{x}^{(t + 1)}$ . The operations inside the ET block are carefully engineered so that the entire network has a 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 is 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 a continuous time differential equation, which describes interactions between these particles. Initially, at  $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  $\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} = \gamma \frac {x _ {i} - \bar {x}}{\sqrt {\frac {1}{D} \sum_ {j} \left(x _ {j} - \bar {x}\right) ^ {2} + \varepsilon}} + \delta_ {i}, \quad \text {w h e r e} \quad \bar {x} = \frac {1}{D} \sum_ {k = 1} ^ {D} x _ {k} \tag {1}
$$
The scalar  $\gamma$  and the vector elements  $\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 \gamma \sqrt {\frac {1}{D} \sum_ {j} \left(x _ {j} - \bar {x}\right) ^ {2} + \varepsilon} + \sum_ {j} \delta_ {j} x _ {j}, \quad \text {s o t h a t} \quad g _ {i} = \frac {\partial L}{\partial x _ {i}} \tag {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  $\alpha = 1\dots Y$  to denote elements of this internal space, and index  $h = 1\dots H$  to denote different heads of this operation. With these notations the energy-based attention operation is described by the following energy function:
$$
E ^ {\mathrm {A T T}} = - \frac {1}{\beta} \sum_ {h = 1} ^ {H} \sum_ {C = 1} ^ {N} \log \left(\sum_ {B \neq C} \exp \left(\beta A _ {h B C}\right)\right) \tag {3}
$$
where the attention matrix  $A_{hBC}$  is computed from query and key tensors as follows:
$$
A _ {h B C} = \sum_ {\alpha} K _ {\alpha h B} Q _ {\alpha h C}, \quad \mathbf {A} \in \mathbb {R} ^ {H \times N \times N}
$$
$$
K _ {\alpha h B} = \sum_ {j} W _ {\alpha h j} ^ {K} g _ {j B}, \quad \mathbf {K} \in \mathbb {R} ^ {Y \times H \times N} \tag {4}
$$
$$
Q _ {\alpha h C} = \sum_ {j} W _ {\alpha h j} ^ {Q} g _ {j C}, \quad \mathbf {Q} \in \mathbb {R} ^ {Y \times H \times N}
$$
and the tensors  $\mathbf{W}^K\in \mathbb{R}^{Y\times H\times D}$  and  $\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$ ) 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 in realistic images. The energy of this sub-block is defined as:
$$
E ^ {\mathrm {H N}} = - \sum_ {B = 1} ^ {N} \sum_ {\mu = 1} ^ {K} G \left(\sum_ {j = 1} ^ {D} \xi_ {\mu j} g _ {j B}\right), \quad \xi \in \mathbb {R} ^ {K \times D} \tag {5}
$$
where  $\xi_{\mu j}$  is a set of learnable weights (memories in the Hopfield Network), and  $G(\cdot)$  is an integral of the activation function  $r(\cdot)$ , so that  $G(\cdot)' = 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(\cdot) = \mathrm{ReLU}$ ), or as a modern continuous Hopfield Network [10, 12, 19] if the activation function is sharply peaked around the memories (e.g.,  $r(\cdot) = \mathrm{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  $\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
$$
\tau \frac {d x _ {i A}}{d t} = - \frac {\partial E}{\partial g _ {i A}}, \quad \text {w h e r e} \quad E = E ^ {\mathrm {A T T}} + E ^ {\mathrm {H N}} \tag {6}
$$
Here  $x_{iA}$  is the token representation (input and output from the ET block), and  $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  $\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
$$
\frac {d E}{d t} = \sum_ {i, j, A} \frac {\partial E}{\partial g _ {i A}} \frac {\partial g _ {i A}}{\partial x _ {j A}} \frac {d x _ {j A}}{d t} = - \frac {1}{\tau} \sum_ {i, j, A} \frac {\partial E}{\partial g _ {i A}} M _ {i j} ^ {A} \frac {\partial E}{\partial g _ {j A}} \leq 0 \tag {7}
$$
The last inequality sign holds if the symmetric part of the matrix
$$
M _ {i j} ^ {A} = \frac {\partial g _ {i A}}{\partial x _ {j A}} = \frac {\partial^ {2} L}{\partial x _ {i A} \partial x _ {j A}} \tag {8}
$$
is positive semi-definite (for each value of index  $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 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):
$$
- \frac {\partial E ^ {\mathrm {A T T}}}{\partial g _ {i A}} = \sum_ {C \neq A} \sum_ {\alpha} W _ {\alpha i} ^ {Q} K _ {\alpha C} \operatorname {s o f t m a x} _ {C} \left(\beta \sum_ {\gamma} K _ {\gamma C} Q _ {\gamma A}\right) + W _ {\alpha i} ^ {K} Q _ {\alpha C} \operatorname {s o f t m a x} _ {A} \left(\beta \sum_ {\gamma} K _ {\gamma A} Q _ {\gamma C}\right)
$$
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 [5] with the value matrix equal to  $\mathbf{V} = (\mathbf{W}^Q)^T\mathbf{K} = \sum_\alpha W_{\alpha i}^Q 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
$$
- \frac {\partial E ^ {\mathrm {H N}}}{\partial g _ {i A}} = \sum_ {\mu = 1} ^ {K} \xi_ {\mu i} G ^ {\prime} \left(\sum_ {j = 1} ^ {D} \xi_ {\mu j} g _ {j A}\right) \tag {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 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
*The code is available: https://github.com/bhoov/energy-transformer-jax.
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$  steps. The token representations after  $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.
Figure 3: Reconstruction examples of our Energy Transformer using images from the ImageNet-1k validation set. Top row: input images where  $50\%$  of the patches are masked with the learned MASK token. Middle row: output reconstructions after 12 time steps. Bottom row: original images.