Frozen-Diffusion Parametrix Preconditioner

Integrate the method into a neural HJB solver at the training-target and residual-evaluation stage, not into the network architecture itself. Use an MLP $V_\theta(t,x)$ and a backward time grid $t_0<\cdots<t_N=T$. Partition sampled states into local patches with centers $(\tau_j,x_j)$; estimate $A_j=\sigma(\tau_j,x_j)\sigma(\tau_j,x_j)^\top$, symmetrize it, and add $\epsilon I$ if needed. For each patch and timestep, compute the network quantities $p=\nabla_xV_\theta(t_{n+1},x)$ and $B=D_x^2V_\theta(t_{n+1},x)$ using autodiff. Evaluate the Hamiltonian $H=\min_u[(f+gu)^\top p+u^\top Ru]$; for quadratic control with invertible $R$, use $u^*=-\frac12R^{-1}g^\top p$ and substitute it into $H$. Form the source correction $q=\ell+H+\frac12(A(t_{n+1},x)-A_j):B$. Apply the frozen Gaussian propagator $P_{\Delta t}^{A_j}$ to $q$ by either Monte Carlo samples $x'=x+L_j\sqrt{\Delta t}\,z$, where $L_jL_j^\top=A_j$, averaging $q(t_{n+1},x')$, or a truncated cosine/Fourier feature approximation using the multiplier $\exp[-\frac12\Delta t\xi^\top A_j\xi]$. Construct the backward target $y=P_{\Delta t}^{A_j}V_\theta(t_{n+1},\cdot)+\Delta tP_{\Delta t}^{A_j}q$ and minimize $\|V_\theta(t_n,x)-y\|^2$ over sampled points, optionally adding a small raw PDE residual term. Run one Picard sweep per epoch, detach the target for stability, and refresh patch centers every few epochs. The paper supplies the frozen-diffusion/variable-coefficient decomposition; Monte Carlo error, patch assignment, and timestep size are empirical choices. First test on a 10-dimensional linear-quadratic diffusion-control problem with known Riccati solution, comparing a vanilla PINN residual baseline against the same MLP with raw backward targets. Measure value-function relative error, gradient error, loss decrease per optimizer step, and sensitivity to $\Delta t$. Success is faster decrease and lower error at equal network evaluations, especially when $A(x)$ varies strongly across the state space.