Diffusion 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\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: (something about integral probability metrics such as Wasserstein/earth-mover’s distance with a concrete computation in one dimension)

Solution:

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()

No description has been provided for this image

$\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

Pseudo-Riemannian Geometry

Posted on February 13, 2026 by wdengquantum.me

Problem: Define the signature of a matrix. Hence, state and prove Sylvester’s law of inertia.

Solution: The signature of an $n\times n$ matrix $A$ is a $3$-tuple $(n_+,n_-,n_0)$ where $n_+$ is the number of positive eigenvalues of $A$ (including multiplicity), $n_-$ is the number of negative eigenvalue of $A$ (including multiplicity), and $n_0=\text{dim}\ker A$ is the multiplicity of the zero eigenvalue; thus, $n_++n_-+n_0=n$.

Let $A,B$ be real symmetric matrices. Then Sylvester’s law of inertia asserts that $A$ and $B$ are congruent matrices iff they have the same signature (which sometimes is called “inertia” because of this invariance under congruence, hence the name).

Proof: It’s easy to see that the nullity $n_0$ is preserved by congruence transformations. If one can can show that $n_+$ is also preserved, then it implies $n_-$ is conserved by virtue of $n_++n_-+n_0=n$. To show this, the idea is to prove by contradiction, assuming $n_+$ is not preserved which by dimension counting would imply a non-zero vector living in the intersection of the subspace spanned by the positive-eigenvalue eigenvectors of $A$ and the congruence-transformed subspace spanned by the non-positive-eigenvalue eigenvectors of $B$.

Since any real, symmetric matrix $A$ is isomorphic to a real quadratic form $Q(\mathbf x):=\mathbf x^TA\mathbf x$, the concepts of signature and Sylvester’s law of inertia can also be reformulated in the language of quadratic forms rather than real symmetric matrices.

Problem: Let $X$ be a smooth $n$-manifold. Explain what it means to place the additional structure of a pseudo-Riemannian metric $g$ on $X$. Then explain how Riemannian and Lorentzian geometry are special cases of pseudo-Riemannian geometry.

Solution: A Riemannian metric $g$ on $X$ is a type $(0,2)$ tensor field that defines an inner product $g_x:T_x(X)^2\to\mathbf R$ at the tangent space $T_x(X)$ of each point $x\in X$. The “tensor field” part is sometimes rephrased as saying that $g_x$ is a bilinear form. Moreover, to flesh out the usual axioms of a real inner product space, the Riemannian metric tensor $g_x$ at each $x\in X$ must be symmetric $g_x(v_x,v’_x)=g_x(v’_x,v_x)$ and positive-definite $v_x\neq 0\Rightarrow g_x(v_x,v_x)>0$.

A pseudo-Riemannian metric relaxes the positive-definite requirement of a Riemannian metric, instead merely requiring non-degeneracy (i.e. at each point $x\in X$, the zero vector $0$ is the only vector orthogonal $g_x(0,v_x)=0$ to all $v_x\in T_x(X)$). More concisely, what this is saying is that $n_0=0$, so the signature of a pseudo-Riemannian metric may be thought of as a pair $(n_+,n_-)$.

Thus, Riemannian metrics are the special subset of pseudo-Riemannian metrics for which $(n_+,n_-)=(n,0)$. Meanwhile, Lorentzian metrics are another special subset of pseudo-Riemannian metrics for which $(n_+,n_-)=(n-1,1)$ (this is the relativists’/ convention). Typically, spacetime $X$ has dimension $n=4$, so e.g. the Minkowski metric is said to have signature $(3,1):=(-,+,+,+)$.

Problem: Let $(X,g)$ be a pseudo-Riemannian manifold. Often, one writes the general expression for an infinitesimal line element on $X$ as $ds^2=g_{\mu\nu}dx^{\mu}dx^{\nu}$; explain the shorthand being used here.

Solution: This is nothing more than choosing some chart $x^{\mu}$ on $X$ and expanding the metric tensor field $g$ in the coordinate basis $\{dx^{\mu}\otimes dx^{\nu}\}$ of type $(0,2)$ tensor fields:

\[g=g_{\mu\nu}dx^{\mu}\otimes dx^{\nu}\]

for some real, symmetric scalar fields $g_{\mu\nu}:X\to\mathbf R$. One then simply writes $ds^2:=g$ and omits the tensor product $\otimes$.

Problem: Let $x(t)\in X$ be a curve on a Riemannian manifold $(X,g)$. By choosing a chart $x^{\mu}$ to cover a suitable region of $X$ spanned by the trajectory $x(t)$, explain how to compute the length $s$ of the curve $x(t)$ using the Riemannian metric $g$.

Solution: Heuristically, it is:

\[s=\int ds=\int\sqrt{g_{\mu\nu}dx^{\mu}dx^{\nu}}=\int dt\sqrt{g_{\mu\nu}(x(t))\dot{x}^{\mu}\dot{x}^{\nu}}\]

Problem: Write down the action $S$ for a non-relativistic particle of mass $m$ moving on a Riemannian manifold $(X,g)$. Hence, derive the geodesic equation. What happens if the particle is relativistic?

Solution: A nonrelativistic free particle of mass $m$ has action $S=\int dt L$ described by the usual nonrelativistic kinetic Lagrangian:

\[L=\frac{m}{2}|\dot{\mathbf x}|^2\]

The catch is that $|\dot{\mathbf x}|^2=g_{ij}(\mathbf x)\dot x^{i}\dot x^{j}$ depends on the Riemannian metric $g$ in which the particle moves. Applying the Euler-Lagrange equations, one finds the equation of motion (called the geodesic equation):

\[\ddot x^i+\Gamma^i_{jk}\dot x^j\dot x^k=0\]

where the Christoffel symbols $\Gamma^i_{jk}(\mathbf x):=\frac{1}{2}g^{i\ell}(\mathbf x)\left(\frac{\partial g_{\ell j}}{\partial x^k}+\frac{\partial g_{\ell k}}{\partial x^j}-\frac{\partial g_{jk}}{\partial x^{\ell}}\right)$ are analogous to “fictitious forces” on the particle due to the curvature of $X$.

By contrast, a relativistic free particle has action $S=-mcs=-mc^2\tau$. This can be recast into non-relativistic form $S=\int dt L$ using the relativistic kinetic Lagrangian:

\[L=-mc\sqrt{g_{\mu\nu}\dot x^{\mu}\dot x^{\nu}}\]

which in flat space $g_{\mu\nu}=\text{diag}(1,-1,-1,-1)$ reduces to the familiar $L=-mc^2/\gamma$ (nb. more generally, the quantity under the square root is always $\geq 0$ for the timelike worldlines of particles). Some comments:

The stationary action principle $\delta S=0$ in which the action $S$ tends to be minimized corresponds to the well-known fact that relativistic free particles travelling in straight line geodesics through spacetime maximize the Minkowski distance $s$, or equivalently Alice experiences more proper time (i.e. ages more) than Bob in the twin paradox.
The relativistic action (unlike its non-relativistic counterpart) is manifestly reparameterization invariant.

Now the Euler-Lagrange geodesic equations for the relativistic Lagrangian are almost identical to the non-relativistic case except there’s an additional term:

\[\ddot x^{\mu}+\Gamma^{\mu}_{\nu\rho}(x)\dot x^{\nu}\dot x^{\rho}=\frac{\dot L}{L}\dot x^{\mu}\]

The additional term vanishes $\dot L=0$ iff $L=L(\tau’)$ is parameterized as any affine transformation $\tau’=a\tau+b$ of the particle’s proper time $\tau$ (prove this by showing $\dot{\tau}’=-aL/mc^2$). Only in this case does the relativistic geodesic equation reduce to the non-relativistic form:

\[\frac{d^2 x^{\mu}}{d\tau’^2}+\Gamma^{\mu}_{\nu\rho}(x(\tau’))\frac{dx^{\nu}}{d\tau’}\frac{dx^{\rho}}{d\tau’}=0\]

Problem: For $X=\mathbf R^2$, calculate all non-vanishing Christoffel symbols in the polar coordinate chart $(\rho,\phi)$.

Solution:

