Original: E
Revised: E

Original: NERGY T RANSFORMER
Revised: NERGY T RANSFORMER

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Transformers have become the de facto models of choice in machine learning, typically leading to impressive performance on many applications.
Revised: Transformers have become the de facto models of choice in machine learning, typically leading to impressive performance on many applications.

Original: At the same time, the architectural development in the transformer world is mostly driven by empirical findings, and the theoretical understanding of their architectural building blocks is rather limited.
Revised: At the same time, the architectural development in the transformer world is mostly driven by empirical findings, and the theoretical understanding of their architectural building blocks is rather limited.

Original: In contrast, Dense Associative Memory models or Modern Hopfield Networks have a well-established theoretical foundation, but have not yet demonstrated truly impressive practical results.
Revised: In contrast, Dense Associative Memory models or Modern Hopfield Networks have a well-established theoretical foundation, but have not yet demonstrated truly impressive practical results.

Original: We propose a transformer architecture that replaces the sequence of feedforward transformer blocks with a single large Associative Memory model.
Revised: We propose a transformer architecture that replaces the sequence of feedforward transformer blocks with a single large Associative Memory model.

Original: Our novel architecture, called Energy
Revised: Our novel architecture, called Energy

Original: Transformer (or ET for short), has many of the familiar architectural primitives that are often used in the current generation of transformers.
Revised: Transformer (or ET for short), has many of the familiar architectural primitives that are often used in the current generation of transformers.

Original: However, it is not identical to the existing architectures.
Revised: However, it is not identical to the existing architectures.

Original: The sequence of transformer layers in ET is purposely designed to minimize a specifically engineered energy function, which is responsible for representing the relationships between the tokens.
Revised: The sequence of transformer layers in ET is purposely designed to minimize a specifically engineered energy function, which is responsible for representing the relationships between the tokens.

Original: As a consequence of this computational principle, the attention in ET is different from the conventional attention mechanism.
Revised: As a consequence of this computational principle, the attention in ET is different from the conventional attention mechanism.

Original: In this work, we introduce the theoretical foundations of ET, explore it’s empirical capabilities using the image completion task, and obtain strong quantitative results on the graph anomaly detection task.
Revised: In this work, we introduce the theoretical foundations of ET, explore it’s empirical capabilities using the image completion task, and obtain strong quantitative results on the graph anomaly detection task.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformers have become pervasive models in various domains of machine learning, including language, vision, and audio processing.
Revised: Transformers have become pervasive models in various domains of machine learning, including language, vision, and audio processing.

Original: Every transformer block uses four fundamental operations: attention, feed-forward multi-layer perceptron (MLP), residual connection, and layer normalization.
Revised: Every transformer block uses four fundamental operations: attention, feed-forward multi-layer perceptron (MLP), residual connection, and layer normalization.

Original: Different variations of transformers result from combining these four operations in various ways.
Revised: Different variations of transformers result from combining these four operations in various ways.

Original: For instance, Press et al.
Revised: For instance, Press et al.

Original: (2019) propose to frontload additional attention operations and backload additional MLP layers in a sandwich-like instead of interleaved way, Lu et al.
Revised: (2019) propose to frontload additional attention operations and backload additional MLP layers in a sandwich-like instead of interleaved way, Lu et al.

Original: (2019) prepend an
Revised: (2019) prepend an

Original: MLP layer before the attention in each transformer block, So et al.
Revised: MLP layer before the attention in each transformer block, So et al.

Original: (2019) use neural architecture search methods to evolve even more sophisticated transformer blocks, and so on.
Revised: (2019) use neural architecture search methods to evolve even more sophisticated transformer blocks, and so on.

Original: Various methods exist to approximate the attention operation, multiple modifications of the norm operation, and connectivity of the block; see, for example, (Lin et al., 2021) for a taxonomy of different models.
Revised: Various methods exist to approximate the attention operation, multiple modifications of the norm operation, and connectivity of the block; see, for example, (Lin et al., 2021) for a taxonomy of different models.

Original: 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.
Revised: 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.

Original: Additionally, the computational role of the four elements remains the subject of discussions.
Revised: Additionally, the computational role of the four elements remains the subject of discussions.

Original: Originally, Vaswani et al.
Revised: Originally, Vaswani et al.

Original: (2017) 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.
Revised: (2017) 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.

Original: On the other hand more recent investigations (Yu et al., 2021) argue that the entire transformer block is important.
Revised: On the other hand more recent investigations (Yu et al., 2021) argue that the entire transformer block is important.

Original: 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.
Revised: 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.

Original: In a seemingly unrelated line of work, Associative Memory models, also known as Hopfield Networks (Hopfield, 1982; 1984), have been gaining popularity in the machine learning community thanks to theoretical advancements pertaining to their memory storage capacity and novel architectural modifications.
Revised: In a seemingly unrelated line of work, Associative Memory models, also known as Hopfield Networks (Hopfield, 1982; 1984), have been gaining popularity in the machine learning community thanks to theoretical advancements pertaining to their memory storage capacity and novel architectural modifications.

Original: Specifically, it has been shown that increasing the sharpness of the activation functions can lead to super-linear (Krotov & Hopfield, 2016) and even exponential (Demircigil et al., 2017) memory storage capacity for these models, which is important for machine learning applications.
Revised: Specifically, it has been shown that increasing the sharpness of the activation functions can lead to super-linear (Krotov & Hopfield, 2016) and even exponential (Demircigil et al., 2017) memory storage capacity for these models, which is important for machine learning applications.

Original: This new class of Hopfield Networks is called Dense Associative Memories or Modern
Revised: This new class of Hopfield Networks is called Dense Associative Memories or Modern

Original: [Figure]
Revised: [Figure]

Original: ET for graphs
Revised: ET for graphs

Original: [Figure]
Revised: [Figure]

Original: Hopfield Networks.
Revised: Hopfield Networks.

Original: Ramsauer et al.
Revised: Ramsauer et al.

Original: (2020) additionally describe how the attention mechanism in transformers is closely related to a special model of this family with the softmax activation function.
Revised: (2020) additionally describe how the attention mechanism in transformers is closely related to a special model of this family with the softmax activation function.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Dense
Revised: Dense

Original: Associative Memory models are designed for completing the incomplete patterns.
Revised: Associative Memory models are designed for completing the incomplete patterns.

Original: For instance, a pattern can be the concatenation of an image and its label, and the model can be trained to predict part of the input (the label), which is masked, given the query (the image).
Revised: For instance, a pattern can be the concatenation of an image and its label, and the model can be trained to predict part of the input (the label), which is masked, given the query (the image).

Original: They can also be trained in a self-supervised way by predicting the occluded parts of the image, or denoising the image.
Revised: They can also be trained in a self-supervised way by predicting the occluded parts of the image, or denoising the image.

Original: There are also high-level differences between the two classes of models.
Revised: There are also high-level differences between the two classes of models.

Original: 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.
Revised: 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.

Original: Transformers are typically not described as dynamical systems at all.
Revised: Transformers are typically not described as dynamical systems at all.

Original: Rather, they are thought of as feed-forward networks built of the four computational elements discussed above.
Revised: Rather, they are thought of as feed-forward networks built of the four computational elements discussed above.

Original: Even if one thinks about them as dynamical systems with tied weights, e.g., (Bai et al., 2019), there is no reason to expect that their dynamics converge to a fixed point attractor (see the discussion in (Lan et al., 2020)).
Revised: Even if one thinks about them as dynamical systems with tied weights, e.g., (Bai et al., 2019), there is no reason to expect that their dynamics converge to a fixed point attractor (see the discussion in (Lan et al., 2020)).

Original: Additionally, a recent study (Yang et al., 2022) uses a form of Majorization-Minimization algorithms (Sun et al., 2016) to interpret the forward path in the transformer block as an optimization process.
Revised: Additionally, a recent study (Yang et al., 2022) uses a form of Majorization-Minimization algorithms (Sun et al., 2016) to interpret the forward path in the transformer block as an optimization process.

Original: This interpretation requires imposing certain constraints on the operations inside the block, and attempting to find an energy function that describes the constrained block.
Revised: This interpretation requires imposing certain constraints on the operations inside the block, and attempting to find an energy function that describes the constrained block.

Original: We take a complementary approach by using intuition developed in Associative Memory models to start with an energy function that is perfectly suited for the problem of interest.
Revised: We take a complementary approach by using intuition developed in Associative Memory models to start with an energy function that is perfectly suited for the problem of interest.

Original: The optimization process and the resulting transformer block in our approach is a consequence of this specifically chosen energy function.
Revised: The optimization process and the resulting transformer block in our approach is a consequence of this specifically chosen energy function.

Original: Concretely, we use the recent theoretical advancements and architectural developments in Dense
Revised: Concretely, we use the recent theoretical advancements and architectural developments in Dense

Original: Associative Memories to design an energy function tailored to route the information between the tokens.
Revised: Associative Memories to design an energy function tailored to route the information between the tokens.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In the image domain, this fixed point corresponds to the completed image with masked tokens replaced by plausible autocompletions of the occluded image patches.
Revised: In the image domain, this fixed point corresponds to the completed image with masked tokens replaced by plausible autocompletions of the occluded image patches.

Original: In the graph domain, the fixed point reveals the anomaly status of a given node given that node’s neighbors; see Figure 1.
Revised: In the graph domain, the fixed point reveals the anomaly status of a given node given that node’s neighbors; see Figure 1.

Original: The energy function in our ET block is designed with the goal to describe the relationships between the tokens .
Revised: The energy function in our ET block is designed with the goal to describe the relationships between the tokens .

