analytic

implementation

Closed-form noise predictors: the optimal ${\epsilon }^{\ast }$ of a Gaussian mixture, with its exact input Jacobian and parameter VJP. A ring of narrow components is a circle relation with an exact score, which makes implicit inference and its backward pass checkable against arithmetic.

Sources: code: analytic.jl

Theory (CT-ML wiki): Statistical Game · Bayesian Inversion

1. The formulas

For $p_0=\frac{1}{J}{\sum }_j\mathcal{N}({\mu }_j,s^{2}I)$ under the VP-SDE, $p_t$ is the mixture with means ${\alpha }_t{\mu }_j$ and variance $v_t={\alpha }_t^2{s}^2+{\sigma }_t^2$. With $u_j=(x-{\alpha }_t{\mu }_j)/v_t$, responsibilities ${\gamma }_j\propto \exp (-\Vert x-{\alpha }_t{\mu }_j{\Vert }^2/2v_t)$ and $\bar{u}={\sum }_j{\gamma }_j{u}_j$:

$${\epsilon }^{\ast }={\sigma }_t\bar{u},{\partial }_x{\epsilon }^{\ast }={\sigma }_t[\frac{1}{v_t}I-{\sum }_j{\gamma }_j{u}_j(u_j-\bar{u}{)}^{\top }],({\partial }_{{\mu }_j}{\epsilon }^{\ast }{)}^{\top }w={\sigma }_t{\alpha }_t{\gamma }_j[-\frac{w}{v_t}+u_j(u_j-\bar{u}{)}^{\top }w].$$

The input Jacobian is ${\sigma }_t$ times ($1/v_t$ minus a covariance of the $u_j$), so it is symmetric, as $-{\sigma }_t{\nabla }^2\log p_t$ must be.

2. Implementation difficulties

• The sign of the covariance term. The first prototype had $+\mathrm{Cov}$; finite differences caught it (error 0.58). The responsibilities’ gradient is ${\gamma }_j(\bar{u}-u_j)$, not ${\gamma }_j(u_j-\bar{u})$.
• Log-sum-exp. Responsibilities are computed from shifted log-weights, because at small $t$ the exponents are of order $1/v_t\sim 10^4$.
• Equal weights, isotropic components. Enough for relations along curves; general weights and covariances would add parameters but no new ideas.

3. Why it lives in src/ and not in test/

It is an oracle, the same role GaussianFactor plays in the parent package. It is also a usable model: a mixture whose means are trained by implicit’s adjoint is an implicit learner with an exact score, which Backpropagation through Implicit Inference §7 uses to learn a parabola from a circle.

4. The derivative interface it implements

epsilon_jacobian and epsilon_vjp_params are defined in this file for every predictor. A closed-form predictor implements them exactly, as above. A NoisePredictor around a Lux network gets them from the AD backend in its ad field, through a package extension; without one the Jacobian falls back to finite differences and the VJP raises an error that names the fix (backends). The mixture’s gradient and Hessian of $\log p_t$ are reused by MixtureProx, the exact proximal operator (proxdm).

Related: implicit, predictor, backends, proxdm, Implicit Diffusion Learners