Note it is usually more efficient to compute Christoffel symbols this way rather than using the explicit formula in terms of partial derivatives of the metric $g$ (though they’re of course equivalent).

Problem: Explain how the presence of a pseudo-Riemannian metric $g$ on a smooth manifold $X$ allows one to identify covectors in $T_x^*(X)$ with tangent vectors in $T_x(X)$ at each point $x\in X$.

Solution: Just as in $X=\mathbf R^n$, one identifies a tangent vector $\mathbf v$ with its covector $\mathbf v\cdot$ via the inner product $\cdot$, so at a given $x\in X$, one uses the pseudo-inner product $g_x$ on $T_x(X)$ to make the exact same identification $v_x\leftrightarrow g_x(v_x,\space)$ (this sometimes called the musical isomorphism induced by $g$ in light of the map $\flat_x:T_x(X)\to T^*_x(X)$ defined by $v_x^{\flat_x}(v’_x):=g(v_x,v’_x)$ and its inverse $\sharp_x=\flat_x^{-1}$).

In some chart $x^{\mu}$, one can write the $1$-form $v^{\flat}=(v^{\flat})_{\mu}dx^{\mu}$ with components $(v^{\flat})_{\mu}=v^{\flat}(\partial_{\mu})=g(v,\partial_{\mu})=g_{\mu\nu}v^{\nu}$ where $v^{\nu}=v(x^{\nu})$ are the components of its musically isomorphic vector field $v=v^{\nu}\partial_{\nu}$ in the same coordinate basis. Physicists typically abuse notation by merely writing $(v^{\flat})_{\mu}=g_{\mu\nu}v^{\nu}$ as just $v_{\mu}=g_{\mu\nu}v^{\nu}$, identifying $v^{\flat}\equiv v$ under the musical isomorphism. But then this lazy notation makes it look as if the metric $g$ performs a trivial mechanical action of just “lowering the index” on $v$.

Problem: Define a tensor field $g^*$ that performs the trivial mechanical action of “raising the index” on a $1$-form $A$.

Solution: Define $g^*_x:T^*_x(X)^2\to\mathbf R$ to be a type $(2,0)$ tensor given by:

\[g^*(A,A’):=g(A^{\sharp}, A’^{\sharp})\]

Then on the one hand, in some chart $x^{\mu}$, one has:

\[g^*(v^{\flat},v’^{\flat})=g(v,v’)=g_{\mu\nu}v^{\mu}v’^{\nu}\]

On the other hand, in that same chart $x^{\mu}$, one has:

\[g^*(v^{\flat},v’^{\flat})=(g^*)^{\mu\nu}(v^{\flat})_{\mu}(v’^{\flat})_{\nu}=(g^*)^{\mu\nu}g_{\mu\rho}v^{\rho}g_{\nu\sigma}v’^{\sigma}=(g^*)^{\rho\sigma}g_{\mu\rho}g_{\nu\sigma}v^{\mu}v’^{\nu}\]

This implies $g_{\mu\nu}=(g^*)^{\rho\sigma}g_{\mu\rho}g_{\nu\sigma}$. Since $g$ is non-degenerate (the $N_0=0$ axiom of a pseudo-Riemannian metric), it follows that $\det g_{\mu\nu}\neq 0$ is invertible, so the only possibility is $(g^*)^{\mu\nu}g_{\nu\rho}=\delta^{\mu}_{\rho}$, i.e. $g^*\sim g^{-1}$. It is common practice to just drop the $*$ and write $g^{\mu\nu}$ in lieu of $(g^*)^{\mu\nu}$. With this lazy notation, $g^{\mu\nu}$ looks like it just “raises the index” on $A^{\mu}=g^{\mu\nu}A_{\nu}$ (where one last abuse of notation $(A^{\sharp})^{\mu}=A^{\mu}$ has been made!). cf. raising lowering operators $a,a^{\dagger}$ in quantum mechanics.

Problem: What is the canonical volume form on a smooth, orientable pseudo-Riemannian $n$-manifold $(X,g)$? Why is it canonical?

Solution: In a chart $x^{\mu}$ in which the pseudo-Riemannian metric has components $g=g_{\mu\nu}dx^{\mu}\otimes dx^{\nu}$, the canonical volume form is $\sqrt{|\det g|}dx^1\wedge…\wedge dx^n$. This is indeed a volume form because it is a top form with $\det g\neq 0$ nowhere vanishing for same reasons as before. Moreover, it is canonical because, despite appearances, it does not depend on the choice of chart $x^{\mu}$ (provided it’s the same orientation) as $\sqrt{|\det g|}$ is a scalar density of weight $1$ whereas $dx^1\wedge…\wedge dx^n$ is a scalar density of weight $-1$.

Problem: Let $(X,g)$ be a smooth, $n$-dimensional, orientable pseudo-Riemannian manifold, and let $\omega\in\Omega^k(X)$ be a differential $k$-form on $X$. Define the Hodge dual $(n-k)$-form $\star\omega\in\Omega^{n-k}(X)$ on $X$.

Solution: The Hodge dual $\star\omega\in\Omega^{n-k}(X)$ is defined to be the unique differential $(n-k)$-form such that for all test differential $k$-forms $\omega’\in\Omega^k(X)$ one has the equality of top forms:

\[\omega’\wedge\star\omega=\widetilde{\langle\omega’,\omega\rangle}\sqrt{|\det g|}dx^1\wedge…\wedge dx^n\]

Here, the operation $\widetilde{\langle\omega’,\omega\rangle}$ is defined for $\omega’=A’_1\wedge…\wedge A’_k$ and $\omega=A_1\wedge…\wedge A_k$ by the Gram determinant:

\[\widetilde{\langle\omega’,\omega\rangle}:=\det\begin{pmatrix}g^*(A’_1,A_1) & \dots & g^*(A’_1,A_k) \\ \vdots & \ddots & \vdots \\ g^*(A’_k,A_1) & \dots & g^*(A’_k,A_k)\end{pmatrix}\]

and extended to arbitrary differential $k$-forms $\omega’,\omega\in\Omega^k(X)$ by bilinearity. It is easy to check that the Hodge star $\star:\Omega^k(X)\to\Omega^{n-k}(X)$ is a linear transformation $\star(\phi_1\omega_1+\phi_2\omega_2)=\phi_1\star\omega_1+\phi_2\star\omega_2$ between $2$ vector spaces of the same dimension $\dim\Omega^k(X)={n\choose k}={n\choose n-k}=\dim\Omega^{n-k}(X)$. Moreover, it is either an involution or behaves like multiplication by $i$ in the sense that $\star^2=(-1)^{k(n-k)+n_-}$.

Problem: For flat Minkowski spacetime $X=\mathbf R^{1+3}$ with the Lorentzian metric $ds^2=c^2dt^2-dx^2-dy^2-dz^2$ and orientation defined by the canonical volume form $cdt\wedge dx\wedge dy\wedge dz$, compute the Hodge dual:

\[\star(e^{-x\cos y}dy\wedge dt+4tz^2 dx\wedge dz)\]

Solution: By linearity, one can focus on computing Hodge duals of the basis $2$-forms $\star (dy\wedge dt)$ and $\star(dx\wedge dz)$. Demonstrating this for $\star (dy\wedge dt)$, the idea is to (roughly speaking) wedge it against itself and exploit the defining property of the Hodge dual:

\[dy\wedge dt\wedge\star(dy\wedge dt)=\det\begin{pmatrix}g^{yy} & g^{yt} \\ g^{ty} & g^{tt}\end{pmatrix}cdt\wedge dx\wedge dy\wedge dz\]

\[=\det\begin{pmatrix}-1 & 0 \\ 0 & 1/c^2\end{pmatrix}cdt\wedge dx\wedge dy\wedge dz\]

\[=-\frac{1}{c}dt\wedge dx\wedge dy\wedge dz\]

\[=-\frac{1}{c}dy\wedge dt\wedge dx\wedge dz\]

In this case, one can directly read off the answer $\star(dy\wedge dt)=-c^{-1}dx\wedge dz$ or equivalently $\star(dy\wedge dct)=dz\wedge dx$. Analogously, one can check $\star(dx\wedge dz)=dy\wedge dct$ (consistent with the earlier identity for $\star^2$), so the final answer is:

\[\star(e^{-x\cos y}dy\wedge dt+4tz^2 dx\wedge dz)=\frac{e^{-x\cos y}}{c}dz\wedge dx+4tz^2dy\wedge dct\]

For more general non-orthogonal metrics $g$, one cannot simply read off the answer like that but should set up an ansatz for the Hodge dual like $\star(dy\wedge dt)=c_{tx}dt\wedge dx+c_{ty}dt\wedge dy+c_{tz}dt\wedge dz+c_{xy}dx\wedge dy+c_{xz}dx\wedge dz+c_{yz}dy\wedge dz$ and obtaining the coefficients by wedging against the relevant basis differential forms.

Problem: Let $(X,g)$ be a smooth, $n$-dimensional pseudo-Riemannian manifold. Just as the metric tensor field $g$ induces a pseudo-inner product between tangent vectors at each point across the manifold $X$, show that $g$ also induces an inner product on the space $\Omega^k(X)$ of differential $k$-forms on $X$.

Solution: Given $2$ differential $k$-forms $\omega,\omega’\in\Omega^k(X)$, their $g$-induced inner product is defined by:

\[\langle\omega,\omega’\rangle:=\int_{X}\omega\wedge\star\omega’\]

(this is well-defined because $\omega\wedge\star\omega’\in\Omega^n(X)$ is indeed a top form as observed earlier).

Problem: Having endowed $\Omega^k(X)$ with an inner product, one can define the adjoint (sometimes called the codifferential) $d^{\dagger}:\Omega^k(X)\to\Omega^{k-1}(X)$ of the exterior derivative operator $d:\Omega^{k-1}(X)\to\Omega^k(X)$ by insisting that for arbitrary differential forms $\omega\in\Omega^{k-1}(X),\omega’\in\Omega^k(X)$:

\[\langle d\omega,\omega’\rangle=\langle\omega,d^{\dagger}\omega’\rangle\]

(the LHS is the inner product on $\Omega^k(X)$ whereas the RHS is the inner product on $\Omega^{k-1}(X)$). Show that if the underlying pseudo-Riemannian manifold $X$ is closed, then:

\[d^{\dagger}=(-1)^{n(k+1)+1+n_-}\star d\space\star\]

Solution: Applying graded integration by parts and Stokes’ theorem:

\[\langle d\omega,\omega’\rangle=\int_{\partial X}\omega\wedge\star\omega’+(-1)^k\int_X\omega\wedge d\star\omega’\]

The first term vanishes because $X$ is assumed to be closed, and by comparing with the expected form:

\[\langle\omega,d^{\dagger}\omega’\rangle=\int_X\omega\wedge\star d^{\dagger}\omega’\]

it is clear this would hold provided:

\[(-1)^kd\star=\star d^{\dagger}\]

Isolating for $d^{\dagger}$ by acting with $\star$ on both sides, using $\star^2=(-1)^{(k-1)(n-k+1)+n_-}$, and exploiting trivial identities like $(-1)^{k^2-k}=1$ and $(-1)^{-n}=(-1)^n$ for $k,n\in\mathbf Z$ leads to the claimed result.

Problem: Using the technology of $d^{\dagger}$, define the Laplace-de Rham operator $\Delta:\Omega^k(X)\to\Omega^k(X)$. For the case of a scalar field $k=0$, show that the action of the Laplace-de Rham operator $\Delta$ coincides with that of the negative Laplacian $-\nabla^2$ (sometimes called the Laplace-Beltrami operator).

Solution: One has the positive semi-definite Laplace-de Rham operator:

\[\Delta:=(d+d^{\dagger})^2=\{d,d^{\dagger}\}\]

thanks to nilpotence $d^2=(d^{\dagger})^2=0$. For a scalar field $\phi\in C^{\infty}(X)$, $d^{\dagger}\phi=0$ so:

\[\Delta\phi=d^{\dagger}d\phi=(-1)^{n_-+1}\star d(\partial_{\mu}\phi\star dx^{\mu})\]

The Hodge dual of the $\mu^{\text{th}}$ basis $1$-form can be checked to be:

\[\star dx^{\mu}=(-1)^{\nu}g^{\mu\nu}\sqrt{|\det g|}dx^0\wedge…\wedge dx^{\nu-1}\wedge dx^{\nu+1}\wedge…\wedge dx^{n-1}\]

Thus:

\[d(\partial_{\mu}\phi\star dx^{\mu})=\partial_{\nu}\left(\sqrt{|\det g|}g^{\mu\nu}\partial_{\mu}\phi\right)dx^0\wedge…\wedge dx^{n-1}\]

And finally, the Hodge dual of a top form returns a scalar field:

\[\star (dx^0\wedge…\wedge dx^{n-1})=\frac{\sqrt{|\det g|}}{\det g}=\frac{(-1)^{n_-}}{\sqrt{|\det g|}}\]

Thus, one obtains:

\[\Delta\phi=-\frac{1}{\sqrt{|\det g|}}\partial_{\nu}\left(\sqrt{|\det g|}g^{\mu\nu}\partial_{\mu}\phi\right)=-\nabla^2\phi\]

Problem: Let $\omega\in\Omega^k(X)$ be a differential $k$-form on a closed, Riemannian manifold $(X,g)$. Explain what it means for $\omega$ to be harmonic. Furthermore, prove that $\omega$ is harmonic iff it is both closed $d\omega=0$ and co-closed $d^{\dagger}\omega=0$ (this equivalence underlies the Hodge decomposition theorem and hence the rest of Hodge theory).

Solution: This is just a generalization of the usual notion of a harmonic scalar field obeying Laplace’s equation $\nabla^2\phi=0$, only now the Laplacian $\nabla^2$ is replaced by the Laplace-de Rham operator $\Delta$ and the scalar field $\phi$ by an arbitrary differential $k$-form $\omega\in\Omega^k(X)$:

\[\Delta\omega=0\]

Because $X$ is closed, one can say:

\[\langle\omega,\Delta\omega\rangle=\langle\omega, dd^{\dagger}\omega\rangle+\langle\omega,d^{\dagger}d\omega\rangle\]

\[=\langle d^{\dagger}\omega,d^{\dagger}\omega\rangle+\langle d\omega,d\omega\rangle\]

If $\Delta\omega=0$ is harmonic, then it follows $\langle d^{\dagger}\omega,d^{\dagger}\omega\rangle=\langle d\omega,d\omega\rangle=0$ because $g$ is Riemannian, so in particular $d\omega=d^{\dagger}\omega=0$. The converse is trivial: $\Delta\omega=dd^{\dagger}\omega+d^{\dagger}d\omega=d0+d^{\dagger}0=0$.

Problem: Let $X$ be a smooth manifold not necessarily equipped with a metric $g$. What does it mean to endow $X$ with the structure of a affine connection $\nabla$?

Solution: An affine connection is any map $\nabla:\mathfrak{X}(X)^2\to\mathfrak{X}(X)$ which is linear in its first tangent vector field argument but acts like a tangent vector (more precisely, like a derivation) in its second tangent vector field argument. Thus, $\nabla$ is not tensorial because it is not linear in that second argument. Nevertheless, given a (possibly non-coordinate) basis $\text{span}_{\mathbf R}\{e_{\mu}\}=\mathfrak{X}(X)$ of tangent vector fields, one can still define affine connection coefficients for $\nabla$:

\[\nabla_{\nu}e_{\rho}=\Gamma^{\mu}_{\nu\rho}e_{\mu}\]

with $\nabla_{\nu}:=\nabla_{e_{\nu}}$ a shorthand. The map $\nabla_v:\mathfrak{X}(X)\to\mathfrak{X}(X)$ for some fixed tangent vector field $v\in\mathfrak{X}(X)$ is called the covariant derivative along $v$.

Problem: Let $X$ be a smooth manifold. For two tangent vector fields $v,v’\in\mathfrak{X}(X)$, given a (possibly non-coordinate) basis $\text{span}_{\mathbf R}\{e_{\mu}\}=\mathfrak{X}(X)$ of tangent vector fields so that $v=v^{\mu}e_{\mu}$ and $v’=v’^{\nu}e_{\nu}$, show that:

\[\nabla_vv’=v^{\nu}(e_{\nu}(v’^{\mu})+\Gamma^{\mu}_{\nu\rho}v’^{\rho})e_{\mu}=v^{\nu}\nabla_{\nu}v’\]

Solution: This is a straightforward application of the properties of affine connections. The only thing to watch out for is that the affine connection is extended to act on scalar fields $\phi\in C^{\infty}(X)$ in the obvious way:

\[\nabla_v\phi:=v(\phi)=\mathcal L_v\phi\]

In particular, $\nabla_{\mu}\phi=e_{\mu}(\phi)$, and if $e_{\mu}=\partial_{\mu}$ happens to be a coordinate basis for $\mathfrak{X}(X)$, then the covariant derivative $\nabla_{\mu}=\partial_{\mu}$ coincides with the partial derivative on scalar fields.

Problem: For tangent vector fields $v,v’\in\mathfrak{X}(X)$, compare the covariant derivative $\nabla_vv’$ with the Lie derivative $\mathcal L_vv’$ (both of which are tangent vector fields).

Solution: Basically, in the case of the Lie derivative $\mathcal L_vv’=[v,v’]$, both entries of the commutator are subject to the product rule, for example $[vv^{\prime\prime},v’]=v[v^{\prime\prime},v’]+[v,v’]v^{\prime\prime}$, so in particular are nonlinear (being additive but inhomogeneous). By contrast, in the covariant derivative $\nabla_vv’$, the $v$-argument is explicitly made to be linear (though the $v’$ argument behaves in the same way as both arguments of the Lie derivative), so this is what allows familiar linear algebraic expressions like $\nabla_v=v^{\mu}\nabla_{\mu}$ to be written (this was proven above in the form $\nabla_vv’=v^{\nu}\nabla_{\nu}v’$).

Problem: Explain how to compute the covariant derivative of a $1$-form $A\in\Omega^1(X)$ along a tangent vector field $v\in\mathfrak{X}(X)$ to end up with another $1$-form $\nabla_v A$.

Solution: As a $1$-form, $A$ maps a tangent vector field $v’\in\mathfrak{X}(X)$ to a scalar field $A(v’)\in C^{\infty}(X)$. But earlier, it was already explained one can take the covariant derivative of a scalar field, so it makes sense to compute $\nabla_v(A(v’))=v(A(v’))$. However, by insisting that the product rule holds in the sense:

\[\nabla_v(A(v’))=(\nabla_vA)(v’)+A(\nabla_vv’)\]

So the desired definition ought to be:

\[(\nabla_vA)(v’):=v(A(v’))-A(\nabla_vv’)\]

or in a basis of $\mathfrak{X}(X)$:

\[(\nabla_vA)(v’)=(\nabla_{\mu}A)_{\nu}v^{\mu}v’^{\nu}\]

with $(\nabla_{\mu}A)_{\nu}=\partial_{\mu}A_{\nu}-\Gamma^{\rho}_{\mu\nu}A_{\rho}$.

Problem: Let $X$ be a smooth manifold equipped with an affine connection $\nabla$. Define the type-$(1,2)$ torsion tensor field $T$ of $\nabla$, and compute the components of $T$ in a coordinate basis.

Solution: For a $1$-form $A\in\Omega^1(X)$ and tangent vector fields $v,v’\in\mathfrak{X}(X)$, one has:

\[T(A,v,v’):=A(\nabla_vv’-\nabla_{v’}v-[v,v’])\]

One should check that $T$ is indeed a tensor, since the affine connection $\nabla$ is not a tensor as mentioned earlier. It is easy to see that $T$ is linear in $A$, and it is linear in $v$ iff it is linear in $v’$ because of antisymmetry $T(A,v’,v)=-T(A,v,v’)$. Furthermore, it is easy to see that $T(A,v_1+v_2,v’)=T(A,v_1,v’)+T(A,v_2,v’)$. The tricky property to prove is that $T(A,\phi v,v’)=\phi T(A,v,v’)$ for any scalar field $\phi\in C^{\infty}(X)$ (it must be a scalar field and not just some scalar $\phi\in\mathbf R$ because one is trying to show that $T$ is a tensor field not just a tensor):

\[T(A,\phi v,v’)=A(\nabla_{\phi v}v’-\nabla_{v’}(\phi v)-[\phi v,v’])\]

\[=A(\phi\nabla_vv’-v'(\phi)v-\phi\nabla_{v’}v-\phi[v,v’]+v'(\phi)v)\]

\[=A(\phi\nabla_vv’-\phi\nabla_{v’}v-\phi[v,v’])\]

\[=\phi A(\nabla_vv’-\nabla_{v’}v-[v,v’])=\phi T(A,v,v’)\]

Thus, in a coordinate basis $\{\partial_{\mu}\}$ of $\mathfrak{X}(X)$, the components $T^{\rho}_{\mu\nu}:=T(dx^{\rho},\partial_{\mu},\partial_{\nu})$ are tensorial:

\[T^{\rho}_{\mu\nu}=\Gamma^{\rho}_{\mu\nu}-\Gamma^{\rho}_{\nu\mu}\]

where Clairaut’s theorem $[\partial_{\mu},\partial_{\nu}]=0$ for a coordinate basis has been used.

Problem: Let $X$ be a smooth manifold equipped with an affine connection $\nabla$. Define the type-$(1,3)$ Riemann curvature tensor field $R$ of $\nabla$, and compute the components of $R$ in a coordinate basis.

Solution: For a $1$-form $A\in\Omega^1(X)$ and tangent vector fields $v,v’,v^{\prime\prime}\in\mathfrak{X}(X)$, one has:

\[R(A,v,v’,v^{\prime\prime}):=A(\nabla_v\nabla_{v’}v^{\prime\prime}-\nabla_{v’}\nabla_vv^{\prime\prime}-\nabla_{[v,v’]}v^{\prime\prime})\]

Similar to the torsion tensor field $T$, one should check that $R$ genuinely deserves its designation as a tensor field. In a coordinate basis $\{\partial_{\mu}\}$ of $\mathfrak{X}(X)$, the (slightly awkward but conventional) tensor components $R^{\sigma}_{\rho\mu\nu}:=R(dx^{\sigma},\partial_{\mu},\partial_{\nu},\partial_{\rho})$ are:

\[R^{\sigma}_{\rho\mu\nu}=\partial_{\mu}\Gamma^{\sigma}_{\nu\rho}-\partial_{\nu}\Gamma^{\sigma}_{\mu\rho}+\Gamma^{\lambda}_{\nu\rho}\Gamma^{\sigma}_{\mu\lambda}-\Gamma^{\lambda}_{\mu\rho}\Gamma^{\sigma}_{\nu\lambda}\]

where again Clairaut’s theorem $[\partial_{\mu},\partial_{\nu}]=0$ for a coordinate basis has been used. Again similar to the case for $T$, the antisymmetry $R(A,v’,v,v^{\prime\prime})=-R(A,v,v’,v^{\prime\prime})$ manifests at the component level as $R^{\sigma}_{\rho\nu\mu}=-R^{\sigma}_{\rho\mu\nu}$.

Problem: (about the Ricci identity connecting $T$ and $R$ together in terms of the commutator of covariant derivatives)

Solution:

Problem: State and prove the fundamental theorem of pseudo-Riemannian geometry.

Solution: Let $(X,g)$ be a smooth, pseudo-Riemannian manifold. Then there exists a unique affine connection $\nabla$ (called the Levi-Civita connection) which is torsion-free $T=0$ and $g$-compatible in the sense that $\nabla_v g=0$ for all tangent vector fields $v\in\mathfrak{X}(X)$.

Proof:

Posted in Blog | Leave a comment

Reinforcement Learning (Part $1$)

Posted on February 1, 2026 by wdengquantum.me

Problem: How does the paradigm of reinforcement learning (RL) fit within the broader context of machine learning?

Solution: It is instructive to compare/contrast reinforcement learning with supervised learning. In this way, it will be seen that RL can in fact be viewed as a generalization of SL:

In SL, the training feature vectors $\mathbf x_i$ manifest as the state vector $\mathbf x_t$ of an agent (similar to how it’s okay to conflate the position vector $\mathbf x(t)$ of a particle in classical mechanics with the particle itself, there’s no harm in conflating the state vector $\mathbf x_t$ with the agent itself, see map-territory relation). The difference is purely a perspective shift reflected in the choice of subscript variable $i=1,2,3,…$ vs. $t=0,1,2,…$.
In SL, the feature vectors $\mathbf x_i$ are i.i.d. draws from the same underlying data-generating distribution. In RL, the initial state vector $\mathbf x_{t=0}$ is drawn from an initial distribution, and subsequent state vectors $\mathbf x_t$ are dependent on the previous state $\mathbf x_{t-1}$ (discrete-time Markov chain).
In SL, the model output $\hat y(\mathbf x|\boldsymbol{\theta})$ is analogous to an action $\mathbf f$ outputted by an agent’s policy function $\pi(\mathbf f|\mathbf x,t)$.
In SL, each training feature vector $\mathbf x_i$ is labelled by its correct output $y_i$. This concept has no analog in reinforcement learning as there is no notion of a “correct” action for the agent to take.
In SL, the feedback mechanism is provided by a loss function $L(\hat y,y)$. By contrast, in RL the feedback mechanism is provided by a reward signal $r_t$.
In SL, the objective is to find optimal parameters $\boldsymbol{\theta}^*=\text{argmin}_{\boldsymbol{\theta}}\sum_{i=1}^{N_{\text{train}}}L(\hat y(\mathbf x_i|\boldsymbol{\theta}),y_i)$ minimizing the training cost over the $N_{\text{train}}$ training examples. In RL, the objective is to find an optimal agent policy $\pi^*=\text{argmax}_{\pi}\langle\sum_{t=0}^T\gamma^tr_t|\pi\rangle$ maximizing the expected return over an episode of horizon $T\leq\infty$ (which depends on if/when the agent reaches a terminal state $\mathbf x_T$).
In SL, there is often a distinction between training, cross-validation, and testing datasets, with a golden rule often being that the model should obviously not be able to see any of the testing data. In RL, it is societally acceptable to train on your test set 🙂

Problem: Imagine the set of all states $\mathbf x$, i.e. the agent’s state space. On this state space, one can impose a scalar field $v_{\pi}(\mathbf x,t)$; what is the meaning of this field? Give a simple example thereof.

Solution: The scalar field $v_{\pi}(\mathbf x,t)$ can roughly be thought of as describing how “valuable” the state $\mathbf x$ is at time $t$. More precisely, if one were to initialize an agent $\pi$ at state $\mathbf x_t:=\mathbf x$, then $v_{\pi}(\mathbf x,t)$ is the expected return:

\[v_{\pi}(\mathbf x,t):=\biggr\langle\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\biggr|\mathbf x,\pi\biggr\rangle\]

Consider a $\pi$-creature (of $3$Blue$1$Brown fame) symbolizing an agent with policy $\pi$. This $\pi$-creature has to complete a maze. Within the RL framework, this can be modelled as an MDP in which each square is one possible state $\mathbf x$ of the agent, the maze structure defines the allowed actions $\mathbf f$ from each state $\mathbf x$, and one can work with an undiscounted $\gamma=1$ return in which the reward is $r_t=-1$ for each action taken. Thus, the optimal policy $\pi^*$ is a time-independent, deterministic policy $\mathbf f=\pi^*(\mathbf x)$ that gets the $\pi$-creature from its initial state to the terminal state in the fewest number of moves. With that in mind, one can label on top of each square $\mathbf x$ its corresponding optimal value $v_{\pi^*}(\mathbf x)$ (note, this image was generated using Nano Banana Pro, some of the calculated values are just wrong but the idea is clear):

Problem: The value function $v_{\pi}(\mathbf x,t)$ is passive; it simply assess how good/bad it is to be at $\mathbf x_t:=\mathbf x$. In order to remedy this, one can look at the quality function $q_{\pi}(\mathbf x,\mathbf f,t)$; explain how this takes on a more active role compared to the passive nature of the value function $v_{\pi}(\mathbf x,t)$.

Solution: Because $q_{\pi}(\mathbf x,\mathbf f,t)$ assesses the quality of the agent choosing the action $\mathbf f_t:=\mathbf f$ while in state $\mathbf x_t:=\mathbf x$. That is, it is a more refined conditional expected return:

\[q_{\pi}(\mathbf x,\mathbf f,t)=\biggr\langle\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\biggr|\mathbf x,\mathbf f,\pi\biggr\rangle\]

Problem: How come the agent’s policy, value function and quality function are sometimes respectively written as $\pi(\mathbf f|\mathbf x),v_{\pi}(\mathbf x)$ and $q_{\pi}(\mathbf x,\mathbf f)$ without the $t$ argument?

Solution: There could be $3$ reasons:

Similar to the above maze example, sometimes it just doesn’t depend on $t$.
If one is seeking to maximize the agent’s expected return over an infinite horizon $T=\infty$; in such a case, it is both intuitively and mathematically clear that $\pi(\mathbf f|\mathbf x),v_{\pi}(\mathbf x)$ and $q_{\pi}(\mathbf x,\mathbf f)$ should all be invariant with respect to time translations.
If the horizon is finite $T<\infty$, then it is standard to package $\mathbf x$ and $t$ together, thereby working with an “augmented” state vector $\mathbf x’:=(\mathbf x,t)$. Then everything can be made to look identical to the infinite horizon $T=\infty$ case by replacing $\mathbf x\mapsto\mathbf x’$.

(cf. distinction between state variables and path variables in thermodynamics).

Problem: State and prove the law of iterated expectation.

Solution: If $X,Y$ are any random variables, then:

\[\langle X\rangle=\langle\langle X|Y\rangle\rangle\]

The proof is easy as long as one is careful to interpret all the expectations correctly. For instance, $\langle X|Y\rangle$ is not a scalar but a random variable with respect to $Y$:

\[\langle X|Y\rangle=\int dx p(x|Y) x\]

Thus, it is clear that the outer expectation must also be with respect to $Y$ alone:

\[\langle\langle X|Y\rangle\rangle=\int dy p(y)\langle X|y\rangle=\int dxdy p(y)p(x|y)x\]

Rewriting in terms of the joint distribution $p(x,y)=p(y)p(x|y)$ and integrating out $\int dy p(x,y)=p(x)$, one finally obtains:

\[=\int dx p(x) x=\langle X\rangle\]

Of course, this also generalizes easily to identities such as the following equality of random variables (w.r.t. $Z$):

\[\langle X|Z\rangle=\langle\langle X|Y,Z\rangle|Z\rangle\]

Problem: Show that the value function $v_{\pi}(\mathbf x)$ satisfied the following identity (to be used as a lemma later):

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

Solution: Denoting the return random variable by:

\[R_t:=\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\]

The fundamental observation is that $R_t$ obeys a simple recurrence relation:

\[R_t=r_{t+1}+\gamma R_{t+1}\]

Taking suitable conditional expectations thereof:

\[\langle R_t|\mathbf x_t,\pi\rangle=\langle r_{t+1}|\mathbf x_t,\pi\rangle+\gamma\langle R_{t+1}|\mathbf x_t,\pi\rangle\]

The LHS is the definition of $v_{\pi}(\mathbf x_t)$. Applying the law of iterated expectation to the $2^{\text{nd}}$ term on the RHS:

\[\langle R_{t+1}|\mathbf x_t,\pi\rangle=\langle\langle R_{t+1}|\mathbf x_{t+1},\mathbf x_t,\pi\rangle|\mathbf x_t,\pi\rangle\]

The Markov property ensures that $\langle R_{t+1}|\mathbf x_{t+1},\mathbf x_t,\pi\rangle=\langle R_{t+1}|\mathbf x_{t+1},\pi\rangle=v_{\pi}(\mathbf x_{t+1})$. One thus obtains the desired result:

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

