Effective Latent Differential Equation Models via Attention and Multiple Shooting
Authors: Germán Abrevaya
Mahta Ramezanian-Panahi
Jean-Christophe Gagnon-Audet
Pablo Polosecki
Irina Rish
Silvina Ponce Dawson
Guillermo Cecchi
Guillaume Dumas
Venue: Accepted by TMLR ()

Abstract: Scientific Machine Learning (SciML) is a burgeoning field that synergistically combines domain-aware and interpretable models with agnostic machine learning techniques. In this work, we introduce GOKU-UI, an evolution of the SciML generative model GOKU-nets. GOKU-UI not only broadens the original model's spectrum to incorporate other classes of differential equations, such as Stochastic Differential Equations (SDEs), but also integrates attention mechanisms and a novel multiple shooting training strategy in the latent space. These modifications have led to a significant increase in its performance in both reconstruction and forecast tasks, as demonstrated by our evaluation on simulated and empirical data. Specifically, GOKU-UI outperformed all baseline models on synthetic datasets even with a training set 16-fold smaller, underscoring its remarkable data efficiency. Furthermore, when applied to empirical human brain data, while incorporating stochastic Stuart-Landau oscillators into its dynamical core, our proposed enhancements markedly increased the model's effectiveness in capturing complex brain dynamics. GOKU-UI demonstrated a reconstruction error five times lower than other baselines, and the multiple shooting method reduced the GOKU-nets prediction error for future brain activity up to 15 seconds ahead. By training GOKU-UI on resting state fMRI data, we encoded whole-brain dynamics into a latent representation, learning a low-dimensional dynamical system model that could offer insights into brain functionality and open avenues for practical applications such as the classification of mental states or psychiatric conditions. Ultimately, our research provides further impetus for the field of Scientific Machine Learning, showcasing the potential for advancements when established scientific insights are interwoven with modern machine learning.

Published in Transactions on Machine Learning Research (02/2024)
Effective Latent Differential Equation Models via Attention and Multiple Shooting
Germán Abrevaya!
gabrevayo uba..r
Mahta Ramezanian-Pannhi.
mahla. mahla.ramecaaiann &maia. quebec
Jean-Christophe Gagnon-Audet
quebec
Pablo Polosecki
pipolose@uus. wwm.iiimoooooooo com
Irina Rish 3
irina. rmmila quebec
Silvina Ponce Dawson
silvina uba.ar
Guillermo Cecchi
gcecchi@@s..bmn com
Guillaume Dumas
Universidad de Buenos Aires, FCRNN, Departamento de Física. Buenos Aires, Argentina.
Mila Queber AT Instttute. Montréal, Québec, Canada.
Université de Montréal. Montréol, Québec, Canada.
IBM Research, T.J. Wotson Research Center. Yorktown Heights, New York, USA.
*CONICET- Universidad de Buenos Aires, Instituto de Física de Buenos Aires (IFIBA). Buenos Aires, Argentina. "CHU Suinle-Justine Research Center. Montréal, Québec, Canada.
Department of Psychiatry and Addictology, Universilé de Montréal. Montvéal, Québee, Canada.
Reviewed on OpenReview: https: // openreview. net/ forun? id= uNNfNaPU1
Abstract
Scientific Machine Learning (SciML) is burgeoning field that synergistically combines domain-aware and interpretable models with agnostic machine learning techniques. In this work, we introduce GOKU-UI, an evolution of the SciML generative model GOKU-nets. GOKU-UI not only broadens the original model's spectrum to incorporate other classes of differential equations, such as Stochastic Differential Equations (SDEs), but also integrates attention mechanisms and novel multiple shooting training strategy in the latent space. These modifications have led to significant increase in its performance in both reconstruc- lion and forecast. tasks, as demonstrated by our evaluation on simulated and empirical data. Specifically, GOKU-UI outperformed all baseline models on synthetic datasets even with training set 16-fold smaller, underscoring its remarkable data efficiency. Furthermore, when applied to empirical human brain data, while incorporating stochastic Stuart-Landan oscil- lators into its dynamical core, our proposed enhancements markedly increased the model's effectiveness in capturing complex brain dynamics. GOKU-UI demonstrated reconstruc- tion error five times lower than other baselines, and the multiple shooting method reduced the GOKU-nets prediction error for future brain activity up to 15 seconds ahead. By train- ing GOKU-UT on resting state fMRI data, we encoded whole-brain dynamics into latent atpresentation, learning a low-dimensional dynamical system model that could offer insights into brain functionality and open avenues for practical applications such as the classification of mental states or psychiatric conditions. Ultimatoly, our rescarch provides further impetus for the field of Scientific Machine Learning, showcasing the potential for advancements when established scientific insights are interwoven with modern machine learning.

---

Published in Transactions on Machine Learning Research (02/2024)
1 Introduction
1.1 Scientific Machine Learning
Scientific Machine Learning (SciML) is an emerging field that, drawing insights from scientific data, seeks to advance data-driven discovery with an approach that can produce interpretable results (Baker et al., 2019). Its synergistie blend of machine learning and scientific computing based on mechanistic models makes it very powerful for addressing complex problems across all STEM areas and beyond (Willard et al., 2022). Using informed priors, its application is already making key contributions in scientific inference, data analysis, and machine learning enhanced modeling (von Rueden et al., 2023). Recent developments in SciML include various approaches to derive dynamical system models from observational data. The sparse identification of nonlinear dynamics (SINDy) algorithm is one such approach which, leveraging recent advances in sparsity techniques, exploits the observation that only a few important terms dominate the dynamics in most physical systems (Brunton et al., 2016). A different well-established candidate for handling model-agnostic systems is deep reinforcement learning, which has been particularly widely adopted by the fluid mechanics community, with applications in flow control and shape optimization 2022). Physics-informed neural networks (PINNs) are another approach in which neural nelworks are trained to solve supervised learning tasks with respect to the laws of physics and can be used to derive data-driven partial differential equations or their solutions (Raissi et al., 2019). The application of PINNs to numerically stiff systems (Wang et al., 2021) is challenging. Universal differential equations (UDEs) are a recent met hod that not only overcomes the stiffness limitation, but also represents a perfect example of the essence of SciML: use all the prior knowledge and scientific insights available for your problem and fill the missing gaps with machine learning (Rackauckas et al., 2020). The simple yet powerful idea behind UDEs involves using traditional differential equation models with some unknown terms that are substituted by universal function approximators, such neural networks. These approximators will be learned simultaneously with the equations parameters by using sensitivity algorithms (Ma et al., 2021a) within the differentiable modelling framework (Shen et al., 2023).
Geurrerr & Lamata, 2021; Viquerat et al.,
The evolution equations that might be derived from the data may not actually correspond to mechanistic models based on first principles. That is, on many occasions, the dynamics take place on a manifold of a lower dimension than the full phase space of the system (Foias et al., 1988). Having reduced equations that describe the evolution on these manifolds is very useful, especially in high-dimensional systems (Lucia et al., 2004). Reduced-order models (ROMs) can be derived from data. (Guo & Hesthaven, 2019). In particular, generative adversarial network (GAN) (Goodfellow et al., 2014) approaches have been used to enhance the application of ROMs to simulations of fluid dynamics (Kim et al., 2019; 2022).
1.2 Neural Differential Equations
Despite the basic ideas behind differential equations parameterized by neural networks and its connection with deep Icarning had older roots in litcrature (Rico-Martinez et al., 1992; 1994; Rico-Martinez & Kevrokidis, 1993; Chang et al., 2018; Weinan, 2017; Haber & Ruthotto, 2017; Lu et al., 2018), the publication of Chen et al. (2018) was a turning point in the young history of SeiMT. Since then, the topic of neural differential equations (neural DEs) has become a field, as stated and evidenced in the comprehensive survey by Kidger (2022). In Chen et al. (2018), by interpreting ResNets (He et al., 2016) as a discrete integration of a vector field with the Fuler method, the authors proposed an infinitesimally layered neural network as its continuous limit and modcled it with an ordinary differential equation (ODE) parameterized by a neural network, giving rise to the Neural Ordinary Differential Equation (NODE) models. They also demonstrated that NODEs can be trained by backpropagating through black-box ODE solvers using the adjoint method (Pearlmutter, 1995), making it a memory-effciiet model.
Furthermore, Chen et al. (2018) introduced the Latent Ordinary Differential Equations (Latent ODEs), a continuous-tim generative model that encodes time series data into a latent space that could potentially capture its undorlying dynamics, which are modeled using a NODE. First, the observed time series are encoded in the latent space using recognition model, typically a Recurrent Neural Network (RNN). The temporal dynamics in the latent space is then modeled using NODEs, and lastly, its solution is decoded
2

---

Published in Transactions on Machine Learning Research (02/2024)
back into the observation space to generate predictions perform other tasks such as anomaly detection imputation of missing values.
By using this approach, Latent ODEs can capture the intricate and potentially nonlinear dynamical systems that underlic many real-world time series data. The continnous nature of NODEs allows the model to handle irregularly sampled data (Rubanova. al., 2019), a common challenge many applications. Additionally, the use of the latent space enables the model to capture complex patterns in high-dimensional data, while still providing compact and interpretable representation of the underlying dynamics.
1.3 GOKU-nets
The work of Linial et al. (2021) builds on the basis of the Latent ODEs model demonstrating that the incorporation of prior knowledge of the dynamics involved the form of a backbone differential equation structure can increase the performance of purely agnostic model. They propose another continuous-time generative model called GOKU-nets (which stands for Generative ODE Modeling with Known Unknowns), which are the focus of this paper. This model incorporates variational autoencoder (VAB) structure with differential equation to model the dynamics in the latent space. However, in this case, a specific form for the ODE is provided while allowing its parameters (the known unknowns) to be inferred. The model is trained end-to-end jointly learning the transformation to the latent space, inferring the initial conditions and parameters of the ODE, which will be integrated afterward; and finally, a last transformation performed to go back to the input space. In the next section, the details of its architecture will be described in detail. Note that the ODE can be integrated further than the input time span in order to generate an extrapoaaiion thus becoming forecast of the future evolution of the time series. The study in Linial et al. (2021) compares GOKU-net with basclines such as LSTM and Latent-ODe in three domains: a video of a pendulum, a vidco of a double pendulum, and a dynamic model of the cardiovascular system. The authors show that their model significantly outperforms the others in reconstruction and extrapolation capabilities, reduces the size of the required training sets for effective learning, and furthermore, has greater interpretability, allowing for the identification of unobserved but clinically meaningful parameters.
The original GOKU-net model was limited to handling only ODEs. In this work, we expand its capabilities by implementing the model in the Julia Programming Language (Bezanson et al., 2017), leveraging its potent SciML Ecosystem (Rackauckas & Nie, 2017) which enables us to utilize a wide spectrum of differential equation classes (including SDEs, DDEs, DAEs) and a diverse suite of advanced solvers and sensitivity algorithms (Rackauckas et al., 2019; Ma et al., 2021a).
As reported in the literature, the identification of nonlinear dynamical systems, and in particular the gradient- descend based training of neural DE models, can be often challenging due to their highly complex loss landscapes, leading to poor local minima training stagnation (Ribeiro et al., 2020; Turan & Jäschke, 2021). We propose an enhancement to the original GOKU-net architecture which adds attention mechanisms to the main part of the model that infers the parameters of the differential equations. Moreover, to overcome the inherent difficulties of training, we developed a novel strategy to train the GOKU-net based on the multiple shooting technique (Bock & Plitt, 1984; Ribeiro et al., 2020; Turan & Jäschke, 2021) in the latent space. We have evaluated our enhanced model and training strategy on simulated data from a network of stochastic oscillators, specifically Stuart-Landau oscillators, as well as empirical brain data derived from resting state human functional Magnetic Resonance Imaging (fMRI). In both cases, the GOKU-net that fuses multiple shooting and attention, labeled GOKU-nets with Ubiquitous Inference (GOKU-UI), outperformed both the base GOKU-net model and the bascline models in terms of reconstruction accuracy, forccasting capability, and data efficiency. The application of GOKU-UI in IMRI data analysis not only demonstrate its superiority in handling complex neural datasets but also highlights its potential in enhancing real-time neuromodulation strategics (Parastarfcizabadi & Kouzani, 2017). We believe that GOKU-UI represents a promising step forward in Scientific Machine Learning, underscoring the rich possibilities that emerge when melding traditional scientific insights with contemporary machine learning techniques.
3

