arXiv CSopen access

FluxDisco: Symbolic Regression for Stoichiometric Dynamical Systems via Monte Carlo Graph Search

· arxiv_cs

arXiv CS · Papers · License: Open Access

machine learning, deep learning, neural networks

FluxDisco: Symbolic Regression for Stoichiometric Dynamical Systems via Monte Carlo Graph Search Cassandra Durr1 , Dr Alvaro Köhn-Luque2 , Prof. Chris Jewell1 , Dr Lloyd A. C. Chapman1 1

Lancaster University

University of Oslo

arXiv:2609.05207v1 [stat.ML] 4 Sep 2026

Abstract Dynamical symbolic regression methods identify governing differential equations from noisy data, balancing interpretability and predictive accuracy. However, standard methods often produce expressions that violate known physical laws. To address this, we propose FluxDisco, a physics-informed framework tailored for flux-based, stoichiometric ODE systems. By leveraging a known stoichiometry, we reduce the expression search space and ensure physical adherence. Our framework adapts the Monte Carlo Graph Search algorithm for the unique challenges associated with joint flux discovery of stoichiometric systems. We evaluate our method across a range of physical and biological systems, demonstrating its ability to accurately recover governing dynamics through interpretable equations. Keywords: Dynamical Symbolic Regression, Monte Carlo Physics-Informed Equation Discovery, Stoichiometric ODE Systems

Graph

Search,

Code: https://github.com/CassandraDurr/FluxDisco

Introduction

Many complex, dynamical systems can be described using a set of ordinary differential equations (ODEs) of the general form dx(t) = ẋ(t) = f (x(t), θ) dt where x(t) describes the state of the system at time t and f describes the governing dynamical equations, parametrised by θ. This mathematical framework is the basis of dynamical system identification which seeks to build models of dynamical systems using observed data. Traditionally, researchers have approached the problem of dynamical system identification by assuming the functional form of the governing equations a priori, and then using observed data to fit the parameters θ. This ‘white-box’ approach is reasonable if we understand the governing dynamics of the system up to a set of unknown constants. However, assuming such knowledge is often unrealistic for real-world systems so using this approach often entails making restrictive and potentially incorrect assumptions. If the assumed form of f mischaracterises the true dynamics, the misspecified model will inevitably result in inaccurate forecasts and flawed inferences. In recent years, the scientific community has increasingly turned to ‘black-box’ methods, driven by advances in the field of machine learning, to overcome the limitations of inflexible, ‘white-box’ approaches. These data-driven methods learn to forecast the states of dynamical systems, or the derivatives thereof, from historical data without providing an explicit mathematical formulation of the governing dynamics [1, 2]. These frameworks are often very capable emulators of complex dynamics, providing accurate predictions of how dynamical 1

systems may evolve over time. However, their impressive predictive abilities often come at the expense of model interpretability. Building a strong understanding of governing system dynamics is often of equal importance to having accurate predictions. In these cases, purely data-driven methods fall short. Additionally, having interpretable model output often fosters greater confidence in the model. ‘Grey-box’ modelling provides an opportunity to combine the strengths of traditional ‘white-box’ methods and data-driven ‘black-box’ approaches. Equation discovery can be considered a grey-box approach as it allow us to learn system dynamics directly from observed data without imposing restrictive assumptions, and it outputs interpretable expressions to describe the governing dynamical equations. The class of equation discovery methods we consider in this work is symbolic regression methods. Symbolic regression methods can be split into two categories: functional symbolic regression and dynamical symbolic regression. The former has been more widely studied and consists of estimating a function f (x) from paired, input-output data, (x, f (x)). In contrast, dynamical symbolic regression is used to infer the form of differential equations, ẋ(t) = f (x(t)), from observed trajectory data (t, x(t)) [3]. Dynamical symbolic regression allows us to simultaneously learn the functional form of the ODEs, f (x(t), θ), and estimate its parameters, θ. While pure equation discovery typically imposes no assumptions on the underlying dynamics, we may not be completely oblivious to the governing dynamical laws. This work focuses on fluxbased, stoichiometric systems where the dynamics can be represented as ẋ(t) = Sv(x(t)).

(1)

In this context, S is a stoichiometry matrix encoding the system’s structure and v(x(t)) is a vector of flux expressions, and we have suppressed dependence of v on the model parameters θ for notational convenience, as we do throughout this paper. In physical systems, the underlying stoichiometry is typically known based on prior physical knowledge, whereas the functional form of the fluxes is often unknown. Traditionally, symbolic regression methods fail to account for coupled dynamics in constrained systems, producing solutions that violate known physical laws. Our work assumes a known stoichiometry and performs flux discovery, estimating the system’s flux expressions jointly through a novel dynamical symbolic regression framework we call FluxDisco. Rather than performing unconstrained equation discovery, we benefit from imposing a known stoichiometry that reduces the search space of possible expressions and ensures better physical-adherence.

1.1

Background

To discover the governing dynamics of flux-based stoichiometric systems, we treat the generation of candidate flux expressions as a search problem. This section establishes and motivates the framework required to navigate this vast search space efficiently. 1.1.1

Markov Decision Process Formulation

Sequential decision-making problems involve taking a series of actions where each choice influences long-term outcomes. This class of problems can be formalised using the mathematical framework of finite Markov Decision Processes (MDPs) [4]. An MDP is defined by a set of states s ∈ S representing possible environment configurations, and a set of actions a ∈ A(s) available from each state. As an agent navigates this environment, it receives feedback on its actions through a reward function, r(s). The ultimate goal of the agent is to learn a policy for selecting actions that maximises expected cumulative reward. 2

x0 (t) (S)

v0 (x(t))

v1 (x(t))

x1 (t) (I)

x2 (t) (R)

Figure 1: Flow diagram of the susceptible-infectious-recovered (SIR) epidemic model with state variables x(t) and fluxes v(x(t)). Illustrative Example: SIR Model To bridge the gap between an abstract MDP and the task of flux discovery, we introduce the susceptible-infectious-recovered (SIR) epidemic ODE system, visually represented in Figure 1, as a running example. The stoichiometry matrix, S ∈ RD×F , and flux vectors, v(x(t)) ∈ RF , of the SIR system can be represented as, ẋ(t) = Sv(x(t))      ẋ0 (t) −1 0 −v0 (x(t)) ẋ1 (t) =  1 −1 v0 (x(t)) = v0 (x(t)) − v1 (x(t)) . v1 (x(t)) ẋ2 (t) 0 1 v1 (x(t)) 

(2)

The vector x(t) = (x0 (t), x1 (t), x2 (t))⊤ represents the system variables—the susceptible, infectious and recovered population proportions—and ẋ(t) denotes their respective time derivatives. Our objective is to jointly discover flux expressions v(x(t)). MDPs for Flux Discovery

We define the components of the flux discovery MDP as follows:

• States (s ∈ S): A state represents a set of complete, or partially-complete, flux expressions. We jointly estimate flux expressions, therefore the states include all fluxes undergoing estimation. Flux expressions are composed of terminal tokens (system variables and constants) and non-terminal placeholder tokens (m) representing mathematical expressions that require further expansion. Terminal states represent complete expressions consisting of terminal tokens, whereas partial states contain non-terminal tokens. For example, a valid state comprising partial expressions for the SIR system (Equation 2) is: v0 (x(t)) = m + m × x0 (t) v1 (x(t)) = m × (x1 (t) − c)

(3)

where c is a constant token. • Actions (a ∈ A(s)): The set of valid actions for a state evolve non-terminal m tokens through the use of a context-free grammar. The context-free grammar defines how m can be replaced by operators, constants, or system variables. For example, a grammar for the aforementioned SIR system may be defined as: m → x0 (t) | x1 (t) | x2 (t) | c m→m+m|m−m|m×m|

√

(terminal actions) m

(non-terminal actions)

(4)

• Rewards (r(s)): The reward function evaluates how closely a state’s expressions match the observed data.

1.1.2

Monte Carlo Tree Search

Monte Carlo Tree Search (MCTS) is a popular search algorithm used to learn policies for sequential decision-making [5]. The search space is explored through iterative construction of a tree structure, with nodes representing MDP states and the connections between nodes representing valid actions, or grammar rules. The search algorithm, or agent, intelligently navigates the search space by balancing exploration, trialling unvisited paths, and exploitation, revisiting paths which have demonstrated high potential. The algorithm builds up the search tree by cycling through the following phases (illustrated in Figure 2a): 1. Selection: Starting at the root node, a selection rule is used to traverse the tree until a leaf node is reached. 2. Expansion: If the selected leaf node does not represent a terminal state (complete flux expressions), a new child node is added to the tree. 3. Rollout: Actions (grammar rules) are sampled sequentially from the newly expanded node until a terminal state is reached, yielding a reward. 4. Backpropagation: The reward obtained is propagated up through the selected nodes, and the estimated values of these nodes are updated to inform future iterations. A disadvantage of tree-based methods is that they fail to identify that different paths can result in identical states. For example, expressions m + c × x0 and m + x0 × c in Figure 2a are mathematically equivalent, but remain distinct states within the tree structure. Without merging these duplicate states, information can only be backpropagated along the path currently being traversed. 1.1.3