Problem: How are $v_{\pi}(\mathbf x_t)$ and $q_{\pi}(\mathbf x_t,\mathbf f_t)$ related to each other?

Solution: Apply the law of iterated expectations:

\[v_{\pi}(\mathbf x_t):=\langle R_t|\mathbf x_t,\pi\rangle=\langle\langle R_t|\mathbf x_t,\mathbf f_t,\pi\rangle|\mathbf x_t,\pi\rangle=\langle q_{\pi}(\mathbf x_t,\mathbf f_t)|\mathbf x_t,\pi\rangle\]

Since $\mathbf x_t$ is being conditioned upon, the expectation must be over $\mathbf f_t$ so one can explicitly write it as:

\[=\sum_{\mathbf f_t}p(\mathbf f_t|\mathbf x_t,\pi)q_{\pi}(\mathbf x_t,\mathbf f_t)\]

But trivially $p(\mathbf f_t|\mathbf x_t,\pi)=\pi(\mathbf f_t|\mathbf x_t)$. Thus, one has an expression for $v_{\pi}(\mathbf x_t)$ in terms of $q_{\pi}(\mathbf x_t,\mathbf f_t)$. Now one would like to proceed the other way, relating $q_{\pi}(\mathbf x_t,\mathbf f_t)$ to $v_{\pi}(\mathbf x_t,\mathbf f_t)$. This can be achieved by fleshing out the expectation from the earlier lemma:

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

\[=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\pi)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

Further expand $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\pi)=\sum_{\mathbf f_t}p(\mathbf f_t|\mathbf x_t,\pi)p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t,\pi)$ and recognize again $p(\mathbf f_t|\mathbf x_t,\pi)=\pi(\mathbf f_t|\mathbf x_t)$ and $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t,\pi)=p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)$ because the action $\mathbf f_t$ is already fixed. Moving the summation $\sum_{\mathbf f_t}$ to the outside so as to compare with the earlier identity expressing $v_{\pi}(\mathbf x_t)$ in terms of $q_{\pi}(\mathbf x_t,\mathbf f_t)$, one can pattern-match:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

(which in hindsight is manifest…)

Problem: Hence, deduce the Bellman equations for the value and quality functions.

Solution: To obtain the Bellman equation for $v_{\pi}(\mathbf x_t)$, substitute the above expression for $q_{\pi}(\mathbf x_t,\mathbf f_t)$ into the expression for $v_{\pi}(\mathbf x_t)$:

\[v_{\pi}(\mathbf x_t)=\sum_{\mathbf f_t}\pi(\mathbf f_t|\mathbf x_t)\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

Similarly, to the get the Bellman equation for $q_{\pi}(\mathbf x_t,\mathbf f_t)$ substitute vice versa:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)\left(r_{t+1}+\gamma\sum_{\mathbf f_{t+1}}p(\mathbf f_{t+1}|\mathbf x_{t+1},\pi)q_{\pi}(\mathbf x_{t+1},\mathbf f_{t+1})\right)\]

Intuitively, the Bellman equations relate the value of every state with the values of states that can be transitioned into, thereby providing a system of simultaneous equations that in principle can be solved to deduce all state values. A key caveat for this dynamic programming approach to work is that one must have complete knowledge of the environment dynamics $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)$ (similar remarks apply to the qualities of different state-action pairs).

Problem: The above Bellman equations apply to a generic agent policy $\pi$; how do they specialize to the case of an optimal policy $\pi^*$?

Solution: Reverting back to the form of the “pre-Bellman” equations and using the key insight that the optimal policy $\pi^*(\mathbf f_t|\mathbf x_t)=\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q_{\pi^*}(\mathbf x_t,\mathbf f_t)}$ should only assign non-zero probability to actions of the best quality, one has:

\[v_{\pi^*}(\mathbf x_t)=\text{max}_{\mathbf f_t}q_{\pi^*}(\mathbf x_t,\mathbf f_t)\]

\[q_{\pi^*}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi^*}(\mathbf x_{t+1}))\]

which lead to the Bellman optimality equations.

Problem: In light of Bellman optimality, discuss how policy evaluation and policy improvement work. Hence, explain the notion of generalized policy iteration (GPI).

Solution: Imagine the space of all $(\pi,v)$ pairs, where $\pi(\mathbf f_t|\mathbf x_t)$ is any policy distribution and $v(\mathbf x_t)$ is the value function of some policy. Within this space, there are $2$ canonical subspaces. One is the set $(\pi,v_{\pi})$ of all pairs where $v_{\pi}$ is specifically the value function associated to policy $\pi$, and the other subspace $(\pi_v,v)$ where $\pi_v(\mathbf f_t|\mathbf x_t):=\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q(\mathbf x_t,\mathbf f_t)}$ is greedy with respect to the value function $v$, where $q(\mathbf x_t,\mathbf f_t)=\langle r_{t+1}+\gamma v(\mathbf x_{t+1})|\mathbf x_t,\mathbf f_t\rangle$ (in particular, notice $q\neq q_{\pi}$; the policy $\pi$ does not appear in the definition of $q$).

Policy evaluation is an algorithm for projecting $(v,\pi)\mapsto (v_{\pi},\pi)$. The idea is to view $v$ as a random initial guess for the underlying true policy value function $v_{\pi}$ (though any terminal states $\mathbf x_T$ should be initialized $v(\mathbf x_T):=0$). Then, sweeping through each state $\mathbf x_t$, one updates the value of $v(\mathbf x_t)$ using the Bellman equation involving the (known) randomly initialized values of the other states:

\[v(\mathbf x_t)\mapsto\sum_{\mathbf f_t}\pi(\mathbf f_t|\mathbf x_t)\sum_{\mathbf x_{t+1},r_{t+1}}p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)(r+\gamma v(\mathbf x_{t+1}))\]

With sufficiently many sweeps over the state space, general theorems guarantee convergence $v\to v_{\pi}$.

Policy improvement is an algorithm for projecting onto the other canonical subspace $(v,\pi)\mapsto (v,\pi_v)$. Essentially just act greedy:
\[\pi(\mathbf f_t|\mathbf x_t)\mapsto\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q_{\pi}(\mathbf x_t,\mathbf f_t)}\]
and is rigorously justified by the policy improvement theorem.

Policy iteration roughly amounts to alternating between the $2$ steps of policy evaluation and policy improvement; however there is some leeway in how one does this, for instance one need not necessarily run the policy evaluation step to completion but rather just perform a single sweep over the state space each time (this more “generalized” version of policy iteration is called value iteration, and can be shown to eventually still funnel/converge onto the optimal policy $\pi^*$ as one can prove that $\pi$ is greedy with respect to its own value function $v_{\pi}$ iff $\pi=\pi^*$).

Problem: Distinguish between model-free vs. model-based methods/algorithms in reinforcement learning.

Solution: In both cases, the fundamental limitation (which previously was taken for granted) is incomplete knowledge of the environment dynamics $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$; instead, one can only sample trajectories $\mathbf x_{t=0},\mathbf f_{t=0},r_{t=0},\mathbf x_{t=1},\mathbf f_{t=1},r_{t=1},…$ when running a policy $\pi$ through the MDP. Recalling the fundamental objective of RL is to compute $\pi^*=\text{argmax}_{\pi}\langle R_t|\pi\rangle$, a model-free method for doing so is one which does not attempt to estimate $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$ (here $p$ is what the word “model” refers to, i.e. a model of the environment dynamics) whereas model-based methods do attempt to estimate $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$ and thereby obtain $\pi^*$ through GPI as outlined earlier (see the OpenAI Spinning Up documentation for a nice diagram of this taxonomy and further comments).

Problem: Explain how Monte Carlo evaluation and Monte Carlo control are both model-free methods that accomplish the same actions as their respective model-based cousins, i.e. policy evaluation and policy improvement.

Solution: Fix $\pi$. The key is to recognize that $v_{\pi}(\mathbf x_t)$ is an expectation of the return random variable starting at $\mathbf x_t$, and Monte Carlo point estimators can always be applied to approximate expectation values:

\[v_{\pi}(\mathbf x_t)=\langle R_t|\mathbf x_t,\pi\rangle\approx\frac{1}{N_{\mathbf x_t}}\sum_{n=1}^{N_{\mathbf x_t}}R_n\]