---

Published in Transactions on Machine Learning Research (02/2024)
Initial valuess
Pattern Extractor
Z
Differential Eq. Solver
Latent out
Latent
Parameters
Feature Extractor
Reconstructor
t
Figure 1: Schematic represenlation of general Lalent Differential Equation model. First, each time frame of the input data x2 is independently processed by a Feature Extructor, then the whole sequence goes through pattern Extractor, which learns the mean and variance of the initial conditions, and possibly also the parameters, of the differential equation that will be next integrated. Finally, the solution goes through last transformation performed by Reconstructor in order to go back to the input space.
2 Methods
2.1 Basic GOKU-nets
GOKU-nets be considered a specific instance within the broader category of models denominate Latent Differential Equation models (Latent DEs). These generative models have structure similar to Variational Autoencoders (VAEs), but with key distinction: they encode time series data into a latent space governed by differential equations. The objective, akin to any autoencoder model, is to ensure that the output closely reconstructs the input by passing through this latent space. In Latent DEs, the input typically comprises complex, high-dimensional time series data, such as sequential brain images, financial records, or biosignals. After encoding the data into its latent space, beyond simply generating an output that corresponds to the time span of the input, these models integrate the differential equation further in time, allowing them to forecast the system's future evolution. This capability is not only crucial for validating model, as is slandard in scientific modeling, but also has diverse practical applications. These applications range from finance and weather forecasting to neuromodulation in the context of brain data, where accurate short-term forocasts can help compensate for delays in closed-loop foodback systems (Parastarfcizabadi & Kouzani, 2017).
The schema in Figure illustrates general Latent DE model. Initially, each temporal frame of the input data 2; is independently processed by Feature Extractor, usually reducing its dimensionality. Following this, the entire sequence is subjected to Pattern Extractor, which aims to learn the distribution of the initial conditions and possibly of the parameters for the differential equation that will be subsequently integrated. Lastly, the solution undergoes a final transformation via Reconstructor, going back to the original input space. The original model is trained as a standard VAE, maximizing the evidence lower bound (ELBO) (Kingma Welling, 2013).
In the case of the Latent ODEs proposed by Chen al. (2018), an RNN is used for the Pattern Extractor, fully connected NN for the Reconstructor, and the differential equation is parametrized with another NN. This means that in this model the specific form of the differential equation is not provided beforehand but is learned simultancously during training. The Feature Extractor and the intermediate layers, Latent in and Lalent oul, are not present so they could be considered as identity operations. On the other hand, the GOKU-net model proposed by Linial et al. (2021) has a ResNet with fully connected NNs as the Feature Extractor, while in the Pattern Extractor, an RNN used learn the initial conditions and a bidirectional LSTM for the ODE parameters. In this case the differential equation is explicitly predefined, allowing incorporate some prior knowledge about the dynamical nature of the system under consideration. The Latent.
4

---

Published in Transactions on Machine Learning Research (02/2024)
in and out layers are fully connected NN and finally, the Reconstructor is another ResNet, similar to the initial one. In the next section, some enhancements to the original GOKU-net model are proposed.
2.2
GOKU-UI
2.2.1
Attention mechanism
The first modification is the addition of a basic attention mechanism (Vaswani et al., 2017) to the Pattern Extractor, specifically in the part associated with the learning of the parameters of the differential equation. Namely, instead of keeping the last element of the bidirectional LSTM (BiLSTM) used in the original GOKU- net model, all of their sequential outputs pass through a dense layer with softmax activation to calculate the atlentional scores that would weight the sum of all the BiLSTM outputs in order to obtain its final output.
2.2.2 Multiple Shooting
When training Neural DE models, gradients have to be calculated through differential equations with respect to its initial conditions and paramcters, by mcans of some scnsitivity algorithm (Ma et al., 2021a). This tends to produce highly complex loss landscapes (Ribeiro et al., 2020; Metz al., 2021). The work of Turan & Jäschke (2021) demonstrates that training Neural ODEs even on very simple oscillatory data could be problematic, showing that the outcome may result in trajectory similar to a moving average of the original data, thus failing to capture responses of higher frequency. They proposed a solution based on the multiple shooling methods, which are widely used in Optimal Control (Bock & Plitt, 1984; Diehl et al., 2006) and Systems Identification (Baake et al., 1992; Ribeiro et al., 2020) to alleviate the problem of high sensitivity to initial conditions and lower the probability of getting trapped at local minima with very poor performance. The basic idea of multiple shooting is to partition the complete time span over which the differential equation would be integrated into smaller time windows, for each of which the initial conditions are inferred in parallel. Afterwards, the multiple segments are joined and continuity constraint is imposed during optimization, in our case, through the penalty method, which in practice simply consists of adding a regularization term to the loss function.
However, applying the multiple shooting method to GOKU-nets is not straightforward. Firstly, in most cases that use this method, such as in Turan & Jäschke (2021), the differential equations are typically directly modeling the observable data, having direct access to the true initial conditions for each window. In the case of GOKU-nets, the dynamics modeled by differential equations occur in the latent space, which is being learned sinultaneously; a result, such true initial conditions are not available. Secondly, it is necessary to determine how the method will behave in relation to the parameters of the differential equation, which in the case of Neural ODEs are implicitly learned as part of their parameterization through the neural network.
Our proposal for exlending the multiple shooting method to GOKU-nets is as follows. After passing through the Feature Extractor, we divide the temporal interval in the latent space in such a way that the Pattern Extractor generates in parallel different initial conditions for each temporal window, but provides single set of parameters for the differential equations that will be shared by all windows. By this strategy, we maintain the potential benefits inherent to the multiple shooting method while leveraging the information available in wider temporal range for the task of parameter inference, which is generally more challenging than estimating initial conditions. As mentioned before, we do not have access to the true initial conditions for grounding the latent trajectories. However, we can strive to achieve continuity across different windows, which is crucial in multiple shooting methods to obtain solutions equivalent to those achieved with single shooting. To this end, these intervals are defined by overlapping the last temporal point of each window with the first one of the following and the goal is to minimize the distance between these points. Specifically, we employ regularization in the cost function when training the model, quadratically penalizing the discrepancy in the latent space of the overlapping points, that is, between the initial condition of each window and the end point of its preceding segment. (Vantilborgh et al., 2022; Turan & Jäschke, 2021).
Our experiments indicated that non-variational GOKU-nets models outperform their variational counterparts significantly (see Supplementary Information C). Therefore, we used aan-vaiiatioall GOKU-nets for all the remaining results in this work. Specifically, instead of sampling from normal distributions in the latent space
5

---

Published in Transactions on Machine Learning Research (02/2024)
shown in Figure 1, we directly take the mean values Pao and Mlg. As result, the cost function associated with these models does not include the KT, divergence term associated with the ELBO, but it does retain the reconstruction term, which is calculated as the mean squared error between the model's output and the input, normalized by the mean absolute value of the imput. Furthermore, when using multiple shooting training, the continuity regularization described in the previous paragraph is included.
2.3 Experiments
In the next sections, we evaluate our proposed attention and multiple shooting enhancements through two highly challenging cases; onc on synthetic data basod on a network of stochastic oscillators known as Stuart- Landau oscillators and the other on empirical human brain data.
We compare the reconstruction and forecast performance of different variations of the GOKU-model (basic with attention) trained in the original single shooting fashion or with the proposed multiple shooting method, as well as some baseline models: LSTM, Latent ODE (Chen el al., 2018), Augmented Neural ODE (ANODE) (Dupont et al., 2019) and a naive model. In the forecast task evaluation, a vector autoregressive (VAR) model (Lütkepohl, 2005) was also included as a baseline, as it is a standard statistical model used for forecasting multivariate time series. For a fair comparison, both the LSTM and Latent ODE models are constructed maintaining the same GOKU-net general architecture and changing only the differential equation layer. Specifically, the Feature Extractor, Pattern Extractor, Latent In, Latent Out, and Reconstructor layers (sec Figure 1) maintain the same architecture and hyperparameters. However, the differential equation layer is substituted with a Neural ODE for the Latent ODE model, and in the other case, it is replaced by recursively executed LSTM. In contrast, the ANODE and VAR directly model the input variables without inferring any latent space representation. The size of the NN parameterizing the differential equation inside the NODE and Latent ODE, as well as its latent state dimensionality and size of the LSTM were selected to match the total number of parameters of their contending GOKU-UT (with attention and multiple shooting). The lag order used in VAR was selected in cach case as the one that minimizes the Akaike's Information Criterion) (AIC) on the training data. The naïve predictors, both for the reconstruction and forecast. task, are historical averages of the input data, i.e. constant predictions with the mean values of the inputs across time.
All models were trained under identical conditions and following the same procedure, with the aim of minimizing reconstruction error. In all instances, the input sequences to the model consisted of 46 time steps. During each training epoch, a random interval of this length was selected within the training data available for each sample in every batch of 64 samples. In the case of the multi-shooting training, the 46-time step sequence was further partitioned in the latent space into five windows, each comprising 10 time steps, with the last point of one window overlapping with the first point of the subsequent window. Under these circumstances, the loss function was augmented by the sum of squared differences between the endpoints of each segment. This sum was normalized by the number of junctions and multiplied by a regularization coofficient to imposc a continuity constraint among the different segments. In the results presented here, the regularization coefficient was set to 2. Comprehensive details of the training process, as well as the specific architecture of the models and the hyperparameters used, can be found in the Supplementary Information.
2.3.1 Simulated data
Stuart-Landau (SL) oscillators, representing the normal form of a supercritical Hopf bifurcation, serve as a fundamental archetype of mathematical models and are extensively used across diverse scientific disciplines to study self-sustained oscillations (Kuramoto, 1984; Kuznetsov et al., 1998). The SL oscillator is often described by the complex-valued valued differential equation
=(a + iw) - =|=|2
(1)
where: 2 = "e = x + iy, a is the bifurcation parameter, and w is intrinsic frequency of the oscillator. The parameter represents the linear growth rate of the system. When a is positive, the amplitude of the oscillation increases, and when is negative, the amplitude of the oscillation decreases. At a 0, a Hopf bifurcation occurs, and the system transitions from stable fixed point to limit, cycle oscillations (or vice
6

