Algorithm 1: Training and inference pseudocode of ET for image reconstruction task
HyperParameters   
\(\alpha\) : Energy descent stepsize   
\(\epsilon\) : Learning rate   
\(p\) : Token mask probability   
b: batch size   
Parameters   
\(\mathbf{W}^{K}\in \mathbb{R}^{Y\times H\times D},\mathbf{W}^{Q}\in \mathbb{R}^{Y\times H\times D}\) : Key, Query kernels of the Energy Attention   
\(\xi \in \mathbb{R}^{M\times D}\) : Kernel of Hopfield Network   
\(\gamma_{\mathrm{norm}}\in \mathbb{R},\delta_{\mathrm{norm}}\in \mathbb{R}^D\) : Scale, bias of LayerNorm   
\(\mathrm{MASK}\in \mathbb{R}^{D}\) : Mask token   
\(\delta_{\mathrm{pos}}\in \mathbb{R}^{N\times D}\) : Position bias, added to each token   
\(\mathbf{W}_{\mathrm{enc}}\in \mathbb{R}^{P\times D},\delta_{\mathrm{enc}}\in \mathbb{R}^{D}\) : Kernel, bias of affine Encoder   
\(\mathbf{W}_{\mathrm{dec}}\in \mathbb{R}^{D\times P},\delta_{\mathrm{dec}}\in \mathbb{R}^{D}\) : Kernel, bias of affine Decoder   
Infer Inputs   
Corrupted image tokens \(\tilde{X}\in \mathbb{R}^{N\times D}\)   
Add position biases: \(\tilde{X}\gets \tilde{X} +\delta_{\mathrm{pos}}\) .   
for timesteps \(t = 1,\dots ,T\) do   
Normalize each token: \(\tilde{g}\gets\) LayerNorm( \(\tilde{X};\gamma_{\mathrm{norm}},\delta_{\mathrm{norm}}\) ); \(\tilde{g}\in \mathbb{R}^{N\times D}\)   
Calculate Energy of tokens: \(E\gets\) EnergyTransformer \((\tilde{g};\mathbf{W}^{K},\mathbf{W}^{Q},\xi)\) .. \(E\in \mathbb{R}\)   
\(\tilde{X}\gets \tilde{X} -\alpha \nabla_{\tilde{g}}E;\)   
return \(\tilde{\bar{x}}\)   
Train Inputs   
Dataset \(S_{train}\) with elements \(X\in \mathbb{R}^{\textit{channels}\times \textit{height}\times \textit{width}}\)   
Initialize   
Randomly initialize from \(\mathcal{N}(0,0.02)\) ..   
\(\mathbf{W}^{K},\mathbf{W}^{Q},\xi ,\mathbf{M A S K},\mathbf{W}_{\mathrm{e n c}},\mathbf{W}_{\mathrm{d e c}},\delta_{\mathrm{pos}}\sim \mathcal{N}(0,0.02)\)   
Set other biases to zero: \(\delta_{\mathrm{enc}},\delta_{\mathrm{dec}},\delta_{\mathrm{norm}}\gets 0\)   
Set LayerNorm scale to one: \(\gamma_{\mathrm{norm}}\gets 1\)   
for epoch \(n = 1,\ldots ,N_{\mathrm{epoch}}\) do   
\(S_{epoch}\gets S_{train}\)   
for batch \(B\subset S_{epoch}\) \(B\in \mathbb{R}^{b\times \textit{channels}\times \textit{height}\times \textit{width}}\)   
do   
Convert image into non-overlapping patches: \(B_{patch}\gets\) Patchify(B); \(B_{patch}\in \mathbb{R}^{b\times N\times P}\)   
Embed image patches into tokens: \(X\gets\) Encode(Bpatch; Wenc, \(\delta_{\mathrm{enc}}\) ); \(X\in \mathbb{R}^{b\times N\times D}\)   
Replace image tokens randomly by MASK: \(\tilde{\bar{x}},I_{\mathrm{mask}}\gets \mathrm{Mask}(X;\mathrm{MASK},p)\) \(\tilde{\bar{x}}\in \mathbb{R}^{b\times N\times D},I_{\mathrm{mask}}\in \left\lbrace 0,1\right\rbrace^{b\times N}\)   
Reconstruct tokens with ET: \(\tilde{\bar{x}}\gets\) Infer( \(\tilde{\bar{x}}\) )   
Decode tokens: \(\hat{B}_{patch}\gets\) Decode( \(\tilde{\bar{x}} [I_{\mathrm{mask}}];\) Wdec, \(\delta_{\mathrm{dec}}\) ); \(B_{patch}\in \mathbb{R}^{b\times N\times P}\)   
Calculate MSE loss on corrupted tokens: \(L\gets\) Mean(|BPatch[Imask] - Bpatch[Imask]|2) \(L\in \mathbb{R}\)   
params \(\leftarrow\) params - e\nov.paramsL   
Sepoch \(\leftarrow\) Sepoch \\(