# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# # Approximating KL using FIM
#
# > Estimate KL using FIM, then estimate FIM using samples 
#
# - toc: true 
# - badges: true
# - comments: true
# - categories: [pml2-funda-stats]
# - image: images/chart-preview.png

# Purpose of notebook: "Mahalanobis distance based on the Fisher information can be viewed as an approximation to the KL divergence between two distributions, as we verify mathematically and empirically below. This is the basis for natural gradient method."

# ## Theory
#
# Result:
#
# $$D_{KL}(p_{\theta}, p_{\theta'}) \approx \frac{1}{2} (\theta' - \theta)^T F(\theta) (\theta' - \theta)$$
#
# Derivation: 
#
# Do second-order Taylor expansion of $D_{KL}(p_{\theta}, q_{\theta'})$ around $\theta' = \theta$. The key is to think of $\theta$ as a constant, and $\theta'$ as the variable.

# $$
# D_{KL}(p_{\theta}, p_{\theta'}) 
# \approx
# \underbrace{D_{KL}(p_{\theta}, p_{\theta})}_{0}
# + (\theta' - \theta)^T \underbrace{\nabla_{\theta'} D_{KL}(p_{\theta}, p_{\theta'}) \mid_{\theta'=\theta}}_{\text{gradient}}
# + \frac{1}{2} (\theta' - \theta)^T \underbrace{\nabla^2_{\theta'} D_{KL}(p_{\theta}, p_{\theta'}) \mid_{\theta'=\theta}}_{\text{hessian}} (\theta' - \theta)
# $$

# Now, let's keep expanding the gradient and the hessian.

# The gradient turns out to be the expected value of the score function, which is zero.
#

# $$
# \begin{aligned}
# \nabla_{\theta'} D_{KL}(p_{\theta}, p_{\theta'}) \mid_{\theta'=\theta}
# & = \nabla_{\theta'} \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [\log p_\theta (x) - log p_{\theta'} (x)] \mid_{\theta'=\theta} \\
# & = \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [0 - \nabla_{\theta'} log p_{\theta'} (x) \mid_{\theta'=\theta} ] \\
# & = - \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [\nabla_{\theta'} log p_{\theta'} (x) \mid_{\theta'=\theta} ] \\
# & = 0
# \end{aligned}
# $$

# The Hessian turns out to be the negative hessian of NLL and hence the Fisher information (which can be approximated using Monte-Carlo): 

# $$
# \begin{aligned}
# \nabla_{\theta'}^2 D_{KL}(p_{\theta}, p_{\theta'}) \mid_{\theta'=\theta}
# & = \nabla_{\theta'}^2 \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [\log p_\theta (x) - log p_{\theta'} (x)] \mid_{\theta'=\theta} \\
# & = \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [0 - \nabla_{\theta'}^2 log p_{\theta'} (x) \mid_{\theta'=\theta} ] \\
# & = - \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} [\nabla_{\theta'}^2 log p_{\theta'} (x) \mid_{\theta'=\theta} ] && \text{negative Hessian of NLL}\\
# & = \mathbb{E}_{\mathbf{x} \sim p_{\mathbf{\theta}}} \left[(\nabla_{\theta'} log p_{\theta'} (x) \mid_{\theta'=\theta}) (\nabla_{\theta'} log p_{\theta'} (x) \mid_{\theta'=\theta})^T \right] && \text{Fisher information}\\
# & = F(\theta)
# \end{aligned}
# $$

# Putting everything together, we have the result stated earlier.

# ## Code

# We would first plot the the true KL divergence function, then approximate it using the result mentioned earlier



# Imports:

# +
import numpy as np
import torch
import torch.nn as nn
import torch.distributions as d
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from functorch import make_functional_with_buffers, vmap, grad
# -

# Create the distribution $p$ that we will be comparing against:

p = d.Independent(
    d.Normal(loc=torch.tensor([0., 0.]),
             scale=torch.tensor([1., 2.])),
    reinterpreted_batch_ndims=1
)

# Create a batch of other distributions $q$:

# +
diff = 10

xs = np.linspace(-diff, diff, 100)
ys = np.linspace(-diff, diff, 100)

xxs, yys = np.meshgrid(xs, ys)
xs, ys = xxs.flatten(), yys.flatten()

loc = torch.from_numpy(np.hstack([xs.reshape(-1, 1), ys.reshape(-1, 1)])).float()
# -

q = d.Independent(
    d.Normal(loc=loc, scale=torch.tensor([1., 2.]).reshape(1, -1)), 
    1
)

# Compute all KL-divergence scores with broadcasting!!

kls = d.kl.kl_divergence(p, q)

# Plot:

fig = plt.figure(figsize=(4,4))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(xxs, yys, kls.numpy().reshape(100, 100), cmap='viridis')
plt.show()



# ## Approxmating using FIM (approximate FIM as well using batches of grad outer product)

class Model(nn.Module):
    
    def __init__(self):
        super().__init__()
        self.loc = nn.Parameter(torch.tensor([0., 0.]))
        self.scale = torch.tensor([1., 2.])
    
    def forward(self, data):
        _inner = d.Normal(loc=self.loc, scale=self.scale)
        _inner._validate_args = False
        dist = d.Independent(_inner, 1)
        return dist.log_prob(data)


fmodel, params, buffers = make_functional_with_buffers(Model())


def compute_loss(params, buffers, sample):
    batch = sample
    loss = fmodel(params, buffers, batch).sum()
    return loss


ft_compute_grad = grad(compute_loss)

ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0))

data = p.sample(sample_shape=(10000, ))
ft_per_sample_grads = ft_compute_sample_grad(params, buffers, data)
grads = ft_per_sample_grads[0]

F_est = (grads.T @ grads) / len(data)
F_est

fig = plt.figure(figsize=(4,4))
ax = fig.add_subplot(111, projection='3d')
# ax.plot_surface(xxs, yys, kls.numpy().reshape(100, 100), cmap='coolwarm')
ax.plot_surface(xxs, yys, 0.5 * torch.diag(loc @ F_est @ loc.T).detach().numpy().reshape(100, 100), cmap='viridis')
plt.show()











StarcoderdataJupyterScriptsDedupFiltered