Polymers

Posted on May 21, 2026 by wdengquantum.me

Problem: Explain why for a random walk in $\mathbf R^d$, the probability distribution of the random vector sum $\mathbf x:=\mathbf x_1+…+\mathbf x_N$ of $N$ i.i.d. random vectors $\mathbf x_i$ each with identical mean $\boldsymbol{\mu}:=\langle\mathbf x_i\rangle$ and identical covariance matrix $\sigma^2:=\langle\mathbf x_i^{\otimes 2}\rangle-\boldsymbol{\mu}^{\otimes 2}$ in the large $N\gg 1$ limit is of the form:

\[p(\mathbf x)=\frac{1}{(2\pi N)^{d/2}\det\sigma}e^{-(\mathbf x-N\boldsymbol{\mu})^T\sigma^{-2}(\mathbf x-N\boldsymbol{\mu})/2N}\]

Solution: Central limit theorem.

Problem: Explain what an ideal polymer chain is.

Solution: A model of polymer chains lacking a concept of excluded volume. The word “ideal” here is strictly weaker than e.g. the “ideal” in “ideal gas” (where it really means no interactions of any kind); here the monomers are permitted to have short-range interactions (e.g. the covalent bonds holding them together, steric hindrance, etc.) but no long-range interactions (in particular, if a later segment of the chain happens to fold back onto an earlier segment, this is allowed within the ideal chain paradigm due to the lack of excluded volume). In particular, one can think of ideal polymer chains as the universality class of polymer chains whose expected end-to-end distance scales as $\sqrt{\langle r^2\rangle}\sim N^{1/2}$.

Problem: Consider the freely-jointed chain model of a polymer in $\mathbf R^3$ which consists of $N\gg 1$ monomers $\mathbf x_i$ each of fixed length $\ell:=|\mathbf x_i|$ connected to each other but are otherwise mutually non-interacting and are also isolated from any external solvent, etc. For the ideal chain, calculate the distribution of the end-to-end vector $\mathbf x:=\sum_{i=1}^N\mathbf x_i$ of the polymer chain.

Solution: In this case, one has $\boldsymbol{\mu}=\mathbf 0$ and writing $\mathbf x:=\ell(\cos\phi\sin\theta,\sin\phi\sin\theta,\cos\theta)$ for $(\cos\theta,\phi)\in [-1,1]\times[0,2\pi]$ uniformly distributed, the covariance matrix is proportional to the identity $\sigma^2=\ell^2/3$ (e.g. $\langle z^2\rangle=\frac{1}{2\times 2\pi}\int_{-1}^1d\cos\theta\int_0^{2\pi}d\phi (\ell\cos\theta)^2=\ell^2/3$). Thus:

\[p(\mathbf x)=\left(\frac{3}{2\pi N\ell^2}\right)^{3/2}e^{-3|\mathbf x|^2/2N\ell^2}\]

which implies the radial distribution function for $r:=|\mathbf x|$:

\[p(r)=\left(\frac{3}{2\pi N\ell^2}\right)^{3/2}4\pi r^2e^{-3r^2/2N\ell^2}\]

with the expected $\langle r^2\rangle=\int_0^{\infty}drp(r)r^2=N\ell^2$

(aside: for certain kinds of ring polymers that close back on themselves, the end-to-end distance is trivially $0$. In that case, an alternative metric for quantifying the polymer chain size is the radius of gyration $r_g^2:=\frac{1}{N}\sum_{i=1}^N|\mathbf r_i-\mathbf R|^2$ where each $\mathbf r_i$ refers to the position of (say, the center of mass, but in the coarse-grained limit it doesn’t matter) the $i^{\text{th}}$ monomer, and $\mathbf R:=\frac{1}{N}\sum_{i=1}^N\mathbf r_i$ is the center of mass of the polymer chain. By Lagrange’s identity, this can be rewritten as the average distance between pairs of monomers on the chain:

\[r_g^2=\frac{1}{2N^2}\sum_{i=1}^N\sum_{j=1}^N|\mathbf r_i-\mathbf r_j|^2=\frac{1}{N^2}\sum_{i<j}|\mathbf r_i-\mathbf r_j|^2\]

This form is easier to work with because (assuming $i<j$) one can express the displacement $\mathbf r_j-\mathbf r_i=\frac{\mathbf x_i+\mathbf x_j}{2}+\sum_{k=i+1}^{j-1}\mathbf x_{i+1}$ (assuming each monomer’s center of mass is at its midpoint, but this is irrelevant in the $N\gg 1$ limit). The $N\gg 1$ coarse-grained result for a freely-jointed chain is:

\[\langle r_g^2\rangle=\frac{\langle r^2\rangle}{6}=\frac{N\ell^2}{6}\]

is smaller than the average end-to-end distance by a factor of $\sqrt{6}$).

Problem: An improvement to the freely-jointed chain model can be made if one assumes that $\mathbf x_i\cdot\mathbf x_{i+1}=\ell^2\cos\theta$ for some fixed bond angle $\theta$ (strictly speaking this is supplementary to what chemists usually call the “bond angle” $\theta_{\text{chemist}}=\pi-\theta$). However, the torsion angle $\phi\in[0,2\pi]$ is assumed to be unrestricted (this model is called the freely rotating chain). In this case, show that for $N\gg 1$:

\[\langle r^2\rangle\approx N\ell^2\frac{1+\cos\theta}{1-\cos\theta}\]

(note: $\frac{1+\cos\theta}{1-\cos\theta}=\cot^2\theta/2$ is monotonically decreasing from a singularity at $\theta=0$ to $1$ at $\theta=\pi/2$ to $0$ at $\theta=\pi$). If furthermore the torsion angle $\phi$ is non-uniform $p(\phi)\neq 1/2\pi$ but instead subject to an achiral potential $V(-\phi)=V(\phi)$ so that $p(\phi)\sim e^{-\beta V(\phi)}$ (e.g. trans conformations being preferred over gauche), show that $\langle r^2\rangle$ receives an additional correction factor of the same form as the above but with $\cos\theta\mapsto\langle\cos\phi\rangle:=\int_0^{2\pi}d\phi p(\phi)\cos\phi$:

\[\langle r^2\rangle\approx N\ell^2\frac{1+\cos\theta}{1-\cos\theta}\frac{1+\langle\cos\phi\rangle}{1-\langle\cos\phi\rangle}\]

(this is called the hindered rotation model of the polymer chain).

Solution: The key lemma is to convince oneself that the average of any vector dotted with a circle of unit vectors in an arbitrary plane is zero. Once one realizes this, it is straightforward to check $\langle\mathbf x_i\cdot\mathbf x_j\rangle=\ell^2\cos^{|i-j|}\theta$, and the geometric series may be summed and terms simplified in the $N\gg 1$ limit to obtain the desired result.

Note that one can rewrite the $2$-point correlator:

\[\langle\mathbf x_i\cdot\mathbf x_j\rangle=\ell^2\cos^{|i-j|}\theta=\ell^2e^{-|i-j|/N_p}\]

where $N_p:=-1/\ln\cos\theta$ is the persistence number, and $\ell_p:=\ell N_p$ is the persistence length. For $\theta\ll 1$, $N_p\approx 2/\theta^2$ in which case one can the Kratky-Porod limit $\theta,\ell\to 0$ in such a way that $\ell_p=2\ell/\theta^2$ is fixed. This effectively replaces the freely rotating model with a worm-like chain model in which the polymer is parameterized by a smooth trajectory $\mathbf x(s)$ of its arc length $s$ with unit tangent vector $d\mathbf x/ds$ obeying:

\[\biggr\langle\frac{d\mathbf x}{ds}\cdot\frac{d\mathbf x}{ds’}\biggr\rangle=e^{-|s-s’|/\ell_p}\]

One can estimate $\ell_p$ by modelling the worm-like chain as an Euler-Bernoulli beam! In that case, the polymer Hamiltonian functional is quadratic:

\[H[\theta(s)]=\frac{1}{2}EI\int_0^Lds\left(\frac{d\theta}{ds}\right)^2\]

hence by the equipartition theorem, the variance $\langle\theta^2\rangle$ in the angular fluctuations is $\frac{1}{2}EI\frac{\langle\theta^2\rangle}{\ell}=2\times\frac{1}{2}k_BT$, so equating $\langle\theta^2\rangle=\theta^2$ from the deterministic freely rotating model, one finds the persistence length is set by the ratio of flexural rigidity to thermal energy:

\[\ell_p=\frac{EI}{k_BT}\]

and the expected end-to-end distance is now $\langle r^2\rangle=\int_0^L ds’\int_0^L ds e^{-|s-s’|/\ell_p}=2\ell_p^2(e^{-L/\ell_p}-1+L/\ell_p)$ which in the limit $\ell_p\ll L$ results in a Kuhn length $\ell_{\text{eff}}=2\ell_p$ twice the persistence length (in the opposite limit $\ell_p\gg L$, one has the obvious $\langle r^2\rangle\to L^2$).

For the hindered rotation model, the math is similar, just need to remember that $\langle\sin\phi\rangle=0$ because $V(-\phi)=V(\phi)$ is even.

Problem: Define the Kuhn length $\ell_{\text{eff}}$ of an ideal polymer chain.

Solution: The Kuhn length of an ideal polymer chain is defined by:

\[\ell_{\text{eff}}:=\frac{\langle r^2\rangle}{L}\]

with $L=N\ell$ the total chain length. This ensures that even if the polymer chain is not freely jointed, it may be approximated by an equivalent freely jointed chain whose “Kuhn monomers” are of length $\ell_{\text{eff}}>\ell$. For instance, $\ell_{\text{eff}}=\ell\cot(\theta/2)$ for the freely rotating chain.

Problem: Describe Flory’s mean-field theory.

Solution: (need to explain why this is a mean-field theory?). In a good solvent, $r\sim N^{3/5}$ (the exact exponent is not $3/5=0.6$ but actually $\nu\approx 0.588$ arising from detailed RG flow analysis). In a $\theta$-solvent, $r\sim N^{1/2}$, and in a poor solvent, $r\sim N^{1/3}$.

Posted in Blog | Leave a comment

Magnetism

Posted on May 16, 2026 by wdengquantum.me

Problem: Define the $2$ words in the phrase “ideal paramagnet“. Show that a classical ideal paramagnet of $N$ spins each with the same fixed magnetic dipole moment $\mu:=|\boldsymbol{\mu}|$ placed in a uniform external magnetic field $B:=|\mathbf B|$ will develop a uniform non-zero induced average magnetization $\langle M(B)\rangle=n\mu L(\beta\mu B)$ along the direction of $\mathbf B$ where the Langevin function $L(x):=\coth(x)-1/x$.

Solution: Ideal means the $N$ spins are non-interacting with respect to each other (hence there’s no need to specify the precise geometric configuration of the $N$ spins such as whether they’re on a lattice $\Lambda$, etc.). Paramagnet imposes conditions on the first $2$ Taylor expansion coefficients of $\langle M(B)\rangle$ about $B=0$:

Non-magnetic/absence of spontaneous magnetization $\langle M(B=0)\rangle=0$.
Positive-definite zero-field magnetic susceptibility $\chi_{\mu}/\mu_0:=\partial \langle M\rangle/\partial (B=0)>0$.

Classically, for a single classical spin, the state space is $(\theta,\phi)\in S^2$ and the $\phi$-independent Hamiltonian is $H(\theta)=-\boldsymbol{\mu}\cdot\mathbf B=-\mu B\cos\theta$, so the single-spin partition function is:

\[Z_1=\int_{-1}^1d\cos\theta\int_0^{2\pi}d\phi e^{-\beta H(\theta)}=4\pi\text{sinhc}\beta\mu B\]

Thus, $F_1=-k_BT\ln Z_1$ and $\langle\mu_x\rangle=\langle\mu_y\rangle=0$ whereas:

\[\langle\mu_z\rangle=\int_{-1}^1d\cos\theta\int_0^{2\pi}d\phi\frac{e^{\beta\mu B\cos\theta}}{Z_1}\mu\cos\theta=\frac{\partial \ln Z_1}{\partial\beta B}=-\frac{\partial F}{\partial B}=\mu L(\beta\mu B)\]

and hence the result follows from $\langle M(B)\rangle=n\langle\mu_z\rangle$ where $n:=N/V$.

Problem: Show that in the high-temperature limit where $x:=\beta\mu B\ll 1$, the Taylor expansion of the Langevin function about $x=0$ is $L(x)=x/3+O_{x\to 0}(x^3)$. Hence, derive Curie’s high-temperature ideal paramagnet law $\chi_{\mu}=C/T$ for the zero-field magnetic susceptibility and state the value of the Curie constant $C>0$.

Solution: One has:

\[L(x)=\coth x-\frac{1}{x}=\frac{\cosh x}{\sinh x}-\frac{1}{x}\approx\frac{1+x^2/2+…}{x+x^3/6}-\frac{1}{x}\]

\[=\frac{1}{x}\left(\frac{1+x^2/2}{1+x^2/6}-1\right)\approx\frac{1}{x}\left((1+x^2/2)(1-x^2/6)-1\right)=x/3+O_{x\to 0}(x^3)\]

Thus, it is straightforward to compute the Curie constant $C=\mu_0 n\mu^2/3k_B$ for a classical ideal paramagnet.

Problem: Repeat the above analysis but for a quantum ideal paramagnet in which all $N$ spins have the same fixed total angular momentum quantum number $j\in\{0,1/2,1,…\}$.

Solution: Now, for a single quantum spin with state vector in the manifold $\mathbf C^{2j+1}$, the spectrum of its Hamiltonian is given by the weak-field Zeeman splitting $E_{|j,m_j\rangle}=g_jm_j\mu_BB$ for Landé $g$-factor $g_j=1+\frac{j(j+1)+s(s+1)-\ell(\ell+1)}{2j(j+1)}$ and the canonical partition function is a Dirichlet kernel:

\[Z_1=\sum_{m_j=-j}^je^{-\beta E_{|j,m_j\rangle}}=\frac{\sinh(j+1/2)g_j\beta\mu_BB}{\sinh(g_j\beta\mu_BB/2)}\]

Repeating the same steps as above, this time one finds:

\[\langle M(B)\rangle=ng_jj\mu_B B_j(\beta g_jj\mu_BB)\]

where the Brillouin function is defined by:

\[B_j(x):=\left(1+\frac{1}{2j}\right)\coth\left(1+\frac{1}{2j}\right)x-\frac{1}{2j}\coth\frac{x}{2j}\]

In particular, as $j\to\infty$, $2j+1\to\infty$ and one recovers the classical continuous angle $\theta\in [0,\pi]$ and $\lim_{j\to\infty}B_j(x)=L(x)$.

This is consistent with the Taylor expansion $B_j(x)=\frac{j+1}{3j}x+O_{x\to 0}(x^3)$ which leads to the quantum Curie constant $C=\mu_0ng_j^2j(j+1)\mu_B^2/3k_B$. Instead of looking at $j\to\infty$, one can also take the quantum limit $j=s=1/2$ and $\ell=0$, in which case $g_j=2$ and:

\[B_{1/2}(x)=2\coth 2x-\coth x=\tanh x\]

so one recovers the familiar $2$-level system average magnetization $\langle M\rangle=n\mu_B\tanh\beta\mu_BB$ with Curie constant $C=\mu_0n\mu_B^2/k_B$.

Problem: Explain why classical statistical mechanics (when applied consistently!) predicts $\langle\mathbf M(\mathbf B)\rangle=\mathbf 0$ for all $\mathbf B$ (this is called the Bohr-van Leeuwen theorem). Since the Langevin derivation used a classical stat mech approach yet was able to predict a nontrivial $\mathbf M(\mathbf B)$, explain why it doesn’t violate the BvL theorem.

Solution: The BvL theorem follows mathematically from the fact that the canonical $N$-particle partition function:

\[Z(\beta,\mathbf B)=\frac{1}{h^{3N}N!}\int d^3\mathbf x_1…d^3\mathbf x_Nd^3\mathbf p_1…d^3\mathbf p_Ne^{-\beta H}\]

with $H=\sum_{i=1}^N\frac{|\mathbf p_i-q_i\mathbf A(\mathbf x_i)|^2}{2m_i}+V(\mathbf x_1,…,\mathbf x_N)$ can instead (via change of variables) be integrated over the kinetic momenta $m_i\mathbf v_i:=\mathbf p_i-q_i\mathbf A(\mathbf x_i)$ rather than the canonical momenta $\mathbf p_i$ without incurring a Jacobian penalty $\partial (m_1\mathbf v_1,…,m_N\mathbf v_N)/\partial(\mathbf p_1,…,\mathbf p_N)=1$, so $Z(\beta,\mathbf B)=Z(\beta,\mathbf B=\mathbf 0)$ is $\mathbf B$-independent and hence $F=-k_BT\ln Z$ is also $\mathbf B$-independent, leading to the BvL theorem $\langle\mathbf M\rangle=-V^{-1}\partial F/\partial\mathbf B=\mathbf 0$.