Original: 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 impaint the missing eye, etc.
Revised: 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 impaint the missing eye, etc.

Original: In the graph domain, these are the relationships between the attributes and the anomaly status of the connected nodes.
Revised: In the graph domain, these are the relationships between the attributes and the anomaly status of the connected nodes.

Original: The optimization procedure during the forward path of ET uses continuous time differential equations, and describes a gradient decent on the specifically chosen energy function.
Revised: The optimization procedure during the forward path of ET uses continuous time differential equations, and describes a gradient decent on the specifically chosen energy function.

Original: The core mathematical principle of the ET block – the existence of the global energy function – dictates strong constraints on the possible operations, the order in which these operations are executed in the forward path, and the symmetries of the weights in the network.
Revised: The core mathematical principle of the ET block – the existence of the global energy function – dictates strong constraints on the possible operations, the order in which these operations are executed in the forward path, and the symmetries of the weights in the network.

Original: As a corollary of this theoretical principle, the attention mechanism of ET is different from the attention mechanism commonly used in feed-forward transformers (Vaswani et al., 2017).
Revised: As a corollary of this theoretical principle, the attention mechanism of ET is different from the attention mechanism commonly used in feed-forward transformers (Vaswani et al., 2017).

Original: In the following section we introduce the global energy function for the ET block and explain the block’s architecture.
Revised: In the following section we introduce the global energy function for the ET block and explain the block’s architecture.

Original: We then explore the inner workings of the ET network for image completion and qualitatively assess the learned representations (the model is trained on ImageNet-1k in a general pipeline similar to Dosovitskiy et al.
Revised: We then explore the inner workings of the ET network for image completion and qualitatively assess the learned representations (the model is trained on ImageNet-1k in a general pipeline similar to Dosovitskiy et al.

Original: Finally, we turn to the graph anomaly detection task, which is conceptually similar to the image completion setting but has a record of strong published benchmarks against which our approach can be quantitatively compared.
Revised: Finally, we turn to the graph anomaly detection task, which is conceptually similar to the image completion setting but has a record of strong published benchmarks against which our approach can be quantitatively compared.

Original: We show that the ET network stands in line with or outperforms the latest benchmarks.
Revised: We show that the ET network stands in line with or outperforms the latest benchmarks.

Original: Although we focus on the computer vision and anomaly detection domains in this paper, we believe that the computational principles developed can be applied to other exciting domains (e.g., NLP, audio, and video) in which conventional transformers have shown promising results.
Revised: Although we focus on the computer vision and anomaly detection domains in this paper, we believe that the computational principles developed can be applied to other exciting domains (e.g., NLP, audio, and video) in which conventional transformers have shown promising results.

Original:  E NERGY T RANSFORMER B LOCK
Revised:  E NERGY T RANSFORMER B LOCK

Original: We now introduce the theoretical framework of the ET network.
Revised: We now introduce the theoretical framework of the ET network.

Original: For clarity of presentation, we use language associated with the image domain.
Revised: For clarity of presentation, we use language associated with the image domain.

Original: For the graph domain, one should think about “image patches” as nodes on the graph.
Revised: For the graph domain, one should think about “image patches” as nodes on the graph.

Original: The overall pipeline is similar to the Vision Transformer networks (ViTs) and is shown in Figure 1.
Revised: The overall pipeline is similar to the Vision Transformer networks (ViTs) and is shown in Figure 1.

Original: An input image is split into non-overlapping patches.
Revised: An input image is split into non-overlapping patches.

Original: 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 .
Revised: 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 .

Original: In the following the indices i, j, k = 1 ...D are used to denote the token vector’s elements, indices A, B, C = 1 ...
Revised: In the following the indices i, j, k = 1 ...D are used to denote the token vector’s elements, indices A, B, C = 1 ...

Original: N are used to enumerate the patches and their corresponding tokens.
Revised: N are used to enumerate the patches and their corresponding tokens.

Original: 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 x A .
Revised: 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 x A .

Original: 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).
Revised: 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).

Original: The ET block is described by a continuous time differential equation, which describes interactions between these particles.
Revised: The ET block is described by a continuous time differential equation, which describes interactions between these particles.

Original: Initially, at t = [-
-] the network is given a set containing two groups of particles corresponding to open and masked patches.
Revised: Initially, at t = [-
-]1 the network is given a set containing two groups of particles corresponding to open and masked patches.

Original: The “open” particles know their identity and location in the image.
Revised: The “open” particles know their identity and location in the image.

Original: The “masked” particles only know where in the image they are located, but are not provided the information about what image patch they represent.
Revised: The “masked” particles only know where in the image they are located, but are not provided the information about what image patch they represent.

Original: 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.
Revised: 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.

Original: This dynamical evolution is designed so that it minimizes a global energy function, and is guaranteed to arrive at a fixed point attractor state.
Revised: This dynamical evolution is designed so that it minimizes a global energy function, and is guaranteed to arrive at a fixed point attractor state.

Original: The identities of the masked particles are considered to be revealed when the dynamical trajectory reaches the fixed point.
Revised: The identities of the masked particles are considered to be revealed when the dynamical trajectory reaches the fixed point.

Original: Thus, the central question is: how can we design the energy function that accurately captures the task that the Energy Transformer needs to solve?
Revised: Thus, the central question is: how can we design the energy function that accurately captures the task that the Energy Transformer needs to solve?

Original: 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.
Revised: 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.

Original: These two pieces of information are described by two contributions to the
Revised: These two pieces of information are described by two contributions to the

Original: ET’s energy function: the energy based attention and the Hopfield Network, respectively, for reasons that will become clear in the next sections.
Revised: ET’s energy function: the energy based attention and the Hopfield Network, respectively, for reasons that will become clear in the next sections.

Original: Below we define each element of the ET block in the order they appear in Figure 2.
Revised: Below we define each element of the ET block in the order they appear in Figure 2.

Original: L AYER N ORM
Revised: L AYER N ORM

Original: Each token is represented by a vector x ∈ R D .
Revised: Each token is represented by a vector x ∈ R D .

Original: At the same time, most of the operations inside the
Revised: At the same time, most of the operations inside the

Original: ET block are defined using a layer-normalized token representation
Revised: ET block are defined using a layer-normalized token representation

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: The scalar γ and the vector δ i are learnable parameters, ε is a small regularization constant.
Revised: The scalar γ and the vector δ i are learnable parameters, ε is a small regularization constant.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: so that
Revised: so that

Original: [Equation]
Revised: [Equation]

Original: See Krotov & Hopfield (2021); Tang & Kopp (2021); Krotov (2021) for a detailed discussion of this property.
Revised: See Krotov & Hopfield (2021); Tang & Kopp (2021); Krotov (2021) for a detailed discussion of this property.

Original: M ULTI -H EAD E NERGY A TTENTION
Revised: M ULTI -H EAD E NERGY A TTENTION

Original: The first contribution to the ET’s energy function is responsible for exchanging information between the particles (patches).
Revised: The first contribution to the ET’s energy function is responsible for exchanging information between the particles (patches).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: Below we use index α = 1 ...Y to denote elements of this internal space, and index h
Revised: Below we use index α = 1 ...Y to denote elements of this internal space, and index h

Original: = 1 ...H to denote different heads of this operation.
Revised: = 1 ...H to denote different heads of this operation.

Original: With these notations the energy-based attention operation is described by the following energy function:
Revised: With these notations the energy-based attention operation is described by the following energy function:

Original: [Equation]
Revised: [Equation]

Original: where the queries and keys tensors are defined as
Revised: where the queries and keys tensors are defined as

Original: [Equation]
Revised: [Equation]

Original: and the tensors W K ∈ R Y × H × D and W Q ∈ R Y × H × D are learnable parameters.
Revised: and the tensors W K ∈ R Y × H × D and W Q ∈ R Y × H × D are learnable parameters.

Original: 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?).
Revised: 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?).

Original: 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.
Revised: 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.

Original: Different heads (index h ) contribute to the energy additively.
Revised: Different heads (index h ) contribute to the energy additively.

Original: H
Revised: H

Original: OPFIELD N ETWORK M ODULE
Revised: OPFIELD N ETWORK M ODULE

Original: 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.
Revised: 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.

Original: The energy of this sub-block is defined as:
Revised: The energy of this sub-block is defined as:

Original: [Equation]
Revised: [Equation]

Original: where ξ µj is a set of learnable weights (memories in the Hopfield Network), and r ( · ) is an activation function.
Revised: where ξ µj is a set of learnable weights (memories in the Hopfield Network), and r ( · ) is an activation function.

Original: Depending on the choice of the activation function this step can be viewed either as a classical continuous Hopfield Network (Hopfield, 1984) if the activation function grows slowly (e.g.,
Revised: Depending on the choice of the activation function this step can be viewed either as a classical continuous Hopfield Network (Hopfield, 1984) if the activation function grows slowly (e.g.,

Original: ReLU), or as a modern continuous Hopfield Network (Krotov & Hopfield, 2016; Ramsauer et al., 2020; Krotov & Hopfield, 2021)
Revised: ReLU), or as a modern continuous Hopfield Network (Krotov & Hopfield, 2016; Ramsauer et al., 2020; Krotov & Hopfield, 2021)

Original: if the activation function is sharply peaked around the memories (e.g. power or softmax).
Revised: if the activation function is sharply peaked around the memories (e.g. power or softmax).

Original: 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 neuron’s space to be the same (transposed matrix) as the weights of the subsequent projection from the hidden space to the token space.
Revised: 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 neuron’s space to be the same (transposed matrix) as the weights of the subsequent projection from the hidden space to the token space.

