Memory and Diffusion

From Memorization to Generalization with Diffusion Models

🚧 Under construction

This notebook is under construction. It will be completed by July 14, 2025.

This notebook is a simplified, step-by-step walkthrough of the 2D toy example from the paper: “Memorization to Generalization: Emergence of Diffusion Models from Associative Memory”.

We will train a score-based diffusion model on a small dataset of points lying on a circle. Our goal is to understand how the model learns the data distribution and to visualize its learned “energy landscape,” which reveals how it behaves like an Associative Memory system initially to later transition into a generative model.

For more details, please read the paper and the code repository.

# --- Essential Libraries ---
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm
from itertools import cycle
import os
from IPython.display import display
import PIL
from pathlib import Path

from copy import deepcopy

# --- SciPy for specific math functions ---
from scipy.special import i0, i1 # Modified Bessel functions for analytical energy
import scipy.integrate as integrate # For ODE solving (likelihood calculation)

# --- Scikit-learn for Clustering ---
from sklearn.cluster import AgglomerativeClustering

# Set a nice plot style
sns.set_theme(style="whitegrid")

# turn off warnings
import warnings
warnings.filterwarnings('ignore')
from collections import OrderedDict

# For caching trained models
CACHE_DIR = Path("./cache/02_diffusion_as_memory")
CACHE_DIR.mkdir(parents=True, exist_ok=True)
CACHE_MODELS = True

Synthetic Data

The paper uses a simple dataset: points sampled from the circumference of a unit circle. This helps us easily visualize how the model learns.

We’ll define a function to generate these points and a PyTorch Dataset class to handle them.

def generate_circle_data(num_samples=50000, radius=1, seed=59):
    """Generates data points that lie on a unit circle."""
    np.random.seed(seed)
    # Sample angles uniformly from 0 to 2*pi
    angles = np.random.uniform(0, 2 * np.pi, num_samples)

    # Convert polar coordinates (angles, radius) to Cartesian (x, y)
    x = radius * np.cos(angles)
    y = radius * np.sin(angles)
    return np.stack([x, y], axis=1)

class CircleDataset(Dataset):
    """A PyTorch Dataset to wrap our circle data."""
    def __init__(self, num_samples=50000, radius=1, seed=9):
        # Generate and store the data as a torch tensor
        self.data = torch.from_numpy(generate_circle_data(num_samples, radius, seed)).float()

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

def create_subset(dataset, sample_size, seed=42):
    """Create a subset of the dataset based on the specified sample size. """
    max_size = len(dataset)
    generator = torch.Generator().manual_seed(seed)
    if not 1 <= sample_size <= len(dataset):
        raise ValueError("Sample size must be between 1 and the size of the dataset inclusive.")
    subset, _ = torch.utils.data.random_split(
        dataset, [sample_size, max_size - sample_size], generator=generator
    )
    return subset

def prepare_datasets(sample_size, train_size=60000, test_size=10000, seed=9):
    dataset = CircleDataset(num_samples=train_size, seed=seed)
    test_dataset = CircleDataset(num_samples=test_size, seed=seed)
    train_subset = create_subset(dataset, sample_size)
    test_subset = create_subset(test_dataset, sample_size)
    return train_subset, test_subset

Diffusion models can learn from very few samples. In the paper, this is referred to as memorizing “patterns”. Let’s create a tiny dataset with just 9 data points (patterns) to train on.

# --- Configuration ---
SAMPLE_SIZE = 9 # The number of data points (patterns) to memorize
BATCH_SIZE = min(500, SAMPLE_SIZE)  # Use all data points in each batch
SEED = 9       # For reproducibility

# Create the full dataset
train_subset, _ = prepare_datasets(SAMPLE_SIZE, seed=SEED)

# Create a DataLoader
train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True)

# Extract the training data points for visualization
patterns = train_subset.dataset[train_subset.indices]

Let’s plot our small dataset. These are the specific points we want our model to learn and remember.

The Diffusion Architecture

The diffusion model is simply a model \(s_\theta(\mathbf{x}_t, t)\) which approximates the score function: \[ s_\theta(\mathbf{x}_t, t) ≈ \nabla_{\mathbf{x}_t} \log p_t (\mathbf{x}_t) \] over a series of timesteps.

In this tutorial, we will be using Variance Exploding (VE) SDE, which defines how data is gradually noised over time ranging from \(t \in [\epsilon, 1]\): \[ \mathrm{d} \mathbf{x}_t = \sigma \mathrm{d} \mathbf{w}_t \] and the corresponding reverse process: \[ \mathrm{d} \mathbf{x}_t = \big [ -\sigma^2 \nabla_{\mathbf{x}_t} \log p_t (\mathbf{x}_t) \big ] \mathrm{d}t + \sigma^2 \mathrm{d} \mathbf{w}_t \] where \(g(t) = \sigma\) is the diffusion coefficient and \(\mathbf{w}_t\) is brownian motion.

