proxdm

implementation

Proximal diffusion models: a diffusion prior that is queried through its proximal operator instead of its score. Implemented here are the interface, an exact oracle, ProxDM’s sampler, proximal inference for the implicit learner, and the training loss.

Sources: Fang, Díaz, Buchanan & Sulam, Beyond Scores: Proximal Diffusion Models, arXiv:2507.08956 (Algorithm 1, Eq. 9); Zhu et al., Denoising Diffusion Models for Plug-and-Play Image Restoration, CVPR 2023 (DiffPIR, the half-quadratic splitting); code: proxdm.jl

Theory: ProxDM and Proximal Alternatives · Deterministic Relaxation · Inference Signatures

Notation. As everywhere in this vault, $z\in Z$ is the joint state and the prior is $p_t$ on $Z$. The paper writes $x$ for the same thing.

1. The object

$${\mathrm{prox}}_{-\lambda \log p_t}(v)=\arg \underset{u}{\min }\frac{1}{2}\Vert u-v{\Vert }^2-\lambda \log p_t(u),u-v=\lambda \nabla \log p_t(u).$$

The optimality condition is a gradient step evaluated at the new point: a backward (implicit) Euler step, where a score is a forward one. That is the whole idea of ProxDM: discretise the reverse SDE backwards, and the network you need is a prox, not a score.

type	what it is
AbstractProximalPredictor	anything with proximal(p, v, t, λ, ps, st) -> (u, st)
ProxNetwork(model, schedule; input)	a Lux model in ProxDM’s parametrisation $f_{\theta }(v;t,\lambda )=v-\sqrt{\lambda }{\epsilon }_{\theta }(v;t,\lambda )$; conditioned on two scalars
MixtureProx(::GaussianMixtureEps)	the exact prox of a Gaussian mixture, by Newton with backtracking; the oracle a perfectly trained ProxNetwork would match

MixtureProx uses $\nabla \log p_t$ and its Hessian directly from the mixture, not $-\epsilon /{\sigma }_t$, so it is valid at $t=0$, where the sampler’s last step needs it.

2. The sampler (Algorithm 1)

On a uniform grid $0=t_0<\cdots <t_N=T$ with ${\gamma }_k={\int }_{t_{k-1}}^{t_k}\beta$:

• PDA (hybrid = false): $z_{k-1}={\mathrm{prox}}_{-\frac{2{\gamma }_k}{2-{\gamma }_k}\log p_{t_{k-1}}}(\frac{2}{2-{\gamma }_k}(z_k+\sqrt{{\gamma }_k}{\xi }_k))$, which needs ${\gamma }_k<2$; proxdm_sample throws otherwise.
• PDA-hybrid (hybrid = true): $z_{k-1}={\mathrm{prox}}_{-{\gamma }_k\log p_{t_{k-1}}}((1+\frac{1}{2}{\gamma }_k)z_k+\sqrt{{\gamma }_k}{\xi }_k)$: the linear drift explicit, the score implicit; no step-size limit.

The last step evaluates the prox at $t_0=0$, so the output is already denoised; there is no separate final denoising step.

Validation (one Gaussian, mean $(0.5,-1)$, std $0.3$, exact prox, 1,500 samples):

steps	20	40	100	400
sample std (true 0.30)	0.248	0.268	0.286	0.293

The mean is right at every step count; the spread converges at first order, as a backward Euler scheme should. The sampler on the circle mixture gives radius $1.002\pm 0.024$.

3. Proximal inference for the implicit learner

prox_infer(p, z₀, ρ, t, ps, st; λ) is the implicit learner’s query with a proximal prior. It splits the relaxed problem

$$\underset{u,w}{\min }-\log p_t(u)+\frac{1}{2\lambda }\Vert u-w{\Vert }^2+\frac{1}{2}\Vert P(w-z_0){\Vert }^2,P=\mathrm{diag}(\rho ),$$

and alternates the two proximal maps (half-quadratic splitting, as in DiffPIR):

$$u\leftarrow {\mathrm{prox}}_{-\lambda \log p_t}(w),w_i\leftarrow \frac{u_i/\lambda +{\rho }_i^2{z}_{0,i}}{1/\lambda +{\rho }_i^2}({\rho }_i<\infty ),w_i\leftarrow z_{0,i}({\rho }_i=\infty ).$$

Hard clamps are exact; the free coordinates (${\rho }_i=0$) take $w_i=u_i$. As $\lambda \to 0$ the coupling becomes $u=w$ and the problem becomes the implicit learner’s (Deterministic Relaxation). It returns an ImplicitSolution with z = w (which satisfies the clamps) and residual $=\Vert u-w\Vert$.

On the circle mixture, the query ”$x=0.6$, which $y$?” from $y_0=\pm 0.5$ returns $y=\pm 0.7936$, symmetric to the last digit, in 20 iterations.

4. Training: proximal matching

proximal_matching_loss is one sample of the paper’s Eq. 9:

$${\ell }_{PM}({\epsilon }_{\theta }(z_t+\sqrt{\lambda }\xi ;t,\lambda ),\xi ;\zeta ),{\ell }_{PM}(a,b;\zeta )=1-\exp (-\Vert a-b{\Vert }^2/(d{\zeta }^2)).$$

The squared loss would train the MMSE denoiser $\mathbb{E}[z\mid v]$; this bounded loss, as $\zeta \to 0$, trains the MAP denoiser, which is the prox. Shrink $\zeta$ over training. The loss is a plain function of ps, so any Lux training loop and any AD backend trains it.

5. Implementation difficulties and open ends

• Non-convex priors make the prox multivalued. For a mixture, $-\log p_t$ is not convex, and $\mathrm{prox}$ can have several minimisers. MixtureProx uses Newton where $I-\lambda {\nabla }^2\log p_t$ is positive definite and a gradient step otherwise, and returns the local minimiser reached from $v$. A learned ProxNetwork makes the same choice implicitly, wherever training put it.
• The PDA step-size limit. ${\gamma }_k<2$ fails for coarse grids near $t=T$ under the default VP schedule. The hybrid scheme removes it; PDA throws with a message saying so instead of returning garbage.
• No adjoint for prox_infer yet. The fixed point satisfies $u=\mathrm{prox}(w(u))$, so the implicit function theorem applies just as for implicit_infer. It needs the prox’s input Jacobian and parameter VJP, which is why ProxNetwork would get an ad field when this is added (backends).
• stable is not assessed by the splitting; it reports true when converged.
• A ProxDM factor inversion (alongside REDDiff and ImplicitProx in factor.jl) is the natural next step and would make ProxDM usable inside a factor graph. It is the first case with two parameter trees (the prior’s ${\epsilon }_{\theta }$ and the prox network), the structure ProxDM and Proximal Alternatives §4 anticipated.

Related: implicit, reddiff, backends, analytic, ProxDM and Proximal Alternatives, Inference Signatures, Deterministic Relaxation