Probability-Flow ODE in Infinite-Dimensional Function Spaces
Kunwoo Na, Junghyun Lee, Se-Young Yun, Sungbin Lim
TL;DR
This work derives an infinite-dimensional probability-flow ODE (PF-ODE) for function-space diffusion models by formulating a weak, measure-based evolution in a Hilbert space with a logarithmic gradient on the Cameron–Martin space. The PF-ODE, $dY_t = \left[B(t,Y_t)- \tfrac12 A(t) \rho_{\mathcal{H}_Q}^{\mu_t}(Y_t)\right] dt$, preserves the law of the original SDE and enables faster sampling for function-generation tasks including PDE solutions by reducing the number of function evaluations. Empirical results on 1D function generation and PDE tasks (diffusion-reaction and heat equations) show that PF-ODE sampling achieves comparable or superior sample quality at lower NFEs compared with time-reversed SDEs, with notable improvements in stability and fidelity. The study provides practical implementation details (e.g., using Fourier Neural Operators and Bessel priors) and highlights the potential of infinite-dimensional diffusion modeling for efficient functional generation and PDE inference, while outlining future work on faster sampling, knowledge distillation, and discretization-error analysis.
Abstract
Recent advances in infinite-dimensional diffusion models have demonstrated their effectiveness and scalability in function generation tasks where the underlying structure is inherently infinite-dimensional. To accelerate inference in such models, we derive, for the first time, an analog of the probability-flow ODE (PF-ODE) in infinite-dimensional function spaces. Leveraging this newly formulated PF-ODE, we reduce the number of function evaluations while maintaining sample quality in function generation tasks, including applications to PDEs.