Original: Thus, the HN module here is an MLP with shared weights that is applied recurrently .
Revised: Thus, the HN module here is an MLP with shared weights that is applied recurrently .

Original: The energy contribution of this block is low when the tokens representations are aligned with some rows of the matrix ξ , which represent memories.
Revised: The energy contribution of this block is low when the tokens representations are aligned with some rows of the matrix ξ , which represent memories.

Original: D YNAMICS OF T OKEN U PDATES
Revised: D YNAMICS OF T OKEN U PDATES

Original: The forward path of the ET network is described by the continuous time differential equation, which minimizes the sum of the two energies described above
Revised: The forward path of the ET network is described by the continuous time differential equation, which minimizes the sum of the two energies described above

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Here x iA is the token representation (input and output from the ET block), and g iA is its layernormalized version.
Revised: Here x iA is the token representation (input and output from the ET block), and g iA is its layernormalized version.

Original: The first energy is low when each patch’s queries are aligned with the keys of its neighbors.
Revised: The first energy is low when each patch’s queries are aligned with the keys of its neighbors.

Original: 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 χ ).
Revised: 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 χ ).

Original: The dynamical system (6) finds a trade-off between these two desirable properties of each token’s representation.
Revised: The dynamical system (6) finds a trade-off between these two desirable properties of each token’s representation.

Original:  To demonstrate that the dynamical system (6) minimizes the energy, consider the temporal derivative dE
Revised: For numerical evaluations equation (6) is discretized in time. To demonstrate that the dynamical system (6) minimizes the energy, consider the temporal derivative dE

Original: [Equation]
Revised: [Equation]

Original: The last inequality sign holds if the symmetric part of the matrix
Revised: The last inequality sign holds if the symmetric part of the matrix

Original: [Equation]
Revised: [Equation]

Original: is positive semi-definite (for each value of index A ).
Revised: is positive semi-definite (for each value of index A ).

Original: The Lagrangian (2) satisfies this condition.
Revised: The Lagrangian (2) satisfies this condition.

Original: R ELATIONSHIP TO M ODERN H OPFIELD N ETWORKS AND C ONVENTIONAL A TTENTION
Revised: R ELATIONSHIP TO M ODERN H OPFIELD N ETWORKS AND C ONVENTIONAL A TTENTION

Original: One of the theoretical contributions of our work is the design of the energy attention mechanism and the corresponding energy function (3).
Revised: One of the theoretical contributions of our work is the design of the energy attention mechanism and the corresponding energy function (3).

Original: Although heavily inspired by prior work on Modern
Revised: Although heavily inspired by prior work on Modern

Original: Hopfield Networks, our approach is fundamentally different from it.
Revised: Hopfield Networks, our approach is fundamentally different from it.

Original: Our energy function (3) may look somewhat similar to the energy function of a continuous Hopfield Network with the softmax activation function.
Revised: Our energy function (3) may look somewhat similar to the energy function of a continuous Hopfield Network with the softmax activation function.

Original: 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).
Revised: 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).

Original: In contrast, in our energy attention network the keys are dynamical variables that evolve in time with the queries.
Revised: In contrast, in our energy attention network the keys are dynamical variables that evolve in time with the queries.

Original: To emphasize this further, it is instructive to write explicitly the ET attention contribution to the update dynamics (6).
Revised: To emphasize this further, it is instructive to write explicitly the ET attention contribution to the update dynamics (6).

Original: It is given by (for clarity, assume only one head of attention):
Revised: It is given by (for clarity, assume only one head of attention):

Original: [Equation]
Revised: [Equation]

Original: In both terms the softmax normalization is done over the token index of the keys, which is indicated by the subscript in the equation.
Revised: In both terms the softmax normalization is done over the token index of the keys, which is indicated by the subscript in the equation.

Original: The first term in this formula is the conventional attention mechanism (Vaswani et al., 2017) with the value matrix equal to V = ( W Q ) T K = (cid:80) α W Qαi K αC .
Revised: The first term in this formula is the conventional attention mechanism (Vaswani et al., 2017) with the value matrix equal to V = ( W Q ) T K = (cid:80) α W Qαi K αC .

Original: The second term is the brand new contribution that is missing in the original attention mechanism.
Revised: The second term is the brand new contribution that is missing in the original attention mechanism.

Original: The presence of this second term is crucial to make sure that the dynamical system (6) minimizes the energy function if applied recurrently.
Revised: The presence of this second term is crucial to make sure that the dynamical system (6) minimizes the energy function if applied recurrently.

Original: This second term is the main difference of our approach compared to the Modern Hopfield Networks.
Revised: This second term is the main difference of our approach compared to the Modern Hopfield Networks.

Original: The same difference applies to the other recent proposals (Yang et al., 2022).
Revised: The same difference applies to the other recent proposals (Yang et al., 2022).

Original: Lastly, we want to emphasize that our ET block contains two different kinds of Hopfield Networks acting in parallel, see Figure 2.
Revised: Lastly, we want to emphasize that our ET block contains two different kinds of Hopfield Networks acting in parallel, see Figure 2.

Original: The first one is the energy attention module, which is inspired by, but not identical to, Modern Hopfield Networks.
Revised: The first one is the energy attention module, which is inspired by, but not identical to, Modern Hopfield Networks.

Original: The second one is the “Hopfield Network” module, which can be either a classical or modern Hopfield Network.
Revised: The second one is the “Hopfield Network” module, which can be either a classical or modern Hopfield Network.

Original: These two should not be confused.
Revised: These two should not be confused.

Original:  Q UALITATIVE I NSPECTION OF THE ET FRAMEWORK ON I MAGE N ET
Revised:  Q UALITATIVE I NSPECTION OF THE ET FRAMEWORK ON I MAGE N ET

Original: We have trained the ET network on the masked image completion task using ImageNet-1k dataset (Deng et al., 2009).
Revised: We have trained the ET network on the masked image completion task using ImageNet-1k dataset (Deng et al., 2009).

Original: Each image was broken into non-overlapping patches of 16x16 RGB pixels, which were projected with a single affine encoder into the token space.
Revised: Each image was broken into non-overlapping patches of 16x16 RGB pixels, which were projected with a single affine encoder into the token space.

Original: Half of these tokens were “masked”, e.g., by replacing them with a learnable MASK token.
Revised: Half of these tokens were “masked”, e.g., by replacing them with a learnable MASK token.

Original: A distinct learnable position encoding vector was added to each token.
Revised: A distinct learnable position encoding vector was added to each token.

Original: Our ET block then processes all tokens recurrently for
Revised: Our ET block then processes all tokens recurrently for

Original: T steps.
Revised: T steps.

Original: The token representations after T steps are passed to a simple linear decoder (consisting of a layer norm and an affine transformation).
Revised: The token representations after T steps are passed to a simple linear decoder (consisting of a layer norm and an affine transformation).

Original: The loss function is the standard MSE loss on the occluded patches.
Revised: The loss function is the standard MSE loss on the occluded patches.

Original: See more details on the implementation and the hyperparameters in Appendix A.
Revised: See more details on the implementation and the hyperparameters in Appendix A.

Original: Examples of occluded/reconstructed images (unseen during training) are shown in Figure 3.
Revised: Examples of occluded/reconstructed images (unseen during training) are shown in Figure 3.

Original: In general, our model learns to perform the task very well, capturing the texture in dog fur (col 3) and understanding meaningful boundaries of objects.
Revised: In general, our model learns to perform the task very well, capturing the texture in dog fur (col 3) and understanding meaningful boundaries of objects.

Original: 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 (col 4) and completing irregular brick patterns in the name of extending the un-occluded borders (last col).
Revised: 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 (col 4) and completing irregular brick patterns in the name of extending the un-occluded borders (last col).

Original: We additionally inspect the positional encoding vectors associated with every token,
Revised: We additionally inspect the positional encoding vectors associated with every token,

Original: Figure 2, where the model learns a locality structure in the image plane that is very similar to the original ViT (Dosovitskiy et al., 2021).
Revised: Figure 2, where the model learns a locality structure in the image plane that is very similar to the original ViT (Dosovitskiy et al., 2021).

Original: The position embedding of each image patch has learned high similarity values to other patches in the same row and column, with similarity values higher for neighboring tokens than distant tokens.
Revised: The position embedding of each image patch has learned high similarity values to other patches in the same row and column, with similarity values higher for neighboring tokens than distant tokens.

Original: 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.
Revised: 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.

Original: This functionality makes it possible to visualize essentially any token representation , weight , or gradient of the energy directly in the image plane.
Revised: This functionality makes it possible to visualize essentially any token representation , weight , or gradient of the energy directly in the image plane.

Original: 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.
Revised: 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.

Original: In Figure 2 this functionality is used for inspecting the learned weights of the HN module directly in the image plane.
Revised: In Figure 2 this functionality is used for inspecting the learned weights of the HN module directly in the image plane.

Original: According to our theory, these weights should represent basis vectors in the space of all possible image patches.
Revised: According to our theory, these weights should represent basis vectors in the space of all possible image patches.

Original: These learned representations look qualitatively similar to the representations typically found in networks trained on image datasets, e.g., Zeiler & Fergus (2014).
Revised: These learned representations look qualitatively similar to the representations typically found in networks trained on image datasets, e.g., Zeiler & Fergus (2014).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Thus, the attention layer 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.
Revised: Thus, the attention layer 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.

Original: 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.
Revised: 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.

Original:  G RAPH A NOMALY D ETECTION
Revised:  G RAPH A NOMALY D ETECTION

