Modern Methods in Associative Memory: A Universal Language for Associative Memory

Modern Methods in Associative Memory: A Universal Language for Associative Memory

Speakers: Dmitry Krotov, Benjamin Hoover, Parikshit Ram

Conference: ICML 2025

YouTube: https://slideslive.com/39043832

Overview

This talk, presented by Benjamin Hoover and Dmitry Krotov at ICML 2025, introduces a groundbreaking "universal language" for describing and building a wide array of associative memory (AM) models, from classical Hopfield networks to modern transformer architectures and diffusion models. The central goal is to establish a core set of rules, terminology, and mathematical frameworks that can unify seemingly disparate energy-based associative memories. By leveraging concepts like Legendre transforms and Lagrangians, the speakers demonstrate how diverse activation functions and network dynamics can be expressed within a single, consistent theoretical framework.

The significance of this work lies in its potential to bridge the gap between theoretical models of memory and the practical application of large-scale deep learning systems. By reframing complex architectures like the Transformer as a form of associative memory (dubbed the Energy Transformer), the authors offer new avenues for interpretability, theoretical guarantees, and potentially more efficient, biologically plausible AI. Furthermore, the talk establishes a deep mathematical correspondence between dense associative memories and diffusion models, revealing how phenomena like "misremembering" and "spurious states" in generative AI can be understood through the lens of memory overload and energy landscapes. This unified perspective promises to unlock novel insights into the computational capabilities and failure modes of advanced AI systems.

Background

▶ Watch: Goal: A universal language for associative memory (0:00)

The field of associative memory has seen a proliferation of models, each with its own specific characteristics and mathematical descriptions. From the foundational Hopfield network, which operates in both binary and continuous states, to dense associative memories with higher-order polynomial or exponential activation functions, and more recent hierarchical associative memories or neuron-astrocyte networks, a common challenge has been the lack of a unified language. This fragmentation makes it difficult to compare, combine, or systematically develop new associative memory architectures.

Two primary problems are addressed in this context. First, associative memories can operate with either discrete (binary) or continuous neuron states. A universal language must accommodate both. The proposed solution involves viewing binary activation functions (like the sign function) as the sharp limit of continuous, monotonically increasing functions such as tanh, controlled by an inverse temperature parameter beta. This allows for a relaxation of hard binary constraints into a continuous gradient flow problem, where the goal of a neuron state is to minimize the network's energy.

Second, the vast array of nonlinear activation functions employed in modern neural networks (e.g., tanh, sigmoid, ReLU, ELU, softmax, layer norm) further complicates unification. The key insight here is that all these functions are monotonically increasing. This property implies the existence of a convex function whose gradient represents the activation function. This convex function is identified as the indefinite integral of the activation function and is termed a Lagrangian. The relationship between the activations (output of the nonlinearity) and pre-activations (input to the nonlinearity) is formalized through a Legendre transform, a powerful tool from physics that allows the Lagrangian to be transformed into an energy function that the system seeks to minimize. This framework establishes a fundamental connection between the activation function and the energy landscape, providing a consistent way to describe regularization effects (such as exponential decay) that arise naturally from the Legendre transform machinery.

Key Findings

▶ Watch: Unifying diverse nonlinearities with convex Lagrangians (2:50)

The talk presents several key findings that significantly advance the understanding and application of associative memory:

1. A Universal Language for Associative Memory: The introduction of a consistent mathematical framework using Legendre transforms, Lagrangians, and the beta parameter successfully unifies diverse associative memory models, allowing for a coherent description of both binary and continuous states, as well as arbitrary monotonic nonlinearities. This framework enables the systematic derivation of energy functions for complex neural network components.

1. The Energy Transformer: The speakers demonstrate how a Transformer architecture can be re-derived as a form of associative memory, termed the Energy Transformer. This involves reformulating the multi-head attention and MLP blocks as gradients of modular energy functions. The resulting architecture exhibits interpretable, attractor-like dynamics, where tokens evolve through layers to minimize a total energy function, effectively performing tasks like image inpainting and generalization.

1. Interpretability by Design: A significant benefit of the Energy Transformer is its inherent interpretability. The learned Hopfield network psi matrix (representing stored patterns) can be directly decoded and visualized, providing a clear insight into the model's knowledge base and the building blocks it uses to construct predictions. This contrasts with many modern black-box AI models, offering a pathway to architectures grounded in theory and transparent in their operation.

