08_discrete_walk_jump_sampling.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [Discrete Walk-Jump Sampling](https://arxiv.org/abs/2306.12360)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Callable, Iterable, Optional, Tuple, Union\n",
    "import functools\n",
    "from tqdm import trange\n",
    "\n",
    "import jax\n",
    "import jax.config\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "import chex\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "from clu import parameter_overview\n",
    "import optax\n",
    "import tqdm\n",
    "from tensorflow_probability.substrates import jax as tfp\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation\n",
    "\n",
    "plt.rcParams[\"animation.html\"] = \"jshtml\"\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Setup\n",
    "\n",
    "Fundamentally, discrete walk-jump sampling (dWJS) is a method for sampling noisy latents from an energy-based model, and denoising them with a neural network.\n",
    "\n",
    "Thus, to understand dWJS, we first need to understand energy-based models and neural empirical Bayes."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Energy-Based Models\n",
    "\n",
    "Energy-based models (EBMs) model the probability of a data point $y$ with an energy function $E_\\theta(y)$ parameterized by $\\theta$:\n",
    "$$\n",
    "p_\\theta(y) = \\frac{\\exp(-E_\\theta(y))}{Z(\\theta)}\n",
    "$$\n",
    "where $Z(\\theta)$ is the normalizing constant (called the partition function):\n",
    "$$\n",
    "Z(\\theta) = \\int \\exp(-E_\\theta(y)) dy\n",
    "$$\n",
    "Low energy samples have high probability, and vice versa."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### How to train an EBM? \n",
    "\n",
    "We can train an EBM by minimizing the KL divergence between the true distribution $p$ and the model distribution $p_\\theta$. This is equivalent to maximizing the expected value of log-likelihood of the data sampled from $p(y)$:\n",
    "$$\n",
    "\\theta^* = \\argmin_\\theta \\text{KL}(p \\ || \\ p_\\theta) = \\argmax_\\theta \\mathbb{E}_{y\\sim p(y)} [\\log p_\\theta(y)] = \\argmin_\\theta \\mathbb{E}_{y\\sim p(y)} [-\\log p_\\theta(y)]\n",
    "$$\n",
    "\n",
    "This is all standard so far. The problem is that $Z(\\theta)$ is intractable to compute:\n",
    "$$\n",
    "\\mathbb{E}_{y\\sim p(y)} [-\\log p_\\theta(y)] = \\mathbb{E}_{y\\sim p(y)} [E_\\theta(y)] + \\log Z(\\theta) \n",
    "$$ \n",
    "If we use gradient descent to optimize $\\theta$, we need to compute:\n",
    "$$\n",
    "\\nabla_\\theta \\mathbb{E}_{y\\sim p(y)} [-\\log p_\\theta(y)] = \\mathbb{E}_{y\\sim p(y)} [\\nabla_\\theta E_\\theta(y)] + \\nabla_\\theta\\log Z(\\theta) \n",
    "$$ \n",
    "The first term is easy to compute, but the second term is (usually) intractable. The common approach is to approximate the second term via MCMC sampling:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\nabla_\\theta\\log Z(\\theta) &= \\frac{\\nabla_\\theta Z(\\theta)}{Z(\\theta)} \\\\\n",
    "&= \\frac{\\nabla_\\theta \\int \\exp(-E_\\theta(y)) dy}{Z(\\theta)} \\\\\n",
    "&= \\frac{\\int \\nabla_\\theta \\exp(-E_\\theta(y)) dy}{Z(\\theta)} \\\\\n",
    "&= \\frac{\\int - \\nabla_\\theta E_\\theta(y) \\exp(-E_\\theta(y)) dy}{Z(\\theta)} \\\\\n",
    "&= \\frac{\\int - \\nabla_\\theta E_\\theta(y) Z(\\theta)  p_\\theta(y) dy}{Z(\\theta)} \\\\\n",
    "&= \\int - \\nabla_\\theta E_\\theta(y) \\ p_\\theta(y) dy \\\\\n",
    "&= \\mathbb{E}_{y\\sim p_\\theta(y)} [- \\nabla_\\theta E_\\theta(y)] \\\\\n",
    "\\\\\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Thus, we seek to minimize:\n",
    "$$\n",
    "\\mathbb{E}_{y\\sim p(y)} [-\\log p_\\theta(y) = \n",
    "\\mathbb{E}_{y\\sim p(y)} [\\nabla_\\theta E_\\theta(y)] - \\mathbb{E}_{y\\sim p_\\theta(y)} [ \\nabla_\\theta E_\\theta(y)]  \n",
    "$$\n",
    "We are seeking to minimize the energy of positive samples (from the data distribution) and maximize the energy of negative samples (from the model distribution). This is why this approach is also called contrastive divergence."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Langevin MCMC\n",
    "\n",
    "We have computed an estimator for the gradient of the negative log-likelihood (NLL) loss.\n",
    "\n",
    "Note that this estimator requires us to sample from the model distribution $p_\\theta(y)$ at each iteration. We perform this sampling from $p_\\theta(y)$ via Langevin MCMC. Langevin MCMC is similar to a noisy version of gradient ascent on the log-likelihood.\n",
    "\n",
    "Initialize a sample $y_0$ randomly.\n",
    "Then, for each iteration $t$, compute:\n",
    "$$\n",
    "y_{t+1} = y_t + \\delta \\nabla_{y_t} \\log p_\\theta(y_t)  + \\sqrt{2\\delta} \\epsilon_t\n",
    "$$\n",
    "where $\\epsilon_t \\sim \\mathcal{N}(0, I)$.\n",
    "Then, as $t \\rightarrow \\infty$, $y_t$ will appear to be sampled from $p_\\theta(y)$.\n",
    "\n",
    "Note that the partition function $Z(\\theta)$ does not show up in the sampling procedure:\n",
    "$$\n",
    "\\nabla_{y_t} \\log p_\\theta(y_t) = - \\nabla_{y_t} E_\\theta(y_t)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@functools.partial(jax.jit, static_argnames=(\"grad_log_prob_fn\", \"num_steps\"))\n",
    "def langevin_sample(grad_log_prob_fn: Callable[[chex.Array], float], init: chex.Array, delta: float, rng: chex.PRNGKey, num_steps: int):\n",
    "    \"\"\"Langevin sampling from a given log probability function.\"\"\"\n",
    "\n",
    "    def one_step_langevin(y_t: chex.Array, rng: chex.PRNGKey):\n",
    "        eps = jax.random.normal(rng, y_t.shape)\n",
    "        y_next = y_t + delta * grad_log_prob_fn(y_t) + jnp.sqrt(2 * delta) * eps\n",
    "        return y_next, y_next\n",
    "\n",
    "    sampling_rngs = jax.random.split(rng, num_steps)\n",
    "    _, samples = jax.lax.scan(one_step_langevin, init, xs=sampling_rngs, length=len(sampling_rngs))\n",
    "    return samples"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Time!\n",
    "\n",
    "To illustrate the math, we will use a simple 1D example. Let's assume that the data distribution is a mixture of two Gaussians centered at -1 and 1, respectively:\n",
    "$$\n",
    "p(y) = \\frac{1}{2} \\mathcal{N}(y; -1, 0.5) + \\frac{1}{2} \\mathcal{N}(y; 1, 0.5)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfd = tfp.distributions\n",
    "px = tfd.Categorical(probs=[0.5, 0.5])\n",
    "p = tfd.Mixture(\n",
    "    cat=px,\n",
    "    components=[\n",
    "        tfd.Normal(loc=-1., scale=0.5),\n",
    "        tfd.Normal(loc=+1., scale=0.5),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Plot the PDF.\n",
    "y = jnp.linspace(-5., 5., int(1e4))\n",
    "plt.grid()\n",
    "plt.plot(y, p.prob(y))\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('p(y)')\n",
    "plt.title('True PDF')\n",
    "plt.show();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the Langevin MCMC sampling process below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(0)\n",
    "grad_log_prob_fn = jax.jit(jax.grad(lambda y: p.log_prob(y).squeeze()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_rng, rng = jax.random.split(rng)\n",
    "delta = 0.1\n",
    "langevin_samples_from_2 = langevin_sample(grad_log_prob_fn, init=2 * jnp.ones((1,)), delta=delta, rng=sampling_rng, num_steps=500)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.grid()\n",
    "ax.set_xlim(-5., 5.)\n",
    "ax.set_ylim(0., 1.)\n",
    "ax.set_xlabel('y')\n",
    "ax.set_ylabel('p(y)')\n",
    "scatter = ax.scatter([], [], lw=2, color='C0')\n",
    "\n",
    "def animate(i: int):\n",
    "    offsets = [langevin_samples_from_2[:i], p.prob(langevin_samples_from_2[:i])]\n",
    "    offsets = jnp.stack(offsets, axis=-1).squeeze()\n",
    "    scatter.set_offsets(offsets)\n",
    "    # Adjust opacity.\n",
    "    if i > 0:\n",
    "        scatter.set_alpha(jnp.arange(i) ** 2 / i ** 2)\n",
    "    ax.set_title(r'Langevin Sampling Starting from 2 with $\\delta={}$: Step {}'.format(delta, i))\n",
    "    return (scatter,)\n",
    "\n",
    "anim = matplotlib.animation.FuncAnimation(fig, animate, frames=100, interval=100, blit=True)\n",
    "plt.close()\n",
    "anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_rng, rng = jax.random.split(rng)\n",
    "langevin_samples_from_neg_2 = langevin_sample(grad_log_prob_fn, init=-2*jnp.ones((1,)), delta=delta, rng=sampling_rng, num_steps=500)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.grid()\n",
    "ax.set_xlim(-5., 5.)\n",
    "ax.set_ylim(0., 1.)\n",
    "ax.set_xlabel('y')\n",
    "ax.set_ylabel('p(y)')\n",
    "scatter = ax.scatter([], [], lw=2, c='C1')\n",
    "\n",
    "def animate(i: int):\n",
    "    offsets = [langevin_samples_from_neg_2[:i], p.prob(langevin_samples_from_neg_2[:i])]\n",
    "    offsets = jnp.stack(offsets, axis=-1).squeeze()\n",
    "    scatter.set_offsets(offsets)\n",
    "    # Adjust opacity.\n",
    "    # scatter.set_sizes(100 * jnp.ones(i))\n",
    "    if i > 0:\n",
    "        scatter.set_alpha(jnp.arange(i) ** 2 / i ** 2)\n",
    "    ax.set_title(r'Langevin Sampling Starting from -2 with $\\delta={}$: Step {}'.format(delta, i))\n",
    "    return (scatter,)\n",
    "\n",
    "anim = matplotlib.animation.FuncAnimation(fig, animate, frames=100, interval=100, blit=True)\n",
    "plt.close()\n",
    "anim"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check that the histogram of the samples from Langevin MCMC match the data distribution somewhat closely:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(langevin_samples_from_2.flatten(), bins=20, density=True, alpha=0.5, color='C0', label='Starting from 2')\n",
    "plt.hist(langevin_samples_from_neg_2.flatten(), bins=20, density=True, alpha=0.5, color='C1', label='Starting from -2')\n",
    "plt.grid()\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('p(y)')\n",
    "plt.legend()\n",
    "plt.title('Histograms of Langevin Samples')\n",
    "plt.show();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our EBM, we will use a simple 2-layer neural network for the energy function:\n",
    "$$\n",
    "E_\\theta(y) = W_1\\text{softplus}(W_0y + b_0) + b_1\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EnergyBasedModel(nn.Module):\n",
    "    \"\"\"A simple energy-based model.\"\"\"\n",
    "    hidden_size: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, y: chex.Array) -> chex.Array:\n",
    "        if len(y.shape) <= 1:\n",
    "            y = jnp.expand_dims(y, axis=0)\n",
    "        y = nn.Dense(self.hidden_size)(y)\n",
    "        y = jax.nn.softplus(y)\n",
    "        y = nn.Dense(self.hidden_size)(y)\n",
    "        y = jax.nn.softplus(y)\n",
    "        y = nn.Dense(1)(y)\n",
    "        y = jnp.squeeze(y, axis=-1)\n",
    "        return y\n",
    "\n",
    "# Initialize the model.\n",
    "model = EnergyBasedModel(hidden_size=10)\n",
    "dummy_input = jnp.ones((1,))\n",
    "init_params = model.init(rng, dummy_input)\n",
    "energy_fn = jax.jit(model.apply)\n",
    "\n",
    "# Overview of the model parameters.\n",
    "print(parameter_overview.get_parameter_overview(init_params))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the unnormlized probability distribution of the EBM below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = jnp.linspace(-10., 10., int(1e4))\n",
    "y_probs = jax.vmap(lambda y: jnp.exp(-energy_fn(init_params, y)))(y)\n",
    "y_probs /= y_probs.sum() * (y[-1] - y[0]) / len(y)\n",
    "plt.grid()\n",
    "plt.plot(y, y_probs, color='C2')\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('p_model(y)')\n",
    "plt.title('(Approximately Normalized) Model PDF')\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_grad_log_prob_fn(params: optax.Params, energy_fn: Callable[[optax.Params, chex.Array], float]) -> Callable[[chex.Array], chex.Array]:\n",
    "    \"\"\"Creates a function that computes the gradient of the log probability under the EBM.\"\"\"\n",
    "    def grad_log_prob_fn(y: chex.Array) -> chex.Array:\n",
    "        return -jax.grad(lambda y: energy_fn(params, y).squeeze())(y)\n",
    "    return grad_log_prob_fn\n",
    "\n",
    "sampling_rng, rng = jax.random.split(rng)\n",
    "langevin_samples_from_model = langevin_sample(create_grad_log_prob_fn(init_params, energy_fn), init=jnp.zeros((1,)), delta=1, rng=sampling_rng, num_steps=10000)\n",
    "plt.hist(langevin_samples_from_model.flatten(), bins=20, density=True, alpha=0.5, color='C2', label='Model Samples')\n",
    "plt.grid()\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('p(y)')\n",
    "plt.legend()\n",
    "plt.title('Histograms of Langevin Samples')\n",
    "plt.show();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now train the EBM using the gradient estimator from above. We use the last sample from Langevin MCMC as the negative sample.\n",
    "We can simply use automatic differentiation to compute the gradient of the loss!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@functools.partial(jax.jit, static_argnames=(\"energy_fn\", \"num_sampling_steps\", \"take_every_sample\", \"burn_in_samples\"))\n",
    "def ebm_loss_fn(params: optax.Params, energy_fn: Callable[[optax.Params, chex.Array], float], y_true_samples: chex.Array, rng: chex.PRNGKey,\n",
    "            delta: float, num_sampling_steps: int, take_every_sample: int, burn_in_samples: int) -> float:\n",
    "    \"\"\"Computes the EBM loss function.\"\"\"\n",
    "    init_rng, rng = jax.random.split(rng)\n",
    "    init = jax.random.normal(init_rng, y_true_samples[0].shape)\n",
    "    sampling_rng, rng = jax.random.split(rng)\n",
    "    y_model_samples = langevin_sample(create_grad_log_prob_fn(params, energy_fn), init=init, delta=delta, rng=sampling_rng, num_steps=num_sampling_steps)\n",
    "    y_model_samples = y_model_samples[burn_in_samples:]\n",
    "    y_model_samples = y_model_samples[::take_every_sample]\n",
    "    # We don't differentiate through the sampling procedure.\n",
    "    y_model_samples = jax.lax.stop_gradient(y_model_samples)\n",
    "    energy_fn_vmapped = jax.vmap(lambda y: energy_fn(params, y))\n",
    "    return energy_fn_vmapped(y_true_samples).mean() - energy_fn_vmapped(y_model_samples).mean()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are ready to train our model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_ebm_model(init_params: optax.Params, rng: chex.PRNGKey, num_training_steps: int, **loss_kwargs) -> optax.Params:\n",
    "    \"\"\"Train the EBM model using the Adam optimizer.\"\"\"\n",
    "    @jax.jit\n",
    "    def train_step(params: optax.Params, opt_state: optax.OptState, y_true_samples: chex.Array, rng: chex.PRNGKey) -> Tuple[optax.Params, optax.OptState, float]:\n",
    "        loss, grad = jax.value_and_grad(ebm_loss_fn, has_aux=False)(params, energy_fn, y_true_samples, rng, **loss_kwargs)\n",
    "        updates, opt_state = tx.update(grad, opt_state)\n",
    "        params = optax.apply_updates(params, updates)\n",
    "        return params, opt_state, loss\n",
    "\n",
    "    tx = optax.adam(5e-4)\n",
    "    opt_state = tx.init(init_params)\n",
    "\n",
    "    params_at_steps = {\n",
    "        0: init_params,\n",
    "    }\n",
    "    \n",
    "    params = init_params\n",
    "    for step in tqdm.trange(num_training_steps):\n",
    "        step_rng, samples_rng, rng = jax.random.split(rng, num=3)\n",
    "        y_true_samples = p.sample(32, seed=samples_rng)\n",
    "\n",
    "        params, opt_state, loss = train_step(params, opt_state, y_true_samples, step_rng)\n",
    "\n",
    "        # Log the training progress.        \n",
    "        if step % 500 == 0:\n",
    "            params_at_steps[step] = params\n",
    "        \n",
    "        if step % 2000 == 0:\n",
    "            print('Step {}: Loss = {}'.format(step, loss))\n",
    "    \n",
    "    return params, params_at_steps\n",
    "\n",
    "train_rng = jax.random.PRNGKey(0)\n",
    "energy_params, energy_params_at_steps = train_ebm_model(init_params, train_rng, num_training_steps=30000, delta=0.1, num_sampling_steps=10, take_every_sample=2, burn_in_samples=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the learnt (unnormalized) probability distribution of the EBM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.grid()\n",
    "ax.set_xlim(-5., 5.)\n",
    "ax.set_ylim(0., 1.)\n",
    "ax.set_xlabel('y')\n",
    "ax.set_ylabel('p_model(y)')\n",
    "line, = ax.plot([], [], lw=2, color='C2')\n",
    "steps = sorted(energy_params_at_steps.keys())\n",
    "\n",
    "def animate(i: int):\n",
    "    y = jnp.linspace(-5., 5., int(1e3))\n",
    "    y_probs = jax.vmap(lambda y: jnp.exp(-energy_fn(energy_params_at_steps[steps[i]], y)))(y)\n",
    "    y_probs /= y_probs.sum() * (y[-1] - y[0]) / len(y)\n",
    "    line.set_data(y, y_probs)\n",
    "    ax.set_title('(Approximately Normalized) Model PDF: Step {}'.format(steps[i]))\n",
    "    return (line,)\n",
    "\n",
    "anim = matplotlib.animation.FuncAnimation(fig, animate, frames=len(steps), interval=100, blit=True)\n",
    "plt.close()\n",
    "anim"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Empirical Bayes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section follows the work of [Saeed Saremi and Aapo Hyvarinen](https://arxiv.org/abs/1903.02334).\n",
    "\n",
    "Consider an observation denoted by the random variable $X \\in \\mathbb{R}^d$, and a noisy observation of $X$ denoted by $Y \\in \\mathbb{R}^d$:\n",
    "$$\n",
    "Y = X + \\epsilon, \\quad \\epsilon \\sim \\mathcal{N}(0, \\sigma^2 I_d)\n",
    "$$\n",
    "Thus, we are given that $p_{Y|X}$ is a Gaussian.\n",
    "\n",
    "Given the observation $Y = y$, the Bayes' least squares estimator for $X$ is:\n",
    "$$\n",
    "\\hat{X}(y) = \\mathbb{E}[X \\ | \\ Y = y]\n",
    "$$\n",
    "Computing this estimator requires knowledge of $p_{X|Y} \\propto p_{Y|X}\\cdot p_X$ by Bayes' rule.\n",
    "Thus, it seems that we need to know $p_X$ to compute this estimator.\n",
    "\n",
    "The trick (figured out by [Robbins](https://link.springer.com/chapter/10.1007/978-1-4612-0919-5_26) and [Miyasawa](https://mit.primo.exlibrisgroup.com/discovery/openurl?institution=01MIT_INST&vid=01MIT_INST:MIT&rft.epage=188&rft_val_fmt=info:ofi%2Ffmt:kev:mtx:journal&rft.stitle=B%20INT%20STATIST%20INST&rft.volume=38&rfr_id=info:sid%2Fwebofscience.com:WOS:WOS&rft.jtitle=BULLETIN%20OF%20THE%20INTERNATIONAL%20STATISTICAL%20INSTITUTE&rft.aufirst=K&rft.genre=article&rft.issue=4&rft.pages=181-188&url_ctx_fmt=info:ofi%2Ffmt:kev:mtx:ctx&rft.aulast=MIYASAWA&url_ver=Z39.88-2004&rft.auinit=K&rft.date=1960&rft.spage=181&rft.atitle=AN%20EMPIRICAL%20BAYES%20ESTIMATOR%20OF%20THE%20MEAN%20OF%20A%20NORMAL%20POPULATION&rft.issn=0074-8609)) turns out that we can compute this estimator without knowing $p_X$!\n",
    "\n",
    "For all $x$ and $y$, we have:\n",
    "$$\n",
    "p_{Y|X}(y|x) = \\mathcal{N}(y; x, \\sigma^2) = \\frac{1}{(2\\pi\\sigma^2)^{\\frac{d}{2}}} \\exp\\left(-\\frac{\\|y - x\\|^2}{2\\sigma^2} \\right)\n",
    "$$\n",
    "so:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\nabla_y p_{Y|X}(y|x) &= -\\frac{y - x}{\\sigma^2} p_{Y|X}(y|x) \\\\\n",
    "\\implies (x - y) p_{Y|X}(y|x) &= \\sigma^2 \\nabla_y p_{Y|X}(y|x) \\\\\n",
    "\\implies \\int (x - y) p_{Y|X}(y|x) p_X(x) dx &= \\sigma^2 \\int \\nabla_y p_{Y|X}(y|x) p_X(x) dx\n",
    "\\end{aligned}\n",
    "$$\n",
    "Note that, by Bayes' rule: $p_{Y|X}(y|x) p_X(x) = p_{X,Y}(x, y) = p_{X|Y}(x|y) p_Y(y)$.  \n",
    "Also, by definition of the marginals: $\\int p_{X,Y}(x, y) dx = p_{Y}(y)$.  \n",
    "For the left-hand side, we have:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\int (x - y) p_{Y|X}(y|x) p_X(x) dx &= \\int x p_{Y|X}(y|x) p_X(x) dx - \\int  y p_{Y|X}(y|x) p_X(x) dx \\\\\n",
    "&= \\int x p_{X,Y}(x, y) dx - \\int  y p_{X,Y}(x, y) dx \\\\\n",
    "&=  p_Y(y) \\left(\\int x p_{X|Y}(x|y) dx - y \\int  p_{X|Y}(x|y) dx\\right) \\\\\n",
    "&=  p_Y(y) \\left(\\mathbb{E}[X \\ | \\ Y = y] - y \\right) \\\\\n",
    "&=  p_Y(y) \\left(\\hat{X}(y) - y \\right)\n",
    "\\end{aligned}\n",
    "$$\n",
    "For the right-hand side, we have:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\sigma^2 \\int \\nabla_y p_{Y|X}(y|x) p_X(x) dx &= \\sigma^2 \\nabla_y \\int p_{Y|X}(y|x) p_X(x) dx \\\\\n",
    "&=  \\sigma^2 \\nabla_y \\int p_{X,Y}(x, y) dx  \\\\\n",
    "&=  \\sigma^2 \\nabla_y p_{Y}(y)\n",
    "\\end{aligned}\n",
    "$$\n",
    "Thus,\n",
    "$$\n",
    "\\begin{aligned}\n",
    "p_Y(y) \\left(\\hat{X}(y) - y \\right) &= \\sigma^2 \\nabla_y p_{Y}(y)\n",
    "\\\\\n",
    "\\implies \\hat{X}(y) &= y + \\sigma^2 \\frac{\\nabla_y p_{Y}(y)}{p_Y(y)}\n",
    "\\\\\n",
    "\\implies \\hat{X}(y) &= y + \\sigma^2 \\nabla_y \\log p_{Y}(y)\n",
    "\\end{aligned}\n",
    "$$\n",
    "Thus, the estimator $\\hat{X}(y)$ can be computed without knowledge of $p_X$, only the knowledge of the score $\\nabla_y \\log p_{Y}(y)$ is required.\n",
    "\n",
    "Now, there are two approaches to learning the score function $\\nabla_y \\log p_{Y}(y)$.\n",
    "The first is to approximate $p_{Y}(y)$ by an EBM:\n",
    "$$\n",
    "p_{Y}(y) \\approx \\frac{\\exp(-E_\\theta(y))}{Z(\\theta)} \\implies \\nabla_y \\log p_{Y}(y) = - \\nabla_y E_\\theta(y)\n",
    "$$\n",
    "The EBM can be learned using contrastive divergence described before.\n",
    "Then, the learned EBM can be used as a denoiser, by denoising $Y$ to obtain an estimate of $X$:\n",
    "$$\n",
    "\\hat{X}(y) = y - \\sigma^2 \\nabla_y E_\\theta(y)\n",
    "$$\n",
    "\n",
    "The second approach, proposed in this Discrete Walk-Jump Sampling paper, is to directly parametrize the score function by a 'denoising' neural network:\n",
    "$$\n",
    "g_\\phi(y) \\approx \\nabla_y \\log p_{Y}(y)\n",
    "$$\n",
    "Denoising $Y$ as before:\n",
    "$$\n",
    "\\hat{X}(y) = y + \\sigma^2 g_\\phi(y)\n",
    "$$\n",
    "The denoising network $g_\\phi$ can be trained from observations of $X$ and adding noise to obtain examples $Y$:\n",
    "* Sample $X_i \\sim p_X$.\n",
    "* Sample $\\epsilon_j \\sim \\mathcal{N}(0, \\sigma^2 I_d)$.\n",
    "* Compute $Y_{ij} = X_i + \\epsilon_j$.\n",
    "* Optimize $\\phi$:\n",
    "$$\n",
    "\\phi^* = \\argmin_\\phi \\sum_{i,j} \\|X_i - (Y_{ij} + g_\\phi(Y_{ij})) \\|^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the score model that predicts the score.\n",
    "class ScoreNetwork(nn.Module):\n",
    "    \"\"\"A simple score neural network.\"\"\"\n",
    "    hidden_size: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, y: chex.Array) -> chex.Array:\n",
    "        if len(y.shape) <= 1:\n",
    "            y = jnp.expand_dims(y, axis=0)\n",
    "        init_dims = y.shape[-1]\n",
    "        y = nn.Dense(self.hidden_size)(y)\n",
    "        y = jax.nn.softplus(y)\n",
    "        y = nn.Dense(self.hidden_size)(y)\n",
    "        y = jax.nn.softplus(y)\n",
    "        y = nn.Dense(init_dims)(y)\n",
    "        return y\n",
    "\n",
    "# Initialize the score model.\n",
    "score_model = ScoreNetwork(hidden_size=10)\n",
    "dummy_input = jnp.ones((1,))\n",
    "init_score_params = score_model.init(rng, dummy_input)\n",
    "score_fn = jax.jit(score_model.apply)\n",
    "\n",
    "# Overview of the model parameters.\n",
    "print(parameter_overview.get_parameter_overview(init_score_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@functools.partial(jax.jit, static_argnames=(\"score_fn\", \"num_noise_samples\"))\n",
    "def score_loss_fn(\n",
    "    params: optax.Params, score_fn: Callable[[optax.Params, chex.Array], chex.Array], x_true_samples: chex.Array, rng: chex.PRNGKey, noise_std: float, num_noise_samples: int) -> float:\n",
    "    \"\"\"Computes the denoising loss.\"\"\"\n",
    "    assert len(x_true_samples.shape) == 2\n",
    "    num_true_samples, num_dims = x_true_samples.shape\n",
    "\n",
    "    noise_rng, rng = jax.random.split(rng)\n",
    "    noise = noise_std * jax.random.normal(noise_rng, (num_noise_samples, num_dims))\n",
    "    y_samples = x_true_samples[:, None, :] + noise[None, ...]\n",
    "    assert y_samples.shape == (num_true_samples, num_noise_samples, num_dims)\n",
    "\n",
    "    predictions = score_fn(params, y_samples)\n",
    "    assert predictions.shape == (num_true_samples, num_noise_samples, num_dims)\n",
    "\n",
    "    l2_loss = jax.vmap(lambda x, ys, preds: jnp.linalg.norm(x - (ys + noise_std ** 2 * preds), axis=-1).mean())(x_true_samples, y_samples, predictions)\n",
    "    assert l2_loss.shape == (num_true_samples,)\n",
    "\n",
    "    return l2_loss.mean()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_score_model(init_params: optax.Params, rng: chex.PRNGKey, num_training_steps: int, **loss_kwargs) -> optax.Params:\n",
    "    \"\"\"Train the score model using the Adam optimizer.\"\"\"\n",
    "    @jax.jit\n",
    "    def train_step(params: optax.Params, opt_state: optax.OptState, x_true_samples: chex.Array, rng: chex.PRNGKey) -> Tuple[optax.Params, optax.OptState, float]:\n",
    "        loss, grad = jax.value_and_grad(score_loss_fn, has_aux=False)(params, score_fn, x_true_samples, rng, **loss_kwargs)\n",
    "        updates, opt_state = tx.update(grad, opt_state)\n",
    "        params = optax.apply_updates(params, updates)\n",
    "        return params, opt_state, loss\n",
    "\n",
    "    tx = optax.adam(5e-4)\n",
    "    opt_state = tx.init(init_params)\n",
    "\n",
    "    params_at_steps = {\n",
    "        0: init_params,\n",
    "    }\n",
    "    \n",
    "    params = init_params\n",
    "    for step in tqdm.trange(num_training_steps + 1):\n",
    "        step_rng, samples_rng, rng = jax.random.split(rng, num=3)\n",
    "        x_true_samples = px.sample(32, seed=samples_rng)\n",
    "        x_true_samples = x_true_samples[:, None]\n",
    "\n",
    "        params, opt_state, loss = train_step(params, opt_state, x_true_samples, step_rng)\n",
    "\n",
    "        # Log the training progress.        \n",
    "        if step % 500 == 0:\n",
    "            params_at_steps[step] = params\n",
    "        \n",
    "        if step % 2000 == 0:\n",
    "            print('Step {}: Loss = {}'.format(step, loss))\n",
    "    \n",
    "    return params, params_at_steps\n",
    "\n",
    "train_rng = jax.random.PRNGKey(0)\n",
    "noise_std = 0.5\n",
    "score_params, score_params_at_steps = train_score_model(init_score_params, train_rng, num_training_steps=30000, noise_std=noise_std, num_noise_samples=10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check that the score function works!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_rng, rng = jax.random.split(rng)\n",
    "noise = noise_std * jax.random.normal(noise_rng, (10, 1))\n",
    "x_true = jnp.asarray([[-1.], [1.]])\n",
    "y = x_true[None, :] + noise[:, None]\n",
    "y = y.transpose((1, 0, 2)).reshape((-1, 1))\n",
    "preds = score_fn(score_params, y)\n",
    "x = y + noise_std ** 2 * preds\n",
    "labels = jnp.where(x < 0.5, 0, 1)\n",
    "plt.grid()\n",
    "plt.scatter(y, x, color=['C0' if label == 0 else 'C1' for label in labels])\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('x')\n",
    "plt.title('Denoising Model Predictions')\n",
    "plt.show();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Walk-Jump Sampling\n",
    "\n",
    "The idea behind Walk-Jump Sampling is that it is easier to walk in the space of noisy observations $Y$ than in the space of clean observations $X$. The noise helps connect different modes of the distribution. Given any noisy observation $Y$, we can always go back to the clean observation $X$ by denoising.\n",
    "\n",
    "* Walk in noisy observation space with Langevin MCMC:\n",
    "$$\n",
    "    y_t = y_{t-1} + \\delta \\nabla_{y_{t-1}} \\log p_Y(y_{t-1})  + \\sqrt{2\\delta} \\epsilon_t\n",
    "$$\n",
    "* Jump to clean observation (at any time $\\tau$):\n",
    "$$\n",
    "    x_\\tau = y_\\tau + \\sigma^2 \\nabla_{y_\\tau} \\log p_Y(y_\\tau)\n",
    "$$\n",
    "\n",
    "Note that both the walk and jump steps need an estimate of the score function $\\nabla_{y} \\log p_Y(y)$.\n",
    "We have choices for how we parametrize these in each step. Here, they find that using an EBM for the walker, and a denoising network for the jumper works best:\n",
    "* EBM:\n",
    "$$\n",
    "p_Y(y) = \\frac{\\exp(-E_\\theta(y))}{Z(\\theta)}\n",
    "$$\n",
    "* Denoiser:\n",
    "$$\n",
    "\\nabla_y \\log p_{Y}(y) \\approx g_\\phi(y)\n",
    "$$\n",
    "\n",
    "Note that unlike diffusion, every single sample from walk-jump sampling is approximately from the data distribution $p_X$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@functools.partial(jax.jit, static_argnames=(\"energy_fn\", \"score_fn\", \"num_steps\", \"noise_std\"))\n",
    "def walk_jump_sampling(energy_fn_params: optax.Params, energy_fn: Callable[[optax.Params, chex.Array], float], score_fn_params: optax.Params, score_fn: Callable[[optax.Params, chex.Array], chex.Array], rng: chex.PRNGKey, delta: float, num_steps: int, noise_std: float):\n",
    "    \"\"\"Performs walk-jump sampling.\"\"\"\n",
    "    grad_log_prob_fn = create_grad_log_prob_fn(energy_fn_params, energy_fn)\n",
    "    noisy_observations = langevin_sample(grad_log_prob_fn, init=jnp.zeros((1,)), delta=delta, rng=rng, num_steps=num_steps)\n",
    "    scores = score_fn(score_fn_params, noisy_observations)\n",
    "    denoised_observations = noisy_observations + noise_std ** 2 * scores\n",
    "    return noisy_observations, denoised_observations\n",
    "\n",
    "\n",
    "walk_jump_sampling_rng, rng = jax.random.split(rng)\n",
    "noisy_observations, denoised_observations = walk_jump_sampling(energy_params, energy_fn, score_params, score_fn, rng=walk_jump_sampling_rng, delta=0.1, num_steps=1000, noise_std=noise_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(noisy_observations.flatten(), bins=20, density=True, alpha=0.5, color='C0', label='Noisy Observations')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.title('Histogram of Noisy Observations')\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(denoised_observations.flatten(), bins=20, density=True, alpha=0.5, color='C1', label='Denoised Observations')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.title('Histogram of Denoised Observations')\n",
    "plt.show();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}