class VESDETerms:
    """Defines the terms for the Variance Exploding SDE."""
    def __init__(self, sigma_max, device=None):
        self.sigma = sigma_max
        self.device = device

    def marginal_prob_std(self, t):
        t = torch.as_tensor(t, device=self.device, dtype=torch.float32)
        return self.sigma * torch.sqrt(t)

    def diffusion_coeff(self, t):
        t = torch.as_tensor(t, device=self.device, dtype=torch.float32)
        return self.sigma * torch.ones_like(t)

Our score network is a simple Multi-Layer Perceptron (MLP). It takes a noisy data point x and a time step t as inputs, and returns the estimated score. The conditioning on time step t is performed via the Fourier embedding, a standard method of time conditioning in diffusion models.

@torch.no_grad()
def update_ema(ema_model, model, decay=0.9999):
    """
    Step the EMA model towards the current model.
    """
    ema_params = OrderedDict(ema_model.named_parameters())
    model_params = OrderedDict(model.named_parameters())

    for name, param in model_params.items():
        if param.requires_grad == True:
            ema_params[name].mul_(decay).add_(param.data, alpha=1. - decay)

class FourierEmbedding(torch.nn.Module):
    """Embeds time `t` into a high-dimensional feature space."""
    def __init__(self, embed_dim, scale=16):
        super().__init__()
        self.register_buffer('freqs', torch.randn(embed_dim // 2) * scale)


    def forward(self, x):
        x = x.ger((2. * torch.pi * self.freqs).to(x.dtype))
        x = torch.cat([x.cos(), x.sin()], dim=1)
        return x

class ScoreNet(nn.Module):
    """The score-based model (a simple MLP)."""
    def __init__(
        self,
        input_dim=2,
        num_layers=4,
        hidden_dim=128,
        embed_dim=128,
        marginal_prob_std=None
    ):
        super().__init__()

        self.act = nn.SiLU()
        self.marginal_prob_std = marginal_prob_std

        # Time embedding
        self.time_embed = nn.Sequential(
            FourierEmbedding(embed_dim=embed_dim),
            nn.Linear(embed_dim, embed_dim),
            nn.SiLU(),
            nn.Linear(embed_dim, embed_dim),
            nn.SiLU()
        )

        # Project combined (x + time-embedding) to hidden dimension
        self.input_proj = nn.Linear(input_dim + embed_dim, hidden_dim)

        # Hidden MLP layers
        layers = []
        for _ in range(num_layers - 1):
            layers.append(nn.Linear(hidden_dim, hidden_dim))

            if _ == num_layers - 2:
                layers.append(nn.LayerNorm(hidden_dim))
                layers.append(nn.SiLU())
            else:
                layers.append(nn.SiLU())

        self.hidden = nn.Sequential(*layers)

        # Final output to 2 dimensions
        self.output = nn.Linear(hidden_dim, 2)

    def forward(self, x, t):
        # Generate time embedding and concatenate with x
        t_emb = self.time_embed(t)
        h = torch.cat([x, t_emb], dim=1)

        # Pass through MLP
        h = self.input_proj(h)
        h = self.hidden(h)
        h = self.output(h)

        # Scale by 1 / marginal_prob_std(t)
        return h / self.marginal_prob_std(t)[:, None]

Training

We use the denoising score matching (DSM) loss. The goal is to train the ScoreNet model so that its output, the score, matches the direction of the noise z that was added to the clean data x at each time step t. \[ \mathcal{L} = \mathbb{E}_{\mathbf{x}_0, \mathbf{x}_t, t} \, \bigg [ \lambda(t) \lVert s_\theta (\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t | \mathbf{x}_0) \rVert^2 \bigg ] \] where \(\mathbf{x}_0\) denotes the clean data point and \(\mathbf{x}_t\) is the perturbed data point.

For example, assume \(\tilde{\mathbf{x}} \sim \mathcal{N} (\tilde{\mathbf{x}} | \mathbf{x}, \sigma^2 \mathbf{I})\) for the simple case of DSM. We have the following: \[ \nabla_{\tilde{\mathbf{x}}} \log p(\tilde{\mathbf{x}} | \mathbf{x}) = \nabla_\tilde{\mathbf{x}} \bigg ( -\frac{1}{2\sigma^2} (\tilde{\mathbf{x}} - \mathbf{x})^2 \bigg ) = -\frac{\tilde{\mathbf{x}} - \mathbf{x}}{\sigma^2} = -\frac{\epsilon}{\sigma} \] as the score function, which we have to learn for a single timestep of denoising. In the case of diffusion, our mean \(\mu_\theta (\mathbf{x}_t, t)\) is conditioned by time, where \(\mu_\theta (\mathbf{x}_0, 0) = \mathbf{x}\) at t = 0 when we are back at the data distribution.

def loss_fn(model, x, marginal_prob_std, eps=1e-5):
    """The denoising score matching loss function."""
    # Sample a random time t
    random_t = torch.rand(x.shape[0], device=x.device) * (1. - eps) + eps

    # Sample a random noise vector
    z = torch.randn_like(x)
    std = marginal_prob_std(random_t)[:, None]

    # Create the noisy data point
    perturbed_x = x + z * std

    # Get the model's score prediction
    score = model(perturbed_x, random_t)

    # Calculate the loss
    #loss = torch.mean(torch.sum((score * std + z)**2, dim=1))
    loss = torch.mean(torch.square(score * std + z))
    return loss

def train_loop(train_loader, vesde, iterations=100_000, lr=1e-4, device='cuda', log_freq=10_000):
    # create our score model
    score_model = ScoreNet(
        marginal_prob_std=vesde.marginal_prob_std
    )

    # create an exponential moving average version of the model
    ema = deepcopy(score_model).to(device)
    score_model = score_model.to(device)

    optimizer = torch.optim.Adam(score_model.parameters(), lr=lr)

    # Use an infinite data loader to cycle through our small dataset
    infinite_loader = iter(cycle(train_loader))

    # --- Training ---
    score_model.train()
    running_loss = 0.
    pbar = tqdm(range(iterations))
    for iteration in pbar:
        # Get a batch of data
        x = next(infinite_loader).to(device)

        # Calculate loss
        loss = loss_fn(score_model, x, vesde.marginal_prob_std)
        running_loss += loss.item()

        # Optimize
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        update_ema(ema, score_model) # udpate ema

        # Log progress
        if iteration % log_freq == 0 and iteration > 0:
            pbar.set_description(f"Loss: {running_loss / log_freq:.4f}")
            running_loss = 0.

    # return the exponential moving average model
    ema.eval()
    return ema

Time to train! The following code takes a few minutes to run, but the results are cached after the first run.

# Train our SDE-based diffusion models for training data sizes: 2, 9, and 1000.
SEED = 9       # For reproducibility
LR = 1e-4
SIGMA_MAX = 1.
ITERATIONS = 50_000

# Instantiate the SDE and the Model
device = "cuda" if torch.cuda.is_available() else "cpu" 
vesde = VESDETerms(sigma_max=SIGMA_MAX, device=device)

CACHE_MODELS = True

def make_cache_name(sample_size, seed, lr, sigma_max, iterations):
    return f"ema_model_{sample_size}_{seed}_{lr}_{sigma_max}_{iterations}.pth"

def get_ema_model(cache_name):
    cache_path = CACHE_DIR / cache_name
    if cache_path.exists():
        model = ScoreNet(marginal_prob_std=vesde.marginal_prob_std)
        state_dict = torch.load(cache_path)
        model.load_state_dict(state_dict)
        return model
    else:
        return None # Will need to train it

data_sizes = [2, 9, 1000] # Takes ~5 min on an M1 Pro CPU
ema_set, pattern_set = [], [] # store our ema models and training patterns into two separate lists

for sample_size in data_sizes:
    ema = None        
    cache_name = make_cache_name(sample_size, SEED, LR, SIGMA_MAX, ITERATIONS)
    if CACHE_MODELS:
        ema = get_ema_model(cache_name)
    if ema is None or not CACHE_MODELS: 
        # Train the model if no cache exists
        batch_size = min(500, sample_size)  # Use all data points in each batch
        train_subset, _ = prepare_datasets(sample_size, seed=SEED)
        train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)

        # ensure the initialization of our model remains the same
        torch.manual_seed(SEED)
        ema = train_loop(train_loader, vesde, ITERATIONS, LR, device=device)
        torch.save(ema.state_dict(), CACHE_DIR / f"{cache_name}")
        
        
    # Extract the training data points for visualization
    patterns = train_subset.dataset[train_subset.indices]
    pattern_set.append(patterns)
    ema_set.append(ema)

Sampling from the Trained Model

To generate new samples, we run the diffusion process in reverse. We start with pure random noise (sampled at t=1) and use our trained score model to guide it back towards the data distribution (towards t=0). This is done using a numerical SDE solver, like the Euler-Maruyama method.

def Euler_Maruyama_sampler(score_model,
                           sde,
                           batch_size=64,
                           num_steps=1000,
                           device='cuda',
                           eps=1e-5):
    """Generate samples from the score-based model using the Euler-Maruyama solver."""
    score_model.eval()
    t_end = torch.ones(batch_size, device=device)

    # Start with orthogonalized random noise ~ N(0, sigma_max^2 * I)
    init_x = torch.randn(batch_size, 2, device=device)
    init_x = init_x  / torch.norm(init_x, dim = (1), keepdim=True)
    init_x = init_x * sde.marginal_prob_std(t_end)[:, None]

    time_steps = torch.linspace(1., eps, num_steps, device=device)
    step_size = time_steps[0] - time_steps[1] #dt
    x = init_x

    with torch.no_grad():
        for time_step in tqdm(time_steps, desc="Sampling"):
            batch_time_step = torch.ones(batch_size, device=device) * time_step
            g = sde.diffusion_coeff(batch_time_step)
            # This is the reverse SDE update step
            mean_x = x + (g**2)[:, None] * score_model(x, batch_time_step) * step_size
            eps = torch.randn_like(x)
            noise = torch.sqrt(step_size) * g[:, None] * eps
            x = mean_x + noise

    score_model.train()
    return mean_x # Return the final denoised sample

We can now sample from the trained models.

generated_set = []

for ema in ema_set:
    generated_samples = Euler_Maruyama_sampler(ema, vesde, batch_size=1_000, device=device)
    generated_samples = generated_samples.detach().cpu().numpy()
    generated_set.append(generated_samples)

Compute the Potential Energy of Generated vs. Data Samples

Recall the relationship between energy and probability denoted by the Boltzmann distribution: \[ p_\theta(\mathbf{x}) = \frac{\exp{(-E_\theta(\mathbf{x}))} }{Z_\theta} \] This indicates that our energy (up to a constant) is obtained by computing the negative log-likelihood: \[ -\log p_\theta(\mathbf{x}) = E_\theta(\mathbf{x}) + C \]

Since we are dealing with a non-equilibrium system, that is our diffusion model, we follow the formulations and codes provided in Song et al. (2021) to compute the log-likelihood: \[ \log p_0(\mathbf{x}_0;\mathbf{\theta}) = \log p_T(\mathbf{x}_T; \theta) + \int_0^T \nabla \cdot \tilde{\mathbf{f}}(\mathbf{x}_t, t) \mathrm{d}t \] where \[ \tilde{\mathbf{f}}(\mathbf{x}_t, t) =-\frac{1}{2}\sigma^2 \nabla_{\mathbf{x}_{t}} \log p_t(\mathbf{x}_{t}; \theta) \] for this setting. Keep in mind, \(\nabla \cdot ()\) denotes the laplacian operation.

To derive the above equation, we start with the Fokker-Planck equation, as did in Chen et al. (2018) and Song et al. (2021), which yields the following general probability flow ODE (derived from the forward process SDE): \[ \mathrm{d} \mathbf{x}_t = \tilde{\mathbf{f}}(\mathbf{x}_t, t)\mathrm{d} t + \tilde{\mathbf{g}} (\mathbf{x}_t, t) \mathrm{d} \mathbf{w}_t \] where \[\tilde{\mathbf{f}}(\mathbf{x}_t, t) = \mathbf{f} (\mathbf{x}_t, t) - \frac{1}{2} \nabla \cdot \big[\mathbf{g} (\mathbf{x}_t, t) \mathbf{g} (\mathbf{x}_t, t)^\top \big] - \frac{1}{2} \big [ \mathbf{g} (\mathbf{x}_t, t) \mathbf{g} (\mathbf{x}_t, t)^\top \big ] \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)\] and \(\tilde{\mathbf{g}} (\mathbf{x}_t, t) = 0\).

