Machine Learning Interatomic Potentials

Posted on July 15, 2026 by wdengquantum.me

Problem: Explain what an MLIP is.

Solution: A better name might have been MLPES (machine-learned potential energy surface), but MLIP (machine-learned interatomic potential) is the standard term (it should not be confused with e.g. a pairwise interaction potential like the Lennard-Jones potential).

That is, within the Born-Oppenheimer approximation, the ground state PES is just an effective potential $V(\mathbf x_1,…,\mathbf x_N)$ defined on the nuclear configuration space $(\mathbf x_1,…,\mathbf x_N)\in\mathbf R^{3N}$ on which the nuclei dynamics are adiabatic. Treating the nuclei as classical point masses, it follows that the $i^{\text{th}}$ nucleus experiences a conservative force field $\mathbf F_i(\mathbf x_1,…,\mathbf x_N)=-\partial V/\partial\mathbf x_i$.

An MLIP is then just some parametric neural network ansatz seeking to approximate $\hat V(\mathbf x_1,…,\mathbf x_N|\boldsymbol{\theta})\approx V(\mathbf x_1,…,\mathbf x_N)$. The corresponding MLFF (machine-learned force field) is thus $\hat{\mathbf F}_i(\mathbf x_1,…,\mathbf x_N|\boldsymbol{\theta}):=-\partial\hat{V}/\partial\mathbf x_i$ (note: the nuclear charges $Z_1,…,Z_N\in\mathbf Z^+$ of the $N$ nuclei are also considered inputs to the MLIP, and for more sophisticated MLIPs one may also input spin information, a global charge, etc.).

To be more precise, not just any estimator $\hat V(\mathbf x_1,…,\mathbf x_N|Z_1,…,Z_N,\boldsymbol{\theta})$ counts as an MLIP. There are certain fundamental requirements (the LAS trio: locality, analyticity, symmetry) that must be fulfilled:

Locality: $\hat V(\mathbf x_1,…,\mathbf x_N|\boldsymbol{\theta})=\sum_{i=1}^N\hat V_i$ can be decomposed as a sum of local potentials experienced by each of the $N$ nuclei, that is $\hat{V}_i$ depends only on the nuclei $\{\mathbf x_j:r_{ij}<r^*\}$ within a radius $r^*$ of $\mathbf x_i$ (strictly speaking, the locality assumption may be relaxed if long-range electrostatics matters). By the triangle inequality, a corollary of this is that $r_{ij}>2r^*\Rightarrow\partial\hat{\mathbf F}_i/\partial\mathbf x_j=0$.
Analyticity: $\hat V(\mathbf x_1,…,\mathbf x_N|\boldsymbol{\theta})$ is smooth with respect to the nuclear coordinates $(\mathbf x_1,…,\mathbf x_N)\in\mathbf R^{3N}$, so in particular the forces $\mathbf F_i$ actually exist! (it should of course also be suitably smooth with respect to the parameters $\boldsymbol{\theta}$ otherwise gradients wouldn’t be able to backprop).
Symmetry: the usual Noetherian symmetries $\hat V(\mathbf x_1+\Delta\mathbf x,…,\mathbf x_N+\Delta\mathbf x|\boldsymbol{\theta})=\hat V(R\mathbf x_1,…,R\mathbf x_N|\boldsymbol{\theta})=\hat V(\mathbf x_1,…,\mathbf x_N|\boldsymbol{\theta})$ for $\Delta\mathbf x\in\mathbf R^3$ and $R\in SO(3)$ and discrete symmetries $\hat V(\mathbf x_{\sigma(1)},…,\mathbf x_{\sigma(N)},Z_{\sigma(1)},…,Z_{\sigma(N)}|\boldsymbol{\theta})=\hat V(-\mathbf x_1,…,-\mathbf x_N,Z_1,…,Z_N|\boldsymbol{\theta})=\hat V(\mathbf x_1,…,\mathbf x_N,Z_1,…,Z_N|\boldsymbol{\theta})$ for $\sigma\in S_N$.

Problem: Explain how the atomic cluster expansion (ACE) enforces the LAS trio.

Solution: Starting from the ansatz of L, it is clear that A and S can be fulfilled provided each of the $N$ local potentials $\hat V_i$ are analytic and symmetric in their neighbouring nuclear coordinates $\{\mathbf x_j:r_{ij}<r^*\}$.

\[\hat V_i=\hat V^{(1)}(Z_i)+\sum_{\langle j,i\rangle}\hat V^{(2)}(\mathbf r_{ij})+\sum_{\langle j<k,i\rangle}\hat V^{(3)}(\mathbf r_{ij},\mathbf r_{ik})+…\]

\[=\text{Point}+\text{Lines}+\text{Triangles}+\text{Pyramids}+…\]

Explicit models of such many-body interactions include the $2$-body Lennard-Jones potential:

\[V^{(2)}(r)=-V_0\left(2\left(\frac{r_0}{r}\right)^6-\left(\frac{r_0}{r}\right)^{12}\right)\]

the $3$-body Stillinger-Weber potential:

\[V^{(3)}(r_1,r_2,\theta)\sim V_0 e^{\sim r_0\left(\frac{1}{r_1-r^*}+\frac{1}{r_2-r^*}\right)}(\cos\theta-\cos\theta_0)^2[r_1<r^*][r_2<r^*]\]

the $4$-body potential:

\[V^{(4)}(\mathbf x_1,\mathbf x_2,\mathbf x_3)\sim\left(\frac{\mathbf x_1\cdot(\mathbf x_2\times\mathbf x_3)}{r_1r_2r_3}\right)^2\]

etc.

To compute the $(\nu+1)$-body interaction term $\hat V^{(\nu+1)}(\mathbf r_1,…,\mathbf r_{\nu})$ in ACE (the use of the symbol $\nu$ as the Greek equivalent of the English letter $n$ is meant to evoke the concept of “neighbours”), one has to sum ${\nu_i}\choose{\nu}$$\in O(\nu_i^{\nu})$ interaction terms which is polynomial in the number $\nu_i:=|\{\mathbf x_j:r_{ij}<r^*\}|$ of $r^*$-neighbours associated to the $i^{\text{th}}$ atom, but exponential in the number of simultaneous $r^*$-neighbours $\nu\leq\nu_i$.

