Improving Surrogate Model Robustness to Perturbations for Dynamical Systems Through Machine Learning and Data Assimilation

Complex networks provide valuable insights into real-world systems, such as the spread of epidemics [1] and the understanding of biochemical and neuronal processes [2]. However, modeling these systems becomes computationally complex when dealing with large state spaces. To address this challenge, methods like Proper Orthogonal Decomposition (POD) ([3], [4], [5], [6]) are employed to capture patterns while operating in reduced dimensions. While dimensionality reduction techniques like POD are effective, their sensitivity to data necessitates robust surrogate models that can maintain accuracy under input perturbations. Consequently, the question arises regarding how to enhance a surrogate model to accommodate perturbed data sets. This study proposes a framework designed to improve the robustness of surrogate models, ensuring their reliability under varying input conditions.

Surrogate models serve as approximations of complex real-world systems, which are often modeled using ordinary differential equations (ODEs) or partial differential equations (PDEs). In many cases, the complexity of these systems makes it difficult to fully capture their underlying processes. To address the challenges of modeling high-dimensional and chaotic systems, several machine learning (ML) approaches have been proposed to construct surrogate models that provide accurate and computationally efficient approximations. Several ML-based approaches have been developed to construct surrogate models that provide accurate and computationally efficient approximations. These include analog models [7], recurrent neural networks (RNNs) [8, 9], residual neural networks that approximate resolvents [10, 11], and differential equation-based models such as in [12, 13, 14, 15]. A notable study [10] integrates machine learning and data assimilation (DA) to improve predictions in scenarios with sparse observations. However, the DA step is computationally expensive, limiting its practical applicability. To address model inaccuracies, researchers have explored alternative approaches such as weak-constraint DA methods [16] and ML-DA frameworks [17].

The contributions of our paper can be summarized as follows:

1. 1.

   We introduce a novel framework (Figure 1) that improves the robustness of surrogate models against input perturbations by integrating data assimilation and machine learning techniques.

2. 2.

   We conduct extensive experiments on dynamical systems represented as graphs, specifically the diffusion system ($\mathbb{D}$) and the chemical Brusselator model ($\mathbb{C}$) discussed in Section 2. For dynamical systems on graphs, we propose a dynamic optimization step, described in Sections  5.2 and  5.3, to obtain a sparse graph and reduce memory complexity.

   1. (a)

      For the diffusion equation ($\mathbb{D}$) on graphs (Section 2.2), we derive conditions (Theorem 3) to guide the constraints in our dynamic optimization step.

   2. (b)

      Section 5.3 presents the dynamic optimization step for reaction-diffusion systems on graphs, detailing the required constraints.

3. 3.

   Section 1 presents experimental results for the framework on both linear and non-linear dynamical systems represented on graphs, using the POD-based surrogate model. Empirical results in Table 1 demonstrate that our framework enhances surrogate model performance under input perturbations.

4. 4.

   In Section 7, we demonstrate how our framework can be integrated into a general setting to reinforce neural ODE-based surrogate models trained on data-driven systems. The improved models exhibit enhanced performance and robustness against input perturbations, outperforming both POD-based models and the original neural ODE solutions, as shown in Figures 3, 4, and 5.

We present the fundamental tools and techniques essential for understanding the methodologies discussed in this paper.

Surrogate modeling techniques, such as the reduced-order model (ROM) discussed in Section 2.3, provide computational efficiency but are highly sensitive to input variations, as indicated in Equation 3 Consequently, when initialized with a new random condition, the surrogate model ($M(\cdot)$) efficiently approximates the system trajectory but may fail to capture deviations from the true state.
The surrogate model is used to obtain compressed representations of state vectors at various timesteps, which are treated as noisy observations. The observations are compressed for better memory efficiency if the model operates on the state dimension, as in the neural ODE surrogate model (Section 7). Let $h:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}$ (where $k<n$) denote the state observation relationship. $e_{mh}(t_{k+1})$ and $e_{ml}(t_{k+1})$ denote the high and low wavelength components of error, capturing deviations between the true state and the state obtained from the surrogate model ($M(x_{k})$) at time $t_{k+1}$. $e_{oh}(t_{k})$ and $e_{ol}(t_{k})$ denote the high and low-wavelength components of error that capture the difference between the observation and the state at time $t_{k}$. $x_{k+1}$ denotes the true state of the system at time $t_{k+1}$. $M(\cdot)$ denotes the forward model. For instance, if the forward model $M(\cdot)$ uses the Euler forward scheme and the surrogate model applies POD to the vector field $\dot{x}=f(x(t))$, then:

where $h$ denotes the step size. The state forward dynamics and state observation relationship are given by:

The key steps of the framework, as illustrated in Figure 1, are as follows. For brevity, the framework is discussed in the context of dynamical systems represented on graphs. For dynamical systems that do not involve graphs, the dynamic optimization step is not required. The framework’s applicability to general modeling scenarios is demonstrated in Section 7, where a neural ODE-based surrogate model is used to learn from data.

1. 1.

   Dynamic optimization. For dynamical systems on graphs, a dynamic optimization problem is formulated to extract a sparse graph from known solution trajectory snapshots. The sparse graph improves memory efficiency in the filtering step (Step 4) of the framework. The original graph can be replaced with the sparse graph, further enhancing memory efficiency. This step outputs the graph Laplacian matrix of the sparse graph, given by $L_{1}=B^{T}\text{diag}(w^{*})B$, where the weights $w^{*}$ for linear and nonlinear dynamical systems on graphs are determined by solving the optimization problems detailed in Sections 5.2 and 5.3.

2. 2.

   Hierarchical clustering. This step quantifies deviations between the surrogate model (e.g., POD) and the actual system trajectories. Given a dataset $\mathcal{D}=\{x_{1},x_{2},\ldots,x_{T}\}$ and surrogate model predictions $\mathcal{P}=\{z_{1},z_{2},\ldots,z_{T}\}$, we apply hierarchical clustering to partition $\mathcal{P}$ into $p$ clusters. The assigned cluster labels are represented as $Y=\{y_{1},y_{2},\ldots,y_{T}\}$, where $y_{t}$ denotes the cluster assignment for data point $z_{t}$. For each cluster $C_{i}$, we construct a set of triplets $\{z_{j},e_{j},\hat{e}_{j}\}_{j=1}^{N_{C_{i}}}$, where $N_{C_{i}}$ is the number of triplets in cluster $C_{i}$. For reduced order model using POD (Section 2.3), we denote $e_{j}=x_{j}-(\rho^{T}{z}_{j}+\bar{x}),\hat{e}_{j}={z}_{j}-\rho(x_{j}-\bar{x})$. The center of cluster $C_{i}$ is the triplet $\{{z}^{C_{i}}_{m},e^{C_{i}}_{m},\hat{e}^{C_{i}}_{m}\}$, where $e^{C_{i}}_{m}=\frac{\sum^{N_{C_{i}}}_{j=1}e_{j}}{N_{C_{i}}},\hat{e}^{C_{i}}_{m}=\frac{\sum^{N_{C_{i}}}_{j=1}\hat{e}_{j}}{N_{C_{i}}},{z}^{C_{i}}_{m}=\frac{\sum^{N_{C_{i}}}_{j=1}{z}_{j}}{N_{C_{i}}}$. As shown in Figure 2, a visualization of the clusters is provided, with $I_{C_{i}}$ representing the indices of data points within cluster $C_{i}$.