Here, \(\mathbf{f}(\mathbf{x}_t, t) = 0\) denotes the drift term which vanishes in the VE setting while \(\mathbf{g} (\mathbf{x}_t, t) = \sigma\) which is constant in this setting. Thus, the probability flow ODE for our setting is simply: \[ \mathrm{d} \mathbf{x}_t = \tilde{\mathbf{f}}(\mathbf{x}_t, t)\mathrm{d} t = -\frac{1}{2}\sigma^2 \nabla_{\mathbf{x}_{t}} \log p_t(\mathbf{x}_{t}) \, \mathrm{d}t \]

# Compute Laplacian and Log-Likelihood
def compute_laplacian(score_fn, x, t):
    """Compute the Laplacian of the score function."""
    laplacian = torch.zeros(x.size(0), device=x.device)
    with torch.enable_grad():
        x.requires_grad_(True)
        score = score_fn(x, t)
        for i in range(x.shape[1]):
            grad_score_i = torch.autograd.grad(score[:, i].sum(), x, create_graph=True)[0][:, i]
            laplacian += grad_score_i
    x.requires_grad_(False)
    return laplacian.detach()


def ode_likelihood_with_laplacian(x, score_model, sde, device='cuda', eps=1e-5):
    """Compute the log-likelihood of x by solving the probability flow ODE."""
    shape = x.shape

    def score_eval_wrapper(sample, time_steps):
        """A wrapper for evaluating the score-based model for the ODE solver."""
        sample = torch.tensor(sample, device=device, dtype=torch.float32).reshape(shape)
        time_steps = torch.tensor(time_steps, device=device, dtype=torch.float32).reshape((sample.shape[0], ))
        with torch.no_grad():
            score = score_model(sample, time_steps)
        return score.cpu().numpy().reshape((-1, shape[1])).astype(np.float64)

    def laplacian_eval_wrapper(sample, time_steps):
        """A wrapper for evaluating the Laplacian of the score function."""
        sample = torch.tensor(sample, device=device, dtype=torch.float32).reshape(shape)
        time_steps = torch.tensor(time_steps, device=device, dtype=torch.float32).reshape((sample.shape[0], ))
        laplacian = compute_laplacian(score_model, sample, time_steps)
        return laplacian.cpu().numpy().reshape((-1,)).astype(np.float64)

    def ode_func(t, x_and_logp):
        time_steps = torch.from_numpy(np.ones((shape[0],)) * t).to(device, torch.float32)
        sample = torch.from_numpy(x_and_logp[:-shape[0]].reshape(shape)).to(device, torch.float32)

        with torch.no_grad():
            g = sde.diffusion_coeff(time_steps).cpu().numpy()

            score = score_eval_wrapper(sample, time_steps)
            laplacian = laplacian_eval_wrapper(sample, time_steps)

        drift = -0.5 * g[:, None]**2 * score
        logp_grad = -0.5 * g**2 * laplacian
        return np.concatenate([drift.flatten(), logp_grad], axis = 0)

    init = np.concatenate([x.cpu().numpy().flatten(), np.zeros((shape[0],))])
    res = integrate.solve_ivp(ode_func, (eps, 1.), init, rtol=1e-5, atol=1e-5, method='RK45')
    zp = torch.tensor(res.y[:, -1], device=device)
    z = zp[:-shape[0]].reshape(shape)
    delta_logp = zp[-shape[0]:]

    sigma_max = sde.marginal_prob_std(torch.tensor(1.))

    prior_logp = -shape[1] / 2. * torch.log(2 * np.pi * sigma_max ** 2)
    prior_logp = prior_logp - torch.sum(z ** 2, dim=-1) / (2 * sigma_max ** 2)
    return (prior_logp + delta_logp)

