Local Fokker--Planck Geometry for Score Estimation: Heat-Ball Mean-Value Representations and Exact High-Dimensional Sampling

2026-06-26

AI Summary

This work addresses the challenge of score estimation for diffusion processes with nonlinear, state-dependent drift, where the marginal density lacks a closed-form solution, leading to significant errors in low-density regions when using conventional global averaging methods. The authors propose a local Fokker–Planck geometric framework that transforms the variable-coefficient Fokker–Planck equation into an inhomogeneous heat equation via a cumulative variance time change. By replacing global conditional expectations with local heat-ball averages, the method enables accurate local estimation of the score function. Key contributions include a heat-ball mean representation theorem, the introduction of a κ-measure, and an unweighted high-dimensional sampler. The framework is shown to be a one-parameter generalization of the global Fokker–Planck residual approach. Experiments on 2D structured data and 256-dimensional MNIST demonstrate its efficacy, with the concentration properties of the κ-measure markedly improving high-dimensional sampling accuracy and score estimation in low-density regions.

Abstract

Score-based generative models and Langevin samplers rely on estimating the score function $\nabla_x\log p_t(x)$ of a forward diffusion. Classically this is tractable when the drift is linear: the marginal density is Gaussian and the score is a global conditional expectation. For a general nonlinear, state-dependent drift the marginal density has no closed form, and existing methods--denoising score matching and global Fokker--Planck residual penalties--resort to global averaging that inflates estimation error in low-density regions precisely where accuracy is most critical. We address this by developing a local Fokker--Planck geometric framework that replaces global conditioning with local parabolic averaging. Our approach rests on three contributions. First, a time change to the cumulative-variance coordinate reduces the variable-coefficient Fokker--Planck equation to a standard inhomogeneous heat equation, on which we extend Evans' classical heat-ball monotonicity method to derive exact local mean-value representations for the score $\nabla_x\log p$ together with the density, log-density, and entropy density; local well-posedness is established under an explicit dimension-dependent drift budget. Second, for high-dimensional Monte Carlo evaluation of the resulting heat-ball integrals, we introduce the $κ$-measure and derive its exact factorized sampler with unit per-sample weight, $χ^2_2$ radial concentration. Third, the $r\to0$ limit of the heat-ball residual recovers the pointwise Fokker--Planck residual, showing that the local framework is a one-parameter generalization of global FP-residual methods, and that the DSM population minimizer is feasible for the heat-ball constraint at every scale. We validate the framework on 2D structured data on 256-dimensional MNIST, and on a dedicated sampler study confirming the concentration laws.

Problem

• score estimation
• Fokker-Planck equation
• nonlinear drift
• low-density regions
• generative modeling

Innovation

• local Fokker--Planck geometry
• heat-ball mean-value representation
• score estimation
• exact high-dimensional sampling
• κ-measure