Original: Having built the theoretical foundation of the ET network and gained an intuition about its inner workings through visualizations, we turn to quantitatively evaluating its performance on the graph anomaly detection problem, a task with plenty of strong and recently published baselines.
Revised: Having built the theoretical foundation of the ET network and gained an intuition about its inner workings through visualizations, we turn to quantitatively evaluating its performance on the graph anomaly detection problem, a task with plenty of strong and recently published baselines.

Original: Anomalies are outliers that significantly deviate in their properties from the majority of the samples.
Revised: Anomalies are outliers that significantly deviate in their properties from the majority of the samples.

Original: Detecting anomalies on graphs has broad applications in cybersecurity (Hong et al., 2014; Pan et al., 2019), fraud detection (Huang et al., 2018;
Revised: Detecting anomalies on graphs has broad applications in cybersecurity (Hong et al., 2014; Pan et al., 2019), fraud detection (Huang et al., 2018;

Original: Pourhabibi et al., 2020), and social networks (Chaker et al., 2017).
Revised: Pourhabibi et al., 2020), and social networks (Chaker et al., 2017).

Original: Generally, there are three types of graph anomalies: node anomaly, edge anomaly, and subgraph anomaly.
Revised: Generally, there are three types of graph anomalies: node anomaly, edge anomaly, and subgraph anomaly.

Original: In this work, we focus on node anomaly detection in attributed graphs.
Revised: In this work, we focus on node anomaly detection in attributed graphs.

Original: This task is perfectly suited for the ET network, since each node’s attributes can be encoded in the latent space and treated as a token (with added learnable positional embeddings).
Revised: This task is perfectly suited for the ET network, since each node’s attributes can be encoded in the latent space and treated as a token (with added learnable positional embeddings).

Original: The network iterates these representations in time, and the outputs can be used for the node anomaly classification task.
Revised: The network iterates these representations in time, and the outputs can be used for the node anomaly classification task.

Original: Graph Convolutional Networks (GCN) (Kipf & Welling, 2016) have been widely used for this task due to their capability of learning high level representations of graph structures and node attributes (Ding et al., 2019; Peng et al., 2020).
Revised: Graph Convolutional Networks (GCN) (Kipf & Welling, 2016) have been widely used for this task due to their capability of learning high level representations of graph structures and node attributes (Ding et al., 2019; Peng et al., 2020).

Original: However, vanilla GCNs suffer from the over-smoothing problem (Wu et al., 2019).
Revised: However, vanilla GCNs suffer from the over-smoothing problem (Wu et al., 2019).

Original: In each layer of the forward pass, the outlier node aggregates information from its neighbors.
Revised: In each layer of the forward pass, the outlier node aggregates information from its neighbors.

Original: This averaging makes the features of anomalies less distinguishable from the features of benign nodes.
Revised: This averaging makes the features of anomalies less distinguishable from the features of benign nodes.

Original: Our approach does not suffer from this problem, since the routing of the information between the nodes is done through the energy based attention, which uses different aggregation procedure depending on whether or not the node is anomalous.
Revised: Our approach does not suffer from this problem, since the routing of the information between the nodes is done through the energy based attention, which uses different aggregation procedure depending on whether or not the node is anomalous.

Original: In order to turn the anomaly detection task on graphs into the ET framework, consider an undirected graph with N nodes.
Revised: In order to turn the anomaly detection task on graphs into the ET framework, consider an undirected graph with N nodes.

Original: Every node has a vector of attributes y A ∈ R F , where F is the number of node’s features.
Revised: Every node has a vector of attributes y A ∈ R F , where F is the number of node’s features.

Original: Additionally, every node has a binary label l A , indicating whether the node is benign or not.
Revised: Additionally, every node has a binary label l A , indicating whether the node is benign or not.

Original: We focus on node anomaly and assume that all edges are trusted.
Revised: We focus on node anomaly and assume that all edges are trusted.

Original: The task is to predict the label of the node given the graph structure and the node’s features.
Revised: The task is to predict the label of the node given the graph structure and the node’s features.

Original: Since there are far more benign nodes in the graph than anomalous ones, anomaly detection can be regarded as an imbalanced node classification task.
Revised: Since there are far more benign nodes in the graph than anomalous ones, anomaly detection can be regarded as an imbalanced node classification task.

Original: First, the feature vectors for every node are converted to a token representation using a linear embedding E and adding a learnable positional embedding λ A
Revised: First, the feature vectors for every node are converted to a token representation using a linear embedding E and adding a learnable positional embedding λ A

Original: [Equation]
Revised: [Equation]

Original: where the superscript 0 indicates the time of the update of the ET dynamics.
Revised: where the superscript t = 1 indicates the time of the update of the ET dynamics.

Original: This token representation is iterated through the ET block for T iterations.
Revised: This token representation is iterated through the ET block for T iterations.

Original: When the retrieval dynamics becomes stable, we have the final representation for each node x t = T A (or more precisely g t = T A , since the outputs are additionally passed through a layer norm operation after the final ET update).
Revised: When the retrieval dynamics becomes stable, we have the final representation for each node x t = T A (or more precisely g t = T A , since the outputs are additionally passed through a layer norm operation after the final ET update).

Original: This output is concatenated with the initial (layer normalized) token to form the final output of the network
Revised: This output is concatenated with the initial (layer normalized) token to form the final output of the network

Original: [Equation]
Revised: [Equation]

Original: Following Tang et al.
Revised: Following Tang et al.

Original: (2022), the node representation g final A is fed into an MLP with the sigmoid activation function to compute the anomaly probabilities p A .
Revised: (2022), the node representation g final A is fed into an MLP with the sigmoid activation function to compute the anomaly probabilities p A .

Original: The weighted cross entropy
Revised: The weighted cross entropy

Original: [Equation]
Revised: [Equation]

Original:  is then minimized to train the whole network.
Revised: 11) is used to train the whole network.

Original: Above, σ is the ratio of anomalous labels ( l A = 1) to the regular labels ( l A = 0).
Revised: Above, σ is the ratio of the regular labels ( l A = 0) to anomalous labels ( l A = 1).

Original: 4. E XPERIMENTAL E VALUATION
Revised: 4. E XPERIMENTAL E VALUATION

Original: Four datasets are used for the graph anomaly detection experiments.
Revised: Four datasets are used for the graph anomaly detection experiments.

Original: YelpChi dataset (Rayana &
Revised: YelpChi dataset (Rayana &

Original: Akoglu, 2015) aims at opinion spam detection in Yelp reviews.
Revised: Akoglu, 2015) aims at opinion spam detection in Yelp reviews.

Original: Amazon dataset is used to detect anomalous users under the Musical Instrument Category on amazon.com (McAuley & Leskovec, 2013).
Revised: Amazon dataset is used to detect anomalous users under the Musical Instrument Category on amazon.com (McAuley & Leskovec, 2013).

Original: T-Finance and T-Social datasets (Tang et al., 2022) are used for anomalous account detection in the transactions and social networks, respectively.
Revised: T-Finance and T-Social datasets (Tang et al., 2022) are used for anomalous account detection in the transactions and social networks, respectively.

Original: 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.
Revised: 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.

Original: The task is to predict the label (anomaly status) of the nodes.
Revised: The task is to predict the label (anomaly status) of the nodes.

Original: For each dataset, either 1% or 40% of the nodes are used for training, and the remaining 99% or 60% are split
Revised: For each dataset, either 1% or 40% of the nodes are used for training, and the remaining 99% or 60% are split

Original: 2 into validation and testing, see Appendix B for details.
Revised: 2 into validation and testing, see Appendix B for details.

Original: We compare with state-of-the-art approaches for graph anomaly detection, which include GraphConsis (Liu et al., 2020), CAREGNN (Dou et al., 2020), PC-GNN (Liu et al., 2021) and BWGNN (Tang et al., 2022).
Revised: We compare with state-of-the-art approaches for graph anomaly detection, which include GraphConsis (Liu et al., 2020), CAREGNN (Dou et al., 2020), PC-GNN (Liu et al., 2021) and BWGNN (Tang et al., 2022).

Original: Additionally, multi-layer perceptrons (MLP) and Graph Transformer (GT) (Dwivedi & Bresson, 2020) are included in the baselines for completeness.
Revised: Additionally, multi-layer perceptrons (MLP) and Graph Transformer (GT) (Dwivedi & Bresson, 2020) are included in the baselines for completeness.

Original: 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 Davis & Goadrich (2006).
Revised: 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 Davis & Goadrich (2006).

Original: See Appendix B for more details on training protocols and the hyperparameters choices.
Revised: See Appendix B for more details on training protocols and the hyperparameters choices.

Original: The results are reported in
Revised: The results are reported in

Original: Table 1.
Revised: Table 1.

Original: Our ET network demonstrates very strong results across all the datasets.
Revised: Our ET network demonstrates very strong results across all the datasets.

Original: [Table]
Revised: [Table]

Original:  D ISCUSSION AND C ONCLUSIONS
Revised:  D ISCUSSION AND C ONCLUSIONS

Original: A lot of recent research has been dedicated to understanding the striking analogy between Hopfield
Revised: A lot of recent research has been dedicated to understanding the striking analogy between Hopfield

Original: Networks and the attention mechanism in transformers.
Revised: Networks and the attention mechanism in transformers.

Original: At a high level, the main message of our work is that the entire transformer block (including feed-forward MLP, layer normalization
Revised: At a high level, the main message of our work is that the entire transformer block (including feed-forward MLP, layer normalization

Original: [Table]
Revised: [Table]

Original: and residual connections) can be viewed as a single large Hopfield Network, not just attention alone.
Revised: and residual connections) can be viewed as a single large Hopfield Network, not just attention alone.