Monte Carlo Graph Search

Leurent and Maillard [6] developed an adaption of MCTS which operates on a graph structure called Monte Carlo Graph Search (MCGS), illustrated in Figure 2b. This method overcomes the aforementioned limitation of tree-based search algorithms in environments where states can be reached through multiple distinct action sequences. By operating directly on a graph, MCGS allows information to flow freely from a node through all of its ancestors to the root node. This global sharing of information sharpens the value estimates of nodes and improves sample efficiency. The benefits of a graph-based search are particularly advantageous in environments with a high-degree of state overlap so merging duplicate states dramatically reduces the possible search space. The use of a graph is particularly beneficial for symbolic regression due to the nature of mathematical expressions. Because there are multiple ways to construct the same mathematical expression, we expect a high degree of overlap in states. Figure 2b provides an example of how two distinct states can evolve into equivalent expressions that require state merging.

1.2

Related Work

There are three main classes of ODE discovery methods: sparse regression, symbolic regression, and deep learning approaches. Sparse regression methods are foundational in this field and include methods such as SINDy (Sparse Identification of Nonlinear Dynamics) and its many variants [7]. These methods use regularised regression to identify key terms from a user-specified library of candidate functions to form part of the governing dynamical equations. In contrast, 4

Selection

Expansion

Rollout

Backpropagation

m m→m×m

m→m+m

m×m

m→m×m

m+m

m+m×m m → x0

m+m×c

m→c

m → x0

m + m × x0

m → x0

m→c

m → x0

m + c × x0

m+m×c

m + c × x0

m + x0 × c

m+m m→c

m→m×m

m+c

m→c

m + m × x0

m+m×c

m→c

m + c × x0

m → x0

m→c

m→m+m

m×m m→c

m+m×m

m+c

m→m×m

m+m

m→m×m

m→c

m+m×m

m→m+m

m×m

m+m

m→m×m

m+c

m→c

m + m × x0

m→m+m

m×m m→c

m→m×m

m+m×m m → x0

m+c

m→c

m + m × x0

m+m×c m → x0

m→c

m + c × x0

m + x0 × c

(a) Monte Carlo Tree Search (MCTS) Selection

Expansion

Rollout

Backpropagation

m→m×m

m×m

m→m+m

m+m

m+m×m m → x0

m + m × x0 m→c

m + c × x0

m×m m→c

m→m×m

m+c

m→c

m+m×c

m→c

m→m+m

m+m

m+m×m

m+c

m → x0

m→c

m + m × x0

m+m×c

m→c

m → x0

m+c

m→c

m+m×c m → x0

m + c × x0

m→m×m

m×m m→c

m→m×m

m→c

m+m×m

m + m × x0

m×m

m+m

m→m×m

m → x0

m→m×m

m→m+m

m+m m→c

m→m×m

m+m×m m → x0

m + m × x0 m→c

m+c

m→c

m+m×c m → x0

m + c × x0

(b) Monte Carlo Graph Search (MCGS)

Figure 2: The four phases of the MCTS and MCGS algorithm applied to symbolic regression. Unlike MCTS, MCGS allows nodes to be reached via multiple paths, enabling more efficient information propagation across the search structure. For simplicity, a single flux expression per state is illustrated. Highlighted nodes indicate the path traversed during a single iteration of the algorithm. Within-node text shows the flux expression, between-node text shows the applied grammar rule, and arrow directions indicate information flow. symbolic regression methods are search-based and involve exploring a vast combinatorial space of mathematical operators, variables, and functions to build equations. Lastly, deep learning approaches include sequence-to-sequence transformers which encode numerical trajectory data and output predicted equations as a sequence of tokens. These methods prioritise generalisability by pre-training on large, synthetic datasets. After pre-training, inference on new ODE systems can be performed almost instantaneously [3, 8, 9].

Symbolic Regression for Stoichiometric Systems A small subset of equation discovery methods are tailored towards stoichiometric, dynamical systems. An early example of such an approach is Reactive SINDy [10, 11] which learns both the stoichiometry matrix and flux expressions for chemical reaction systems using a user-specified library of possible reactions. This work operates on the assumption of mass action kinetics whereby reaction rates are proportional to the product of the concentrations of reactants raised to the power of their stoichiometric coefficients in the balanced chemical equations. This assumption restricts candidate flux expressions to polynomials. For example, suppose we have a system of N reactions between M chemicals [12, 13], which can be written as: M X

aij Xi →

i=1

M X

bij Xi ,

j ∈ {1, . . . , N }

i=1

where (aij ) and (bij ) represent the coefficients of reactants and products (i = {1, . . . , M }) respectively in reaction j, and Xi represents a molecule of the ith chemical. Assuming mass action kinetics, the reaction rates for this system are polynomials, rj = kj

M Y

xi (t)aij ,

j ∈ {1, . . . , N }

i=1

where (xi (t))i={1,...,M } are the numbers of molecules or concentrations of the reactants, with resulting ODEs ẋ(t) = Sv(x(t)) (S)ij = bij − aij (v(x(t)))j = rj ẋi (t) =

N X

(bij − aij )rj .

j=1

Reactive SINDy requires the specification of candidate reactions, therefore, the algorithm only needs to identify the subset of correct reactions and learn the rates kj ∀ j. A significant limitation of Reactive SINDy is the manual pre-specification of candidate reactions. The Stoichiometrically-Informed Symbolic Regression (SISR) method [14] overcomes this limitation by formulating equation discovery as a search problem, utilising a genetic algorithm to explore the expression space. While our approach is also search-based, we use a grammar-based graph search approach rather than an evolutionary algorithm. Like Reactive SINDy, SISR operates on the assumption of mass action kinetics, restricting the discovered fluxes to polynomials. However, unlike prior methods, SISR discovers the polynomial terms alongside the rate parameters kj and the stoichiometric matrix S. The final equation discovery method for stoichiometric systems that we consider is KinFormer [15]. This method utilises a transformer architecture and is pre-trained on a large dataset of synthetic chemical reaction systems. KinFormer is similar to the generalised ODEFormer model [3], although the authors specialise the method for chemical reaction systems. While both ODEFormer and KinFormer leverage comprehensive pre-training, ODEFormer considers a wide range of mathematical operators and functions, whereas KinFormer restricts the discovered dynamics to polynomials based on mass action kinetics. Our philosophy differs fundamentally from the existing equation discovery methods for flux-based, stoichiometric systems. While our method assumes a known stoichiometry, it allows 6

for the identification of complex, non-polynomial dynamics. In contrast, the existing methods are able to uncover stoichiometries, but restrict the discovered dynamics to polynomials based on mass action kinetics. This assumption is often unrealistic in real-world settings where systems frequently exhibit non-polynomial dynamics. To our knowledge, among equation discovery methods specifically designed for stoichiometric systems, ours is the first to support non-polynomial flux discovery without relying on the restrictive assumption of mass action kinetics.

Methods

2.1

Search Algorithm 0. Conditional Expansion (If node has no children) Expand children. Warm start rollouts. Update & propagate stats.

State Merging (On expansion) Merge identical states.

Next iteration

Next episode

1. Selection Choose best child via selection rule.

2. Rollout Perform rollouts to complete ODEs.

3. Reward Evaluation Optimise constants, solve ODEs, and compute reward.

4. Stopping Conditions Check Terminal state or max iterations.

5. Propagation (On stopping) Propagate info from last visited node. Move to root node & reset.

Figure 3: Workflow of FluxDisco, our modified MCGS algorithm. This section describes the technical details of our modified version of the MCGS algorithm, FluxDisco. Our key contribution lies in how we have both adapted and specialised MCGS for the purpose of symbolic regression, specifically joint flux discovery of stoichiometric ODE systems. FluxDisco, illustrated in Figure 3, provides an iterative procedure to intelligently navigate the expansive search space of governing equations for dynamical systems. In this framework, the search algorithm acts as an ‘agent’ that learns to generate governing dynamical equations by trialling various combinations of mathematical operators, states and constants. Each node in the graph traversed by the agent represents a set of partial or complete flux expressions. The learning is episodic in that the agent starts each cycle of the search at the root node with 7

