<jupyter_start><jupyter_text>Demonstration of the instability of likelihood-based training<jupyter_code>%matplotlib inline

import numpy as np
from scipy.stats import norm
import matplotlib
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import plot_settings as ps
from matplotlib.animation import FuncAnimation

ps.setup()<jupyter_output><empty_output><jupyter_text>"Training data"<jupyter_code>n = 1000
x = np.vstack((np.zeros(n), np.random.normal(size=n))).T<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>def model(x, alpha, std):
    # Rotation matrix
    c, s = np.cos(alpha), np.sin(alpha)
    r = np.matrix([[c, s], [-s, c]])
    
    # Encode to latent space
    u, v = np.einsum("ij,ni->jn", r, x)
    
    # Density in latent space
    log_prob = np.mean(norm(loc=0, scale=std).logpdf(u))
    
    # Decode onto manifold
    uv_reco = np.vstack((u, np.zeros_like(v))).T
    x_reco = np.einsum("ji,ni->nj", r, uv_reco)
    
    # Reconstruction error
    delta_x = np.mean(np.sum((x_reco - x)**2, axis=1)**0.5)
    
    return log_prob, delta_x<jupyter_output><empty_output><jupyter_text>Evaluate model over a grid of parameters<jupyter_code>res = 100

alpha_range = np.linspace(0., np.pi, res)
std_range = np.linspace(1.e-9, 1.5, res)
alpha_grid, std_grid = np.meshgrid(alpha_range, std_range)
alpha_grid = alpha_grid.flatten()
std_grid = std_grid.flatten()

log_prob_grid, delta_x_grid = [], []
for alpha, std in zip(alpha_grid, std_grid):
    log_prob, delta_x = model(x, alpha, std)
    log_prob_grid.append(log_prob)
    delta_x_grid.append(delta_x)
    
log_prob_grid = np.array(log_prob_grid).reshape(res, res)
delta_x_grid = np.array(delta_x_grid).reshape(res, res)<jupyter_output><empty_output><jupyter_text>Plot setup<jupyter_code>def make_pdf_data(sigma, alpha, extent, pdf_res):
    x_pdf = np.linspace(-4., 4., pdf_res) * np.cos(alpha)
    y_pdf = np.linspace(-4., 4., pdf_res) * np.sin(alpha)
    z_pdf = norm(scale=sigma).pdf(np.linspace(-4., 4., pdf_res))
    z_pdf[~np.isfinite(z_pdf)] = 0.
    z_pdf[z_pdf > 10.] = 10.

    extent_cut = (x_pdf**2 < extent**2) * (y_pdf**2 < extent**2)
    x_pdf = x_pdf[extent_cut]
    y_pdf = y_pdf[extent_cut]
    z_pdf = z_pdf[extent_cut]

    x_pdf = [x_pdf[0]] + list(x_pdf) + [x_pdf[-1]]
    y_pdf = [y_pdf[0]] + list(y_pdf) + [y_pdf[-1]]
    z_pdf = [3.e-3] + list(z_pdf) + [3.e-3]
    
    return x_pdf, y_pdf, z_pdf
alpha = 1.2
sigma = np.sin(alpha)
pdf_res = 200
extent = 2.8

unit_vec = np.array([np.cos(alpha), np.sin(alpha)])
unit_vec_3d = [np.cos(alpha), np.sin(alpha), 0.]

x_pdf, y_pdf, z_pdf = make_pdf_data(sigma, alpha, extent, pdf_res)
x_pdf0, y_pdf0, z_pdf0 = make_pdf_data(1., np.pi/2, extent, pdf_res)

x_std = sigma * np.cos(alpha)
y_std = sigma * np.sin(alpha)
z_std = norm(scale=sigma).pdf(sigma)

n_data = 30

fig, ax = ps.figure(cbar=False, make3d=True, large_margin=0.2, small_margin=0.2)