Visualizing the Energy

Since we are generating quite a lot of synthetic data points, we use hierarchical clustering to get a sense of where the concentrations of these new points are at. To be more informative, we are also displaying the energy profile of these concentrations alongside that of the training data points.

At K = 2, we can see that the concentrations of generated points are pretty much surrounding the data points and their energy profile are similar to that of the training data points.

Meanwhile, at K = 9, we now see local minima of the energy that devitate drastically from the training points. These new local minima of the energy are called spurious patterns.

Finally, when K = 1000, the energy now very closely matches that of the DenseAM’s derived exact energy, see below.

def to_energy(loglikelihood, normalize=True):
    # nll = E + C
    energy = -loglikelihood
    if normalize:
         # normalize energy by its minimum
        return energy - energy.min()
    return energy


def plot_combined_landscape(fig, ax, score_model, sde, patterns, samples, labels, centers, t_eval=1e-5, device='cpu', annotate=True):
    """Visualize the energy landscape with annotated points."""
    score_model.eval()

    # 1. Calculate Energy for Patterns and Centers
    # Energy for original patterns
    patterns_tensor = patterns.to(device)
    logp_patterns = ode_likelihood_with_laplacian(patterns_tensor, score_model, sde, device, eps=t_eval)
    energy_patterns = to_energy(logp_patterns).cpu().numpy()

    # Energy for found cluster centers
    centers_tensor = torch.from_numpy(centers).float().to(device)
    logp_centers = ode_likelihood_with_laplacian(centers_tensor, score_model, sde, device, eps=t_eval)
    energy_centers = to_energy(logp_centers).cpu().numpy()

    # 2. Create a grid of points
    bounds=(-1.5, 1.5); resolution=75
    x_ = torch.linspace(bounds[0], bounds[1], resolution)
    y_ = torch.linspace(bounds[0], bounds[1], resolution)
    X, Y = torch.meshgrid(x_, y_, indexing='ij')
    grid_tensor = torch.stack([X.flatten(), Y.flatten()], dim=1).to(device)

    # 3. Calculate the energy and score for the grid
    logp_grid = ode_likelihood_with_laplacian(grid_tensor, score_model, sde, device, eps=t_eval)
    energy_grid = to_energy(logp_grid).cpu().numpy().reshape(resolution, resolution)

    with torch.no_grad():
        vec_t = torch.ones(grid_tensor.shape[0], device=device) * t_eval
        scores = score_model(grid_tensor, vec_t).cpu().numpy()

    # 4. Create the comprehensive plot
    #fig, ax = plt.subplots(1, 1, figsize=(12, 8))
    ax.set_title(f'K = {len(patterns)}', fontsize=16)
    ax.set_aspect('equal'); ax.grid(False)

    # Plot energy contour
    contour = ax.contourf(X.cpu(), Y.cpu(), energy_grid, levels=100, cmap='inferno', zorder=0)

    # Plot the unit circle
    theta = np.linspace(0, 2 * np.pi, 200)
    ax.plot(np.cos(theta), np.sin(theta), color='white', linestyle='--', alpha=0.6, zorder=1, label='Unit Circle')

    # Plot score field
    ax.quiver(grid_tensor[:, 0].cpu(), grid_tensor[:, 1].cpu(),
              scores[:, 0], scores[:, 1], color='white', alpha=0.5,
              width=0.003, headwidth=3, zorder=2)

    # Plot original patterns
    ax.scatter(patterns[:, 0].cpu(), patterns[:, 1].cpu(), marker="*", alpha=0.75,
               s=400, color="deeppink", label="Original Patterns", edgecolor='black', zorder=5)

    if annotate: # turn off annotation and plotting center since there are too many patterns at this point
        # Plot generated samples
        ax.scatter(samples[:, 0], samples[:, 1], c=labels, cmap='viridis',
                   s=15, alpha=0.2, zorder=3, label='Generated Samples')

        # Plot found cluster centers
        ax.scatter(centers[:, 0], centers[:, 1], marker='X', s=250,
                    color='aqua', edgecolor='black', zorder=4, label='Cluster Centers')

        # Add energy annotations for patterns
        for i, p in enumerate(patterns.cpu().numpy()):
            ax.annotate(f'{energy_patterns[i]:.2f}', (p[0], p[1]),
                        xytext=(15, -15), textcoords='offset points', color='deeppink', fontsize=10,
                        weight='bold', bbox=dict(boxstyle="round,pad=0.2", fc="black", ec="lime", lw=1, alpha=0.6))

        # Add energy annotations for cluster centers
        for i, c in enumerate(centers):
            ax.annotate(f'{energy_centers[i]:.2f}', (c[0], c[1]),
                        xytext=(-20, 15), textcoords='offset points', color='aqua', fontsize=10,
                        weight='bold', bbox=dict(boxstyle="round,pad=0.2", fc="black", ec="yellow", lw=1, alpha=0.6))

    ax.set_xlabel('x'); ax.set_ylabel('y')
    ax.set_xlim(bounds); ax.set_ylim(bounds)
    return contour