blank expressions and traverses the nodes of the graph building up the flux expressions through a series of applied grammar rules until it produces complete flux expressions at the terminal states. Within an episode, each iteration follows a structured sequence of selection, simulation, and evaluation to refine the agent’s understanding of which functional forms best describe the observed data. An episode of the algorithm proceeds via the following steps (Figure 3), each explained in more detail in the sub-sections below (with additional details provided in the Supplementary Material, Section S1): 0 Conditional Expansion (Section 2.1.2): This pre-selection step determines if the current node, starting with the root node at the beginning of each episode, has existing children. If the node has no children, all valid child nodes are generated and compared against the graph’s existing nodes to merge mathematically identical states. To prevent a ‘cold start’ where the search algorithm has no data to differentiate between the children, a set of initial rollouts is performed to establish baseline rewards. 1 Selection (Section 2.1.1): Once child nodes are available, the selection rule is used to traverse the graph from the root towards the most promising terminal states. 2 Rollout (Section 2.1.3): Rollouts are conducted from selected nodes to complete partial flux expressions. 3 Reward Evaluation (Section 2.1.4): The ODE systems corresponding to rollouts are solved and then evaluated against observed data to produce a set of rewards based on equation parsimony and goodness-of-fit. 4 Stopping Condition Check: This check determines if the recently evaluated node is terminal, representing complete equations, or if the computational budget for the episode has been exhausted. These conditions trigger information propagation and the end of the episode. 5 Propagation (Section 2.1.5): Reward information learnt during the episode is propagated throughout the graph from the last node visited to all ancestor nodes to inform the next episode’s search. After propagation, the next episode begins from the root node. 2.1.1

Selection

The selection rule ensures the search agent navigates the graph efficiently, identifying promising nodes for expansion. The algorithm begins selection from the root node, corresponding to the empty state where fluxes can evolve into any expression supported by the grammar, and traverses the graph towards terminal states. Many graph nodes will correspond to partially-complete flux expressions that cannot be directly evaluated. Therefore, we use Monte Carlo rollouts to provide intermediate feedback at non-terminal states, employing the stochastic reward variant of MCGS [6]. To account for reward stochasticity, confidence intervals (l(s′ ), u(s′ )) for the mean reward E(r|s′ ) associated with state s′ are constructed (see ‘Reward Confidence/Credible Intervals’ paragraph below). The original MCGS selection rule selects the next state s′ ∈ S(s) according to s ← arg ′max u(s′ ) + γU(s′ ) s ∈S(s)

where U (s′ ) represents the state’s upper value bound, discounted by γ (Section 2.1.5). 8

(5)

This selection rule balances local reward, u(s′ ), with future potential reward, γU(s′ ). The local component of the selection rule prioritises nodes providing the highest average reward. While appropriate in most stochastic reward environments, maximising the average reward is unsuitable for equation discovery, where the objective is to identify a single optimal functional form rather than one that performs well on average. Therefore, we propose the following selection rule: s ← arg ′max {max (u(s′ ), r∗ (s′ )) + γU(s′ )} s ∈S(s)

(6)

where r∗ (s′ ) denotes the maximum reward sampled from state s′ . The local reward component is now bounded below by the maximum sampled reward, preventing the premature soft-pruning of nodes with high observed performance. Reward Confidence/Credible Intervals Leurent and Maillard [6] recommend using the Binary Kullback-Leibler (BKL) divergence to construct the mean reward confidence intervals (l(s), u(s)). However, this formulation is excessively conservative in expansive search spaces like those of equation discovery tasks, resulting in an almost uniform state space search (Supplementary Material Section S1.2). To encourage faster convergence of the confidence intervals, we derive a new formulation based on the theory underlying Thompson sampling [16]. We use the generalisation of Thompson sampling to scalar, bounded rewards to obtain a posterior distribution for the mean empirical reward of a state s: ! n n X X ri ri , β = 1 + n − (Θ|R = r) ∼ Beta α = 1 + i=1

i=1

where r = (r1 , . . . , rn )⊤ denotes the state’s sampled rewards. Under this Bayesian framework, we can derive quantile-based credible intervals in place of confidence intervals for the mean empirical reward −1 (l(s), u(s)) = Q−1 (Θ|R = r), Q (Θ|R = r) (7) α/2 1−α/2 denotes the p-quantile of the posterior Beta distribution. While frequentist where Q−1 p approaches, such as the BKL framework, generate confidence intervals, credible intervals serve a similar purpose by defining a plausible region for the location of the true mean reward. The BKL confidence intervals and Beta credible intervals are compared in Figure 4. Most of the BKL upper confidence bounds remain at their initialised value of one, whereas significantly fewer Beta credible interval bounds remain as initialised. This behaviour suggests that under the proposed formulation, the sampled rewards are utilised more effectively to differentiate between nodes, allowing for a more targeted search within the expansive state space. Having less conservative intervals raises the risk that the true mean falls outside of the intervals, however, this is a worthwhile risk to avoid a uniform exploration of the graph. 2.1.2

Conditional Expansion

Before the selection phase can begin, the potential child nodes are created and initialised. Initial rollouts are then performed from each child node to ‘warm start’ their mean reward credible interval bounds (l(s), u(s)), and upper and lower value bounds (L(s), U (s)). However, due to the high branching factor associated with symbolic regression, initialising every possible child node introduces significant computational overhead. As a consequence, we limit the number 9

Figure 4: BKL confidence intervals compared with Beta credible intervals (Equation 7). Both panels illustrate that the more frequently a node is visited, the more its intervals contract towards the reference line, l(s) = u(s). of child nodes expanded in a single step to the rightmost non-terminal m token for all fluxes. Consequently, multiple iterations may be required to consider all possible children of a state. Expansion is further constrained by the hyperparameter κ, the maximum expression depth. Following the approach in [17], this hyperparameter applies a hard constraint on the complexity of the generated expressions by controlling the number of applied grammar rules per flux. Fluxes with κ applied rules have limited expansion—instead of applying all possible grammar rules, only terminal actions are considered. This constraint concentrates the search on parsimonious expressions, following the inductive bias that physical systems generally comprise few terms [7, 17]. This parsimony constraint is applied in both the expansion phase, and in the rollout phase where fluxes reaching κ grammar rules are force-completed using terminal rules. Our approach differs from that of [17], which awards a zero reward to paths that do not reach a terminal state within κ applied rules. In contrast, we opt to complete expressions using terminal rules once κ grammar rules have been applied. The benefit of our approach is that a dense and meaningful reward signal is preserved to better guide the search agent. If the grammar allows a wide array of non-terminal actions, [17]’s method will rarely complete any expressions, leaving most nodes unevaluated and degrading the search. State Merging To prevent the re-creation of existing graph nodes during expansion, we utilise a transposition table which maps state IDs to the corresponding graph nodes. If the state s′ already exists, the parent s is connected to the existing child s′ , otherwise a new node is created and added to the transposition table. The method of serialising states (creating state IDs) should ensure that mathematically equivalent expressions are correctly mapped to the same state ID. To ensure accurate state merging, we apply a constant folding procedure. This process simplifies and canonicalises the form of flux expressions, and also reduces the number of constants requiring optimisation. There are two forms of constant folding applied: • Additive folding merges constants which are coefficients to the same terms, e.g., c0 x0 (t)+ c1 x0 (t) → k0 x0 (t). • Multiplicative folding merges constants that are multiplied together, e.g., c0 c1 x0 (t) → k0 x0 (t). 10

These procedures are applied recursively to equations to catch redundancies within nested subexpressions, such as root terms or fractions. Reducing the number of redundant constants aids constant optimisation, improving its stability and increasing the likelihood of finding suitable constants. After expressions undergo constant folding, their terms are re-ordered to ensure that terms appear in a consistent, lexicographical order. 2.1.3

Rollout

Post-selection, independent rollouts are conducted from the chosen node to complete its flux expressions. In completing these expressions, grammar rules are sampled according to predefined probabilities. To incorporate domain-specific knowledge and guide the search towards plausible candidates, we allow user-specified sampling probabilities to be assigned to the grammar rules on a per-flux basis. These sampling probabilities also provide an additional constraint during the expansion phase; assigning a zero probability to a specific grammar rule prevents the creation of the associated child node. Grammar rules are repeatedly sampled until either a terminal set of equations is formed, or the fluxes reach κ applied grammar rules. If κ is reached, the expression undergoes ‘forcecompletion’, where only terminal rules are sampled (Section 2.1.2). Rollouts provide the opportunity to ‘probe’ the potential of a node. After the rollout phase, the rewards associated with completed flux expressions are evaluated to provide samples from the node’s underlying reward distribution. These rewards inform a node’s credible interval bounds (l(s), u(s)) (Section 2.1.1) and value bounds (L(s), U(s)). Consequently, nodes with high-reward rollouts will be prioritised during selection, while the value of nodes with poor rollouts are allowed to decay naturally. 2.1.4

Reward Evaluation

During the evaluation phase, the reward for each sampled flux expression is calculated through four main steps: 1. Constant folding: To simplify and canonicalise sampled expressions, a constant folding procedure is employed (Section 2.1.2). This procedure eliminates redundant constants, simplifying the subsequent constant optimisation process. 2. ODE system construction: The simplified expressions are then combined to form a system of ODEs using the pre-specified stoichiometry of the system. 3. Constant optimisation: Given the system of ODEs, the unknown constants are optimised iteratively by minimising the error between the observed data and numerical integrations of the ODEs. 4. Reward calculation: Finally, the reward is calculated using the optimised constants. These steps are detailed in Algorithm S1 of the Supplementary Material. The reward calculation is a computational bottleneck of this algorithm due to its use of ODE solving and iterative, multi-dimensional constant optimisation. Many symbolic regression methods avoid ODE solving by using smoothed finite-differencing to estimate state derivatives, which are then used as the search’s target. While this approach is more computationally efficient, it has a number of downsides. Finite-differencing greatly amplifies data noise, and while smoothing methods can mitigate this effect, they run the risk of distorting true system dynamics and require the specification of additional hyperparameters. By integrating candidate

