Robust Symbolic Regression for Dynamical System Identification
Authors: Ramzi Dakhmouche
Ivan Lunati
Hossein Gorji
Venue: Accepted by TMLR ()

Abstract: Real-world complex systems often miss high-fidelity physical descriptions and are typically subject to partial observability. Learning the dynamics of such systems is a challenging and ubiquitous problem, encountered in diverse critical applications which require interpretability and qualitative guarantees.Our paper addresses this problem in the case of sparsely observed probability distribution flows, governed by ODEs. Specifically, we devise a {\it white box} approach -dubbed Symbolic Distribution Flow Learner (\texttt{SDFL})- leveraging  symbolic search with a Wasserstein-based loss function, resulting in a robust model-recovery scheme which naturally lends itself to cope with partial observability.   
Additionally, we furnish the proposed framework with theoretical guarantees on the number of required {\it snapshots} to achieve a certain level of fidelity in the model-discovery.
We illustrate the performance of the proposed scheme on the prototypical problem of Kuramoto networks and a standard benchmark of single-cell RNA sequence trajectory data. The numerical experiments demonstrate the competitive performance of \texttt{SDFL} in comparison to the state-of-the-art.

Published in Transactions on Machine Learning Research (03/2025)
Robust Symbolic Regression for Dynamical System Identification
Ramzi Dakhmouche
ramaz.dahkmoouchheeeppp...c
Institute of Mathematics, EPFL Laboratory for Computational Engineering, EMPA
Ivan Lunati
lean. Sonnllunttiinnnn Sempa.ch
Laboratory for Computational Engineering, EMPA
Hossein Gorji
Laboratory for Compulational Engineering, EMPA
Reviewed on OpenReview: https: // openreview. net/ wttp:://dd/oooo/oaa/aaaaooooaaaa -esc
Abstract
Real-world complex systems often miss high-fidelity physical descriptions and are typically subject to partial observability. Learning the dynamics of such systems is challenging and ubiquitous problem, encountered in diverse critical applications which require interpretability and qualitative guarantees. Our paper addresses this problem in the case of sparsely observed probability distribution flows, governed by ODEs. Specifically, we devise a white box approach -dubbed Symbolic Distribution Flow Learner (SDFL)- leveraging symbolic search with Wasserstein-based loss function, resulting robust model-recovery scheme which naturally lends itself to cope with partial observability. Additionally, furnish the proposed framework with theoretical guarantees the number of required snopshots to achieve certain level of fidelity in the model-discovery. We illustrate the performance of the proposed scheme on the prototypical problem of Kuramoto networks and a standard benchmark of single-cell RNA sequence trajectory data. The numerical experiments demonstrate the competitive performance of SDFL in comparison to the state-of-the-art.
1 Introduction
Complex systems often emerge as networks of physical and societal interactions, with examples ranging from epidemics to consensus dynamics, and from power grids to biological organisms. The pursuit of accurate symbolic expressions that describe the evolution of such dynamical systems, is thus of paramount importance in many areas of science and engineering, Indeed, such parsimonious model descriptions offer several advantages, including compactness, explicit interpretations and high-fidelity generalization. Unlike Macck-or approaches, they give explicit insights on the underlying processes. Most notably, they allow for safety related qualitative guarantecs such as asymptotic stability, which are crucial for critical applications.
Yet complex systems prone to partial observability, uncertainty and presence of noise, which render deriving explicit representations particularly challenging, Hence, devising robust and efficient learning methods to uncover equations directly from data -also known as symbolic regression- is of great interest.
Symbolic regression has been extensively studied over recent years in both algebraic and differen- tial equation discovery contexts. Common approaches include: sparse regression (Brunton al. 2016 Rudy Kubalík et al.] (023), sequence-to-sequen deep neural network modeling
et al. 2017 Chen al. 2021
[Becker et al.] 2022 Vastl et al. 2022 Biggio et. al. 2021), and symbolic search-based formulations (Cornforth

---

Published in Transactions on Machine Learning Research (03/2025)
& Lipson 2012 Gaucel et al. 2014 Cazenave 2013 Lu et al. 2021 Sun et al. 2023
One important
setting, which has not been studied much though, is the discovery of dynamical systems governing probability distribution flows using only a few sampled screen shots, which is particularly relevant for many applications ranging from epidemics modeling to cellular evolution prediction (Bunne et al.) 2022). Consequently, we propose to address this question, in the case of netk flows, by designing a suitable equation recovery
framework and illustrating its numerical performance on both synthetic and real-world data. This setup of symbolic learning of probability flows in networks entails several challenges:
Snapshols, arriving at random times, have to be described in a continuous-time probabilistic repre- sentation.
A robust inference algorithm enabling efficient search for an accurate solution -under partial observability- has to be designed.
Permutation invariance, which is inherent in network structures, should be incorporated into the symbolic search, in particular to reduce the computational cost.
These challenges are addressed by our technical contributions:
1) We propose trajectory inference (Bunne et al. 2022 approach based symbolic search, yielding thote-boz models with the permutation invariance property,
2) We devise a suitable loss function that leverages the robustness properties of the Wasserstein distance Villani 2009), while taking into account the limited observability of the system under study,
3) We back our devised algorithm with suitable sample complexity theoretical results,
4) We demonstrate the performance of the proposed approach on the prototypical problem of Ku- ramoto networks and a standard benchmark of single-cell population trajectory data. Our approach handles settings with low sampling frequency of the order of 100 times less for the Kuramoto model (sec Appendix
E.4 than what required by the current symbolic regression state-of-the-arl (Qian et al., 2022
Sun el al.
2023
Gao & Yan 2022
The rest of the paper is organized follows. In section 2 review related works and partially motivate our design choices. In section 3 we introduce the notions upon which the proposed algorithm is based. In section 4 we outline the main contributions of this work and provide the theoretical results that guided the design. In sections 5 we present numerical evaluations, followed by concluding romarks in section
2 Related works
Symbolic search. Symbolic regression was initially formulated as a discrete optimization problem (Augusto & Barbosa 2000; Smits & Kotanchek 2005 Cornforth & Lipson] 2012 Cazenave 2013), where the goal is to find the most accurate mathematical expression based on a predefined set of elementary operations and functions (e.g. +,-, sin, exp, ...). The mathematical expressions were represented through one-to-one correspondence with pre-order traversal trees. For the resolution, Genetic Programming (GP) heuristics [Koza 1994) were used to recover the equations underlying the training data. This method subsequently inspired applications in population evolution modeling (Bongard & Lipson 2007). prediction of solar power production (Quade et al. 20166), and Eulerian fluid flow hidden paramoterization discovery (Vaddireddy et. al. 2020). Yel, GP sulfers from a number of issues including over-fitling, brittleness to noise and poor-scalability (Brunton et al.) 2016 Petersen et al. 2019 Lu et al., 2021). A more recent approach Petersen et al., (019)) uses a Deep Reinforcement Learning (DRL) model to solve the optimization problem and generally outperforms GP-based models. However, deep learning models require large amounts of data in addition to their relative lack of full-automation, since architecture hyper- parameters have to be
2

---

