Stochastic Process Learning via Operator Flow Matching

Authors: Yaozhong Shi, Zachary E Ross, Domniki Asimaki, Kamyar Azizzadenesheli
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Neural Operator, Flow Matching, Stochastic Process, Uncertainty Quantification

Abstract:
Using neural operators, we propose a novel framework for stochastic process learning across arbitrary domains. In particular, we develop operator flow matching for learning stochastic process priors on function spaces. Operator flow matching provides the probability density of any finite collection of points, and enables mathematically tractable functional regression at new points with mean and density estimation. Our method outperforms state of the art models at stochastic process learning, functional regression, and prior learning.

Decision: reject
Reviewer scores: 3, 5, 3, 8

Reviews and discussion:
<thread>
creation date:
1729124112377

thread_type:
official_review

<thread_object>
creator: reviewer_rqau

<content>
summary:
This paper propose operator flow matching (OFM) for stochastic process learning, which generalize the FM to learn functional vector field and hence we can obtain samples from infinite-dimensional function space. They also provide an efficient way for computing probability density for any finite collection of points, which enables trctable functional regression at new points. In the experiments part, they illustrate their OFM and compare it to other methods, using 1D and 2D examples.

strengths:
1. The paper is well written. The idea is clear and easy to follow.
2. Extending FM to functional space can be be useful even beyond stochastic process learning.

weaknesses:
Maybe the experiements part is too simple, as the dimension (in terms of both $n$ and dimension of u) is too low? There may be some computational issue for practice. Anway, this paper provides a framework for functional space extension, and these experiements are enough to illustrate the idea. Therefore, I don't take points off for this point.

questions:
- Q1. the experiement setting: let's take 1D-GP experiment for example. Here $n = 6$, and use the same $(x_1,...,x_6)$ for each of 20000 trainings? If I'm correct, the current OFM can only apply to repeated trials in scientific settings (and in 2D image, apply to samples from the same location)? 

- Q2. In practice, if there are many observations ($n$ is large) and dimension of $u\in \mathbb{R}^d$ is high, will FNO still feasible? Since now the dimension for OFM should be $nd$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729827555314

thread_type:
official_review

<thread_object>
creator: reviewer_pmpj

<content>
summary:
The authors explore flow matching in infinite-dimensional spaces and propose a training objective to optimize a neural network-parameterized vector field (or conditional vector field). They further introduce a resolution-free Bayesian universal functional regression framework. This approach is then applied to a range of GP and non-GP regression tasks.

strengths:
* The paper is well-motivated and interesting.
* The development of algorithms for infinite-dimensional flow matching and conditional flow matching is interesting.
* The authors introduce a Bayesian functional regression algorithm.

weaknesses:
* In Sections 4.1 to 4.2, the authors introduce an infinite-dimensional flow matching approach, termed OFM, along with its conditional flow matching extension. Given that recent work has tackled flow matching for the conditional optimal transport problem formulated in Hilbert spaces (and hence in infinite dimensions) [1], the primary novelty here seems to center on the Bayesian UFR discussed in Section 4.3. If the proposed methods may differ from [1], further discussion compares with [1] would be appreciated.
* Despite the definition of the proposed OFM on infinite-dimensional spaces, the authors formulate the UFR problem over a finite collection of data points, assuming a density over this finite set. I understand that real-world data is always a finite number of evaluations; however, having established an infinite-dimensional framework, justifying a density function due to finite data points raises some questions. Would it not be more consistent to define the distribution as an approximation based on finite data points within this infinite-dimensional setting? In such a case, since Lesbesgue measure is not available in infinite-dimensional spaces, further discussion on defining a density function under these circumstances would be helpful. Furthermore, if this is the case, I am also curious whether the Hutchinson trace estimator proposed by the authors is well-defined in Hilbert spaces.
* In Section 5, the authors compare several methods, including GP and NP approaches, for GP regression. However, regarding the NP experiments, it appears that some baselines may be outdated. In particular, comparing with models like NDP [2], which also propose infinite-dimensional generative models, seems crucial. Demonstrating strong performance against these more recent models would be beneficial. Moreover, recent NP models tackle not only 1D problems but also more complex tasks, such as predicting full RGB-channel images from partially observed 2D domains. This image completion task is challenging, as it lacks the smooth properties typical of PDEs, yet these models show promising performance. I am curious to know if the proposed approach can also handle such complex image functions effectively.