Flux: (x0 (t) × x1 (t)) + c + ×

Flux complexity = 5 Operations:

Substitutions:

←

Figure 5: Illustration of the complexity calculation of a flux, where the flux expression is visualised as an expression tree. Blue nodes represent non-terminal tokens, orange nodes represent applied terminal actions, and grey nodes represent applied non-terminal actions. ODEs, our approach is more robust to noise and preserves the original data as the target of the search. These benefits justify the additional computational expense. The reward function is designed to balance equation parsimony and goodness-of-fit, and takes the form: MSEtotal (8) R = η C exp − τ where η C forms the parsimony penalty and the exponential term forms the goodness-of-fit factor. To encourage interpretable solutions, we apply a penalty for non-parsimonious equations, η C [17]. The hyperparameter η (≤ 1) is the parsimony coefficient and determines the degree to which complex equations should be penalised. The variable C represents the system’s complexity, where the complexity of a single flux is defined as the number of operations and substitutions required to build its mathematical representation. The system’s complexity is taken to be the maximum complexity of the individual fluxes: Ci = complexity(flux i), i ∈ {1, . . . , F } C = max(C1 , . . . , CF ). For reference, Figure 5 illustrates the complexity calculation for a simple flux. The flux complexities are calculated after constant folding to reflect the simplified, final forms of the expressions. The larger the value of C, or the lower the value of η, the greater the applied penalty. The exponential term of the reward function, MSEtotal exp − , τ quantifies the goodness-of-fit between observations and the candidate solution. The total mean squared error (MSEtotal ) can be decomposed as: ! D−1 N T 2 X 1 X X (n) (n) MSEtotal = x̂ (ti ) − xd (ti ) N × T n=1 i=1 d d=0

(n)

where x̂d (ti ) denotes the model’s estimate for state xd at time ti for realisation n, and xd (ti ) reflects the corresponding observed data. 12

The temperature hyperparameter τ scales the error to the range of the data, preventing a flat and uninformative reward landscape. Without this scaling factor, the goodness-of-fit component of the reward may be unable to distinguish between solutions, allowing the parsimony penalty to dominate the reward function. Constant Optimisation The reward function evaluates candidate ODEs by comparing their integrated trajectories to observed data. Since symbolic constant tokens in the expressions prevent direct numerical integration, the reward is calculated iteratively by proposing constants, integrating the ODEs, and then computing the reward. We optimise constants in the flux expressions using non-linear least squares estimation. We use the trust region reflective algorithm for this on account of its robustness which is essential when evaluating a broad range of candidate ODEs [18, 19]. The constant optimisation process is repeated using multiple initialisations run in parallel to avoid getting stuck in local minima of the error landscape. Candidate expressions are evaluated by numerically integrating systems using the fourth-order Runge–Kutta method (RK4). To promote numerical stability, predicted trajectories are bounded to a range wider than the true data range. This prevents trajectories from exploding towards infinite values, causing integration errors. Furthermore, allowing the trajectories to extend beyond the range of the data improves the stability of the constant optimisation procedure, as informative gradients are maintained even for unsuitable constants. 2.1.5

Propagation

Graph propagation occurs at the end of an episode once a terminal state is reached, or a maximum number of iterations have been performed. During this phase, the upper and lower value bounds (L(s), U(s)) of nodes are updated according to the modified Bellman equations, U(s) ← ′max {max (u(s′ ), r∗ (s′ )) + γU(s′ )}

(9)

L(s) ← ′max {l(s′ ) + γL(s′ )} ,

(10)

s ∈S(s)

which are constructed to align with the selection rule (Equation 6). The bounds are only ever tightened as per the monotonicity condition laid out in [6]. Propagation begins at the last node visited during the episode and propagates through all ancestor nodes, including parent nodes not visited during the episode. By updating all paths stemming from a visited node, sample efficiency is improved. The bounds of terminal states, representing complete expressions, collapse to the true reward to reflect the elimination of uncertainty: L(s) = r(s) = U(s). To manage the high computational cost of updating bounds over the full graph, we employ an efficient, queue-based implementation of the bound updates proposed by Leurent and Maillard [6] (Algorithm 4 ). While the authors recommend propagation after every iteration in an episode, we perform propagation once at the end of each episode. The drawback of this approach is that node selection is based on slightly outdated information within an episode. However, this adaption significantly reduces the total number of updates. Empirically, this approach provides a significant reduction in computational cost with no noticeable performance degradation.

2.2

Simulation Study

We evaluate the performance of our framework, FluxDisco, in identifying governing equations from simulated data across a variety of flux-based dynamical systems. To rigorously test our method, we benchmark it against simulated data from a range of systems exhibiting complex characteristics such as fast and slow dynamics, bifurcations, and oscillations. The benchmarked systems selected for this evaluation are described in Table 1. Observations from real-world systems are typically noisy, therefore we simulate data for each ODE system under noiseless, low noise, and high noise conditions. Further details of the data simulation process can be found in Section S3 of the Supplementary Material. SIR The first system we consider is the SIR epidemic model introduced in Section 1.1.1. We consider three different disease transmission behaviours to test the ability of our method to recover different functional forms of the fluxes. To incorporate prior knowledge of the system, we exclude specific grammar rules from being sampled via zeroed sampling probabilities. The grammar rule m → x2 (t) is excluded from the flux v0 , representing the infection process, as the proportion of recovered individuals x2 (t) does not impact the rate of new infections. Similarly, m → x0 (t) is excluded from flux v1 representing the recovery process, as the proportion of susceptible individuals x0 (t) has no impact on disease recovery. Outside of these exclusions, the remaining sampling probabilities are balanced uniformly. Lotka-Volterra The Lotka-Volterra system is a well-studied ecological model that describes how populations of interacting predator and prey groups fluctuate over time, producing oscillatory behaviours. In this formulation, the states x0 (t) and x1 (t) represent the population density of the prey and predators, respectively. The grammar rule m → x1 (t) was excluded from v0 , as prey birth rates are independent of predator population density, and m → x0 (t) was excluded from v2 , as the natural death rate of predators is independent of prey availability. Brusselator The Brusselator is a theoretical chemical reaction system which models the concentration of interacting chemicals x0 (t) and x1 (t) whose reaction is driven by constant inputs c0 and c1 : • Flux v0 represents an autocatalytic reaction that consumes x0 (t) and x1 (t) to produce more x0 (t). • Flux v1 represents a reaction between chemical x0 (t) and input c0 which outputs x1 (t). • Flux v2 represents two reactions: the creation of x0 (t) from input c1 , and the removal, or decay, of x0 (t). While this composite term could be split into separate fluxes, combining them reduces the dimensionality of the search. The Brusselator exhibits a Hopf bifurcation at the threshold c0 = 1 + c21 . When c0 < 1 + > 1 + c21 , the system produces oscillations characterised by an attracting limit cycle [20]. Additionally, this system exhibits fast and slow dynamics [20]. In the unstable, oscillatory regime, there are phases where the chemical concentrations spike rapidly and others where the concentrations change more gradually. To capture both dynamical behaviours, we considered constants from both the stable and the unstable regimes. The grammar rule m → x1 (t) was excluded from fluxes v1 and v2 as the chemical x1 (t) is not a reactant in the corresponding reactions. c21 , the system stabilises towards a fixed point, and when c0

15 1 0

Fairen-Velarde

−1 −1

−1 1

1 −1

−1 1

v0 − v 1 v1 − v 2

v0 − v 1 + v 2 −v0 + v1

−v0 + v1 −v0 + v2

1 0

0 1

 −v0 v0 − v1  v1



ẋ

0 −1

 0 −1 1

Brusselator

1 0

 −1 1 0

Lotka-Volterra

SIR

System

p(m → x1 |v2 ) = 0

Non-terminal actions:

m→m+m|m−m|m×m|

m → x0 | x1 | c

v2 = c2

Terminal actions:

m→m+m|m−m|m×m

p(m → x1 |v1 ) = 0 m → x0 | x1 | c

Terminal actions:

Non-terminal actions:

v0 =

m m

p(m → x0 |v2 ) = 0

Non-terminal actions: m→m+m|m−m|m×m

p(m → x1 |v0 ) = 0

m → x0 | x1 | c

Terminal actions:

x0 x1 1 + c0 x20 v1 = c1 − x0

v2 = c1 − x0

v1 = c0 x0

v0 = x20 x1

v2 = c2 x1

v1 = c1 x0 x1

v0 = c0 x0