def plot_energy_surface_3d(fig, ax, X, Y, energy, view_angle=(60, -60), title=None):
    """
    Creates a 3D surface plot of the energy landscape.

    Args:
        X (np.ndarray): Meshgrid for X coordinates.
        Y (np.ndarray): Meshgrid for Y coordinates.
        energy (np.ndarray): 2D array of energy values.
        view_angle (tuple): Tuple of (elevation, azimuth) for the plot's camera angle.
    """
    # Plot the 3D surface
    surface = ax.plot_surface(X, Y, energy, cmap='inferno', rstride=1, cstride=1,
                              linewidth=0, antialiased=True, alpha=0.9)

    # Set labels and title
    ax.set_xlabel('x', fontsize=12)
    ax.set_ylabel('y', fontsize=12)
    ax.set_title(title, fontsize=18)

    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])

    # Set a nice viewing angle
    ax.view_init(elev=view_angle[0], azim=view_angle[1])
    return surface

# @title Visualization of the Energy Landscape Across K training Sizes in 2D
# This threshold determines how close points need to be to be considered in the same cluster.
# You may need to tune this value based on your results.
dist_thresholds = [3, 2, 0.45]
annotates = [True, True, False]
t_evals = [0.15] * len(dist_thresholds)

fig, axs = plt.subplots(1, 3, figsize=(18, 8), constrained_layout=True)
for i, (patterns, generated_samples, ema, threshold, t_eval, annotate) in enumerate(zip(pattern_set, generated_set, ema_set, dist_thresholds, t_evals, annotates)):
    # Perform Agglomerative (Hierarchical) Clustering
    clustering = AgglomerativeClustering(
        n_clusters=None,                                      # We let the algorithm find the clusters based on the threshold
        distance_threshold=threshold
    ).fit(generated_samples)

    # Find the center of each identified cluster
    cluster_labels = clustering.labels_
    n_clusters_found = len(np.unique(cluster_labels))
    #print(f"Found {n_clusters_found} clusters!")

    cluster_centers = np.array([
        generated_samples[cluster_labels == i].mean(axis=0)
        for i in range(n_clusters_found)
    ])

    contour = plot_combined_landscape(
        fig, axs[i],
        ema,
        vesde,
        patterns,
        generated_samples,
        cluster_labels,
        cluster_centers,
        device=device,
        t_eval=t_eval,
        annotate=annotate,
    )