```
[1] Kerrigan et al., Dynamical Conditional Optimal Transport through Simulation-Free Flows
[2] Dutordoir et al., Neural Diffusion Processes
```

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730636426122

thread_type:
official_review

<thread_object>
creator: reviewer_mbxf

<content>
summary:
This paper proposes a framework for learning stochastic processes parameterized by neural operators by generalizing flow-matching to function spaces. They do so by generalizing the conditional flow-matching objective to function spaces. This enables tractable density estimation on any new collection of points.

strengths:
1. I think the problem of generalizing flow-matching to function spaces is an interesting one.

weaknesses:
## Clarity

1. I found the notation throughout the paper to be extremely unclear. One example of this would be that throughout section 3, $\mathcal{G}$ is defined to be an operator but always applied to a dataset instead of a function. And there are several others throughout the paper. I am okay with some abuse of notation where it helps with readability but in my opinion, this is not the case here. I found it quite hard to keep track of when we are talking about objects in a function space and when we are talking about objects in a domain. I believe the readability can be improved significantly.

2. The authors don't make any effort to disentangle their novel contribution from what is already known in the literature. They have a relatively long background section and ideally, the distinction between the methods section and the background section would have served to make it clear what is new in this paper and what is already known in the literature. But throughout the methods section, they refer to results known in the literature without clarifying whether they are being applied, adapted, or extended, making it hard to identify the paper’s unique contributions.

## Novelty

The objective proposed in this paper is a relatively straightforward generalization of the conditional flow-matching objective. While the extension to function spaces may offer some theoretical appeal, it doesn’t represent a particularly groundbreaking advancement beyond established work as other works have already generalized flow-matching to function space.

## Experiments
Basically, all the details of the experimental setup are missing. There are no details about the FNO used to parameterize the operator, no wall-clock times for training, or the number of Hutchinson samples used from equation 22. Table 1 is also missing error bars.

questions:
See weaknesses. Additionally, why is the following equation in page 3 true:
$$ \mathbb{P}( \{ u(x_1) \dots, u(x_m) \} ) = J\mathcal{G}\Bigr\rvert_{ \{ a(x_1)\dots,a(x_n) \} }\mathbb{P} \{ a(x_1)\dots,a(x_n) \} $$
</content>
</thread_object>

</thread>
<thread>
creation date:
1730647461653

thread_type:
official_review

<thread_object>
creator: reviewer_mzqu

<content>
summary:
Using neural operators, this paper extends flow matching based continuous normalizing flows to infinite dimensions, to learn stochastic processes. By virtue of finite-dimensional densities being tractable, these stochastic processes can be used as prior distributions for subsequent Bayesian inference, which authors propose doing using stochastic gradient Langevin dynamics.

The approach is empirically tested on synthetic and more realistic datasets, showing improvements.

strengths:
- I think that the overall idea of using infinite-dimensional flow matching for learning stochastic process priors is a good one.
- The empirical advantage over GP baselines is convincing in the black hole and MNIST setups.
- Whenever authors use some ideas, they mention the sources of those ideas (provide citations).

weaknesses:
The paper brings together numerous components from modern literature to create something new. In my opinion, it's not hard to imagine that something along the lines of this paper would be possible. Working out the details, again in my view, doesn't require a lot of new ideas but might demand significant effort. I believe that papers of this nature can be quite useful and make worthy publications. However, they should (1) provide high-quality code, so future researchers don't have to replicate the work themselves, and (2) clearly present the "diff": the differences compared to existing methods and tools, highlighting the details that proved challenging, and aspects potential users should be aware of. Unfortunately, I think this paper falls short in these areas:
- Writing is sloppy (a far-from-exhausting list of examples is given below in the “Questions” section, as evidence). It feels like nobody even tried to proofread the text before submission.
- There is no clear “diff”, instead, the text concentrates on reviewing the prior work or something that only marginally differs from the prior work.
- There is no code in the supplementary.

I don’t think that the method is unworthy of publication. However, I think that the paper, in today's form, is. The text should be majorly improved, the code should be provided.