m→m+m|m−m|m×m|

v1 = c1 x1

p(m → x0 |v1 ) = 0

Non-terminal actions:

√

p(m → x2 |v0 ) = 0

Exclusions

m → x0 | x1 | x2 | c

Terminal actions:

Grammar

True flux expressions   c0 x0 x1 , (standard) v0 = c0 x0 x21 , (squared)  √  c0 x0 x1 , (square root)

Table 1: Details of the coupled ODE systems considered in this study, including their stoichiometries S, differential equations ẋ, selected grammar rules, and grammar rule exclusions via sampling probabilities.

Fairen-Velarde The Fairen-Velarde model describes an oscillatory bacterial respiration system [21]. The dynamical governing equations capture the amount of oxygen, x0 (t), and nutrients, x1 (t), in the system over time. Unlike previous systems considered, the Fairen-Velarde system includes a rational polynomial flux, v0 , which can prove challenging for symbolic regression methods. Furthermore, the functional form of the rational polynomial is significantly more complex than previously considered fluxes, with a minimum complexity of C = 11. To compensate for the increased system complexity, we assume that the functional form of the oxygen and nutrient supplies, v1 and v2 , are known a priori. Therefore, we aim to fully discover the functional form of the consumption flux, v0 , and only estimate the constants [c1 , c2 ] from fluxes v1 and v2 . This assumption is reasonable as the oxygen and nutrient supplies are typically controlled and monitored in an experimental setup. 2.2.1

Performance Evaluation

To measure the performance of symbolic regression methods, generated expressions can be evaluated symbolically or numerically. Symbolic evaluation compares the ground truth mathematical expressions with the top-scoring, candidate expressions. This evaluation can indicate whether the primary objective of symbolic regression, discovering dynamics from data, is met. In comparison, numerical evaluation involves comparing the noiseless trajectory data with the integrated candidate ODE systems, x̂(t). We utilise both numerical and symbolic evaluation to holistically evaluate model performance. We also consider computational runtime and efficiency in our evaluation where efficiency measures the balance between computational expense and the predictive error of generated expressions. This measure standardises the comparison between methods by quantifying whether models make great sacrifices in speed or accuracy in favour of the other. Numerical Evaluation We evaluate our method’s performance quantitatively by computing the normalised mean squared error (NMSE), NMSE(x̂(t), xtrue (t)) =

PT T 2 X 100 2 i=1 (xtrue (ti ) − x̂(ti )) (x (t ) − x̂(t )) = 100 PT true i i 2 2 T × σtrue i=1 i=1 (xtrue (ti ) − x̄true )

between the true, noiseless trajectories xtrue and our estimates x̂(t) (Figure 7). The NMSE measures the model’s error relative to a baseline error obtained by predicting the true data’s mean x̄true across all time points. If the method naïvely predicts the mean of the true data across all time steps, then a NMSE of 100 is attained. We anticipate that well-performing symbolic regression methods should fit the data at least as well as the constant mean x̄true , however, the NMSE can grow very large if the predictions are highly inaccurate. Symbolic Evaluation Because our work focuses on coupled, stoichiometric systems, we assess ability to identify the coupled expressions of each system (Table 2). This evaluation is more lenient than full-system identification, as it does not account for uncoupled expressions and does not penalise the inclusion of incorrect terms. To determine if our method is able to identify a coupled expression, we evaluate its top predicted equations using the reward function (Equation 8) as the criterion for selecting topperforming equations. If the coupled term is present in the relevant state derivatives and the coefficients of this term are sufficiently close, the term is correctly identified.

Table 2: Coupled expressions within each ODE system considered ODE System

ẋ 

SIR (Standard)

SIR (Squared)

SIR (Square root) Lotka-Volterra Brusselator Fairen-Velarde

Number of coupled expressions 

−c0 x0 x1 c0 x0 x1 −c1 x1  c1 x1   −c0 x0 x21 c0 x0 x21 −c1 x1  c1 x1   √ −c0 x0 x1 √ c0 x0 x1 −c1 x1  c1 x1 c0 x0 −c1 x0 x1 c1 x0 x1 − c2 x2 2 x0 x1 −c0 x0 + c1 − x0 −x20 x1 +c0 x0 # " xx − 1+c0 0 1x2 + c1 − x0 0 x0 x1 − 1+c 2 + c2 0x

2 1 2 1

Efficiency Evaluation To facilitate a fair comparison of the tradeoff between speed and accuracy, we define a composite efficiency metric, Efficiency =

1 Runtime0.5 × NMSE0.5

which assigns equal weight to computational expense and predictive error. This metric penalises both methods that achieve marginal accuracy gains at great computational costs, and those that run quickly but fail to produce accurate solutions. 2.2.2

Method Comparison

To contextualise our results, we benchmark our approach against a diverse range of popular equation discovery methods using the evaluation metrics outlined in Section 2.2.1. Key characteristics of the benchmarked methods are detailed in Table 3. Of the methods considered, ours is the only method that accounts for known stoichiometries and estimates system fluxes as opposed to the system ODEs. To the best of our knowledge, there are no existing dynamical symbolic regression approaches which allow a stoichiometry to be imposed. While Reactive SINDy [10, 11], SISR [14], and KinFormer [15] allow stoichiometries and flux expressions to be estimated, they are confined to the restrictive assumption of mass action kinetics which does not hold for all of the ODE systems considered in this study. Furthermore, these methods estimate stoichiometries as opposed to imposing them. In the absence of directly comparable physics-informed methods for flux-based symbolic regression, we benchmark our approach against a wide range of top-performing generalised approaches. To ensure a fair comparison against these baselines, we evaluate two variants of our framework: our standard stoichiometry-aware method (denoted ‘Ours’) and an ablation without imposed stoichiometries (‘Ours (No Stoich.)’). In the Supplementary Materials, Section S4, we explore two additional 17

Table 3: Overview of methods included in our comparative study. Traits include whether finite differencing of state trajectories are required (F.D.), whether the method estimates all ODEs in a system jointly or independently (Joint Disc.), whether the method accounts for the flux-based nature of the dynamical systems (Fluxes), and what general category the method falls under. Method

F.D.

Joint Disc.

Fluxes

Category

Ref

Ours Ours (No Stoic.) SPL SINDy PySR ProGED ODEFormer

✗ ✗ ✓ ✓ ✓ ✗ ✗

✓ ✓ ✗ ✓ ✗ ✓ ✓

✓ ✗ ✗ ✗ ✗ ✗ ✗

MCGS MCGS MCTS Sparse regression Genetic programming Monte Carlo, grammar-based Pretrained transformer

[17] [7, 22] [23] [24] [3]

ablations that determine the effects of state merging and the grammar rule exclusions presented in Table 1. We evaluate each method using its default or recommended hyperparameters. To prevent a biased evaluation, we applied our framework’s grammar rules to the grammar-based methods, SPL and ProGED. While PySR is not grammar-based, it allows the specification of operations and functions which we configured to match our grammar. We were, however, unable to apply our framework’s associated grammar rule sampling probabilities to these alternative methods as they do not allow sampling at the flux-based level. We selected third-order polynomial features for SINDy’s library of candidate functions, including square-root and rational terms for SIR and Fairen-Velarde respectively to better mimic these systems. Unfortunately, SINDy cannot multiply or nest features, so it is unable to fully reconstruct the square root SIR ODEs and the Fairen-Velarde ODEs using the available feature set. As a pre-trained, zero-shot inference method, ODEFormer required no additional specification.

Results

3.1

Performance of FluxDisco

Figure 6 allows us to visually compare FluxDisco’s estimated state trajectories against the noiseless and observed trajectories across all ODE systems considered in this study. Across all systems and regimes, the estimated trajectories closely align with the ground truth trajectories, suggesting that our method accurately recovers system dynamics, even in the presence of noise. Our method also demonstrates high numerical accuracy across all systems, with the maximum observed trajectory reconstruction error remaining below 2%. As expected, Figure 7 shows that as noise is introduced, the error generally increases. The Fairen-Velarde system produces the highest errors across all systems considered, reflecting the complex nature of its ODEs. After Fairen-Velarde, the unstable Brusselator yields the highest errors, likely due to the presence of both fast and slow dynamics. The ablation studies in Section S4 of the Supplementary Materials illustrate that two of FluxDisco’s key features, state merging and grammar rule exclusions, improve numerical accuracy and efficiency, but these gains diminish as the computational budget of the search increases. The results of the ablation without an imposed stoichiometry, shown in Section 3.2, demonstrate that imposing known stoichiometries improves both numerical and symbolic performance, as well as

Figure 6: Comparison of estimated, noiseless, and observed trajectories.

Figure 7: Normalised MSE (NMSE) across all benchmarked systems and levels of noise. Errors are averaged across fluxes per system and noise level. computational efficiency.

3.2

Method Comparison