1. Associative Memory as Generative AI: A profound correspondence is established between dense associative memories and diffusion models. By considering dense associative memories overloaded with training data beyond their storage capacity, they transition from memory retrieval systems to generative AI models. This "misremembering" phenomenon, where multiple training data merge to form novel, plausible but non-existent memories, is presented as a reflection of creativity in generative models.

1. Prediction and Detection of Spurious States: The associative memory theory predicts the existence of a "third regime" in diffusion models, termed spurious states. These are point-like attractors in the energy landscape that do not coincide with any training sample but capture a basin of attraction encompassing multiple samples. This theoretical prediction has been empirically validated in conventional diffusion models, offering a new conceptual tool for understanding generalization and novelty generation in generative AI beyond simple memorization or manifold learning.

Technical Deep Dive

▶ Watch: Legendre transform defines energy from Lagrangians (4:00)

The core of this work lies in developing a unified mathematical framework for associative memory. This framework addresses two major challenges: handling both continuous and binary neuron states, and incorporating the wide variety of nonlinear activation functions found in modern neural networks.

Unifying Binary and Continuous Dynamics:

The talk introduces a beta parameter, analogous to inverse temperature, to bridge the gap between binary and continuous activations. A hard sign function (for binary states) is approximated by a sharply rising tanh function as beta approaches infinity. This allows the discrete energy flip dynamics of binary networks to be understood as the limit of a continuous gradient flow problem, where the neuron state x evolves to minimize the network's energy: $dx/dt = - \nabla E(x)$. Here, x hat represents the activations (output of the nonlinearity) and x the pre-activations (internal state).

Generalized Nonlinearities via Legendre Transform:

The commonality among all monotonic activation functions (e.g., tanh, ReLU, softmax, layer norm) is exploited. For any such function, there exists a convex function, called a Lagrangian, whose gradient is that activation function. Formally, if $x̂ = \nabla \mathcal{L}(x)$, where $x̂$ are activations and $x$ are pre-activations. The relationship between $x$ and $x̂$ is further formalized as conjugate variables via a Legendre transform. The Legendre transform of the Lagrangian, denoted $\mathcal{T}(\mathcal{L})$, yields an energy function: $E(x̂) = \sum_i x̂_i x_i - \mathcal{L}(x)$. This energy term acts as a regularization, introducing exponential decay ($dx/dt = -x$) when no stored patterns are present, stemming directly from the mathematical properties of the Legendre transform.

The Energy Transformer Architecture:

The talk re-frames the Transformer as an Energy Transformer, deriving its components from modular energy functions. The key idea is to express the outputs of multi-head attention and MLP blocks as gradients of corresponding energy modules.

The dynamic token states, referred to as "blue neurons," evolve through layers, with the layer norm acting as the nonlinear activation. The layer norm function is monotonic, allowing its derivation from a Lagrangian. Crucially, the Energy Transformer updates are structured such that attention and MLP operate in parallel, not sequentially as in traditional Transformers, and their outputs are gradients of modular energies. The residual connection operates on the dynamic states. A significant departure from standard Transformers is that the weights are shared across layers, as the system takes repeated steps down a single, consistent energy landscape.

The total energy of an Energy Transformer block is the sum of three components:

1. Layer Norm Energy: Derived from the Legendre transform of the layer norm's Lagrangian. It provides an exponential decay regularization.
2. Multi-Head Attention Energy: This energy component drives query tokens to align with key tokens. It involves learnable parameters WQ and WK (value matrix is a function of these for weight sharing). The energy calculation involves einsum for efficient tensor contraction and the logsumexp function, which is the integral of a softmax. Taking the gradient of this energy with respect to activations produces a softmax and value matrix, resembling the standard attention update.
3. Hopfield Network Energy (for MLP): The MLP block, which projects out and projects back in, is reinterpreted as a Hopfield network where the "projecting back in" matrix is the transpose of the "projecting out" matrix. The Hopfield energy is a similarity function (dot product) between layer-normalized activations and stored patterns (the memory matrix $\Psi$), combined with a separation function (integral of ReLU, or quadratic for classical Hopfield).