Original: 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.
Revised: 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.

Original: When used in the encoder-decoder setting, an appealing feature of our network is that any state, weight
Revised: When used in the encoder-decoder setting, an appealing feature of our network is that any state, weight

Original: [Table]
Revised: [Table]

Original: or state update can be mapped directly into the data domain.
Revised: or state update can be mapped directly into the data domain.

Original: This provides the possibility to inspect the inner workings of the whole network, contributing to its interpretability.
Revised: This provides the possibility to inspect the inner workings of the whole network, contributing to its interpretability.

Original: The attention mechanism in our network contains an important extra term compared to conventional attention.
Revised: The attention mechanism in our network contains an important extra term compared to conventional attention.

Original: We have tested the ET network on the image completion task (qualitatively) and node anomaly detection on graphs (quantitatively).
Revised: We have tested the ET network on the image completion task (qualitatively) and node anomaly detection on graphs (quantitatively).

Original: The qualitative investigation reveals the perfect alignment between the theoretical design principles of our network and its empirical computation.
Revised: The qualitative investigation reveals the perfect alignment between the theoretical design principles of our network and its empirical computation.

Original: The quantitative evaluation demonstrates strong results, which stand in line or exceed the methods recently developed specifically for this task.
Revised: The quantitative evaluation demonstrates strong results, which stand in line or exceed the methods recently developed specifically for this task.

Original: Although we have only tested ET on two tasks, we intentionally picked two entirely different data domains (images and graphs).
Revised: Although we have only tested ET on two tasks, we intentionally picked two entirely different data domains (images and graphs).

Original: We believe that the proposed network will be useful for other tasks and domains and deserves a comprehensive investigation in line with other popular variants of transformers.
Revised: We believe that the proposed network will be useful for other tasks and domains and deserves a comprehensive investigation in line with other popular variants of transformers.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: Section does not count towards page limit
Revised: Section does not count towards page limit

Original: In the experiments presented in this paper we have taken several steps to ensure reproducibility of our results.
Revised: In the experiments presented in this paper we have taken several steps to ensure reproducibility of our results.

Original: Namely, all the training protocols and implementation details including the hyperparameter selection are described in Appendix A and Appendix B. All code ison GitHub and awaiting the paper’s acceptance to be released (we can also release it during the discussion period for ther use by the reviewers).
Revised: Namely, all the training protocols and implementation details including the hyperparameter selection are described in Appendix A and Appendix B. The model and training code is released anonymously.

Original: The code for image reconstruction is written in JAX Bradbury et al.
Revised: The code for image reconstruction is written in JAX Bradbury et al.

Original: (2018) with a single entry script to launch the training process, with defaults set to the configuration that produced the models used in this paper.
Revised: (2018) with a single entry script to launch the training process, with defaults set to the configuration that produced the models used in this paper.

Original: The training script sets a seed that can recreate the exact same training setup as we had, ensuring the exact same weight initialization and random data augmentation provided default arguments.
Revised: The training script sets a seed that can recreate the exact same training setup as we had, ensuring the exact same weight initialization and random data augmentation provided default arguments.

Original: No additional training data was used beyond the training set (ImageNet-1k 2012), which is publicly available.
Revised: No additional training data was used beyond the training set (ImageNet-1k 2012), which is publicly available.

Original: The code for Graph Anomaly Detection is written in PyTorch.
Revised: The code for Graph Anomaly Detection is written in PyTorch.

Original: Given the nature of the anomaly detection problem (random splits into training, validation, testing sets) all our results on graphs are reported with mean and standard deviations, describing typical variability in the performance.
Revised: Given the nature of the anomaly detection problem (random splits into training, validation, testing sets) all our results on graphs are reported with mean and standard deviations, describing typical variability in the performance.

Original: A D ETAILS OF T RAINING ON I MAGE N ET
Revised: A D ETAILS OF T RAINING ON I MAGE N ET

Original: We trained the ET network on a masked-image completion task on the ImageNet-1k (IN1K) dataset.
Revised: We trained the ET network on a masked-image completion task on the ImageNet-1k (IN1K) dataset.

Original: We treat all images in IN1K as images of shape 224 × 224 that are normalized according to standard
Revised: We treat all images in IN1K as images of shape 224 × 224 that are normalized according to standard

Original: IN1K practices (mean 0, variance 1 on the channel dimension) and use data augmentations provided by the popular timm library (Wightman, 2019) (See Table 2).
Revised: IN1K practices (mean 0, variance 1 on the channel dimension) and use data augmentations provided by the popular timm library (Wightman, 2019) (See Table 2).

Original: Following the conventional ViT pipeline (Dosovitskiy et al., 2021), we split these images into non-overlapping patches of 16x
Revised: Following the conventional ViT pipeline (Dosovitskiy et al., 2021), we split these images into non-overlapping patches of 16x

Original: 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.
Revised: 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.

Original: We proceed to randomly and uniformly assign 100 of these tokens as “occluded” which are the only tokens considered by the loss function.
Revised: We proceed to randomly and uniformly assign 100 of these tokens as “occluded” which are the only tokens considered by the loss function.

Original: “Occluded” tokens are designated as follows: of the 100 tokens, 90 tokens are replaced with a learnable MASK token of dimension D and 10 we leave untouched (which we find important for the HN to learn meaningful patch representations).
Revised: “Occluded” tokens are designated as follows: of the 100 tokens, 90 tokens are replaced with a learnable MASK token of dimension D and 10 we leave untouched (which we find important for the HN to learn meaningful patch representations).

Original: To all tokens we then add a distinct learnable position bias.
Revised: To all tokens we then add a distinct learnable position bias.

Original: These tokens are then passed to our Energy Transformer block which we recur for T steps (the “depth” of the model in conventional Transformers).
Revised: These tokens are then passed to our Energy Transformer block which we recur for T steps (the “depth” of the model in conventional Transformers).

Original: At each step, the feedback signal (the sum of the energy gradients from our attention block and HN block) is subtracted from our original token representation with a scalar step size α = dtτ which we treat as a non-learnable hyperparameter in our experiments.
Revised: At each step, the feedback signal (the sum of the energy gradients from our attention block and HN block) is subtracted from our original token representation with a scalar step size α = dtτ which we treat as a non-learnable hyperparameter in our experiments.

Original: The token representations after T steps are passed to a simple linear decoder (consisting of a layer norm and an affine transformation) to project our representations back into the image plane.
Revised: The token representations after T steps are passed to a simple linear decoder (consisting of a layer norm and an affine transformation) to project our representations back into the image plane.

Original: We then use the standard MSE Loss between the original pixels and reconstructed pixels for only the 100 occluded patches.
Revised: We then use the standard MSE Loss between the original pixels and reconstructed pixels for only the 100 occluded patches.

Original: We allow self attention as in the following formula for the energy of multiheaded attention.
Revised: We allow self attention as in the following formula for the energy of multiheaded attention.

Original: [Equation]
Revised: [Equation]

Original: We give details of our architectural choices in Table 2.
Revised: We give details of our architectural choices in Table 2.

Original: In the main paper we present our Energy
Revised: In the main paper we present our Energy

Original: Transformer with a configuration similar to the standard base Transformer configuration (e.g., token dimension 768 , 12 heads each with Y = 64 , softmax’s β = 1 √ Y , ...), with several considerations learned from the qualitative image evaluations:
Revised: Transformer with a configuration similar to the standard base Transformer configuration (e.g., token dimension 768 , 12 heads each with Y = 64 , softmax’s β = 1 √ Y , ...), with several considerations learned from the qualitative image evaluations:

Original: • The dtτ (step size) of 1 implicitly used in the traditional transformer noticeably degrades our ability to smoothly descend the energy function.
Revised: • The dtτ (step size) of 1 implicitly used in the traditional transformer noticeably degrades our ability to smoothly descend the energy function.

Original: We find that a step size of 0 .
Revised: We find that a step size of 0 .

Original: 1 provides a smoother descent down the energy function and benefits the image reconstruction quality.
Revised: 1 provides a smoother descent down the energy function and benefits the image reconstruction quality.

Original: • We observe that our MSE loss must include some subset of un-occluded patches in order for the HN to learn meaningful filters.
Revised: • We observe that our MSE loss must include some subset of un-occluded patches in order for the HN to learn meaningful filters.

Original: • Values of β in the energy attention that are too high prevent our model from training.
Revised: • Values of β in the energy attention that are too high prevent our model from training.

Original: This is possibly due to vanishing gradients in our attention operation from a softmax operation that is too spiky.
Revised: This is possibly due to vanishing gradients in our attention operation from a softmax operation that is too spiky.

Original: • Without gradient clipping, our model fails to train at the learning rates we tried higher than 1e-4.
Revised: • Without gradient clipping, our model fails to train at the learning rates we tried higher than 1e-4.

Original: We observe that gradient clipping not only helps our model train faster at the trainable learning rates, it also allows us to train at higher learning rates.
Revised: We observe that gradient clipping not only helps our model train faster at the trainable learning rates, it also allows us to train at higher learning rates.

Original: Our architecture and experiments for the image reconstruction task were written in JAX (Bradbury et al., 2018) using Flax (Heek et al., 2020).
Revised: Our architecture and experiments for the image reconstruction task were written in JAX (Bradbury et al., 2018) using Flax (Heek et al., 2020).

Original: This engineering choice means that our architecture definitions are quite lightweight, as we can define the desired energy function of the ET and use
Revised: This engineering choice means that our architecture definitions are quite lightweight, as we can define the desired energy function of the ET and use