# Model pdf
vertices = [list(zip(x_pdf, y_pdf, z_pdf))]
poly = Poly3DCollection(
    vertices,
    facecolors=[ps.add_transparency(ps.COLORS[1], alpha=0.25)],
)
ax.add_collection3d(poly, zdir='z')
ax.plot(x_pdf, y_pdf, z_pdf, c=ps.COLORS[1], zorder=4.)

# Manifold
ax.plot(x_pdf, y_pdf, np.zeros_like(z_pdf), c=ps.COLORS[2])

# Label sigma
ax.plot([-x_std, x_std], [-y_std, y_std], [z_std, z_std], c=ps.COLORS[1])
ax.text(0.,0.,1.1*z_std,r"$\sigma$", unit_vec_3d, c=ps.COLORS[1], ha="center", va="center")

# Label alpha
ax.plot([0, 1.2], [0, 0], [0,0], c=ps.COLORS[2], lw=1.)
angles=np.linspace(0, alpha, 100)
ax.plot(np.cos(angles), np.sin(angles), np.zeros_like(angles), c=ps.COLORS[2], lw=1.)
ax.text(0.6, 0.4, 0, r"$\alpha$", "x", c=ps.COLORS[2], fontsize=9, ha="center", va="center")

# Data projection
for x_ in x[:n_data]:
    if x_[0]**2 < extent**2 and x_[1]**2 < extent**2:
        ax.scatter([x_[0]], [x_[1]], [0.], alpha=1., s=10., c=[ps.COLOR_NEUTRAL1])
        proj = [unit_vec.dot(x_) * np.cos(alpha), unit_vec.dot(x_) * np.sin(alpha)]
        ax.plot([x_[0], proj[0]], [x_[1], proj[1]], np.zeros(2), c=ps.add_transparency(ps.COLOR_NEUTRAL1, 0.25))
        #ax.scatter([proj[0]], [proj[1]], [0.], alpha=0.25, s=8., c=[ps.COLORS[1]])

# Axis ranges and view
ax.set_xlim3d(-extent, extent)
ax.set_ylim3d(-extent, extent)
ax.set_zlim3d(0., 0.5)
ax.view_init(30, -60)

# Labels
ax.set_xlabel("$x_0$", labelpad=-10.)
ax.set_ylabel("$x_1$", labelpad=-10.)
ax.set_zlabel(r"$p(x)$", labelpad=-10., c=ps.COLORS[1])
ax.set_xticklabels([""]*5)
ax.set_yticklabels([""]*5)
ax.set_zticklabels([""]*5)

plt.savefig("../figures/instability_demo_setup.pdf")<jupyter_output><empty_output><jupyter_text>Animate setup<jupyter_code>def update_alpha(alpha_, remove_labels=False):
    # Calculus
    sigma_ = max(0.001, np.abs(np.sin(alpha_)))
    unit_vec_ = np.array([np.cos(alpha_), np.sin(alpha_)])
    x_pdf_, y_pdf_, z_pdf_ = make_pdf_data(sigma_, alpha_, extent, pdf_res)
    vertices_ = [list(zip(x_pdf_, y_pdf_, z_pdf_))]
    
    # Remove labels at some point
    if remove_labels:
        sigma_line.remove()
        sigma_label.remove()
        alpha_line1.remove()
        alpha_line2.remove()
        alpha_label.remove()
        
    # Move other things around
    poly.set_verts(vertices_)
    pdf.set_data_3d(x_pdf_, y_pdf_, z_pdf_)
    manifold.set_data_3d(x_pdf_, y_pdf_, np.zeros_like(z_pdf_))
    for proj_line, x_ in zip(proj_lines, x[:n_data]):
        proj_ = [unit_vec_.dot(x_) * np.cos(alpha_), unit_vec_.dot(x_) * np.sin(alpha_)]
        proj_line.set_data_3d([x_[0], proj_[0]], [x_[1], proj_[1]], np.zeros(2))