---

Published in Transactions on Machine Learning Research (02/2024)
versa). Despite its apparent simplicity, the SL model can exhibit wide range of behaviors, including limit cycles, chaotic oscillations, and quasiperiodicity, making it versatile tool in the study of nonlinear dynamics and a good candidate for evaluating the capabilities of the GOKU-net models.
In particular, we generate the simulated data with a network of coupled stochastic Stuart-Landan oscillators that has been widely used to model brain dynamics for resting state IMRT (Jobst. al., 2017; Deco al., 2017; Donnelly-Kehoe et al., 2019; Ipiña al., 2020), which will also be used in the empirical evaluation on brain data, described in the next section. The dynamics of the i-th node within network of N oscillators is given by the following equation:
8, = Re(2)) [a,- - 10) wAys + +G Σ
(z mz) (ty
N
ú = Im(2;)= =
= (aj - =3 ve - wjzs + G Σ
Cj(lv - 83) + Bny(t)
(2)
where Cij is a connectivity matrix between all nodes in the network, G is a global coupling factor, while nj represents additive Gaussian noise. Note that independent bifurcation parameters ay and frequencies w are used for each node.
During the construction of our dataset, we perform dimensionality augmentation on the network of oscil- lators, which are utilized as latent dynamics. Specifically, apply fixed random linear transformation, f RR2 - D to the latent trajectories of each sample, where the dimension D is much larger than 2N. All the synthetic data experiments were performed using N = 3 stochastic oscillators with a high dimension D = 784. Experiments with various sizes of training sets were performed, with training sets ranging from 75 to 4800 different samples. Each sample corresponds to a different set of initial conditions and parameters for the SL diffcrential equations. A validation set of another 200 samples was used for the training termination criteria, and separate testing set containing 900 different samples was employed for the evaluation. All details of the implementation and hyperparameters be found in the Supplementary Information.
2.3.2 Empirical evaluation
Next, in an effort to evaluate both our proposed models in a challenging empirical dataset and to provide an example of the advantages of incorporating prior scientific insights into agnostic Al models, we focus on one of the most complex systems in nature: the human brain.
We used the resting state MMR data from 153 subjects, sourced from the Track-On HD study (Klöppel et al., 2015). The data was pre-processed as described in Polosecki et al. (2020), followed by a 20-component Canon- ical ICA (Varoquaux et al., 2010). Of the original 20 components, 9 were identified as artifacts and therefore eliminated, leaving 11 components for further analysis. The contribution of each subject was captured in two separate visits, yielding an aggregate of 306 data samples. These samples, with 160 time points each, were acquired at a temporal resolution of 3 seconds. In contrast to the simulated data scenario, the empirical data was considerably limited. Consequently, we optimized our data splitting approach. To avoid 'double- dipping', we ensured that the training and testing datasets comprised entirely distinct samples. However, for the training and validation sets, we utilized the same samples, implementing a temporal split. Specifically, approximately 20% of the data samples (n=60) were reserved for testing, while ensuring a balanced represen- tation of the sex, condition, and measurement site. The rest of the data samples (n=246) were allocated for training and validation purposes. Within this subset, the initial 114 time points were designated for model training. The remaining data served for validation and for an criterion early training termination. For a visual representation of the dataset partitioning, please refer to Figure 5 in the Supplementary Information.
Our proposed GOKU-UI model was trained using 20 Stuart-Landau stochastic oscillators (Eq. 2) in the latent space. A comparison of the model's performance using various numbers of oscillators and thorough description of the training procedure are presented in the Supplementary Information section.
7

---

Published in Transactions on Machine Learning Research (02/2024)
b
0.8
ftwc with Sirple
Shoet ing Shooklngg
Reomstructinn NRMSE
GORU Actention with Single
LSTH
flnetc
with SaiiSiple shooting
Shouting
(GOKU-UE)
Forecaes NRMSE
0.8
GERU Bexsc
LSTH
Attencion
DDE
with Na3lspla
Shooting (00R)-00)
0.4
Netve
AMBLE
VAI
022
70
130
900
1230
2400 4800
100
3001
500
1200
2400
4903
Training samples
Training samples
Figure 2: Comparison the reconstruction (a) and forecast (b) performances on the synthetic Stuart- datasct mcasured by the median normalized RMSE while incrcasing the number of samples in the training set. The averages were taken with respect the input dimensions and time span, while the shaded regions correspond standard errors derived from various training runs with multiple random seeds. The forecast. was evaluated on 20 time-steps horizon.
3 Results
We evaluated the performance of four GOKU-net variants across both reconstruction and forecast tasks. These variants encompass single and multiple shooting methods, either with or without the use of the attention mechanism. For the reconstruction task, we compared them against four baseline models: LSTM, Latent ODEs, ANODE, and a naïve predictor. In addition to these, we also included VAR in the set of baseline models for the forecast task.
All models were exclusively trained on the reconstruction task and their forecasts were generated during the evaluation stage. The prediction between the target ground truth and either the reconstruction or the forecast.
quantified using the normalized root mean square error (NRMSE)
3.1 Simulated data
Figure 2a shows the performance of the different models in the reconstruction task on a synthetic datasct generated with the latent dynamics of three stochastic Stuart-Landau oscillators. Tn this context, the GOKU- net variants employing the multiple shooting method demonstrated significantly lower errors compared other models. Notably, GOKU-nets that utilized the attention mechanism consistently showed better out- comes relative to their corresponding basic models. The single shooting variant particularly benefited from the attention mechanism. In the multiple shooting scenario, the performance boast due to attention pronounced when using fewer training samples. However, when the number of training samples exceeded 600, the performance of both models converged to much closer values. However, GOKU-UI, which incorpo- rates both the attention mechanism and multiple shooting, was still significantly the best-performing model. Wilcoxon signed rank tests, comparing GOKU-UT with each of the other models, for each number of training samples independently, yielded all p-values 0.02 after Holm correction.
In contrast, the Latent ODEs, whose training instances consistently became trapped in local minima sociated with the mean of the time scries, markedly underperformed and exhibited errors akin to those of the naïve predictor. A similarly poor, though slightly better performance was obtained by ANODE. On the other hand, the LSTMs outperformed the basic GOKU-nets trained with the single shooting method. GOKU-UI, when trained on more 300 samples, yielded a reconstruction performance comparable that achieved with 4800 samples. Moreover, when trained on just 150 unique samples, GOKU-UI outperformed all other single shooting models, those trained on dataset. 32 times larger.
8

---

Published in Transactions on Machine Learning Research (02/2024)
b
3.6
12
Rosmstrucion KRMSE
Forecaes NRMSE
0.9
I 1241 Atbention
Basic
with Single
Steecing Shocting
0.3
024U AtSention with Multiple Shooling (G0RU-U1)
LSTN
ANTEL
0.45
Mazve
0.15
GUIKU Bhe
BOOI ANL
DOKE BM
DROSU Art
15TM Later: ORE
ANODE
Nitw
5
30
15
20
25
30
dustee
Time (6)
Figure 3: Comparison the reconstruction (a) and forecast (b) performances on testing fMRI data, measured by the normalized RMSE. The averages were taken with respect the input dimensions and time span, while the shaded regions correspond to standard errors derived from various training runs with multiple random seeds.
When forecasting 20 temporal points beyond the reconstruction limit, observed similar trend, as depicted in Figure 2b. However, in this lask, GOKU-nets employing single shooting with the atlention mechanism outperformed LSTMs. Except at the lowest end of the training data size, where VAR recorded the lowest forecasting error, GOKU-nets using multiple shooting consistently demonstrated superior overall perfor- mance. is expected that simpler statistical methods, such as VAR, can sometimes surpass more complex machine learning models in scenarios with scarce training data. Nevertheless, GOKU-UI significantly outper- formed other models when using 150 and 300 training samples (all p-values < 0.05, according Wilcoxon signed-rank tests conducted independently for each sample size, after Holm correction). Beyond this, its performance was statistically indistinguishabl from that of the basic GOKU model with multiple shooling (p-values > 0.05, Wilcoxon signed-rank tests, Holm corrected).
3.2 Empirical data
In a separate analysis, we trained GOKU-net models on human brain recordings, employing 20 coupled stochastic Stuart-Landan oscillators to govern the dynamics in their latent space. These models were trained alongside baseline models on 246 samples of 11-component ICA time series obtained from [MRT data. As shown in Figure 3a, unlike the simulated data scenario, the addition of the attention mechanism did not improve the reconstruction performance of the single shooting GOKU-net models. However, the multiple shooting training method again resulted in significant improvement factor (p 0.001, Wilcoxon signed- rank test.. Remarkably, the GOKU-UT model, which integrates both multiple shooting and the attention mechanism, achieved a median reconstruction етrоr that was five times lower than that of single shooting baseline models. Specifically, GOKU-UI also exhibited a significantly lower reconstruction NRMSE compared the basic GOKU model that employed multiple shooting (p < 0.04, Wilcoxon signed-rank test.. For a qualitative performance comparison between GOKU-UI and the original GOKU-net, refer to representative reconstruction plots in the Supplementary Information E.
Models trained exclusively for reconstruction were also assessed in forecast task, which consisted of pre- dicting the time scries evolution immediately beyond the reconstruction limit. Figure 3b illustrates the performance Irends of the different models as the forecast. horizon extends. Notably, VAR exhibited the best short-term prediction accuracy, but its effectiveness diminished rapidly over time. This trend, in line with findings from Abrevaya et al. (2021), ; can be attributed to the relative slowness of the fMRI signal. The hemo- dynamic response function (HRF) in fMRI, which convolves underlying neural activity, typically operates over timescale of about 10 seconds. Given the time resolution of (MRI is around 3 seconds, this temporal

---

Published in Transactions on Machine Learning Research (02/2024)
b
d
0000
3.00
Figure 4: Comparison of the averaged errors in whole brain reconstructions from GOKU Basic with single shooting (a), GOKU Attention with single shooting (b), GOKU Basic with multiple shooting (c) and GOKU Attention with multiple shooting (d). The color scale represents the RMSE normalized by the maximum voxel value in the original brain data. The averages were calculated across all testing samples, time and random seeds.
setup implies that genuine signal fluctuations are unlikely to manifest within less than three to four time points. Furthermore, IMRI noise is characterized by strong autocorrelation, as noted in Arbabshirani et al. (2014). Additionally during preprocessing, data undergo bandpass filtering within the range of 0.01-0.16 16 Hz, process designed to enhance the neural signal-to tolltto-oooe noise ratio while aligning with the HRF's natural filtering properties, but this also inherently limits the observation of rapid changes in the fMRI signal.
On the other hand, examining the performance of the GOKU-net models reveals that the multiple shooting method was again a key factor for improvement. Initially, GOKU-UI models tended to start their forecasts with larger error margin than VAR but exhibited slower degradation, maintaining better performance than other models for up to 15 seconds. Beyond this point, GOKU-UT was oulperformed by other baseline models, though it continued to exhibit marginally better performance compared to the basic GOKU-nets with single shooting, Nonetheless, all models demonstrated limited efficacy in long-range forecasts, a foreseeable outcome considering the complexity of the system under study. Given their exclusive training on reconstruction tasks, these models present opportunities for further refinement and enhancement.
In addition evaluating performance on temporal ICA components, our study extends to reconstructing the full brain by linearly combining the components of ICA spatial support, each weighted by their corre- sponding time courses. Figure 4 demonstrates this approach by comparing whole brain reconstruction errors across the four GOKU-net variants. Each variant's performance is quantified using the RMSE, normalized by the maximum voxel value, and is visualized through a glass brain representation from three distinct inct. perspectives. This comparison distinctly highlights the beneficial impact of the multiple shooting training approach in enhancing model performance. The visualization effectively conveys the quantitative differences in reconstruction errors and also offers an intuitive understanding of the spatial distribution of these errors across the brain. For additional insights, including visual representation of the TCA spatial support and analysis of whole brain forecasts, readers are directed to the Supplementary Information F.
4 Discussion and Limitations
Our research focused on two key enhancements to the baseline GOKU-net model: the addition of a basic attention mechanism and the implementation of a multiple shooting method. Independently, each these modifications tended to improve the model's performance in both reconstruction and forecast. tasks. The use of multiple shooting yielded the most substantial improvement. In particular, GOKU-UI, a composite of both enhancements, exhibited the best overall performance. During the evaluation the synthetic dataset, GOKU-UI demonstrated remarkable efficiency with respect to the training data. With mere 150 training samples, it outperformed all other single shooting bascline models, even when those had employed training sets 32 times larger (Figure 2). On average, for the synthetic dataset, GOKU-UT achieved a threefold reduction in reconstruction errors and a twofold reduction in forecast errors when compared to the original GOKU-net. In the case of the empirical dataset, compared to the original GOKU-nets, GOKU-UT managed to accurately capture the complex brain dynamics, achieving fivefold reduction in reconstruction error. The GOKU-UT, while an extension of GOKU-nets rather than completely new model, significantly enhances
10

---

Published in Transactions on Machine Learning Research (02/2024)
performance by integrating attention and multiple shooting. This has resulted in large effect sizes across various applications, as detailed in Supplementary Tables 1 and 2. The notable exception is in IMRT future predictions beyond 15 seconds, where these enhancements are less effective.
Furthermore, we implemented GOKU-nets in the Julia Programming Language, which broadened the capa- bilities of the model. It helped overcome the initial limitation to Ordinary Differential Equations (ODEs) and facilitated the use of a wider range of differential equation classes (e.g., SDEs), as well as alternative advanced solvers and sensitivity algorithms. The unique SciML ecosystem of the Julia Language proved to be a potent and effective tool for research at the intersection of dynamical systems and machine learning.
Although the Stuart-Landau model has appeared in prior literature modeling brain dynamics (Deco et al., 2017; Domnelly-Kehoe et al., 2019; Ipiña et al., 2020), it has generally been employed with fixed coupling between oscillators. In these empirical studies, structural connectivity estimates were derived from Diffusion Tensor Imaging (DTT). Other parameters were adjusted to maximize some goodness-of-fi critcrion between the simulated and empirical averaged functional connectivities. In contrast to those studies, the methodology employed in the current work targets the minimization of residuals between the simulated and empirical time series themselves. On the other hand, to our knowledge, this is the first instance of employing Stuart-Landar oscillators to model latent brain data dynamics while simultancously inferring oscillator connectivity and learning the nonlinear transformation into the latent space. This approach differs from previous works, such as Abrevaya et al. (2021), which fitted the more general van der Pol oscillators to the latent representation of brain data. Specifically, their learning process involved two separate procedures: encoding in the latent space and parameter estimation of the differential equations. In contrast, our GOKU-UI model integrates these processes into a. single end-to-end training regime. As discussed in Ramezanian- this approach provides the unique advantage of integrating the known dynamics that govern the system, thus facilitating the interpretability and potential applicability with a smaller training set, while still maintaining the flexibility to learn nonlinear dependencies.
et al. (2022),
We have demonstrated these capabilities by training the GOKU-UI on fMRI data and encoding whole-brain dynamics into a latent representation. This representation's 's temporal evolution is effectively modeled by a low-dimensional, interpretable dynamical system, which can yield profound insights into brain functionality, such as the inference of functional connectivity. Beyond thcoretical understanding, the model also holds potential for applied usage, including the classification of mental states or psychiatric conditions. These applications could either leverage the parameters of the differential equations or exploit higher-level features of the latent system, such as its attractor topology.
Despite its versatility as a tool in the SciML toolbox, GOKU-UT's main advantage may somelimes also be its primary limitation: unlike traditional, more agnostic machine learning models, GOKU-UT requires a preliminary differential equation model hypothesized to govern the data's intrinsic temporal dynamies. This requirement may be challenging to meet in many cases. For example, with latent ODEs, one can bypass this task, allowing another neural network to learn the differential equation. However, the significant complexity of the system under investigation, as evidenced by our experiments, could potentially hinder the efficacy of this method.
To the uninitiated in the field of dynamical systems, the process of proposing a specific differential equation to model the data's intrinsic but not immediately evident dynamics might seem like a guessing game. However, this approach has been the successful foundation of the field of physics since the era of Newton The application of GOKU-UI to new problem might not be straightforward as general-purpose black-box neural network model. Still, when guided by the vast theory of dynamical systems, it is not only possible, but potentially highly fruitful.
Acknowledgments
Guillaume Dumas was supported by the Institute for Data Valorization, Montreal (IVADO; CF00137433 & PRF3) Professor Startup & Operational Funds, the Fonds de la Recherche en Santé du Québec (FRQ; 295289; 295291) Junior salary award, the Natural Sciences and Engineering Research Council of Canada (NSERC; (HECR-2002300089)) the Brain Canada Foundation (2022 Future Leaders in Canadian Brain Research program), and the Azrieli Global Scholars Fellowship from the Canadian Institute for Advanced Research
11

---

Published in Transactions on Machine Learning Research (02/2024)
(CIFAR) in the Brain, Mind, & Consciousness program. Irina Rish, Mahta Ramezanian Panahi and Jean- Christophe Gagnon-Auder acknowledge the support from the Canada CTFAR AT Chair Program, the Canada Excellence Research Chairs (CERC) program. Furthermore, Mahta Ramezanian Panahi acknowledges the UNIQUE Center support. Finally, the work of Silvina Ponce Dawson and Germán Abrevaya was funded by UBA (UBACyT 20020170100482BA) and ANPCYT (PICT-2018-0202)) PICT-2001-I-A-00122, PICT-2021- ITTA-00091).. The computational resources used in this work were provided (in part) by the HPC center DIRAC, funded by Instituto de Física de Buenos Aires (UBA-CONICET) and part of the SNCAD-MinCyl initiative, Argentina.
References
Germán Abrevaya, Guillaume Dumas, Alcksandr Y. Aravkin, Peng Zheng, Jean-Christophe Gagnon-Audet, James Kozloski, Pablo Polosecki, Guillaume Lajoie, David Cox, Silvina Ponce Dawson, Guillermo Cecchi, and Irina Rish. Learning brain dynamics with coupled low-dimensional nonlinear oscillators and deep recurrent networks. Neural Computation, 33(8))2087-2122) 2021.
Mohammad R. Arbabshirani, Eswar Damaraju, Ronald Phlypo, Sergey Plis, Elena Allen, Sai Ma, Danicl Mathalon, Adrian Preda, Jatin G. Vaidya, Tülay Adali, and Vince D. Calhoun. Impact of autocorrelation on functional connectivity. Neuroimage, 102:294 308, 2014.
Ellen Baake, Michael Baake, HG Bock, and KM Briggs. Fitting ordinary differential equations to chaotic data. Physical Review A, 45(8):5524, 1992.
Nathan Baker, Frank Alexander, Timo Bremer, Aric Hagberg, Yannis Kevrekidis, Habib Najm, Manish Parashar, Abani Patra, James Sethian, Stefan Wild, Karen Willcox, and Steven Leo. Workshop Report on Basic Research Needs for Scientific Machine Learning: Core Technologies for Artificial Intelligence. Technical report, USDOE Office of Science (SC), Washington, DC (United States), 2019.
Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B Shah. Julia: A fresh approach to numerical computing. SIAM review, 59(1):65-98, 2017.
Hans Georg Bock and Karl-Josef Plitt. A multiple shooting algorithm for direct solution of optimal control problems. IFAC Proceedings Volumes, 17(2):1633-1660 1608, 1984.
Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15):3932 -937, 2016. doi: 10.1073/ 21/0///nas.155 151338411..
Bo Chang, Lili Meng, Eldad Haber, Lars Ruthotto, David Begert, and Elliot Holtham. Reversible architec- tures for arbitrarily deep residual neural networks. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018
Ricky T. Q. Chen, Yulia. Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018.
Jacob Cohen. Statislical power analysis for the behavioral sciences. Academic press, 2013.
George Datseris, Jonas Isensce, Sebastian Pech, and Tamás Gál. DrWatson: the perfect sidekick for your scientific inquiries. Journal of Open Source Software, 5(54):2673, 2020.
Gustavo Deco, Morten L Kringelbach, Viktor K Jirsa, and Petra Ritter. The dynamics of resting fluctuations in the brain: metastability and its dynamical cortical core. Scientific reports, 7(1):3095, 2017.
Moritz DichI, Hans Georg Bock, Holger Diedam, and P-B Wieber. Fast direct multiple shooting algorithms for optimal robot control. Fast motions in biomechanics and robotics: optimization and feedback control, pp. 65 93, 2006.
12