Fortunately, there exists a simple trick that dramatically simplifies this computation from $O(\nu_i^{\nu})\to O(\nu_i)$ time. The trick consists of $2$ steps.

Step #$1$: given an arbitrary single-particle basis $\psi_{n\ell m}(\mathbf r):=R_{n}(r)Y_{\ell}^m(\hat{\mathbf r})$, expand $\hat V^{(\nu+1)}(\mathbf r_1,…,\mathbf r_{\nu})$ in that basis:

\[\hat V^{(\nu+1)}(\mathbf r_1,…,\mathbf r_{\nu})=\sum_{}\]

Step #$2$:

via a trivial factorization:

Lemma: given an arbitrary collection of $\nu$ scalar fields $\psi_1(\mathbf r),…\psi_{\nu}(\mathbf r)$, one can efficiently evaluate separable sums of the form $\sum_{j_1,…,j_{\nu}}\psi_1(\mathbf r_{j_1})…\psi_{\nu}(\mathbf r_{j_{\nu}})$ by factorizing them (aka inverting the distributive property!) as:

\[\sum_{j_1,…,j_{\nu}}\psi_1(\mathbf r_{j_1})…\psi_{\nu}(\mathbf r_{j_{\nu}})=\sum_j\psi_1(\mathbf r_j)…\sum_j\psi_{\nu}(\mathbf r_j)\]

Problem: One prominent example of an MLIP is MACE, developed by the group of Gábor Csányi. Walk through a single forward pass of MACE at inference time. Then, explain how MACE is trained.

Solution: MACE is an MLIP built on a variant of the MPNN architecture subclass of GNNs. In addition to the standard LAS inductive biases of most MLIPs, MACE itself introduces several further innovations.

The first misconception worth clarifying is that the actual graph input to most MLIPs (including MACE) does not necessarily coincide with how its structural formula might be drawn by a chemist. Instead, the edges are defined solely by the hyperparametric cutoff radius $r^*$.

With the input undirected graph defined, the rest is just ML (no more physics/chemistry). Just as transformers begin by embedding tokens into latents, here “tokens” can be thought of as vertex features (the atomic numbers $Z_i$) and edge features (the translationally-invariant displacement vectors $\mathbf r_{ij}:=\mathbf x_i-\mathbf x_j$ between vertices). They are learnably and independently embedded as latents:

i) The vertex features $Z_i$ are embedded by first one-hot encoding $Z_i\mapsto\hat{\mathbf e}_i$ w.r.t. the size of the periodic table “vocabulary”. These unit one-hots are then projected $\hat{\mathbf e}_i\mapsto\mathbf h_i\in\mathbf R^d$ via a linear layer into a latent space of some hyperparametric hidden dimension $d$ ($d=128$ is the MACE default). This latent $\mathbf h_i$ is then extruded along a further dimension of size $\ell^*+1$ indexed by $\ell=0,1,…,\ell^*$, and for each such $\ell$ a further extrusion along a dimension of size $2\ell+1$ indexed by $m=-\ell,…,\ell$ results in a feature “double-staircase” whose “steps” $\ell$ encode the $SO(3)$ irreps:

A staircase $(\mathbf h_i)_{k,(\ell,m)}$ does not have a well-defined tensor shape, so in practice one concatenates along the $m$-axis $\sum_{\ell=0}^{\ell^*}(2\ell+1)=(\ell^*+1)^2$ to represent the staircase as a PyTorch tensor of shape $(d,(\ell^*+1)^2)$.

ii) The edge features $\mathbf r_{ij}$ factorize into magnitude and direction features $\mathbf r_{ij}=r_{ij}\hat{\mathbf r}_{ij}$ which lead directly to a correspondingly-factorized scalar latent $h_{ij}:=R(r_{ij})Y_{\ell}$. The magnitude/radial function is defined (for $n=1,…,n^*$ with MACE default $n^*=8$) by the sinc/$j_0$ spherical Bessel function $R_n(r):=\sqrt{\frac{2}{r^*}}\frac{\sin (n\pi r/r^*)}{r^*}f_p(r)$ (cf. the eigenstates $\psi_n(x)=\sqrt{\frac{2}{L}}\sin\frac{n\pi x}{L}$ for $n=1,2,…$ of a quantum particle in the box $x\in[0,L]$) modulated by a degree-$(p+2)$ polynomial envelope cutoff $f_p(r):=\left(1-\frac{(p+1)(p+2)}{2}\left(\frac{r}{r^*}\right)^p+p(p+2)\left(\frac{r}{r^*}\right)^{p+1}-\frac{p(p+1)}{2}\left(\frac{r}{r^*}\right)^{p+2}\right)[r<r^*]$ engineered to obey $f_p(0)=1$ and $f_p(r^*)=f’_p(r^*)=f^{\prime\prime}_p(r^*)=0$ (MACE default is $p=5$) and the angular factor is the usual spherical harmonic $\mathbf Y_{\ell}(\hat{\mathbf r})\in\mathbf C^{2\ell+1}$.

The rotationally equivariant property of the spherical harmonics is that they transform under a specific unitary $(2\ell+1)$-dimensional $SO(3)$-irrep $D^{(\ell)}$ known as the Wigner-$D$ representation, that is for $R\in SO(3)$:

\[\mathbf Y_{\ell}(R^{-1}\hat{\mathbf r})=D^{(\ell)}(R)\mathbf Y_{\ell}(\hat{\mathbf r})\]

where $D^{(\ell)}(R)\in U(2\ell+1)$. Intuitively, a rotation $\hat{\mathbf r}\mapsto R^{-1}\hat{\mathbf r}$ on $S^2$ induces a corresponding “rotation” $\mathbf Y_{\ell}\mapsto D^{(\ell)}(R)\mathbf Y_{\ell}$ in $\mathbf C^{2\ell+1}$.

Posted in Blog | Leave a comment

