## The Classical Hopfield Network
Let’s revisit our task to store *K* binary patterns each of dimension
*D* into an energy function. Let’s keep things simple and fast for the
first part of this notebook and focus on storing and retrieving *K* = 2
patterns: an eevee and pichu, where each `(48,48)` image is rasterized
to a vector dimension of *D* = 2304.
``` python
desired_names = ["eevee", "pichu"]
eevee_pichu_idxs = [poke_names.index(name) for name in desired_names]
Xi = data[eevee_pichu_idxs]
fig, ax = show_im(Xi, figsize=(6,3));
ax.set_title("Stored patterns")
plt.show()
print(f"K={Xi.shape[0]}, D={Xi.shape[1]}")
```
K=2, D=2304
The Classical Hopfield Network (CHN) \[1\] defines an energy function
for this collection of patterns, putting the *μ*-th stored pattern
*ξ**μ* at a *low* value of energy. The CHN energy is a
quadratic function described by dot-product correlations:
$$
E\_\text{CHN}(\sigma) = -\frac{1}{2} \sum\_\mu \left(\sum\_{i} \xi^\mu_i \sigma_i\right)^2 = -\frac{1}{2} \sum\_{i,j} T\_{ij} \sigma_i \sigma_j.
\qquad(2)$$
We see the familiar equation for CHN energy on the RHS if we expand the
quadratic function, where
$T\_{ij} := \sum\_{\mu=1}^K \xi^\mu_i \xi^\mu_j$ is the matrix of
symmetric synapses. Learned patterns *ξ**μ* are stored in *T*
via a simple, Hebbian learning rule.
The CHN can be easily implemented in code via
``` python
class CHN(BinaryAM):
def energy(
self,
sigma: Float[Array, "D"] # Possibly noisy query pattern
):
"Quadratic energy function for the CHN"
return -0.5 * jnp.sum((self.Xi @ sigma)**2, axis=0)
chn = CHN(Xi)
```
The asynchronous update rule of
Equation 1 uses the
energy difference of a flipped bit to determine whether to keep the flip
or not. That update rule is equivalent to the following, arguably more
familiar update rule, which describes the next state based on the sign
of the *total input current* to the neuron *σ**i*.
$$
\begin{align\*}
\sigma_i^{(t+1)} &\leftarrow \text{sgn}\left(\sum\_{\mu} \xi^\mu_i \sum\_{j \neq i} \left(\xi^\mu_j \sigma_j^{(t)}\right) \right)\\
\text{sgn}(x) &:= \begin{cases}
1 & \text{if } x \geq 0 \\
-1 & \text{if } x \< 0
\end{cases}\quad.
\end{align\*}
$$
This update rule also ensures the network always moves toward lower
energy states. Because the *E*CHN is bounded from below, the
network will eventually converge to a local minimum that (ideally)
corresponds to one of the stored patterns.
Let’s observe the recall process! We’ll start with a noisy version of
the first pattern and see if we can recover it.
``` python
def flip_some_bits(key, x, p=0.1):
"Flip `p` fraction of bits in `x`"
prange = np.array([p, 1-p])
return x * jr.choice(key, np.array([-1, 1]), p=prange, shape=x.shape)
sigma_og = Xi[0]
sigma_noisy = flip_some_bits(jr.PRNGKey(0), sigma_og, 0.2)
show_im(jnp.stack([sigma_og, sigma_noisy]), figsize=(6, 3));
```
For the pedagogical purpose of this notebook, we’ll cache the recall
process and results so we don’t have to run it every time.
``` python
@delegates(BinaryAM.async_recall)
def cached_recall(am, cache_name, sigma_noisy, key=jr.PRNGKey(0), save=True, **kwargs):
"Cache the recall process using key `cache_name`"
npz_fname = Path(CACHE_DIR) / (cache_name + '.npz')
if npz_fname.exists() and CACHE_RECALL:
npz_data = np.load(npz_fname)
sigma_final, frames, energies = npz_data['sigma_final'], npz_data['frames'], npz_data['energies']
print("Loading cached recall data")
else:
sigma_final, (frames, energies) = am.async_recall(sigma_noisy, key=key, **kwargs)
if save: jnp.savez(npz_fname, sigma_final=sigma_final, frames=frames, energies=energies)
return sigma_final, frames, energies
cache_name = 'basic_hopfield_recovery'
sigma_final, frames, energies = cached_recall(chn, cache_name, sigma_noisy, nsteps=12000, key=jr.PRNGKey(5))
```
Loading cached recall data
We can animate the recall process to view the “thinking” process of the
CHN.
### Retrieving “inverted” images
If we initialize a query with *too much* noise, it’s possible to
retrieve the negative of a stored pattern or an “inverted image”.
Because the energy is quadratic, both *σ* and −*σ* produce the same
small value of energy. Whether we retrieve the original *σ* or the
inverted −*σ* is dependent on whether we initialize our query closer to
the original or inverted pattern.
$$
E\_\text{CHN}(-\sigma) = -\frac{1}{2} \left(\sum\_{\mu} \xi^\mu_i (-\sigma_i)\right)^2 = E\_\text{CHN}(\sigma)
$$
Loading cached recall data
Accidentally retrieved the inverted pattern!
### Memory retrieval failure
Unfortunately, the CHN is terrible at storing and retrieving multiple
patterns. If we add even four more patterns into the synaptic memory,
our network will fail to retrieve our eevee.
``` python
Xi = data[eevee_pichu_idxs]
Xi = jnp.concatenate([Xi, jr.choice(jr.PRNGKey(10), data, shape=(4,), replace=False)])
fig, ax = show_im(Xi, figsize=(6, 4));
ax.set_title(f"Stored patterns (K={Xi.shape[0]})")
plt.show()
```
Loading cached recall data
CHN failed to retrieve the correct pattern!
## Dense Associative Memory
The CHN has a *quadratic* energy, which is a special case of a more
general class of models called **Dense Associative Memory** (DenseAM)
\[2\]. If we increase the degree of the polynomial used in the energy
function, we strengthen the coupling between neurons and can store more
patterns into the same synaptic matrix.
The new energy function, written in terms of polynomials of degree *n*
and using the same notation for stored patterns
*ξ**i**μ*, is
$$
\begin{align\*}
E\_\text{DAM}(\sigma) &= -\sum\_{\mu=1}^K F_n\left(\sum\_{i=1}^D \xi^\mu_i \sigma_i\right),\\
\text{where}\\F_n(x) &= \begin{cases} \frac{x^n}{n} & \text{if } x \geq 0 \\ 0 & \text{if } x \< 0 \end{cases}.
\end{align\*}
\qquad(3)$$
A higher degree polynomial gives us more **storage capacity**, which
means that it is easier to retrieve the patterns we have stored in the
network. Note that the higher the degree *n*, the narrower the basins of
attraction, which makes it easier to pack more patterns into the energy
landscape.
## Gotta catch ’em all!
Let’s try to store and retrieve all `1024` pokemon patterns into our
network (though we will only show retrieval for a subset of them for
computational reasons). To do this, we’ll need very large values of *n*,
which is bad for numeric overflow (computers don’t like working in
really really large numbers i.e., `inf` energy regimes).
We’ll implement an exponential version of the DenseAM \[3\].
Specifically, we will use a numerically stable `logsumexp` version
\[4\].
$$
\begin{align\*}
E\_\text{eDAM}(\sigma) &= -\log \sum\_{\mu=1}^K \exp \left(\beta \sum\_{i=1}^D \xi^\mu_i \sigma_i\right)
\end{align\*}
\qquad(5)$$
where increasing the inverse temperature *β* has a similar effect to
increasing *n* in the DenseAM polynomial energy function. Because the
`log` is a monotonically increasing function, the energy minima of the
original energy function are preserved, while simultaneously making the
energy function more numerically stable.
``` python
class ExponentialDenseAM(BinaryAM):
Xi: jax.Array # (K, D) Memory patterns
beta: float = 1.0 # Temperature parameter
def energy(self, sigma):
return -jax.nn.logsumexp(self.beta * self.Xi @ sigma, axis=-1)
```
``` python
# Show larger batch retrieval
Xi = data[:1024]
Nshow = 255
Xi_show = jnp.concatenate([data[eevee_pichu_idxs], jr.choice(jr.PRNGKey(10), Xi, shape=(Nshow - len(eevee_pichu_idxs),), replace=False)])
fig1, ax1 = show_im(Xi_show, figsize=(8,8));
ax1.set_title(f"Random sample of {Nshow} stored patterns")
print(f"Storing {Xi.shape[0]} patterns")
```
Storing 1024 patterns
