Stochastic Process Learning via Operator Flow Matching

We appreciate the reviewer for the thorough review, detailed summary, and constructive suggestions. We provide our responses below

W1. Dense presentation and clarity

We appreciate the reviewer’s observation that presenting a rigorous mathematical framework spanning several disciplines within the space constraints of the main text is challenging.

To further improve clarity for readers from diverse backgrounds, we have (i) added an appendix section covering neural operators (operator learning), flow matching, Gaussian measures on function spaces, and the Cameron–Martin theorem to provide readers with essential background; and (ii) included additional explanatory text and proof sketches to better convey intuition.

W2. Novelty relative to recent work

Great comment. Following the reviewer’s suggestion, we have carefully revised the text to better communicate our contributions. The updated text now more clearly highlights the specific novelty of OFM, which clarifies how our work builds upon, yet is distinct from, recent literature. Indeed it is subtle but in fact a very important difference.

W3. Theoretical assumption and scope

Thank you for raising this point. We now state explicitly in the main text that “we assume the target measure has full support on the same space as the reference GP.”  This absolute‑continuity assumption is standard in functional generative‑model literature (see Kerrigan et al., 2023; Lim et al., 2023; Seidman et al., 2023; Schneider et al., 2024), and we follow that convention here. For verifying whether processes lie in the Cameron-Martin space of the GP prior, we refer to the detailed proof and discussion in Appendix C of Lim et al. (2023, see refs below). We have brought this point into the main text.

Regarding flow matching, we emphasise that although the conditional path $\mu_t(\cdot\|z)$ is Gaussian, the marginal path obtained after integrating over $z$ is a mixture of Gaussians (weighted by the target data distribution) and therefore need not be Gaussian. This mirrors the finite‑dimensional case: in standard flow matching (Lipman et al., 2022, Eq. 11) and optimal‑transport flow matching (Tong et al., 2023, Thm. 2.1), a Gaussian conditional is paired with an unconstrained—and often highly non‑Gaussian—target distribution, such as images, audio, or video datasets. The resulting Gaussian mixture faithfully captures the complex structure of the target law. We have added an appendix section that reviews flow matching and elaborates this point.

Reference :

Kerrigan et al “ Functional Flow Matching” 2023

Lim et al, “Score-based Diffusion Models in Function Space” 2023

Seidman et al “Variational Autoencoding Neural Operators” 2023

Schneider et al “An Unconditional Representation of the Conditional Score in Infinite Dimensional Linear Inverse Problems” 2024

Lipman et al, “Flow Matching for Generative Modeling” 2022

Tong et al, “Improving and generalizing flow-based generative models with minibatch optimal transport” 2023

W4. Empirical evaluation details

Following Reviewer 2aqv’s suggestion, we strengthened the 1D GP regression experiments by regenerating the training and test sets with more complex kernels—Gibbs / non‑stationary, Rational Quadratic kernel. OFM still outperforms all baselines on these cases as shown in the table (replied to Reviewer 2aqv)

We acknowledge the current limitation in statistical robustness and plan a full multi‑seed validation in future work. For the baseline set‑up, Appendix A (Q5) explains why NDP is omitted from the quantitative tables and provides a detailed theoretical comparison instead in Appendix N

On computational cost, the revised draft includes wall‑clock and peak‑memory numbers, together with a discussion of practical mitigations—e.g. mixed‑precision FNO kernels (Tu et al., 2023) or low‑step ODE solvers such as DPM‑Solver (Lu et al., 2022)—that can substantially reduce GPU time and memory.

Reference.

Tu et al, “Guaranteed Approximation Bounds for Mixed-Precision Neural Operators” 2023

Lu et al, “DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps” 2022.

W5. Lack of the ablation in main text

To address this concern, we have reordered several figures for a clearer narrative flow and inserted a new ablation table in the main body.

W6. Evaluation on realistic tasks

Thank you for this valuable suggestion. We agree that this is a critical direction for future work. We intend to collaborate with domain experts to build a comprehensive benchmark of real-world datasets. This benchmark will specifically include challenging settings with non-i.i.d. and non-Gaussian noise. Developing this benchmark and extending our framework accordingly is a significant undertaking that warrants its own dedicated investigation. We are confident that this future work will enable more rigorous validation and help spur progress across many scientific and engineering fields.

Q1. Discuss the limitation in the main body rather than in the Appendix N

Thanks for the suggestion, we have moved the discussion of the model's limitations from the appendix into the main text in the updated draft to ensure it is more prominent.