Reinforcement Learning (Part $2$)

Posted on July 6, 2026 by wdengquantum.me

Problem: Let $p(r_{t+1},s_{t+1}|s_t)$ be a Markov reward process. Compare constant-$\alpha$ Monte Carlo (MC) evaluation and $n$-step constant-$\alpha$ temporal difference (TD) evaluation.

Solution: Essentially, TD is a bootstrapped version of MC, and indeed it turns out $\lim_{n\to\infty}\text{TD}_{(n,\alpha)}=\text{MC}_{\alpha}$. That is, both TD and MC are model-free methods for evaluating the value function $V(s_t)$, but whereas MC patiently waits for each episode to reach a terminal state before the value functions of the visited states are simultaneously updated:

\[\Delta V(s_t)=\alpha(\sum_{\text{visits }i\text{ to } s_t}R_i-V(s_t))\]

$n$-step TD is more impatient and instead replaces the empirical return $R_i\mapsto r_{t+1}+…+\gamma^{n-1}r_{t+n}+\gamma^nV(s_{t+n})$ with a bootstrapped estimate depending on the estimated value $V(s_{t+n})$ (this is not to be confused with the RL definition of dynamic programming (DP) which like TD also relies on recursive bootstrapping, but unlike TD is model-based rather than model-free. Also, don’t confuse the concept of DP in RL with DP in CS which simply means recursion with memoization).

Problem: Explain why, in practice, TD outperforms MC at value function estimation.

Solution: A useful analogy is to imagine a carnival (MRP) with $2$ games (states) $A,B$. Game $A$ always returns reward $r_A=0$ but also asserts the player has to proceed to game $B$ (i.e. $p(r_{t+1}=0,s_{t+1}=B|s_t=A)=1$), and game $B$ can return reward $r_B\in\{0,1\}$ with (say) $50:50$ odds but after that the player must go home. Imagine a batch of $8$ carnival players/”episodes”, $2$ of which are initialized at game $A$, and the other $6$ are initialized at game $B$. It turns out that both of the players who begin at game $A$ then receive $r_A=0$ as expected, and then proceed to game $B$ and happen to receive $r_B=0$ and go home. On the other hand, of the $6$ players who begin at game $B$, $4$ of them get $r_B=1$ while the remaining $2$ get $r_B=0$. Assuming undiscounted return $\gamma=1$, the Monte Carlo approach would look at the $2$ episodes that passed through game $A$, and conclude that game $A$ has no value $V(A)=0$, whereas game $B$ seems to have value $V(B)=4/8=1/2$. By contrast, a TD approach would say that both games $A$ and $B$ have value $1/2$ because anyone who plays game $A$ has to play game $B$. Thus, TD is about finding the maximum likelihood for the MRP, whereas MC is about minimizing mean-squared error.

(aside: to illustrate this more clearly, consider $n=1$-step TD where $\Delta V(s_t)=\alpha(r_{t+1}+\gamma V(s_{t+1})-V(s_t))$ for each visit to state $s_t$. Then the total bootstrap on the value over all times $t$ when $s_t$ was visited would be given by: $\sum_t\Delta V(s_t)=0$ at equilibrium, so rearranging, one finds a Bellman-like equation $V(s_t)=\hat r_{t+1}+\gamma\hat{V}(s_{t+1})$. This result is sometimes called Sutton’s theorem).

Problem: TD learning is solely for policy evaluation (i.e. bootstrapping $V(s_t)$). By contrast, SARSA, expected SARSA, and $Q$-learning are GPI algorithms. Explain how each of them works.

Solution: In all $3$ cases, the proportionality constant $\alpha>0$ is a hyperparameter.

For $1$-step SARSA, the idea is that a particular instantiation $s_t\to a_t\to (r_{t+1}, s_{t+1})\to a_{t+2}$ of the MDP will lead to the $Q$-evaluation:

\[\Delta Q(s_t,a_t)\propto r_{t+1}+\gamma Q(s_{t+1},a_{t+1})-Q(s_t,a_t)\]

(which can be trivially extended to $n$-step SARSA). Expected ($1$-step, though again the generalization to $n\geq 1$ is trivial) SARSA replaces $Q(s_{t+1},a_{t+1})$ associated to whatever specific action $a_{t+1}$ was drawn from the $(1-\varepsilon)$-greedy policy $\pi(a_{t+1}|s_{t+1})$ with, as its name suggests, the expectation value $\langle Q(s_{t+1},a_{t+1})\rangle_{a_{t+1}}:=\sum_{a_{t+1}}\pi(a_{t+1}|s_{t+1})Q(s_{t+1},a_{t+1})$. Finally, ($1$-step) $Q$-learning is the special case of expected SARSA where the policy is maximally greedy ($\varepsilon:=0$) and hence $\langle Q(s_{t+1},a_{t+1})\rangle_{a_{t+1}}=\max_{a_{t+1}}Q(s_{t+1},a_{t+1})$. In all cases, armed with $Q$, one can then perform standard policy improvement by being greedy with respect to $Q$, hence closing the GPI loop.

Problem: Explain how the “MODERN” acronym provides a bird’s eye-view over the field of RL.

Solution:

M: Model-free/based

O: On/off-policy

D: Distribution

E: Evaluation

R: Representation

N: $n$-step

Problem: Explain function approximation in the context of RL.

Solution: As the name suggests, function approximation is just RL’s glorified way of saying “use a neural network”. Exactly which function one is approximating with a neural network depends on the application, but could be e.g. the value function $\hat{V}(s|\boldsymbol{\theta})$, the action function $\hat{Q}(s,a|\boldsymbol{\theta})$, the policy itself $\hat{\pi}(a|s,\boldsymbol{\theta})$, etc.

Gradient Monte Carlo approximates the true value function $V(s|\boldsymbol{\theta})$ by a neural network $\hat{V}(s|\boldsymbol{\theta})$ via stochastic gradient descent on the loss:

\[L(\boldsymbol{\theta})=\frac{1}{2N(s)}\sum_t(\hat{V}(s_t|\boldsymbol{\theta})-R_t)^2\]