questions:
Questions/major comments (some are probably consequences of sloppy writing):
- In Figure 1, it looks like a Gaussian process is just chosen poorly. First of all, it does not look smooth enough (=> e.g., use RBF instead of Matérn, use a larger length scale). Second, since here you obviously have a diagonal trend, you could consider a matrix ARD: when instead of one length scale parameter $l$ that transforms $x \to x/l$, you learn a matrix $A$ that transforms $x \to A x$.
- Line 091, you say that GP parameters in deep GPs are “manually tuned”. This is plainly not true, they are typically optimized during the variational inference procedure.
- Lines 092-095. “Warped GPs (Kou et al., 2013) and transforming GP (Maroñas et al., 2021) methods use historical data to learn a pointwise transformation of GP values and outperform deep GP type methods).” This is a strong and controversial claim without any evidence to support it.
- Lines 096-097, I don’t think neural processes are a variational inference approach, this is certainty not how they were originally introduced.
- You are not careful when going into the infinite-dimensional setting, and going infinite-dimensional is perhaps the most important contribution of the paper. For example, in Equation (11) you use the gradient, but you don’t talk about the precise meaning of this operation in infinite dimension.

Examples of sloppy writing:
- The content goes over the right margin on multiple occasions: in the display-style formula at the bottom of page 3, in Equation (24), in Table 1.
- Figures contain barely readable font sizes.
- Lines 042-044, the clause “allowing us to learn probability priors over the stochastic process, ergo, sampling of any collection points with their associated density.” You are not learning a prior over stochastic process, you are learning a prior over functions, which _is_ a stochastic process. Also, the word “sampling” typically means “random sampling” throughout the paper, I don’t think here you mean “random sampling”.
- “Other works in this area have following the success of flow matching and diffusion models”. I guess it was supposed to be “followed”, the typo makes this sentence very hard to read.
- In the second paragraph you talk about the “universal functional regression” of Shi et al. that makes it feel like a classical approach, whereas it is a paper from 2024. Also you vaguely overload the standard term “functional analysis”.
- Line 115. The notation for a probability space and a measurable space is introduced, which I believe is never subsequently used.
- Throughout the text: it is very strange that you use the curly braces {} for vectors, instead of the standard () or somewhat less standard [], especially since curly braces are standard for something else, for sets.
- Lines 151-154 The sentence “More specifically, for m ≥ n points at which the function is to be inferred, … display-style formula …” does not make sense, grammatically.
- At the time Figure 1 is first referred to, “OFM” is an undefined abbreviation.
- Line 049, “deep GP” -> “deep GPs”.

The list could go on and on.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732243902689

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Changes made in the revised draft

comment:
We would like to thank all reviewers for their insightful and valuable feedback, which has significantly helped us enhance the clarity and quality of our work. Below, we outline the key modifications made during the first revision. Changes in the main text are highlighted in blue in the revised draft for easier identification.

- Carefully revised the text, adjusted the font size for all figures, and uploaded the code to the supplementary materials. Additionally, we added a paragraph in the introduction to further clarify our contribution.

-  Added detailed analysis and comparison with related models or baselines in Appendix A.8, highlighting our contribution and potential limitations

-  Provided regression results on finer grids (currently 64x64, previous 32x32) for Navier-Stokes, Black hole, MNIST-SDF examples.
-  Added detailed experimental setup and parameters used in Appendix A.7
-  Added the SGLD algorithm used in this study in Appendix A.5 
-  Added one co-domain regression experiment in Appendix A.4
-  Added more posterior sampling results in Appendix A.3
-  Added the background and proof for generalizing flow matching to stochastic process with Kolmogorov extension theorem, provided in Appendix A.2
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a flow matching method in infinite-dimensional function spaces. The proposed Operator Flow Matching (OFM) framework extends standard flow matching methods to infinite-dimensional spaces by constructing a framework for learning functional vector fields. The suggested stochastic process enables the computation of finite-dimensional densities, making it applicable as a prior distribution for Bayesian inference. The authors propose inference using Stochastic Gradient Langevin Dynamics (SGLD). While the paper presents an important and highly intriguing approach, many reviewers pointed out mathematical issues and presentation problems. Due to numerous typographical errors and unclear points in the mathematical formulations, substantial revisions are necessary before the paper can be considered for publication. I strongly encourage the authors to refine these aspects and resubmit the paper.