Published in Transactions on Machine Learning Research (03/2025)
tuned by human expert. On the other hand, parallel well performing approach with the potential to achieve full-automation is Monte-Carlo Tree Search (MCTS). Based on the tree representation used in GP formulations, MCTS builds upon ipploration-exploitatiiio trade-off insights from sequential learning (Munos ct al. 2014p. It is designed for a stochastic sctting where data-points are costly to obtain and therefore naturally handles noisy input and scare-data settings well. Some works have applied MCTS successfully to
simple (Cazenave 2013; Islam et al., 2018 Lu et al.] 2021p and broader (Sun et al. 2023) symbolic regression problems. The former used MCTS to uncover non-linear expressions in a supervised setting, whereas the latter also applied it to Ordinary Differential Equation (ODE) discovery when the initial conditions are deterministie. In the case of network structured problems, (Shi et al. 2023 Cranmer et al., 2020) are the only works -to the best of our knowledge- which distill explicit equations. That is done by training a Graph Neural Network (GNN) and distilling equations for the message-passing operators. Additionally, only deterministic physical systems for which the whole trajectory is observed are considered. In contrast, we address a setting, where the goal is to identify the governing equations of stochastic network dynamics (with randomness in the initial condition) from the observation of a reduced number of screen shots across time.
Sparse regularization. One common property that emerges when modelling many natural phe- nomena and engineering problems is sparsity. Leveraging this fact, a number of works (Brunton et al. 2016
Schaeffer 2017 Rudy et al.] 2017 Loiseau & Brunton 2018) formulate (deterministic) dynamics discovery as a 4-regularized linear regression problem, over a predefined dictionary of basis functions. The resulting
convex optimization is then solved using a sequential thresholding scheme. In particular, [Brunton et al. 2016
demonstrate the performance of the scheme they propose, dubbed SINDy, uncovering the chaotic Lorenz
system and fluid vortex shedding ODE. Subsequently, [Boninsegna et al. [ Huang et al. (2022) extend
the approach to the discovery of Stochastic Differential Equations (SDEs), using the Kramers- Moyal formula.
Besides that, Chen et al. (2021) addresses scarce-data setting by training a neural net as a surrogate, combining data from PDEs with various boundary conditions. Once the neural network is pretrained, it is used to discover the underlying PDE using an altornating direction optimization (ADO) algorithm to solve the é--regllarized linear regression problem. From a theoretical perspeclive, [Schaeller et al. (2018) study the question of minimal number of screen-shots required to recover multivariate quadratic ODEs, in the case of random initial condition. However, they assume velocities to be known which is inconsistent with the small and scarce data setting. Overall, although usefulness of pparsity-promoting approaches has been extensively demonstrated, they still rely on prior knowledge to define their main component, namely the basis function library. Whereas, if a library of massive size is chosen, the algorithm empirically fails to hold the sparsity constraint (Sun et al.) 2023). In comparison, we apply MCTS for our symbolic search part, which is not bound by such constraints.
Sequence to sequence models. The most recent approach to symbolic regression leverages the success of deep learning in sequential data modeling, such as in natural language processing (Devlin et al.] 2018). More precisely, (Biggio et al. 2021: Kamienny et al., 2022 Vastl et al. 2022) propose a transformer-basee based architecture, which is trained to output mathematical expressions based on data-sets of fealure-prediction pairs. That is, for each input dala-set (i) the model outputs an expression e corresponding to a function fe satisfying Vi E [1,m], yi ~ fx(). One challenge though, is to generate a N which is rich enough to represent parsimonious equations that training data-sct {(((=), (=)()) j=1
are frequently encountered in practice. To achieve that, equations are generated as binary trees, following the work of (Lample & Charton] 2020 Kusner et al.] 2017) and previously mentioned symbolic search approach representations, e.g. (Cazenave 2013). On the other hand, (Li et al. 2019) combine recurrent neural network with MCTS to enforce asymptotic constraints on the learned expression, while some works
2018 Costa et al. 2020 Kubalík et al. 2023) propose to use different elementary mathematical operators (e.g. +, x, coc...etc) as activation functions for neural net architectures, while imposing sparsity on the parameters to extract interpretable analytic expressions. However, these approaches rely on the generation of large data-sets, which is prohibitively costly for high-dimensional ODF/PDEs.
Martius & Lampert 2016 Sahoo et al.
Neural ODEs. First introduced CChen et al. 2018 as a continuous counterpart to residual net-
3

---

Published in Transactions on Machine Learning Research (03/2025)
works for generative modeling, neural ODEs have demonstrated success in several tasks such as density estimation, image generation and variational inference (Grathwohl et. al., 2018). They are based on modeling the vector field defining an ODE as neural network and leveraging the adjoint sensitivity method
PPnntyygiin 2018) for efficient back-propagation. Subsoquently, they have been extended to stochastic
and latent variable settings allowing for an increased expressive power (Dupont et al. 2019 Kidger 2022
Besides, notable parallel enhancements included extensions to non-Euclidean spaces (Lou et al. 2020 Jeong
et al. 2023) by devising a manifold adjoint method. However, one of the main drawbacks of neural ODEs is the fact that they do not allow for qualitative analyses of the models, which is often needed for physical and biological systems.
3 Problem Formulation and Background
3.1 General setup
Fix T> > 0 as positive time horizon, and consider the random variable zo E Rd to be the initial condition of the state variable 2; which evolves according the ODE system
ý = f()
(1)
in time & E [0, T]. In our setting, f encodes the interactions between different components of x; over time, which are distributed according to known and fixed network topology G. We denote by Pt: the probability measure induced by xa i.e. Ty ~ Pt, and by pt the corresponding probability density. The latter evolves f4 is absolutely continuous with respect to the Lebesgue measure- according to the equation
(xzz) logPo(zo) - f Tr (f=()))x)))))
along the trajectory (#s))e(o.r] (see e-g. Chen et al. (2018)) For n 2 2 measurement times (lo, li..... (n), assume m 2 samples of each of the distributions Mro:Me :e.............................. Mr, are observed, resulting in snapshots reprosented by the corresponding empirical probability measures .f.,, Pr.,,. The qucstion addressed in this paper is that of the recovery of the function defining the ODE that governs the given dynamics.
That is, trajectory inference (Hashimoto et al. 2016; Tong et al. 2020 Bunne et al. 2022 Huguet et al.) 2022), through an explicit closed form ODE. For that matter, the solution of the ODE can be seen as the image f by the operator given by
F(): : (1, - (z())z))
where 1 (zo) is the solution of the Cauchy problem:
1 = f(v) y = 25
(2)
for given initial condition 2o E Rdd, Note that, the problem of recovering f from in general, and further assumptions are needed. Such assumptions should represent prior information that can act form of regularization rendering the learning task at hand more feasible. In case, the required prior information comes from the fact that the estimator of f is constructed by combining analytic expressions from a fixed pre-selected set. We make the assumption that the obscrvation instants fulfill Vi E (O...., n), 4 E (iT/n, (i 1)T/n. A straight-forward extension, though, can be obtained for uniformly sampled time instants (tx)ee (o....n) based on a Quasi-Monte Carlo scheme convergence argument (Niederreiter) 1978.. To translate this sctting to an optimization framework, first, we review some concepts about distances in probability spaces upon which a suitable goodness-of-fit measure is proposed.
ill- dll-po
3.2 Wasserstein guidance
Considering probability distribution flow modeling, we will need to compare predicted distributions with observed ones. Popular measures of disparity between probability distributions include the Kullback-Leiba
4

---

Published in Transactions on Machine Learning Research (03/2025)
(KL) divergence in the computational context, the total variation distance in the theoretical one, and the Wasserstein distance (Villani 2009) in both. Indeed, due to its singular theoretical properties, e.g. weak topology metrization, convexity and robustness (see e.g. MMhaajerin Esfahani & Kuhn (2018) for the latter), but also its amenability to efficient computation, the 2- 2-Wasserstein distance W2 represents a natural choice for the design of a robust trajeclory inference algorithm. Consider two measure spaces (X, μ1) and (), and denote by ((,,)) v) the set of their couplings, then W2 reads
(e(p, = min / ||a -112 (a(a,)) *ell(u.)) Xx'
where || ll is the Euclidean norm. If we had access to (JL))E the goodness-of-fit of a candidate estimator f could be suitably defined as:
L()- - I," Wa(F"() pa) de
(3)
where # po is the push-forward of the initial probability measure by the partial flow map 2o x (x)) In other words, L(f) is the time aggregated Wasserstein distance between the measure resulting from the inferred trajectory and the one resulting from the ground-truth dynamics. However, since μr is only partially known, we work with an approximation of L(f) based on the observed snapshols, as discussed in the following, after presenting the basic building blocks of MCTS. Note that, the KT-divergence could not be used here given that it assumes that the compared distributions have the same support.
3.3 Monte-Carlo Tree Search for Symbolic Search
Unlike genetic programming, which suffers from over-fitting and brittleness to noise, MCTS handles stochastic settings by design in a data-efficient way (see for e.g. (Sun et al., 2023 Lattimore & Szepesvári 2020 ). The latter property also sets it forth compared to reinforcement learning approaches, which require large amounts of data. Due the combinatorial nature of the state of possible mathematical expressions, it crucial reduce the number of required evaluations. To cope with that, MCTS targeted to symbolic search, relies tree representations of mathematical expressions, where the nodes express unitary (sin(-), cos(.)....) binary operations (t..) and leafs express variables and constants (æ,, C.......) This is formalized using the notion of context-free grammars (Kusner et al. 2017 Sun et al. 2023). Given such representation, the intuitive idea behind MCTS is to explore as few operators as possible in building the tree, yet to identify the optimum with a high probability. To achieve that, leverages insights from
Bound (UCB) algorithms in sequential learning (Lattimore & Szepesvárij 2020), where the goal translates optimizing a function which measures both accuracy and -nper-exploration of candidate operation. We refer to (Sun et al. 2023) for detailed pseudo-code, and recall the main components translated to context below:
Stochastic roll-outs. To complete partially built tree i.e. a partial expression denoted by s, the value of each potential additional node a (representing an elementary operation) needs to be obtained, in order to choose the one with the highest value. For that matter, for cach choice of node, random roll-outs of polential completions are performed. Each roll-out will either lead to a complete expression f(x), whose corresponding score function can be computed; or an expression whose length exceeds a predefined maximal length M leading to score of 0. The largest of the obtained scores is then stored in V(s, a) to be used in the next step.
Exploration-aware selection. Once all operations have been visited at least once, select the next one by maximizing a performance measure that takes into account under-exploration, and which can be taken in practice as
UCT(s,a) (c/h)))) c>
For the sample complexity result presented in prop. hold though, the scheme should based on the more involved expression for UCT described in (Shah al. 2020
5

---

Published in Transactions on Machine Learning Research (03/2025)
where as is the maximal value of operation a, given the operations constituting the current tree i.c. (4o: 01:-................. ................................ Ap-1]- Furthermore, (ss,a) is the number of times that operation was chosen at tree state with (s)) = EncA(s) (s,, a), where A(s) is the set of available operations at tree state s. Note that the square root term v (ss)/s/)/N(s. quantifies the under-exploration of operation a at state 8.
4 Technical Approach
To minimize the introduced loss function, given by equation 3 replace the neural net estimator in the neural ODE framework of (Chen et al. 2018 Grathwohl et al. 2018 Kidger 2022 with search for an explicit analytic formula for f. This is achieved through symbolic regression algorithm instantiated by MCTS (Sun et al. (202)). The devised workflow -of inferring the network dynamics from given snapshots- is illustrated Fig. 6 It is based on selecting a given ODE candidate, evaluating its performance, and using that evaluation inform future selections, as presented below. From a guarantec perspective, in addition bound on the approximation error of L(f) by its empirical estimate Lm,n(f) for given f, we quatify the sample complexily of MCTS, using regret bound-based results from bandit theory (Lattimore & Szepesvári 2020
4.1 Discrete Loss function
The convergence of the discrete version of the loss function, defined for continuously differentiable f: Rd Edd
by
Lm (f) = - Σ Wa(k"s (f) Ae
(4)
the continuous one i.e. L, is controlled by the number of snapshots and the size of the snapshot sample set. We build upon results of Fournier & Guillin (2015): Bonnans & Shapiro (2013) about finite-sample rates
of convergence of the empirical measure in the Wasserstein space, and regularity of the Wasserstein distance to obtain the following theorem, where the proof is postponed to Appendix A
Theorem 1.
Let (pe)220 have compact support and f: IRd - Rd be differentiable function. Assume that t + Me is differentiable. Then, for all m, 2 2,
1 T [D) = =o (mooaa +
E
where fi.,m for E (1....,n.) n), denotes the empirical measure corresponding to the observed realizations of each of the probability distributions (Hs,)
Proof. We sketch the proof here and postpone the more detailed version to Appendix A
Without loss
of generality, we assume T = By theorem 1 in (Fournier & Guillin 2015), given that To + F² F'o))) is contimuous and for all E {l.....7 n): He, has compact support on the Polish space Rd, there exists C2 C2 > such that
C
E
(F(F F (f)
1id/2-1)
C2 mld/2-1)
and
E
Wa(Deamm
Therefore, thanks to the triangular inequality satisfied by W2, get
E
1 mf.
L
Σ
C1 C2 mld/2+1)
n
n
6

---

Published in Transactions on Machine Learning Research (03/2025)
On the other hand, by the stability property of optimal transport as shown in Appendix leading
A
we get
C3 m
L(I-
n
4 L (f) -L(f) =0 ( m (m n
n
4.2 Symbolic flow discovery algorithm
Given the discrete loss Lm the model-recovery algorithm is based computing predicted snapshots corresponding a candidate estimate f and comparing them the observed ones, The predictions will be determined by solving the obtained ODE, i.e. equation [ where f is replaced by f, through numerical integration scheme, such as Runge- Kutta solvers. In parallel, the optimization over candidate estimates is
Salect
Compune Fiff)
siuh
Snapshats
Figure 1: Building blocks of the proposed algorithm
realized using MCTSS- symbolic search. For that matter, the goal of MCTS will be to maximize score function defined for a differentiable f. Rd - Rd by
v( 1+ Lm,
To ensure permutation invariance in the obtained expressions, which is characteristic of network systems, restrict the search space by applying permutation invariant aggregation operation (e.g. sum) to each expression before each evaluation. That is, we restrict our search to permutationninvarian expressions by design. We report an ablation study corresponding to the removal of the permutation invariance module in subsection 5.2 The pseudo-code of the algorithm is summarized in Algorithn where the main steps are illustrated in Fig. with more details Appendix D As for the implementation of Lm,,s it reduces when comparing two empirical distributions fi m D 0a and D m - 83,, to solving the following linear program
min (P.M) PEU().0)
where M (Iaa B)(l2)ss is the cost matrix and U the set of couplings between and &. Hence, its computational cost is in O(dm² + m²²), where m is the number of samples and d, the dimension of the system.
Remark 4.1
The restriction of the setup to vector fields f such that the ODE system ÿ = f(y) has well-defined solution, is implicitly made by restricting our attention states + D, defined on [0. 7]. From practical perspective, if field f leading discarded by the search algorithm, since its loss function will be very large.
solution that explodes before >>> 0 is selected, then such field will be automatically
Remark 4.2
Н represents here the equivalent of number of epochs in training a deep learning model. As for N, it represents the number runs, equivalent to training a neural network for different seeds and keeping the best model. Hence, the larger these parameters the more likely a better performance obtained. As for the potential for over-fitting that is encountered training a neural network for too long, it does not occur here because the optimization is discrete, hence the best performing model is saved anyway.
7

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm 1 Symbolic Distribution Flow Learner
Inputs: Number of episodes N, number of roll-outs H, screen-shots (ftr,m)im at (4:)osisn
Initialization: Estimate the value of each operation (+, -+,-, x, sin.... as root node; for 1.... N do
Randomly select root node and build an expression tree as follows:
if Tree is complete then
Evaluate the corresponding estimate f by computing V(f), Back-propagate the obtained value else
Run H roll-outs, Store the best estimate, Back-propagatt the corresponding value
Select the operation maximizing UCT(ss a) where is the current state of the tree end if
end for
Return: Most accurate j over the N episodes
4.3 Monte- Morte-Carlo Tree Search sample complexity
A crucial question in MCTS-based algorithms is to estimate the number of required episodes, executed by the algorithm, in order to honor certain tolerance in the obtained solution of the target optimization problem. We apply non-asymptotic error analysis result by (Shah et al. 2020)) to determine the minimal number of episodes -also known as sample complexity- that ought to be used.
Proposition 1.
The average number of score evaluations E required for the MCTS algorithm to find an s-optima solution where & > 0, is at most given by:
E E= (-0(9---(4t))) -log 2
where M is the maximum allowed expression length and the size of the chosen elementary function set.
Proof. The proof is postponed to Appendix
5 Numerical Experiments
First, we illustrate the performance of the proposed algorithm on the Kuramoto network system of ODEs, used across the biological, chemical and electrical domains to model circadian oscillators, pacemaker cells in the heart and electrical power networks among other applications (Discacciati Hesthaven 2021 Dörfler & Bullo 2014). Then, we conduct an evaluation on real-world dataset of embryoid stem cell trajectories (Moon et al., 2019). We provide comparisons of our algorithm with the current trajectory inference state-of-the-art algorithms, namely TrajectoryNet (Tong et al. 2020) and JKOnet (Bunne et al.] 2022). TrajectoryNet relies on a (neural net-based) continuous normalizing flow (Grathwohl et al. 2018)) formulation, augmented with relevant regularizations such as growth rate and velocity penalization. JKOnet builds upon the celebrated JKO scheme (Jordan et al., 1998) describing energy gradient flows, where it parameterizes the energy function and the Monge potential using Input Convex Neural Networks (Amos al. 2017). In addition to being based on black- box models -in this instance, neural nets- these approaches contrast with SDFL in that they require (extra))lyyper-paranee tuning.
For the experimental comparison, we retrain the models with the architectures and hyper-parammeters
proposed by the respective authors (Tong et al. 2020; Bunne et al. 2022); however, we employ early-stopping to avoid over-fitting to the smaller data-sets. For JKOnet, we use a small regularization parameter & = 0.001 to make its target closer to the Wasserstein distance, which our evaluation metric. We conduct experiments in the small and noisy data regime with varying Training Sample Size (TSS) per snapshot, illustrating the
With UCT defined (Shah al. 2020).
3 An e-optimal solution of an oplimzation problem min.gE ((x) value ae E F. that satisfies g(xz) S min g(r) +
8

---

Published in Transactions on Machine Learning Research (03/2025)
suitability of SDFL for practical costly data-collection conditions. Similar to previous studies, numerical evaluations are conducted based on the Wasserstein distance between the predicted empirical distributions and the ground-truth distributions (averaged over 3 runs) on unseen data. The computational time of the different methods reported in section 5.3 For completeness, we report a performance comparison of SDFL against. the symbolic regression slate-oftthe-rl methods by Qian al. (2022) and Sun el al. (2023) Appendix E.4 although they represent less challenging competitors.
5.1
Kuramoto system of ODEs
We investigate the recovery of the Kuramoto system with the state variable #; (8.(s)) following
1 (t) wi + d ==1
Σ Kis sin (0,(1) (t()),
where (wi)isisd are the corresponding natural frequencies. Moreover G (Kiy)ssusa is the graph weight matrix. Gaussian initial condition (with mean and unity variance) employed. For simplicity, following Discacciati & Hesthaven 2021). we assume fully connected uniformly weighted graph (iee. we take K₃ for all i,j (1,...,d]) d})- We consider 15 snapshots, in d = dimensions, over the time horizon T 30. set to 0.01. To highlight the accuracy and robustness of our approach, Figure 2 illustrates the inferred trajectory of state variable (using SDFL), and Figure 3 the robustness of Wasserstein distance with respect to the noise. In the latter figure, the distances are computed after training. Note that to reduce the SDFL running time, it was run in two steps where the global function structure is
Furthermore the natural frequencies
3,70
1.2
Empirical distance
Wasserstein with non-noisy data
State
3.55
0.8
Wasserstein with noisy data
Euclidean distance with noisy data
3.50
0.6
Euclidean distance with non-noisy data
Estimated
3.45
Estimated curve dim
Estimated curve dim
3.40
Ground truth
Ground truth dim
0.2
Ground truth dim
3.35
20
Time
Time
Figure/ 2: Inferred and ground truth state trajec- tory
Figure 3: Distance between inferred and reference
distributions, by different metrics
fitted first, followed by the estimation of regression coefficients (see Appendix C for details). The results (representing the test error for sample size of 50) suggest SDFL is competitive with the stato-of-tho-art. The system of equations obtained by SDFL with 150 samples per screen-shot is given by
reported in table I
Table 1: Prediction error comparison in the Wasserstein metric for the Kuramoto model
TSS
SDFL
JKONet
TrajectoryNet
50
0.57
0.04
0.86 0.22
4.33 = 0.13
100
0.57 = 0.06
0.76 0.18
4.83 = 0.10
150
0.46 + 0.05
0.76 0.22
3.15 + 0.11

---

Published in Transactions on Machine Learning Research (03/2025)
81 0.0087 + 0.3293 * (sin 01) + sin(0s 01)) 0g 0.0087 + 0.3293 * (sin(0) - 02) + sin(03 02))
0g 0.0087 + 0.3293 (sin(02- - 03) + sin(01 03))
Additionally, we illustrate the predicted distributions by SDFL against the ground truth for snapshot time = 2 for two different dimensions, with Gaussian mixture as initial condition. Figure translates the Wasserstein prediction errors to more intuitive and visual comparison.
Screan-shat
Predicted distiibution
Groud tuuth
Hicagram
15
10
Figure 4: Un-normalized empirical density comparison for the 1st dimension
5.2 Single- Cell Population Dynamics
We consider the problem of learning the evolution of embryonic stem cells based single-cell RNA sequencing data measurements over a period of 27 days, where the data was collected at 5 different snapshots (sec Moon el al. (2019)) Apart from the high cost. of obtaining a large dataset, a key difficulty in this setting is that cells are (usually) destroyed during measurement (Bunne et al., 2022). Hence, there is a need for schemes which prodict the distribution evolution across time, rather than individual trajectorios, from limited obscrvations. The dataset was pre-processed using the recent dimensionality reduction technique PHATE (Moon al., 2019), which is specifically designed to preserve maximum variablity in the low-dimensional space while allowing for intuitive visualization, and therefore interpretation. With this goal in mind, natural choice for the reduced dimension is d = 3.
Table 2: Prediction error comparison in the Wasserstein metric for embryoid stem cell trajectory inference
TSS
SDFL
JKONet
TrajectoryNet
100
0.97 0.14
3.18 + 0.54
6.78 : 0.26
200
0.87
3.72 0.51
6.65 = 0.06
300
0.79 + 0.16
3.03 + 0.70
6.25 + 0.11
Similar to the previous case of Kuramoto network, SDFL outperforms the competitor state-of-the-art methods, shown in table 2 with test sample size m 50. This is encouraging, specifically since for the cell population dynamics, the datasct is not produced based on an underlying mathematical model. Most importanlly, it confirms the relevance of SDFL as a dynamic regression tool for real-world data subject to noise and limited observability. It should be emphasized though, that neural net based approaches (TrajectoryNet and JKONet) require sufficiently large amounts of data to reach their optimal performance. As for JKONet, one explanation for why it performs better than Trajectorylet, is that it enforces an inductive bias through the JKO scheme (Bunne et al. 2022). We note that the performance of the SDFL model is stable across
10

---

Published in Transactions on Machine Learning Research (03/2025)
training sample sizes. The system of ODEs obtained by SDFL with 50 samples per snapshot is given by
1 (22 21 + cos( +
222 cos( * 22 + (23)) 22
23 (18( 33 + (2r1) 23
It is worth noting that network modeling is most relevant for scRNAseq trajectory inference, as extensively demonstrated in static settings [Wang et al., 2021 Van Dijk et al. 2018; Van de Sande et al. 2020). For SDFL, we demonstrate the importance of enforcing permutation invariance through a simple ablation study. More precisely, we report in table 8 below the performance metric of the scheme with and Without the Permutation Invariance (WPT) module, for the scRNA-seq evolution modeling task, without fitted regression parameters. We observe drop in performance given the same scarch horizon (number of episodes).
Table 3: Prediction loss for the scRNA-soq evolution modeling task
Method
SDFL
SDFL-WPI
W2 metric 2.0822
2.3247
5.3 Computational time
We present in tables 4 computational time comparison between SDFL, JKOnet and TrafectoryNet. For fair comparison, all the reported running times are obtained on an Intel(R) Core(TM) i7-7500U CPU. The reported times correspond to training on a sample of m 50 per screen-shot for the Kuramoto
Table 4: Average running time (hours) for the Kuramoto system and the scRNA-seq evolution modeling
Task
SDFL
JKONet
TrajectoryNet
Kuranoto
5.8166
0.7044 0.4696
5.7709
scRNA-seq
3.8777
3.8472
system modeling task and m 100 per screen-shot for the scRNA-seq evolution modeling one. It worth emphasizing that the reported running times for JKONet and TrajectoryNet do not take into account hyper-parameter tuning. And, since SDFL does not require the latter, we believe it is still competitive (for low dimensions).
Remark 5.1
It worth noting that, even for dow-dimensional settings, extracting explicit models under the sparse data regime (one measured distribution every few days) could not be achieved manually.
Remark 5.2
To give estimate of the corresponding computational time for higher dimensional systems, considering the same expression length for all dimensions, note that the computational time is of the order of O(g²), where is the size of the set. of elementary functions. Hence, if we consider a system of dimension instead of 3, the computational time SDFL would approximately be multiplied by 50, with the chosen elementary function set. We further discuss this limitation in soction t below. As for opportunities for optimization, there are least two of them: one is parallelize the stochastic roll-outs and the other is to parallelize the runs of different episodes. This can be done easily given the structure of the scheme.
11

---

Published in Transactions on Machine Learning Research (03/2025)
6 Explainability & Limitations
Given that the models obtained by SDFL are based on compact symbolic expressions, they allow for the extraction of interpretable insights (Qian et al. 2022 Menara et al., 2022). For instance, the obtained model for the single-cell population dynamics suggests complex oscillating patterns that are not captured by a mere linear interaction between the componenis, as illustrated by the plot comparisons in figure below and the accuracy comparison in table (Appendix E Additionally, such explicit models can allow domain experts to conduct sensitivity and parametric analyscs, revealing critical parameters that result in jumps significant changes in the system. One specific example of such behaviors is the critical coupling strength of the Kuramoto model (Discacciati & Hesthaven 2021 Specifically, if K = Kij 4 then
d
e (3)
= -
d
jml
denoting incoherent behavior. On the other hand, if > 1 then R 2 1, denoting complete synchronization. However, such advantages come at the cost of challenging scalability higher dimensions (Virgolin & Pissis 2022 as well the need to solve candidate ODEs the fly which can be computationally intensive.
0025
Ground-truth
Cround-truth
SDRL Estimate
SDFL Estimate
Linear Estimate
0.03
0.02
Statt
State
0.01
0.00
0.01
-0.02
30
12
Figure 5: Inferred trajectories for different samples
7 Conclusion
In this work, we investigated model-discovery of dynamic systems, subject to randomness. Specifically, taking a symbolic regression perspective, we considered setup in which only a few random observations of the system are provided in time and space. We proposed SDFL, a scheme incorporating several innovative ideas, introduced an appropriate
in order to tackle this multifaceted dynamic inference problem. In particular, measure for goodness-offfit by devising time integrated Wasserstein loss. This design choice turned out to be powerful by offering robustness of the model-discovery in presence of noise/unceet /uncertainty in the data. In addition, we derived theoretical guarantees on the scheme sample complexity. Motivated by the performance of SDFL in our numerical experiments, we believe these developments contribute to robust model-discovery approaches from noisy and limited data, and bridge the rich literatures on probabilistic modeling and symbolic regression. Computational scaling higher dimensions remains an interesting open problem for future studies.
NNoe that, the estimates based normalized data.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Acknowledgments
This work was supported by the Swiss National Science Foundation under grant No. 212876. We acknowledge computational resources from the Swiss National Supercomputing Centre CSCS.
References
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155. PMLR, 2017.
Douglas Adriano Augusto and Helio JC Barbosa. Symbolic regression via genetic programming. In Procecdings. Vol. 1. Sirth Brazilian Symposium on Neural Networks, pp. 173-178. IEFE, 2000.
Sören Becker, Michal Klein, Alexander Neitz, Giambattista Parascandolo, and Niki Kilbertus. Discovering ordinary differential equations that govern time-series. arXiv preprint arXit:2211002880 2022.
Luca Biggio, Tommaso Bendinelli, Alexander Neitz, Aurelien Lucchi, and Giambattista Parascandolo. Neural symbolic regression that scales. In International Conference on Machine Learning, pp. 936-945. PMLR,
2021.
Josh Bongard and Hod Lipson. Automated reverse engineering of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 104(24):9943 -99499999999,,,, 2007.
Lorenzo Boninsegna, Feliks Nüske, and Cecilia Clementi. Sparse learning of stochastic dynamical equations. The Journal of chemical physics, 148(24):241723, 2018.
J Frédéric Bonnans and Alexander Shapiro. Perturbation analysis of oplimization problems. Springer Science & Business Media, 2013.
Steven Brunton, Joshua Proctor, and J Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113 (15):3932-300) 3937, 2016.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport modcling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511-6528. PMLR, 2022.
Tristan Cazenave. Monte-carlo expression discovery. International Journal on Artificial Intelligence Tools, 22 (01):1250035, 2013.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Zhao Chen, Yang Liu, and Hao Sun. Physics-informin learning of governing equations from scarce data. Nature communicattons,, 12(1):6136, 2021.
Earl A Coddington, Norman Levinson, and T Teichmann. Theory of ordinary differential equations, 1956.
Theodore Cornforth and Hod Lipson. Symbolic regression of multiple-time-scale dynamical systems. In Proceedings of the 14th annual conference on Genetic and evolutionary computation, pp. 735-742, 2012.
Allan Costa, Rumen Dangovski, Owen Dugan, Samuel Kim, Pawan Goyal, Marin Soljačié, and Joseph Jacobson. Fast neural models for symbolic regression at scale. arXiv preprint arXiv:2007.10784, 2020.
Miles Cranmer, Alvaro Sanchez Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases. Advances in Neural Information Processing Systems,
17442, 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint ar 180:1810.0488 2018.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Niccolò Discacciati and Jan S Hesthaven. Modeling synchronization in globally coupled oscillatory systems using model order reduction. Choos: An Interdisciplinary Journal of Nonlinear Science, 31), 2021.
Florian Dörfler and Francesco Bullo. Synchronization in complex networks of phase oscillators: A survey. Automatica, 50(6):1539-156; 2014.
Emilien Dupont, Arnand Doucet, and Yee Whye Teh. Augmented neural odes. Advances in neural information processing systems, 32, 2019.
Nicolas Fournier and Arnaud Guillin, On the rate of convergence in wasserstein distance of the empirical measure. Probability theory and related fields, 162(3-4):707-788 2015.
Ting-Ting Gao and Gang Yan. Autonomous inference of complex network dynamics from incomplete and noisy data. Nature Computational Science, 2(3):160-168, 2022.
Sébastien Gaucel, Maarten Keijzer, Evelyne Lutton, and Alberto Tonda. Learning dynamical systems using standard symbolic regression. In Genetic Programming: 17th European Conference, EuroGP 2014, Granada, Spain, April 23-25, 2014, Revised Selected Papers 17, pp. 25 35-36. Springer, 2014.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arviv preprint 1rXiu:1810.01300 2018.
Tatsumori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rmns. In International Conference on Machine Learning, pp. 2417-2426. PMI., 2016.
Yunfei Huang, Youssef Mabrouk, Gerhard Gompper, and Benedikt Sabass. Sparse inference and active learning of stochastic differential equations from data. Scientific Reports, 12(1):21691, 2022.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. Advances in Newral Information Processing Systems, 35:29705-290 297118 2022.
Mohiul Islam, Nawwaf N Kharma, and Peter Grogono. Expansion: A novel mutation operator for genetic programming. In LJCCI, pp. 55-66, 2018.
Seungwoo Jeong, Wonjun Ko, Ahmad Wismu Mulyadi, and Heung-Il Suk. Deep efficient continuous manifold learning for time series modeling. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the licker-planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Pierre-Alexandre Kamienny, Stéphane d' A'soooli, Guillaume Lample, and Francois Charton, End-to-end symbolic regression with transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (ods.), Advances in Newral Information Processing Systems, 2022. URL https:// net/
forum forum?idd-GouuI+HHGG
Patrick Kidger. On neural differential equations. arXiv preprint arXix:2202.n 02435, 2022.
John R Koza. Genetic programming as means for programming computers by natural selection. Statistics and computing, 4:87-112, 1994.
Jiří Kubalík, Erik Deruer, and Robert Babuška. Neural networks for symbolic regression. arXiv preprinl arXiv:2302.0077 2023.
Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobaloo Grammar variational autoencoder. In International conference on machine learning, pp. 1945-1954 PMLR, 2017.
Guillaume Lample and Frangois Charton. Deep learning for symbolic mathernttics. In Internalional Conference on Learning Representations, 2020. URL htttp:///peeeeeeeee.ee//rrrrrrss
Tor Lattimore and Csaba Szepesvári. Bandil algorithmms. Cambridge University Press, 2020.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Li Li, Minjie Fan, Rishabh Singh, and Patrick Riley. Neural-guided symbolic regression with asymptotic constraints. arXie preprint wrXiv:1901.07 2019.
Jean-Christophe Loiseau and Steven L Brunton. Constrained sparse galerkin regression. Journol of Fluid Mechanics, 838:42-68 67, 2018.
Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Neural manifold ordinary differential equations. Advances in Neural Information Processing Systems, 33: 17548-17555 2020.
Qiang Lu, Fan Tao, Shuo Zhou, and Zhiguang Wang. Incorporating actor-critic in monte carlo tree search for symbolic regression. Neural Computing and Applications, 33:8495 8511, 2021.
Georg Martius and Christoph H Lampert. urXiv:1610.0290 2016.
Extrapolation and learning equations. arXiv preprint
Tommaso Menara, Giacomo Baggio, Dani Bassett, and Fabio Pasqualetti. Functional control of oscillator networks. Nature communications, 43(1), 2022.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: performance guarantees and tractable reformulations. Mathematical Programming, 171 11-2):115-166, 2018
Kevin R Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37(12)) 1492, 2019.
Rémi Munos et al. From bandits to monte-carlo tree search: The optimistic principle applied to optimization and planning. Foundations and Trends& in Machine Learning, 7(1):1-129, 2014.
Harald Niederreiter. Quasi-monte carlo methods and pscudo-random numbers. Bulletin of the Americon mathematical sociely, 84(6):957 1041, 1978.
Brenden K Petersen, Mikel Landajuela, T Nathan Mundhenk, Claudio P Santiago, Soo K Kim, and Joanne T Kim. Deep symbolic regression: Recovering mathematical expressions from data via risk-secking policy gradients. arxvv preprint 0rXiv:191..0487,, 2019.
Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 2018.
Zhaozhi Qian, Krzysztof Kacprzyk, and Mihacla van der Schaar. D-code: Discovering closed-form odes from observed trajectories. In International Conference on Learning Represenlations, 2022.
Markus Quade, Markus Abel, Kamran Shafi, Robert K Niven, and Bernd R Noack. Prediction of dynamical systems by symbolic regression. Physical Review E, (4(1):012214, 2016.
Samuel H Rudy, Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. Data-driven discovery of partial differential equations. Science advances, 364):16602614 2017.
Subham Sahoo, Christoph Lampert, and Georg Martius. Learning equations for extrapolation and control. In International Conference on Machine Learning, pp. 4442-4450. PMTR, 2018.
Hayden Schaeffer. Learning partial differential equations via data discovery and sparse optimization. Pro- ceedings of the Royal Society A: Mathematical, Physicall and Engineering Sciences, 473(2199):201604446 2017.
Hayden Schaeffer, Giang Tran, and Rachel Ward. Extracting sparse high-dimensional dynamics from limited data. SIAM Journal on Applied Mathematics, 3876))327733295, 2018.
Devavrat Shah, Qiaomin Xic, and Zhi Xu. Non-asymptotic analysis of monte carlo tree scarch. In Abstracts of the 2020 SIGMETRICS/ //TTRICC//Pefoonncce Joint International Conference on Measurement and Modeling of Computer Systems, pp. 32, 2020.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Hongzhi Shi, Jingtao Ding, Yufan Cao, Li Liu, Yong Li, et al. Learning symbolic models for graph-structured physical mechanism. In The Eleventh International Conference on Learning Regresenlations, 2023.
Guido F Smits and Mark Kotanchek. Pareto-front exploitation in symbolic regression. Genetic programming theory and practice II, pp. 283 299, 2005.
Fangzheng Sun, Yang Liu, Jian-Xun Wang, and Hao Sun. Symbolic physies learner: Discovering governing equations via monte carlo tree search. In The Eleventh International Conference on Learning Representations, 2023. URL https:://peerrevee wtt.://wwwweeeoooooddddeaaaaa htte:/foormiidddii1- 2TK3Sef88
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526 9536. PMLR, 2020.
Harsha Vaddireddy, Adil Rasheed, Anne E Staples, and Omer San. Feature engineering and symbolic regression methods for detecting hidden physics from sparse sensor observation data. Physics of Fluids, 32 (1):015113, 2020.
Bram Van de Sande, Christopher Flerin, Kristofer Davic, Maxime De Waegeneer, Gert Hulselmans, Sara Aibar, Ruth Seurinck, Wouter Saelens, Robrecht Cannoodt, Quentin Rouchon, et al. A scalable scenic workflow for single-cell gene regulatory network analysis. Nature Protocols, (5(7):2247-22)
2020.
David Van Dijk, Roshan Sharma, Juozas Nainys, Kristina Yim, Pooja Kathail, Ambrose J Carr, Cassandra Burdziak, Kevin R Moon, Christine I. Chaffer, Diwakar Pattabiraman, et al. Recovering gene interactions from single-cell data using data diffusion. Cell, 174(3):716 729, 2018.
Martin Vastl, Jonáš Kulhánek, Jirí Kubalík, Erik Derner, and Robert Babuška. Symformer: End-to-end symbolic regression using transformer-bart based architecture. arXiv preprint arXiv:2205.1576, 2022.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Marco Virgolin and Solon P Pissis. Symbolic regression is np-hard. arXiv preprint arXi::2207001018, 2022.
Juexin Wang, Anjun Ma, Yuzhou Chang, Jianting Gong, Yuexu Jiang, Ren Qi, Cankun Wang, Hongjun Fu, Qin Ma, and Dong Xu. segnn novel graph neural network framework for single-cell rna-seq analyses. Nature communications, 12(1):1882, 2021.
16

---

Published in Transactions on Machine Learning Research (03/2025)
A Proof of Theorem
Proof. Recall that we would like to prove that there are positive constants M1, > such that:
W2(F"1 (f) pe.m) T /
T Wz (2(F' 14) de - m (d/2+1)
Mi
M2
E n
+
n
provided m, n 2 2. Without loss of generality, we assume T = 1. By theorem in Fournier & Guillin 2015
given that zo + F (xo) is continuous and for all i {l...., n.) pt, has compact support on the Polish space Rd, there exists C₁, C2 > such that
2 W2 (f) Fra (f) mmro)
Ci
E
mld/2-1)
n
C2 m(d/2-1)
and
E
Wz (D Mei)
Therefore, thanks to the triangular inequality satisfied by W2, get
E
(2 (f)
Wz (F)to
M
n
C1 + mdd/2+1)
E
W2(F": (f) Fr. F (f)
+E
Wal
(5)
On the other hand, denoting by M(Rd x IRd) the Banach space finite measures on Rdd Rd equipped with the total variation norm, the mapping 71) E [0, M(R4 x Rd) + Jæayz FC (z) yll dr(æ, y) is differentiable. Moreover, by compactness the set of couplings given two fixed marginals and by theorem 5.20 in [Villani 2009). conditions (ii) and (idi) of theorem 4.24 in (Bonnans & Shapiro 2013) are satisfied
leading to the fact that the mapping + (f(FF())) pe) continuously differentiable on [0, 5 Hence, there exists a positive constant C3 > such that
S M2(F'() pe) de W Mes)
i/n
C3
n
for i nI, which yields
n L W2 (22(F* pe) de
n Σ (
C3 n
S
Consequently, by summation with 5 we get
m Wz(F' Pi Peamm S'' (2(FF' pe) dt s
Ci + C2 mld/2+1)
C2
E
+
n
n
B Proof of Proposition
Proof. Denote by (uu)eea,,y the ground-truth probability flow and let f, g be the continuously differentiable functions defined by two analytic expressions from the selection space of MCTS. Recall that the goal is to maximize the functional given by
V:f+ 1+ (f)
Givven that the optimal coupling unique and that the directional derivative continuous with respect 1] -as supremum differentiable parametric family convex functions.
17

---

Published in Transactions on Machine Learning Research (03/2025)
The result is obtained by showing that the conditions of theorem 2 in (Shah et al. 2020) are satisfied in our setting. For that matter, it sulfices to show that
Lif S (2(FF pe) dt
1 is Lipschitz with respect to the L or 11 norms; since the derivative of h + 1+x implying that h is Lipschitz. Furthermore, we have
is bounded on IR+
[S(f) - 5(9)) / [F) pe) (g) pe) dt
$ J. W2(F F (0) dt
because W2 is distance. Additionally, since both distributions have compact support, and by regularity of (J),, F (g), there exists constants C²C² > 0 such that for all ( E [0,T],
()(F(()) (g) S C- (IFF No: F #))
SC. sup Lip(e)
q ((FFF(())
/ (9)#())
/ 9o F'( duo / go (') dup
<C- sup Lip(a)
SC. sup
qo F*) oo F'(g)) duo
Lip(e)
$ C. IF'(f) - -F'(F)-F'(()) F'()))
duo
T C" -1f -
where the second inequality justified by the dual representation of W1, the third by the change of variable formule and the fifth by the Lipschitz property of q. The last inequality is justified by the fact that the solutions t (z) and + (z) of the differential equations ÿ == =f(y) and ÿ = g(y) respectively, for given initial condition E Rd, are fixed points of the Picard operato which is Lipschitz.
Consequently, integrating over / E [0, on both sides, we obtain the Lipschitz property of + I#LU) Finally, since L(J) can be approximated by (J) with arbitrary accuracy, get the sample complexity upper-bound.
C Implementation details & Reproducibility
For the implementation of SDFL, we set the building operations consisting of (+, -, x, cos, sin, exp) with a maximum of L 20 operations per expression, with number of episodes of 500 to 1000. For the recovery of
the Kuramoto system, we use 15 snapshots with time-stamps I; 2i for 1 15, and we set K 1/3. The fitted parameters correspond to 50 sample points per screenshot. For the cellular dynamics data, after pre-processing step using PHATE (Moon et al. 2019), the dimension is reduced to = 3 then standard Gaussian noise samples were added to the training set. The following governing ODE system is obtained (after 50 episodes):
Φ cos( (22) 21 + cos( Z1
2/2 cos x1) * x2 + cos (o)) * x2
2/3 cos (22) 23 + cos Ø3
Note that, since the considered state space is finite, we do not need construct explicit covering. Note that o [*(f) denotes the composition of by F*((). 8The Picard operator introduced in the proof of the existence (local) solutions ODEs (see e.g. Coddington al. 1956
18

---

Published in Transactions on Machine Learning Research (03/2025)
Given the inductive bias of searching for explicit ODE model, note that SDFL results in model with similar accuracy across sample sizes (of the same order, those illustrated in table 2 For comparison, we used the publicly available implementations of JKOnet and TrajectoryNet, from (Bunne et al. 2022 and
(Tong et al. 2020 respectively. we retrain the competitor models with the architectures and hyper-parametog parameters proposed by the respective authors (Tong et al. 2020; Bunne et al., 2022]: however, we employ early-stopping to avoid over-fitting to the smaller data-sots 2 For JKOnet, we use a small regularization parameter E = 0.001
to make its larget closer to the Wasserstein distance. Additionally, to foster reproducibiiity, Python implementation of SDFL has been made public at hhttps://gitae htt.a//...ccmmaamaaoooo/////////
D Detailed pseudo-code
Below, we provide more detailed pseudo-code of SDFL while featuring more clearly the MCTS components.
Algorithm Symbolic Distribution Flow Learner [extended description
1: Inputs: Number episodes N, number of roll-outs H, maximal expression length M, elementary functions set (+: x, sin....... screen-shots (ft,.m)i, at (44)ogign
Initialization:
+ Estimate the value of each operation (t,,,x,sin,,, as root node through H stochastic roll-outs
5: + Store these values in V(O, a), where = O represents the empty tree state and the chosen root operation
6: + Define Smax == max a))
for I,..,,, do
Randomly select root node and build expression tree follows:
10:
if Tree is complete then
11:
Evaluate the corresponding estimate of by computing S(f)
12:
if S(f)= > Smax then
13:
Smax S(f)
Back-propagate the obtained value s(ʃ) by updating the values of (V(sp: 4)) where (8p)221 [og, .,,.....,.............. is the finite sequence of encountered tree states before completion end if
14:
15:
16:
else
17:
Run H roll-outs by randomly selecting operations to extend and complete the tree Assign a value of to trees resulting in an inconsistent mathematical expression
18:
19:
Store the best estimate in V(s, az) where a, is the corresponding best operation
20:
Back- Back-propase propagate the obtained value of ((s, ab) to the encountered tree states (sp)p21
(an, ai....,-p-11 Select the operation a maximizing UCT(s, a) where the current state of the tree
21:
22: end if
23: end for
24: Return: Most accurate f over the N episodes
Number iterations used
1000 for JKKOnet and 1500 for TrajectoryNet
19

---

Published in Transactions on Machine Learning Research (03/2025)
Remarks:
The back-propagation, such as in line 14, consists in updating the values of the tree states that have been encountered until that step, in that specific episode. Note that the tree states consist of the sequence of operations which have been selected to constitute the chosen expression (until that step).
The idea is to keep track of the operations (nodes) leading to a good score, and to give them higher chance of being selected the next rounds.
prediction
F' (f)
Monte-Cario Tree Search
S(f)
Initial state
(f)
snapshot
Figure 6: The devised model-discovery setup.
E Additional numerical results
E.1 Robustness in higher dimensions
We demonstrate numerically the robustness of the the proposed loss function for higher dimensions. As shown in Figures
and g the robustness property extends higher dimensions.
3.5
3.0
Empirica distance
2.5
Empirical distance
Wasserstein with non-noisy data
Wasserstein with non non-noisy data
2.0
Wasserstein with noisy data
Wasserstein with noisy data
Euclidean distance with noisy data
Euclidean distance with noisy data
Euclidean distance non no-noisy data
Euclidean distance with non-noisy data
20
25
5
10
15
20
25
Time
Time
Figure Distance between inferred and reference distributions, for system dimension d = 5
Figure 8: Distance between inferred and reference distributions, for system dimension d 10
20

---

Published in Transactions on Machine Learning Research (03/2025)
Empirical distance
non-nolsy
Wasserstein with noisy data
Euclidean distance with neisy data
Euclidean non-nolsy
Time
Figure 9: Distance between inferred and reference distributions, for system dimension d 20
E.2 Heterogeneous Kuramoto system
We run the proposed scheme on data obtained from 50 trajectories, measure at time points to 0, and t+1 = 1 + tú for 1 i 15, for the Kuramoto system given by
01 0.01 + 1.333 (sin 01) sin -01)
02 0.01 + 1.222 (sin(0) 02) sin(O3 02))
0g = 0.01 + 1.111 (sin(02 02) + sin(O1 - -0))
We note that the scheme recovers the global structure of the system as expected. As for the parameters, given the limited amount of data and non-convexity of the estimation, the true parameters have not been successfully recovered.
E.3 Comparison against symbolic baselines
We provide numerical comparison of the performance of the model obtained by SDFL in comparison with (symbolic) baselines, over 3 runs following [Togg et al.) 2020 Bunne et al. 2022 As illustrated in table we note that the model obtained by SDFL offers significant improvement, suggesting its usefulness for domain experts, with interest for explicit representations of dynamic processes.
Table 5: Prediction loss for the scRNA-seg evolution modeling task
Model
Error
Obtained by SDFL
2.02 + 0.34
Linear
3.99 + 1.23
Independent components
3.41 + 0.76
Table displays the Lest error for the models, where linear regression parameters were estimated based on the general symbolic structure obtained previously. The results suggest once again that the model obtained by SDFL captures the non-lincar aspoct of the system behavior.
E.4 Comparison against symbolic regression state-of-the-..
We compare SDFL estimates against those obtained by the state-offtthearr methods SPL (Sun al., 2023) and D-Code QQaa et al. 2022 for varying sampling frequencies, using the publicaly available code by the
21

---

Published in Transactions on Machine Learning Research (03/2025)
Table Prediction loss for the scRNA--eq evolution modeling task, with additional regression parameters
Model
Error
Obtained by SDFL
0.94 + 0.41
Lincar
1.28 0.55
Independent components
3.41 0.76
respective authors 10 The point of this section is to show numerically that, under similar sample sizes, SDFL needs much fewer data-points across time than the current state-of-the-art. We use the Kuramoto model as a benchmark and note that it would not be possible to run the same comparisons on the scRNAseg dataset, since in that case we do not have the possibility to change the sampling frequency. We compare the SDFL estimate with a sampling frequency of freq = 0.5, that is corresponding to t,+1 - t, = 2 where t, are the screenshot times, against SPL and D-Code estimates with a sampling frequency of freq 100, that is f+11 t 0.01. We report in tables 7 and the prediction loss, where the Kuramoto system has been solved with a maximum integration step of 0.1 and 0.01 respectively. Additionally, we report in table 9 below the accuracy evolution across sampling frequencies of the competitor methods. We note that SDFL requires 100 times less sampling frequency to reach comparable accuracy.
Table 7: Normalized prediction loss for the Kuramoto system, with integration step of 0.1
Method
Wasserstein Error
SDFL
0.34 = 0.24 2.54 = 0.62
SPL
D-Code
0.34 + 0.05
Table 8: Prediction loss for the Kuramoto system, with integration step of 0.01
Method Wasserstein Error
SDFL
1.39 0.59
SPL
61.13 + 5.47
D-Code
188.48 + 52.74
Table 9: Normalized prediction loss evolution across sampling frequencics
Frequency
100
10
SPL
2.54 + 0.62
7.45 + 3.57
10.36 + 1.00
D-Code
0.34 + 0.05
6.53 + 0.92
23.72 4 2.66
We observe that the accuracy of the different methods increases with the sampling frequency expected. Note that this remark goes in the same direction the findings of the papers Sun et. al. 20223 and Qian et al. (022)). that proposed SPL and D-Code respectively, which limit their evaluations to freq = 100 (i.e. 100
points per unit of time) or notice drop in performance for lower sampling frequencies. As for the important drop of performance in the case of a smaller integration step size, it could be explained by the fact that the
10We
optimized by the respective authors with running horizon H given by HSDEL 200, Her 500
and Hp--ode 1500.
22

---

Published in Transactions on Machine Learning Research (03/2025)
estimates only recovered the global behavior of the Kuramoto system state and trying to match them to the true state at fine scale leads to error compounding.
23


Reviews:
[{"review_id": "2W3QLmQ1jA", "reviewer": "Reviewer_JcUL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strenghts\n- The motivation of the work, its scope, as well as the writing is clear. I also think it is a timely problem since the limits of black-box function estimators for ODE systems is already reached and white-/grey-box approaches have a lot of potential to explain system dynamics.\n- Using Wasserstein distance (rather than the standard MSE or evidence lower bound based losses) for loss computation is a solid idea, given that the aim is modeling probability distribution flows.\n- I did not see any flaw in the presented approach although some details are to be clarified.\n\n### Weaknesses\n- The biggest weakness is the limited experimental setup: ablations and comparisons with more methods are needed (see the next section). These additional experiments are not to show the superiority of SDFL but to see if benchmarking is done right (in other words, to see the robustness and limits of the method).\n- Presentation needs to be improved (see the next section again)\n- Unless I'm missing anything, the estimation of initial value is not described. Also, to my understanding, it is the only source of uncertainty. I would be happy to hear from authors how they estimated it and how it impacts the overall performance.\n\n### Minor comments/weaknesses/questions\n- I am not sure if the presented approach is suitable _in particular_ to network inference. It seems to me as a quite generic approach; hence, I do not understand why authors focused on \"network modeling\". If the focus is indeed on network modeling, then I would be glad to see experiments on more complicated networks\n- What is the main message of theorem 5.1? As I understand, it justifies the use of a discrete loss function to estimate a continuous-time system but I guess this is the only possible way to estimate the system?\n- What is the interpretation of the equations inferred in the single-cell population dynamics section?\n- The reported losses/errors are in Wasserstein distance, which is the objective the model minimizes. What are the corresponding MSEs (or are they the values reported in Tables 5-9)?\n- Concerning figure 5: if we only look at state trajectories, then neural ODEs are also equally informative. I feel the merits of this approach lie in the estimated closed-form equations, which authors can demonstrate better", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a method called Symbolic Distribution Flow Learner (SDFL) to learn equations underlying probabilistic network flows. Unlike the rather recent, black-box ODE inference techniques, SDFL is a white-box approach that uses symbolic search to recover differential equations. SDFL employs a Wasserstein-based loss function that is compatible with probabilistic real-world systems. The authors also provide a sample complexity result, quantifying the number of required snapshots for reliable model discovery. The findings on Kuramoto networks and single-cell RNA sequencing trajectories show the efficacy of the presented approach.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Additional experiments\n- Comparison against standard neural ODE baseline, which is actually a probabilistic approach as the initial value is sampled from a distribution.\n- Model parameters/assumptions need to be ablated (initial value distribution, problem dimensionality, $n$, $m$, $L$, etc.)\n- It would be interesting to see how the model performs if we had larger set of unitary and binary expressions, e.g., absolute value and powers.\n- I would also like to see if the estimation would still be fine if we had a different Kuramoto network (with different $K$s.)\n- An analysis on how the method scales to high-dimensional spaces is necessary. Now, it seems that the experiments are conducted only on 3D systems. What would be the computational cost of the Wasserstein loss in, e.g., 10 or 50 dimensions?\n\n### Presentation\n- It would be great to have a conceptual/summarizing figure early in the text. Figure 1 and 6 are the candidates but it is not clear to me what they communicate.\n- Algorithm 1 can be detailed, I find it too generic now\n- Does proposition 1 apply the presented method or is it a generic one?\n- Beginning of section 4 can be written more verbosely, it is now difficult to follow (for example, terms that are not defined are used.)\n- It would be good to see how exactly the loss is computed as it is rather abstract now.\n- In Fig3, what does the \"Euclidean distance\" refer to? Is it computed after training or is it a way of measuring the difference between the estimated trajectories and ground truth?", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IxDe2ZJzss", "reviewer": "Reviewer_7KWH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Strong and convincing theoretical foundations\n- Comprehensive empirical evaluation on both synthetic and real data\n- Handles practical challenges like partial observability and sparse measurements\n\nWeaknesses:\n\n- My main concerns were about scalability, which is not necessary for a proof of concept publication as intended here.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces SDFL (Symbolic Distribution Flow Learner), a method for discovering interpretable ODEs that govern probabilistic flows in network systems from sparse/partial observations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "As authors likely have anticipated, I have mainly questions about the runtimes/applicability beyond a proof of concept. I think the work satisfies the criteria for acceptance without these, and I did really appreciate the fact that authors were very direct about the shortcomings. \n\n- Even though this is left as a future work, I would still like to see at least a simple example showcasing how the computational complexity would scale for higher dimensional systems? Even a simple back of the envelope calculation would be valuable.\n- The runtime comparison suggests SDFL is significantly slower than baselines. Are there opportunities for optimization?\n\nThe paper presents a valuable contribution to the field of symbolic regression. Despite some potential limitations in scalability analysis and computational efficiency, the theoretical foundations are strong and the empirical results demonstrate clear advantages over existing methods. I believe the claims are well balanced and the topic is of interest, so I vote for acceptance.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PTUDu9InTf", "reviewer": "Reviewer_mq2Y", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. It proposes a novel trajectory inference approach based on symbolic search and Wasserstein distance.\n2. Replacing neural network by the MCTS search is an interesting approach which takes better observability of the system\n\nWeaknesses:\n1. It is not clear how the inference is sensitive to some hyper-parameters in SDFL such as N and H.\n2. It is a bit wired that the SOTA TrajectoryNet performs so worse on Kuramoto model and cell trajectory. What is the neural model size for those experiments? Some discussion may be useful.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes a white box approach-dubbed Symbolic Distribution Flow Learner (SDFL). This approach leverages the symbolic search with a Wasserstein-based loss function. It also provides some theoretical guarantees of the model-discovery and complexity. Finally, it compares to the state-of-the-art on two dynamic systems.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Some discussion on why authors apply MCTS-based algorithm to optimize and search for the trajectory? Discussion the benefits of symbolic search for optimization would be helpful.\n2. Some ablation study for the metric selection would be interesting to explore as well.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers were unanimous in recommending acceptance of this work and I agree with this assessment.
Decision: Accept (poster)