Thus, gradient descent is the usual $\boldsymbol{\theta}\mapsto\boldsymbol{\theta}-\alpha\frac{\partial L}{\partial\boldsymbol{\theta}}$ with:

\[\frac{\partial L}{\partial\boldsymbol{\theta}}=\frac{1}{N(s)}\sum_t(\hat{V}(s_t|\boldsymbol{\theta})-R_t)\frac{\partial\hat{V}(s_t|\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}\]

(or some stochastic or mini-batch version).

Semi-gradient TD instead uses MSE loss against a (biased) estimator $r_{t+1}+\gamma\hat{V}(s_{t+1}|\boldsymbol{\theta})$ of the return $R_t$.

\[L(\boldsymbol{\theta})=\frac{1}{2N(s)}\sum_t(\hat{V}(s_t|\boldsymbol{\theta})-r_{t+1}-\gamma\hat{V}(s_{t+1}|\boldsymbol{\theta}))^2\]

Here, although the true gradient $\partial L/\partial\boldsymbol{\theta}$ would take into account the $\boldsymbol{\theta}$-dependence present in both $\hat{V}(s_t|\boldsymbol{\theta})$ and $\hat{V}(s_{t+1}|\boldsymbol{\theta})$, it turns out to be better to only take into account the $\boldsymbol{\theta}$-dependence present in $\hat{V}(s_t|\boldsymbol{\theta})$, hence the term “semi-gradient”:

\[\frac{\partial L}{\partial\boldsymbol{\theta}}\approx \frac{1}{N(s)}\sum_t(\hat{V}(s_t|\boldsymbol{\theta})-r_{t+1}-\gamma\hat{V}(s_{t+1}|\boldsymbol{\theta}))\frac{\partial\hat{V}(s_t|\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}\]

Problem: on-policy & off-policy control with function approximation…

Solution:

Problem: Explain how the REINFORCE algorithm works, what its flaws are, and introducing a baseline can help remedy this.

Solution: (policy gradient method)

Posted in Blog | Leave a comment

Graphics Processing Units (GPUs)

Posted on June 16, 2026 by wdengquantum.me

Problem: Define the arithmetic intensity $I$ of an algorithm/computation.

Solution: It is basically what it sounds like: the more “intense” an algorithm’s computations (“arithmetic”) are, the greater its arithmetic intensity $I$:

\[I:=\frac{W}{Q}\]

where $W$ is the algorithm’s compute/work (measured in floating-point operations (FLOPs) not to be confused with floating-point operations per second (FLOPS)) whereas $Q$ is the algorithm’s (quantity of) memory (measured in bytes). Unlike the efficiency $\eta:=W/Q_{\text{in}}<1-T_C/T_H$ of a heat engine which is bounded by the Carnot limit, here the arithmetic intensity $I$ can be too large (in which case the algorithm is compute-bound) or it can be too small (in which case the algorithm is memory-bound). This is often “formalized” via a toy roofline model (which conjures up unpleasant memories of supply-demand graphs in economics) of the performance “power” $\dot W$ (measured in FLOPS, not FLOPs) as a function of arithmetic intensity $I$:

Posted in Blog | Leave a comment

Screening

Posted on June 4, 2026 by wdengquantum.me

Problem: State the Poisson-Boltzmann mean-field equation for the electrostatic potential $\phi(\mathbf x)$ at temperature $T$ in a dielectric of isotropic permittivity $\varepsilon$ containing mobile and immobile charge carriers. Explain what it means to apply Debye-Huckel linearization to the Poisson-Boltzmann mean-field equation and state the conditions under which such an approximation is valid.

Solution: It is just Poisson’s electrostatic equation $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\phi=-\rho_f/\varepsilon$ with the free charge density $\rho_f=\rho_{\text{mobile}}+\rho_{\text{immobile}}$ and the mobile charge carrier density $\rho_{\text{mobile}}(\mathbf x)=\sum_i\rho^{\infty}_ie^{-\beta q_i\phi (\mathbf x)}$ given by a mean-field Boltzmann distribution where $\rho^{\infty}_i=q_in^{\infty}_i$ is the bulk $|\mathbf x|,1/\phi(\mathbf x)\to\infty$ density of the $i^{\text{th}}$ species of mobile charge carriers.

To apply Debye-Huckel linearization means Taylor expanding the Boltzmann factors $e^{-\beta q_i\phi (\mathbf x)}\approx 1-\beta q_i\phi(\mathbf x)$ for $\beta q_i\phi(\mathbf x)\ll 1$. This results in the Debye-Huckel equation:

\[\left(\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2-\frac{1}{r^2_D}\right)\phi(\mathbf x)=-\frac{\rho_{\text{immobile}}(\mathbf x)+\sum_i\rho_i^{\infty}}{\varepsilon}\]

with Debye screening length $r_D:=\sqrt{\varepsilon k_BT/\sum_iq_i^2n_i^{\infty}}$ (note: the Debye screening length is sometimes denoted $\lambda_D$ which is misleading because it is purely the length scale of an exponential decay, corresponding at best to an imaginary wavelength. Also, it’s a bit strange that the screening potential $\phi$ is named after both Debye and Huckel but its screening length $r_D$ is named after Debye only…).

Problem: Show that for a uniform jellium background of charge density $-\rho^{\infty}$ with an impurity charge $Q$ at the origin $r=0$, the classical Debye-Huckel screening potential $\phi$ around this impurity takes the isotropic Yukawa form:

\[\phi(r)=\frac{Q}{4\pi\varepsilon_0 r}e^{-r/r_D}\]

where the Debye screening length may be written $r_D:=v_T/\omega_p$ with thermal velocity $v_T:=\sqrt{k_BT/m}$ and plasma frequency $\omega_p:=\sqrt{n^{\infty}q^2/\varepsilon m}$.

Solution: In this case, for a single charge carrier species $q$ (e.g. electrons in a plasma or metal):

\[\rho_{\text{immobile}}(\mathbf x)=Q\delta^3(\mathbf x)-\rho^{\infty}\]

So it boils down to finding the (physical) Green’s function of the Helmholtz operator:

\[\left(\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2-\frac{1}{r^2_D}\right)\phi(\mathbf x)=-\frac{Q\delta^3(\mathbf x)}{\varepsilon}\]

which yields the answer. Incidentally, this shows the general solution of the Debye-Huckel equation may be written as a convolution with this Yukawa kernel:

\[\phi(\mathbf x)=\frac{1}{4\pi\varepsilon}\int d^3\mathbf x’\frac{\rho_{\text{immobile}}(\mathbf x)+\sum_i\rho_i^{\infty}}{|\mathbf x-\mathbf x’|}e^{-|\mathbf x-\mathbf x’|/r_D}\]

Problem: Consider a $T=0$ fully degenerate ideal Fermi gas. Often, one is taught that the Fermi energy $E_F:=\mu(T=0)$ is just the zero-temperature chemical potential. Now, suppose one applies a uniform external potential $V$ across the entire ideal Fermi gas, so that the Hamiltonian now reads $H=\sum_i\frac{|\mathbf p_i|^2}{2m}+V$. How do the new Fermi energy $E’_F$ and zero-temperature chemical potential $\mu'(T=0)$ of this gauge-shifted Hamiltonian compare to the original $V=0$ case.

Solution: The Fermi energy is gauge invariant, so in particular $E’_F=E_F$ (think “Fermi kinetic energy”). By contrast, the chemical potential is gauge covariant, so $\mu'(T=0)=\mu(T=0)+V$.

Problem: A key property of the Debye screening length is that $r_D\propto (n^{\infty})^{-1/2}$. However, this was a classical treatment of screening. By performing an analogous quantum treatment of screening, show that the corresponding Thomas-Fermi screening length $r_{TF}\propto (n^{\infty})^{-1/6}$.

Solution: One has to make the following assumptions:

(Fully degenerate Fermi gas) $T=0$ and the charge carriers have to actually be fermions (usually electrons).
(Local Density Approximation) The potential $\phi(\mathbf x)$ varies on length scales much longer than the de Broglie wavelength $\lambda_F$ of fermions on the Fermi surface, i.e. $|\partial\phi/\partial\mathbf x|\ll\phi(\mathbf x)/\lambda_F$. Note $\lambda_F=\sqrt{\pi}\lambda_{T_F}$ is almost but not quite the same thing as the thermal de Broglie wavelength $\lambda_{T_F}$ at the Fermi temperature $T_F$). Intuitively, this means one can approximate the potential $\phi(\mathbf x)$ by a suitable “staircase” where the size of each stair is $\gg\lambda_F(\mathbf x)\propto n^{-1/3}(\mathbf x)$ but also $\ll\phi(\mathbf x)/|\partial\phi/\partial\mathbf x|$.
(Linear Response) $|q\phi(\mathbf x)|\ll E_F^{\infty}$
(Mean-Field) Same as Poisson-Boltzmann (namely, neglecting electron-electron many-body exchange correlations, just working with the mean, self-consistent field $\phi(\mathbf x)$).
(Jellium Background) Same as the above Debye-Huckel model.

Then, the idea is that the electron number density will be non-uniform $n=n(\mathbf x)$. By the local density approximation, one can coarse grain the domain into boxes of suitable size (which will be modelled as point-like $\mathbf x$ but not thought of like that), such that each box $\mathbf x\in\mathbf R^3$ is a locally uniform ideal Fermi gas in the presence of an external mean-field potential $\phi(\mathbf x)$ (this should be read as saying “each box $\mathbf x$ receives its own tailored uniform potential $\phi(\mathbf x)$”). However, because these boxes are all connected to each other and can thus exchange particles, at equilibrium there will be some global uniform value of chemical potential $\mu=E_F^{\infty}$ across the entire system. The local density $n(\mathbf x)$ competes with the local potential $\phi(\mathbf x)$ in such a way that their sum is conserved:

\[E_F^{\infty}=E_F(\mathbf x)+q\phi(\mathbf x)\]

Whereas for Poisson-Boltzmann one had $n(\phi)=n^{\infty}e^{-\beta q\phi}$, here the analogous relation is thus $n(\phi)=\frac{1}{3\pi^2}\left(\frac{2m(E^{\infty}_F-q\phi)}{\hbar^2}\right)^{3/2}$, and whereas applying Debye-Huckel linearization to Poisson-Boltzmann meant linearizing the exponential, here it means linearization via the binomial expansion $n(\phi)\approx n^{\infty}\left(1-\frac{3q}{2E^{\infty}_F}\phi\right)$ which is valid thanks to the linear response assumption above. Note the $3$-dimensional density of states per unit volume at the Fermi surface is $g_V(E_F^{\infty})=3n^{\infty}/2E_F^{\infty}$ s0 one can equivalently write $n(\phi)\approx n^{\infty}-qg_V(E_F^{\infty})\phi$, and from here the rest of the analysis is identical to the above (e.g. $\rho_{\text{immobile}}(\mathbf x)=Q\delta^3(\mathbf x)-\rho^{\infty}$ is the exact same thing), and one ends up with the same Yukawa form for the Thomas-Fermi screened potential:

\[\phi(r)=\frac{Q}{4\pi\varepsilon_0 r}e^{-r/r_{TF}}\]

only now the Thomas-Fermi screening length $r_{TF}:=\sqrt{\varepsilon/q^2g_V(E_F^{\infty})}=\left(\frac{\pi^4\hbar^6\varepsilon^3}{3q^6m^3n^{\infty}}\right)^{1/6}$. This is related to the Debye screening length by:

\[\frac{r_D^2}{r_{TF}^2}=\frac{3}{2}\frac{T}{T^{\infty}_F}\]

Hence for typical metals at room temperature $T\ll T^{\infty}_F\sim 10^3\text{ K}$, quantum mechanics (specifically the Pauli exclusion principle) weakens the effect of screening $r_{TF}\gg r_D$ compared to classical predictions.

(aside: to see more naturally why the Fermi surface density of states $g_V(E_F^{\infty})$ appeared earlier, one can recast the Debye-Huckel linearization step in the form $n(E_F)\approx n(E_F^{\infty})+\frac{\partial n}{\partial E_F^{\infty}}(E_F-E_F^{\infty})=n^{\infty}+g_V(E_F^{\infty})(-q\phi)$ thanks to the equilibrium condition $E_F^{\infty}=E_F(\mathbf x)+q\phi(\mathbf x)$. This also makes it clear that the form with the density of states is more general in that it holds in arbitrary dimension).

Problem:

Posted in Blog | Leave a comment

Schwarzschild Metric, Orbits & Black Holes

Posted on May 24, 2026 by wdengquantum.me

Problem: State the Schwarzschild metric solution to the Einstein field equations and the assumptions underlying it.

Solution: The Schwarzschild metric is given by:

\[ds^2=\left(1-\frac{r_s}{r}\right)d(ct)^2-\frac{dr^2}{1-\frac{r_s}{r}}-r^2d\Omega^2\]

(where $r_s:=2GM/c^2$ is the Schwarzschild radius) and is (by Birkhoff’s theorem) the unique isotropic solution of the vacuum Einstein field equations in the absence of a cosmological constant $\Lambda:=0$.

Problem: For a test particle moving along a geodesic in the Schwarzschild spacetime, show that its radial coordinate $r(\tau)$ obeys:

\[\frac{1}{2}\dot{r}^2+\phi_{\text{eff}}(r)=\frac{c^2(\gamma^2-1)}{2}\]

where $\dot r:=\frac{dr}{d\tau}$ and the effective potential is given by:

\[\phi_{\text{eff}}(r)=\phi_{\text{eff,0}}(r)-\frac{GM\ell^2}{c^2r^3}\]

where the Newtonian effective potential is just $\phi_{\text{eff,0}}(r):=-\frac{GM}{r}+\frac{\ell^2}{2r^2}$.

Solution: Reading off the Schwarzschild metric $g$ and assuming an orbit in the equatorial plane $\theta=\pi/2$ w.l.o.g., the Lagrangian may be taken as:

\[L=\left(1-\frac{r_s}{r}\right)(c\dot{t})^2-\frac{\dot{r}^2}{1-\frac{r_s}{r}}-r^2\dot{\phi}^2\]

Since $dL/d\tau=0$ is not explicitly dependent on the affine parameter $\tau$, the Beltrami identity asserts:

\[c^2=\frac{\partial L}{\partial\dot t}\dot t+\frac{\partial L}{\partial\dot r}\dot r+\frac{\partial L}{\partial\dot\phi}\dot\phi-L\]

But fortunately, $t$ and $\phi$ are ignorable coordinates so one has the conserved quantities $\frac{\partial L}{\partial\dot t}=2(1-r_s/r)c^2\dot{t}:=2c^2\gamma$ and $\frac{\partial L}{\partial\dot\phi}=-2r^2\dot{\phi}:=-2\ell$ which can be substituted into the Beltram identity to obtain the desired answer.

Alternatively, one can write $ds^2=c^2d\tau^2$ and divide through the metric by $d\tau^2$ to isolate the terms of the form $\dot t^2,\dot r^2,\dot\phi^2$, etc. which has the advantage of explaining that $c^2$ term. For massless particles such as photons, the geodesics are null $ds^2=0$, leading to $V_{\text{eff}}(r)=\frac{\ell^2}{2r^2}\left(1-\frac{r_s}{r}\right)$ missing the Newtonian potential term $-GM/r$. This leads to concepts such as the photon sphere $r=3r_s$ and the capture impact parameter $\rho^*=\sqrt{27}r_s/2$:

Posted in Blog | Leave a comment

Polymers

Posted on May 21, 2026 by wdengquantum.me

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

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

Solution: Central limit theorem.

Problem: Explain what an ideal polymer chain is.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Problem: Describe Flory’s mean-field theory.

Solution: The free energy is $F=\langle E\rangle-TS$. The energy $E$ due to collisions of monomers with excluded volume $V_{\text{ex}}$ can be modelled by a pseudopotential $V(\mathbf x)=V_{\text{ex}}k_BT\delta^3(\mathbf x)$ as this can be checked (after Taylor expanding the Mayer $f$-function $f(\mathbf x)=e^{-V_{\text{ex}}\delta^3(\mathbf x)}-1\approx -V_{\text{ex}}\delta^3(\mathbf x)$) to have the same second virial coefficient $\tilde{B}_2(T)=-\int d^3\mathbf x f(\mathbf x)=V_{\text{ex}}$ as a corresponding hard-sphere monomer of excluded volume $V_{\text{ex}}=8V_{\text{monomer}}$. Then for number density $n(\mathbf x)$ of monomers:

\[E=\frac{1}{2}\sum_{i\neq j}V(\mathbf x_i-\mathbf x_j)\approx \frac{V_{\text{ex}}k_BT}{2}\int d^3\mathbf x d^3\mathbf x’n(\mathbf x)n(\mathbf x’)\delta^3(\mathbf x-\mathbf x’)\]

\[=\frac{V_{\text{ex}}k_BT}{2}\int d^3\mathbf x n^2(\mathbf x)\]

Thus, the ensemble average $\langle E\rangle$ is evaluated by a mean-field approximation:

\[\langle E\rangle=\frac{V_{\text{ex}}k_BT}{2}\int d^3\mathbf x \langle n^2(\mathbf x)\rangle\]

\[\approx \frac{V_{\text{ex}}k_BT}{2}\int d^3\mathbf x \langle n(\mathbf x)\rangle^2\]

If one furthermore assumes monomers $\langle n(\mathbf x)\rangle=N/V$ are uniformly distributed, then this simplifies to $E(N,r)=V_{\text{ex}}N^2k_BT/2V=3V_{\text{ex}}N^2k_BT/8\pi r^3$. Meanwhile, the entropy is simply (up to some irrelevant baseline $S_0$) the log of the normal distribution $p(\mathbf x)$ for the freely-jointed Gaussian chain from earlier:

\[S(N,r)=S_0-\frac{3k_Br^2}{2N\ell^2}\]

Thus, by setting $\partial F/\partial r:=0$, one obtains the equilibrium relation:

\[r=\left(\frac{3V_{\text{ex}}\ell^2}{8\pi}\right)^{1/5}N^{3/5}\]

By fortunate cancellation, it turns out in a good solvent this simple mean-field exponent $\nu=3/5$ is close the true exponent $\nu\approx 0.588$ arising from a detailed RG flow analysis.

Problem: In the naive analysis above, it was assumed that the excluded volume $V_{\text{ex}}>0$ is a literal volume, and must therefore be positive. What’s wrong with that?

Solution: In practice this is not the case, and indeed one can run through the standard van der Waals correction (this time including some long-range potential such as Lennard-Jones on top of the hard sphere excluded volume) and find a $T$-dependent effective excluded volume which is just another way of saying the second virial coefficient:

\[\tilde B_2(T)=V_{\text{ex}}\left(1-\frac{T_B}{T}\right)\]

where the interaction-dependent Boyle temperature $T_B$ is also called the “theta temperature” in this context. It is conventional to define the dimensionless Curie-like Flory interaction parameter $\chi(T):=T_B/2T$ in terms of which $\tilde{B}_2(T)=V_{\text{ex}}(1-2\chi(T))$. Thus, there exist $3$ regimes:

$T>T_B$ means $B_2(T)>0$ and $\chi(T)<1/2$ in which case the polymer chain swells as $r\propto N^{3/5}$ (“good solvent”).
$T=T_B$ means $B_2(T)=0$ and $\chi(T)=1/2$ in which case the polymer chain behaves as if ideal/freely-jointed $r\propto N^{1/2}$ (“theta solvent”).
$T<T_B$ means $B_2(T)<0$ and $\chi(T)>1/2$ in which the polymer coil transitions into a globule $r\propto N^{1/3}$ (“poor solvent”).

Problem: State the Flory-Huggins free energy density per lattice site $\Delta F_{\text{mix}}(\phi)/N$ as a function of the polymer volume fraction $\phi:=N_mN_p/N$ (i.e. each of the $N_p$ polymer chains contains $N_m$ monomer subunits, and all $N=N_s+N_mN_p$ lattice sites are assumed to be occupied by either one of the $N_mN_p$ monomers or one of the $N_s$ solvent molecules).

Solution: The Flory interaction parameter for coordination number $z\in\mathbf Z^+$ is given by $\chi(T):=z\Delta E/k_BT$ where $\Delta E:=E_{sm}-(E_{ss}+E_{mm})/2$:

\[\Delta F_{\text{mix}}(\phi)/N=k_BT\left(\frac{\phi}{N_m}\ln\phi+(1-\phi)\ln(1-\phi)+\chi\phi(1-\phi)\right)\]

(aside: how is this definition of $\chi$ related to the previous one that was used when discussing the coil-to-globule transition? After all, this one seems like it can be positive or negative but the other one can only be positive right…).

If instead of polymers dissolved in a solvent, it was one polymer (of chain length $N_m$) dissolved in another polymer (of chain length $N’_m$) (known as a polymer blend), the Flory-Huggins free energy density would naturally be modified to:

\[\Delta F_{\text{mix}}(\phi)/N=k_BT\left(\frac{\phi}{N_m}\ln\phi+\frac{1-\phi}{N’_m}\ln(1-\phi)+\chi\phi(1-\phi)\right)\]

If in addition the polymer blend is in fact a polymer melt (i.e. even though $N_m\neq N’_m$, the monomer subunits are chemically identical), then $\chi=0$ and hence the polymers behave like random walks $r\sim N_m^{1/2}$. Think hot glue or $3$D printer filament immediately after extrusion.

Problem: Compute the osmotic pressure $p_{\text{os}}$ and the chemical potential $\mu$ from the Flory-Huggins interaction parameter.

Solution: Take the relevant partial derivatives…

Problem: The dilute limit is $\phi\to 0$, whereas the polymer melt is the limit $\phi\to 1$. It turns out one can define a critical volume fraction $\phi^*\in [0,1]$ representing the threshold for transition from the dilute regime $\phi<\phi^*$ to a “semi-dilute” regime $\phi^*<\phi\ll 1$. Show that $\phi^*\propto N_m^{-4/5}$ in a good solvent.

Solution: Similar to Bose-Einstein condensation, the critical point occurs when the “phase space density” $N^*_pr^3/V\propto 1\Rightarrow N^*_p/N\propto r^{-3}\propto N_m^{-9/5}$. The corresponding critical volume fraction $\phi^*=N_mN^*_p/N\propto N_m^{-4/5}$.

Problem: Compute the spinodal curve $\chi(\phi)$ and state its physical significance.

Solution: The spinodal curve is given by $\partial^2(\Delta F_{\text{mix}}/N)/\partial\phi^2=0$ which amounts to:

\[\chi(\phi)=\frac{1}{2}\left(\frac{1}{N_m\phi}+\frac{1}{1-\phi}\right)\]

This is maximized $\partial\chi/\partial(\phi=\phi_c)=0$ at:

\[\phi_c=\frac{1}{1+\sqrt{N_m}}\]

with corresponding Flory interaction parameter slightly larger than $\chi_c>1/2$:

\[\chi(\phi_c):=\chi_c=\frac{1}{2}\left(1+\frac{1}{\sqrt{N_m}}\right)^2\]

Posted in Blog | Leave a comment

Magnetism

Posted on May 16, 2026 by wdengquantum.me

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

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

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

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

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

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

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

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

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

Solution: One has:

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

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

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

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

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

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

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

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

where the Brillouin function is defined by:

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

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

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

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

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

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

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

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

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

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

Problem: State and prove the Stoner criterion $\frac{g_V(E_F)}{2}g>1$ for itinerant ferromagnetism.

Solution: The Stoner criterion is derived within a toy model of itinerant ferromagnetism known fittingly as the Stoner model (a mean-field version of the Fermi-Hubbard model). That is, given a fixed number density $n=n_{\uparrow}+n_{\downarrow}$ of spin-$1/2$ fermions, each of these can be considered to form their own Fermi sea with $E_{F,\uparrow}=\hbar^2(3\pi^2n_{\uparrow})^{2/3}/2m$ and similarly for $E_{F,\downarrow}$. The total energy density $\mathcal E$ of the system is then:

\[\mathcal E(n_{\uparrow},n_{\downarrow})=\frac{3}{5}n_{\uparrow}E_{F,\uparrow}+\frac{3}{5}n_{\downarrow}E_{F,\downarrow}+gn_{\uparrow}n_{\downarrow}-\mu(n_{\uparrow}+n_{\downarrow}-n)\]

Setting $\partial\mathcal E/\partial n_{\uparrow}=\partial\mathcal E/\partial n_{\downarrow}=\partial\mathcal E/\partial\mu=0$ using the fact that $nE_F\propto n^{5/3}\Rightarrow\partial (nE_F)/\partial n=5E_F/3$ leads to the system of equations:

\[E_{F,\uparrow}+gn_{\downarrow}=E_{F,\downarrow}+gn_{\uparrow}\]

\[n_{\uparrow}+n_{\downarrow}=n\]

The simplest (and in most circumstances, the unique) solution to these equations is the unpolarized paramagnetic occupation numbers $n_{\uparrow}=n_{\downarrow}=n/2$ (and thus $E_{F,\uparrow}=E_{F,\downarrow}:=E_F$) as if the Stoner interaction $g=0$ didn’t exist. So $(n_{\uparrow},n_{\downarrow})=(n/2,n/2)$ is always a stationary point of the function $\mathcal E(n_{\uparrow},n_{\downarrow})$, but what about its stability? One can compute the Hessian matrix:

\[\begin{pmatrix}\partial^2\mathcal E/\partial n_{\uparrow}^2&\partial^2\mathcal E/\partial n_{\uparrow}\partial n_{\downarrow} \\ \partial^2\mathcal E/\partial n_{\uparrow}\partial n_{\downarrow} & \partial^2\mathcal E/\partial n_{\downarrow}^2\end{pmatrix}=\begin{pmatrix}\partial E_{F,\uparrow}/\partial n_{\uparrow} & g \\ g & \partial E_{F,\downarrow}/\partial n_{\downarrow}\end{pmatrix}=\begin{pmatrix} 2/g_V(E_F) & g \\ g & 2/g_V(E_F) \end{pmatrix}\]

using the fact that the total density of states per unit volume per spin $g_{V,\uparrow}(E)=g_{V,\downarrow}(E)=g_V(E)/2$ are equal and sum to the total density of states per unit volume and moreover $g_V(E_F)=\partial n/\partial E_F$. The eigenvectors of the Hessian are clearly $(1,1)$ and $(1,-1)$ from which one deduces the eigenvalues $\frac{2}{g_V(E_F)}+g$ and $\frac{2}{g_V(E_F)}-g$. Since both $g_V(E_F),g>0$, the stationary point $(n_{\uparrow},n_{\downarrow})=(n/2,n/2)$ is thus an unstable saddle point iff the eigenvalue $\frac{2}{g_V(E_F)}-g<0$ which rearranges to the standard form of the Stoner criterion $\frac{g_V(E_F)}{2}g>1$.

Because $\mathcal E(n_{\uparrow},n_{\downarrow})$ is a function of only $2$ variables, the second partial derivative test asserts that this

Posted in Blog | Leave a comment

Physics-Informed Neural Networks

Posted on May 5, 2026 by wdengquantum.me

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

Solution:

report

Spectral Bias of Physics-Informed Neural Networks¶

1. Introduction and Background¶

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

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

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

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

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

The project is split into two experiments:

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

2. Algorithms and Implementation¶

2.1 Network Architecture¶

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

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

2.2 Training¶

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

2.3 Hardware¶

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

3. Experiment 1: Van der Pol Oscillator¶

3.1 Problem Setup¶

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

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

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

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

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

In [1]:

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

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

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

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

Using device: cuda

3.2 Results: Time Series and Phase Portraits¶

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

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

In [2]:

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


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

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

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

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

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

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

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

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

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

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

        plt.tight_layout()
        plt.show()

In [3]:

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

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

No description has been provided for this image

4. Experiment 2: Fokker-Planck Equation¶

4.1 Problem Setup¶

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

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

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

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

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

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

4.2 Cost Function¶

The cost function has three terms:

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

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

4.3 Results¶

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

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

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

In [4]:

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

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

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

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

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

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

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

In [5]:

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

5. Discussion and Extensions¶

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

Several extensions were considered but not implemented:

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

References¶

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

Posted in Blog | Leave a comment

Diffusion & Flow-Matching Models

Posted on April 12, 2026 by wdengquantum.me

Problem: State and prove Tweedie’s formula.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

which intertwines convolution and multiplication).

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

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

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

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

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

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

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

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

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

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

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

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

Or equivalently:

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

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

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

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

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

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

Hence:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Posted in Blog | Leave a comment

Information Geometry

Posted on March 27, 2026 by wdengquantum.me

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

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

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

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

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

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

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

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

Solution: One has:

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

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

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

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

Solution:

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

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

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

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

Solution:

Posted in Blog | Leave a comment

Quantumplations

Machine Learning Interatomic Potentials

Reinforcement Learning (Part \(2\))

Graphics Processing Units (GPUs)

Screening

Schwarzschild Metric, Orbits & Black Holes

Polymers

Magnetism

Physics-Informed Neural Networks

Spectral Bias of Physics-Informed Neural Networks¶

1. Introduction and Background¶

2. Algorithms and Implementation¶

2.1 Network Architecture¶

2.2 Training¶

2.3 Hardware¶

3. Experiment 1: Van der Pol Oscillator¶

3.1 Problem Setup¶

3.2 Results: Time Series and Phase Portraits¶

4. Experiment 2: Fokker-Planck Equation¶

4.1 Problem Setup¶

4.2 Cost Function¶

4.3 Results¶

5. Discussion and Extensions¶

References¶

Diffusion & Flow-Matching Models

Information Geometry

Archives

Meta