---

Published in Transactions on Machine Learning Research (02/2024)
Patricio Donnelly-Kchoe, Victor M Saenger, Nina Lisofsky, Simone Kühn, Morten L Kringelbach, Jens Schwarzbach, Ulman Lindenberger, and Gustavo Deco. Reliable local dynamics in the brain across sessions are revealed by whole-brain modeling of resting state activity. Human Brain Mapping, 40(10):2967 2980,
2019.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural ODEs. Advonces in neural infor- malion processing systems, 32, 2019.
Ciprian Foias, George Sell, and Roger Temam. Inertial manifolds for nonlinear evolutionary equations.
Journal of Differential Equations, 73(2):309 -333,0933333-00 1988.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27.
Curran Associates, Inc., 2014.
Mengwu Guo and Jan S. Hesthaven. Data-driven reduced order modeling for time-dependent problems. Computer Methods in Applied Mechanics and Engineering. 345:75-99 99, 2019. ISSN 0045-7825.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1): 014004, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Michael Innes, Elliot Saba, Keno Fischer, Dhairya Gandhi, Marco Concetto Rudilosso, Neethu Mariya Joy, Tejan Karmali, Avik Pal, and Viral Shah. Fashionable Modelling with Flux. arXie:1811.01457, 2018.
Ignacio Perez Ipiña, Patricio Donnelly Kehoe, Morten Kringelbach, Helmut Laufs, Agustín Thañez, Gustavo Deco, Yonatan Sanz Perl, and Enzo Tagliazucchi. Modeling regional changes in dynamic stability during sleep and wakefulness. Neurolmage, 215:116833, 2020.
Beatrice M Jobst, Rikkert Hindriks, Helmut Laufs, Enzo Tagliazucchi, Gerald Hahn, Adrián Ponce-Alvarez, Angus Stevner, Morten I, Kringelbach, and Gustavo Deco. Increased stability and breakdown of brain effec- tive connectivity during slow-wave sleep: mechanistic insights from whole-brain computational modelling. Scientific reports, 7(1):1 16, 2017.
Patrick Kidger. On neural differential equations. arXiv:2202.024456 2022.
Byungsoo Kim, Vinicius C. Azevodo, Nils Thucrcy, Theodore Kim, Markus Gross, and Barbara Solenthalor. Deep Fluids: Generative Network for Parameterized Fluid Simulations. Computer Graphics Forum, 38 (2):59-70 70, 2019.
Youngkyu Kim, Youngsoo Choi, David Widemann, and Tarek Zohdi. A fast and accurate Dhysics-informen neural network reduced order model with shallow masked autoencoder. Journal of Computational Physics, 451:110841, 2022.
Diederik P Kingma and Max Welling. Auto-cncoding variational bayes. arXi:: 2013.
Stefan Klöppel, Sarah Gregory, Elisa Scheller, Lora Minkova, Adeel Razi, Alexandra Durr, Raymund AC Roos, Blair R Leavitt, Marina Papoutsi, G Bernhard Landwehrmeyer, al. Compensation in preclinical Huntington's 's discase: evidence from the track-on HD study. EBioMedicine, 2210):142- 2015.
Yoshiki Kuramoto. Chemical Oscillations, Waves and Turbulence. Springer-Verlag, 1984.
Yuri A Kuznetsov, Iu A Kuznetsov, and Y Kuznetsov. Elements of applied bifurcation theory, volume 112. Springer, 1998.
Ori Linial, Neta Ravid, Danny Eytan, and Uri Shalit. Generative ODE modeling with known unknowns. In Proceedings of the Conference on Health, Inference, and Learning. ACM, apr 2021.
13