Original: JAX’s autograd to automatically calculate the desired update.
Revised: JAX’s autograd to automatically calculate the desired update.

Original: All training code and software will be released upon the paper’s acceptance.
Revised: All training code and software will be released upon the paper’s acceptance.

Original: [Table]
Revised: [Table]

Original: A. E XPLORING THE H OPFIELD M EMORIES
Revised: A. E XPLORING THE H OPFIELD M EMORIES

Original: A distinctive aspect of our network is that any variable that has a vector index i of tokens can be mapped into the data domain by applying the decoder network to this variable.
Revised: A distinctive aspect of our network is that any variable that has a vector index i of tokens can be mapped into the data domain by applying the decoder network to this variable.

Original: This makes it possible to inspect all the weights in the model.
Revised: This makes it possible to inspect all the weights in the model.

Original: For instance, the concept of “memories” is crucial to understanding how Hopfield networks function.
Revised: For instance, the concept of “memories” is crucial to understanding how Hopfield networks function.

Original: The memories within the HN module represent the building blocks of all possible image patches in our data domain, where an encoded image patch is a superposition of a subset of memories.
Revised: The memories within the HN module represent the building blocks of all possible image patches in our data domain, where an encoded image patch is a superposition of a subset of memories.

Original: The complete set of memory vectors from the HN module is shown in Figure 5.
Revised: The complete set of memory vectors from the HN module is shown in Figure 5.

Original: The same analysis can be applied to the weights of the ET-attention module.
Revised: The same analysis can be applied to the weights of the ET-attention module.

Original: In Figure 6, we show all the weights from this module mapped into the image plane.
Revised: In Figure 6, we show all the weights from this module mapped into the image plane.

Original: A. B IAS C ORRELATIONS
Revised: A. B IAS C ORRELATIONS

