Flow Matching Neural Processes

TL;DR

FlowNP solves the neural process conditioning problem with flow matching, using a transformer to predict flow velocities for all target points in parallel and enabling ODE-based sampling and likelihood estimation. By modeling conditional distributions directly and amortizing conditioning, FlowNP achieves state-of-the-art results on synthetic GP benchmarks, 2D image datasets, and real-world ERA5 weather data, while offering a simple, implementable alternative to autoregressive and diffusion-based NP methods. The approach preserves exchangeability, approximates consistency through NP training, and provides a controllable trade-off between accuracy and runtime via the number of ODE steps. Overall, FlowNP delivers fast, coherent conditional sampling and likelihood computation for function-valued data, with practical impact across domains requiring flexible conditional uncertainty modeling.

Abstract

Neural processes (NPs) are a class of models that learn stochastic processes directly from data and can be used for inference, sampling and conditional sampling. We introduce a new NP model based on flow matching, a generative modeling paradigm that has demonstrated strong performance on various data modalities. Following the NP training framework, the model provides amortized predictions of conditional distributions over any arbitrary points in the data. Compared to previous NP models, our model is simple to implement and can be used to sample from conditional distributions using an ODE solver, without requiring auxiliary conditioning methods. In addition, the model provides a controllable tradeoff between accuracy and running time via the number of steps in the ODE solver. We show that our model outperforms previous state-of-the-art neural process methods on various benchmarks including synthetic 1D Gaussian processes data, 2D images, and real-world weather data.

Figures (9)

• Figure 1: FlowNP: we model a probability flow of a stochastic process where at each step our model takes the values of observed context points from the given function (ctx), and an intermediate value of the target points (tgt) at time t and predicts the velocities of the target points. With an ODE solver this can be used as a model of $p\left( y^\mathtt{tgt} | x^\mathtt{tgt}, \{ x^\mathtt{ctx}, y^\mathtt{ctx} \} \right)$ to generate samples or compute likelihoods.
• Figure 2: Samples from models trained on an RBF kernel (left) and a Matern-$\frac{5}{2}$ kernel (right). ANP captures the uncertainty mostly through its predicted local variance (shaded area) while TNP, NDP and FlowNP can generate coherent samples that cover the global uncertainty. In contrast to NDP, FlowNP generates conditional samples directly without needing auxiliary conditioning methods such as guidance. In contrast to TNP that generates the samples point-by-point, FlowNP generates all points in parallel, resulting in faster and smoother sampling.
• Figure 3: Conditional EMNIST samples generated by TNP, NDP and FlowNP that were trained only on subsets of pixels. For each model we show 4 samples. FlowNP generates sharp and diverse samples faster than other models.
• Figure 4: Visualization of temperature and wind prediction using conditional samples generated by FlowNP on two held-out data points from ERA5. The model generates coherent samples based on the context points. As more context points are given (bottom row) predictions become more accurate.
• Figure 4: Comparing different ODE solvers on the RBF GP data.