Learning High-dimensional Ionic Model Dynamics Using Fourier Neural Operators

TL;DR

The paper tackles learning the full, high-dimensional dynamics of stiff ionic ODEs governing excitable cells by using Fourier Neural Operators to map applied current $I_{app}$ to the state triplet $(V,\vec{w},\vec{c})$ across 2-, 4-, and 41-variable models. It introduces NOs and the Fourier variant, detailing how $R_{\theta_t}(k)$ parameterizes the kernel in the frequency domain to capture global dependencies, while maintaining mesh-resolution independence. Through automated HPC-based hyperparameter tuning, the authors compare unconstrained and constrained parameterizations, finding similar accuracy across FitzHugh–Nagumo, Hodgkin–Huxley, and O'Hara–Rudy models, with unconstrained networks converging faster (fewer epochs). The results demonstrate FNOs' capability to accurately reproduce multiscale, high-dimensional ionic dynamics and suggest scalable applications for personalized, HPC-enabled cardiac and neural simulations, with future work extending to spatial PDE formulations.

Abstract

Ionic models, described by systems of stiff ordinary differential equations, are fundamental tools for simulating the complex dynamics of excitable cells in both Computational Neuroscience and Cardiology. Approximating these models using Artificial Neural Networks poses significant challenges due to their inherent stiffness, multiscale nonlinearities, and the wide range of dynamical behaviors they exhibit, including multiple equilibrium points, limit cycles, and intricate interactions. While in previous studies the dynamics of the transmembrane potential has been predicted in low dimensionality settings, in the present study we extend these results by investigating whether Fourier Neural Operators can effectively learn the evolution of all the state variables within these dynamical systems in higher dimensions. We demonstrate the effectiveness of this approach by accurately learning the dynamics of three well-established ionic models with increasing dimensionality: the two-variable FitzHugh-Nagumo model, the four-variable Hodgkin-Huxley model, and the forty-one-variable O'Hara-Rudy model. To ensure the selection of near-optimal configurations for the Fourier Neural Operator, we conducted automatic hyperparameter tuning under two scenarios: an unconstrained setting, where the number of trainable parameters is not limited, and a constrained case with a fixed number of trainable parameters. Both constrained and unconstrained architectures achieve comparable results in terms of accuracy across all the models considered. However, the unconstrained architecture required approximately half the number of training epochs to achieve similar error levels, as evidenced by the loss function values recorded during training. These results underline the capabilities of Fourier Neural Operators to accurately capture complex multiscale dynamics, even in high-dimensional dynamical systems.

Figures (14)

• Figure 1: Visual representation of a Fourier Neural Operator.
• Figure 2: FHN model: FNO training relative $L^2$ (blue) loss and relative test $L^1$ (green), $L^2$ (orange), and $H^1$ (red) loss functions. (a) Unconstrained FNO model, (b) constrained FNO model.
• Figure 3: FHN model: FNO performance comparison. Figure \ref{['fig:barplot_fhn_unconstrained']} shows the bar plot of the relative $L^2$ error for the unconstrained FNO. Figure \ref{['fig:barplot_fhn_constrained']} shows the bar plot of the relative $L^2$ error for the constrained FNO. Figure \ref{['fig:boxplot_fhn']} is a box plot illustrating the distribution of relative $L^2$ errors for both the constrained and unconstrained architectures.
• Figure 5: FHN model: examples of FNO performance. Each column represents a single example from a subset of the test dataset (Table \ref{['table:fhn_dataset']}). Within each column, the rows illustrate: applied current $I_{app}$, voltage $V$, pointwise error for the voltage $V$, recovery variable $w$, pointwise error for the recovery variable $w$, and phase space.
• Figure 6: HH model: FNO training relative $L^2$ (blue) loss and test relative $L^1$ (green), $L^2$ (orange), and $H^1$ (red) loss functions. (a) Unconstrained FNO model, (b) constrained FNO model.