Energy Transformer

Highlights

• Attention that descends an energy. ET’s attention layers are not designed first and explained later. They are derived as the gradient of a specifically engineered energy function, so a forward pass through the network is literally energy minimization.
• One block, applied recurrently. Rather than stacking distinct layers, ET applies a single block repeatedly, and the token representations follow the update \(\mathbf{x}^{t+1} = \mathbf{x}^{t} - \alpha \nabla_{\mathbf{g}} E^{t}\) until the energy converges.
• An interpretable system that converges. Every update is the gradient of a scalar energy, so each module’s contribution can be decoded back into the input space and inspected at any step. Because the energy is non-increasing along the trajectory, those updates shrink and the tokens settle at a fixed point, rather than merely stopping after a fixed depth.
• A bridge between three traditions. Attention supplies the empirical power, energy-based models supply the principled objective, and Dense Associative Memory supplies a well-founded recipe for designing the energy itself.
• Works across modalities. The same architecture completes masked image patches and delivers strong quantitative results on graph anomaly detection and graph classification.

Attention as the Gradient of an Energy

A single block holds two modules that read the same layer-normalized tokens. Energy Attention contributes a term that is low when queries align with keys, so attention here is not a hand-designed operation but a sum in the energy across all attention heads and positions. The Hopfield Network contributes a second term that is low when tokens resemble the stored memory vectors, which is what makes the tokens behave like memories. Their sum, \(E = E^{\text{ATT}} + E^{\text{HN}}\), is the entire objective, and the block updates the tokens by descending its gradient, either in continuous time or in discrete steps of size \(\alpha\). Because both terms are differentiable scalars, the resulting implementation is short: the update is nothing more than automatic differentiation of the energy.

An Interpretable System That Converges

Because every update is the gradient of an energy, the contribution of each module can be decoded back into the input space and simply looked at. The figure follows one image completion across the descent, where at each step the token state receives a Hopfield Network update and an Energy Attention update, and each can be rendered on its own. Early in the descent the two updates are large and coarse as they rough in the masked patches. Late in the descent they have shrunk and the reconstruction has stopped changing, which is what convergence looks like from the outside: the energy has flattened and the tokens have settled at a fixed point.

Node Anomaly Detection

Mean and standard deviation over 5 runs with different train/dev/test splits, following the protocol of Tang et al. Standard deviations are shown only where prior work reports them. Bold marks the best result in each row.

Baselines. GraphConsis (Liu et al., SIGIR 2020) · CARE-GNN (Dou et al., CIKM 2020) · PC-GNN (Liu et al., WWW 2021) · BWGNN (Tang et al., 2022) · GT (Dwivedi & Bresson, 2020). MLP is a plain multilayer perceptron with no graph structure.

Graph Classification

Mean and standard deviation from 100 runs of 10-fold cross validation, following the TUDataset protocol. Entries unavailable in prior work are marked –. The shaded cell in each column is the strongest published baseline for that dataset, and every other entry carries its difference against it: ▲ above, ▼ below.

Baselines. WKPI (Zhao & Wang, NeurIPS 2019) · Spec-GN / Norm-GN (Yang et al., ICML 2022) · GWL-WL (Orsini et al., IJCAI 2015) · HGP-SL (Zhang et al., TKDE 2021) · DSGCN (Balcilar et al., 2020) · U2GNN (Nguyen et al., WWW 2022) · NDP (Bianchi et al., TNNLS 2020) · ASAP (Ranjan et al., 2019) · EvoG (Domingue et al., WNYISPW 2019).

Related Work

The Associative Memory foundations that ET is built on are surveyed in Modern Methods in Associative Memory, and the same energy-based perspective is applied to generative models in Memorization to Generalization: Emergence of Diffusion Models from Associative Memory and Language Diffusion Models are Associative Memories Capable of Retrieving Unseen Data. For background, see my blog post on energy-based modeling.