We appreciate the reviewer’s thoughtful feedback. Our responses below seek to clarify all concerns and potential misunderstandings

W1. The paper lacks mathematical rigor and clarity in presentation. a) it is not immediately clear that Eq. 14 is a solution to Eq. 11. b) unclear likelihood estimation tool being referenced

Thank you for these suggestions. Indeed, this paper presents a rigorous framework that combines several diverse topics to tackle the important problem of learning stochastic processes. As Reviewer sYHe also noted, the main text length limit of the paper has made it particularly challenging to communicate the relevant details from each of these fields. Nonetheless, we have tried our best to address this issue as follows. To further improve clarity for readers from diverse backgrounds, we have (i) added an appendix section covering neural operators (operator learning), flow matching, Gaussian measures on function spaces, and the Cameron–Martin theorem to provide readers with essential background; and (ii) included additional explanatory text and theorem-style statements to better convey intuition.

To address the reviewer’s specific concerns:

a) To clarify, Eq. 11 is the training objective, which requires an explicit form for $\mathcal{G}_t(u_t|z)$. Eq. 14 is then used to provide the explicit $\mathcal{G}_t(u_t|z)$ based on (derived from) the constructed probability path. We provided the full proof and derivation for Eq.14 in Appendix G (as mentioned in line 177, above Eq 14). To further avoid any confusion, we have added a proposition before Eq. 14 to clarify this exactly in the updated draft.

b) The phrase "the likelihood estimation tool described above" refers to the Hutchinson trace estimator discussed in Section 3.2. We have revised the text to remove this ambiguity.

W2. The paper lacks a high-level explanation of flow matching

To address this, we have added a dedicated section to the appendix providing a background and conceptual overview of flow matching.

In brief, flow matching is a state-of-the-art generative paradigm that learns to transport samples from a simple distribution to a complex data distribution. Its strong performance and scalability have led to its adoption in highly challenging domains, including recent large-scale models for video generation, in-context image generation (e.g., FLUX.1 Kontext, 2025), and protein ensemble generation (Jing et al., 2024). We selected flow matching as the foundation for our prior learning due to its numerous advantages. It has deep connections to physics (continuity equation) and offers simulation-free training with a simple learning objective. This generally leads to faster training and sampling compared to diffusion models, while also being fully invertible and highly scalable.

Reference:

“FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space” 2025

Jing et al, “AlphaFold Meets Flow Matching for Generating Protein Ensembles” 2024

W3. While an ablation study is included, a broader sensitivity analysis would improve the robustness of the empirical evaluation

We thank the reviewer for this constructive suggestion and agree that a broad sensitivity analysis would further strengthen our empirical evaluation. Given the extensive scope of this paper, which already spans multiple fields and involves a two‑phase framework with many interacting components, we prioritized demonstrating the core viability and performance of our approach across different tasks. Therefore, while we acknowledge its importance, a comprehensive sensitivity analysis across all components was beyond the scope of this initial investigation. We consider this a valuable and important direction for future work.

W4. The paper only evaluates on synthetic GP datasets, real-world dataset such as MNIST and CelebA are missing

We appreciate the opportunity to clarify our experimental evaluation.

While our paper features GP datasets, the evaluation also includes several highly non-Gaussian benchmarks: $**2D Navier-Stokes**$, high-resolution simulations of $**black holes**$, and the $**MNIST-SDF**$ dataset (signed‑distance functions extracted from MNIST digits). For these non-Gaussian functional datasets, the standard likelihood metric is not applicable for OFM since the constant evidence term is intractable (as detailed in Appendix A, Q5), so we provide qualitative comparisons. These visualizations show that posterior samples from OFM align closely with the ground truth and preserve key features (e.g. density and swirling patterns of the blackhole) far better than those from baselines.

To further strengthen the empirical section, and following Reviewer 2aqv’s suggestion, we have also added GP regression tasks using more complex kernels (Gibbs (non-stationary) kernel and Rational Quadratic kernel (RQK)). As shown in the following table, these new results confirm that OFM consistently outperforms all comparative methods across all metrics.

Q1, how Eq. 16 is computed in practice and why its complexity is claimed to be O(m)

The divergence term in Eq. 16, which is the trace of a Jacobian, is computationally expensive to calculate exactly. We compute it in practice using the Hutchinson trace estimator, a standard, unbiased technique for neural ODE, first used by Grathwohl et al. (2018), which requires only one reverse‑mode autodiff pass through the network, so the time complexity is linear w.r.t cardinality O(m). Below is the PyTorch‑style pseudocode used in our implementation; the full script is included in the supplementary ZIP.