3. 3.

   Inverse Problem for Estimating Weights Between Clusters. For an initial condition $x_{0}$ at time step $t_{0}$, we obtain a compressed representation of the data point $z_{0}$. For example, when using the POD method, we get the compressed representation of $x_{0}$ as ${z}_{0}=\rho(x_{0}-\bar{x})$. From the clusters formed in Step 3, we assign the cluster to the point $z_{0}$ based on the following,

   $$\arg\min_{f\in\{1,2,\dots,p\}}\left\|z^{C_{f}}_{m}-{z}_{0}\right\|_{2}.$$

   (4)

   The initial cluster distribution $p_{0}\in\mathbb{R}^{p}$ is a one-hot vector, where the index corresponding to $f\in\{0,1,\ldots,p\}$ is set to 1. The cluster distribution at time $t$ is denoted by $p_{t}$. For successive transitions $\hat{x}_{1}\rightarrow\hat{x}_{2}\rightarrow\ldots\rightarrow\hat{x}_{T}$, it is essential to efficiently determine the cluster of each approximate state $\hat{x}_{i}$ generated by the surrogate model at time $t_{i}$. This step is performed for the estimation of high wavelength components of errors described in Section 3. The transitioning of cluster distribution ($p_{t}$) is modelled as a Markovian process discussed in [25]. We consider the nodes of the clusters from Step 2 as nodes of graph $\mathcal{G}_{I}=(\mathcal{V}_{I},\mathcal{E}_{I},q)$, $|\mathcal{V}_{I}|=p$, where the weights of the graph $q$ are unknown. The topology of the graph is considered complete with self-loops, meaning there is an edge between each node of the cluster and an edge from each node to itself. We use a supervised learning approach to determine the weights $q$ within the graph, where the labels $y_{t}$ from Step 2 are used in the optimization objective function with the cross-entropy loss (See P).

   The transitions of the cluster distribution is modeled as Markovian, based on the assumption that the states of the ODE exhibit Markovian properties, which is explained below. The general explicit Runge-Kutta numerical scheme using $n$ slopes ([26]) for obtaining solution of the differential equation $\dot{x}=f(x(t))$ is given by the following relation

   $$x_{t+1}=x_{t}+hb_{1}k_{1}+hb_{2}k_{2}+\ldots hb_{n}k_{n},$$

   $$k_{i}=hf(x_{i}+h\sum_{j=1}^{i-1}a_{ij}k_{j}).$$

   From this formulation, it follows that the solution transitions of the ODE adhere to a discrete-time Markov chain:

   $$p(x_{t+1}\mid x_{0},x_{1},\dots,x_{t})=p(x_{t+1}\mid x_{t}).$$

   In a graph $\mathcal{G}_{I}$, a walk is represented as a sequence of vertices $(v_{0},v_{1},\dots,v_{s})$, where each consecutive pair of vertices $(v_{i-1},v_{i})$ is connected by an edge in $\mathcal{G}_{I}$. In other words, there is an edge between $v_{i-1}$ and $v_{i}$ for all $1\leq i\leq s$. A random walk is characterized by the transition probabilities $P(u,v)$, which define the probability of moving from vertex $u$ to vertex $v$ in one step. Each vertex $u$ in the graph can transition to multiple neighboring vertices $v$. If we consider the transitions between the clusters as Markovian, the transition matrix $P=D^{-1}A$, where $D$ is the diagonal matrix of degrees and $A$ is the adjacency matrix of the graph. For each vertex $u$,

   $$\sum_{v}P(u,v)=1.$$

   $$P\left(u,v\right)=\;\left\{\begin{array}[]{ll}\frac{1}{d_{u}}&\;\text{if u and v are adjacent},\\ 0&\mathrm{otherwise}\end{array}\right.$$

   If we consider a complete graph on 3 nodes with self-loops as shown in Figure 2, the transition matrix,

   $$P=\left[\begin{array}[]{ccc}\frac{1}{q_{11}+q_{12}+q_{13}\;}&0&0\\ 0&\frac{1}{q_{22}+q_{23}+q_{12}}&0\\ 0&0&\frac{1}{q_{13}+q_{23}+q_{33}}\end{array}\right]\left[\begin{array}[]{ccc}q_{11}&q_{12}&q_{13}\\ q_{12}&q_{22}&q_{23}\\ q_{13}&q_{23}&q_{33}\end{array}\right].$$

   

   We now pose a constrained optimization problem P to estimate the weights ($q$) responsible for governing the transitions of data points.

   $$\left.\begin{aligned} &\textbf{minimize\;\;}_{q\,\in\,\mathbb{R}^{n(n+1)/2}}\;\;&&J=\sum_{t=1}^{T}\sum_{k=1}^{p}-({y}_{t})\log({p}_{t}(k))-(1-{y}_{t})\log(1-({p}_{t}(k)))\\ &\textbf{subject\;\; to \;\;}&&{p}_{t}=P(q)\;{p}_{t-1}\\ &&&q_{ij}\geq 0\;\;i,j=1,2,\ldots p.\\ \end{aligned}\right\}$$

   (P)

   ${y}_{t}$ represents the cluster index of data point $z_{t}$. The optimization problem above is solved to find the optimal $q^{*}$, which can be efficiently obtained using the adjoint method for data assimilation, as detailed in [23]. Given the cluster distribution ($p_{k}$) at time $t_{k}$, we estimate components $A_{k},R_{k}$ (Eq. 5) for the high wavelength component of errors $e_{mh}(t_{k}),e_{oh}(t_{k})$ as follows: $A_{k}=\text{diag(}{\sum_{i=1}^{p}p_{k}(i){e}_{m}^{C_{i}}}),R_{k}=\text{diag(}{\sum_{i=1}^{p}p_{k}(i)\hat{e}_{m}^{C_{i}}}).$
   