def update(t):
    if t < beginning:
        update_alpha(alpha)
    elif t == beginning:
        update_alpha(alpha, remove_labels=True)
    elif beginning < t < n_frames - end:
        update_alpha(alpha + np.pi * (t - beginning) / (n_frames - beginning - end))
    else:
        update_alpha(alpha)

n_data = 30
n_frames = 200
interval = 40
beginning = 50
end = 12

fig, ax = ps.figure(cbar=False, make3d=True, large_margin=0.2, small_margin=0.2)

# Model pdf
vertices = [list(zip(x_pdf, y_pdf, z_pdf))]
poly = Poly3DCollection(
    vertices,
    facecolors=[ps.add_transparency(ps.COLORS[1], alpha=0.25)],
)
ax.add_collection3d(poly, zdir='z')
pdf = ax.plot(x_pdf, y_pdf, z_pdf, c=ps.COLORS[1], zorder=4.)[0]

# Manifold
manifold = ax.plot(x_pdf, y_pdf, np.zeros_like(z_pdf), c=ps.COLORS[2])[0]

# Label sigma
sigma_line = ax.plot([-x_std, x_std], [-y_std, y_std], [z_std, z_std], c=ps.COLORS[1])[0]
sigma_label = ax.text(0.,0.,1.1*z_std,r"$\sigma$", unit_vec_3d, c=ps.COLORS[1], ha="center", va="center")

# Label alpha
alpha_line1 = ax.plot([0, 1.2], [0, 0], [0,0], c=ps.COLORS[2], lw=1.)[0]
angles=np.linspace(0, alpha, 100)
alpha_line2 = ax.plot(np.cos(angles), np.sin(angles), np.zeros_like(angles), c=ps.COLORS[2], lw=1.)[0]
alpha_label = ax.text(0.6, 0.4, 0, r"$\alpha$", "x", c=ps.COLORS[2], fontsize=9, ha="center", va="center")

# Data projection
proj_lines = []
for x_ in x[:n_data]:
    if x_[0]**2 < extent**2 and x_[1]**2 < extent**2:
        ax.scatter([x_[0]], [x_[1]], [0.], alpha=1., s=10., c=[ps.COLOR_NEUTRAL1])
        proj = [unit_vec.dot(x_) * np.cos(alpha), unit_vec.dot(x_) * np.sin(alpha)]
        proj_lines.append(
            ax.plot([x_[0], proj[0]], [x_[1], proj[1]], np.zeros(2), c=ps.add_transparency(ps.COLOR_NEUTRAL1, 0.25))[0]
        )

# Axis ranges and view
ax.set_xlim3d(-extent, extent)
ax.set_ylim3d(-extent, extent)
ax.set_zlim3d(0., 0.5)
ax.view_init(30, -60)

# Labels
ax.set_xlabel("$x_0$", labelpad=-10.)
ax.set_ylabel("$x_1$", labelpad=-10.)
ax.set_zlabel(r"$p(x)$", labelpad=-10., c=ps.COLORS[1])
ax.set_xticklabels([""]*5)
ax.set_yticklabels([""]*5)
ax.set_zticklabels([""]*5)

anim = FuncAnimation(fig, update, frames=np.arange(0, n_frames), interval=interval)
anim.save('../figures/instability_demo_setup.gif', dpi=100, writer='imagemagick')<jupyter_output><empty_output><jupyter_text>Plot loss fns<jupyter_code>label_kwargs={"ha":"left", "va":"top", "x":0.1, "y":1.43, "c":"white", "fontsize":9., "fontweight":"500"}

fig, gs = ps.grid(2, 2)