By assuming that one could speak of a “fixed magnetic dipole moment $\mu$” for all the atoms, Langevin was implicitly quantizing the system since in hindsight $\mu=g_jj\mu_B$ and $j$ is quantized and fixed (indeed, making the replacement $\mu\mapsto g_jj\mu_B$ in the Langevin magnetization maps in the $j\to\infty$ limit directly onto the Brillouin magnetization). As a result, it is better to regard the Langevin result for $\langle M(B)\rangle$ as a semi-classical formula rather than strictly belonging to classical physics (otherwise it would violate the BvL theorem!).

Posted in Blog | Leave a comment

Physics-Informed Neural Networks

Posted on May 5, 2026 by wdengquantum.me

Problem: Train a physics-informed neural network (PINN) on both the van der Pol oscillator and the drift-free Fokker-Planck diffusion equation.

Solution:

report

Spectral Bias of Physics-Informed Neural Networks¶

1. Introduction and Background¶

The physics-informed neural network (PINN) paradigm is a form of unsupervised machine learning first proposed in Raissi et al. (2019). At inference time, one freezes the model’s internal parameters and queries the PINN at an arbitrary point of the relevant domain (typically some spacetime event $(\mathbf{x}, t)$). The PINN outputs its prediction $\hat{\phi}(x,t)$ for the actual value $\phi(x,t)$ of the PDE solution at that input point, for some fixed initial/boundary conditions and constant parameters in the PDE.

This motivates the training stage to be based off an MSE cost function that simultaneously enforces both the initial/boundary conditions and the PDE constraint itself. At training time, one randomly samples a batch of $N_{\partial}$ points on the boundary and $N_{\text{int}}$ points in the interior of the domain. For each point, the PINN prediction $\hat{\phi}(x_i, t_i)$ is computed, and automatic differentiation (the same chain rule algorithm that powers backpropagation) is applied to obtain the PINN’s predictions for the relevant partial derivatives $\partial\phi/\partial x, \partial\phi/\partial t, \partial^2\phi/\partial x^2, \ldots$ as dictated by the PDE, which is assumed to take the kernelized form $\text{PDE}(\phi(x,t))=0$. The total cost function is then:

$$C(\theta) = \frac{\lambda_{\partial}}{2N_{\partial}}\sum_{i=1}^{N_{\partial}}(\hat{\phi}_{\theta}(x_i,t_i) – \phi(x_i,t_i))^2 + \frac{\lambda_{\text{int}}}{2N_{\text{int}}}\sum_{i=1}^{N_{\text{int}}}\text{PDE}^2(\hat{\phi}_{\theta}(x_i,t_i))$$

The parameters $\lambda_{\partial}, \lambda_{\text{int}}$ are hyperparameters meant to emphasize one component of the loss more heavily (e.g. to prevent pathological underfitting such as $\hat{\phi}=0$, one should choose $\lambda_{\partial} \gg \lambda_{\text{int}}$).

In this project, I explore and quantify the difficulties that PINNs have at modelling high-frequency spatial and temporal Fourier components of PDE solutions, following work of Wang et al. (2021) which in turn follows from pioneering early work by Rahaman et al. (2018) and Xu et al. (2019) on the spectral bias and frequency principle of deep neural networks outside the specialized context of PINNs. It is hypothesized that the smoothness of the activation function (which in the experiments below are always $\tanh$ activations) could induce a smoothness prior on the PINNs responsible for this spectral bias.

The project is split into two experiments:

Van der Pol oscillator (ODE): probing spectral bias in the time domain via relaxation oscillations.
Fokker-Planck equation (PDE): probing spectral bias in the spatial domain via increasingly narrow initial conditions.

2. Algorithms and Implementation¶

2.1 Network Architecture¶

Both experiments use the same core architecture: a fully-connected multilayer perceptron (MLP) with 5 hidden layers of 64 neurons each, $\tanh$ activations, and Xavier normal weight initialization. Xavier initialization draws initial weights from $\mathcal{N}(0, \sigma^2)$ with $\sigma^2 = 2/(n_{\text{in}} + n_{\text{out}})$, which is the optimal scaling for networks with symmetric, saturating activations like $\tanh$, as it ensures the initial inputs to each activation fall into the active linear region, avoiding vanishing or exploding gradients at the start of training.

For the van der Pol experiment, the input is 1-dimensional ($t \mapsto \hat{x}(t)$). For the Fokker-Planck experiment, the input is 2-dimensional ($(x,t) \mapsto \hat{p}(x,t)$).

2.2 Training¶

All models were trained with the Adam optimizer using a learning rate schedule consisting of a linear warmup phase followed by cosine annealing decay. Collocation points were re-sampled uniformly at random each epoch.

2.3 Hardware¶

All training was performed on a single NVIDIA GPU via PyTorch with CUDA. Each van der Pol model (15,000 epochs) took roughly 1-2 minutes; each Fokker-Planck model (15,000 epochs) took roughly 3-4 minutes.

3. Experiment 1: Van der Pol Oscillator¶

3.1 Problem Setup¶

The van der Pol equation is a second-order stiff nonlinear ODE first proposed by electrical engineer Balthasar van der Pol as a model of the voltage $x(t)$ in certain vacuum tube circuits:

$$\ddot{x} + \mu(x^2 – 1)\dot{x} + x = 0$$

For $\mu = 0$ this reduces to the simple harmonic oscillator $\ddot{x} + x = 0$. As $\mu$ increases, the system develops a stable limit cycle that becomes increasingly contorted in $(x, \dot{x})$ phase space, leading to relaxation oscillations with sharp, high-frequency transitions. The period grows asymptotically as $T \sim (3 – 2\ln 2)\mu$ for $\mu \gg 1$.

I first learned about this dynamical system in high school while reading Strogatz’s textbook on Chaos and Nonlinear Dynamics. Upon doing initial research into PINNs, I learned that they had a known deficiency with respect to high-frequency signals, and upon searching for a simple toy ODE to illustrate this, I recalled the van der Pol oscillator and how its relaxation oscillations could serve as a natural probe of the spectral bias.

In all experiments, I used the initial conditions $x(0) = 1$, $\dot{x}(0) = 0$, and varied the damping coefficient $\mu = 0, 0.1, 0.2, \ldots, 1.0, 2.0, \ldots, 10.0$. Training from scratch with Xavier initialization, only the $\mu = 0$ PINN was able to learn the sinusoidal SHM solution. To address this, I adopted a curriculum training strategy: the $\mu = 0$ checkpoint was used to initialize the $\mu = 0.1$ PINN, whose checkpoint in turn initialized the $\mu = 0.2$ PINN, and so forth.

In [1]:

import sys, os
import numpy as np
import matplotlib.pyplot as plt
import torch

# Add module directories to path
sys.path.insert(0, os.path.join(os.getcwd(), 'vanderpol_PINN'))
sys.path.insert(0, os.path.join(os.getcwd(), 'fokkerplanck_PINN'))

from vanderpol_PINN.model import VanillaPINN
from vanderpol_PINN.vanderpol_data import generate_reference_solution
from fokkerplanck_PINN.model import VanillaPINN2D
from fokkerplanck_PINN.fokkerplanck_data import analytical_fokker_planck

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {DEVICE}")

Using device: cuda

3.2 Results: Time Series and Phase Portraits¶

The figures below show the curriculum training results for $\mu = 0$ through $\mu = 6$. The left column compares the PINN prediction (blue dashed) against the reference Radau solution (black solid), alongside the training loss curve. The right column shows the corresponding phase portrait $(x, \dot{x})$ with the nullclines of the dynamical system and a background vector field.

Even with the curriculum training approach, by about $\mu \approx 5$ the PINN prediction starts deviating from the Radau solution. This is evidence for spectral bias in the time domain: the PINN cannot represent the sharp, high-frequency transitions of the relaxation oscillations.

In [2]:

def evaluate_pinn_with_derivatives(model, t_eval):
    # Evaluate a trained PINN, returning x(t) and x_dot(t) via autodiff.
    model.eval()
    t_tensor = torch.tensor(t_eval, dtype=torch.float32, device=DEVICE).unsqueeze(1)
    t_tensor.requires_grad_(True)
    x_pred = model(t_tensor)
    x_dot_pred = torch.autograd.grad(
        x_pred, t_tensor, grad_outputs=torch.ones_like(x_pred), create_graph=False
    )[0]
    return (x_pred.detach().cpu().numpy().flatten(),
            x_dot_pred.detach().cpu().numpy().flatten())


def plot_vdp_comparison(mu_values, ckpt_dir, T_DOMAIN=17.0):
    # Load checkpoints and plot time series + phase portraits side by side.
    for mu in mu_values:
        mu_str = f"{mu}".replace('.', 'p')
        ckpt_path = os.path.join(ckpt_dir, f'mu_{mu}_vanilla_mu_{mu_str}.pt')
        if not os.path.exists(ckpt_path):
            print(f"Checkpoint not found: {ckpt_path}, skipping.")
            continue

        checkpoint = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
        model = VanillaPINN()
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(DEVICE)

        # Reference solution
        t_ref, x_ref, xdot_ref = generate_reference_solution(mu, t_span=(0, T_DOMAIN), num_points=2000)

        # PINN prediction
        x_pinn, xdot_pinn = evaluate_pinn_with_derivatives(model, t_ref)

        # Training loss history
        loss_hist = checkpoint.get('loss_history', None)

        fig, axes = plt.subplots(1, 2, figsize=(14, 4.5))

        # --- Time series ---
        ax = axes[0]
        ax.plot(t_ref, x_ref, 'k-', lw=2, label='Reference (Radau)')
        ax.plot(t_ref, x_pinn, 'b--', lw=1.5, label='Vanilla PINN', alpha=0.8)
        ax.set_xlabel('$t$'); ax.set_ylabel('$x(t)$')
        ax.set_title(f'Time Series, $\\mu = {mu}$')
        ax.legend(fontsize=9); ax.grid(True, alpha=0.3)

        # Inset for loss curve if available
        if loss_hist is not None:
            ax_inset = ax.inset_axes([0.55, 0.55, 0.4, 0.4])
            losses = loss_hist if isinstance(loss_hist, list) else loss_hist.get('total', loss_hist)
            ax_inset.semilogy(losses, 'r-', lw=0.5)
            ax_inset.set_xlabel('Epoch', fontsize=7); ax_inset.set_ylabel('Loss', fontsize=7)
            ax_inset.tick_params(labelsize=6)

        # --- Phase portrait ---
        ax = axes[1]
        # Nullcline: x_dot = -x / (mu*(x^2 - 1))
        if mu > 0:
            x_nc = np.linspace(-3, 3, 500)
            nc_color = '#2ca02c'
            for mask in [x_nc < -1, (x_nc > -1) & (x_nc < 1), x_nc > 1]:
                seg = x_nc[mask]
                if len(seg) > 0:
                    ax.plot(seg, -seg / (mu * (seg**2 - 1)), color=nc_color,
                            ls='--', lw=1.2, alpha=0.6)

        ax.plot(x_ref, xdot_ref, 'k-', lw=2, label='Reference (Radau)')
        ax.plot(x_pinn, xdot_pinn, 'b--', lw=1.5, label='Vanilla PINN', alpha=0.8)
        ax.plot(x_ref[0], xdot_ref[0], 'ro', ms=7, label='IC')
        ax.set_xlabel('$x$'); ax.set_ylabel('$\\dot{x}$')
        ax.set_title(f'Phase Portrait, $\\mu = {mu}$')
        ax.legend(fontsize=8); ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()

In [3]:

# Plot van der Pol results for the curriculum training experiment
mu_curriculum = [0.0, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
plot_vdp_comparison(mu_curriculum, 'vanderpol_PINN/checkpoints')

c:\Users\weidu\AppData\Local\Programs\Python\Python313\Lib\site-packages\torch\autograd\graph.py:823: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\pytorch\aten\src\ATen\cuda\CublasHandlePool.cpp:180.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass

No description has been provided for this image

4. Experiment 2: Fokker-Planck Equation¶

4.1 Problem Setup¶

The second experiment steps up in complexity from an ODE to a PDE. We consider the drift-free Fokker-Planck equation (i.e. the diffusion/heat equation) for a probability distribution $p(x,t)$ in one spatial dimension:

$$\frac{\partial p}{\partial t} = D \frac{\partial^2 p}{\partial x^2}$$

for diffusion constant $D$. When the initial condition is an infinite spike at the origin, $p(x,0) = \delta(x)$, the analytical solution (the heat kernel/Green’s function) is a Gaussian whose width grows with variance $2Dt$:

$$p(x,t) = \frac{1}{\sqrt{4\pi D t}} \exp\!\left(-\frac{x^2}{4Dt}\right)$$

Because a $\delta$-function is a high-frequency spike (its Fourier transform $\int_{-\infty}^{\infty}dx\, e^{-ikx}\delta(x) = 1$ contains all spatial frequencies uniformly), PINNs should have difficulty modelling it. To quantify when this failure occurs, we set $D = 1$ and initialize the PINN with $p(x,0) = \frac{1}{\sqrt{4\pi t_0}}\exp(-x^2/4t_0)$ for some $t_0 > 0$. As $t_0 \to 0$, the initial condition approaches $\delta(x)$ and should become harder for the PINN to represent.

We truncate the spatial domain to $x \in [-L, L]$ with $L = 10$ and impose Dirichlet boundary conditions $p(\pm L, t) = 0$, which is a reasonable approximation since the Gaussian tails decay exponentially. The PINN architecture is the same MLP as before but with a 2-dimensional input $(x, t)$: five hidden layers of 64 neurons each with $\tanh$ activations and Xavier normal initialization.

4.2 Cost Function¶

The cost function has three terms:

$$C(\theta) = \frac{1}{N_{\text{PDE}}}\sum_{i=1}^{N_{\text{PDE}}}\!\left(\frac{\partial\hat{p}_\theta}{\partial t} – D\frac{\partial^2\hat{p}_\theta}{\partial x^2}\right)^{\!2}_{\!(x_i,t_i)} + \frac{\lambda_{\text{IC}}}{N_{\text{IC}}}\sum_{j=1}^{N_{\text{IC}}}\!\left(\hat{p}_\theta(x_j,0) – p(x_j,t_0)\right)^{\!2} + \frac{\lambda_{\text{BC}}}{N_{\text{BC}}}\sum_{k=1}^{N_{\text{BC}}}\hat{p}_\theta(\pm L, t_k)^2$$

with $N_{\text{PDE}} = 2500$, $N_{\text{IC}} = 1000$, $N_{\text{BC}} = 1000$, $\lambda_{\text{IC}} = 10$ (set higher to prevent the trivial $\hat{p} \equiv 0$ solution from dominating), and $\lambda_{\text{BC}} = 1$. All collocation points are re-sampled uniformly at random each epoch. The model was trained for $t_0 \in \{1.0, 0.1, 0.05, 0.01, 0.001, 0.0001\}$, with 15,000 epochs per run.

4.3 Results¶

The figures below show the PINN’s predicted profile $\hat{p}(x,t)$ (blue dashed) compared against the analytical Green’s function (black solid) at six time snapshots $t = 0, 0.2, 0.4, 0.6, 0.8, 1.0$, for decreasing values of $t_0$.

As $t_0$ decreases, the initial condition sharpens into a narrower and taller spike, whose Fourier spectrum becomes progressively more broadband (in the limit $t_0 \to 0$, the spectrum is flat since $\hat{\delta}(k) = 1$). This is where the spectral bias of standard MLPs should manifest most clearly.

For $t_0 = 1.0$ and $t_0 = 0.1$, the PINN achieves near-perfect agreement with the analytical solution at all time snapshots, since the initial Gaussian is broad and smooth (dominated by low spatial frequencies that the PINN can represent without difficulty). By $t_0 = 0.01$, the initial peak is sharp enough that small discrepancies appear in the $t = 0$ panel, though the PINN still captures the subsequent diffusive broadening reasonably well. The failure becomes catastrophic at $t_0 = 0.001$ and $t_0 = 0.0001$: the PINN cannot represent the near-$\delta$-function initial condition at all, and this failure propagates forward in time, producing profiles that are asymmetric, too broad, and shifted away from the origin. The PINN’s predictions at later times ($t = 0.6, 0.8, 1.0$) remain wrong even though the analytical solution has by then diffused into a smooth, low-frequency Gaussian. This suggests that spatial spectral bias corrupts the PINN’s internal representation of the solution globally, not just at the initial time.

In [4]:

def plot_fokkerplanck_results(t0_values, ckpt_dir, t_snapshots, D=1.0, x_range=(-10, 10)):
    # Load FP checkpoints and plot PINN vs analytical at multiple time snapshots.
    for t0 in t0_values:
        t0_str = str(t0).replace('.', 'p')
        ckpt_path = os.path.join(ckpt_dir, f'fp_t0_{t0_str}.pt')
        if not os.path.exists(ckpt_path):
            print(f"Checkpoint not found: {ckpt_path}, skipping.")
            continue

        checkpoint = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
        model = VanillaPINN2D()
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(DEVICE)
        model.eval()

        x = np.linspace(x_range[0], x_range[1], 500)
        n_snaps = len(t_snapshots)
        fig, axs = plt.subplots(2, 3, figsize=(16, 9))
        axs = axs.flatten()

        for i, t in enumerate(t_snapshots):
            ax = axs[i]
            exact_p = analytical_fokker_planck(x, t + t0, D)

            xt_grid = np.stack([x, np.full_like(x, t)], axis=1)
            xt_tensor = torch.tensor(xt_grid, dtype=torch.float32, device=DEVICE)
            with torch.no_grad():
                p_pred = model(xt_tensor).cpu().numpy().flatten()

            ax.plot(x, exact_p, 'k-', lw=2, label='Analytical')
            ax.plot(x, p_pred, 'b--', lw=1.5, label='Vanilla PINN')
            ax.set_title(f'$t = {t}$', fontsize=13)
            ax.set_xlabel('$x$')
            if i % 3 == 0:
                ax.set_ylabel('$p(x, t)$')
            ax.legend(fontsize=9)
            ax.grid(True, alpha=0.3)

        fig.suptitle(f'Fokker-Planck: PINN vs Analytical ($t_0 = {t0}$, $D = {D}$)',
                     fontsize=14, fontweight='bold')
        plt.tight_layout()
        plt.show()

In [5]:

# Plot Fokker-Planck results for all t0 values
t0_values = [1.0, 0.1, 0.05, 0.01, 0.001, 0.0001]
t_snapshots = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
plot_fokkerplanck_results(t0_values, 'fokkerplanck_PINN/checkpoints', t_snapshots)

5. Discussion and Extensions¶

These two experiments probe spectral bias along complementary axes. The van der Pol experiment demonstrates bias in the temporal domain: as $\mu$ increases, the solution $x(t)$ develops sharp transitions in time that the PINN cannot capture. The Fokker-Planck experiment demonstrates bias in the spatial domain: as $t_0$ decreases, the initial condition $p(x, 0)$ develops sharp spatial structure that the PINN cannot represent. In both cases, the fundamental issue is the same: standard MLPs with smooth activation functions tend to learn low-frequency components of the target function first and struggle with high-frequency components.

Several extensions were considered but not implemented:

Parametric PINN for the van der Pol oscillator: rather than training a separate PINN for each $\mu$, one could train a single network that accepts $(t, x(0), \dot{x}(0), \mu)$ as inputs, learning the entire family of oscillators at once.
Transfer learning: following Mustajab et al. (2024), one could first train a PINN on low-frequency solutions and then transfer the weights to better learn higher-frequency solutions.
Normalization penalty for Fokker-Planck: although a normalized initial distribution $\int dx\, p(x,0) = 1$ implies normalization for all subsequent $t$ (provided the probability current vanishes sufficiently fast at the boundaries), for the smaller values of $t_0$ the PINN predictions were no longer even normalized. One way to address this would be to incorporate an explicit normalization penalty proportional to $(1 – \int dx\, p(x,t))^2$ into the cost function, estimated via Monte Carlo integration over the sampled interior points.
Fourier transform penalties: one could add terms to the cost function that penalize discrepancies in the Fourier domain, particularly at high frequencies (e.g. via a Wasserstein distance).

References¶

Raissi, M., Perdikaris, P. and Karniadakis, G.E. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378, pp.686-707.
Wang, S., Teng, Y. and Perdikaris, P. (2021). Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal on Scientific Computing, 43(5), pp.A3055-A3081.
Rahaman, N. et al. (2018). On the spectral bias of neural networks. ICML 2019.
Xu, Z.J. et al. (2019). Frequency principle: Fourier analysis sheds light on deep neural networks. Communications in Computational Physics, 28(5), pp.1746-1767.
Karniadakis, G.E. et al. (2021). Physics-informed machine learning. Nature Reviews Physics, 3(6), pp.422-440.
Mustajab, A.H. et al. (2024). Transfer learning enhanced physics informed neural network for phase-field modeling of fracture. Theoretical and Applied Fracture Mechanics, 130, p.104321.

Posted in Blog | Leave a comment

Diffusion & Flow-Matching Models

Posted on April 12, 2026 by wdengquantum.me

Problem: State and prove Tweedie’s formula.

Solution: Tweedie’s formula asserts that if $p(\mathbf x|\boldsymbol{\mu},\sigma)=\frac{1}{\det(\sqrt{2\pi}\sigma)}e^{-(\mathbf x-\boldsymbol{\mu})^T\sigma^{-2}(\mathbf x-\boldsymbol{\mu})/2}$ is normally distributed, then without needing to know anything about the prior $p(\boldsymbol{\mu}|\sigma)$ on the mean random vector $\boldsymbol{\mu}$, one has the following Bayesian point estimate for it:

\[\langle\boldsymbol{\mu}|\mathbf x,\sigma\rangle=\mathbf x+\sigma^2\frac{\partial\ln p(\mathbf x|\sigma)}{\partial\mathbf x}\]

At first glance, one might think that $\langle\boldsymbol{\mu}|\mathbf x,\sigma\rangle\approx\mathbf x$, but Tweedie’s formula provides the empirical Bayes correction $\sigma^2\frac{\partial\ln p(\mathbf x|\sigma)}{\partial\mathbf x}$ to the naive estimate. The proof amounts to a brute force computation of the score vector field:

\[\frac{\partial\ln p(\mathbf x|\sigma)}{\partial\mathbf x}=\frac{1}{p(\mathbf x|\sigma)}\frac{\partial p(\mathbf x|\sigma)}{\partial\mathbf x}=\frac{1}{p(\mathbf x|\sigma)}\int d\boldsymbol{\mu}p(\boldsymbol{\mu}|\sigma)\frac{\partial p(\mathbf x|\boldsymbol{\mu},\sigma)}{\partial\mathbf x}\]

\[=-\frac{\sigma^{-2}}{p(\mathbf x|\sigma)}\int d\boldsymbol{\mu}p(\boldsymbol{\mu}|\sigma)(\mathbf x-\boldsymbol{\mu})p(\mathbf x|\boldsymbol{\mu},\sigma)\]

\[=-\frac{\sigma^{-2}}{p(\mathbf x|\sigma)}\left(\mathbf xp(\mathbf x|\sigma)-\int d\boldsymbol{\mu}\boldsymbol{\mu}p(\mathbf x|\sigma)p(\boldsymbol{\mu}|\mathbf x,\sigma)\right)\]

\[=\sigma^{-2}(\langle\boldsymbol{\mu}|\mathbf x,\sigma\rangle-\mathbf x)\]

Problem: Show that the set of Gaussians is closed under both convolution and multiplication.

Solution: The convolution of two normalized Gaussians is also a normalized Gaussian representing the probability distribution of the sum of the two independent normal random variables (i.e. it is a kind of orthonormal bilinear transformation):

\[p(\mathbf x|\boldsymbol{\mu}_1,\sigma_1)*p(\mathbf x|\boldsymbol{\mu}_2,\sigma_2)=p(\mathbf x|\boldsymbol{\mu}_{1*2},\sigma_{1*2})\]

where $\boldsymbol{\mu}_{1*2}=\boldsymbol{\mu}_1+\boldsymbol{\mu}_2$ and $\sigma^2_{1*2}=\sigma^2_1+\sigma^2_2$ (thus, addition of independent normal random variables is isomorphic to vector addition in $(\boldsymbol{\mu},\sigma^2)$-space).

The product of two normalized Gaussians is an (unnormalized) Gaussian (but whose normalization constant is itself has a Gaussian form):

\[p(\mathbf x|\boldsymbol{\mu}_1,\sigma_1)p(\mathbf x|\boldsymbol{\mu}_2,\sigma_2)=p(\boldsymbol{\mu}_1|\boldsymbol{\mu}_2,\sigma_{1*2})p(\mathbf x|\boldsymbol{\mu}_{12},\sigma_{12})\]

where $\sigma^{-2}_{12}\boldsymbol{\mu}_{12}=\sigma^{-2}_1\boldsymbol{\mu}_1+\sigma^{-2}_2\boldsymbol{\mu}_2$ and $\sigma^{-2}_{12}=\sigma^{-2}_1+\sigma^{-2}_2$.

(aside: the properties above are of course tied by the fact that Gaussians are (roughly speaking) eigenfunctions of the Fourier transform:

\[\int d\mathbf xe^{-i\mathbf k\cdot\mathbf x}p(\mathbf x|\boldsymbol{\mu},\sigma)=e^{-i\boldsymbol{\mu}\cdot\mathbf k-\mathbf k^T\sigma^{2}\mathbf k/2}\]

which intertwines convolution and multiplication).

Problem: What is the fundamental problem that diffusion models solve?

Solution: Given some prior data distribution $p(\mathbf x_0)$ (e.g. the distribution of images or audio), and given some latent $\mathbf x_T$ which does not lie in the support of $p$ in the sense that $p(\mathbf x_T)\approx 0$, a diffusion model is any (possibly non-deterministic) algorithm for transporting $\mathbf x_T$ onto the submanifold of latents $\mathbf x_0$ for which $p(\mathbf x_0)>0$ is supported. Roughly speaking, a diffusion model is any parametric ansatz for the score vector field $\partial\ln p(\mathbf x_0)/\partial\mathbf x_0$ of the data distribution, so one can get from $\mathbf x_T\mapsto\mathbf x_0$ by taking steps along the score vector field. Diffusion models are commonly used as generative models in the sense that the latent $\mathbf x_0$ represents a novel sample from $p(\mathbf x_0)$ though other non-generative applications (e.g. reconstruction) do exist.

Problem: Explain the problem that denoising diffusion probabilistic models (DDPMs) solve, and contrast how they are used during inference vs. how they are trained.

Solution: DDPMs are generative diffusion models. At inference time, one begins by sampling some noisy latent $\mathbf x_T$ where $p(\mathbf x_T)\approx 0$ and iteratively $\mathbf x_T,…,\mathbf x_0$ denoising $\mathbf x_T$ towards some submanifold latent $\mathbf x_0$ where $p(\mathbf x_0)>0$, thereby generating the sample $\mathbf x_0$. More explicitly, the sequence $\mathbf x_T,…,\mathbf x_0$ will turn out to be a discrete-time Markov chain, and the operation of “denoising” $\mathbf x_t\mapsto\mathbf x_{t-1}$ will really just amount to sampling from the Markov chain’s transition kernel $p(\mathbf x_{t-1}|\mathbf x_t)=\int d\mathbf x_0p(\mathbf x_0|\mathbf x_t)p(\mathbf x_{t-1}|\mathbf x_0,\mathbf x_t)$. As it is, this is intractable, but focus on that second term in the integrand:

\[p(\mathbf x_{t-1}|\mathbf x_0,\mathbf x_t)\propto p(\mathbf x_t|\mathbf x_{t-1},\mathbf x_0)p(\mathbf x_{t-1}|\mathbf x_0)\]

Both of these are essentially transition kernels for the forward Markov chain $\mathbf x_0,…,\mathbf x_T$. At this point, the DDPM paper decides to use a Gaussian forward transition kernel:

\[p(\mathbf x_t|\mathbf x_{t-1})\sim e^{-|\mathbf x_t-\sqrt{1-\sigma_t^2}\mathbf x_{t-1}|^2/2\sigma_t^2}\]

where $0<\sigma^2_1<\sigma^2_2<…<\sigma^2_T\ll 1$ are just $T$ fixed hyperparameters (called a variance schedule, where the horizon $T$ is also a hyperparameter). Equivalently, using the reparameterization trick, adding Gaussian noise looks like:

\[\mathbf x_t=\sqrt{1-\sigma_t^2}\mathbf x_{t-1}+\sigma_t\boldsymbol{\varepsilon}\]

where $\boldsymbol{\varepsilon}$ is drawn from an isotropic standard normal $p(\boldsymbol{\varepsilon})\sim e^{-|\boldsymbol{\varepsilon}|^2/2}$. The $\sqrt{1-\sigma_t^2}$ factor is used to scale down the previous state $\mathbf x_{t-1}$ in order to prevent the variance (i.e. diagonal entries of the covariance matrix) of $\mathbf x_t$ from exploding (and indeed, would be exactly variance-preserving iff the data distribution covariance $\sigma^2_{\mathbf x_0}=1$); cf. Brownian motion of a particle attached to the origin by a spring.

Rather than iteratively stepping from $\mathbf x_0$ to $\mathbf x_t$ in $O(t)$ samples, one can jump in $O(1)$ time from $\mathbf x_0\mapsto\mathbf x_t$ via just a single $\boldsymbol{\varepsilon}$-sample:

\[\mathbf x_t=\sqrt{(1-\sigma_1^2)…(1-\sigma_t^2)}\mathbf x_0+\sqrt{1-(1-\sigma^2_1)…(1-\sigma^2_t)}\boldsymbol{\varepsilon}\]

Or equivalently:

\[p(\mathbf x_t|\mathbf x_0)\sim e^{-|\mathbf x_t-\sqrt{(1-\sigma_1^2)…(1-\sigma_t^2)}\mathbf x_0|^2/2(1-(1-\sigma^2_1)…(1-\sigma_t^2))}\]

Thus, the “oracle posterior” $p(\mathbf x_{t-1}|\mathbf x_0,\mathbf x_t)\sim e^{-|\mathbf x_{t-1}-\boldsymbol{\mu}_{t-1}|^2/2\tilde{\sigma}^2_{t-1}}$ is an isotropic Gaussian with:

\[\tilde{\sigma}^{-2}_{t-1}=\frac{1-\sigma_t^2}{\sigma_t^2}+\frac{1}{1-(1-\sigma_1^2)…(1-\sigma^2_{t-1})}\Rightarrow \tilde{\sigma}^2_{t-1}=\frac{1-(1-\sigma_1^2)…(1-\sigma_{t-1}^2)}{1-(1-\sigma^2_1)…(1-\sigma^2_t)}\sigma^2_t\approx\sigma^2_t\]

\[\tilde{\sigma}^{-2}_{t-1}\boldsymbol{\mu}_{t-1}=\frac{1-\sigma_t^2}{\sigma_t^2}\frac{\mathbf x_t}{\sqrt{1-\sigma_t^2}}+\frac{\sqrt{(1-\sigma_1^2)…(1-\sigma^2_{t-1})}}{1-(1-\sigma^2_1)…(1-\sigma^2_{t-1})}\mathbf x_0\]

\[=\frac{\sqrt{1-\sigma_t^2}}{\sigma_t^2}\mathbf x_t+\frac{\sqrt{(1-\sigma_1^2)…(1-\sigma^2_{t-1})}}{1-(1-\sigma^2_1)…(1-\sigma^2_{t-1})}\frac{1}{\sqrt{1-\sigma^2_t}}\left(\mathbf x_t-\sqrt{1-(1-\sigma^2_1)…(1-\sigma^2_t)}\boldsymbol{\varepsilon}\right)\]

\[=\frac{1}{\tilde{\sigma}^2_{t-1}\sqrt{1-\sigma^2_t}}\mathbf x_t-\frac{\sigma_t^2}{\tilde{\sigma}^2_{t-1}\sqrt{(1-\sigma_t^2)(1-(1-\sigma_1^2)…(1-\sigma_t^2))}}\boldsymbol{\varepsilon}\]

Hence:

\[\boldsymbol{\mu}_{t-1}=\frac{1}{\sqrt{1-\sigma^2_t}}\left(\mathbf x_t-\frac{\sigma_t^2}{\sqrt{1-(1-\sigma_1^2)…(1-\sigma_t^2)}}\boldsymbol{\varepsilon}\right)\]

Although by design the forward Markov chain has Gaussian transition kernels $p(\mathbf x_t|\mathbf x_{t-1})\sim e^{-|\mathbf x_t-\sqrt{1-\sigma_t^2}\mathbf x_{t-1}|^2/2\sigma_t^2}$, the reverse transition kernel $p(\mathbf x_{t-1}|\mathbf x_t)$ is in general non-Gaussian. However, a theorem of Anderson guarantees that it is in fact Gaussian at the level of stochastic differential equations, and so it is not a bad approximation to simply sample $\mathbf x_{t-1}$ from the Gaussian $p(\mathbf x_{t-1}|\mathbf x_t,\mathbf x_0)$ as a proxy for sampling from the true, non-Gaussian $p(\mathbf x_{t-1}|\mathbf x_t)$. Reparameterized, this means that inference looks like unannealed Langevin dynamics (ULD):

\[\mathbf x_{t-1}=\boldsymbol{\mu}_{t-1}+\tilde{\sigma}_{t-1}\tilde{\boldsymbol{\varepsilon}}\]

\[=\frac{1}{\sqrt{1-\sigma^2_t}}\left(\mathbf x_t-\frac{\sigma_t^2}{\sqrt{1-(1-\sigma_1^2)…(1-\sigma_t^2)}}\boldsymbol{\varepsilon}\right)+\sqrt{\frac{1-(1-\sigma_1^2)…(1-\sigma_{t-1}^2)}{1-(1-\sigma^2_1)…(1-\sigma^2_t)}}\sigma_t\tilde{\boldsymbol{\varepsilon}}\]

The only problem here is that the Gaussian noise $\boldsymbol{\varepsilon}$ that was injected to get from $\mathbf x_0\mapsto\mathbf x_t$ is unknown.

This is where training comes in. Training consists of $3$ distinct sampling steps and assembling the samples:

Sample $\mathbf x_0$
Sample $t\in\{1,…,T\}$
Sample $\boldsymbol{\varepsilon}$
Hence, compute $\mathbf x_t=\sqrt{(1-\sigma_1^2)…(1-\sigma_t^2)}\mathbf x_0+\sqrt{1-(1-\sigma^2_1)…(1-\sigma^2_t)}\boldsymbol{\varepsilon}$

Architecturally, the diffusion model itself is then a ResNet-like noise predictor $\hat{\boldsymbol{\varepsilon}}(\mathbf x_t,t|\boldsymbol{\theta})$ (e.g. a U-Net or a transformer, etc.) that seeks to predict the sampled Gaussian noise $\boldsymbol{\varepsilon}$ that was added to $\mathbf x_0$ to obtain $\mathbf x_t$. This is enforced by the MSE loss function $L(\hat{\boldsymbol{\varepsilon}},\boldsymbol{\varepsilon})=|\hat{\boldsymbol{\varepsilon}}-\boldsymbol{\varepsilon}|^2/2$ with corresponding cost function over the training set:

\[C_{\text{tr}}(\boldsymbol{\theta})=\frac{1}{N_{\text{tr}}}\sum_{i=1}^{N_{\text{tr}}}L(\hat{\boldsymbol{\varepsilon}}(\mathbf x_{t_i},t_i|\boldsymbol{\theta}),\boldsymbol{\varepsilon}_i)\]

The upshot is that the inference loop Markov chain dynamics are governed (for $t=T,T-1,…,2$) by:

\[\mathbf x_{t-1}=\frac{1}{\sqrt{1-\sigma^2_t}}\left(\mathbf x_t-\frac{\sigma_t^2}{\sqrt{1-(1-\sigma_1^2)…(1-\sigma_t^2)}}\hat{\boldsymbol{\varepsilon}}(\mathbf x_t,t|\boldsymbol{\theta})\right)+\sqrt{\frac{1-(1-\sigma_1^2)…(1-\sigma_{t-1}^2)}{1-(1-\sigma^2_1)…(1-\sigma^2_t)}}\sigma_t\tilde{\boldsymbol{\varepsilon}}\]

and for $t=1$, the final generated sample $\mathbf x_0=\frac{1}{\sqrt{1-\sigma^2_1}}\left(\mathbf x_1-\sigma_1\hat{\boldsymbol{\varepsilon}}(\mathbf x_1,1|\boldsymbol{\theta})\right)$ should not have any noise added to it.

Problem: In DDPMs, explain how estimating the noise $\hat{\boldsymbol{\varepsilon}}(\mathbf x_t,t|\boldsymbol{\theta})$ is synonymous with estimating the score vector field of the data distribution $\partial\ln p(\mathbf x_t)/\partial\mathbf x_t$.

Solution: Because $\mathbf x_t$ is normally distributed about $\sqrt{(1-\sigma^2_1)…(1-\sigma^2_t)}\mathbf x_0$ with covariance $1-(1-\sigma_1^2)…(1-\sigma^2_t)$, Tweedie’s formula asserts that:

\[\langle\sqrt{(1-\sigma^2_1)…(1-\sigma^2_t)}\mathbf x_0|\mathbf x_t\rangle=\mathbf x_t+(1-(1-\sigma_1^2)…(1-\sigma^2_t))\frac{\partial\ln p(\mathbf x_t)}{\partial\mathbf x_t}\]

Comparing this with $\mathbf x_t=\sqrt{(1-\sigma^2_1)…(1-\sigma^2_t)}\mathbf x_0+\sqrt{1-(1-\sigma^2_1)…(1-\sigma^2_t)}\hat{\boldsymbol{\varepsilon}}(\mathbf x_t,t|\boldsymbol{\theta})$, one concludes that:

\[\frac{\partial\ln p(\mathbf x_t)}{\partial\mathbf x_t}=-\frac{\hat{\boldsymbol{\varepsilon}}(\mathbf x_t,t|\boldsymbol{\theta})}{\sqrt{1-(1-\sigma^2_1)…(1-\sigma^2_t)}}\]

Problem: Explain how the DDPM architecture can be reformulated in terms of a stochastic differential equation, and (following work of Song et al.) deduce its distributionally equivalent probability flow ODE.

Solution: Recall the DDPM forward Markov chain transition kernel (rewritten using $i\in\{0,…,T\}$ instead of $t$ because in a moment $t\in [0,T]$ will be reserved for a continuous variable).

\[\mathbf x_i=\sqrt{1-\sigma_i^2}\mathbf x_{i-1}+\sigma_i\boldsymbol{\varepsilon}\]

Taking the continuous time limit of this discrete difference equation amounts to letting $T\to\infty$ while $\sigma^2_i\to 0$ in such a way that their product $T\sigma^2_i:=\sigma^2(t)$ is fixed (cf. the dipole limit in electrostatics which takes $q\to\infty$ and $\Delta\mathbf x\to\mathbf 0$ such that $\boldsymbol{\pi}:=q\Delta\mathbf x$ is fixed). Hence, writing $dt:=1/T$:

\[\mathbf x_i=\sqrt{1-\sigma^2(t)dt}\mathbf x_{i-1}+\sigma(t)\sqrt{dt}\boldsymbol{\varepsilon}\]

So after binomial expanding, isolating $d\mathbf x:=\mathbf x_i-\mathbf x_{i-1}$, and recognizing the Wiener process $d\mathbf w:=\sqrt{dt}\boldsymbol{\varepsilon}$:

\[d\mathbf x=-\frac{\sigma^2(t)}{2}\mathbf xdt+\sigma(t)d\mathbf w\]

This particular SDE is also an instance of an Ornstein-Uhlenbeck process, but is special in that it is also variance-preserving. The corresponding Fokker-Planck probability current density is $-\frac{\sigma^2(t)}{2}\left(p(\mathbf x,t)\mathbf x+\frac{\partial p}{\partial\mathbf x}\right):=p(\mathbf x,t)\mathbf v_{\text{eff}}(\mathbf x,t)$ which is distributionally equivalent to the ODE $d\mathbf x=\mathbf v_{\text{eff}}dt$ with effective velocity field:

\[\mathbf v_{\text{eff}}(\mathbf x,t):=-\frac{\sigma^2(t)}{2}\left(\mathbf x+\frac{\partial \ln p}{\partial\mathbf x}\right)\]

Hence, the SDE admits a corresponding probability flow ODE $d\mathbf x/dt=\mathbf v_{\text{eff}}(\mathbf x,t)$ which can be integrated during inference using standard numerical algorithms (e.g. Euler, Runge-Kutta); recall the score vector field $\partial\ln p/\partial\mathbf x$ is what the diffusion model estimates via $\hat{\boldsymbol{\epsilon}}(\mathbf x,t|\boldsymbol{\theta})$.

Problem: Hence, in light of the above discussion, explain and motivate flow matching models, and describe the simplest instance of it (i.e. rectified flow).

Solution: Roughly speaking, diffusion models are trained to stochastically map signal to noise. At inference, one is therefore obliged to denoise along these stochastic trajectories in order to map from noise back to signal. But this is a bit like shooting oneself in the foot, unnecessarily making one’s life difficult with no ostensible gain. Flow matching models are (loosely speaking) what one gets after cutting through the fluff of diffusion models with Occam’s razor, replacing SDEs with ODEs (indeed, this is w.l.o.g. thanks to the probability flow ODE construction).

Intuitively, the simplest possible map $\mathbf x_0\mapsto\mathbf x_1$ one could engineer is a simple linear interpolation for $t\in [0,1]$:

\[\mathbf x_t=t\mathbf x_1+(1-t)\mathbf x_0\]

(aside: in the flow matching literature, the convention about whether $\mathbf x_0$ represents signal vs. noise is sometimes reversed compared with the diffusion literature, but here it is being assumed that $\mathbf x_0$ represents signal as in diffusion models). The relevant velocity field to be learned is thus $d\mathbf x_t/dt=\mathbf x_1-\mathbf x_0$. A flow-matching model thus no longer seeks to estimate a score vector field, but rather a velocity vector field $\hat{\mathbf v}(\mathbf x,t|\boldsymbol{\theta})$ via an MSE loss function $L(\hat{\mathbf v},\mathbf v):=|\hat{\mathbf v}-\mathbf v|^2/2$ and training cost function:

\[C_{\text{tr}}(\boldsymbol{\theta})=\frac{1}{N_{\text{tr}}}\sum_{i=1}^{N_{\text{tr}}}L(\hat{\mathbf v}(\mathbf x_{t_i},t_i|\boldsymbol{\theta}),\mathbf x_1-\mathbf x_0)\]

A key advantage of learning such a simple rectified flow (also called optimal transport) is that inference is very easy, and can be achieved with very few steps since one just has to step along a straight line.

Posted in Blog | Leave a comment

Information Geometry

Posted on March 27, 2026 by wdengquantum.me

Problem: Let $\boldsymbol{\Theta}$ be a smooth statistical manifold, and let $D:\boldsymbol{\Theta}^2\to [0,\infty)$ be a smooth function. What does it mean for $(\boldsymbol{\Theta},D)$ to be a “divergence manifold“?

Solution: The notion of a divergence manifold relaxes the axioms of a metric space, specifically still demanding $D(\boldsymbol{\theta} || \boldsymbol{\theta}’)\geq 0$ and $D(\boldsymbol{\theta} || \boldsymbol{\theta}’)=0\Leftrightarrow\boldsymbol{\theta}’=\boldsymbol{\theta}$ for all $\boldsymbol{\theta},\boldsymbol{\theta}’\in\boldsymbol{\Theta}$ but no longer enforcing symmetry $D(\boldsymbol{\theta} || \boldsymbol{\theta}’)=D(\boldsymbol{\theta}’ || \boldsymbol{\theta})$ or the triangle inequality $D(\boldsymbol{\theta} || \boldsymbol{\theta}^{\prime\prime})\leq D(\boldsymbol{\theta} || \boldsymbol{\theta}’) + D(\boldsymbol{\theta}’ || \boldsymbol{\theta}^{\prime\prime})$.

Problem: Explain how any divergence manifold $(\boldsymbol{\Theta},D)$ enjoys the free gift of being automatically equipped with a canonical Riemannian metric tensor field $g_D(\boldsymbol{\theta}):T_{\boldsymbol{\theta}}(\boldsymbol{\Theta})^2\to\mathbf R$ induced by the divergence function $D$.

Solution: The so-called Fisher information metric $g_D(\boldsymbol{\theta})=(g_D)_{ij}(\boldsymbol{\theta})d\theta_i\otimes d\theta_j$ on the statistical manifold $\boldsymbol{\Theta}$ induced by the divergence $D$ is basically just its Hessian:

\[(g_D)_{ij}(\boldsymbol{\theta}):=\left(\frac{\partial^2 D(\boldsymbol{\theta}||\boldsymbol{\theta}’)}{\partial\theta’_i\partial\theta’_j}\right)_{\boldsymbol{\theta}’=\boldsymbol{\theta}}\]

Intuitively, one can think of the divergence $D(\boldsymbol{\theta}||\boldsymbol{\theta}’)$ of $\boldsymbol{\theta}’\in\boldsymbol{\Theta}$ from $\boldsymbol{\theta}\in\boldsymbol{\Theta}$ as holding the “ground truth” distribution $\boldsymbol{\theta}$ fixed while sniffing around for a proxy distribution $\boldsymbol{\theta}’$ to approximate $\boldsymbol{\theta}$. By the axioms of a divergence manifold, the global minimum value $D(\boldsymbol{\theta}||\boldsymbol{\theta}’)=0$ is attained at $\boldsymbol{\theta}’=\boldsymbol{\theta}$, so Taylor expanding about this global minimum ($(\partial D(\boldsymbol{\theta}||\boldsymbol{\theta}’)/\partial\boldsymbol{\theta}’)_{\boldsymbol{\theta}’=\boldsymbol{\theta}}=\mathbf 0$), one has the local quadratic form:

\[D(\boldsymbol{\theta}||\boldsymbol{\theta}+d\boldsymbol{\theta})\approx\frac{1}{2}(g_D)_{ij}(\boldsymbol{\theta})d\theta_i d\theta_j\]

Problem: Let $f:(0,\infty)\to\mathbf R$ be a nonlinear convex function with a zero at $f(1)=0$. Define the family of $f$-divergences $D_f:\Theta^2\to [0,\infty)$, prove that they do indeed satisfy the axioms of a divergence manifold, and exhibit some examples of functions $f$ and the corresponding $f$-divergence $D_f$.

Solution: One has:

\[D_f(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\int d\mathbf x p(\mathbf x|\boldsymbol{\theta}’)f\left(\frac{p(\mathbf x|\boldsymbol{\theta})}{p(\mathbf x|\boldsymbol{\theta}’)}\right)\]

By rewriting this as an expectation $D_f(\boldsymbol{\theta}||\boldsymbol{\theta}’)=\langle f\left(p(\mathbf x|\boldsymbol{\theta})/p(\mathbf x|\boldsymbol{\theta}’)\right)\rangle_{\mathbf x\sim p(\mathbf x|\boldsymbol{\theta}’)}$, applying Jensen’s inequality, using $\int d\mathbf x p(\mathbf x|\boldsymbol{\theta})=1$ and finally using $f(1)=0$, one establishes non-negativity $D_f(\boldsymbol{\theta}||\boldsymbol{\theta}’)\geq 0$. The same $f(1)=0$ condition also ensures $\boldsymbol{\theta}=\boldsymbol{\theta}’\Rightarrow D_f(\boldsymbol{\theta}||\boldsymbol{\theta}’)=0$, and the converse is proven by recalling that Jensen’s inequality becomes an equality iff $f$ is linear (forbidden by hypothesis) or its argument $p(\mathbf x|\boldsymbol{\theta})/p(\mathbf x|\boldsymbol{\theta}’)=\text{constant}$. But since the expectation of this constant in $\mathbf x\sim p(\mathbf x|\boldsymbol{\theta}’)$ was $1$, the constant itself must be $1$, Q.E.D.

$f(x):=x\ln x$ generates the asymmetric Kullback-Leibler (KL) divergence \[D_{\text{KL}}(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\int d\mathbf x p(\mathbf x|\boldsymbol{\theta})\ln\frac{p(\mathbf x|\boldsymbol{\theta})}{p(\mathbf x|\boldsymbol{\theta}’)}\]
$f(x):=-\ln x$ generates the dual KL divergence $D_{\text{KL}}(\boldsymbol{\theta}’||\boldsymbol{\theta})$ (in general, $D_{xf(1/x)}(\boldsymbol{\theta}||\boldsymbol{\theta}’)=D_{f(x)}(\boldsymbol{\theta}’||\boldsymbol{\theta})$)
$f(x):=|x-1|/2$ generates the total variation divergence \[d_{\text{TV}}(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\frac{1}{2}\int d\mathbf x|p(\mathbf x|\boldsymbol{\theta})-p(\mathbf x|\boldsymbol{\theta}’)|\] (indeed, TVD is actually symmetric and obeys the triangle inequality so is a metric in the metric space sense, yet does not admit a Fisher information metric due to the non-differentiability of the absolute value).
$f(x):=(\sqrt{x}-1)^2/2$ generates the symmetric squared Hellinger divergence \[d^2_H(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\frac{1}{2}\int d\mathbf x\left(\sqrt{p(\mathbf x|\boldsymbol{\theta})}-\sqrt{p(\mathbf x|\boldsymbol{\theta}’)}\right)^2=1-\int d\mathbf x\sqrt{p(\mathbf x|\boldsymbol{\theta})p(\mathbf x|\boldsymbol{\theta}’)}\]
$f(x):=(x-1)^2$ generates the Pearson $\chi^2$ divergence \[D_{\chi^2}(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\int d\mathbf x\frac{(p(\mathbf x|\boldsymbol{\theta})-p(\mathbf x|\boldsymbol{\theta}’))^2}{p(\mathbf x|\boldsymbol{\theta}’)}\] and $f(x)=(x-1)^2/x$ generates the dual divergence (called the Neyman $\chi^2$ divergence) in which one replaces $p(\mathbf x|\boldsymbol{\theta}’)\mapsto p(\mathbf x|\boldsymbol{\theta})$ in the denominator.
$f(x):=\frac{1}{2}\left(x\ln x-(x+1)\ln\frac{x+1}{2}\right)$ generates the symmetric Jensen-Shannon divergence \[d^2_{\text{JS}}(\boldsymbol{\theta}||\boldsymbol{\theta}’):=\frac{D_{\text{KL}}(p_{\boldsymbol{\theta}}||\frac{p_{\boldsymbol{\theta}}+p_{\boldsymbol{\theta}’}}{2})+D_{\text{KL}}(p_{\boldsymbol{\theta}’}||\frac{p_{\boldsymbol{\theta}}+p_{\boldsymbol{\theta}’}}{2})}{2}\]

Problem: What is a fundamental limitation in the family of $f$-divergences $D_f$? How do integral probability metrics such as the Wasserstein (a.k.a. earth-mover’s) distance address this shortcoming?

Solution:

\[W_{n}(p,p’)=\text{min}_{p(\mathbf x,\mathbf x’)|\int d\mathbf xp(\mathbf x,\mathbf x’)=p(\mathbf x’)\text{ and }\int d\mathbf x’p(\mathbf x,\mathbf x’)=p'(\mathbf x)}\left(\int dp(\mathbf x,\mathbf x’)|\mathbf x-\mathbf x’|^n\right)^{1/n}\]

where the joint probability is $dp(\mathbf x,\mathbf x’)=d\mathbf xd\mathbf x’p(\mathbf x,\mathbf x’)$.

(include a concrete computation in $1$ dimension).

Problem: (something about Bregman divergence, comment on how KL div is the only simultaneous f and Bregman divergence)

Solution:

Posted in Blog | Leave a comment

Autoencoders

Posted on March 26, 2026 by wdengquantum.me

VAE

$\textbf{Problem}$: Apply the PRAC-DTDT-ID workflow to a (vanilla) $\textit{autoencoder}$.

$\textbf{Solution}$:

Problem: to learn a lower-dimensional latent manifold representation of the input manifold (this is justified by the $\textit{submanifold hypothesis}$, namely that the data-generating distribution $p(\mathbf x)$ is essentially supported only for $\mathbf x\in \mathbf R^{<d}$ in some lower-dimensional submanifold $\mathbf R^{<d}$ of the input manifold $\cong\mathbf R^d$).
Representation: assume the abstract data manifold has already been coordinatized via some atlas of charts into feature vectors $\mathbf x\in\mathbf R^d$.
Architecture: an encoder neural network $\mathbf e_{\boldsymbol{\phi}}:\mathbf R^d\to\mathbf R^{<d}$ that maps input feature vectors $\mathbf x\in\mathbf R^d$ to a lower-dimensional latent manifold $\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\in\mathbf R^{<d}$. This is then passed through a decoder neural network $\mathbf d_{\boldsymbol{\theta}}:\mathbf R^{<d}\to\mathbf R^d$ that attempts to reconstruct the original input feature vector $\mathbf d_{\boldsymbol{\theta}}(\mathbf e_{\boldsymbol{\phi}}(\mathbf x))\approx\mathbf x\in\mathbf R^d$ from its encoded latent manifold representation $\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\in\mathbf R^{<d}$. In this context, the latent manifold $\mathbf R^{<d}$ is also called the $\textit{bottleneck}$ as it’s the narrowest pinch point (a hidden layer of $<d$ neurons) in the autoencoder. If it were not $<d$ but rather $\geq d$, then the autoencoder could just memorize the input data without learning the relevant submanifold structure, so by forcing the information to flow through a bottleneck (a deliberately lossy form of compression), the autoencoder has to encode efficiently to survive the bottleneck.
Cost: since the overall goal of the autoencoder (by virtue of its etymology!) is to essentially implement an identity function $\mathbf x\mapsto\mathbf x$ (“reconstruction of $\mathbf x$”), the reconstruction loss function may taken to be the usual MSE: $$L(\hat{\mathbf x},\mathbf x)=\frac{|\hat{\mathbf x}-\mathbf x|^2}{2}$$ so the corresponding cost function is just the expected reconstruction loss over the empirical approximation of $p(\mathbf x)$ due to $N$ i.i.d. samples $\mathbf x_1,…,\mathbf x_N\in\mathbf R^d$ from $p(\mathbf x)$: $$C(\boldsymbol{\theta},\boldsymbol{\phi})=\frac{1}{N}\sum_{i=1}^NL(\mathbf d_{\boldsymbol{\theta}}(\mathbf e_{\boldsymbol{\phi}}(\mathbf x_i)),\mathbf x_i)$$
Data: one of the advantages of autoencoders is that it is a data-agnostic framework which is thus completely general and can be applied to learn more compact representations of essentially any kind of data (e.g. images, text, audio, …).
Train: use the Adam optimizer.
Dev: use to fine-tune hyperparameters like the bottleneck depth $<d$.
Test: self-explanatory.
Iterate: self-explanatory.
Deploy: self-explanatory.

$\textbf{Problem}$: Having fleshed out the theory of autoencoders, implement a vanilla autoencoder in PyTorch to learn a $16$-dimensional latent manifold representation of the MNIST dataset (where each $28\times 28$ image is flattened and normalized into a feature vector $\mathbf x\in\mathbf [0,1]^{784}$).

$\textbf{Solution}$:

In [ ]:

import torch
import torch.nn as nn

class VanillaAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(in_features=784, out_features=128),
            nn.ReLU(),
            nn.Linear(in_features=128, out_features=16),
        )

        self.decoder = nn.Sequential(
            nn.Linear(in_features=16, out_features=128),
            nn.ReLU(),
            nn.Linear(in_features=128, out_features=784),
            nn.Sigmoid() # because MNIST pixel intensities will be normalized [0, 1]
        )    

    def forward(self, x):
        # x is a batch of grayscale images with shape (batch_size, 1, 28, 28)
        batch_size = x.shape[0]
        x = x.reshape(batch_size, 784) # or x.reshape(-1, 784)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        x_hat = x_hat.reshape(batch_size, 1, 28, 28)
        return x_hat

In [20]:

from torch.optim import Adam

mnist_autoencoder = VanillaAutoencoder()
loss_func = nn.BCELoss() # rather than MSE b/c MNIST images will have [0, 1] normalized pixel intensities
optimizer = Adam(mnist_autoencoder.parameters(), lr=1e-3)

In [21]:

from datasets import load_dataset
from torchvision import transforms
from torch.utils.data import DataLoader

hf_mnist = load_dataset("mnist")
normalize_tensor = transforms.ToTensor() 

def normalize_batch(batch):
    # Apply the [0, 1] normalization to every image in the current batch
    batch["image"] = [normalize_tensor(img) for img in batch["image"]]
    return batch

hf_mnist.set_transform(normalize_batch)
train_loader = DataLoader(
    dataset=hf_mnist["train"],
    batch_size=64,
    shuffle=True
)
test_loader = DataLoader(
    dataset=hf_mnist["test"],
    batch_size=256,
    shuffle=False
)

num_epochs = 10
C_test = []

for epoch in range(num_epochs):
    mnist_autoencoder.train()
    for batch in train_loader:
        x = batch["image"]

        optimizer.zero_grad()
        x_hat = mnist_autoencoder(x)
        loss = loss_func(x_hat, x)
        loss.backward()
        optimizer.step()

    mnist_autoencoder.eval()
    epoch_loss = 0
    with torch.inference_mode():
        for batch in test_loader:
            x = batch["image"]
            x_hat = mnist_autoencoder(x)
            batch_loss = loss_func(x_hat, x)
            epoch_loss += batch_loss.item()
            
    # Calculate the average loss over all test batches
    avg_loss = epoch_loss / len(test_loader)
    C_test.append(avg_loss)
    
    print(f"Epoch [{epoch+1}/{num_epochs}] | Test Loss: {avg_loss:.4f}")

Epoch [1/10] | Test Loss: 0.1403
Epoch [2/10] | Test Loss: 0.1294
Epoch [3/10] | Test Loss: 0.1249
Epoch [4/10] | Test Loss: 0.1220
Epoch [5/10] | Test Loss: 0.1197
Epoch [6/10] | Test Loss: 0.1182
Epoch [7/10] | Test Loss: 0.1170
Epoch [8/10] | Test Loss: 0.1159
Epoch [9/10] | Test Loss: 0.1152
Epoch [10/10] | Test Loss: 0.1145

$\textbf{Problem}$: Following up on the above, seeing the test cost $C_{\text{test}}(\boldsymbol{\phi},\boldsymbol{\theta})$ decreasing as both $\boldsymbol{\phi},\boldsymbol{\theta}$ are being updated by gradient descent is one metric, but another test one can do is to directly inspect the geometry of the latent manifold $[0,\infty)^{16}$. Specifically, define a linear interpolation function that takes $2$ distinct images $\mathbf x\neq\mathbf x’\in [0, 1]^{784}$ and, for $10$ discrete time steps $t=0, 0.1, 0.2,…,1$, returns the corresponding $10$ reconstructed images $\mathbf d_{\boldsymbol{\theta}}((1-t)\mathbf e_{\boldsymbol{\phi}}(\mathbf x)+t\mathbf e_{\boldsymbol{\phi}}(\mathbf x’))\in [0,1]^{784}$.

$\textbf{Solution}$:

In [22]:

def interpolate_latent_manifold(model, x, x_prime, steps=10):
    interpolated_imgs = []
    model.eval()
    with torch.no_grad():
        x = x.reshape(1, 784)
        x_prime = x_prime.reshape(1, 784)
        z = model.encoder(x)
        z_prime = model.encoder(x_prime)
        for t in torch.linspace(0, 1, steps):
            z_t = (1 - t) * z + t * z_prime
            x_hat_t = model.decoder(z_t)
            interpolated_imgs.append(x_hat_t.reshape(1, 1, 28, 28))
    return interpolated_imgs

In [24]:

import matplotlib.pyplot as plt

data_iterator = iter(test_loader)
batch = next(data_iterator)
images = batch["image"] 

x = images[0]       
x_prime = images[1]

steps = 10
interpolated_sequence = interpolate_latent_manifold(
    model=mnist_autoencoder, 
    x=x, 
    x_prime=x_prime, 
    steps=steps
)

fig, axes = plt.subplots(nrows=1, ncols=steps, figsize=(15, 3))

for i, img_tensor in enumerate(interpolated_sequence):
    img_array = img_tensor.squeeze().cpu().numpy()
    
    axes[i].imshow(img_array, cmap='gray')
    axes[i].set_title(rf"$t = {i/(steps-1):.2f}$")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

$\textbf{Problem}$: Hence, comment on the key shortcoming of the vanilla autoencoder architecture, and explain how a variational autoencoder architecture addresses this shortcoming.

$\textbf{Solution}$: The shortcoming of the vanilla autoencoder architecture is that the latent manifold representation of the data it learns looks like grouping each of the $10$ digits $0, 1, …, 9$ into topologically disconnected “islands” in the latent manifold. Thus, forcing the decoder $\mathbf d_{\boldsymbol{\theta}}$ to “swim” between $2$ such islands leads to a blurry interpolation (e.g. the $7$ is gradually turning off as the $2$ is gradually turning on, leading to something that awkwardly almost looks like a $3$).

Variational autoencoders (VAEs) address this shortcoming of vanilla autoencoders by injecting stochasticity into both the encoder and decoder, transforming them from $\textit{deterministic}$ functions into $\textit{random}$ vectors.

Instead of a deterministic encoder $\mathbf x\in\mathbf R^d\mapsto\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\in\mathbf R^{<d}$, VAEs use a $\textit{probabilistic}$ encoder composed of $2$ steps:

(Deterministic) Map each feature vector $\mathbf x\in\mathbf R^d\mapsto (\langle\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\rangle,\ln\boldsymbol{\sigma}^2_{\boldsymbol{\phi}}(\mathbf x))\in\mathbf R^{2<d}$ to some Gaussian cloud $p_{\boldsymbol{\phi}}(\mathbf e|\mathbf x)=e^{-(\mathbf e-\langle\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\rangle)^T\sigma^{-2}_{\boldsymbol{\phi}}(\mathbf x)(\mathbf e-\langle\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\rangle)/2}/\det(\sqrt{2\pi}\sigma_{\boldsymbol{\phi}}(\mathbf x))$ in the latent manifold (diagonal covariance matrix $\sigma^2_{\boldsymbol{\phi}}(\mathbf x)=\text{diag}(\boldsymbol{\sigma}^2_{\boldsymbol{\phi}}(\mathbf x))\in\mathbf R^{<d\times<d}$)
(Random) Sample a concrete latent vector $\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\in\mathbf R^{<d}$ from $p_{\boldsymbol{\phi}}(\mathbf e|\mathbf x)$. In order to ensure the gradient $\partial\mathbf e_{\boldsymbol{\phi}}(\mathbf x)/\partial\boldsymbol{\phi}$ exists, apply the trick of sampling a random vector $\boldsymbol{\varepsilon}\in\mathbf R^{<d}$ from a standard isotropic normal distribution $p(\boldsymbol{\varepsilon})=e^{-|\boldsymbol{\varepsilon}|^2/2}/(2\pi)^{<d/2}$ and hence $\textit{reparameterizing}$ $\mathbf e_{\boldsymbol{\phi}}(\mathbf x)=\langle\mathbf e_{\boldsymbol{\phi}}(\mathbf x)\rangle+\boldsymbol{\varepsilon}\odot\boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf x)$, treating $\boldsymbol{\varepsilon}$ as if it were a constant vector during backpropagation.

Similar to the probabilistic encoder, VAEs use a $\textit{probabilistic}$ decoder composed of $2$ analogous steps:

(Deterministic) Map each latent vector $\mathbf e\in\mathbf R^{<d}\mapsto d_{\boldsymbol{\theta}}(\mathbf e)\in\mathbf R^d$ to a Gaussian cloud back in the input manifold with fixed unit covariance $p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)=e^{-|\hat{\mathbf x}-d_{\boldsymbol{\theta}}(\mathbf e)|^2/2}/(2\pi)^{d/2}$.
(Deterministic) Although one’s natural instinct might be to sample from $p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)$ to generate an actual decoded output, it turns out empirically it’s better to just pluck out the mean $d_{\boldsymbol{\theta}}(\mathbf e)$ and call it a day (equivalently, to sample from $p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)=\delta^d(\hat{\mathbf x}-d_{\boldsymbol{\theta}}(\mathbf e))$).

Finally, there’s one other aspect, specifically the C in PRAC-DTDT-ID, on which variational autoencoders differ from vanilla autoencoders, and which also explains the name “variational”. Recall that cost functions in ML often arise from maximum-likelihood estimation (or some variant like MAP). In this case, one would like to maximize the likelihood $p_{\boldsymbol{\theta}}(\hat{\mathbf x}):=p(\hat{\mathbf x}|\boldsymbol{\theta})$ that the $\textit{decoder}$ generates an output $\hat{\mathbf x}\in\mathbf R^d$ with respect to its parameters $\boldsymbol{\theta}$. By taking the log, reformulating the decoder’s generative process $p_{\boldsymbol{\theta}}(\hat{\mathbf x})=\int d^{<d}\mathbf ep(\mathbf e)p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)$ in terms of the same standard normal latent prior $p(\mathbf e)=e^{-|\mathbf e|^2/2}/(2\pi)^{<d/2}$ that was used earlier when sampling $\boldsymbol{\varepsilon}$, replacing the difficult posterior $p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)$ with the Gaussian encoder distribution $p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})$ via importance sampling, and applying Jensen’s inequality to the concave $\ln$ function to obtain an $\textit{evidence lower bound}$ (ELBO) on the log-likelihood consisting of distinct $\textit{reconstruction}$ and $\textit{regularization}$ contributions:

$$\ln p_{\boldsymbol{\theta}}(\hat{\mathbf x})\geq\langle\ln p_{\boldsymbol{\theta}}(\hat{\mathbf x}|\mathbf e)\rangle_{\mathbf e\sim p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})}-D_{\text{KL}}(p(\mathbf e)|p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})):=-L_{\text{ELBO}}(\boldsymbol{\phi},\boldsymbol{\theta},\hat{\mathbf x})$$

where the Kullback-Leibler divergence is: $$D_{\text{KL}}(p(\mathbf e)|p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})):=\langle\ln p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})/p(\mathbf e)\rangle_{\mathbf e\sim p_{\boldsymbol{\phi}}(\mathbf e|\hat{\mathbf x})}=-\frac{n}{2}-\text{Tr}\ln\sigma_{\boldsymbol{\phi}}(\hat{\mathbf x})+\frac{|\langle\mathbf e_{\boldsymbol{\phi}}(\hat{\mathbf x})\rangle|^2+|\boldsymbol{\sigma}_{\boldsymbol{\phi}}(\hat{\mathbf x})|^2}{2}$$

Then, similar to e.g. variational principles in quantum mechanics, the VAE cost function to be $\textit{minimized}$ is the empirical expectation: $$C(\boldsymbol{\phi},\boldsymbol{\theta})=\frac{1}{N}\sum_{i=1}^NL_{\text{ELBO}}(\boldsymbol{\phi},\boldsymbol{\theta},\hat{\mathbf x}_i)$$