4. 4.

   State Estimation with Filtering. The goal of filtering is to estimate system states over time using noisy measurements and state observation relations, as discussed in Section 2.4. The POD surrogate model defines the evolution of state dynamics and the state-measurement relationship as follows:

   $$\begin{cases}x_{k+1}=M(Pf({x}_{k}))+A_{k+1}v_{k+1}+w_{k+1},&\\ z_{k}=\rho(x_{k}-\bar{x})+R_{k}\beta_{k}+\mu_{k},&\end{cases}$$

   (5)

   The terms $A_{k+1}v_{k+1}$ and $w_{k+1}$ represent the high- and low-wavelength components of model error at time $t_{k+1}$, respectively. Similarly, $R_{k}\beta_{k}$ and $\mu_{k}$ characterize the high- and low-wavelength components of the state-observation error at time $t_{k}$. $P=\rho^{T}\rho$ denotes the projection matrix. At a particular time step, the data point at time $t$ $({z}_{t})$ is obtained using the ROM $\dot{{z}}_{t}=f_{a}({z}_{t},t)=\rho f(\rho^{T}{z}_{t}+\bar{x},t)$, and the state vector after projection ${x}_{t}=\rho^{T}{z}_{t}+\bar{x}$. The matrices $A_{k+1}$ and $R_{k}$ are computed following the methodology outlined in Step 3 of the framework. The term $v_{1}$ is the solution to the linear system $A_{1}v_{1}=x_{1}-M(Pf(x_{0}))$. For dynamical systems represented on graphs, the term $x_{1}$ is computed using an explicit numerical scheme with the sparse graph obtained from the dynamic optimization Step 1. In the case of a diffusion equation represented on graphs $(\mathbb{D}$ in Section 2.2) with Euler discretization step size $h$, the sparse Laplacian matrix $L_{1}$ from the output of Step 1 is used instead of matrix $L$ to obtain $x_{1}=x_{0}-hL_{1}x_{0}$. $\beta_{1}$ is the solution to the linear system $R_{1}\beta_{1}=z_{1}-\rho(x_{1}-\bar{x})$. $w_{k+1}$ and $\mu_{k}$ are modelled as Gaussian with means $\mathbf{0}_{n\times 1},\mathbf{0}_{k\times 1}$ and variances $\zeta_{x}=\sigma_{x}I_{n\times n},\zeta_{y}=\sigma_{y}I_{k\times k}.$ The distribution governing the low wavelength components of errors ($w_{k+1},\mu_{k}$) is deliberately shaped to exhibit minimal variability, with a centered mean of $\mathbf{0}_{n\times 1}$. This assumption is made with the expectation that the high wavelength component will accurately capture the true nature of the error. There exist linear and non-linear filters for state estimation as described in [23], [24]. When the uncertainty in the state prediction exceeds a threshold, vectors $v_{k+1}$ and $\beta_{k}$ are updated as solutions of the linear systems defined below:

   	$\displaystyle A_{k+1}\;v_{k+1}=x_{k+1}-M(Pf(x_{k})),$		(6)
   	$\displaystyle R_{k}\;\beta_{k}=z_{k}-\rho(x_{k}-\bar{x}).$		(7)

   The computation of $x_{k+1}$ is performed using an explicit numerical scheme that leverages the state vector obtained from the filter at time step $k$ using the sparse graph.

This section outlines the constraints necessary for the dynamic optimization step (Step 1 in Section 4) applied to a linear dynamical system $\mathbb{D}$ represented on graphs (Section 2.2). Section 5.2 details the dynamic optimization process for linear diffusion systems on graphs. Section 5.3 presents the optimization framework for reaction-diffusion systems on graphs, including the corresponding constraint formulations.
Existing techniques for approximating linear dynamical systems on graphs include [27], [28], and [29]. The approach in [29] constructs spectral sparsifiers by sampling edges according to their resistance $R_{e}$. The probability of sampling an edge $e$ to construct a sparse approximation $\bar{G}$ of graph $G$ is given by

where the edge resistance is defined as

See Algorithm 5 for further details.

Computing bounds on the edge weights of the sparse graph approximation $\bar{G}$, using effective resistances or local edge connectivities, is computationally intensive. To overcome this challenge, we propose Algorithm 5.1, which leverages Theorems 2 and 3 to efficiently compute upper bounds on edge weights and node degrees in the sparse graph. The following discussion introduces Bernstein’s Inequality [30], which is utilized in the derivation of Theorem 2.

Theorem 2 provides upper bounds on the edge weights of the sparse graph $\bar{G}$ generated by Algorithm 5. Furthermore, it is necessary to establish bounds on the node degrees in the sparse graph $\bar{G}$. To derive these bounds, we introduce Theorem 3, which utilizes the established relationship between the eigenvalues of the Laplacian matrix and the node degrees, as discussed in Section 5.1.
Notation: Let $G$ be a simple weighted graph, i.e., a graph without self-loops or multiple edges with vertices $\{v_{1},v_{2},\ldots,v_{n}\}$. We define two subsets of vertices in $G$ based on their degrees. The set $l_{m}$ contains the $n-m+1$ vertices with the highest degrees, while $s_{m}$ consists of the $m$ vertices with the lowest degrees. The subgraphs $G_{l_{m}}$ and $G_{s_{m}}$ are defined as the subgraphs of $G$ induced by the vertex sets $l_{m}$ and $s_{m}$, respectively.

Here $d_{v_{1}}=d_{1}\leq d_{v_{2}}=d_{2}\leq\ldots\leq d_{v_{n}}=d_{n}.$