SineNet: Learning Temporal Dynamics in Time-Dependent Partial Differential Equations

TL;DR

It is demonstrated that increasing the number of waves in SineNet while maintaining the same number of parameters leads to a monotonically improved performance, highlighting the effectiveness of SineNet and the potential of the approach in advancing the state-of-the-art in neural PDE solver design.

Abstract

We consider using deep neural networks to solve time-dependent partial differential equations (PDEs), where multi-scale processing is crucial for modeling complex, time-evolving dynamics. While the U-Net architecture with skip connections is commonly used by prior studies to enable multi-scale processing, our analysis shows that the need for features to evolve across layers results in temporally misaligned features in skip connections, which limits the model's performance. To address this limitation, we propose SineNet, consisting of multiple sequentially connected U-shaped network blocks, referred to as waves. In SineNet, high-resolution features are evolved progressively through multiple stages, thereby reducing the amount of misalignment within each stage. We furthermore analyze the role of skip connections in enabling both parallel and sequential processing of multi-scale information. Our method is rigorously tested on multiple PDE datasets, including the Navier-Stokes equations and shallow water equations, showcasing the advantages of our proposed approach over conventional U-Nets with a comparable parameter budget. We further demonstrate that increasing the number of waves in SineNet while maintaining the same number of parameters leads to a monotonically improved performance. The results highlight the effectiveness of SineNet and the potential of our approach in advancing the state-of-the-art in neural PDE solver design. Our code is available as part of AIRS (https://github.com/divelab/AIRS).

Figures (26)

• Figure 1: Illustration of the misalignment issue in U-Nets. The left side shows the one-step prediction of a trained U-Net, with each intermediate layer's feature maps averaged over the channel dimension displayed in the top row. On the right side, the misalignment between $\ell=0$ and $\ell=3$ is demonstrated. Specifically, the feature map of $\ell=3$ is upsampled and then concatenated with $\ell=0$. The time-evolving effect of the preceding U-Net layers results in a misalignment of the corresponding physical features, as indicated by the bounding boxes (a) and (b). This misalignment is particularly problematic for convolutions. Since the kernel is localized, information from misaligned high frequency features, such as those visualized in (a) and (b), cannot be optimally integrated in updating the feature map. Mitigating this misalignment is key for improving the performance of U-Nets.
• Figure 2: Illustration of the proposed SineNet for learning temporal dynamics in PDEs. Multiple U-Net waves are composed to perform one-step prediction, with the output of each wave averaged over the channel dimension displayed in the top row, demonstrating the time-evolving process from $t$ to $t+1$. The orange sinusoidal line illustrates propagation between resolutions and is not part of the model architecture. Feature maps in this figure are from SineNet without wave residuals for clarity and with transposed convolutions for upsampling. In Appendix \ref{['app:fmvis']}, we visualize the feature maps from the SineNet-8 presented in Section \ref{['sec:results']}
• Figure 3: Number of waves $K$ versus test error. Solid line is rollout error, dashed line is 1-step error. Numerical results are presented in Appendix \ref{['app:numres']}.
• Figure 4: Feature map responses to noise injection for trained U-Net-128 and SineNet-8. As opposed to the U-Net, in SineNet, the perturbation influence propagates gradually from the first to the last feature map, demonstrating the reduced latent evolution managed by each wave.
• Figure 5: Comparison of SineNet, U-Net and Deeper U-Net architectures.