Table of Contents
Fetching ...

Stochastic Thermodynamics of Associative Memory

Spencer Rooke, Dmitry Krotov, Vijay Balasubramanian, David Wolpert

Abstract

Dense Associative Memory networks (DenseAMs) unify several popular paradigms in Artificial Intelligence (AI), such as Hopfield Networks, transformers, and diffusion models, while casting their computational properties into the language of dynamical systems and energy landscapes. This formulation provides a natural setting for studying thermodynamics and computation in neural systems, because DenseAMs are simultaneously simple enough to admit analytic treatment and rich enough to implement nontrivial computational function. Aspects of these networks have been studied at equilibrium and at zero temperature, but the thermodynamic costs associated with their operation out of equilibrium are largely unexplored. Here, we define the thermodynamic entropy production associated with the operation of such networks, and study polynomial DenseAMs at intermediate memory load. At large system sizes and intermediate and low load, we use dynamical mean field theory to characterize out-of-equilibrium properties, work requirements, and memory transition times when driving the system with corrupted memories. We characterize a failure mode of higher order networks not observed at zero temperature. Further, we develop a method for calculating work and power costs in the mean field limit. Finally, we find tradeoffs between entropy production, memory retrieval accuracy, and operation speed.

Stochastic Thermodynamics of Associative Memory

Abstract

Dense Associative Memory networks (DenseAMs) unify several popular paradigms in Artificial Intelligence (AI), such as Hopfield Networks, transformers, and diffusion models, while casting their computational properties into the language of dynamical systems and energy landscapes. This formulation provides a natural setting for studying thermodynamics and computation in neural systems, because DenseAMs are simultaneously simple enough to admit analytic treatment and rich enough to implement nontrivial computational function. Aspects of these networks have been studied at equilibrium and at zero temperature, but the thermodynamic costs associated with their operation out of equilibrium are largely unexplored. Here, we define the thermodynamic entropy production associated with the operation of such networks, and study polynomial DenseAMs at intermediate memory load. At large system sizes and intermediate and low load, we use dynamical mean field theory to characterize out-of-equilibrium properties, work requirements, and memory transition times when driving the system with corrupted memories. We characterize a failure mode of higher order networks not observed at zero temperature. Further, we develop a method for calculating work and power costs in the mean field limit. Finally, we find tradeoffs between entropy production, memory retrieval accuracy, and operation speed.
Paper Structure (15 sections, 48 equations, 6 figures)

This paper contains 15 sections, 48 equations, 6 figures.

Figures (6)

  • Figure 1: Memories ($\bm{\xi}^\mu$) are stored as energy minimizing network configurations in an energy landscape. We consider two modes of operation: (A) We initialize the network in a partial memory ($\bm{\zeta}$), let it relax under Glauber dynamics, then do work to reinitialize the network into the next partial memory; (B) We do direct continuous work on the system through the control fields ${\bf h}$. We restrict ${\bf h}$ to be a linear combination of corrupted memories.
  • Figure 2: The free energy landscape of single memory polynomial DenseAM network as a function of memory alignment for (A)$k=2$ (Hopfield) (B)$k=3$, (C)$k=4$, (D)$k=6$ networks, and various temperatures (lighter colors = higher temperature (smaller $\beta$)). For $k=2$, the free energy landscape is identical to that of the mean field Ising model. In this case, at low temperature (large $\beta$) the free energy has aligned and anti-aligned ($\phi=\pm 1$) minima, and an unligned ($\phi=0$) maximum, while at high temperature (small $\beta$) the only minimum is unaligned ($\phi=0$). For $k>2$ there is always a local minimum of the free energy at zero alignment for any finite temperature, leading to a spurious stored memory. However, the minima associated with true memory alignment are closer to $\phi = \pm 1$ for the higher order networks at comparable temperature, implying that the memory is more accurately stored in the free energy basin. The walls of the energy valley surrounding the stored memory are steeper for larger $k$; so dynamics that drives an initial state to a free energy minimum will be able to correct a narrower range of errors in alignment of the initial state with the true memory.
  • Figure 3: (A-C) Phase Portraits associated with two alignments for DenseAM networks storing two memories, with relaxation dynamics given by Eq. \ref{['eq:RelaxationDynamics']}. (A) The quadratic (Hopfield) network at low temperature ($\beta=2.0$). (B-C) The cubic network at (B) low temperature ($\beta=2.0$) and (C) intermediate temperature ($\beta=1.0$). Given an initial state $\vec{\phi}(t=0)$, colors indicate which which attractor the dynamics drive the state towards. These correspond to partial alignment (or anti alignment for $k$ even) with each memory, and zero alignment for $k>2$. For $p>2$, additional attractors associated with linear combinations of memories also appear. In black are single trajectories associated with finite $N$ Glauber simulations. (D) The reconstruction error $1-\phi_{eq}$ after relaxation as a function of $\beta$ for DenseAM networks with varying nonlinearities, assuming relaxation is successful. Higher order networks reconstruct memories with greater fidelity when reconstruction is successful. (E) Time taken to relax to within $\epsilon=10^{-4}$ of dynamic fixed points for DenseAM networks with varying nonlinearities, as a function of initial state corruption and at two different temperatures. This relaxation time grows approximately logarthmically in $\gamma$ and in $\epsilon$. (Top) At intermediate temperatures $\beta=.75$, higher order networks relax more quickly in the regime where relaxation is successful. As temperature decreases (Bottom), relaxation times become similar, as the tanh term in Eq. \ref{['eq:RelaxationDynamics']} approaches a step function. (F) Plot of the maximum amount of corruption that DenseAMs can correct at various temperatures. Lower order networks can correct patterns that are more highly corrupted, although at lower fidelity as in (D).
  • Figure 4: The change in free energy density in the (A)$k = 2$, (B)$k = 3$, and (C)$k = 4$ DenseAM networks as they relax from a corrupted pattern to the equilibrium distribution around the reconstructed pattern, when such reconstruction is successful, as a function of inverse temperature $\beta$ and initial corruption $\gamma$. Multiplying by $\beta$ reproduces Eq. \ref{['eq:EPDensity']}.
  • Figure 5: Numerical Demonstration of mean field theory for $k=3$ networks with 3 memories. Glauber simulations for (A)$N=128$ and (B)$N=1024$ Neurons under corrupted driving strategy (C) ($\gamma = .25$) (plotted are alignments with each of the three memories). The mean and variance of trajectories for each are shown in (D) and (E), with the mean field trajectory shown in (F). As $N$ increases, we expect variances in these trajectories to shrink like $1/\sqrt{N}$. The power density consumed and its variances for each of the three cases are shown in (G-I). Integrating this gives the work divided by $N$. Over any closed cycle, the integral of this quantity must be positive. The mean field work density calculated from Eq. \ref{['eq:PowerDensity22']}(I) agrees with that found from simulation at finite $N$(G,H).
  • ...and 1 more figures