Numerical Results Figure 8 shows the NMSE for the benchmarked methods across all considered ODE systems and noise levels. The bars indicate the mean error across all noise levels, and the error bars represent the full range of these errors. Our method produces the lowest errors across all systems (< 2%), remaining well below the baseline error. The version of our framework which does not impose stoichiometry demonstrates considerably worse performance than the stoichiometry-aware variant, illustrating the benefit of enforcing known system structure in ODE discovery. While we expect well-performing methods to yield errors lower than the naïve baseline, all methods besides ours produced errors greater than the baseline for some experiments. Some of the errors exceed the baseline by orders of magnitude, illustrating that these approaches are prone to divergent behaviour and fail to reconstruct system dynamics through their estimated expressions. Symbolic Results Figure 9 presents the identification accuracy of the coupled expressions for all benchmarked methods. Our method achieves high accuracy across all systems except the Fairen-Velarde system which remains entirely unidentified by all methods. Our ablation without an imposed stoichiometry performs as poorly as the remaining benchmarked systems, demonstrating that enforcing known stoichiometries allows for better identification of coupled dynamics. Efficiency Results Although our method achieves the lowest error and highest identification rate among the benchmarked methods (Figure 8 and Figure 9), this accuracy comes at the expense of computational speed. As shown in Figure 10a, the runtime for both variants of our method exceeds that of alternative methods by orders of magnitude. Box-plots of the efficiency scores (Figure 10b) show that SPL, SINDy, PySR, and ODEFormer are highly volatile, with lower whiskers extending close to zero. While ProGED 20

Figure 8: Comparison of Normalised Mean Squared Error (NMSE) across benchmarked models. Error lines indicate the range of NMSE across varying noise levels. The vertical line denotes the baseline error associated with predicting the data mean across all time points (NMSE = 100%).

Figure 9: Coupled term identification accuracy in top-predicted equations across benchmarked methods for each ODE system. Results are averaged across all noise levels. and the variant of our framework without imposed stoichiometry exhibit lower volatility, these methods are unable to achieve competitive efficiency scores. Our stoichiometrically-aware method and ODEFormer (Opt) exhibit both tight interquartile ranges and high efficiency scores. The discrepancy in efficiency between the two variants of our framework demonstrates that imposing a known stoichiometry successfully constrains the search space, allowing for better identification of dynamics. Although ODEFormer (Opt) has a higher median efficiency score than our framework, our approach demonstrates a higher peak efficiency. We argue that the efficiency of ODEFormer and ODEFormer (Opt) shown in Figure 10b is overstated given that it does not account for the multi-day pre-training regime that the model underwent prior to inference. Should this pre-training be taken into account, the ODEFormer results would seem less competitive than our approach.

Discussion

In benchmarking our framework against a range of top-performing symbolic regression methods (Table 3), we find that our approach produces the lowest trajectory reconstruction errors and the highest identification rate of coupled dynamics. The benchmarked methods show inconsistent relative performance, although all fall short of our approach. While our framework yields the longest runtimes, its increased accuracy justifies this trade-off as equation discovery is typically an offline task and is rarely subject to strict runtime constraints. To transparently establish the boundaries of our approach, we have included a challenging benchmark, the Fairen-Velarde system, which has comparatively complex dynamics. Our 22

(a) Runtime (seconds)

(b) Efficiency score

Figure 10: Comparison of runtime and efficiency across benchmarked models across ODE systems and noise levels. The efficiency score equally balances error (NMSE) and runtime (seconds).

method discovered expressions for this system that produce well-fitting trajectories, but do not precisely match the ground truth expressions. This example highlights a common limitation of symbolic regression methods, with none of the benchmarked approaches successfully recovering its governing equations. Identifying exact mathematical expressions becomes increasingly difficult in complex systems, both due to an exploding combinatorial search space and the fact that different functional forms can produce similar trajectories. One of the greatest limitations of this work is its poor ability to scale to high-dimensional, complex ODE systems. This drawback can largely be attributed to joint flux estimation. While critical for producing physically-adherent equations, joint estimation requires the algorithm produce the correct expressions across all fluxes within the same graph node. As the number of fluxes, or their complexity, increases, the probability of sampling a full set of correct expressions diminishes rapidly. This limitation is shared by most symbolic regression methods—as expression complexity or system dimensionality increases, the expression search space explodes. A further limitation of our work is that our method imposes a hard constraint on the system’s stoichiometry. Consequently, stoichiometry misspecification will lead to inaccurate estimation of flux expressions. While our framework is designed for scenarios where the stoichiometry is known a priori, this constraint can be relaxed. Firstly, as demonstrated in our ablation study, our method can be used without modification by setting the stoichiometry matrix to the identity matrix. In this setting, the full system of ODEs become the target of estimation, though the coupling between the ODEs is lost and they are estimated independently. Alternatively, our method could be modified to use a soft constraint, allowing both the stoichiometry and flux expressions to be estimated while penalising deviations from an expected stoichiometry within the reward function. To our knowledge, no existing methods for flux-based systems employ this type of soft constraint. Existing literature focuses on estimating stoichiometry and flux expressions within the confines of mass action kinetics and low-order reactions [10, 11, 14, 15].

Conclusion

This study addresses the challenge of extracting physically meaningful governing equations from noisy, multi-dimensional data in stoichiometric dynamical systems. Most existing dynamical symbolic regression methods are generalised and are therefore not suitable for the constraints associated with flux-based systems [3, 7, 17, 23, 24]. To overcome these challenges and identify physically-adherent governing equations, we propose a method that jointly estimates the functional form of system fluxes while imposing a known system stoichiometry. To do this, we adapt and specialise the MCGS algorithm [6] for the purpose of joint flux discovery. We demonstrate the robustness of our method across a diverse set of physical systems and varying levels of observation noise. Our results illustrate that our proposed method, FluxDisco, accurately recovers governing dynamics where existing approaches fall short. A promising avenue for future work is relaxing the assumption of full system observability. This assumption is often violated when considering data from real-world physical systems. Therefore, extending our framework to handle latent states would greatly improve its practical utility. This may involve latent state trajectory reconstruction using neural networks [25, 26], or traditional approaches such as time-delay embeddings [27]. Furthermore, addressing the scalability challenges associated with joint flux estimation is critical for applying our method to high-dimensional, complex ODE systems.

References [1] [2] [3] [4] [5] [6] [7]

[8] [9] [10] [11] [12]

[13] [14] [15]

[16]

[17] [18]

[19] [20] [21] [22]

Stephen A Billings. Nonlinear System Identification: NARMAX Methods in the Time, Frequency, and Spatio–Temporal Domains. John Wiley & Sons, 2013. isbn: 9781118535561. Ricky T. Q. Chen et al. “Neural Ordinary Differential Equations”. In: Advances in Neural Information Processing Systems. Ed. by S. Bengio et al. Vol. 31. Curran Associates, Inc., 2018. Stéphane d’Ascoli et al. “ODEFormer: Symbolic Regression of Dynamical Systems with Transformers”. In: arXiv preprint arXiv:2310.05573v1 (2023). Richard S Sutton and Andrew G Barto. Reinforcement Learning: An Introduction. 2nd ed. MIT Press Cambridge, 2018, pp. 159–190. Maciej Świechowski et al. “Monte Carlo Tree Search: A Review of Recent Modifications and Applications”. In: Artificial Intelligence Review 56.3 (2023), pp. 2497–2562. Edouard Leurent and Odalric-Ambrym Maillard. “Monte-Carlo Graph Search: the Value of Merging Similar States”. In: Asian Conference on Machine Learning. PMLR. 2020, pp. 577–592. Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. “Discovering governing equations from data: Sparse identification of nonlinear dynamical systems”. In: Proceedings of the National Academy of Sciences 113.15 (2016), pp. 3932–3937. Sören Becker et al. “Predicting Ordinary Differential Equations with Transformers”. In: International Conference on Machine Learning. PMLR. 2023, pp. 1978–2002. Martin Vastl et al. “Symformer: End-to-end symbolic regression using transformer-based architecture”. In: IEEE Access 12 (2024), pp. 37840–37849. Moritz Hoffmann, Christoph Fröhner, and Frank Noé. “Reactive SINDy: Discovering governing reactions from concentration data”. In: The Journal of Chemical Physics 150.2 (2019). Richard Jiang et al. “Identification of dynamic mass-action biochemical reaction networks using sparse Bayesian methods”. In: PLOS Computational Biology 18.1 (2022), e1009830. Jason Frank. “Lecture 2: Reaction kinetics”. University of Amsterdam. Introduction to Applied Mathematics and Modelling. Lecture notes accessed 2026. 2012. url: https://webspace.science. uu.nl/~frank011/Classes/modelling/L2.pdf. Darren J Wilkinson. Stochastic Modelling for Systems Biology. Chapman and Hall/CRC, 2018. Manuel Palma Banos et al. “Stoichiometrically-informed symbolic regression for extracting chemical reaction mechanisms from data”. In: Digital Discovery 5.5 (2026), pp. 2325–2341. Jindou Chen et al. “KinFormer: Generalizable Dynamical Symbolic Regression for Catalytic Organic Reaction Kinetics”. In: The Thirteenth International Conference on Learning Representations. 2025. Shipra Agrawal and Navin Goyal. “Analysis of Thompson Sampling for the Multi-armed Bandit Problem”. In: Conference on Learning Theory. JMLR Workshop and Conference Proceedings. 2012, pp. 39–1. Fangzheng Sun et al. “Symbolic Physics Learner: Discovering governing equations via Monte Carlo tree search”. In: arXiv preprint arXiv:2205.13134 (2022). Mary Ann Branch, Thomas F. Coleman, and Yuying Li. “A Subspace, Interior, and Conjugate Gradient Method for Large-Scale Bound-Constrained Minimization Problems”. In: SIAM Journal on Scientific Computing 21.1 (1999), pp. 1–23. doi: 10.1137/S1064827595289108. Y Li and T Coleman. “On the Convergence of Reflective Newton Methods for Large-Scale Nonlinear Minimization Subject to Bounds”. In: Math. Program 67.1-3 (1994), pp. 189–224. Maximilian Engel and Guillermo Olicón-Méndez. “A singular perturbation analysis for the Brusselator”. In: Journal of Differential Equations 457 (2026), p. 113866. V Fairen and MG Velarde. “Time-Periodic Oscillations in a Model for the Respiratory Process of a Bacterial Culture”. In: Journal of Mathematical Biology 8.2 (1979), pp. 147–157. Alan A. Kaptanoglu et al. “PySINDy: A comprehensive Python package for robust sparse system identification”. In: Journal of Open Source Software 7.69 (2022), p. 3994. doi: 10.21105/joss. 03994. url: https://doi.org/10.21105/joss.03994.