ax = plt.subplot(gs[0])
zmin, zmax = -5., 1.
im = plt.imshow(
    np.clip(log_prob_grid, zmin, zmax),
    extent=(0., np.pi, 1.e-6, 1.5),
    origin="lower",
    cmap=ps.CMAP,
    norm=matplotlib.colors.Normalize(zmin, zmax),
    interpolation='nearest',
    aspect="auto"
)
plt.scatter([0.5*np.pi],[1.], marker="*", s=30., c=ps.COLOR_NEUTRAL1)
plt.text(s="Naive likelihood", **label_kwargs)
plt.xlim(0., np.pi)
plt.xticks([0., np.pi/2, np.pi], ["", "", ""])
plt.yticks([0.5, 1, 1.5])
plt.ylim(1.e-6, 1.5)
plt.ylabel(r"$\sigma$")

ax = plt.subplot(gs[1])
zmin, zmax = 0.,0.8
im = plt.imshow(
    np.clip(delta_x_grid, zmin, zmax),
    extent=(0., np.pi, 1.e-6, 1.5),
    origin="lower",
    cmap=ps.CMAP_R,
    norm=matplotlib.colors.Normalize(zmin, zmax),
    interpolation='nearest',
    aspect="auto"
)
plt.scatter([0.5*np.pi],[1.], marker="*", s=30., c=ps.COLOR_NEUTRAL1)
plt.text(s="Reconstruction error", **label_kwargs)
plt.xlim(0., np.pi)
plt.ylim(1.e-6, 1.5)
plt.xticks([0., np.pi/2, np.pi], ["", "", ""])
plt.yticks([0.5, 1, 1.5], ["", "", ""])

ax = plt.subplot(gs[2])
zmin, zmax = 0.2,1.2
im = plt.imshow(
    np.clip(delta_x_grid - 0.3 * log_prob_grid, zmin, zmax),
    extent=(0., np.pi, 1.e-6, 1.5),
    origin="lower",
    cmap=ps.CMAP_R,
    norm=matplotlib.colors.Normalize(zmin, zmax),
    interpolation='nearest',
    aspect="auto"
)
plt.scatter([0.5*np.pi],[1.], marker="*", s=30., c=ps.COLOR_NEUTRAL1)
plt.text(s="Combined loss", **label_kwargs)
plt.xlim(0., np.pi)
plt.xticks([0., np.pi/2, np.pi], ["$0$", r"$\pi/2$", r"$\pi$"])
plt.yticks([0.5, 1, 1.5])
plt.ylim(1.e-6, 1.5)
plt.xlabel(r"$\alpha$")
plt.ylabel(r"$\sigma$")

ax = plt.subplot(gs[3])
zmin, zmax = -0.2, 0.
im = plt.imshow(
    np.clip(log_prob_grid - np.max(log_prob_grid, axis=0)[np.newaxis, :], zmin, zmax),
    extent=(0., np.pi, 1.e-6, 1.5),
    origin="lower",
    cmap=ps.CMAP,
    norm=matplotlib.colors.Normalize(zmin, zmax),
    interpolation='nearest',
    aspect="auto"
)
plt.scatter([0.5*np.pi],[1.], marker="*", s=30., c=ps.COLOR_NEUTRAL1)
plt.text(s="Likelihood relative to \n" + r"MLE for each $\alpha$", **label_kwargs)
plt.xlim(0., np.pi)
plt.ylim(1.e-6, 1.5)
plt.xticks([0., np.pi/2, np.pi], ["$0$", r"$\pi/2$", r"$\pi$"])
plt.yticks([0.5, 1, 1.5], ["", "", ""])
plt.xlabel(r"$\alpha$")

plt.savefig("../figures/instability_demo_losses.pdf", dpi=600.)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.20706733974664254, 'fraction_numerical': 0.038956091056291074, 'mean_word_length': 4.343511450381679, 'pattern_counts': {'":': 7, '<': 38, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '180256', 'license': 'MIT', 'n_tokens_mistral': 4871, 'n_tokens_neox': 4711, 'n_words': 777, 'path': 'experiments/notebooks/instability_demo.ipynb', 'repo_name': 'selflein/manifold-flow', 'size': '10499'}