where $N_{\mathbf x_t}$ is the number of times the MDP passes through state $\mathbf x_t$ in however many trajectories, and $R_n$ is the return following the $n^{\text{th}}$ pass through $\mathbf x_t$ within some MDP trajectory.

Lemma: the arithmetic mean $\hat{\mu}_n:=\frac{1}{n}\sum_{i=1}^nx_i$ obeys the recurrence relation $\hat{\mu}_n=\hat{\mu}_{n-1}+\frac{1}{n}(x_n-\hat{\mu}_{n-1})$

Proof: trivial, but instructive to highlight the convex linear combination nature of the recurrence $\hat{\mu}_n=\left(1-\frac{1}{n}\right)\hat{\mu}_{n-1}+\frac{1}{n}x_n$, so for large $n\gg 1$, the majority $1-1/n$ of the weight is given to the prior value of the mean $\hat{\mu}_{n-1}$, only a tiny sliver $1/n$ is given to the new sample $x_n$.

Thus, rather than waiting for a bunch of trajectories to finish before applying the Monte Carlo estimate above, one can introduce an intermediate estimator initialized at $\hat v_{\pi}(\mathbf x_t):=0$ that applies (at the termination of each trajectory) the following iterative rule for each time $n=1,2,…,N_{\mathbf x_t}$ that $\mathbf x_t$ was visited:

\[\hat v_{\pi}(\mathbf x_t)\mapsto\hat v_{\pi}(\mathbf x_t)+\alpha_n(R_n-\hat v_{\pi}(\mathbf x_t))\]

where $\alpha_n:=1/n$. Since this is just a mathematical decomposition of the original Monte Carlo estimate, it is an equally valid model-free method for $v_{\pi}$ estimation. Now comes a curveball: one common modification is to replace $\alpha_n\mapsto\alpha$ with an $n$-independent constant $\alpha\in[0,1]$ (this is called constant-$\alpha$ MC).

If one were presented with the recurrence relation $\hat{\mu}_n=\hat{\mu}_{n-1}+\frac{1}{n}(x_n-\hat{\mu}_{n-1})$ and initial condition $\hat{\mu}_0=0$, one could claim the explicit solution to be $\hat{\mu}_n=\frac{1}{n}\sum_{i=1}^nx_i$. However, if instead one were presented with the recurrence relation $\hat{\mu}’_n=\hat{\mu}’_{n-1}+\alpha(x_n-\hat{\mu}’_{n-1})$ and same initial condition $\hat{\mu}’_0=0$, what would be the solution for $\hat{\mu}’_n$? It’s easy to check $\hat{\mu}’_n=\alpha\sum_{i=1}^n(1-\alpha)^{n-i}x_i$; thus, recent samples are given exponentially more weight, and it’s biased towards underestimating $\langle\hat{\mu}’_n\rangle-\mu=-\mu(1-\alpha)^n$, though is asymptotically unbiased $\lim_{n\to\infty}\langle\hat{\mu}’_n\rangle=\mu$. However, unlike the asymptotically vanishing variance $\sigma^2_{\hat{\mu}_n}=\sigma^2/n$ of the standard mean estimator $\hat{\mu}_n$, the variance $\sigma^2_{\hat{\mu}’_n}=\frac{\alpha\sigma^2}{2-\alpha}(1-(1-\alpha)^{2n})\to\frac{\alpha\sigma^2}{2-\alpha}\approx\alpha\sigma^2/2$ of $\hat{\mu}’_n$ is bounded from below, hence the estimate $\hat{\mu}’_n$ will forever oscillate about $\mu$. It should make sense $\sigma^2_{\hat{\mu}’_n}\propto\alpha$, i.e. a larger step size $\alpha$ means bigger oscillations, though it also means faster convergence of $\hat{\mu}’_n\to\mu$ because of the $(1-\alpha)^n$.

This all begs the question: why bother with constant-$\alpha$ MC? (come back to this…)

Having used constant-$\alpha$ MC to estimate $v_{\pi}(\mathbf x_t)$ in a model-free way for some fixed agent policy $\pi$, one can naturally ask about model-free estimation of $q_{\pi}(\mathbf x_t,\mathbf f_t)$. It turns out to just be:

\[\hat q_{\pi}(\mathbf x_t,\mathbf f_t)\mapsto\hat q_{\pi}(\mathbf x_t,\mathbf f_t)+\alpha(R_n-\hat q_{\pi}(\mathbf x_t,\mathbf f_t))\]

Conceptually, the standard way to convince oneself of this is that because $\pi$ was fixed, it could be thought of as part of the incomplete knowledge of the environment dynamics $p$ so that the agent never takes any actions but just lets the environment push its state around. This perspective shift from MDP trajectories $\mathbf x_{t=0},\mathbf f_{t=0},r_{t=0},\mathbf x_{t=1},\mathbf f_{t=1},r_{t=1},…$ to MRP trajectories $\mathbf x_{t=0},r_{t=0},\mathbf x_{t=1},r_{t=1},…$ amounts to viewing the MRP states $\mathbf x_t$ as corresponding to augmented state-action pairs $(\mathbf x_t,\mathbf f_t)$ in the MDP. But then it is mathematically apparent that $v_{\pi}(\mathbf x_t)$ in the MRP is just $q_{\pi}(\mathbf x_t,\mathbf f_t)$ in the MDP, so constant-$\alpha$ MC works for $q_{\pi}$-evaluation! Thus, Monte Carlo control chooses to be greedy with respect to $q_{\pi}$, and since $q_{\pi}$ could be estimated in a model-free way, this therefore implements model-free policy improvement.

Problem: Explain the exploration-exploitation tradeoff in reinforcement learning and how a soft policy $\pi$ can help to address this.

Solution: Being greedily exploiting high value states all the time could lead to the agent settling into a “locally” optimal but not globally optimal policy $\pi^*$; the only way to know for sure is to explore all state-action sequences, for instance by ensuring that one’s policy is soft $\pi(\mathbf f_t|\mathbf x_t)>0$ for all states $\mathbf x_t$ and actions $\mathbf f_t$. A crude way to guarantee softness during Monte Carlo control is to use an $\varepsilon$-greedy policy (a better name would be either a $(1-\varepsilon)$-greedy policy or an $\varepsilon$-soft policy), though there are other more refined ways to do so.

Problem: Having distinguished between model-based and model-free methods in RL, one can make an orthogonal distinction between on-policy and off-policy methods. Explain and motivate this taxonomy.

Solution: The constant-$\alpha$ MC $q_{\pi}$-evaluation followed by $\varepsilon$-greedy MC control is a simple model-free, on-policy RL algorithm for discovering $\pi^*$. This is because the trajectory samples used to evaluate $q_{\pi}$ came from the same policy $\pi$ that is subsequently being improved. However, in practice the trajectory samples may come from some (known) behavior policy $\pi_b$ even though the goal is to evaluate $q_{\pi}$ so as to improve another, different (known) target policy $\pi\neq\pi_b$. This is where off-policy methods come in (cf. the distinction between on-shell vs. off-shell particles in physics). The obvious way to bridge $\pi_b$ and $\pi$ is via importance sampling:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\langle R_t|\mathbf x_t,\mathbf f_t,\pi\rangle\]

\[=\frac{1}{N_{\mathbf x_t}}\sum_{n=1}^{N_{\mathbf x_t}}\frac{\pi(\mathbf f_{t+1}|\mathbf x_{t+1})…\pi(\mathbf f_T|\mathbf x_T)}{\pi_b(\mathbf f_{t+1}|\mathbf x_{t+1})…\pi_b(\mathbf f_T|\mathbf x_T)}R_n\]

where it is essential to emphasize that the trajectories in this sum are sampled under the behavior policy $\pi_b$! This is model-free because all the transition probabilities have cancelled out in the importance sampling quotient.

Posted in Blog | Leave a comment

Density Functional Theory

Posted on January 28, 2026 by wdengquantum.me

Problem: In one sentence, what is the essence of DFT?