The system performs discrete gradient descent over layers/time steps. An initial state of tokens is iteratively updated by taking the automatic gradient of the total energy (sum of layer norm, attention, and Hopfield energies) with respect to the dynamic states. This process guarantees continuous energy decrease and convergence to a fixed point, a hallmark of associative memories.

Associative Memory and Diffusion Models:

Dmitry Krotov establishes a direct mathematical correspondence between dense associative memories and diffusion models. He argues that diffusion models can be viewed as "non-equilibrium dense associative memories" where the beta parameter (inverse temperature) is time-dependent.

The probability distribution of noisy samples in a diffusion model, $p(x_t)$, can be expressed as a sum of Gaussian kernels centered around training data points. Taking the logarithm of this probability defines an "energy of the diffusion model." This energy function is shown to be mathematically identical to the energy function of a commonly used dense associative memory model, provided beta is considered time-dependent, the similarity function is the negative Euclidean distance, and the separation function is exponential.

This equivalence allows phenomena in generative AI, such as misremembering (generation of novel, plausible but non-existent memories), to be explained as a failure of memory retrieval in an overloaded associative memory, where multiple training data points merge into a single basin of attraction. The concept of spurious states emerges as a prediction from this theory: local minima in the energy landscape that do not correspond to any original training sample but arise from the generalization process when the memory is overloaded.

Experimental Setup & Results

▶ Watch: Introducing the Energy Transformer architecture (20:00)

The practical application of the Energy Transformer and the theoretical correspondence with diffusion models are demonstrated through specific experiments.

Energy Transformer for Image Inpainting:

The Energy Transformer framework was scaled and applied to a Vision Transformer setting, operating on image patches. After training, the model's associative memory capabilities were tested using corrupted inputs at time T=0. For instance, 50 masked image tokens (patches) that the model had never seen during training were provided as initial states.

The system then runs its dynamics using the previously described gradient descent process. The results show that the total energy continuously decreases, and as it does, the model successfully inpaints all the masked tokens, reconstructing the original images. This demonstrates the Energy Transformer's ability to perform error correction and pattern completion on real-world data, effectively acting as an associative memory at an unprecedented scale for this type of architecture.

Interpretability through Stored Patterns:

A key benefit highlighted is the inherent interpretability. The learned Hopfield network psi matrix (the memory matrix) from the trained Energy Transformer can be decoded. Each row of this matrix, when reshaped, corresponds to an image patch. By visualizing these decoded patches, one can see the collection of stored patterns that the model has learned during training. These patches represent the model's "knowledge base" or "building blocks." Each patch acts as an attractor; the model constructs and predicts new patches as linear combinations of these attractors. This provides a direct, transparent view into the model's internal representations and decision-making process, a strong argument for "interpretability by design."

Diffusion Models and Energy Landscapes (Toy Example):

To illustrate the connection between dense associative memories and diffusion models, a simple toy model was used: data points living on a two-dimensional circle ($x^2 + y^2 = 1$). A simple diffusion model was trained on varying numbers of samples from this manifold.

• Two Samples: When trained with only two samples, the reconstructed energy profile showed two distinct local minima, directly corresponding to the two training samples. In this regime, the diffusion model behaves like a standard dense associative memory, perfectly memorizing the data.
• Nine Samples: With nine training data points, the energy landscape still showed local minima at the training points. However, new, additional local minima appeared that did not coincide with any training data. Crucially, the overall manifold of the data (the circle) began to emerge, albeit with some imperfections (energy bumps in regions without training data).
• One Hundred to One Thousand Samples: As the number of training samples increased significantly (e.g., 100 to 1000), the model successfully reconstructed the entire manifold of the 2D circle. In this "overloaded" regime, the model no longer functions primarily as a memory retrieval system but as a generative model, capturing the underlying data distribution rather than just memorizing individual points.

Detection of Spurious States:

The theory of associative memories predicts a "third phase" in diffusion model behavior, distinct from pure memorization (point-like attractors at data points) and generalization (memorizing the data manifold). This phase consists of spurious states: point-like attractors where the training sample is not at the bottom of the attractor, and multiple training samples fall within that basin of attraction. This phenomenon is considered the "first sign of generalization." The talk states that this theoretical prediction has been empirically validated in conventional diffusion models, confirming that these models can generate samples that are not direct copies or simple interpolations but novel formations arising from the "misremembering" process.

Practical Implications

▶ Watch: Transformer components as gradients of energy (21:00)

