Interpretable Spatial-Temporal Fusion Transformers: Multi-Output Prediction for Parametric Dynamical Systems with Time-Varying Inputs

TL;DR

This work extends the Temporal Fusion Transformer to a multi-output, interpretable framework (iSTFT) for predicting parametric dynamical systems with time-varying inputs. By reshaping data into a spatial-temporal sequence and employing block-wise masked interpretable attention, iSTFT captures both temporal dynamics and inter-output correlations while preserving interpretability through variable importance and attention patterns. Across Lorenz-63, FitzHugh-Nagumo, and Ferrocyanide oxidation models, iSTFT achieves accurate, single-pass multi-output predictions with MAE generally outperforming MSE, and provides insight into how parameters, inputs, and past outputs influence future QoIs. The approach reduces surrogate-model training effort, handles nonlinear and high-dimensional parameter spaces, and offers practical impact for fast, interpretable surrogates in complex engineering systems.

Abstract

We explore the promising performance of a transformer model in predicting outputs of parametric dynamical systems with external time-varying input signals. The outputs of such systems vary not only with physical parameters but also with external time-varying input signals. Accurately catching the dynamics of such systems is challenging. We have adapted and extended an existing transformer model for single output prediction to a multiple-output transformer that is able to predict multiple output responses of these systems. The multiple-output transformer generalizes the interpretability of the original transformer. The generalized interpretable attention weight matrix explores not only the temporal correlations in the sequence, but also the interactions between the multiple outputs, providing explanation for the spatial correlation in the output domain. This multiple-output transformer accurately predicts the sequence of multiple outputs, regardless of the nonlinearity of the system and the dimensionality of the parameter space.

Figures (12)

• Figure 1: The structure of TFT for parametric output prediction. The main structure is a copy of Fig. 2 in mlLimALetal21. Only the notation of the TFT input data, TFT prediction values and parameters (static metadata in mlLimALetal21) are different. We use TFT to predict the actual output values, referred to as point forecast in mlLimALetal21.
• Figure 2: Structure of the masked self-attention mechanism proposed in mlVasSP17 and used in the TFT.
• Figure 3: The structure of the block-wise masked attention weight matrix $\widetilde{\boldsymbol{A}}$ in a single iSTFT (right) compared to the normal masked attention weight matrix $\overline{\boldsymbol{A}}$ (left) in three separate TFT models for three outputs. Here, we use $n_t = 3$ as an example.
• Figure 4: Lorenz-63 model: the predicted solution (the MSE loss) and the reference solution.
• Figure 5: Lorenz-63 model: the predicted solution (the MAE loss) and the reference solution.

Theorems & Definitions (2)

• Theorem 3.1
• proof