# samples: target sample repeated n_noise times along the batch dimension.
# choose Hutchinson noise

if hutchinson_type == "Gaussian":
    noise = torch.randn_like(samples)              
elif hutchinson_type == "Rademacher":
    noise = (torch.randint_like(samples, 0, 2).float() * 2 - 1)  

with torch.enable_grad():
    samples = samples.detach().requires_grad_(True)
    
    # time condition
    t_vec = torch.full((samples.size(0),), t, device=samples.device)

    # drift gθ(t, x)
    drift = model(t_vec, samples)                       

    # Hutchinson estimate of divergence
    f_eps = torch.sum(drift * noise)              
    grad = torch.autograd.grad(f_eps, samples,
                               retain_graph=True)[0]   
    div_estimate = torch.sum(grad * noise,
                             dim=tuple(range(1, grad.dim())))

Reference:

Grathwohl et al., "FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models." 2018.

Q2. It would be interesting to see how the flow matching framework performs on denoising tasks

This is a good point. We discussed the technical connections between these two approaches in Appendix A (Q3). While we provide this comparison, we agree that a comprehensive empirical evaluation on denoising-specific tasks is a valuable next step, which we leave for future work

We appreciate the reviewer’s valuable comments and positive take on our work. We provide our feedback as follows

Weakness 1 & Suggestion 1. The presentation is dense and contains very advanced material. The suggestion is to add a basic introduction to neural operators and flow matching in the appendix, focusing the main paper on building intuition for the model and its usage.

Thank you for these excellent suggestions. Indeed, this paper presents a rigorous framework that combines several rather diverse topics to tackle the important problem of learning stochastic processes. As Reviewer sYHe also noted, the main text length limit of the paper has made it particularly challenging to communicate the relevant details from each of these fields. Nonetheless, we have tried our best to address this issue as follows.

To further improve clarity for readers from diverse backgrounds, we have (i) added an appendix section covering neural operators (operator learning), flow matching, Gaussian measures on function spaces, and the Cameron–Martin theorem to provide readers with essential background; and (ii) included additional explanatory text and proof sketches to better convey intuition.