axs[0].legend(loc='lower left', fontsize=12)
# Set the ticks to only be at the min and max
cbar = fig.colorbar(contour, ax=axs[-1], label='Energy (Lower is Better)', shrink=0.8)
vmin, vmax = contour.get_clim()
cbar.set_ticks([vmin, vmax])
cbar.set_ticklabels(['Low', 'High'], fontsize=14)
plt.show()
plt.show()

# @title Visualization of the Energy Landscape Across K training Sizes in 3D
# This code is slow to run, so we cache the figure
CACHE_FIG = True

bounds = (-1.5, 1.5)
resolution = 100
X_grid, Y_grid = torch.meshgrid(
    torch.linspace(bounds[0], bounds[1], resolution),
    torch.linspace(bounds[0], bounds[1], resolution),
    indexing='ij'
)
grid_tensor = torch.stack([X_grid.ravel(), Y_grid.ravel()], dim=1).to(device)

fig_fname = CACHE_DIR / "slow_fig.png"
if CACHE_FIG and fig_fname.exists():
    img = PIL.Image.open(str(fig_fname))
    display(img)
else:
    # 1. Create a grid of points (same as before)

    # Make Figure
    fig, axs = plt.subplots(1, 3, figsize=(20, 8), constrained_layout=True, subplot_kw={"projection": "3d"})

    for i, (patterns, score_model) in enumerate(zip(pattern_set, ema_set)):
        # 2. Calculate the energy for the grid (this is the slow step)
        logp_grid = ode_likelihood_with_laplacian(grid_tensor, score_model, vesde, device=device, eps=0.05)
        energy_grid = to_energy(logp_grid).cpu().numpy()
        energy_grid = energy_grid.reshape(resolution, resolution)

        # 3. Create the plot
        contour = plot_energy_surface_3d(fig, axs[i], X_grid.cpu().numpy(), Y_grid.cpu().numpy(), energy_grid, title=f'K = {len(patterns)}')

    # Set the ticks to only be at the min and max
    cbar = fig.colorbar(contour, ax=axs[-1], label='Energy (Lower is Better)', shrink=0.75, pad=0.25)
    vmin, vmax = contour.get_clim()
    cbar.set_ticks([vmin, vmax])
    cbar.set_ticklabels(['Low', 'High'], fontsize=14)
    plt.savefig(fig_fname)
