Squint, and the Transformer looks like a dynamical system.
This notebook is under construction. It will be completed by July 14, 2025.
At its core, the transformer is a stack of \(L\) transformer blocks that takes a length \(N\) sequence of input tokens \(\{\mathbf{x}^{(0)}_1, \ldots, \mathbf{x}^{(0)}_N\}\) and outputs a length \(N\) sequence of output tokens \(\{\mathbf{x}^{(L)}_1, \ldots, \mathbf{x}^{(L)}_N\}\). Each token \(\mathbf{x}^{(l)}_i \in \mathbb{R}^D\) is a vector of dimension \(D\).
When blocks are stacked, the residual connections form a “residual highway” that consists entirely of normalizations and additions from Attention and MLP operations.
Associative Memory (AM) requires a global energy function, where each computation minimizes the total energy of the system. Our goal is to derive an energy function whose gradient looks as much like the Transformer block as possible.
We will now build a kind of associative memory called the “Energy Transformer” (Hoover et al. 2024) that turns the familiar transformer operation into an energy minimization. Energy Transformer (ET) defines a single energy on an \(\mathbf{x} \in \mathbb{R}^{N \times D}\) collection of tokens, where we can think of each token \(\mathbf{x}_B\) as a “particle” that knows some information about itself and needs to figure out what it should become. Some particles (unmasked tokens) already know their identity, while others (masked tokens) only know their position and must discover their identity by interacting with their neighbors.
Minimizing the energy of the Energy Transformer (ET) is a recurrent process. The entire transformer consists of a single Transformer block, and each “layer” of the transformer becomes a gradient descent step down the energy. This gradient descent step looks remarkably like a standard transformer block, complete with attention, MLP-like operations, layer normalizations, and residual connections.
The global energy combines two intuitive ideas: (1) attention energy that encourages masked tokens to align with relevant unmasked tokens, and (2) memory energy that ensures all tokens look like realistic patterns the model has learned. The gradient of each of these energies look like a self-attention and MLP, respectively, with some shared weight constraints.
This is one of those situations where the code ends up being significantly simpler than the equations. We write the equations for completeness, but feel free to skip to Section 2.3 for succinct code.
We describe the energy of a multi-headed attention with \(H\) heads, where the \(h\)-th head of attention is parameterized by \(\mathbf{W}_h^Q, \mathbf{W}_h^K \in \mathbb{R}^{D \times Y}\), where \(Y\) is the “head dimension”. The input to the attention is the normalized token vectors \(\hat{\mathbf{x}} \in \mathbb{R}^{N \times D}\). In the math that follows, we index the heads by \(h=1\ldots H\), the head dimension by \(\alpha=1\ldots Y\), tokens by \(A,B,C=1 \ldots N\), and each token vector by \(i,j=1\ldots D\).
We find it convenient to use Einstein notation for the math, since it maps 1:1 to the einops operations we’ll use in the code. If you aren’t familiar with the notation, check out this awesome tutorial. But fair warning, the equations at first look pretty complicated with all the indices.
One tip for reading equations with lots of indices: you don’t need to remember the shape or order of tensors, just remember the meaning of the indices. The number of subscripts is the number of dimensions of the tensor, and the meaning of each dimension is captured in the index name. For example, let \(B=1\ldots N\) index the token position in a sequence, and let \(i=1\ldots D\) index into each token vector. \(x_{Bi}\) is an element of a 2-dimensional tensor capturing the sequence length \(N\) and token dimension \(D\). Transposes don’t have meaning since things are named, so \(x_{Bi} = x_{iB}\). So long as you know the index semantics, you can read always read the equation. Everything is just scalar multiplication and addition.
The familiar queries and keys are computed as normal linear transformations:
\[ \begin{split} K_{h \alpha B} &= \sum\limits_j W^K_{h \alpha j}\; \hat{x}_{Bj}, \qquad \mathbf{K} \in \mathbb{R}^{H \times Y \times N} \\ Q_{h \alpha C} &= \sum\limits_j W^Q_{h \alpha j}\; \hat{x}_{Cj}, \qquad \mathbf{Q} \in \mathbb{R}^{H \times Y \times N} \end{split} \]
Our familiar “raw attention scores” (pre-softmax) are still the dot-product correlations between each query and key:
\[ A_{hBC} = \sum_{\alpha} K_{h\alpha B} Q_{h\alpha C} \]
Now for the different part: we describe the energy of the attention as the negative log-sum-exp of the attention scores. We will use the \(\beta\) as an inverse-temperature hyperparameter to scale the attention scores.
\[ E^\text{ATT} = -\frac{1}{\beta} \sum_{h=1}^H \sum_{C=1}^N \log \left( \sum_{B \neq C} \exp(\beta A_{hBC}) \right) \tag{1}\]
As we saw in a previous notebook, the negative log-sum-exp is an exponential variation of the Dense Associative Memory. The cool thing is that the gradient of the negative log-sum-exp is the softmax, which is what we’d like to see in the attention update rule.
You may recall that traditional attention also has a value matrix. When we take the gradient of Equation 1, we lose the flexibility to include an independently parameterized values: the values must be a function of the queries and the keys.
In traditional transformers, the MLP (without biases) can be written as a two-layer feedforward network with a ReLU on the hidden activations. The MLP is parameterized by two weight matrices \(\mathbf{V}, \mathbf{W} \in \mathbb{R}^{M \times D}\) where \(M\) is the size of the hidden layer (\(M=4D\) is often viewed as the default expansion factor atop token dimension \(D\)). Let’s again use Einstein notation, where \(\mu=1\ldots M\) indexes the hidden units, \(i,j=1\ldots D\) index the token dimensions, and \(B=1\ldots N\) indexes each token.
\[ \text{MLP}(\hat{\mathbf{x}})_{Bi} = \sum_\mu W_{\mu i} \; \text{ReLU}\left(\sum_j V_{\mu j} \hat{\mathbf{x}}_{Bj}\right) \tag{2}\]
If we assume weight sharing between \(\mathbf{V} = \mathbf{W} = \boldsymbol{\xi}\), this is a gradient descent step down the energy of a Hopfield Network
\[ E^{\text{HN}}(\hat{\mathbf{x}}) = - \sum_{B, \mu} F\left(\sum_j \xi_{\mu j} \hat{\mathbf{x}}_{Bj}\right) \]
with rectified quadratic energy \(F(\cdot) := \frac12 \text{ReLU}(\cdot)^2\). If we say \(f(\cdot) := F'(\cdot) = \text{ReLU}(\cdot)\), the negative gradient of the energy is
\[ -\frac{\partial E^{\text{HN}}(\mathbf{\hat{x}})}{\partial \hat{x}_{Bi}} = \sum_\mu \xi_{\mu i} \; f\left(\sum_j \xi_{\mu j} \hat{\mathbf{x}}_{Bj}\right), \]
which is identical to the MLP operation in Equation 2 with a weight sharing constraint.
It is perfectly reasonable to consider other convex functions \(F\) for use in the energy. Polynomials of higher degree \(n\) or exponential functions are both valid and will yield Dense Associative Memory. However, because traditional Transformers use a ReLU activation, we use a rectified quadratic energy.
Let’s implement the attention energy in code. We will use jax and equinox for our code.
import jax, jax.numpy as jnp, jax.random as jr, jax.tree_util as jtu, jax.lax as lax
import equinox as eqx
from dataclasses import dataclass
from typing import *
import matplotlib.pyplot as plt
import numpy as np
from fastcore.basics import *
from fastcore.meta import *
import matplotlib.pyplot as plt
from jaxtyping import Float, Array
import functools as ft
from nbdev import show_doc
from einops import rearrangeThe EnergyTransformer class captures all the token processing in the entire transformer. There are maybe 7 lines of code that perform the actual energy computation. This single energy function, when paired with a layer-norm, is analogous to the full computation across all layers of a traditional transformer. The only things missing are some some token and position embedding matrices to make it work on real data, but we will do that in the following section.
class ETConfig(eqx.Module):
D: int = 768 # token dimension
H: int = 12 # number of heads
Y: int = 64 # head dimension
M: int = 3072 # MLP size
beta: Optional[float] = None # Inverse temperature for attention, defaults to 1/sqrt(Y)
def get_beta(self): return self.beta or 1/jnp.sqrt(self.Y)
smallETConfig = ETConfig(D=12, H=2, Y=6, M=24)
mediumETConfig = ETConfig(D=128, H=4, Y=32, M=256)
fullETConfig = ETConfig(D=768, H=12, Y=64, M=3072, beta=1/jnp.sqrt(64))
class EnergyTransformer(eqx.Module):
config: ETConfig
Wq: Float[Array, "H D Y"] # Query projection
Wk: Float[Array, "H D Y"] # Key projection
Xi: Float[Array, "M D"]
def __init__(self, key, config: ETConfig=fullETConfig):
self.config = config
key1, key2, key3 = jr.split(key, 3)
self.Wq = jr.normal(key1, (config.H, config.D, config.Y)) / config.Y
self.Wk = jr.normal(key2, (config.H, config.D, config.Y)) / config.Y
self.Xi = jr.normal(key3, (config.M, config.D))
def attn_energy(self, xhat: Float[Array, "N D"]):
beta = self.config.get_beta()
K = jnp.einsum("kd,hdy->khy", xhat, self.Wk)
Q = jnp.einsum("qd,hdy->qhy", xhat, self.Wq)
A = jax.nn.logsumexp(beta * jnp.einsum("khy,qhy->hqk", Q, K), -1)
return -1/beta * A.sum()
def hn_energy(self, xhat: Float[Array, "N D"]):
"ReLU energy of a Hopfield Network"
hid = jnp.einsum("nd,md->nm", xhat, self.Xi)
return -0.5 * (hid.clip(0) ** 2).sum()
def energy(self, xhat: Float[Array, "N D"]):
"Total energy of the Energy Transformer"
return self.attn_energy(xhat) + self.hn_energy(xhat)Note that the xhat inputs above are all layer-normalized tokens. However, like other AMs, we restrict ourselves to using non-linearties that are gradients of a convex Lagrangian function. We will just show this in code, but our “special layernorm” is the same as the standard layer normalization except that we need our learnable scale parameter to be a scalar instead of a vector of shape D.
class EnergyLayerNorm(eqx.Module):
"""Define our primary activation function (modified LayerNorm) as a lagrangian with energy"""
gamma: Float[Array, ""] # Scaling scalar
delta: Float[Array, "D"] # Bias per token
use_bias: bool = False
eps: float = 1e-5
def lagrangian(self, x):
"""Integral of the standard LayerNorm"""
D = x.shape[-1]
xmeaned = x - x.mean(-1, keepdims=True)
t1 = D * self.gamma * jnp.sqrt((1 / D * xmeaned**2).sum() + self.eps)
if not self.use_bias: return t1
t2 = (self.delta * x).sum()
return t1 + t2
def __call__(self, x):
"""LayerNorm. The derivative of the Lagrangian"""
xmeaned = x - x.mean(-1, keepdims=True)
v = self.gamma * (xmeaned) / jnp.sqrt((xmeaned**2).mean(-1, keepdims=True)+ self.eps)
if self.use_bias: return v + self.delta
return v
def energy(self, x):
"""Compute the energy of this Lagrangian through the Legendre Transform"""
return (self(x) * x).sum() - self.lagrangian(x)That’s it! We rely on autograd to do the energy minimization, or the “inference” pass through the entire transformer.
Let’s check that the energies of both attention and memory monotonically decreases and is bounded from below.
key = jr.PRNGKey(11)
et = EnergyTransformer(key, config=smallETConfig)
lnorm = EnergyLayerNorm(gamma=1., delta=jnp.zeros(et.config.D))
def energy_recall(Efn, x_init, nsteps, step_size):
"Simple gradient descent to recall a memory"
@jax.jit
def gd_step(x, i):
energy, grad = jax.value_and_grad(Efn)(lnorm(x))
x_next = x - step_size * grad
return x_next, energy
xhat_init = lnorm(x_init)
final_x, energy_history = jax.lax.scan(
gd_step,
xhat_init,
jnp.arange(nsteps)
)
return final_x, energy_history
x_init = jr.normal(key, (100, et.config.D)) # Layer normalized tokens
final_x, energy_history = energy_recall(et.energy, x_init, nsteps=3000, step_size=0.5)
To make the Energy Transformer described above work on real data, we need to add some necessary addendums to work with image data: the token and position embedding matrices, and some data processing code.
Energy Transformer was originally trained on ImageNet. We will load some validation images of the same expected shape of ImageNet to test the performance of ET.
We build a Patcher class to patchify and unpatchify images, which is mostly a simple wrapper around the rearrange function from einops.
class Patcher(eqx.Module):
"Patchify and unpatchify an image."
image_shape: Iterable[int] # (C, H, W) Image shape
patch_size: int # Square patch size
kh: int # Number of patches in the height direction
kw: int # Number of patches in the width direction
@property
def patch_shape(self): return (self.image_shape[0], self.patch_size, self.patch_size)
@property
def num_patch_elements(self): return ft.reduce(lambda a, b=1: a * b, self.patch_shape)
@property
def num_patches(self): return self.kh * self.kw
def patchify(self, img):
"Turn an image (possibly batched) into a collection of patches."
return rearrange(
img,
"... c (kh h) (kw w)-> ... (kh kw) c h w",
h=self.patch_size,
w=self.patch_size,
)
def unpatchify(self, patches):
"Turn a collection of patches (possibly batched) back into an image."
return rearrange(
patches, "... (kh kw) c h w -> ... c (kh h) (kw w)", kh=self.kh, kw=self.kw
)
def patchified_shape(self):
"The expected shape of a patchified image"
return (self.num_patches, *self.patch_shape)
@classmethod
def from_img(cls, img, patch_size):
"Create a Patcher from an example image."
return cls.from_img_shape(img.shape, patch_size)
@classmethod
def from_img_shape(cls, img_shape, patch_size):
"Create a patcher from a specified image shape."
height, width = img_shape[-2:]
assert (height % patch_size) == 0
assert (width % patch_size) == 0
kh = int(height / patch_size)
kw = int(width / patch_size)
return cls(img_shape, patch_size, kh, kw)It lets us do things like:
# show_doc(Patcher.patchify)# show_doc(Patcher.unpatchify)Let’s create a full ET, complete with embeddings, model that can be used for masked-image inpainting. We say that each image has \(N\) total patches/tokens, where each patch as \(Z = c \times h \times w\) pixels when rasterized. We will use linear embeddings (with biases) to embed and unembed rasterized image patches to tokens.
For interoperability with the original ViT (Dosovitskiy et al. 2020), we will add a CLS token and a MASK token as parameters to the model. Again, most of this code is initializing parameters.
class ImageETConfig(eqx.Module):
image_shape: Tuple[int, int, int] = (3, 224, 224) # (C, H, W) Image shape
patch_size: int = 16 # Square patch size
et_conf: ETConfig = fullETConfig
class ImageEnergyTransformer(eqx.Module):
patcher: Patcher
W_emb: Float[Array, "Z D"]
b_emb: Float[Array, "D"]
W_unemb: Float[Array, "D Z"]
b_unemb: Float[Array, "Z"]
pos_embed: Float[Array, "(N+1) D"] # Don't forget the CLS token!
cls_token: jax.Array
mask_token: jax.Array
et: EnergyTransformer
lnorm: EnergyLayerNorm
config: ImageETConfig
def encode(self, x):
"Turn x from img patches to tokens"
x = self.patcher.patchify(x) # (..., N, Z)
return x @ self.W_emb + self.b_emb # (..., N, D)
def decode(self, x):
"Turn x from tokens to img patches"
x = self.lnorm(x) # (..., N, D)
x = x @ self.W_unemb + self.b_unemb # (..., N, Z)
return self.patcher.unpatchify(x) # (..., C, H, W)
def corrupt_tokens(self, x: jax.Array, mask: jax.Array, n_masked: int=100):
"""Corrupt tokens with MASK tokens wherever `mask` is 1.
`n_masked` needs to be known in advance for JAX JIT to work properly
"""
maskmask = jnp.nonzero(mask == 1, size=n_masked, fill_value=0)
return x.at[maskmask].set(self.mask_token) # (..., N, D)
def prep_tokens(self, x, mask):
"Add CLS+MASK tokens and POS embeddings"
x = self.corrupt_tokens(x, mask)
x = jnp.concatenate([self.cls_token[None], x]) # (..., N+1, D)
return x + self.pos_embed # (..., N+1, D)
@delegates(energy_recall)
def __call__(self, x: jnp.ndarray, mask: jnp.ndarray, **kwargs):
"""A complete pipeline for masked image modeling in ET using `energy_recall`
`kwargs` are passed to `energy_recall`
"""
x = self.patcher.patchify(x)
x = self.encode(x)
x = self.prep_tokens(x, mask) # (..., N+1, D)
final_x, energy_history = energy_recall(self.et.energy, x, **kwargs)
x = x[1:] # Discard CLS token for masked inpainting
xhat = self.lnorm(x)
x = self.decode(xhat)
return self.patcher.unpatchify(x) # (..., C, H, W)
@classmethod
def rand_init(cls, key, config=ImageETConfig()):
key1, key2, key3, key4, key5, key6, key7, key8 = jr.split(key, 8)
patcher = Patcher.from_img_shape(config.image_shape, config.patch_size)
W_emb = jr.normal(key1, (patcher.num_patch_elements, config.et_conf.D)) / config.et_conf.D
b_emb = jr.normal(key2, (config.et_conf.D,))
W_unemb = jr.normal(key3, (config.et_conf.D, patcher.num_patch_elements)) / patcher.num_patch_elements
b_unemb = jr.normal(key4, (patcher.num_patch_elements,))
pos_embed = jr.normal(key5, (patcher.num_patches, config.et_conf.D)) / config.et_conf.D
cls_token = 0.002 * jr.normal(key6, (config.et_conf.D,))
mask_token = 0.002 * jr.normal(key7, (config.et_conf.D,))
pos_embed = 0.002 * jr.normal(key8, (1 + patcher.num_patches, config.et_conf.D)) / config.et_conf.D
return cls(
patcher=patcher,
W_emb=W_emb,
b_emb=b_emb,
W_unemb=W_unemb,
b_unemb=b_unemb,
pos_embed=pos_embed,
cls_token=cls_token,
mask_token=mask_token,
et=EnergyTransformer(key7, config.et_conf),
lnorm=EnergyLayerNorm(gamma=1., delta=jnp.zeros(config.et_conf.D)),
config=config
)imageET = ImageEnergyTransformer.rand_init(key, ImageETConfig())ET has publicly available pretrained weights that can be used for masked-image inpainting. Let’s load those weights into our model.
We train ET on a simple dataset, should work on CPU.
The Transformer is a very flexible computing paradigm that can be used for the two major approaches of modern language modeling: masked token prediction (e.g., BERT and diffusion-style transformers) where you predict the fraction of input tokens that are MASKed using information from the unmasked tokens, and autoregressive language modeling (e.g., GPT-style models), where each token in the input sequence is transformed into the next prediction token.