This work has several significant practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML space:

1. Enhanced Interpretability and Trust: The "interpretability by design" feature of energy-based associative memories, particularly the Energy Transformer, offers a compelling alternative to current black-box models. Being able to decode and visualize the learned memory patterns (e.g., image patches) provides direct insight into what the model "knows" and how it constructs its predictions. This can be crucial for applications requiring high levels of trust, debugging, and understanding model biases.

1. Principled Architecture Design: The universal language provides a theoretical foundation for designing new neural network architectures. By expressing components as modular energy functions, developers can build systems with provable properties, such as guaranteed energy minimization and convergence. This moves beyond heuristic design choices towards a more principled, theory-driven approach, potentially leading to more robust and predictable models.

1. New Training Paradigms: The energy-based framework opens doors to a broader range of training methods not readily accessible to standard feedforward networks. The existence of an explicit energy function allows for techniques like contrastive learning, score matching, and other energy-based model training approaches, potentially leading to more efficient or effective learning.

1. Understanding Generative AI Mechanisms: The mathematical correspondence between dense associative memories and diffusion models provides a powerful theoretical lens to understand the inner workings of generative AI. Concepts like "misremembering" and "spurious states" offer explanations for how diffusion models generate novelty and generalize beyond their training data, informing future improvements in generative model design and control.

1. Potential for Energy Efficiency and Neuromorphic Computing: While current implementations on GPUs might be computationally intensive (due to gradient descent over multiple steps per layer), the theoretical foundation of energy minimization aligns well with the principles of analog systems and neuromorphic hardware (e.g., Loihi 2). The dream is that if such hardware can physically implement energy minimization, these models could operate with significantly reduced energy consumption, addressing a major challenge in scaling AI.

1. Tradeoffs and Limitations: The talk also acknowledges practical challenges. Designing these energy-based models introduces new complexities, particularly in balancing the energies of different modular components to create effective architectures. This requires careful consideration and parameter tuning. Furthermore, while an analogy was drawn between multi-head attention and Mixture of Expert (MoE) layers, the current Energy Transformer model lacks an explicit, parameterized router found in modern MoE implementations, which is a point for future development. The generalizability of the "Shakespeare sonnet about transformers" example is presented as a hypothesis, indicating ongoing research into the full creative capabilities of these models.

Key Takeaways

• Unified Framework: A universal language using Legendre transforms and Lagrangians unifies diverse associative memory models, from classical Hopfield networks to modern deep learning components.
• Energy Transformer: The Transformer architecture can be re-derived as an associative memory, where attention, MLP, and layer norm are modular energy functions that drive token dynamics towards an energy minimum.
• Inherent Interpretability: Energy-based models offer "interpretability by design," allowing direct visualization of learned knowledge (e.g., image patches from the Hopfield memory matrix) that guides the model's predictions.
• AM as Generative AI: Dense associative memories, when "overloaded" with data, behave like generative AI models, explaining phenomena like "misremembering" as a form of creative generalization arising from memory failure.
• Spurious States: The theory predicts and empirical evidence confirms the existence of "spurious states" in diffusion models—attractors that don't match training data but represent a form of generalization.
• Future Potential: This theoretical grounding offers pathways for more principled AI architecture design, novel training methodologies, and potentially more energy-efficient implementations on neuromorphic hardware.

About the Speaker(s)

Benjamin Hoover is a key presenter in this talk, focusing on the development of the universal language for associative memory and the detailed technical aspects of the Energy Transformer. His work, as described, involves bridging theoretical physics concepts like Legendre transforms with practical neural network architectures, demonstrating a deep understanding of both fields. He shares that developing this unified language has been a long-standing dream throughout his PhD.

Dmitry Krotov introduces the topic and later delves into the profound connection between dense associative memories and generative AI models, specifically diffusion models. He highlights the concept of "misremembering" as a form of creativity and presents the theoretical prediction and empirical validation of "spurious states." His research interests clearly lie at the intersection of associative memory theory and the computational mechanisms underlying modern generative AI.

Parikshit Ram is mentioned as the upcoming speaker for the next section of the tutorial, which will delve further into the specific dense associative memory model discussed by Dmitry Krotov and its connection to kernel density estimation. This suggests his expertise in the mathematical underpinnings and practical applications of these advanced memory models.