[23] [24] [25]

[26] [27]

Miles Cranmer. “Interpretable Machine Learning for Science with PySR and SymbolicRegression.jl”. In: arXiv preprint arXiv:2305.01582 (2023). Nina Omejc et al. “Probabilistic grammars for modeling dynamical systems from coarse, noisy, and partial data”. In: Machine Learning 113.10 (2024), pp. 7689–7721. Gevik Grigorian, Sandip V George, and Simon Arridge. “Learning Governing Equations of Unobserved States in Dynamical Systems”. In: Physica D: Nonlinear Phenomena 472 (2025), p. 134499. Wei Liu et al. “Structured Kolmogorov-Arnold Neural ODEs for Interpretable Learning and Symbolic Discovery of Nonlinear Dynamics”. In: arXiv preprint arXiv:2506.18339v3 (2026). Steven L Brunton et al. “Chaos as an intermittently forced linear system”. In: Nature Communications 8.1 (2017), p. 19.

Supplementary Material S1

Methodology Details

This section expands on core components of our methodology. Specifically, we justify our use of the stochastic reward variant of the Monte Carlo Graph Search (MCGS) algorithm (Section S1.1), detail our formulation of the mean reward credible intervals required in the stochastic reward setting (Section S1.2), and present the complete reward calculation procedure (Section S1.3).

S1.1

Deterministic vs Stochastic Rewards

The MCGS algorithm supports both deterministic and stochastic reward settings [6]. While either formulation can be applied to the task of symbolic regression, we opt to use the stochastic formulation. In the deterministic setting of MCGS, rewards are only generated at terminal states corresponding to complete mathematical expressions. When flux expressions are partially-complete, we cannot evaluate the error between the partial expressions and the data. Therefore, non-terminal states receive a zero reward, resulting in a sparse-reward MDP. In such environments, an agent cannot easily distinguish between promising and unpromising partial expressions, resulting in an exhaustive exploration of the search space that scales poorly with ODE dimension and grammar size. This inefficiency is compounded by the optimism inherent to the original MCGS algorithm [6]. The deterministic selection rule selects the next state s′ ∈ S(s) according to: s ← arg ′max r(s′ ) + γU(s′ ) s ∈S(s)

where r(s′ ) is the deterministic reward yielded by state s′ and U(s′ ) is the upper bound on the value of state s′ based on Bellman optimality [6]. Because non-terminal states yield zero rewards, selection is driven solely by the optimistically initialised upper bounds of values of child states, U (s′ ) ∀s′ ∈ S(s). The optimistic initialisation of unexplored nodes suppresses signals from discounted terminal states and prevents them from meaningfully guiding the search. Therefore, the search algorithm prioritises unexplored nodes, exploring the state space almost uniformly until a significant proportion of terminal states are reached. Reaching a large proportion of terminal states may be infeasible in large state spaces, such as the ones we are attempting to search. Without sufficient coverage of the state space, the search strategy remains undirected and uniform in deterministic settings. To overcome the undirected nature of the search in sparse-reward settings, we employ the stochastic reward variant of MCGS (Main Text Section 2.1.1). Rather than assigning a zero reward to non-terminal nodes, the flux expressions are completed using a rollout policy to obtain an estimate of a node’s potential reward. By introducing dense, sampled rewards, the agent can better discriminate between neighbouring nodes leading to a more efficient navigation of the search space.

S1.2

Reward Confidence/Credible Intervals

The stochastic formulation of MCGS utilises confidence intervals for the mean empirical reward of all states (l(s), u(s))∀s to account for reward stochasticity. Leurent and Maillard [6] construct 27

these intervals using the Binary Kullback-Leibler (BKL) divergence. The tightness of the BKL bounds depends on the ratio between the total number of rewards sampled across the graph and the number of local samples drawn. In large state-action spaces, states are rarely revisited so this ratio remains high. Therefore, the confidence bounds remain as initialised throughout the search, resulting in little differentiation between nodes during selection. Consequently, this formulation is overly conservative for equation discovery due to its expansive search space. As introduced in the main text, we leverage Thompson sampling [16] to overcome the conservatism of BKL bounds. While Thompson sampling is designed for binary reward systems, the framework can be generalised to the scalar, bounded reward setting. The scheme assumes an uninformative prior on the mean of the empirical reward distribution of a node. When rewards are sampled, the prior is updated based on the Beta-Bernoulli conjugate prior: Prior : Θ ∼ Beta(α = 1, β = 1)

(uniform distribution)

Likelihood : (R|Θ = θ) ∼ Bernoulli(θ) Posterior : (Θ|R = r) ∼ Beta α = 1 +

n X i=1

ri , β = 1 + n −

n X

! ri

i=1

When rewards are scalar, we use a heuristic to maintain the Bayesian update framework. A scalar reward ri ∈ [0, 1] can be considered as ri units of success and 1 − ri units of failure. Therefore, when a reward ri is sampled, the success parameter α of the posterior distribution can be incremented by the reward, ri , and the failure parameter can be incremented by 1 − ri . The resulting posterior distribution is: ! n n X X Beta α = 1 + ri , β = 1 + (1 − ri ) ≡ Beta(α = 1 + nr̄, β = 1 + n(1 − r̄)) (S1) i=1

i=1

While the conjugate prior framework technically assumes integer parameters, due to the Bernoulli likelihood function, this relaxation serves as a useful heuristic to derive a pseudo posterior distribution for the mean reward. The credible intervals (l(s), u(s)) are then extracted via the quantiles of this distribution, as detailed in Equation 7 of the main text. The primary motivation for this approach is that it does not rely on the ratio between the total number of node visits across the graph to the current node’s visit count, improving its suitability for large search spaces. Furthermore, the Beta credible intervals are immediately responsive to sampled rewards, whereas the BKL intervals require a significant number of samples before they respond to the reward signal.

S1.3

Reward Calculation

Algorithm S1 details the complete procedure for evaluating the reward associated with a sampled set of candidate flux expressions post-rollout.

Hyperparameters

The hyperparameter choices for our framework, FluxDisco, are outlined in Table S1.

Data Simulation

The data used in this study are simulated. Algorithm S2 details the simulation process, and Table S2 outlines the simulation parameters of each benchmarked ODE system. 28

Algorithm S1 Reward calculation Require: Candidate flux expressions: v̂(x(t)) = [v̂0 (x(t)), . . . , v̂F −1 (x(t))] Initial conditions: x(0) = [x0 (0), . . . , xD−1 (0)] Stoichiometry matrix: S ∈ RD×F

Algorithm: 1: Constant folding: Apply constant folding to candidate fluxes v̂(x(t)) to obtain simplified expressions, v̂fold (x(t)). 2: Compute complexity: Calculate the system complexity C from the folded expressions. 3: Apply stoichiometry: Join the folded flux expressions using the stoichiometry matrix: ˆ = Sv̂fold (x(t)) ẋ(t) 4: repeat 5: Propose and apply constants: Propose constants c via optimisation algorithm and

ˆ substitute them into the system equations, ẋ(t). 6: Solve ODEs: ˆ x̂(t) = ODESolver ẋ(t), x(0) 7:

Compute optimisation error: Derive the sum of squared errors (SSE) between estimates and observed data to guide constant optimisation: SSE(x̂(t), x(t)) =

D−1 N X T XX