Original: The relationships between our position embeddings exhibit similar behavior to the position correlations of the original ViT in that they are highly susceptible to choices of the hyperparameters (Figure 10 of Dosovitskiy et al.
Revised: The relationships between our position embeddings exhibit similar behavior to the position correlations of the original ViT in that they are highly susceptible to choices of the hyperparameters (Figure 10 of Dosovitskiy et al.

Original: In particular, we consider the effect of weight decay and the β pa- rameter that serves as the inverse temperature of the attention operation (see Equation 3).
Revised: In particular, we consider the effect of weight decay and the β parameter that serves as the inverse temperature of the attention operation (see Equation 3).

Original: The lower the temperature (i e., the higher the value of β ), the spikier the softmax distribution.
Revised: The lower the temperature (i.e., the higher the value of β ), the spikier the softmax distribution.

Original: By using a lower β , we encourage the attention energy to distribute its positional embeddings across a wider range of patches in the model.
Revised: By using a lower β , we encourage the attention energy to distribute its positional embeddings across a wider range of patches in the model.

Original: A. O BSERVING E NERGY D YNAMICS
Revised: A. O BSERVING E NERGY D YNAMICS

Original: We include as part of our supplemental submission a video showing the dynamics of one of our trained models through time together with the corresponding energy at every step.
Revised: We include as part of our supplemental submission a video showing the dynamics of one of our trained models through time together with the corresponding energy at every step.

Original: From the video, it is clear that the image progressively improves in quality as the energy decreases up to the point when the token representations are passed to the decoder.
Revised: From the video, it is clear that the image progressively improves in quality as the energy decreases up to the point when the token representations are passed to the decoder.

Original: At the same time, we found it challenging to find the time constants and number of training steps so that the energy of each image reaches the energy minimum (fixed point) when the loss is computed.
Revised: At the same time, we found it challenging to find the time constants and number of training steps so that the energy of each image reaches the energy minimum (fixed point) when the loss is computed.

Original: For instance, for the image shown in the video, its quality actually starts to degrade when the dynamics is allowed to run for longer than what was used at training (while the energy is still decreasing).
Revised: For instance, for the image shown in the video, its quality actually starts to degrade when the dynamics is allowed to run for longer than what was used at training (while the energy is still decreasing).

Original: Additionally, when training models at greater depth the gradients can vanish since many recurrent applications of the ET block are necessary.
Revised: Additionally, when training models at greater depth the gradients can vanish since many recurrent applications of the ET block are necessary.

Original: We hope to comprehensively investigate these questions/limitations in future work.
Revised: We hope to comprehensively investigate these questions/limitations in future work.

Original: We include a screenshot of the video in Figure 8 and encourage readers to watch the full video for the dynamics.
Revised: We include a screenshot of the video in Figure 8 and encourage readers to watch the full video for the dynamics.

Original: We step through time using a time step of dt = 0 . 1 and record the total energy of our system on the image tokens as the black descending the blue energy curve.

Original: dashed vertical shows the point in the energy curve where the representations were passed to the loss function when training the model, whereas the horizontal red dashed line shows the “fixed at the end of the simulated dynamics (in reality, the energy still descends slightly after that).

Original: We

Original: (masked reconstructed time the time energy trajectory slightly

Original: B D ETAILS OF ET TRAINING ON A NOMALY D ETECTION T ASK
Revised: B D ETAILS OF ET TRAINING ON A NOMALY D ETECTION T ASK

Original: Graph anomaly detection refers to the process of detecting outliers that deviate significantly from the majority of the samples.
Revised: Graph anomaly detection refers to the process of detecting outliers that deviate significantly from the majority of the samples.

Original: Neural network based methods are very popular due to their capability of learning sophisticated data representations.
Revised: Neural network based methods are very popular due to their capability of learning sophisticated data representations.

Original: DOMINANT (Ding et al., 2019) utilizes an auto-encoder framework, using a GCN as an encoder and two decoders for structural reconstruction and attribute reconstruction.
Revised: DOMINANT (Ding et al., 2019) utilizes an auto-encoder framework, using a GCN as an encoder and two decoders for structural reconstruction and attribute reconstruction.

Original: ALARM (Peng et al., 2020) aggregates the encoder information from multiple view of the node attributes.
Revised: ALARM (Peng et al., 2020) aggregates the encoder information from multiple view of the node attributes.

Original: Another study Zhao et al.
Revised: Another study Zhao et al.

Original: (2020), propose a novel loss function to train graph neural networks for anomaly-detectable node representations.
Revised: (2020), propose a novel loss function to train graph neural networks for anomaly-detectable node representations.

Original: In Ding et al.
Revised: In Ding et al.

Original: (2021) generative adversarial learning is used to detect anomaly nodes where a novel layer is designed to learn the anomaly-aware node representation.
Revised: (2021) generative adversarial learning is used to detect anomaly nodes where a novel layer is designed to learn the anomaly-aware node representation.

Original: Recently, Tang et al.
Revised: Recently, Tang et al.

Original: (2022) pointed out that anomalies can lead to the “rightshift” of the spectral energy distribution – the spectral energy concentrates more on the high frequencies.
Revised: (2022) pointed out that anomalies can lead to the “rightshift” of the spectral energy distribution – the spectral energy concentrates more on the high frequencies.

Original: They designed a filter that can better handle this phenomenon.
Revised: They designed a filter that can better handle this phenomenon.

Original: We propose a new anomaly detection model from the perspective of Associative Memory (pattern matching), which does not have the over-smoothing problem often faced by GCNs, and has better model interpretability (outliers should be far from the common pattern).
Revised: We propose a new anomaly detection model from the perspective of Associative Memory (pattern matching), which does not have the over-smoothing problem often faced by GCNs, and has better model interpretability (outliers should be far from the common pattern).

Original: B. D ETAILED M ODEL S TRUCTURE FOR THE G RAPH A NOMALY D ETECTION
Revised: B. D ETAILED M ODEL S TRUCTURE FOR THE G RAPH A NOMALY D ETECTION

Original: First, we compute the features that are passed to our energy-based transformer.
Revised: First, we compute the features that are passed to our energy-based transformer.

Original: Each node’s features y A ∈ R F are mapped into the token space x
Revised: Each node’s features y A ∈ R F are mapped into the token space x

Original: A ∈ R D , using a linear projection E .
Revised: A ∈ R D , using a linear projection E .

Original: Learnable positional embeddings λ
Revised: Learnable positional embeddings λ

Original: A are added to this token at t = 0 ,
Revised: A are added to this token at t = 1 ,

Original: [Equation]
Revised: [Equation]

Original: At each time step the input to the ET-block is layer normalized:
Revised: At each time step the input to the ET-block is layer normalized:

Original: [Equation]
Revised: [Equation]

Original: Let W Q ∈ R Y × H × D
Revised: Let W Q ∈ R Y × H × D

Original: and
Revised: and

Original: W K ∈ R Y × H × D be the query and key weight matrices, respectively.
Revised: W K ∈ R Y × H × D be the query and key weight matrices, respectively.

Original: Here Y is the projection dimension in the attention operation, H is the number of heads.
Revised: Here Y is the projection dimension in the attention operation, H is the number of heads.

Original: We define
Revised: We define

Original: [Equation]
Revised: [Equation]

Original: If we let h indicate the index of the head, we have
Revised: If we let h indicate the index of the head, we have

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Here β controls the temperature of the softmax, N A stands for the neighbors of node A – a set of all the nodes connected to node A , r is the ReLU function.
Revised: Here β controls the temperature of the softmax, N A stands for the neighbors of node A —a set of all the nodes connected to node A , r is the ReLU function.

Revised: Restriction of the attention operation to the neighborhood of a given node is similar to that used in the Graph Attention Networks (GAT), see (Veliˇckovi´c et al., 2017).

Original: Finally, we have residual connection
Revised: Finally, we have residual connection

Original: [Equation]
Revised: [Equation]

Original: Intuitively, the first term considers the influence (attention score) of the neighbor nodes with respect to the target node, the second term considers the influence of the target node with respect to each of its neighbor, and the third term is the contribution of the Hopfield Network module.
Revised: Intuitively, the first term considers the influence (attention score) of the neighbor nodes with respect to the target node, the second term considers the influence of the target node with respect to each of its neighbor, and the third term is the contribution of the Hopfield Network module.

Original: It can be shown that the forward pass of our energy-based transformer layer minimizes the following energy function:
Revised: It can be shown that the forward pass of our energy-based transformer layer minimizes the following energy function:

Original: [Equation]
Revised: [Equation]

Original: This energy function will decrease as the forward pass progresses until it reaches a local minimum.
Revised: This energy function will decrease as the forward pass progresses until it reaches a local minimum.

Original: After T iterations when the retrieval is stable, we have the final representation for each node g final A as g final A =
Revised: After T iterations when the retrieval is stable, we have the final representation for each node g final A as [Equation]

Original: [-g t =
-]
Revised: [-g t =
-][Equation]

Original: A || g t = T A where || is the concatenation sign.
Revised: where || is the concatenation sign.

Original: Following Tang et al.
Revised: Following Tang et al.

Original: , we treat anomaly detection as semisupervised learning task in this work.
Revised: , we treat anomaly detection as semisupervised learning task in this work.

Original: The node representation g final A is fed to another MLP with the sigmoid function to compute the abnormal probability p A , weighted cross entropy is then used to train the network.
Revised: The node representation g final A is fed to another MLP with the sigmoid function to compute the abnormal probability p A , weighted log-likelihood is then used to train the network.

Original: The loss function is as follow:
Revised: The loss function is as follow:

Original: [Equation]
Revised: [Equation]

Original: where σ is the ratio of anomaly labels ( l A = 1) to normal labels ( l A = 0).
Revised: where σ is the ratio of normal labels ( l A = 0) to anomaly labels ( l A = 1).

Original: B. E XPERIMENTAL D ETAILS
Revised: B. E XPERIMENTAL D ETAILS

Original: 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-F1 on the validation set for reporting the final results on the test set.
Revised: 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-F1 on the validation set for reporting the final results on the test set.

Original: Following Tang et al.
Revised: Following Tang et al.

Original: , we use training ratios 1% and 40% respectively (randomly select 1% and 40% nodes of the dataset to train the model, and use the remaining nodes for the validation and testing).
Revised: , we use training ratios 1% and 40% respectively (randomly select 1% and 40% nodes of the dataset to train the model, and use the remaining nodes for the validation and testing).

Original: These remaining nodes are split 1:2 for validation:testing.
Revised: These remaining nodes are split 1:2 for validation:testing.

Original: The statistics of the datasets are listed in Table 3.
Revised: The statistics of the datasets are listed in Table 3.

Original: 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 Tang et al.
Revised: 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 Tang et al.

Original: [Table]
Revised: [Table]

Original: We report the average performance of 5 runs on the test datasets.
Revised: We report the average performance of 5 runs on the test datasets.

Original: The hyperparameters of our model are tuned based on the validation set, selecting the best parameters within 100 epochs.
Revised: The hyperparameters of our model are tuned based on the validation set, selecting the best parameters within 100 epochs.

Original: To speedup the training process, for the large graph datasets T-Finance and T-Social, we sample a different subgraph to train for each epoch (subgraphs have 5% of the nodes with respect to the whole training data).
Revised: To speedup the training process, for the large graph datasets T-Finance and T-Social, we sample a different subgraph to train for each epoch (subgraphs have 5% of the nodes with respect to the whole training data).

Original: The hyperparameters include the number of hidden dimensions in ET-attention Y , the number of neurons K in the hidden layer within the Hopfield Network Module, the number of time iterations T , and the number of heads H .
Revised: The hyperparameters include the number of hidden dimensions in ET-attention Y , the number of neurons K in the hidden layer within the Hopfield Network Module, the number of time iterations T , and the number of heads H .

Original: The weights are learned via backpropagation, which includes embedding projection E , positional embedding λ A , softmax inverse temperature parameter β , ET-attention weight tensors W Q and W K .
Revised: The weights are learned via backpropagation, which includes embedding projection E , positional embedding λ A , softmax inverse temperature parameter β , ET-attention weight tensors W Q and W K .

Original: The optimal hyperparameters used in Table 1 are reported in Table 4.
Revised: The optimal hyperparameters used in Table 1 are reported in Table 4.

Original: The last row in that table summarizes the range of the hyperparameter search that was performed in our experiments.
Revised: The last row in that table summarizes the range of the hyperparameter search that was performed in our experiments.

Original: In general, we have observed that for small datasets (Yelp,
Revised: In general, we have observed that for small datasets (Yelp,

Original: Amazon, T-Finance) a 1 or 2 applications of our network is sufficient for achieving strong results, for larger datasets (T-Social) more iterations (3) are necessary.
Revised: Amazon, T-Finance) a 1 or 2 applications of our network is sufficient for achieving strong results, for larger datasets (T-Social) more iterations (3) are necessary.

Original: For even bigger dataset (ImageNet) our network needs about 12 iterations.
Revised: For even bigger dataset (ImageNet) our network needs about 12 iterations

Original: [Table]
Revised: [Table]

Revised: all the datasets.

Original: C N OTATIONS
Revised: C N OTATIONS

Original: Table 5 lists all the notations used in this paper.
Revised: Table 5 lists all the notations used in this paper.

Original: [Table]
Revised: [Table]

Revised: D A BLATION S TUDY FOR A TTENTION AND H OPFIELD N ETWORK M ODULES

Revised: As we described in the main text the the ET network consists of two modules processing the tokens in parallel: the attention module (ATT) and the Hopfield Network module (HN).

Revised: The ATT module is responsible for routing the information between the tokens, while the HN module is responsible for reinforcing the token representation to be consistent with the general expectation about the particular data domain.

Revised: It is interesting to explore the contribution that these two subnetworks produce on the task performed by the network.

Revised: In this section we ablate the ET architecture by dropping each of the two subnetworks and measuring the impact of the ablation on the performance.

Revised: D. O N G RAPHS

Revised: The results on graphs are reported in Table 6.

Revised: From this table it is clear that most of the computation is performed by the ATT block on this task, which pools the information about other tokens to the token of interest.

Revised: 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 to the full ET are within the statistical error bars).

Revised: 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.

Revised: In this regime the network can only use the features of a given node for that node’s anomalous status prediction.

Revised: This results in a more significant drop in performance, which is about 5% on average.

Revised: D. O N I MAGES

Revised: The ablation results for image reconstruction are shown in Table 7.

Revised: Each experiment was trained using the same hyperparameter settings as shown in Table 2.

Revised: After training the model on IN1K, we calculate the average MSE on 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.

Revised: We make several conclusions from these ablation studies.

Revised: [Table]

Revised: • We gain several insights regarding the use of “self-attention” in our ET (when a token patch query is allowed to consider itself as a key in the attention weights).

Revised: When both selfattention and HN are present (ET-Full+Self), there is no noticeable benefit over ET-Full for a token to attend to itself.

Revised: In fact, preventing the ATTN energy module from attending to itself slightly improves the performance.

Revised: However, when the HN is removed (ET-NoHN*), we notice that allowing self-attention (ET-NoHN+Self) outperforms the version that prevents self-attention (ET-NoHN).

Revised: • On its own, allowing self-attention (ET-NoHN+Self) in the ATTN module performs nearly as well as the full ET at a fraction of the total parameters.

Revised: However, MSE is a forgiving metric for blurry reconstructions.

Revised: While ATTN can capture the global structure of the image quite well, it does so at the expense of image sharpness (Figure 4).

Revised: • 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.

Revised: Figure 9 shows our best performing model (ET-Full) on the qualitative image reconstructions corresponding to the largest errors across IN1K validation images, averaged across all masking seeds.

Revised: Likewise, Figure 10 shows the lowest errors across IN1K validation images and masking seeds.

Revised: In general, image reconstructions that require ET to produce sharp, high frequency, and high contrast lines negatively impact MSE performance.

Revised: Model

Revised: Has ATTN?

Revised: Allow self-

Revised: [Table]

Revised: [Table]

Revised: ET-Full (Ours) ✓ ✗ ✓

Revised: [Table]

Revised: M

Revised: [Table]

Revised: ET-Full+Self ✓ ✓ ✓

Revised: [Table]

Revised: M

Revised: [Table]

Revised: ET-NoHN+Self ✓ ✓ ✗ 1 .

Revised: 3 M 0 .

Revised: [Table]

Revised: ±

Revised: [Table]

Revised: ET-NoHN ✓ ✗ ✗ 1 .

Revised: 3 M

Revised: [Table]

Revised: ±

Revised: [Table]

Revised: ET-NoATT ✗ ✗ ✓ 2 .

Revised: 5 M

Revised: [Table]

Revised: 825 ±

Revised: [Table]

Revised: E ET FOR HETEROGENEOUS G RAPH

Revised: In this section, we show the performance of our ET model in the heterogeneous graph case.

Revised: For the heterogeneous graph case (where more than one type of edges exist in the graph), we first run our

Revised: ET model for different subgraphs (corresponding to different types of edges), and then aggregate the final representations using two methods.

Revised: We have tried max pooling and concatenation for the aggregation step.

Revised: Max pooling means to pick the largest value across all the subgraph representations, and concatenation means to concatenate the representations obtained from different subgraphs.

Revised: Tableshows the comparison between these two variants of our model and BWGNN (heterogeous case).

Revised: ET performs better than heterogeneous BWGNN on Amazon, but loses to heterogeneous BWGNN on Yelp.

Revised: Interestingly, BWGNN in the heterogeneous setting performs worse than BWGNN in the homogeneous setting on Amazon.

Revised: F G RAPH C LASSIFICATION WITH ET

Revised: As mentioned prior, GNNs have emerged as a popular approach to handle graph-related tasks due to their effective and automatic extraction of graph structural information.

Revised: There have been attempts of leveraging Transformer into the graph domain, but only certain key modules, such as feature aggregation, are replaced in GNN variants by the softmax attention

Revised: Ying et al.

Revised: However, the Transformer model has yet to achieved competitive performance on popular leader boards of graph-level prediction compared to mainstream GNN variants Ying et al.

Revised: In this section we explain how ET can be used for graph classification.

Revised: [Table]

Revised: F. D ETAILS OF G RAPH C LASSIFICATION ET M ODEL

Revised: Given a graph G =

Revised: ( V, A )

Revised: [Table]

Revised: where V = { v 1 , v 2 , .

Revised: , v N } and A ∈ { 0 , 1 } N × N is the adjacency matrix of the graph

Revised: [Table]

Revised: each feature vector x B ∈ R F corresponding to node v B is first projected to the token space ˜ x B ∈ R D via a linear embedding.

Revised: Then, the CLS token ˜ x CLS is concatenated to the set of tokens resulting in ˜ X ∈ R ( N +1)

Revised: × D and the positional embedding ˜ λ ∈ R ( N +1)

Revised: × D is added afterwards.

Revised: To obtain the positional embedding ˜ λ

Revised: [Table]

Revised: the adjacency matrix A is first padded in the upper left corner with ones resulting in ˜ A ∈ { 0 , 1 } ( N +1) × ( N +1) .

Revised: This particular step provides the CLS token full connectivity with all of the nodes in the graph.

Revised: The top k smallest eigen-vectors λ ∈ R ( N +1) × k are then obtained from the eigen-value decomposition of the unnormalized Laplacian matrix ˜

Revised: L

Revised: [Equation]

Revised: and projected to the token space via a linear embedding to form the positional embedding ˜ λ ∈

Revised: R ( N +1)

Revised: × D .

Revised: Meanwhile, the attention in ET is modified to take in

Revised: ˆA ∈ R H

Revised: × ( N +1) × ( N +1)

Revised: the parameterized adjacency tensor, which acts as the weighted ‘attention mask’ that enables the model to consider the graph structural information.

Revised: To obtain ˆA , a 2D-convolutional layer with H filters equals to the number of heads in the attention block, ‘SAME’ padding, and a stride of 1 is performed on the outer product of ˜ X to itself.

Revised: The result is then multiplied with ˜

Revised: A element-wise (denoted by ⊙ ) via broadcasting.

Revised: [Equation]

Revised: Altogether, the resulting energy equation is

Revised: [Equation]

Revised: Additionally, in this implementation, the overall model is consisted of S vertically stacked ET blocks, where each block shares the same number of T depth and has a different LayerNorm.

Revised: Similarly, the token representation ˜ X t,ℓ at each dynamic step t corresponding to a block ℓ is first layer-normalized.

Revised: Keep in mind, ˜ X

Revised: = ˜

Revised: X t =1 , ℓ =1 is the initial token representation.

Revised: [Equation]

Revised: Following the dynamic equations 6, we inject a small amount of noise ϵ t, ℓ ∈ (0 , 1) , generated from a normal distribution with standard deviation of 0 .

Revised: 02 and zero mean, into the gradient of energy function to produce ˜ X t +1 , ℓ , the new token representation of block ℓ .

Revised: The premise of this noise injection is to ‘robustify’ the model and help push it towards a local minimum of the energy function.

Revised: [Equation]

Revised: Once stability is reached in the retrieval dynamics of block ℓ , the final representation X t = T, ℓ is then passed on to the next block ℓ + 1 and the whole process is repeated again.

Revised: When the final token representation ˆ Y = ˜

Original: {+X t = T, ℓ = S is computed by the last block S , the resultant CLS token ˆ Y 
+}
Revised: {+X t = T, ℓ = S is computed by the last block S , the resultant CLS token ˆ Y 
+}

Revised: ∈ R D

Revised: ′ extracted from ˆ Y is utilized as the predictor of the current graph G .

Revised: F. E XPERIMENTAL E VALUATION

Revised: Eight datasets of the TUDataset Morris et al.

Revised: (2020) collection are used for experimentation.

Revised: NCI1, NCI109,

Revised: MUTAG, MUTAGENICITY, and FRANKENSTEIN are a common class of graph datasets consists of small molecules with class labels representing toxicity or biological activity determined in drug discovery projects

Revised: Morris et al.

Revised: Meanwhile, DD, ENZYMES, and PROTEINS represent macromolecules.

Revised: The task for both DD and PROTEINS is to classify whether a protein is an enzyme.

Revised: Lastly, for ENZYMES, the task is to assign enzymes to one of the 6 EC-top-level classes, which reflect the catalyzed chemical reaction Morris et al.

Revised: We compare ET with the current state-of-the-art approaches for the mentioned datasets, which include WKPI- kmeans Zhao & Wang (2019), WKPI-kcenters Zhao & Wang (2019), DSGCN Balcilar et al.

Revised: HGP-SL

Revised: Zhang et al.

Revised: U2GNN Nguyen et al.

Revised: and

Revised: Evolution Graph Classifier (EvoG) Domingue et al.

Revised: Additionally, approaches that are close to the baselines are included to further contrast the performance of our model.

Revised: Following the 10-fold cross validation process delineated in Morris et al.

Revised: (2020), accuracy score is used as the evaluation metric and reported in Table 9.

Revised: [Table]

Revised: In general

Revised: [Table]

Revised: we have observed that the modified ET demonstrates consistent performance across all datasets that is near the current state-of-the-art approaches.

Revised: Based on the statistics of the experimental datasets in table 10,

Revised: ET performs extremely well when trained on large graph datasets (e

Revised: [Table]

Revised: g., MUTAGENICITY and FRANKEN- STEIN).

Revised: However, with respect to NCI1, NCI

Revised: [Table]

Revised: and DD datasets, there remains an open question on which graph characteristics (e.g., assortativity and density) impair the performance of the model.

Revised: F.

Revised: E XPERIMENTAL

Revised: D ETAILS

Revised: In the graph domain, it is common to concatenate all of the feature vectors of all graphs in a batch together.

Revised: However, in order for ET to work, we form the batch dimension by separating the feature vectors of all graphs in a given batch and utilize the largest node count to pad all graphs such that they all share the same number of nodes.

Revised: Additionally, we set a limit on the number of nodes (set as 400) to prevent out-of-memory error.

Revised: Specifically, if a graph has a node count exceeding the limit, the number of utilized nodes is equal to the limit.

Revised: Hence, a portion of the graph structural information is ignored as a result.

Revised: However, it is worth mentioning such a graph is rare in the experimental datasets.

Revised: [Table]

Revised: We train all models for 200 epochs using AdamW Loshchilov & Hutter (2017).

Revised: The best model is selected based on its performance obtained from the 10-fold cross validation process delineated in Morris et al.

Revised: Since the task is classification, all models are trained with the cross-entropy loss function with no temperature.

Revised: Additionally, the cosine-annealing with warm-up learning rate scheduler is utilized, where the initial and peak learning rates are set as 5 e − 8 and 0 .

Revised: 001 , respectively.

Revised: The number of warm-up steps is set to 30 epochs while the batch size is 64 for all datasets with the exception of the DD dataset, which requires a batch size of 256 for faster training time.

Revised: The whole experiment is implemented using JAXBradbury et al.

Revised: Flax (Heek et al., 2020), Jraph Godwin* et al.

Revised: (2020), and PyTorch Geometric Fey & Lenssen (2019) packages.

Revised: Lastly, we report the average performance of 10 runs on the 10-fold cross validation process with random seed- ing.

Revised: The hyperparameters of our model are tuned based on the performance of the cross validation, selecting within 100 epochs.

Revised: The optimal hyper-parameters are reported in table 11 and the statistics of the used datasets are reported in table 10.

Revised: [Table]

Revised: G F ORMAL A LGORITHM

Revised: We describe the algorithm for the training and inference of ET in algorithm 1, assuming backpropagation through time using SGD.

Revised: Symbols not defined in the algorithm itself are reported in Table 5.

Revised: We define the

Revised: Infer function to operate independently over each item in a batch.

Revised: H P ARAMETER COMPARISON

Revised: The energy function enforces symmetries in our model, which means ET has fewer parameters than its ViT counterparts.

Revised: In particular, ET has no “Value Matrix” W V in the attention mechanism, and the HN module has only one of the two matrices in the standard MLP of the traditional Transformer block.

Revised: We report these differences in Table 12.

Revised: We take the ET configuration used in this paper, which has an architecture fully comparable to the original ViT-base (Dosovitskiy et al., 2021) with patch size=16 , and report the number of parameters against ViT-base and an “ALBERT” version of ViT (Lan et al., 2020) where a single ViT block is shared across layers.

Revised: We saw no benefit when including biases in ET, so we also exclude the 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.

Revised: [Table]

Revised: 6 Parameters

Revised: Network γ

Revised: δ

Revised: Inputs

Revised: Randomly do height width do

Revised: Convert image patches:

Revised: Embed patches

Revised: Replace

Revised: × b × N

Revised: Reconstruct

Revised: ET:

Revised: return