plt.show()

Exact Energy from Dense Associative Memory

Consider the typical DenseAM’s energy function which involves the logsumexp function: \[ E^\text{AM}(\mathbf{x}) = -\beta^{-1} \log \bigg[\sum\limits_{\mu=1}^K \exp\Big(- \beta \lVert \mathbf{x} - \boldsymbol{\xi}^\mu \rVert^2_2\Big) \bigg] \] which is related to the energy of the diffusion model, derived in our work: \[ E^\text{DM}(\mathbf{x}_t, t) = -2 \sigma^2 t \log\bigg[\sum\limits_{\mu=1}^K \exp \Big(- \frac{\lVert \mathbf{x}_t - \boldsymbol{\xi}^\mu \rVert^2_2}{2 \sigma^2 t}\Big) \bigg] \]

Both of these energies expressed competitions among the memories (or stored data patterns). But, the main difference is in the value of the inverse temperature \(\beta\). In the case of diffusion models, it is alternating over time, i.e., \(\beta_t = \frac{1}{2 \sigma^2 t}\), but for DenseAM, this variable is fixed. Nonetheless, although their dynamical trajetories are slightly different, the fixed points (obtained at \(t \approx 0\)) of both equations are still the same.

Using DenseAM energy, for the case of \(K = 2\), we have the following: \[ E^\text{AM}(\mathbf{x}) = -\beta^{-1} \log \Big[\exp \Big(- \beta \lVert \mathbf{x} - \boldsymbol{\xi}^1 \rVert^2_2\Big) + \exp \Big(- \beta \lVert \mathbf{x} - \boldsymbol{\xi}^2 \rVert^2_2 \Big) \Big] \]

