Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problems

Thank you for acknowledging our novelty. We are sorry that the paper is difficult to follow. If you can point out specific parts that are unclear, we will try our best to explain them in detail.

The running $L$ is usually assumed to be convex, which makes the problem well-posed. An alternative assumption is that the Hamiltonian is convex in the control variable. The mean-field control problem is not well-posed without this assumption, and the solution is not finite.

We agree that comparing our higher-order ODE approach with other methods would provide valuable insights, particularly regarding computational efficiency and performance trade-offs. Our method's strength lies in addressing second-order structures inherent in control problems, mainly when applied to mean field control problems. We are conducting additional numerical experiments and will include a comparative analysis in the revision.

The LQ framework encompasses many problems, as illustrated in our paper. While our current examples focus on LQ problems, our method can handle non-LQ problems. The challenge lies in finding reliable reference solutions for non-LQ problems, which makes quantitative error evaluation difficult.

When we stated that it is probably more important to train the neural network properly, we meant that our method's primary source of error often stems from suboptimal neural network training rather than from the numerical discretization scheme itself. In our experiments, we observed that improving the neural network training (e.g., using better regularization techniques) had a greater impact on performance than switching from a forward Euler scheme to higher-order numerical schemes. This suggests that the neural network's capacity to approximate the underlying dynamics accurately is a critical factor in reducing error. We will clarify this point in the revised paper and provide more details on the training strategies we found most compelling.

Thank you for your suggestion. We agree that comparing our method with existing approaches is essential for evaluating its performance. We will compare scalability, accuracy, and computational efficiency with other relevant methods, including Gaussian process regression, kernel-based methods, or neural networks used for differential equations.

Adding theoretical analysis is an important direction. We aim to tackle these theoretical aspects as part of ongoing research and will include initial theoretical insights in future revisions when applicable.

Thank you for your thoughtful feedback. We have addressed your points (in Weakness) below.

1. Our primary contribution is the novel application of neural ODEs to the mean field control (MFC) problem, an area where, to our knowledge, no prior work has applied these techniques. Additionally, we introduce an HJB regularization mechanism tailored explicitly to our modified MFC problem (see Equation 17). This regularization, designed for second-order score functions, is a key innovation that enhances the accuracy of our method in solving MFC problems. In this regularization, studying evolution in second-order score functions is essential.

2. We acknowledge that a more detailed comparison with existing works on score function computation and MFC solutions would clarify how our method relates to the state-of-the-art. In the revision, we will include a focused discussion on key works, including classical methods for computing score functions and machine learning approaches relevant to second-order MFC.

3. We agree that empirical comparisons are crucial to demonstrate the performance of our approach. While we cannot access the codes for most of the methods referenced in the related works, we can provide a more detailed qualitative comparison in the revision. We are also planning to integrate some baseline methods for future experiments, allowing us to offer quantitative comparisons in subsequent paper versions.

4. You are correct that the dimensionality of our current experiments is limited (maximum 10 dimensions). Solving more challenging high-dimensional problems is a natural next step and part of our ongoing work. While the current paper focuses on foundational examples to illustrate our method’s core features, our framework is well-suited for higher-dimensional problems, including those typically encountered in generative models. We will clarify this direction in the revised paper and include preliminary results.

5. Although we lack access to runtime measurements for competing methods, we will include empirical runtime data for our proposed approach in the revised manuscript. Most of our examples are not on large scales, which takes 1 to 10 minutes to run a single training.

(Regarding your additional question in Questions.) As we mentioned in Section 6, one of the potential applications is the generative model, which has higher dimensions and is more challenging. The connection between generative AI and mean field control is at the mathematical formulation level. In the sense of optimization or optimal control problems, one can represent several state-of-the-art methods, such as neural ODEs and generative adversary networks, time-reversible diffusions, into mean field control problems. However, in numerical implementations with data samples, as reviewers suggested, one needs a simple formulation to compute the control conditional on samples. This is the future research direction we are working on. We would like to collaborate with related experts in this direction.

1. Thank you for mentioning the related work. We will include https://arxiv.org/pdf/2206.00860 when we introduce the ODE system. The authors derives the ODE systems up to third order and apply fix point iteration to solve the Fokker--Planck equation, with error analysis in second order Sobolev norm. The paper https://arxiv.org/abs/2210.04296 differs from our formulation, but we will add it to our related work. We have already mentioned https://arxiv.org/abs/2206.04642 and their equation (13) in our article. Although we are not the first to propose these equations, our main contribution lies in applying this formulation to approximate and compute second-order mean field control (MFC) problems. In this area, current state-of-the-art machine learning numerical methods apply neural ODE-based approaches for first-order mean field control problems. Score-based neural ODEs are essential in computing second-order mean field control problems. For example, the proposed method helps approximate the kernel formula in the regularized Wasserstein proximal operator; see example 5.1. This is a context that allows us to leverage these equations effectively within our proposed numerical scheme. In the revision, we will emphasize this point and explain second-order mean field control problems and their relations with score functions more clearly.

2. We agree that the one-hidden-layer neural network may seem simplistic compared to more complex architectures. However, our architecture offers both conceptual clarity and good performance in the current setting. We remain open to experimenting with more advanced architectures in the future, particularly as we scale our experiments to more complex problems. This choice has proven effective for our current numerical examples, and balancing simplicity with computational efficiency is a primary goal. We leave the design of neural network structures for high-order score-based neural ODEs in the future. It should depend on solutions to mean field control problems.

3. We understand the concern about scalability, especially the cost of differentiation through Equation (11) in Algorithm 1. However, computing the score function is unavoidable in the context of mean field control problems. Our score-based normalizing flow formulation addresses this issue by efficiently computing these quantities, which we believe makes the approach scalable, even in higher-dimensional settings. While our current results are based on smaller-scale problems, we are optimistic about extending the method to larger, more complex scenarios.

4. Proposition 3 differs from the standard forward-backward PDEs typically seen in the literature, such as those in {\it Mean field games and mean field type control theory} Chapter 4. The HJB equation is a modified version tailored to incorporate the score transformation (Equation 16). This modification introduces the second-order derivative of $\log\rho$ into the HJB equation, which is not standard. Equation (3d) becomes crucial for simulating this derivative, and to the best of our knowledge, we are the first to use this modified HJB equation as a regularizer to enhance accuracy. This regularization improves performance even in the Gaussian case, providing significant value to our approach.

5. We agree that the current examples in Sections 5.1 and 5.2, which involve analytically solvable or factorizable cases (such as the linear OU process), are relatively simple. Please note that we have more complicated examples in the Appendix, such as the double moon flow matching and the double well example. Additionally, we are actively working on applying our method to more challenging examples, including those mentioned in your review. We believe our framework can be adapted to handle more complex problems and will include these results in future work.