2 (n) (n) x̂d (ti ) − xd (ti )

d=0 n=1 i=1

8: until Convergence met or max iterations reached.

9: Compute total error: Using the optimised constants, compute total error.

MSEtotal = 10: Compute reward:

SSE(x̂(t), x(t)) N ×T

MSEtotal R = η exp − τ C

return Reward R, folded flux expressions v̂fold (x(t)), optimised constants c.

Algorithm S2 Data Simulation Require: Noise level: ϵ Ground truth ODEs: ẋ(t) Initial conditions: x(0) = [x0 (0), . . . , xD−1 (0)] Constant values: c ODE solver: ODESolver

Algorithm: 1: Get noiseless trajectories: Solve specified ODEs: x(t) = ODESolver(ẋ(t), x(0)) 2: if ϵ > 0 then 3:

Calculate variable ranges: For each system variable d ∈ {0, . . . , D − 1}, compute the range: ∆xd = max(xd (t)) − min(xd (t)) ∀d t

Generate Gaussian noise: Randomly sample noise from the standard normal distribution to produce Z ∈ RT ×D where each element is z[t,d] ∼ N (0, 1)

(i.i.d.)

Scale and apply noise: Apply scaled noise to the noiseless trajectories: x̃d (t) = xd (t) + ϵ · ∆xd · z[t,d]

∀t, d

x̃(t) = [x̃0 (t), . . . , x̃D−1 (t)] 6:

Ensure non-negativity: The systems considered (SIR, Lotka-Volterra, Brusselator, and Fairen-Velarde) cannot take on negative values, so correct for any negativity induced by noise: x̃(t) = max(x̃(t), 0)

7: 8:

if system is SIR then Normalise proportions: Normalise data within each time step to ensure variables sum to 1: x̃d (t) x̃d (t) = PD−1 ∀t, d d=0 x̃d (t)

end if

10: else if ϵ = 0 then 11: Return clean trajectories: If no noise is specified, return the noiseless trajectories:

x̃(t) = x(t) 12: end if

return trajectories x̃(t).

Table S1: Selected hyperparameter values Symbol

Description

Value

Fairen-Velarde Modifications

Discount factor for future rewards

0.9

Accuracy measure for propagation (Algorithm 4, [6])

0.01

Temperature in reward function

0.005

Increased to 0.05 to increase reward scale.

Parsimony coefficient in reward function

0.99

Increased to 0.995 to allow for more complex expressions.

Maximum flux expression depth

Increased to 12 to allow for more complex expressions.

Episodes

100

Decreased to 40 to compensate for increased expression depth.

Max iterations per episode

4κ + 2

Rollouts per node post-selection

Warm start rollouts per node

Number of data simulations, or realisations, per ODE system

Table S2: Simulation parameters for the benchmarked ODE systems, including initial conditions x(0), constant values c, time spans, sampling frequencies, and noise levels ϵ. System

x(0)

Time Span

Time Points

Noise

SIR (Standard) SIR (Squared) SIR (Square root) Lotka-Volterra Brusselator (Stable) Brusselator (Unstable) Fairen-Velarde

[0.98, 0.02, 0] [0.88, 0.12, 0] [0.98, 0.02, 0] [10, 5] [0.5, 2] [0.5, 2] [10, 25]

[0.4, 0.1] [1.3, 0.08] [0.15, 0.06] [1, 0.1, 1.5] [1.5, 1] [0.5, 2] [0.5, 15, 10]

[0, 80] [0, 80] [0, 80] [0, 30] [0, 30] [0, 30] [0, 60]

80 80 80 150 150 150 150

[0, 0.02, 0.05] [0, 0.02, 0.05] [0, 0.02, 0.05] [0, 0.03, 0.07] [0, 0.03, 0.07] [0, 0.03, 0.07] [0, 0.03, 0.07]

Because the ODE systems considered have unit constraints, we need to process the noisy trajectories to ensure these constraints are adhered to. All benchmarked systems require observations to be non-negative, therefore any negative values caused by the applied noise are set to zero. While this processing step introduces a slight positive bias, it is necessary to produce data that is physically realistic. The SIR epidemic model has an additional unit constraint. The system variables represent proportions and should sum to one. Therefore, we normalise the data per time step after zeroing any negative values to meet this extra constraint. This normalisation step introduces additional bias by correlating the noise across the system variables. Table S2 shows that the number of simulated time points and noise levels for the SIR system deviate from the remaining systems. For the SIR systems, initial conditions and constants 31

values were selected to ensure that across all regimes, infections rose to similar peaks around approximately the same time period for fair comparison. Based on these parameters, a time span of [0, 80] days was chosen to ensure that a full epidemic is realised while avoiding an excessively long tail. While the other dynamical systems are simulated across 150 time points, the epidemic models are restricted to 80 time points. This choice reflects the practicalities of epidemiological data collection, where data is typically reported at most daily. The noise levels for the SIR system were reduced to compensate for the lower sampling frequency and the bounded data range. SIR trajectories lie close to the boundaries of the data range for extended periods, therefore they are very susceptible to the biases introduced by the post-processing, unit constraint corrections. Lowering the noise for this system mitigates the effect of this bias. Whilst our framework allows for multiple realisations from the dynamical systems, we focus on the more challenging single-realisation setup, using only a single initial condition per system.

Ablation Studies

To evaluate the efficacy of key features of FluxDisco and validate our design decisions, we conduct various ablation studies. We demonstrate the impact of imposing known stoichiometries in Section 3 of the main text, where the ablation without imposed stoichiometries shows markedly worse numerical and symbolic performance. Here, we present two additional ablation studies that isolate the effects of state merging, and excluding grammar rules by zeroing select sampling probabilities.

S4.1

State Merging

To evaluate whether state merging improves symbolic regression performance, we compare our method, FluxDisco, against an ablation which does not merge states, structuring the search space as a tree as opposed to a graph. To assess performance relative to the search budget, we recorded key metrics every five episodes. Figure S1 presents our results, evaluating performance across four key metrics: efficiency score, normalised mean squared error (NMSE), computational runtime, and the proportion of missed state merges. For this analysis, we exclude results from the Fairen-Velarde system. The Fairen-Velarde system is our most challenging benchmark with complex dynamics necessitating deeper search structures. To compensate for the increased computational demands of this search, fewer episodes were conducted for this system (Table S1). Since the Fairen-Velarde system has unique hyperparameters, including episode count, it is excluded from this ablation study to ensure a uniform comparison. Our graph search method achieves higher median efficiency scores and lower median errors (NMSE) than the tree search ablation without state merging. While there is stochastic variation in the individual ODE system results (Figure S1, right column), the aggregated median scores (Figure S1, left column) illustrate a clear trend. The performance gap is most noticeable at lower computational budgets (episode counts), where the improved sample efficiency and wider information propagation of our graph-based method allow for faster discovery of high-quality solutions. This advantage is supported by the steep increase in the proportion of missed state merges by the tree search method in early stages of the search. As the computational budget increases, the median performance of the two methods converge, illustrating that both methods find good solutions when provided sufficient search budget. While state merging introduces overheads that increase runtime, the difference in median runtimes is marginal. Runtimes for the individual ODE systems vary substantially and in some

Figure S1: Graph search vs. tree search ablation study results. The left panels show median scores aggregated across all conditions, and the right panels break down median scores by ODE system. 33

cases, the long-run tree search runtimes exceed those of the graph searches (e.g. Lotka-Volterra, unstable Brusselator).

S4.2

Grammar Rule Exclusions

Our framework is physics-informed in that we enforce known system structures by constraining system stoichiometry. Furthermore, we allow prior domain knowledge to be incorporated through user-specified grammar rule sampling probabilities (Section 2.1.3). In the main text, we zero specific sampling probabilities for fluxes from the SIR, Lotka-Volterra, and Brusselator systems to reflect prior physical knowledge. To determine the effect of these exclusions, we run an ablation study using uniform sampling probabilities across all defined grammar rules. Figure S2 compares the search performance with and without grammar rule exclusions. The excluded grammar rules, shown in Table 1 of the main text, all correspond to terminal actions. Consequently, the ablation results produce graphs with more terminal nodes and shorter selection paths. Since search episodes terminate once either a terminal set of equations are constructed, or a maximum number of iterations are reached, having more terminal actions will terminate episodes earlier, reducing the overall runtime. Conversely, excluding these terminal grammar rules should yield longer selection paths, larger graphs with more nodes, and comparatively fewer terminal nodes. While the ablation corresponds to faster runtimes due to a higher density of terminal nodes, its efficiency scores are lower than the main method’s scores due to poorer candidate solutions being identified. However, as the search progresses, the ablation is able to find solutions with sufficiently low errors causing the efficiencies of both approaches to converge.

Figure S2: Search results with grammar rule exclusions vs. ablation without exclusions. The left panels show median scores aggregated across all conditions, and the right panels break down median scores by ODE system.

Record · ID 660814 · SHA-256 f92b4bace21ba67e
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.