For small finite values of \(\beta\), it is possible for a minimum to exist: \[ \boldsymbol{\eta} = \underset{\mathbf{x}}{\text{arg min}} \, E^\text{AM}(\mathbf{x}) \] such that \(\boldsymbol{\eta} \neq \xi^1\) and \(\boldsymbol{\eta} \neq \xi^2\). This minimum is the spurious pattern.

Assume that the empirical data distribution is \(p(\mathbf{y}) = \frac{1}{K}\sum\limits_{\mu=1}^K \delta^{(N)}(\mathbf{y} - \boldsymbol{\xi}^\mu)\) where \(\boldsymbol{\xi}^\mu\) represents an individual data point (with data size \(K\)).

When the training data size \(K \rightarrow \infty\), this empirical data distirbution becomes a continuous density of states: \[ p(\mathbf{y}) = \frac{1}{\pi} \delta\big(y_1^2+y_2^2-1\big) \] The probability of the generated data is then proportional (up to terms independent of the state \(\mathbf{x}\)) to \[ p(\mathbf{x}) \sim \int\limits_{-\infty}^{+\infty} dy_1 dy_2\ p(\mathbf{y})\ e^{- \beta \lVert \mathbf{x} - \mathbf{y} \rVert^2_2} = e^{-\beta (R^2+1)} I_0(2\beta R) \] where \(I_0(\cdot)\) is a modified Bessel function of the first kind and \(R\) is the radius of the circle. Then, the exact energy of our toy model is the following: \[ E^\text{AM}(R, \phi) = R^2+1 - \frac{1}{\beta} \log\big[I_0(2\beta R)\big] \underset{\beta\rightarrow\infty}{\approx} (R - 1)^2 \] given the polar angle \(\phi\).

# @title The Empirical Energy and Score Function for the Toy Model
def cartesian_to_polar(samples):
    x, y = samples[:, 0], samples[:, 1]
    r = np.sqrt(x ** 2 + y ** 2)
    angles = np.arctan2(y, x)
    return r, angles

def energy_am(samples, beta, normalize=True):
    """
    Computes the energy function
        E^AM(R, phi) = R^2 + 1 - (1 / beta) * log(I_0(2 * beta * R))

    See Eq. (13) in the paper.
    """
    r, _ = cartesian_to_polar(samples)
    energy = r**2 + 1 - (1 / beta) * np.log(i0(2 * beta * r))

    # Shift so that the lowest energy is
    if normalize:
        energy = energy - energy.min()
    return energy

def score_am(samples, beta, epsilon=1e-6):
    """
    Computes the score function
        S^AM(R, phi) = -2 * R - (2 / beta) * I_1(2 * beta * R) / I_0(2 * beta * R)

    See Eq. (18) in the paper.
    """
    r, _ = cartesian_to_polar(samples)

    # Ensure r is not zero to avoid division by zero
    r = np.clip(r, epsilon, np.inf)

    bessel_ratio = i1(2 * beta * r) / i0(2 * beta * r)

    score_r = 2 * (bessel_ratio - r)
    score = score_r[:, None] * samples / r[:, None]

    return score

# @title Visualization of the Empirical Energy Landscape and its Score Function
# Inverse Temperature -- Beta
beta = 1 / 0.05

grid_square = torch.stack([X_grid, Y_grid], dim=1)
scores = score_am(grid_square, beta)
energy = energy_am(grid_square, beta)
energy_grid = energy.reshape(resolution, resolution)

fig = plt.figure(figsize=(15, 8), constrained_layout=True)

ax = fig.add_subplot(121)
ax.set_title('Exact', fontsize=16)
ax.set_aspect('equal'); ax.grid(False)

# Plot energy contour
contour = ax.contourf(X_grid.cpu(), Y_grid.cpu(), energy_grid, levels=100, cmap='inferno', zorder=0)

# Plot the unit circle
theta = np.linspace(0, 2 * np.pi, 200)
ax.plot(np.cos(theta), np.sin(theta), color='white', linestyle='--', alpha=0.6, zorder=1, label='Unit Circle')

# Plot score field
ax.quiver(grid_square[:, 0].cpu(), grid_square[:, 1].cpu(),
          scores[:, 0], scores[:, 1], color='white', alpha=0.5,
          width=0.003, headwidth=3, zorder=2)

ax.set_xlabel('x', fontsize=12)
ax.set_ylabel('y', fontsize=12)

ax = fig.add_subplot(122, projection="3d")
plot_energy_surface_3d(fig, ax, X_grid, Y_grid, energy_grid, view_angle=(60, -60), title='Exact')

cbar = fig.colorbar(contour, ax=ax, label='Energy (Lower is Better)', shrink=0.8, pad=0.08)
vmin, vmax = contour.get_clim()
cbar.set_ticks([vmin, vmax])
cbar.set_ticklabels(['Low', 'High'], fontsize=14)
plt.show()