The probabilistic decoder $d_{\boldsymbol{\theta}}$ of VAEs are often used in isolation (i.e. discarding the probabilistic encoder) as generative models by sampling from the latent manifold’s prior $p(\mathbf e)=e^{-|\mathbf e|^2/2}/(2\pi)^{<d/2}$ and simply forward passing $\mathbf e\in\mathbf R^{<d}\mapsto d_{\boldsymbol{\theta}}(\mathbf e)\in\mathbf R^d$.

$\textbf{Problem}$: Implement a variational autoencoder on the MNIST dataset with a $16$-dimensional latent manifold.

$\textbf{Solution}$:

In [ ]:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from datasets import load_dataset
from torchvision import transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

class VariationalAutoencoder(nn.Module):
    def __init__(self, latent_dim=16):
        super().__init__()
        
        self.encoder_shared = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
        )
        
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_logvar = nn.Linear(128, latent_dim)
        
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid() # constrain outputs to [0, 1]
        )

    def encode(self, x):
        h = self.encoder_shared(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        batch_size = x.shape[0]
        x_flat = x.view(batch_size, -1)
        
        mu, logvar = self.encode(x_flat)
        z = self.reparameterize(mu, logvar)
        x_hat_flat = self.decoder(z)
        
        x_hat = x_hat_flat.view(batch_size, 1, 28, 28)
        
        return x_hat, mu, logvar

def elbo_loss(x_hat, x, mu, logvar):
    BCE = F.binary_cross_entropy(x_hat, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    
    return BCE + KLD

hf_mnist = load_dataset("mnist")
to_tensor = transforms.ToTensor() 

def format_batch(batch):
    batch["image"] = [to_tensor(img) for img in batch["image"]]
    return batch

hf_mnist.set_transform(format_batch)

train_loader = DataLoader(dataset=hf_mnist["train"], batch_size=128, shuffle=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae = VariationalAutoencoder(latent_dim=16).to(device)
optimizer = Adam(vae.parameters(), lr=1e-3)

num_epochs = 15

print("Constructing the Generative Universe...")
for epoch in range(num_epochs):
    vae.train()
    epoch_loss = 0
    
    for batch in train_loader:
        x = batch["image"].to(device)
        
        optimizer.zero_grad()
        x_hat, mu, logvar = vae(x)
        
        loss = elbo_loss(x_hat, x, mu, logvar)
        loss.backward()
        epoch_loss += loss.item()
        optimizer.step()
        
    avg_loss = epoch_loss / len(train_loader.dataset)
    print(f"Epoch [{epoch+1}/{num_epochs}] | Average ELBO Loss: {avg_loss:.4f}")

print("\nScaffolding removed. Hallucinating novel digits...")
vae.eval()

num_samples = 10

with torch.inference_mode():
    z_random = torch.randn(num_samples, 16).to(device)
    generated_flat = vae.decoder(z_random)
    generated_images = generated_flat.view(num_samples, 28, 28).cpu().numpy()

fig, axes = plt.subplots(1, num_samples, figsize=(15, 2))
for i in range(num_samples):
    axes[i].imshow(generated_images[i], cmap='gray')
    axes[i].axis('off')

plt.suptitle("Hallucinated Digits Sampled from N(0, I)", fontsize=16)
plt.tight_layout()
plt.show()

Constructing the Generative Universe...
Epoch [1/15] | Average ELBO Loss: 187.7829
Epoch [2/15] | Average ELBO Loss: 137.8783
Epoch [3/15] | Average ELBO Loss: 124.7550
Epoch [4/15] | Average ELBO Loss: 118.2380
Epoch [5/15] | Average ELBO Loss: 114.7054
Epoch [6/15] | Average ELBO Loss: 112.2805
Epoch [7/15] | Average ELBO Loss: 110.6472
Epoch [8/15] | Average ELBO Loss: 109.4060
Epoch [9/15] | Average ELBO Loss: 108.4367
Epoch [10/15] | Average ELBO Loss: 107.5782
Epoch [11/15] | Average ELBO Loss: 106.9144
Epoch [12/15] | Average ELBO Loss: 106.3025
Epoch [13/15] | Average ELBO Loss: 105.8472
Epoch [14/15] | Average ELBO Loss: 105.4330
Epoch [15/15] | Average ELBO Loss: 104.9902

Scaffolding removed. Hallucinating novel digits...

Posted in Blog | Leave a comment

Graph Neural Networks

Posted on March 4, 2026 by wdengquantum.me

Problem: Give a broad sketch of the current state of the field of research in graph neural networks.

Solution:

Problem: Okay, so now explain what a graph neural network (GNN) actually is.

Solution: A GNN is basically any neural network whose input is any (undirected/directed/mixed/multi) graph (e.g. molecules, social networks, citation networks, etc.). In order to be sensible, the GNN output (whatever it is, e.g. a binary classifier for molecule toxicity) has to genuinely be an intrinsic function of the graph structure alone, and in particular not depend on any arbitrary choice of “ordering” with which one might index the graph vertices and edges (i.e. the output must be either permutation equivariant or permutation invariant depending on its nature, cf. tensors vs. tensor components in some basis).

Problem: Explain the subclass of GNNs known as message-passing neural networks (MPNNs).

Solution: An MPNN, being a certain category of GNNs, starts its life by taking as input some graph $(V,E)$. More precisely, this looks like some feature vector $\mathbf x_v$ (e.g. mass, charge, atomic number for atoms) for each vertex $v\in V$ and possibly also a feature vector $\mathbf x_e$ (e.g. bond length, bond energy, etc.) for each edge $e\in E$. The idea is that, in a manner similar to a head of self-attention, each vertex $v\in V$ wants to update its current state $\mathbf x_v$ into some new state $\mathbf x’_v$ by soaking in context from its neighbours, and in an analogous manner each edge $e\in E$ also wants to update its current state $\mathbf x_e\mapsto\mathbf x’_e$ based on its “neighbours” (thus it’s not quite the same as self-attention in which a token doesn’t just look at its nearest neighbour tokens, but at all the tokens in the context). In the general MPNN framework, this can be roughly broken down into $3$ conceptual steps:

Message phase: from a sender perspective, each vertex $v\in V$ “broadcasts” a “personalized” message vector $\mathbf m_{vv’}$ along the edge $(v,v’)\in E$ connecting it to a neighbouring vertex $v’\in V$. This message vector $\mathbf m_{vv’}$ is any (learnable) function of its current state $\mathbf x_v$, the current state of the receiving neighbour vertex $\mathbf x_{v’}$, and the current edge feature vector $\mathbf x_{vv’}$ connecting them.
Aggregation phase: simultaneously, from a receiver perspective, each vertex $v\in V$ receives the broadcasted signals from its neighbouring vertices. From this perspective, it then takes all the received message vectors and synthesizes them into a single “message summary” vector $\mathbf m_v$ which in practice is any permutation invariant function of the message vectors $\mathbf m_{vv’}$ it received from neighbouring vertices $v’\in V$ (e.g. their average).
Update phase: Finally, the vertex $v\in V$ updates its own current state $\mathbf x_v$ to some new state $\mathbf x’_v$ using another (learnable) function of its current state $\mathbf x_v$ and the message summary $\mathbf m_v$.

This $3$-step process represents a single forward pass through $1$ message-passing layer; several composed together define an MPNN.

Problem: Now that the general framework of MPNN architectures has been defined, walk through the following specific examples of MPNN architectures:

Graph convolutional networks (GCNs)
Graph attention networks (GATs)

Solution:

Posted in Blog | Leave a comment

Renormalization Group

Posted on March 1, 2026 by wdengquantum.me

Problem: Consider a Landau-Ginzburg statistical field theory involving a single real scalar field $\phi(\mathbf x)$ for $\mathbf x\in\mathbf R^d$ governed by the canonically normalized free energy density:

\[\mathcal F(\phi,\partial\phi/\partial\mathbf x,…)=\frac{1}{2}\biggr|\frac{\partial\phi}{\partial\mathbf x}\biggr|^2+\frac{\phi^2}{2\xi^2}+…\]

Explain what the $+…$ means, explain which terms have temperature $T$-dependence, and explain for which such terms does that $T$-dependence actually matter?

Solution: The $+…$ includes any terms (each with their own coupling constants) compatible with the golden trinity of constraints: locality, analyticity, and symmetry (e.g. a quartic $g\phi^4$ coupling). The part of the free energy density $\mathcal F$ before the $+…$ should be compared to the Lagrangian density $\mathcal L$ of Klein-Gordon field theory:

\[\mathcal L=\frac{1}{2c^2}\left(\frac{\partial\phi}{\partial t}\right)^2-\frac{1}{2}\biggr|\frac{\partial\phi}{\partial\mathbf x}\biggr|^2-\frac{\phi^2}{2\bar{\lambda}^2}\]

with $\bar{\lambda}=\hbar/mc$ the reduced Compton wavelength playing a role analogous to the correlation length $\xi\sim 1/\sqrt{|T-T_c|}$; indeed, this $T$-dependence in $\xi=\xi(T)$ is (usually) the only $T$-dependence that matters, even though generically all the other coupling constants will also have $T$-dependence.

Problem: Define the (non-standard) notion of “theory space”.

Solution: Roughly speaking, “theory space” is the space of all Landau-Ginzburg statistical field theories $(\mathcal F,k^*)$ (notice it is defined not only by the free energy density $\mathcal F$ but also the baggage of the UV cutoff $k^*$; it is an effective field theory). The $\mathcal F$ part can equivalently be parameterized as a countably $\infty$-tuple $(\xi,g,…)$ of the LAS-permitted coupling constants in $\mathcal F$.

Problem: In broad strokes, describe the sequence of $3$ steps that comprise a $\mathbf k$-space $\zeta$-renormalization semigroup transformation from one effective Landau-Ginzburg statistical field theory $(\mathcal F,k^*)\mapsto (\mathcal F_{\zeta},k^*)$ to another with the same UV cutoff $k^*$ but a new Wilsonian effective free energy $\mathcal F_{\zeta}$.

Solution: For $\zeta\in [1,\infty)$, the corresponding $\mathbf k$-space $\zeta$-RG transformation of $(\mathcal F,k^*)$ is given by the $3$-step recipe:

Coarse-graining $k^*\mapsto k^*/\zeta$ (blocking in real space/integrating out shells in momentum space)
Rescaling $\mathbf k’:=\zeta\mathbf k$ to recover the original UV cutoff $k^*/\zeta\mapsto k^*$ (this leads to a reciprocal “zooming out” of space $\mathbf x\mapsto\mathbf x/\zeta$).
Rescale fields $\phi’:=\zeta^{\Delta}\phi$ to make $\mathcal F_{\zeta}$ canonically normalized with respect to $\mathcal F$.

Problem: Consider an effective Landau-Ginzburg statistical field theory of a single real scalar field $\phi(\mathbf x)\in\mathbf R$ with $\mathbf x\in\mathbf R^d$ whose Fourier transform $\phi_{\mathbf k}=\int d^d\mathbf x e^{-i\mathbf k\cdot\mathbf x}\phi(\mathbf x)$ is supported only on a ball of radius $k^*$ (the theory’s UV cutoff). The free energy density corresponds to a free (no pun intended) field:

\[\mathcal F(\phi,\partial\phi/\partial\mathbf x)=\frac{1}{2}\biggr|\frac{\partial\phi}{\partial\mathbf x}\biggr|^2+\frac{\phi^2}{2\xi^2}\]

Perform a $\mathbf k$-space $\zeta$-renormalization of this theory to find the corresponding Wilsonian effective free energy density $\mathcal F_{\zeta}$.

Solution: Work with the free energy $F=\int d^d\mathbf x\mathcal F$ itself instead of just its density $\mathcal F$:

\[F[\phi]=\frac{1}{2}\int_{|\mathbf k|<k^*}\frac{d^d\mathbf k}{(2\pi)^d}\left(|\mathbf k|^2+\frac{1}{\xi^2}\right)|\phi_{\mathbf k}|^2\]

Partition the support $|\mathbf k|<k^*$ of $\phi_{\mathbf k}$ into $|\mathbf k|<k^*/\zeta$ and $k^*/\zeta<|\mathbf k|<k^*$ and based on this $\zeta$, piecewise decompose $\phi_{\mathbf k}=\phi^{<}_{\mathbf k}+\phi^{>}_{\mathbf k}$. Then one has an instance of the freshman’s dream (thanks to the disjoint supports of $\phi^{<}_{\mathbf k}$ and $\phi^{>}_{\mathbf k}$):

\[|\phi_{\mathbf k}|^2=|\phi^{<}_{\mathbf k}+\phi^{>}_{\mathbf k}|^2=|\phi^{<}_{\mathbf k}|^2+|\phi^{>}_{\mathbf k}|^2\]

So:

\[F[\phi]=\frac{1}{2}\int_{|\mathbf k|<k^*/\zeta}\frac{d^d\mathbf k}{(2\pi)^d}\left(|\mathbf k|^2+\frac{1}{\xi^2}\right)|\phi^{<}_{\mathbf k}|^2+\frac{1}{2}\int_{k^*/\zeta<|\mathbf k|<k^*}\frac{d^d\mathbf k}{(2\pi)^d}\left(|\mathbf k|^2+\frac{1}{\xi^2}\right)|\phi^{>}_{\mathbf k}|^2\]

\[=F[\phi^<_{\zeta}]+F[\phi^>_{\zeta}]\]

In this case the partition function factorizes:

\[Z=\int\mathcal D\phi e^{-\beta F[\phi]}=\int\mathcal D\phi^{<}_{\zeta} e^{-\beta F[\phi^{<}_{\zeta}]}\int\mathcal D\phi^{>}_{\zeta}e^{-\beta F[\phi^{>}_{\zeta}]}=Z^{>}_{\zeta}\int\mathcal D\phi^{<}_{\zeta} e^{-\beta F[\phi^{<}_{\zeta}]}\]

where the measures are $\mathcal D\phi^{<}_{\zeta}=\prod_{|\mathbf k|<k^*/\zeta}d\phi^{<}_{\mathbf k}$ and $\mathcal D\phi^{>}_{\zeta}=\prod_{k^*/\zeta<|\mathbf k|<k^*}d\phi^{>}_{\mathbf k}$. The constant $Z^>_{\zeta}$ doesn’t affect the physics, being absorbed as a constant shift into the Wilsonian effective free energy $F_{\zeta}[\phi^{<}_{\zeta}]=F[\phi^{<}_{\zeta}]-k_BT\ln Z^{>}_{\zeta}$.

2. Rescaling $\mathbf k’:=\zeta\mathbf k$, the Wilsonian effective free energy becomes:

\[F_{\zeta}[\phi^{<}_{\zeta}]=\frac{1}{2}\int_{|\mathbf k’|<k^*}\frac{d^d\mathbf k’}{(2\pi)^d}\zeta^{-d}\left(\zeta^{-2}|\mathbf k’|^2+\frac{1}{\xi^2}\right)|\phi^{<}_{\mathbf k’/\zeta}|^2\]

3. Rescaling $\phi^{<\prime}_{\mathbf k’}:=\zeta^{\Delta}\phi^{<}_{\mathbf k’/\zeta}$, the Wilsonian effective free energy becomes:

\[F_{\zeta}[\phi^{<}_{\zeta}]=\frac{1}{2}\int_{|\mathbf k’|<k^*}\frac{d^d\mathbf k’}{(2\pi)^d}\zeta^{-d}\left(\zeta^{-2}|\mathbf k’|^2+\frac{1}{\xi^2}\right)\zeta^{-2\Delta}|\phi^{<\prime}_{\mathbf k’}|^2\]

so in order to canonically normalize the gradient term, one requires $\Delta=-(d+2)/2$. This leads to the desired Wilsonian effective free energy density:

\[\mathcal F_{\zeta}(\phi^{<\prime}_{\zeta},\partial\phi^{<\prime}_{\zeta}/\partial\mathbf x)=\frac{1}{2}\biggr|\frac{\partial\phi^{<\prime}_{\zeta}}{\partial\mathbf x}\biggr|^2+\zeta^2\frac{(\phi^{<\prime}_{\zeta})^2}{2\xi^2}\]

Thus, by construction the gradient coupling is marginal (i.e. $\zeta$-independent) while the quadratic coupling is relevant because $\zeta^2\to\infty$ as $\zeta\to\infty$.

Posted in Blog | Leave a comment

Convolutional Neural Networks

Posted on March 1, 2026 by wdengquantum.me

CNNs_Part_1

$\textbf{Problem}$: Write functions that take an arbitrary grayscale image and convolve them with a Sobel edge detection kernel. Apply both functions to a grayscale image of your choice.

$\textbf{Solution}$:

In [1]:

import numpy as np
import matplotlib.pyplot as plt

sobel_horizontal_kernel = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
sobel_vertical_kernel = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])

def sobel_horizontal(img):
    convolved_image = np.zeros(shape=(img.shape[0]-2,img.shape[1]-2))
    for i in range(1,img.shape[0] - 1):
        for j in range(1,img.shape[1] - 1):
            convolved_image[i-1, j-1] = np.sum(img[i-1:i+2, j-1:j+2] * sobel_horizontal_kernel)
    return convolved_image

def sobel_vertical(img):
    convolved_image = np.zeros(shape=(img.shape[0]-2,img.shape[1]-2))
    for i in range(1,img.shape[0] - 1):
        for j in range(1,img.shape[1] - 1):
            convolved_image[i-1, j-1] = np.sum(img[i-1:i+2, j-1:j+2] * sobel_vertical_kernel)
    return convolved_image

np.random.seed(42)
random_grayscale_img = np.random.randint(0, 256, size=(10, 10))

print("Original Image")
plt.imshow(random_grayscale_img, cmap="gray")
plt.show()

print("Horizontal Sobel")
plt.imshow(sobel_horizontal(random_grayscale_img), cmap="gray")
plt.show()

print("Vertical Sobel")
plt.imshow(sobel_vertical(random_grayscale_img), cmap="gray")
plt.show()

Original Image

Horizontal Sobel

Vertical Sobel

$\textbf{Problem}$: Above, the $3\times 3$ Sobel edge detection kernels were convolved with an $n\times m$ grayscale image, leading to an $(n-2)\times (m-2)$ edge map. However, although for $n, m \gg 1$ this $-2$ shrinking effect is not particularly noticeable, over many layers of convolutions it can compound undesirably. Thus, one solution is to use $\textit{padding}$ to prevent this shrinking effect. Show how this works.

$\textbf{Solution}$: (by the way, the above unpadded version is sometimes called a $\textit{valid}$ convolution, whereas the padded version is sometimes called a $\textit{same}$ convolution).

In [2]:

padded_random_grayscale_img = np.pad(random_grayscale_img, ((1, 1), (1, 1)), 'constant')
print("Padded random grayscale image shape: ", padded_random_grayscale_img.shape)
plt.imshow(padded_random_grayscale_img, cmap='gray')
plt.show()

print("Horizontal Sobel")
plt.imshow(sobel_horizontal(padded_random_grayscale_img), cmap="gray")
plt.show()

print("Vertical Sobel")
plt.imshow(sobel_vertical(padded_random_grayscale_img), cmap="gray")
plt.show()

Padded random grayscale image shape:  (12, 12)

Horizontal Sobel

Vertical Sobel

$\textbf{Problem}$: Generate a random $15\times 15$ grayscale image, a random $5\times 5$ kernel, and perform a $\textit{stride}$ convolution of the image with the kernel using a stride size of $2$.

$\textbf{Solution}$: The dimension of the output image in this (unpadded) case is $(15-5)/2+1=6$.

In [3]:

np.random.seed(42)
rnd_gray_img = np.random.randint(0, 256, size=(15, 15))
rnd_kernel = np.random.randint(0, 10, size=(5, 5))
plt.imshow(rnd_gray_img, cmap='gray')
plt.show()
plt.imshow(rnd_kernel, cmap='gray')
plt.show()

In [4]:

def stride_convolution(img, kernel, stride_size, padding):
    output_img = np.zeros(shape=(int((img.shape[0] + 2 * padding - kernel.shape[0])/stride_size + 1), int((img.shape[1] + 2 * padding - kernel.shape[1])/stride_size + 1)))
    for i in range(output_img.shape[0]):
        for j in range(output_img.shape[1]):
            output_img[i, j] = np.sum(img[i * stride_size:i * stride_size + kernel.shape[0], j * stride_size:j * stride_size + kernel.shape[1]] * kernel)
    return output_img

plt.imshow(stride_convolution(rnd_gray_img, rnd_kernel, 2, 0), cmap='gray')

Out[4]:

<matplotlib.image.AxesImage at 0x164ca230b90>

$\textbf{Problem}$: Suppose one has a $28\times 28\times 192$ volume at some point inside a convolutional neural network. Explain how a same, unpadded, unit-stride convolution can be performed on this volume to obtain an output volume of dimensions $28\times 28\times 32$.

$\textbf{Solution}$: Use $32$ filters of size $1\times 1\times 192$ with stride $1$ and no padding; this is a common technique to reduce channel number while preserving spatial dimensions.

$\textbf{Problem}$: Implement simplified versions of the LeNet-5, AlexNet, and VGG-16 convolutional neural networks with randomly initialized weights and biases (here simplified means that we will not implement batch normalization, dropout, or any activation functions). Run a forward pass through it without bothering to train the parameters.

$\textbf{Solution}$:

In [18]:

class ConvLayer():
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        
        # Initialize weights and biases
        self.weights = np.random.randn(out_channels, in_channels, kernel_size, kernel_size)
        self.biases = np.random.randn(out_channels)
    
    def forward(self, x):
        # x shape: (batch_size, in_channels, height, width)
        batch_size, in_channels, height, width = x.shape
        
        # Calculate output dimensions
        out_height = (height + 2 * self.padding - self.kernel_size) // self.stride + 1
        out_width = (width + 2 * self.padding - self.kernel_size) // self.stride + 1
        
        # Initialize output
        out = np.zeros((batch_size, self.out_channels, out_height, out_width))
        
        # Apply padding if needed
        if self.padding > 0:
            x = np.pad(x, ((0, 0), (0, 0), (self.padding, self.padding), (self.padding, self.padding)), 'constant')
        
        # Convolution operation
        for b in range(batch_size):
            for c in range(self.out_channels):
                for h in range(out_height):
                    for w in range(out_width):
                        # Extract the receptive field
                        receptive_field = x[b, :, h*self.stride:h*self.stride+self.kernel_size, 
                                            w*self.stride:w*self.stride+self.kernel_size]
                        
                        # Perform convolution
                        out[b, c, h, w] = np.sum(receptive_field * self.weights[c]) + self.biases[c]
        
        return out

class PoolingLayer():
    def __init__(self, kernel_size, stride=2, mode="max"):
        self.kernel_size = kernel_size
        self.stride = stride
        self.mode = mode
    
    def forward(self, x):
        # x shape: (batch_size, channels, height, width)
        batch_size, channels, height, width = x.shape
        
        # Calculate output dimensions
        out_height = height // self.stride
        out_width = width // self.stride
        
        # Initialize output
        out = np.zeros((batch_size, channels, out_height, out_width))
        
        # Max pooling operation
        for b in range(batch_size):
            for c in range(channels):
                for h in range(out_height):
                    for w in range(out_width):
                        # Extract the receptive field
                        receptive_field = x[b, c, h*self.stride:h*self.stride+self.kernel_size, 
                                            w*self.stride:w*self.stride+self.kernel_size]
                        
                        # Perform max pooling
                        if self.mode == "max":
                            out[b, c, h, w] = np.max(receptive_field)
                        elif self.mode == "avg":
                            out[b, c, h, w] = np.mean(receptive_field)
        
        return out

class FCLayer():
    def __init__(self, in_features, out_features):
        self.in_features = in_features
        self.out_features = out_features
        
        # Initialize weights and biases
        self.weights = np.random.randn(out_features, in_features)
        self.biases = np.random.randn(out_features)
    
    def forward(self, x):
        # x shape: (batch_size, in_features)
        batch_size = x.shape[0]
        
        # Initialize output
        out = np.zeros((batch_size, self.out_features))
        
        # Perform matrix multiplication
        for b in range(batch_size):
            out[b] = np.dot(self.weights, x[b]) + self.biases
        
        return out

In [19]:

#Le-Net 5 Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 1, 32, 32))

# Create layers
layer1 = ConvLayer(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0)
layer2 = PoolingLayer(kernel_size=2, stride=2, mode="avg")
layer3 = ConvLayer(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0)
layer4 = PoolingLayer(kernel_size=2, stride=2, mode="avg")
layer5 = FCLayer(in_features=400, out_features=120) # both fully connected layers assume the input image dimensions are (h,w,c) = (32, 32, 1)
layer6 = FCLayer(in_features=120, out_features=84)
layer7 = FCLayer(in_features=84, out_features=1) #y_hat is int b/w 0 and 9 to classify handwritten digit, nowadays prefer softmax

# Forward pass
output = layer1.forward(batch_rand_gray_imgs)
output = layer2.forward(output)
output = layer3.forward(output)
output = layer4.forward(output)
output = output.reshape(num_batches, -1)
output = layer5.forward(output)
output = layer6.forward(output)
output = layer7.forward(output)

print(output)

[[-14004126.31782551]
 [-11998811.82382246]
 [-13235583.18652075]]

In [ ]:

# AlexNet Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 3, 227, 227))

# Create layers
layer1 = ConvLayer(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=0)
layer2 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer3 = ConvLayer(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2) # SAME convolutional layer
layer4 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer5 = ConvLayer(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1)
layer6 = ConvLayer(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1)
layer7 = ConvLayer(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1)
layer8 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer9 = FCLayer(in_features=9216, out_features=4096) #9216 = 256 * 6 * 6
layer10 = FCLayer(in_features=4096, out_features=4096)
layer11 = FCLayer(in_features=4096, out_features=1000) 

# Forward pass
output = layer1.forward(batch_rand_gray_imgs)
output = layer2.forward(output)
output = layer3.forward(output)
output = layer4.forward(output)
output = layer5.forward(output)
output = layer6.forward(output)
output = layer7.forward(output)
output = layer8.forward(output)
output = output.reshape(num_batches, -1)
output = layer9.forward(output)
output = layer10.forward(output)
output = layer11.forward(output)
# normally would apply softmax on this 1000-feature vector output

for i in range(num_batches):
    plt.semilogy(output[i, :])
    plt.show()

In [32]:

# VGG-16 Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 3, 224, 224))

# Create layers
# Block 1
layer1_1 = ConvLayer(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
layer1_2 = ConvLayer(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
pool1 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 2
layer2_1 = ConvLayer(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
layer2_2 = ConvLayer(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)
pool2 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 3
layer3_1 = ConvLayer(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)
layer3_2 = ConvLayer(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
layer3_3 = ConvLayer(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
pool3 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 4
layer4_1 = ConvLayer(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1)
layer4_2 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer4_3 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
pool4 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 5
layer5_1 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer5_2 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer5_3 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
pool5 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Fully Connected Layers
fc1 = FCLayer(in_features=25088, out_features=4096) # 25088 = 512 * 7 * 7
fc2 = FCLayer(in_features=4096, out_features=4096)
fc3 = FCLayer(in_features=4096, out_features=1000)

# Forward pass
output = layer1_1.forward(batch_rand_gray_imgs)
output = layer1_2.forward(output)
output = pool1.forward(output)

output = layer2_1.forward(output)
output = layer2_2.forward(output)
output = pool2.forward(output)

output = layer3_1.forward(output)
output = layer3_2.forward(output)
output = layer3_3.forward(output)
output = pool3.forward(output)

output = layer4_1.forward(output)
output = layer4_2.forward(output)
output = layer4_3.forward(output)
output = pool4.forward(output)

output = layer5_1.forward(output)
output = layer5_2.forward(output)
output = layer5_3.forward(output)
output = pool5.forward(output)

output = output.reshape(num_batches, -1)
output = fc1.forward(output)
output = fc2.forward(output)
output = fc3.forward(output)

for i in range(num_batches):
    plt.semilogy(output[i, :])
    plt.show()

$\textbf{Problem}$: Why is a CNN a better choice of architecture than an MLP for computer vision tasks?

$\textbf{Solution}$: As a simple example, given a convolutional layer of a CNN in which a square $n\times n\times c$ image is convolved (without striding or padding) with $N$ kernels each of size $k\times k$, producing an output volume of dimensions $(n-k+1)\times (n-k+1)\times N$, the number of parameters is expected to be $N(k^2+1)$ (each of the $N$ kernels has $k^2$ weights and a single bias term). By contrast, if one were to instead implement a fully connected layer of dense connections between the flattened input image of size $cn^2$ and the flattened output volume of size $N(n-k+1)^2$, the total number of parameters would be $cn^2N(n-k+1)^2$ (much larger!). Essentially, this is due to the $\textit{parameter sharing}$ property of CNNs, the sparsity of their connections. Finally, the convolutional structure is better at capturing $\textit{translational invariance}$ of the image (unlike the MLP where pixels are just flattened arbitrarily, losing their spatial information).

$\textbf{Problem}$: Explain what a $\textit{residual block}$ is and how several residual blocks may be composed together to define a $\textit{residual neural network}$ (ResNet).

$\textbf{Solution}$: A residual block may be viewed as a simple additive perturbation to a standard $2$-layer feedforward neural network. Recall that, given an input feature vector $\mathbf x$, a feedforward layer $(W,\mathbf b)$ may be viewed as the composition of a $\textit{linear layer}$ $\mathbf x\mapsto W\mathbf x + \mathbf b$ with a nonlinear activation function $\boldsymbol{\sigma}$. Thus, the overall action of a $2$-layer feedforward looks something like: $$\mathbf x\mapsto\boldsymbol{\sigma}(W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2)$$

By contrast, in a residual block, the overall action looks like: $$\mathbf x\mapsto\boldsymbol{\sigma}(W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2 + \mathbf x)$$

(assuming compatible dimensions for vector addition with $\mathbf x$). Thus, although the $1^{\text{st}}$ internal layer of the residual block is no different from a standard feedforward, the difference lies in the insertion of a so-called $\textit{skip connection}$ $+\mathbf x$ in between the linear layer and the activation function of the $2^{\text{nd}}$ layer. A ResNet is then literally just a bunch of residual blocks composed together!

$\textbf{Problem}$: Give some hand-wavy intuition/justification for what makes ResNets interesting/useful.

$\textbf{Solution}$: Consider the following $\textit{thought experiment}$. A standard feedforward neural network classifier with $20$ layers, post-training, has a misclassification rate of $15\%$ $\textit{on the training set}$ itself. Clearly, this seems to be a bias/underfitting problem rather than a variance/overfitting problem because it’s on the $\textit{training set}$. So one reasons that, to make the network more expressive and improve its performance on the training set, one might decide to append another $36$ feedforward layers on top of the $20$ layers already there, thus creating a $56$-layer feedforward. Theoretically, it’s manifest that the optimal training set error can only get lower because one could simply have the first $20$ layers doing exactly what they were doing before, and then have the remaining $36$ layers implement an identity function $\mathbf x\mapsto\mathbf x$. In practice however, it turns out that, paradoxically, training set error actually increases!

The problem is that it turns out to be surprisingly delicate to learn an identity mapping $\mathbf x\mapsto\mathbf x$ due to exploding/vanishing gradient problems during cost function minimization. So rather than struggling so hard to learn the identity, why not redesign the architecture such that the identity $\mathbf x\mapsto\mathbf x$ is the $\textit{default}$ mapping, and to then merely learn whatever perturbation (a.k.a. $\textit{residual}$ hence the network’s name) $\boldsymbol{\varepsilon}(\mathbf x)$ to this identity mapping is needed to learn the actual underlying map of interest $\mathbf x\mapsto \mathbf x+\boldsymbol{\varepsilon}(\mathbf x)$. Roughly speaking, in the earlier notation, one has parameterized the residual by $\boldsymbol{\varepsilon}(\mathbf x)=W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2$, assuming the outer activation function $\boldsymbol{\sigma}$ may be neglected (or if it happens to be a ReLU, this argument is even more appealing!). Notice than that (loosely) $W_1=W_2=\mathbf b_1=\mathbf b_2=\mathbf 0$ gives $\boldsymbol{\varepsilon}(\mathbf x)=\mathbf 0$, hence learning an exact identity map $\mathbf x\mapsto\mathbf x$. But it’s very easy to coerce weights and biases towards $\mathbf 0$ using standard regularization techniques.

In [33]:

# ResNets were first used in computer vision, so the W_1, W_2 matrices above really correspond to convolutional layers in a CNN structure
class ResBlock():
    def __init__(self, in_channels, out_channels, stride=1):
        self.conv1 = ConvLayer(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.conv2 = ConvLayer(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        
        # Identity x path (Projection Shortcut if dimensions change)
        self.projection = None
        if stride != 1 or in_channels != out_channels:
            # 1x1 convolution to match dimensions
            self.projection = ConvLayer(in_channels, out_channels, kernel_size=1, stride=stride, padding=0)
            
    def relu(self, x):
        return np.maximum(0, x)

    def forward(self, x):
        identity = x
        
        # 1. First transformation + activation
        out = self.conv1.forward(x)
        out = self.relu(out)
        
        # 2. Second transformation
        out = self.conv2.forward(out)
        
        # 3. Match dimensions of identity if needed
        if self.projection is not None:
            identity = self.projection.forward(x)
            
        # 4. Skip Connection Addition
        out += identity
        
        # 5. Final activation
        out = self.relu(out)
        
        return out

$\textbf{Problem}$: Just like a ResNet is made from composing a sequence of residual blocks together, an inception neural network (InceptionNet) is made from composing a sequence of inception blocks together (mostly true; in practice there may also some other side branches, pooling layers, etc). Explain the architecture of an inception block.

$\textbf{Solution}$: Whereas typical CNN architectures require one to choose a filter size, an inception block asks “why not try several filter sizes?”. To this effect, a naive inception block implementation might look something like:

In [34]:

class NaiveInceptionBlock():
    def __init__(self, in_channels, out_channels):
        self.conv1 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=3, stride=1, padding=1) #same convolution
        self.conv3 = ConvLayer(in_channels, out_channels, kernel_size=5, stride=1, padding=2) #same convolution
        self.pool = PoolingLayer(kernel_size=2, stride=2, mode="max")

    def forward(self, x):
        x1 = self.conv1.forward(x)
        x2 = self.conv2.forward(x)
        x3 = self.conv3.forward(x)
        x4 = self.pool.forward(x)
        return np.concatenate([x1, x2, x3, x4], axis=1)

$\textbf{Problem}$: Needless to say, one of the drawbacks to asking “why not try several filter sizes?” is that, by having to literally try several filter sizes, the naive inception block implementation above can get computationally expensive in terms of the number of multiplications and additions required per forward pass. Explain how the introduction of $1\times 1$ $\textit{bottleneck layers}$ can help reduce computational burden without sacrificing too much accuracy.

$\textbf{Solution}$: The idea is to insert $1\times 1$ convolutional layers before the more computationally expensive $3\times 3$ and $5\times 5$ convolutions. These $1\times 1$ convolutions act as “bottlenecks” by reducing the number of channels (and thus the computational cost) before the larger convolutions are applied. For example, instead of applying a $5\times 5$ convolution directly to 256 input channels, we can first apply a $1\times 1$ convolution to reduce the channels to, say, 64, and then apply the $5\times 5$ convolution to the reduced-channel tensor. This significantly reduces the number of parameters and computations while maintaining most of the representational power of the larger filters.

In [35]:

class InceptionBlockWithBottleNeck(NaiveInceptionBlock):
    def __init__(self, in_channels, num_1x1, num_3x3_reduce, num_3x3, num_5x5_reduce, num_5x5, num_pool_proj):
        super().__init__(in_channels, num_1x1, num_3x3, num_5x5, num_pool_proj)
        self.conv_3x3_reduce = ConvLayer(in_channels, num_3x3_reduce, kernel_size=1)
        self.conv_5x5_reduce = ConvLayer(in_channels, num_5x5_reduce, kernel_size=1)
        self.conv_3x3_reduce_out = ConvLayer(num_3x3_reduce, num_3x3, kernel_size=3, padding=1)
        self.conv_5x5_reduce_out = ConvLayer(num_5x5_reduce, num_5x5, kernel_size=5, padding=2)

    def forward(self, x):
        out_1x1 = self.conv_1x1(x)
        out_3x3_reduce = self.conv_3x3_reduce(x)
        out_3x3 = self.conv_3x3_reduce_out(out_3x3_reduce)
        out_5x5_reduce = self.conv_5x5_reduce(x)
        out_5x5 = self.conv_5x5_reduce_out(out_5x5_reduce)
        out_pool = self.pool(x)
        out_pool_proj = self.pool_proj(out_pool)
        return np.concatenate([out_1x1, out_3x3, out_5x5, out_pool_proj], axis=1)

$\textbf{Problem}$: Explain how the MobileNet architecture replaces the standard convolutional layer with a depthwise separable convolutional layer.

$\textbf{Solution}$: The idea is to sort of break down the standard convolutional layer into two parts: a $\textit{depthwise convolutional layer}$ and a $\textit{pointwise convolutional layer}$. The depthwise convolutional layer is a convolutional layer that only convolves with the depth (channels) dimension, while the pointwise convolutional layer is a convolutional layer that only convolves with the height and width dimensions.

Posted in Blog | Leave a comment

Hamilton’s Optics-Mechanics Analogy

Posted on February 19, 2026 by wdengquantum.me

Problem: Deduce the Hamilton-Jacobi equation of classical mechanics.

Solution: Instead of viewing the action $S=S[\mathbf x(t)]$ as a functional of the particle’s trajectory $\mathbf x(t)$, it can be viewed more simply as a scalar field $S(\mathbf x,t)$ in which the initial point in spacetime $(t_0,\mathbf x_0)$ is fixed and one simply takes the on-shell trajectory from $(t_0,\mathbf x_0)$ to $(t,\mathbf x)$. Then the total differential $dS=\mathbf p\cdot d\mathbf x$ (follows from the usual Noetherian calculation) so in particular:

\[\mathbf p=\frac{\partial S}{\partial\mathbf x}\]

Intuitively, this is saying that the particle moves in a direction (the direction of the momentum $\mathbf p$) orthogonal to the contour surfaces of the action field $S$, i.e. such isosurfaces can be viewed as “wavefronts”. Then the total time derivative is:

\[\dot S=L\]

But $\frac{\partial S}{\partial t}+\frac{\partial S}{\partial\mathbf x}\cdot\dot{\mathbf x}=\frac{\partial S}{\partial t}+\mathbf p\cdot\dot{\mathbf x}$. Thus, isolating for $H=\mathbf p\cdot\dot{\mathbf x}-L$ yields the Hamilton-Jacobi nonlinear $1^{\text{st}}$-order PDE for $S(\mathbf x,t)$:

\[-\frac{\partial S}{\partial t}=H\left(\mathbf x,\frac{\partial S}{\partial\mathbf x},t\right)\]

Problem: When $\partial H/\partial t=0$, the Hamiltonian is conserved with energy $H=E$, so this motivates the additive separation of variables, $S(\mathbf x,t):=S_0(\mathbf x)-Et$ for some constant $E$. What does the Hamilton-Jacobi equation simplify to in this case? For a single non-relativistic particle of mass $m$ moving in a potential $V(\mathbf x)$, what does this look like? What about in $1$ dimension?

Solution: \[H\left(\mathbf x,\frac{\partial S_0}{\partial\mathbf x}\right)=E\]

which for $H(\mathbf x,\mathbf p)=|\mathbf p|^2/2m+V(\mathbf x)$ looks like:

\[\frac{1}{2m}\biggr|\frac{\partial S_0}{\partial\mathbf x}\biggr|^2+V(\mathbf x)=E\]

and in $1$ dimension is integrable to the explicit solution:

\[S_0(x)=\pm\int ^xdx’\sqrt{2m(E-V(x’))}\]

In particular, the usual trajectory $x(t)$ can be obtained by treating $S_o=S_0(x,t;E)$ as a family of solutions parameterized by the energy $E$; this works because $S_0$ can be a viewed as a particular generating function of a canonical transformation $(\mathbf x,\mathbf p,H)\mapsto (\mathbf x’,\mathbf p’,H’)$ in which the “boosted” Hamiltonian vanishes $H’=0$.

\[\frac{\partial S_0}{\partial E}=-t_0\Rightarrow t-t_0=\pm\sqrt{\frac{m}{2}}\int_{x_0}^{x(t)}\frac{dx’}{\sqrt{E-V(x’)}}\]

Problem: Above, the static field $S_0(\mathbf x)$ was introduced to simplify the Hamilton-Jacobi equation when the energy $E$ was conserved. However, if one pulls back to the level of functionals rather than fields, one can define an analogous abbreviated action functional $S_0[\mathbf x]$ which depends only on the path $\mathbf x$ taken rather than the trajectory $\mathbf x(t)$. Define $S_0[\mathbf x]$, and moreover show that when the energy $E$ is conserved, the on-shell path is a stationary point of $S_0$ (this is called Maupertuis’s principle).

Solution: The abbreviated action for a single particle of mass $m$ and (non-relativistic) energy $E=|\mathbf p|^2/2m +V(\mathbf x)$ is:

\[S_0[\mathbf x]:=\int d\mathbf x\cdot\mathbf p=\int ds |\mathbf p|=\int ds\sqrt{2m(E-V(\mathbf x))}\]

(modifying this to $S_0[\mathbf x]:=\int ds |\mathbf p|=\int ds\sqrt{2(E-V(\mathbf x))}$ allows for interpretation as an $N$-particle system in configuration space $\mathbf x\in\mathbf R^{3N}$ with the Riemannian “mass metric” $ds^2=m_1|d\mathbf x_1|^2+…+m_N|d\mathbf x_N|^2$).

To find the stationary paths of $S_0[\mathbf x]$ subject to the constraint $H(\mathbf x,\mathbf p)=E$, one can implement a Lagrange multiplier $\gamma(\tau)$ to perform unconstrained extremization of:

\[S[\mathbf x(\tau)]:=S_0[\mathbf x(\tau)]-\int d\tau\gamma(\tau)(H(\mathbf x,\mathbf p)-E)=\int d\tau (\mathbf p\cdot\dot{\mathbf x}-\gamma(\tau)(H(\mathbf x,\mathbf p)-E))\]

The Euler-Lagrange equations lead to Hamilton’s equations:

\[\frac{d\mathbf x}{d\tau}=\gamma\frac{\partial H}{\partial\mathbf p}\]

\[\frac{d\mathbf p}{d\tau}=-\gamma\frac{\partial H}{\partial\mathbf x}\]

provided the Lagrange multiplier $\gamma=dt/d\tau$ encodes reparameterization invariance; with this choice it’s clear that the integrand in the functional $S$ was nothing more than the Lagrangian $L=\mathbf p\cdot\dot{\mathbf x}-H$ (plus an unimportant constant $E$) so Maupertuis’s principle reduces to the usual Hamilton’s principle.

Problem: What does Fermat’s principle in ray optics assert? Hence, derive the ray equation.

Solution: The time functional $T=T[\mathbf x(s)]$ of a ray trajectory $\mathbf x(s)$ is stationary on-shell. That is:

\[cT[\mathbf x(s)]=\int ds n(\mathbf x(s))\]

This is reparameterization invariant, since one can arbitrarily parameterize $\mathbf x=\mathbf x(t)$ and replace $ds=dt|\dot{\mathbf x}|$. The corresponding Euler-Lagrange equations are:

\[\frac{d}{dt}\left(n(\mathbf x)\frac{\dot{\mathbf x}}{|\dot{\mathbf x}|}\right)=|\dot{\mathbf x}|\frac{\partial n}{\partial\mathbf x}\]

But by choosing the natural parameterization $t:=s$ one has $|d\mathbf x/ds|=1$, hence the ray equation:

\[\frac{d}{ds}\left(n\frac{d\mathbf x}{ds}\right)=\frac{\partial n}{\partial\mathbf x}\]

This can also be written in terms of the curvature vector $\boldsymbol{\kappa}=d^2\mathbf x/ds^2$:

\[\boldsymbol{\kappa}=\left(\frac{\partial\ln n}{\partial\mathbf x}\right)_{\perp d\mathbf x}\]

Problem: Starting from an arbitrary Cartesian component $\psi(\mathbf x,t)=\psi_0(\mathbf x)e^{i(k_0cT(\mathbf x)-\omega t)}$ of either the $\mathbf E$ or $\mathbf B$ fields of a light wave (here $\omega=ck_0$ with $k_0=2\pi/\lambda_0$ is the free space wavenumber), make the eikonal approximation to the dispersionless wave equation obeyed by $\psi$ in order to obtain the (scalar) eikonal equation. By defining light rays as the integral curves of the eikonal field $cT(\mathbf x)$ (a kind of local optical path length), reproduce the vector eikonal equation from Fermat’s principle above.

Solution: The ansatz $\psi(\mathbf x,t)=\psi_0(\mathbf x)e^{i(k_0cT(\mathbf x)-\omega t)}$ is easy to justify; the $e^{-i\omega t}$ is a just a Fourier transform factor that reduces the wave equation $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi=\frac{n^2}{c^2}\frac{\partial^2\psi}{\partial t^2}$ to a Helmholtz equation $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi=-n^2k_0^2\psi$. The remaining piece is just a polar parameterization of an arbitrary $\mathbf C$-valued spatial field $\psi_0(\mathbf x)e^{ik_0cT(\mathbf x)}$. One obtains:

\[\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|^2=n^2+\frac{1}{k_0^2\psi_0}\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi_0+\frac{2i}{k_0\psi_0}\frac{\partial\psi_0}{\partial\mathbf x}\cdot\frac{\partial cT}{\partial\mathbf x}+\frac{i}{k_0}\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2cT\]

The eikonal approximation amounts to taking the ray optics limit $k_0\to\infty$ (in practice, the wavelength $2\pi/k_0$ has to be much shorter than all other length scales), and yields the (scalar) eikonal equation:

\[\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|=n\]

A light ray is thus a trajectory $\mathbf x(s)$ with unit tangent vector:

\[\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial cT}{\partial\mathbf x}\]

The rest is an application of the chain rule:

\[\frac{d}{ds}=\frac{\partial}{\partial\mathbf x}\cdot\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial}{\partial\mathbf x}\cdot\frac{\partial cT}{\partial\mathbf x}\]

followed by the identity:

\[\left(\frac{\partial cT}{\partial\mathbf x}\cdot\frac{\partial}{\partial\mathbf x}\right)\frac{\partial cT}{\partial\mathbf x}=\frac{1}{2}\frac{\partial}{\partial\mathbf x}\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|^2\]

to deduce the (vector) eikonal equation of motion for ray trajectories just as Fermat’s principle predicts.

Problem: Hence, what is Hamilton’s optics-mechanics analogy?

Solution: In a nutshell, the isomorphism proceeds as:

\[(n(\mathbf x), cT)\leftrightarrow (|\mathbf p(\mathbf x)|,S_0)\]

Problem: Use Hamilton’s optics-mechanics analogy to solve the brachistochrone problem (this was how Johann Bernoulli originally solved it).

Solution: By energy conservation, the speed of the particle at distance $y>0$ below its initial dropping height is $v=\sqrt{2gy}$. By Fermat’s principle, minimizing the time functional then amounts to treating the particle as a light ray with $n(\mathbf x)=c/v(y)$. So the question becomes how do light rays bend in a horizontally stratified medium with $n(y)\propto y^{-1/2}$? The answer is given by the ray equations:

\[\frac{d}{ds}\begin{pmatrix} y^{-1/2}dx/ds \\ y^{-1/2}dy/ds\end{pmatrix}=\begin{pmatrix}0 \\ y^{-1/2}/2\end{pmatrix}\]

The horizontal component expresses Snell’s law since $dx/ds=\sin\theta$ (it expresses momentum conservation along the homogeneous $\partial n/\partial x=0$ direction). Using the tangent vector constraint $ds^2=dx^2+dy^2$ gives the ODE of a cycloid:

\[\frac{dy}{dx}=\sqrt{\frac{\text{const}}{y}-1}\]

(the vertical component ODE has an analytical solution $y(s)=-s^2/8R+s$ which is contained in the cycloid, so is redundant information).

Problem: How did Hamilton’s optics-mechanics analogy inspire Schrodinger to propose his famous equations of quantum mechanics?

Solution: Essentially, Schrodinger asked: ray optics is to wave optics as classical mechanics is to what? In other words, one imagines there exists a wave theory of particles/matter and one would like to take the “inverse eikonal limit” of classical mechanics (here, “inverse eikonal limit” is usually called quantization):

Just as light rays propagate parallel to their phase fronts:

\[\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial cT}{\partial\mathbf x}\]

Particles propagate parallel to their “action fronts” exactly according to Hamilton’s analogy:

\[\frac{d\mathbf x}{ds}=\frac{1}{|\mathbf p(\mathbf x)|}\frac{\partial S_0}{\partial\mathbf x}\]

Already this suggests that the action should have some phase interpretation. More precisely, it should be the phase of the particle’s de Broglie wave in units of $\hbar$. It’s also not obvious that particle’s should be described by a scalar wave field rather than e.g. the electromagnetic vector wave fields of light. Schrodinger simply guessed it looked like the equivalent of “scalar diffraction theory” with a single wavefunction $\psi(\mathbf x,t)=\psi_0(\mathbf x,t)e^{iS(\mathbf x,t)/\hbar}$. This gives the Madelung equations of quantum hydrodynamics, one of which is just a continuity equation (giving credence to the Born interpretation of $|\psi|^2$) and the other is a quantum Hamilton-Jacobi equation which in the limit $\hbar\to 0$ (analogous to the eikonal limit $\lambda_0\to 0$) simplifies to the classical Hamilton-Jacobi equation.

Posted in Blog | Leave a comment

Quantumplations

Polymers

Magnetism

Physics-Informed Neural Networks

Spectral Bias of Physics-Informed Neural Networks¶

1. Introduction and Background¶

2. Algorithms and Implementation¶

2.1 Network Architecture¶

2.2 Training¶

2.3 Hardware¶

3. Experiment 1: Van der Pol Oscillator¶

3.1 Problem Setup¶

3.2 Results: Time Series and Phase Portraits¶

4. Experiment 2: Fokker-Planck Equation¶

4.1 Problem Setup¶

4.2 Cost Function¶

4.3 Results¶

5. Discussion and Extensions¶

References¶

Diffusion & Flow-Matching Models

Information Geometry

Autoencoders

Graph Neural Networks

Renormalization Group

Convolutional Neural Networks

Hamilton’s Optics-Mechanics Analogy

Archives

Meta