implicit

implementation

Implicit inference and its backward pass for a diffusion model: a deterministic residual field, a solver that only reports roots it trusts, and an adjoint that turns one linear solve into gradients for parameters, inputs, anchors and precisions. The theory is in Implicit Diffusion Learners and Backpropagation through Implicit Inference.

Sources: code: implicit.jl

Theory (CT-ML wiki): Least Fixed Point · Statistical Game · Bayesian Lens · Reverse Derivative Category

1. What is in the file

name	is
FieldNodes, field_nodes, noisefree_nodes	the fixed $(t_k,{\epsilon }_k,w_k)$ that make the field deterministic
ImplicitDiffusion	the relation: a NoisePredictor plus a node set plus λ
prior_field, prior_jacobian	$g_{\theta }(z)$ and $\partial g/\partial z$
implicit_residual	$r=g+{\rho }^2\odot (z-z_0)$ on the free coordinates
implicit_infer → ImplicitSolution	a root, with residual, iters, converged, stable
implicit_pullback	the adjoint: (z₀ = …, ρ = …, ps = …)

2. Design decisions

• Deterministic by construction. RED-Diff redraws $(t,\epsilon )$ every step; here the nodes are fixed, so inference is a function and the implicit function theorem applies. field_nodes uses antithetic pairs so that the $-\epsilon$ control variate cancels exactly.
• Polarity as a precision vector, the same encoding as precision_vector(f, polarity): Inf is a hard clamp, imposed by projection, never as a penalty.
• The solver trusts Newton only where it should. A Newton step is taken when the symmetric part of $J_{FF}$ is positive definite, with backtracking on $\Vert r_F\Vert$; otherwise a fixed descent step along the field. Newton is attracted to every root, saddles and maxima included; restricting it to locally convex regions keeps it on stable roots.
• Two flags, not one. converged (the residual is small) and stable ($\mathrm{sym}{J}_{FF}\succ 0$) are separate, because a converged unstable root is on the zero set but not in the relation one wants.
• The pullback refuses non-converged states. The IFT says nothing about a non-root, and silently returning a gradient there is what made the first prototype diverge.
• No AD dependency. The input Jacobian has a finite-difference fallback (epsilon_jacobian); the parameter VJP is an interface method (epsilon_vjp_params) that closed-form predictors implement exactly and a Lux network gets from the AD backend in its ad field (backends). A model with no parameters skips it.

3. Implementation difficulties

1. The tolerance floor. The field of a narrow mixture at small $t$ has a floating-point noise floor around $10^{-11}$ (it scales like $1/{\sigma }_t$). An absolute tolerance of $10^{-11}$ made a third of the solves “fail”; the tolerance is relative, tol * (1 + ‖z‖).
2. The prototype’s solver. A version with a fixed descent step in non-convex regions and an energy line search dropped about 60% of solves during training, and its training stalled at RMSE 0.15. The package’s solver (Newton where convex, residual backtracking) dropped none and reached 0.012.
3. Smoothing levels. With RED-Diff’s training range of noise levels, the deterministic field of a circle has its only root at the centre. The default levels stop at $t=0.05$; the trade-off is measured in Implicit Diffusion Learners §5.
4. Dense Jacobians. prior_jacobian forms $J$ densely, which is fine for the small $n$ of factor blocks but not for images. Large $n$ needs a matrix-free GMRES in both the Newton step and the adjoint, each iteration costing one VJP per node.

4. How it is tested

test/runtests.jl, against the closed-form mixture of analytic and the Gaussian oracle:

• the field equals $\nabla \Phi$ and its Jacobian is symmetric;
• both branches of the circle are found, and they merge near the branch point;
• the smoothing bias is monotone in the levels and matches $R-(s^2+{\sigma }^2/{\alpha }^2)/(2R)$;
• noise-free roots are Tweedie fixed points;
• the Gaussian root and its derivative match the closed form;
• every adjoint cotangent matches finite differences of the full solve;
• training through inference reduces the loss by more than 50×;
• implicit_roots finds both branches of the circle, and the Laplace covariance of a Gaussian query equals the closed form $1/(\kappa +{\rho }^2)$.

5. All answers, and how sure (answers.jl)

Two functions build on implicit_infer without changing it.

implicit_roots solves from z₀ and from perturbed starts (normal on the free coordinates, the clamped ones fixed), keeps converged and stable answers, and removes duplicates. For an energy-parametrised predictor the answers are ordered by the query’s energy. On the circle, ”$y$ given $x=0.6$” returns both branches; near $x=\pm 1$ they merge into one. It finds what its starts reach: no guarantee of completeness (Open Problems in Implicit Diffusion Learning T1).

implicit_laplace inverts the symmetric part of $J_{FF}$, the residual’s Jacobian on the free coordinates at the answer, which is the curvature of the query’s energy. The subtle part is the units. With $x_k={\alpha }_{k}z+{\sigma }_k{\epsilon }_k$ and ${\epsilon }_{\theta }=-{\sigma }_t{\nabla }_x\log p_t$,

$$g(z)={\nabla }_z\left[\sum _k{w}_k\lambda \frac{{\sigma }_k^2}{{\alpha }_k^2}\right(-\log p_{t_k}(x_k)\left)\right]+(\text{terms linear in }z),$$

so the prior term is a sum of smoothed negative log-densities with weights $\lambda {\sigma }_k^2/{\alpha }_k^2$. They sum to one exactly when $\lambda =1/{\sum }_k{w}_k{\sigma }_k^2/{\alpha }_k^2$, density_lambda(schedule, nodes), and then the energy is a proper negative log-density in the units of the clamp’s Gaussian likelihood, and the covariance is in data units. The linear terms come from the noise draws and cancel for antithetic nodes. For queries with only hard inputs and free outputs, λ scales the curvature and leaves the answers unchanged.

On the circle with this λ, the standard deviation of $y$ is 0.135 at $x=0$ (the smoothed ring’s width: data width 0.05 combined with the noise levels gives about 0.12), 0.17 at $x=0.6$ (the vertical line crosses the ring obliquely), and it grows towards the branch points (0.31, then 0.41) and off the circle (0.53). Calibration is to the density smoothed at the field’s noise levels, so it is wider than the data’s own spread, the same smoothing that biases the answers inwards (Implicit Diffusion Learners §5).

What is still missing: the covariance is returned as a matrix, not a GaussianBelief, because that type lives above lib/ (Belief Algebra §6); and λ chosen for calibration is not always the λ one wants for balancing soft evidence.

Related: analytic, reddiff, factor, energy, Deterministic Relaxation, Inference Signatures, The Implicit Diffusion Factor as a Statistical Game