---

Published in Transactions on Machine Learning Research (02/2024)
Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference on Machine Learning, pp- 3276 3285. PMLR, 2018.
David J. Lucia, Philip S. Beran, and Walter A. Silva. Reduced-order modeling: new approaches for compu- tational physics. Progress in Acrospace Sciences, 40(1):51-117, 2004.
Helmut Lütkepohl. New introduction to multiple time series analysis. Springer Science & Business Media, 2005.
Yingbo Ma, Vaibhav Dixit, Michael J Innes, Xingjian Guo, and Chris Rackauckas. A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions. In 2021 IEEE High Performance Extreme Computing Conference (HPEC), pp. 1-9. IEEE, 2021a.
Yingbo Ma, Shashi Gowda, Ranjan Anantharaman, Chris Laughman, Viral Shah, and Chris Rack- auckas. ModelingToolkit: A Composable Graph Transformation System For Equation-Basee Based Modeling.
arXiv:2103.05244 2021b.
José D Martfn-Guerrero and Lucas Lamata. Reinforcement learning and physics. Applied Sciences, 11( 8589, 2021.
Luke Metz, Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv:2111.05803 2021.
Diganta Misra. Mish: A self regularized non-monotonic activation function. arXic:190..0800 2019.
Mahboubeh Parastarfeizabadi and Abbas Z Kouzani. Advances in closed-loop deep brain stimulation devices.
Journal of neuroengineering and rehabilitation, 14(1):1-20 2017.
Barak A Pearhmutter. Gradient calculations for dynamic recurrent neural networks: survey. Transactions on Neural networks, 6(5):1212-1228, 1995.
IEEE
Pablo Polosecki, Eduardo Castro, Irina Rish, Dorian Pustina, John Н Warner, Andrew Wood, Cristina Sampaio, and Guillermo Cecchi. Resting-state connectivity stratifies premanifest Huntington's discase by longitudinal cognitive decline rate. Scientific reports, 10(1):1-15, 2020.
Chris Rackauckas, Mike Innes, Yingbo Ma, Jesse Bettencourt, Lyndon White, and Vaibhav Dixit. 02376, 2019.
Diffeqflux.j.. julia library for neural differential equations.
Christopher Rackauckas and Qing Nic. Differentaalcquationn.--- performant and feature-rich ecosystem for solving differential equations in julia. Journal of Open Research Software, 5(1):15, 2017.
Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Do- minic Skinner, Ali Ramadhan, and Alan Edelman. Universal Differential Equations for Scientific Machine Learning. arXiv:2001..0 009358 2020.
M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning frame- work for solving forward and inverse problems involving nonlinear partial differontial equations. Journal of Computational Physics, 378:686-707, 2019.
Mahta Ramesanian-Panahhi Germán Abrevaya, Jean-Christophe Gagnon-Audet, Vikram Voleti, Irina Rish, and Guillaume Dumas. Generative models of brain dynamics. Frontiers in Artificial Intelligence, pp. 147, 2022.
Antônio H Ribeiro, Koen Tiels, Jack Umenberger, Thomas B Schön, and Luis A Aguirre. On the smoothness of nonlinear system identification. Automatica, 121:109158, 2020.
R Rico-Martinez, JS Anderson, and IG Kevrekidis. Contimuous-time nonlincar signal processing: a neural network based approach for gray box identification. In Proceedings of IEEE Workshop on Neural Networks for Signal Processing, pp. 596 605. IEFE, 1994.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Ramiro Rico-Martinez and Ioannis G Kevrekidis. Continnous time modeling of nonlinear systems: A neural network-based approach. In IEEE International Conference on Newral Nelworks, pp. 1522-1555 1525.. TEEE, 1993.
Ramiro Rico-Martinez, K Krischer, TG Kevrekidis, MC Kube, and JI. Hudson. Discrete-vs. contimuuus-time nonlinear signal processing of Cu electrodisooution data. Chemical Engineering Communications, 118(1): 25 48, 1992.
Yulia Rubanova, Ricky T. Q. Chen, and David K Duvenaud. Latent ordinary differential equations for irregularly-samplee time series. In Advances in Neural Informalion Processing Systems, volume 32, 2019.
Chaopeng Shen, Alison P Appling, Pierre Gentine, Toshiyuki Bandai, Hoshin Gupla, Alexandre Tartakovsky, Marco Baity-Jesi, Fabrizio Fenicia, Daniel Kifer, Li Li, et al. Differentiable modelling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Enwironment, pp. 1-16, 2023.
Evren Mert Turan and Johannes Jäschke. Multiple shooting for training neural differential equations on time series. IEEE Control Systems Letters, 6:1897 1902, 2021.
Victor Vantilborgh, Tom Lefebvre, and Guillaume Crevecoeur. Efficient ODE Substructure Identification of the Acrobot under Partial Observability using Neural Networks and Direct Multiple Shooting. In 2022 IEEE/ASME International Conference on Advanced Intelligent Mechatronics (AIM), pp. 1268.. JEÉE, 2022.
Gaäl Varoquaux, Sepideh Sadaghiani, Philippe Pinel, Andreas Kleinschmidt, Jean-Baptiste Poline, and Bertrand Thirion. A group model for stable multi-subject ICA fMRI datasets. Neuroimage, 51(1): 288-299, 2010.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
J. Viquerat, P. Meliga, A. Larcher, and F. Hachem. A review on deep reinforcement learning for fluid mechanics: An update. Physics of Fluids, 34(11):111301, 11 2022. URL https:///01.org/10..1 0128446.
Laura von Rueden, Sebastian Mayer, Katharina Beckh, Bogdan Georgiev, Sven Gicsselbach, Raoul Heese, Birgit Kirsch, Julius Pfrommer, Annika Pick, Rajkumar Ramamurthy, Michal Walczak, Jochen Garcke, Christian Bauckhage, and Jannis Schuecker. Informed Machine Learning - A Taxonomy and Survey of Inte- grating Prior Knowledge into Learning Systems. IEEE Transactions on Knowledge and Data Engineering, 35(1):614-633) 2023.
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and Mitigating Gradient Flow Pathologies Physics-Informed Neural Networks. STAM Journal on Scientific Computing, A3055)33055-A3000 2021.
Ee Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 155)11-11, 2017.
Susan WhitfelldGGabrie Joseph M Moran, Alfonso Nieto-Castanón, Christina Triantafyloou, Rebecca Saxe, and John DE Gabrieli. Associations and dissociations between default and self-reference networks in the human brain. Neuroimage, 55(1):225- 232, 2011.
Jared Willard, Xiaowei Jia, Shaoming Xu, Michael Steinbach, and Vipin Kumar. Integrating Scientific Knowledge with Machine Learning for Engineering and Environmental Systems. ACM Compul. Sure., 55 (4), 2022.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Supplementary Information
A Models architectures
Referring to the diagram in Figure 1, the specific architecture used for the different models, for both simulated and empirical data experiments, as follows:
A.1 Basic GOKU-nets
Feature Extractor
ResNet with 4 fully-connected layers, each with 200 neurons and using mish activation functions (Misra, 2019). Input dim dimension of input. data. Output dim 128.
Pattern Extractor
Initial values path: an RNN with layers and 64 neurons in each with ReLU activations. Input dim 128. Output dim 64.
Parameters path: Bidirectional LSTM with 2 layers and 64 neurons in each. Input dim 128. Output dim 128. Note that the dimension of the output of the forward LSTM and the backward LSTM are 64 but when concatenating them, the resulting output dimension is the given one.
Latent in
Initial values path: single-layered fully connccted NN. Input dim 64. Output dim 64.
Parameters path: fully connected NN with layer. Input dim 128. Output dim 128.
Latent out
Initial values path: fully connected NN with 2 layers and 200 neurons in the hidden layer, using no activation function (identity). Input dim 64. Output dim number of state variables of the differential equation.
Parameters path: fully connected NN with 2 layers and 200 neurons in the hidden layer, using sigmoid acti- vation function. The parameters projected from the interval [0, 1] to the desired range when integrating the differential equation. Input dim = 128. Output dim = number of parameters of the differential equation.
Differential Equation layer
The predefined differential equation is solved numerically for each of the sets of parameters and initial conditions provided by the previous layer. The output is the trajectorics at time points equivalent to the input data.
Roconstructor
ResNet similar to the one the Feature Extractor, except that in this case the input dimension is the number of state variables of the differential equation and the output dimension is the one corresponding the input data.
A.2 GOKU-nets with attention
With the exception of the Pattern Extractor, the rest of the layers in the GOKU-nets with attention model remain identical to those in the basic GOKU-nets.
Pattern Extractor
Initial values path: LSTM with 1 layer. Input dim 128. Output dim = 128.
Parameters path: Bidirectional LSTM (BILSTM) with 1 layer. Input dim 128. Output dim 128. A fully connected NN with input and output dimensions of 128 is used for the attention mechanism. This
16

