Flow Matching Neural Processes

TL;DR

This paper introduces a flow matching neural process model that simplifies implementation, enables efficient conditional sampling with an ODE solver, and outperforms previous neural process methods across multiple data benchmarks.

Contribution

The paper presents a novel flow matching-based neural process model that improves simplicity, efficiency, and performance over prior neural process models.

Findings

Outperforms previous neural process models on synthetic, image, and weather data benchmarks.

Allows controllable tradeoff between accuracy and computation time via ODE steps.

Enables sampling from conditional distributions without auxiliary conditioning methods.

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.