Transformer-based Koopman Autoencoder for Linearizing Fisher's Equation

TL;DR

This work tackles the nonlinear Fisher's reaction-diffusion PDE and proposes a Transformer-based Koopman autoencoder to learn a globally linear latent representation, enabling efficient prediction and analysis. By embedding the state into a latent space of dimension $r$ (here $r=21$) where dynamics are governed by a linear operator $K$, the authors achieve accurate Fisher's equation forecasting and show generalization to other PDEs such as Kuramoto–Sivashinsky and Burger's. The architecture combines outer/inner encoders, a Transformer block, and a five-term loss to enforce autoencoding, latent linearity, and consistent coordinate transformations, with robust performance demonstrated against baselines like DenseRes and ConvRes blocks. Overall, the data-driven approach offers a scalable pathway to linearize nonlinear PDE dynamics without explicit governing equations, with potential applicability to a broad class of reaction-diffusion and other complex PDEs.

Abstract

A Transformer-based Koopman autoencoder is proposed for linearizing Fisher's reaction-diffusion equation. The primary focus of this study is on using deep learning techniques to find complex spatiotemporal patterns in the reaction-diffusion system. The emphasis is on not just solving the equation but also transforming the system's dynamics into a more comprehensible, linear form. Global coordinate transformations are achieved through the autoencoder, which learns to capture the underlying dynamics by training on a dataset with 60,000 initial conditions. Extensive testing on multiple datasets was used to assess the efficacy of the proposed model, demonstrating its ability to accurately predict the system's evolution as well as to generalize. We provide a thorough comparison study, comparing our suggested design to a few other comparable methods using experiments on various PDEs, such as the Kuramoto-Sivashinsky equation and the Burger's equation. Results show improved accuracy, highlighting the capabilities of the Transformer-based Koopman autoencoder. The proposed architecture in is significantly ahead of other architectures, in terms of solving different types of PDEs using a single architecture. Our method relies entirely on the data, without requiring any knowledge of the underlying equations. This makes it applicable to even the datasets where the governing equations are not known.

Figures (6)

• Figure 1: Neural network architecture. (a) Autoencoder comprise an outer encoder/decoder, an inner encoder/decoder and a matrix K, which governs the dynamics of the system. (b) Transformer block that identifies the set of intrinsic coordinates of a linearized dynamical system.
• Figure 2: A comparison of exact solution and neural network predictions for Fisher's equation. The initial conditions utilized include (a) white noise, (b) sine waves, and (c) square waves, which are the same as the training dataset.
• Figure 3: A comparison of exact solution and neural network predictions of Fisher's equation. The initial conditions utilized include (a) Gaussian waves, (b) triangle waves, (c) sawtooth waves and (d) pulse waves which are not found in the training dataset.
• Figure 4: Eigen values of the approximated Koopman operator, $\textit{K}$.
• Figure 5: Relationship between inference time and error as a function of the number of attention heads.