---

Published in Transactions on Machine Learning Research (02/2024)
attention NN processes all the output sequences of the BiLSTM, after which softmax is applied across the time dimension in order to obtain the attentional scores that will be used in the weighted sum of all the time steps returned by the BiLSTM.
A.3 LSTM baseline model
The whole architecture is the same as in the basic GOKU-net, except for the Differential Equation layer, which is replaced by an LSTM:
LSTM layer
We used single-layered LSTM with input and output dimensions set 2 dim. This value is determined in each experiment to ensure that the total number of parameters in the LSTM model closely matches that of the corresponding GOKU-UI. For the simulated dataset experiments, the empirical dataset experiments, dim 105. The LSTM operates recursively. It takes as its first input the value equivalent to the initial condition in differential equations. Subsequently, the model feeds back its last output as the new input, continuing this process until the nunber of time steps matches that of the model's input.
dim 42. In the
of
A.4 Latent ODE baseline model
The whole architecture is the same as in the basic GOKU-net, except for the Differential Equation layer, which is replaced by a Neural ODE:
Neural ODE layer
Neural ODE parametrized by a fully connected NN with 3 layers and node__iidees hiddend dim neurons in each. The input and output dimensions are given by =_dim, which is the number of state variables. In the case of the simulated dataset experiments, the number of state variables was selected to match the true latent dimension _im = 6 and the number of neurons in each layer was adjusted so that the total number of parameters in the model matched as closely as possible that of the corresponding GOKU-UI, resulting in sode__idden____ _diiddiim-- = 137. On the other hand, in the case of the fMRI experiments, the number of and node hidden dim 317, also matching the total number of
state variables was set to z dim parameters of the corresponding GOKU-UI model.
A.5 ANODE
In the case of the ANODE model, there is no latent representation, so it doesn't use the same Latent DE structure as the previous models. The neural ODE in this case is parametrized by layers of fully connected NN, with their sizes adjusted match the size of the GOKU-UI: 380 neurons in the cases of the simulated dataset and 675 neurons in the case of the empirical brain data. In both cases, the augmentation was of 100 dimensions.
A.6 VAR
The VAR model is also applied directly to the input data without any latent representation inference. The implementation from the Python package statsmodels used. The lag order selected in each case such that minimized the AIC in the training datasets.
B Comprehensive description of experiments
B.1 Simulated dataset generation
The high-dimensional simulated dataset used for training the model was constructed based on the simulations of coupled Stuart-Landau ascillators (Eqs. 2) with different random sets of parameters. Each set of parameters corresponds to a different training sample. Whenever we used the Stuart-Landan model in our experiments (both when generating the dataset and when using it inside the GOKU-nets). the time
17

---