Weakness 3. The experiment section is a little weak, add more complex examples (e.g., more complex GP kernels would be easy to test.

Following your suggestion, we have brought the important experimental findings from the appendix into the main text, and have subsequently expanded our 1D Gaussian Process (GP) experiments with more complex kernels for generating training datasets and regression tasks. Specifically, we now include results for the Gibbs (non-stationary) kernel and Rational Quadratic kernel (RQK) , As shown in the updated table below, our OFM model consistently outperforms all baselines on these tasks.

Suggestions 2. The figures and tables in the experiments section are placed rather chaotically, and table 1 is never mentioned, which is all somewhat distracting. It would help the reader if figures and tables were placed on the same page as where they are mentioned.

Thanks for pointing out ways to improve the flow of the paper. To address this, we have moved Table 1 to the appendix and have carefully reordered and resized several figures to place them closer to where they are discussed in the text. We also respectfully note that Table 1 was in fact referenced in Section 3.1 (Line 190).

Q1. You mention that you extend neural operators to map between collections of points. Can you clarify how this is different from a standard flow?

Thank you for this question. A standard flow‑matching model learns a transport map between two distributions defined on a fixed grid (e.g., a pixel lattice). Consequently, it can only generate samples at that specific resolution

In contrast, OFM learns a map between two stochastic processes in function space. This endows it with several properties that standard flow matching lacks:

(i) Resolution agnosticism. OFM learns a transport map between two distributions defined on any given collection of points, without regard for the number of points, or their locations in the domain. This enables capabilities like zero-shot generation without retraining.

(ii) Stochastic process consistency: OFM respects the metric of the underlying space, ensuring that points close to each other in the input domain have appropriately correlated values in the output distribution. This is part of learning a valid stochastic process, which also satisfies crucial theoretical properties like the Kolmogorov extension theorem (i.e, consistency under marginalization)

(iv) Convergence to a Continuous Function: As the collection of query points becomes denser within the domain, the output of OFM converges to the underlying continuous function.

(v) Backwards Compatibility: When evaluated on a fixed set of points, OFM recovers the behavior of a conventional (finite-dimensional) flow matching up to a linear transformation.

Q2. In Eq (11). What is the argument for replacing the suprema in Eq. (8) with expectations? When is this valid?

The supremum in Eq. (8) represents a worst-case error over the entire function space, which is computationally intractable to optimize directly. We therefore relax this hard constraint by replacing the supremum with an expectation, as formulated in Eq. (11). This is a common empirical consideration.

Instead of minimizing the worst-case error, our objective becomes minimizing the tractable average-case error across the distribution. Minimizing the error on average provides a strong practical incentive for the model to perform well across the entire function space, thereby effectively reducing the worst-case error. Such replacement is always valid under the weaker goal. The validity of this approach as a tractable proxy is further confirmed by our empirical results, which show it successfully guides the model to learn the intended functional mapping (Appendix K). In the revised manuscript, we have also added text to explicitly acknowledge the limitations of this relaxation.

Q3. OFM posterior samples in figure 8, 10, 11 look noisy, any ideas why this could be and how one might mitigate it

Posterior sampling in our framework relies on the unbiased Hutchinson trace estimator: we average the score under $n_{noise}$ (Appendix M) Gaussian noise draws to approximate the log‑determinant term. While the estimator is inexpensive in time, it is memory‑bound because all noise realizations must be processed in parallel. Owing to GPU memory limits we used only  $n_{noise}$=8 in Figures 8, 10, 11, which introduces noticeable variance and in turn yields larger stochastic gradients in each SGLD step (Algorithm 1), manifesting as visual noise..

There are several clear paths to mitigate this:

(i) Carefully adjusting the step size ($\eta_t$) and temperature ($T$) in the SGLD algorithm to prevent overly large gradient updates.

(ii) Running the trace estimator across multiple GPUs or using mixed‑precision FNO for memory saving (Tu et al, 2023).

(iii) Using more efficient ODE solvers. Our current implementation uses a high-order ODE solver (dopri5) that requires over 100 steps for sampling. Switching to a more modern and efficient solver, such as the DPM-Solver (Lu et al., 2022), could potentially reduce the number of evaluations, freeing up significant GPU resources that could be allocated to increasing $n_{noise}$.

Reference:

Tu et al, “Guaranteed Approximation Bounds for Mixed-Precision Neural Operators” 2023

Lu et al, “DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps” 2022.

Q4. The GP experiments, the training set seems to be constructed from samples from the sample GP. Is OFM able to learn from samples from GPs with different hyperparameter

There may be a slight misunderstanding of our experimental setup, and we appreciate the opportunity to clarify. OFM is designed to learn the distribution of a single target stochastic process (GP or non-GP). The training set and test examples must therefore come from the same underlying process.

The model first learns this specific process during the prior learning phase. It can then perform inference for that learned process, but it cannot generalize to a new process without being retrained.

Q5. Appendix K that you sample the training set at 256 input locations, but table 3 says use a batch size of 1024 during training

Thank you for the question. These two numbers refer to distinct stages of our framework and don't contradict one another. The batch size of 1024 in Table 3 specifies the number of independent function samples used in each training step during the prior learning phase. In contrast, the 256 input locations mentioned in Appendix K refers to the number of points used to discretize a single underlying function during the posterior sampling phase. In short, 1024 is the number of functions in a training batch, while 256 is the resolution at which we evaluate an individual function.

We appreciate the reviewer’s valuable comments and positive take on our work. We provide our feedback as follows.

W1. Important insights are pushed to the appendix.

Thanks for the suggestion, we have moved the overview of ConvCNP from the appendix into the main text so that readers can understand its role without consulting supplementary material.

W2. Excessive Data requirements

This is an important observation. We now make the trade-off explicit in the main text by adding: “Although our approach attains state-of-the-art accuracy, it relies on substantially larger training datasets than are customary for GP-based models.”

Q1. Issues related to Scalability

The primary bottleneck is the general difficulty of learning operators for functions defined on high-dimensional domains (e.g., 3D space plus time). This is a foundational challenge in the operator learning field, which is still evolving to address both computational limitations and dataset availability

Applying OFM to more complex problems will require key advancements on two fronts. First is the continued development of more scalable operator architectures that are both memory- and compute-efficient in higher dimensions. Second, and equally important, is the creation of large-scale, curated datasets for high-dimensional scientific problems, which are essential for training these more advanced operators.

It's important to emphasize, however, that the current 1D and 2D settings are not merely toy problems; they encompass many critical and valuable real-world applications, such as time-series forecasting, audio processing, and tomography. Our work provides a robust foundation for these domains, while extending the framework to higher dimensions remains a key future goal dependent on these broader advances in the field.