Solution: To replace $\Psi\mapsto n$, where the number density of a system of $N$ identical quantum particles (usually electrons) $n(\mathbf x)$ is:

\[n(\mathbf x):=N\int d^3\mathbf x_2…d^3\mathbf x_N |\Psi(\mathbf x,\mathbf x_2,…,\mathbf x_N)|^2\]

in terms of its $N$-body wavefunction $\Psi(\mathbf x,\mathbf x_2,…,\mathbf x_N)$. Thus, $\int d^3\mathbf x n(\mathbf x)=N$ is normalized.

Problem: Write down the non-relativistic Hamiltonian $H$ of a molecule. Attempt to express the expected energy $E=\langle\Psi|H|\Psi\rangle$ in many-body state $|\Psi\rangle\in\bigwedge^NL^2(\mathbf R^3\to\mathbf C)\otimes\mathbf C^2$ in terms of $n(\mathbf x)$ defined above; this functional $E=E[n]$ of the density $n$ explains the name “density functional theory“.

Solution: (the Born-Oppenheimer approximation is implicitly used whereby the electrons move on an effective/external potential energy surface $V_{\text{ext}}(\mathbf x)$ defined by a collection of stationary nuclei)

\[H=\sum_{i=1}^N\frac{|\mathbf P_i|^2}{2m}+V_{\text{ext}}(\mathbf X_i)+\frac{1}{2}\sum_{1\leq i\neq j\leq N}\frac{\alpha\hbar c}{|\mathbf X_i-\mathbf X_j|}\]

One can check that:

\[E=\frac{\hbar^2}{2m}\sum_{i=1}^N\int d^3\mathbf x_1…d^3\mathbf x_N\biggr|\frac{\partial\Psi}{\partial\mathbf x_i}\biggr|^2+\int d^3\mathbf x n(\mathbf x)V_{\text{ext}}(\mathbf x)+\frac{\alpha\hbar c}{2}\int d^3\mathbf x d^3\mathbf x’\frac{n_2(\mathbf x,\mathbf x’)}{|\mathbf x-\mathbf x’|}\]

where $n_2(\mathbf x,\mathbf x’)=N(N-1)\int d^3\mathbf x_3…d^3\mathbf x_N|\Psi(\mathbf x,\mathbf x’,\mathbf x_3,…,\mathbf x_N)|^2$ so in particular $(N-1)n(\mathbf x)=\int d^3\mathbf x’ n_2(\mathbf x,\mathbf x’)$.

Problem: Write down the Thomas-Fermi approximation for $E[n]$, and explain why it sucks. Also derive the von Weizsäcker correction to the Thomas-Fermi energy functional.

Solution: The Thomas-Fermi functional treats the electrons as a locally ideal Fermi gas, thus having the usual Pauli pressure $p\propto n^{5/3}$ in $\mathbf R^3$ and thus (kinetic) energy density $3p/2$ scaling likewise. Furthermore, the exchange-correlation functional is taken to vanish $E_{\text{xc}}[n]=0$ (thus, Thomas-Fermi is often said to be semiclassical):

\[E_{\text{TF}}[n]=\frac{3\hbar^2(3\pi^2)^{2/3}}{10m}\int d^3\mathbf x n^{5/3}(\mathbf x)+\int d^3\mathbf x n(\mathbf x)V_{\text{ext}}(\mathbf x)+\frac{\alpha\hbar c}{2}\int d^3\mathbf x d^3\mathbf x’\frac{n(\mathbf x)n(\mathbf x’)}{|\mathbf x-\mathbf x’|}\]

One has:

\[\frac{\delta}{\delta n(\mathbf x)}\left(E_{\text{TF}}-\mu\left(\int d^3\mathbf xn(\mathbf x)-N\right)\right)=\frac{\hbar^2k^2_F(\mathbf x)}{2m}+V_{\text{eff}}(\mathbf x)-\mu=0\]

where $k_F(\mathbf x):=(3\pi^2n(\mathbf x))^{1/3}$ and $V_{\text{eff}}(\mathbf x)=V_{\text{ext}}(\mathbf x)+\alpha\hbar c\int d^3\mathbf x’\frac{n(\mathbf x’)}{|\mathbf x-\mathbf x’|}$.

(aside: although the Thomas-Fermi model is meant to be applied to molecules, a result of Teller showed that $E_{\text{TF molecule}}>2E_{\text{TF atom}}$ so Thomas-Fermi theory predicts the molecular state is unstable with respect to dissociation. In a central external potential such as due to a single nucleus $V_{\text{ext}}(r)=-Z\alpha\hbar c/r$, the effective potential $V_{\text{eff}}(r)$ and electron density $n(r)$ are both isotropic as well. Starting from Poisson’s equation:

\[\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\frac{V_{\text{eff}}(r)}{-e}=-\frac{-en(r)}{\varepsilon_0}\]

with $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2f(r)=\frac{1}{r^2}\frac{d}{dr}\left(r^2\frac{df}{dr}\right)=\frac{1}{r}\frac{d^2}{dr^2}(rf(r))$ (use the latter) one can check that by defining $f(r):=\left(\frac{8}{3\pi a_0}\right)^2rk^2_F(r)$, one recovers the Thomas-Fermi ODE:

\[\frac{d^2 f}{dr^2}=\frac{f^{3/2}(r)}{\sqrt{r}}\]

The von Weizsäcker term is a correction to the kinetic energy functional in the Thomas-Fermi model that accounts for density gradients:

\[T_{\text{vW}}[n]=\frac{\hbar^2}{8m}\int d^3\mathbf x\frac{|\partial n/\partial\mathbf x|^2}{n(\mathbf x)}\]

Heuristically, it can be justified by writing the expected kinetic energy of a single electron $\frac{\hbar^2}{2m}\int d^3\mathbf x |\partial\psi/\partial\mathbf x|^2$ with wavefunction $\psi(\mathbf x)\in\mathbf R$ and replacing $\psi(\mathbf x)\mapsto\sqrt{n(\mathbf x)}$.

Problem: Prove the two Hohenberg-Kohn theorems.

Solution: First, it is useful to prove:

Lemma: If $V_{\text{ext}}(\mathbf x)\neq V’_{\text{ext}}(\mathbf x)$ represent $2$ physically distinct external potentials (i.e. they differ by more than just some additive constant), then the corresponding molecular Hamiltonians $H\neq H’$ do not share any non-zero eigenstates.
Proof: Suppose for sake of contradiction that there exists $|\Psi\rangle\neq 0$ such that
\[H|\Psi\rangle=E|\Psi\rangle\]
\[H’||\Psi\rangle=E’|\Psi\rangle\]
Subtracting, one obtains:
\[(H’-H)|\Psi\rangle=(E’-E)|\Psi\rangle\]
Inspecting the form of the molecular Hamiltonian $H$ given earlier, one sees that the “universal” part (i.e. the kinetic energy and electron-electron Coulomb repulsion terms) cancel, leaving only the non-universal residue $H’-H=\sum_{i=1}^NV’_{\text{ext}}(\mathbf X_i)-V_{\text{ext}}(\mathbf X_i)$. But the RHS $E’-E$ is $\mathbf X_i$-independent for all $i=1,…,N$; this means $V’_{\text{ext}}(\mathbf x)-V_{\text{ext}}(\mathbf x)=\text{const}$ which is a contradiction.

With this in mind, the first HK theorem follows by specializing to the respective (distinct by the above lemma) ground states $|\Psi_0\rangle\neq|\Psi’_0\rangle$ of $H$ and $H’$. Apply the variational bound both ways:

\[\]

Problem: Demonstrate the first Hohenberg-Kohn theorem for $2$ electrons in a $1$-dimensional harmonic trap.

Solution:

Problem: Derive the Kohn-Sham equations.

Solution: Fictitious ideal electron gas with density $n(\mathbf x)$.

Problem:

Posted in Blog | Leave a comment

Quantumplations

Diffusion Models

Information Geometry

Autoencoders

Graph Neural Networks

Renormalization Group

Convolutional Neural Networks

Hamilton’s Optics-Mechanics Analogy

Pseudo-Riemannian Geometry

Reinforcement Learning (Part \(1\))

Density Functional Theory

Archives

Meta