Published in Transactions on Machine Learning Research (02/2024)
rescaled by multiplying the right-hand side of Eqs. by 20. Thus, when integrating the equations with the used di = 0.05, the input sequences of length 46 time steps contain few oscillations. The parameters w and C' were sampled from uniform distributions within the following ranges
a [-0.2. w E [0.08T, 0.14T]; Ce [0.0.2]
while G = and 1 = 0.02. On the other hand, the initial conditions for the six state variables were sampled from uniform distributions within the ranges 0.43. For each set of parameters and initial conditions, the system is integrated with the SOSRI solver, Stability-optimized adaptive strong order 1.5 and weak order 2.0 for mongonal/scalas Tto SDEs, from the Differentia//quatioonn complete time span of the integration is 35 units of time and the trajectories are saved every 0.05, resulting in 700 time points. The first 100 time steps are trimmed, in order to remove possible initial transients. Afterwards, a fixed random linear transformation independently applied to each of the 600 remaining time steps, in order obtain 784 dimensions. In other words, every state vector of length from each sample is multiplied by the same 784 matrix, initialized by randomly sampling from uniform distribution in the range [-1, 1]. A training dataset created with 4800 samples, which serves as the source for the different training instances using various sizes of training sets, ranging from 75 to 4800 samples. Two other datasets were generated: one with 200 samples for validation and another with 900 samples for use in the subsequent model evaluations. A schematic representation of the data splits is presented in the left pane of Figure 5. Keep mind that each sample corresponded to a different set of initial conditions and parameters of the differential equations, ensuring that the three datasets are different.
Julia package (Rackauckas & Nie, 2017). The
B.2 Empirical dataset generation
We used resting state fMRI data from 153 participants, obtained from the Track-On HD study (Klöppel et al., 2015). The data underwent pre-processing, described in Polosecki et al. (2020), and 20-component Canonical ICA (Varoquaux et al., 2010) was performed. Upon inspecting the resulting 20 components, 9 were identified as artifacts and thus discarded, leaving 11 components for further analysis in our experiments. Each subject contributed data from two visits, accumulating a total of 306 data samples. Each sample comprised 160 time points, obtained at a temporal resolution of 3 seconds.
In contrast to the simulated data scenario, the empirical data was considerably limited. Consequently, we optimized our data splitting approach. To avoid 'double-dipping,, we ensured that the training and testing datasets comprised entirely distinct samples. However, for the training and validation sets, we utilized the same samples, implementing a temporal split. Specifically, approximately 20% of the data samples (n-60) were reserved for testing, while ensuring a balanced representation of the sex, condition, and measurement site. The rest of the data samples (n=246) were allocated for training and validation purposes. Within this subset, the initial 114 time points were designated for model training. The remaining data served for validation and for an criterion early training termination. Finally, the training, validation, and lest splits were all normalized by the standard deviation of the training set. The right pane of Figure 5 displays a diagram of the data partition employed for the fMRI dataset.
800 time points
114 time points
48 time points
78/1503000_..4BB0
samples
training
248 samples
training
validation
samples
validasion
900 samplee
feeling
samples
testing
160 time points 3Hz
Figure 5: Graphical representation of the data splits of the [MRT dataset (right pane) into distinct sets for training, validation, and testing ensures evaluations of models unseen data during the training phase.
synthetic dataset (left pane) and
18

---

Published in Transactions on Machine Learning Research (02/2024)
B.3
Training settings
All the experiments underwent the same training procedure with identical hyperparameters, which will be described here.
The input sequence length for the models was 46 time steps, and the batch size was set at 64. As described above, the full length of each sample in the training sets was 600 time steps for the synthetic dataset and 114 for the fMRI dataset. The procedure for generating a batch of training data is follows: First, 64 samples that have not been used proviously in the current training epoch randomly selected. Then, for each sample, 46 Line-step-long interval is randomly chosen within the 600 or 114 time sleps available in the full training sample length. Refer to Figure for visual representation of the data partitioning. The GOKU-net based models, contain the same Stuart-Landau differential equations as described above, however, the allowed ranges of parameters differ from the ones used during the generation of the synthetic dataset. In order be closer to real world use-case we allow for a wider range of parameters than those actually used for generating the data, since in principle one would not know the range:
[-1,1]; E []
while keeping, the other parameters the same except of the connectivity in the empirical fMRI training, in which case it was allowed to be negative: C' E [-0.2, 0.2]. The differential equations definitions were optimized for higher computational performance with the help of ModelingToolliit, ji (Ma et al., 2021b). During training, they were solved with the SOSRI solver, a Stability-opimmizd adaptive strong order 1.5 and weak order 2.0 for diagonal/seen /scllar Ito SDEs, from the DiffcrentialEquuatons]]] Julia package (Rackauckas & Nic, 2017). The ForwardDiffSensttivi
sensitivily algorithm used et al., 2020). The models were defined and trained within the deep learning framework of the Flux.jl package (Innes et al., 2018). The experiments were managed using Dr DrWatson.jl package (Datscris et al., 2020).
from the SeiM[Sensitivitg.]] package (Rackauckas
0.8
Reconstruction NRMSE
0.2
0.1
variationa
variational attention
non-variational basic
non-variational attention
basic
Figure 6: Box plot of the normalized RMSE on the test 8ynthetic dataset for GOKU-nets trained with multiple shooting either in their basic or attention variants and in a variational or non-variational version. All hyperparameters are kept as the main paper experiments.
The model was trained with Adam with a weight decay of 10-10, and the learning rate was dynamically determined by the following schedule. The learning rate begins with linear growth (also referred to
19

---

Published in Transactions on Machine Learning Research (02/2024)
0.095
Reconstruction NRMSE
0.090
0.085
0.080
0.075
0.25
0.5
1.0
2.0
4.0
B.0
continuity regularization coefficient
Figure 7: Reconstruction normalized RMSE of GOKU-UI on the test synthetic datasct for different values of the regularization coefficient, in the loss function regarding the continuity constraint. All other hyperpa- rameters are kept as in the main paper experiments.
learning rale warm-up) from 10-7, escalating up 0.005251 across 20 epochs. Afterwards, maintains that value until the validation loss stagnates (has not achieved lower value for 50 epochs), at which point it starts sinusoidal schedule with an exponentially decreasing amplitude.
For the multiple shooting training, all the presented experiments used time window length of 10, therefore partitioning d6-time-steps-long sequences into windows with their endpoints overlapping. The regulariza- tion coefficient in the loss function for the continuity constraint had value of
Since we found that models with variational auloencoders underperform their non-variational versions (see Figure 6), all the results presented in this work were obtained using non- non-variate variiaiiai GOKU-nets. This is, instead of sampling from normal distributions in the latent space as depicted in Figure 1, we pass forward the mean values Mag and Пa. Thus, the associated loss function does not have the KL divergence term associated with the ELBO but retains the reconstruction loss given by the mean squared error between the output of the model and the input, normalized by the mean absolute value of the input. addition, when multiple shooting training employed, the extra term regarding the continuity constraint is inchuded in the loss function. This extra term consists of the mean squared differences between the last point of window and the initial from the next one, divided by the number of junctions and multiplied by regularization coefficient. Please, note that this continuity regularization is performed in the state space of the differential equation and not in the imput space.
C Further exploration of the model
In Figure 6, box plot is presented, comparing the reconstruction performances on the synthetic dataset when the models are variational versus when they are not.
Figures and 8 present the performance of the GOKU-UI on the synthetic dataset when trained using different hyperparameters. As a reminder, the base values used in our experiments for the continuity reg- ularization coefficient, window length, and input sequence length are 2, 10 and 46 respectively. In Figure all hyperparameters are kept the same, except for the continuity coefficient. Similarly, Figure shows variations with respect to the window length and Figure respect to the input sequence length.
20

---

Published in Transactions on Machine Learning Research (02/2024)
0.10
Reconstruction NRMSE
0.08
0.06
0.04
2.0
5.0
10.0
15.0
20.0
window length
Figure 8: Reconstruction normalized RMSE of GOKU-UI on the synthetic dataset for different values of time window length for the multiple shooting partition. All other hyperparameters kept as in the main paper experiments.
0.092
0.090
Reconstruction NRMSE
0.088
0.086
0.084
0.082
0.080
0.078
28.0
46.0
100.0
190.0
370.0
input sequence length
Figure Reconstruction normalized RMSE of GOKU-UI on the test synthetic dataset for different values of time input sequence length. All other hyperparameters kept as in the main paper experiments.
As evidenced by these results, there cxist other sets of hyperparameters that produce higher performances than hhose obtained with the base hyperparameters used our experiments. The set of hyperparameters had been selected based on a grid search but while using a different dynamical system. In any case, the results presented in this paper serve as proof of concept, showing that even when using sub-optimal set of hyperparameters, GOKU-UI demonstrates significantly better reconstruction and forecast performances with respect to the baselines models and other GOKU-nets variants tested.
21

---

Published in Transactions on Machine Learning Research (02/2024)
6
-0.9
Median Reconstruction NRMSE
-0.8
5
-0.
model N
-0.6
-0.5
3
2
-0.3
-0.2
00.1
true N
Figure 10: Reconstruction errors of GOKU-UI using different number of oscillators (model N) synthetic Stuart-Tandau datasets constructed with true N number oscillators.
Motivated by the question of whether would be possible to identify the latent dimensionality some data using the COKU-UT model, the following experiments were performed. GOKU-UI models with different. numbers of Stuart-Landau oscillators were trained on synthetic datasets generated from distinct numbers not evident.
of latent oscillators but all with the same input size of 784, the latent dimensionality Furthermore, in the training settings from our other experiments, the allowed parameter ranges were wider in the GOKU-UT than when generating the datasets. The results of such experiments are presented in Figure 10, where color represents testing reconstruction errors for GOKU-UIs with model N oscillators on datasets built with latent true N oscillators. We see that with an increasing number oscillators in the model, the error progressively diminishes until reaching the true number of latent oscillators in the data,
1.0
0.8
Reconstruction NRMSE
0.6
0.4
0.2
5
10
15
20
Number of oscillators in the model
Figure 11: Reconstruction normalized RMSE [MRT test data, while changing the number of Stuart. Landau oscillators inside the GOKU-UI model.
22

---

Published in Transactions on Machine Learning Research (02/2024)
and from there, the error gets abruptly reduced. Notably, when using more oscillators in the model than the true latent in the data, the model still learns equally well how to reconstruct it. However, the most salient feature is that in this simulated scenario, it is possible to identify the true latent dimensionality of the data as such of the number of oscillators in GOKU-UI for which the reconstruction error gets abruptly reduced.
A similar attempt was made to infer the latent dimensionality of the IMRT data. However, as seen Figure (where the y-axis is in logarithmic scale), there is no abrupt descent in the reconstruction error for any number of oscillators inside GOKU-UI that tried. Nevertheless, inflection point can be identified for N 17. We concluded that any N 2 17 was adequate to use and chose N 20 to perform the experiments presented in this study.
D Effect sizes
In this section we present the Cohen's d effect sizes comparing the performances (measured in RMSE) from original GOKU-nets (without attention and trained with single shooting) against the GOKU-UI. The effect sizes regarding the evaluations of the simulated Stuart-Landau experiments can be found in Table 1, while the corresponding for the empirical fMRI in Table 2. From these effect sizes, we can conclude that the improvements due to the incorporation of attention and multiple shooting into GOKU-nets are large (d > 0.8, according to the criterion suggested by Cohen (2013)) in all cases except for fMRI future predictions beyond 15 scconds.
Training Samples
Cohen's d Effect Size
Reconstruction Task
Forecast Task
75
16.79 [11.51, 22.07]
6.38 [4.22, 8.54]
150
7.8 [.288, 10.37]
8.75 [5.9, 11.6]
300
7.82 [5.24, 10.4]
6.16
[4.06, 8.26
600
4.75 [3004, 6.47]
4.08 22544, 5.62
1200
5.52 (366, 7.45]
5.26
[3.41, 7.12
2400
3.29 11.4, 4.63]
2.91
[1.65, 4.16
4800
3.86 2.38, 5.35)
3.27
[1.93, 4.61]
Table 1: Cohen's d effect sizes with 95% confidence intervals for lesting performance comparison, measured in RMSE, between the original GOKU-net and GOKU-UT. This table corresponds to the experiments on simulated Stuart-Landau oscillators for both reconstruction and forecast tasks, and the effect sizes are calculated independently for each size the training set.
Time (s)
Cohen's d Effect Size
Forecast Task
3
32.01 [20.51, 43.5]
6
19.91 [12.71, 27.11]
Training Samples
Cohen's Effect Size
9
10.4 66.4,, 14.26)
Reconstruction Task 59.18 [37.98, 80.38]
12
4.57 [2.64, 6.49)
15
1.44 [0.3, 2.58]
246
18
0.41 [-0.61, 1.44]
21
0.31 [-0.72, 1.33]
24
0.48 [00.5, 1.51
27
0.72 [-0.32, 1.77
30
0.97 [-0.1, 2.04]
Table 2: Cohen's effect sizes with 95% confidence intervals for testing performance comparison, measured in RMSE, between the original GOKU-net and GOKU-UI for empirical fMRI data. The table presents results for both reconstruction and forecast tasks, with effect sizes calculated independently for each forecast horizon in the latter.
23

---

Published in Transactions on Machine Learning Research (02/2024)
E Reconstruction plots
To provide a visual representation of the model's performance, this section presents trajectories from both the syuthotic and empirical fMRI test sets, along with thoir corresponding reconstructions by GOKU-UI and the original GOKU-nets (lacking attention mechanisms and trained with single shooting). The x-axis represents time steps in all cases. To display representative cases, samples were selected based on their mean reconstruction RMSE being closest the median error across all samples. For the synthetic data, 11 components were randomly selected for display in Figures 12 and 13, due to the impracticality of displaying all 784 components. Each figure displays results from different instances of models, all trained with 4800 samples but each initialized with unique random seed. For the fMRI data, all 11 ICA components are displayed in Figures 14 and 15.
0.10
0.30
0.08
00.5
-0.10
10
20
30
40
10
20
30
40
10
20
30
40
0.15
0.10
0.2
0.0
0.2
-0.05
-0.15
-0.10
-0.4
00.2
10
20
30
40
0
10
20
30
40
10
20
30
40
0.6
0.10
0.4
0.05
0.3
-0.4
-0.0
0.3
0.6
10
20
30
0
10
30
40
10
20
30
40
0.4
0.2
0.03
0.0
0.00
ground truth original GOKU-net GOKU-UI
0.2
-0.03
-044
20
30
10
30 40
Figure 12: Representative example of 46-time-step input sequence from the synthetic test sot, accompanied by its reconstructions from both GOKU-UI and the original GOKU-nets (lacking attention mechanisms and trained with single shooting). The sample selected that its RMSE was the closest to the median across all samples. 11 randomly selected components out of the 784 displayed.
24

---

Published in Transactions on Machine Learning Research (02/2024)
0.050
0.050
0.02
0.025
0.00
0.000
-0.02
-0.025
10
20
30
40
10
20
30
40
10
20
30
40
0.06
075
0.04
0.050
0.02
3.00
0.02
0.025
0.02
0.000
0.06
-0.02
0.025
10
20
30
40
10
20
30
10
20
30
40
0.04
0.050
0.02
0.025
0,050
0.000
0.025
0.002
0.025
0.04
-0.050
-0.025
10
40
10
20
30
10
20
30
40
0.04
0.02
0.02
0.00
ground truth original GDKU-net
-0.02
0.02
-0.04
GOKU UI
10
10 20
Figure 13: Representative example of 46-time-step input sequence from the synthetic test set, accompanicd by its reconstructions from both GOKU-UT and the original GOKU-nets (lacking allention mechanisms and trained with single shooting). The sample was selected that its RMSE was the closest to the median error across all samples. 11 randomly selected components out of the 784 are displayed. This figure is similar to the previous one but presents results from different instances of the trained models, each initialized with a unique random seed.
1.0
0.5
-0.5
-1.0
-05
-15
-1.5
1.0
20
30
40
10
30
10
20
30
40
1.0
1.0
0.5
0.5
0.5
0.0
0.0
-0.5
-0.5
-0.5
10
20
30
40
10
20
30
40
10
20
30
40
0.5
0.5
0.0
-0.5
0.0
-055
-1.0
-0.5
10
20
30
40
10
20
30
40
10
20
30
40
-2
GOKU-net GOKU-UI
-3
0
10
20
30
40
10 30 40
Figure 14: Representative example of a 46-time-steps input sequence for all considered ICA components from the empirical fMRI test set, accompanied by its reconstructions from both GOKU-UT and the original GOKU-nets (lacking attention mechanisms and trained with single shooting). The sample was selected that its RMSE closest to the median across all samples. The x-axis represents time steps, each corresponding to 3 seconds.
25

---

Published in Transactions on Machine Learning Research (02/2024)
0.5
0.5
0.0
-05
-1.0
-1.0
-15
-1.5
10
20
30
40
10
20
30
40
0
10
20
30
40
1.0
0.0
0.3
-0.5
-0.3
20
30
40
10
30
40
10
20
30
40
1.0
1.0
0.5
0.0
-1
0.0
-1.0
-2
20
40
10
30
10
20
40
1.0
0.5
ground truth onggnnal GOKU-nett
-0.5
1
GOKU UI
-2
20
10 30
Figure 15: Reprosentative example of 46-time-steps input sequence for all considered ICA components from the empirical [MRT lest set, accompanied by its reconstructions from both GOKU-UT and the original GOKU-nets (lacking attention mechanisms and trained with single shooting). The sample was selected that its RMSE was closest to the median error across all samples. This figure is similar to the previous one but presents results from different instances of the trained models, each initialized with unique random seed. The x-axis represents time steps, each corresponding to 3 seconds.
F
ICA spatial support and whole brain forecasts
The spatial support the Canonical ICA used in the current study is shown in Figure 16. Since each spatial component extends across the whole brain, each component is displayed in a glass brain representation, from three different perspectives.
2
3
4
5
6
7
8
9
Figure 16: Spatial support of the Canonical ICA analysis employed in this study.
26

---

Published in Transactions on Machine Learning Research (02/2024)
In Figure 17 comparison of the error in future brain activity forecast between the original GOKU-net and COKU-UT is presented. Color code represents the RMSF between the actual brain activity and the forecasts beyond the reconstruction limit generated by these models. In this full brain visual representation it is also noticeable that GOKU-UI generates better forccasts than its original GOKU-net counterpart. However, the deterioration of the forecast. across the brain not homogeneous forecasts and reconstructions (sce Figure 4) we notice that an accentuated error tends to be present the region of precuneus and posterior cingulate cortex. This area a critical hub of the default mode network and is particularly associated with reflective self-conscious thought and internally focused cognitive activities (Whitfield-Gabriel et al., 2011). In this sense, the neural dynamics of the precuneus and posterior cingulate cortex region are less trivially entrained by the environment because they play a central role in the trafficking of information between all the other regions of the brain (Donnelly-Kehen et al., 2019). As a consequence, they are key parts of the "dynamical core" of the brain, making their activity particularly difficult to predict, especially during resting state (Deco et al., 2017).
either case. Actually, both in the
GOKU-UI
Original GOKU-net
6
9
NRMIS
0.16
18
0.078
27
(8)
Figure 17: Time evolution of the whole brain activity, comparing the forecasts from the original GOKU-net respect to GOKU-UI. The color represent the RMSE normalized by the maximun voxel value and averaged over all testing samples and training instances associated different random seeds.
27


Reviews:
[{"review_id": "cbO2x4vNUJ", "reviewer": "Reviewer_jRpQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper has several strengths. It comprehensively explored and evaluated the GOKU-UI model for reconstructing and forecasting time series data. The writing of the paper is very clear, and the detailed information such as the model architecture, training procedures and hyperparameter settings are clearly given. And the experiment includes both synthetic and empirical fMRI datasets.  \n\nWeaknesses:\n\n1. Compared with GOKU-net, the proposed enhancements are quite incremental. The authors updated the GOKU network backbone by attention mechanism and incorporated multiple shooting training strategies for GOKU-net training. The attention mechanism and multiple shooting training strategy are all well-established in Neural ODE areas. Although some improvements may come from better backbone and training strategy,  the modifications are minor and no valuable insights are provided.\n\n2. The authors claimed that a penalty term has been added into the training objective to regularize the continuity of trajectories across different windows for multiple shooting. The penalty term needs to be better motivated and clarified, as this is the core difference between the proposed training strategy and the original multiple shooting training strategy. From my current understanding,  the continuity may be distributed during the training,  with these soft loss constraints.\n\n3. Experiment results lack enough comparison with GOKU-net.  In GOKU-net [Linial et al. (2021)],  they have conducted comprehensive evaluations on realistic data, including single pendulum reconstruction from pixels, double pendulum reconstruction from pixels and cardiovascular system parameters identification. In this paper, the method was only evaluated on one fMRI forecast task. We cannot draw a confident conclusion whether this method is actually effective. \n\n4. The paper compares the proposed method with GOKU and naive latent Neural ODE. However, it does not compare the performance of the framework with other NODE models specifically designed for complex time-series data. For example, ANODE, which augmented the ODE with a vector of zeros, is a very effective baseline in tackling complex time series data. Comparison with those methods would provide a better understanding of the strengths and weaknesses of the proposed framework in relation to similar approaches.\n\n5. The proposed enhancements can not demonstrate consistent improvements in forecasting with different time intervals. The prediction error becomes more sensitive to the prediction length compared with latent ODE and GOKU baselines. Also latent ODE and naive LSTM can outperform the proposed method by a large margin for 15 - 20 seconds prediction of brain activity.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a new SciML generative model GOKU-UI, building on the basis of the Latent ODEs, for modeling scientific time-series data.  Compared with previous methods,  GOKU-UI incorporates attention mechanisms and multiple shooting training strategy in the latent space, leading to improved empirical performance in reconstruction and forecast tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to above detailed comments in weakness.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Axaj4ntg1I", "reviewer": "Reviewer_ZBuz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is overal well written with convincing experiments and literature\nreview, it is however not very amenable to a large audience not already\nfamiliar with neural ODEs and GOKU-like methods. Indeed the paper is\nvery short and comes without background to make the contribution more\napproachable by people interested in the application of these methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The proposed method extends GOKU-nets with the multiple shooting method\nimprovind forecast of non-linear dynamical system. The method is evaluated\non simulated Stuart-Landau (SL) oscillators and on fMRI rest data showing a gain in\nreconstruction and forecast scores. Results show that the proposed is more\nstable to train and is able to learn with moderate amount of data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Results on fMRI data are potentially interesting yet results do not\noffer any insights on the data and neuroscience interpretation of the\nresults (eg where in the brain does the model perform better? etc.)\nThe results in appendix D show very good reconstruction plots but in\nICA space and we do not know what components look like etc.", "broader_impact_concerns": "The paper is overal well written with convincing experiments and literature\nreview, it is however not very amenable to a large audience not already\nfamiliar with neural ODEs and GOKU-like methods.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EyqYaELyM3", "reviewer": "Reviewer_CTWL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n1.  The paper begins with a well-structured introduction that seamlessly integrates both contemporary and historical references, providing a comprehensive overview of the field.\n2.  The authors have done an excellent job in presenting their research; the paper at large is well-articulated and easy to understand.\n3.  The approach to empirical evaluation is commendable, the authors begin with simulations and then transition into real data application. This succession is usually desirable when one is working with neuroimaging data where the ground truth is unknown.\n\n# Weaknesses\n\n1.  The empirical evaluation as it stands is unacceptable.\n    1.  The key issue lies in evaluating the reconstruction while partitioning the identical time sequences (or subjects in the fMRI example) between the train and test sets. This 'double-dipping' results in overly optimistic results, as the model and ODE parameters are being optimized based on the very same time series data. **Please, separate the training and test subjects as well as simulated sequences.** I may be missing the point here and the propose model is literally an approach to fit an ODE to a specific dataset, but then the paper is putting forward quite different claims.\n    2.  While the forecasting experiments may not be as vulnerable to the previous issues, the paper lacks a thorough evaluation of its forecasting performance.  The models that the paper does compare with are closely related to the type of model utilized. Meanwhile, time series forecasting is a vibrant field complete with numerous benchmarks. However, no comparison with the more standard forecasting models is provided.\n2.  The paper lacks a convincing justification of the fMRI forecasting experiments. The use-case of this exercise neither explained nor explored. I believe a reader would appreciate more motivation in this section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper offers two amendments to the original Generative ODE Modeling with Known Unknowns (GOKU) model. First, it introduces changes to the model's structure by substituting the final biLSTM elements in the Pattern Extractor with an attention weighted blend of all outputs for each sequence element. Second, it modifies the training process by adopting a 'multiple shooting' approach to sequence generation. The effectiveness of these alterations is evaluated against the original GOKU model and other alternatives, utilizing both simulated data with established ground truth and an fMRI dataset.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.  Redo experiments reported in Figures 2 and 3 using a proper cross validation, when time-series are not split in parts and then distributed across training and test data. No part of any training time series should be used in the test.\n2.  IMO, to ensure that the entire class of presented models isn't underperforming on the tasks demonstrated, the paper **must** include other models for baseline reference, such as the following [package for forecasting models](https://github.com/Nixtla/statsforecast/tree/main)\n3.  Minor editorial problems I have noticed:\n    -   Section 2.3.2 is called Empiric data. You must mean empirical evaluation?\n    -   It would greatly enhance the accessibility of the referenced literature, at which the paper excels, if all the citations in the PDF were turned into clickable hyperlinks.\n    -   First line of section 2.1 \"could be thought as\" -> \"could be thought of as\". Consider rewriting this sentence entirely - it is awkwardly phrased.", "broader_impact_concerns": "no concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper utilizes attention mechanisms and multiple shooting training strategy in the latent space to build on GOKU-net. The resulting GOKU-UI is benchmarked on reconstruction and forecasting tasks. The authors were able to resolve the concerns that reviewers had with the initial submission, to argue that GOKU-UI is a substantial improvement over GOKU-nets and the evaluation benchmarks are performed without label-leakage.
Decision: Accept (poster)