Paper: SineNet: Learning Temporal Dynamics in Time-Dependent Partial Differential Equations (ICLR-2024)

Reviewer comment (Evaluation):
One advantage U-Nets have is that they are storage efficient since the highest resolution is only used in the first and last block. I am not sure how SineNets handle this - probably through P and Q networks? It would be great to have some information on this. Toward this end I am wondering if the number of waves is increasing the memory footprint and by how much?

Author response:
We summarize time and memory usage in Table 4, Appendix D, where we have additionally added training statistics. We can see that although SineNet uses more high-resolution feature maps than a standard U-Net, the forward memory will not increase much during inference, as despite the deep architecutre, activations do not need to be stored. Additionally, we can adjust the channel multiplier (detailed in Table 5, Appendix D.3) to control the number of channels in the down and up blocks. As a result, although the training memory increases more with number of waves compared to inference since more feature activations need to be stored for backpropagation, we are able to constrain the training memory to grow less than 5GB per additional wave with a batch size of 32.

Reviewer comment (Novelty):
The iterative refinement process is conceptually similar to PDE Refiner (Lippe et al.) - with some major differences to my understanding. SineNets iteratively refines the temporal update in an one end-to-end approach whereas PDE Refiner looks at the input several times and refines different frequency components. It might be interesting to contrast those two approaches in the discussion.

Author response:
Thank you for bringing this work to our attention -- we agree that this connection is interesting, and thus have compared and contrasted the two approaches in our updated related work section.

Reviewer comment (Experiments):
The paper strongly references the original paper of Gupta & Brandstetter. This paper put some emphasis on conditioning on equation parameters as well as flexible lead time predictions. It would be interesting to have statements / experiments on this as well. Similar FiLM like encodings would work for SineNet I suppose?

Author response:
Thank you for the suggestion. We have added results on the conditional task in Appendix F.1. We agree that FiLM encodings are applicable to SineNet, although we only considered injecting the embedding of parameters/lead times to feature maps via addition along the channel dimension, and not with the Hadamard product. Interestingly, we observe that the margin between SineNet and the baseline U-Net-Mod is larger for longer lead times, where there is a greater degree of misalignment between the input and target.

Reviewer comment (Evaluation):
The misalignment made me realize why ViTs or Swin Transformers might be better fit for PDE modeling compared to standard U-Nets. Can you comment on this?

Author response:
Indeed, while both U-Nets and attention-based architectures include multi-scale processing mechansisms, allowing them to achieve a global receptive field, Transformers do not include such long range skip-connections and therefore do not encounter the misalignment issue at all. However, we suggest that the advantage of using a U-Net for PDE simulation could lie in sample complexity. From the paragraph starting at the bottom of page 1 of [4]:   

> When trained on mid-sized datasets such as ImageNet without strong regularization, [ViTs] yield modest accuracies of a few percentage points below ResNets of comparable size. This seemingly discouraging outcome may be expected: Transformers lack some of the inductive biases inherent to CNNs, such as translation equivariance and locality, and therefore do not generalize well
when trained on insufficient amounts of data...However, the picture changes if the models are trained on larger datasets (14M-300M images).

Although further experimentation is needed to determine whether this trend holds for PDE modeling, if it does, it could explain why U-Nets are preferrable for PDE modeling -- it is generally infeasible to obtain such large quantities of dynamics data due to the cost of simulation. One prominent exception has been climate modeling and weather forecasting, where vision transformers have been commonly applied [5,6,7]. The dataset used in each of these works, ERA5 [8], contains hourly measurements dating back 40 years, and is therefore several orders of magnitude larger than those considered in the majority of the PDE modeling literature.    

Swin Transformer [9] improves the computational efficiency of ViT by partitioning feature maps into smaller local windows and only computing self-attention inside each window. To model long-range dependencies, the window size is increased deeper into the network. However, since interactions between windows are not modeled, Swin achieves a global receptive field only in the final layers. Although this is not a limiting factor for image tasks since patterns are mainly localized in space, for dynamics modeling, there are stronger long-range dependencies due to processes such as diffusion. As a result, we may need to have global information as early as possible so that local processing can have sufficient global context. To mitigate this issue, we could mirror the Swin Transformer to construct a U-Shaped Swin Transformer that achieves a global receptive field in the middle of the network. Moreover, while limiting interactions to be modeled inside each local window is key for improving efficiency, communication between neighboring windows may be needed to account for advection and diffusion processes in dynamics modeling. Consequently, further modifications to the Swin architectures may be needed to adapt it for dynamics modeling.

Reviewer comment (Novelty):
Stacked U-net have already been developed for different tasks as detailed in the article (section Stacked U-Nets in the related work). The implementation here differs in some minor points and thus is more effective for PDE prediction but this is a marginal improvement over the previous works.

Author response:
One of our primary contributions is the identification of the misalignment issue encountered by U-Nets when faced with temporal dynamics. As many works have considered U-shaped architectures for dynamics forecasting due to their strengths in multi-scale processing, e.g. [1,2,3], this is a far-reaching issue. Although our proposed solution of a stacked U-Net is not by itself a novel methodology, it improves upon state-of-the-art neural solvers by a substantial amount on 3 challenging datasets. Furthermore, through identification of the issue, our contribution has opened the door for future works on adapting U-Nets in other ways beyond stacking to address the misalignment issue.

Reviewer comment (Experiments):
The experiments are not extensive. First, more baselines should be used, with at least DeepONet [1], MP-PDE ([2], cited but not implemented) or Neural ODE [3], which closely resembles the formulation of SineNet. Second, the evolution of the error with the time is not shown, only at a given number of rollout steps for each datasets, thus limiting the impact of the results. Third, no ablation on the influence of the number of historical time steps is performed.

Author response:
### Neural ODE

Thank you for pointing out the connection between SineNet and Neural ODE. We state this in Section 3.3 of our revisions and present experiments and an in-depth discussion in Appendix F.2. We find that adapting SineNet into the neural ODE framework gives results which are competive with SineNet-8 on INS. However, due to the high number of function evaluations per training example, SineNet-Neural-ODE is expensive to train and inference, which we quantify in Table 4. The [code for this experiment](https://anonymous.4open.science/r/SineNet-03B2/pdearena/pdearena/modules/sinenet_neural_ode.py) can be found on our anonymous GitHub.

### DeepONet and MP-PDE

We implemented a DeepONet with a CNN branch net using strided convolution downsampling based on code ([from the official GitHub](https://github.com/lu-group/deeponet-fno/blob/main/src/darcy_rectangular_pwc/deeponet.py) and [the DeepXDE package](https://deepxde.readthedocs.io/en/latest/_modules/deepxde/nn/pytorch/deeponet.html#DeepONet)). Hyperparameter choices (depth 4 for the trunk net, 512 basis functions, constant hidden width of 512) were based on Table C.1 of [7]. We present results for this model on INS below.

Our MP-PDE implementation used code from [from the official GitHub](https://github.com/brandstetter-johannes/MP-Neural-PDE-Solvers/blob/master/experiments/models_gnn.py) adapted to 2D graphs. As the PDEs we considered were discretized on a 128$\times$128 grid, the resulting graph was very large, with >16K nodes and >145K edges. Therefore, we did not modify the default hyperparameters used in the official repo (e.g., depth and width), which gave 638K parameters. Due to the large graph, even this small model did not fit on an 80GB A100 GPU with batch size 32, and we therefore reduced the batch size to 16. We also found injecting noise at level $\sigma=0.01$ to improve performance for this model.  

Results for both models are presented below on INS. As can be seen, further work is needed to achieve better results on this dataset for these models. For DeepONet, we hypothesize that because the model does not exploit the regular grid on which the target function is discretized, and instead can evaluate the target function on any point in its domain, its capacity for modeling regularly-grided functions is naturally less than that of a CNN-type architecture which exploits this structure. Similarly, MP-PDE has the flexibility to model mappings between irregularly-meshed functions, which we discuss as a limitation of SineNet and other CNN-based architectures in Appendix E. However, in the context of modeling on a square regular grid, this flexibility may be a limiting factor. Furthermore, MP-PDE does not include mechanisms for modeling long-range dependencies, thereby limiting the receptive field. This could also explain the performance, and is a limitation addressed in follow-up work [8].  

Our [implementation of MP-PDE](https://anonymous.4open.science/r/SineNet-03B2/pdearena/pdearena/modules/mp_pde.py) can be found on our anonymous GitHub, as can [our implementation of DeepONet](https://anonymous.4open.science/r/SineNet-03B2/pdearena/pdearena/modules/deepONet.py).

| Model | \# Par. (M) | 1-Step (\%) | Rollout (\%) |
|---|---|---|---|
| SineNet-8 | 35.5 | **1.66**| **19.25** |
| DeepONet | 38.5 | >20.00 | >40.00 |
| MP-PDE | 0.6 | >20.00 | >40.00 |


Thank you for the suggestion. We have included plots in Appendix F.4 (Figure 9) to show the error evolution with time on all three datasets.

Thank you for the suggestion. We have trained new models on the incompressible Navier Stokes (INS) dataset conditioning on 1, 2 or 3 historical step(s), and present results in Appendix F.6. We found that our choice in the main text of conditioning on 4 historical steps was not optimal for SineNet. However, it is not clear whether this trend would hold for other baselines, and thus, we choose to adhere to the number of conditioning steps used by the respective benchmarks on each dataset (4 for INS, 10 for CNS and 2 for SWE).

Reviewer comment (Theory):
This is more minor. As with other auto-regressive architectures, the forecasting process is done at regular steps, so the architecture cannot produce values at any time $t$, but only at multiples of $\Delta t$.

Author response:
In the conditional task considered by [1], models are trained to generalize over time step size and forcing terms by conditioning the model on a time step size and the forcing term. Although in this task, models are only supervised on time step sizes which are multiples of 0.375 seconds (the temporal resolution of the data), it is possible to condition the model to advance time by any time step size. However, since this would be out-of-distribution, further data generation and experimentation is needed to assess model performance in such a setting. Nontheless, we have added results on this conditional task to Appendix F.1.

Reviewer comment (Experiments):
Have you tried with a high number of rollout steps? It would be interesting to see how stable the method is when pushed to its limits.

Author response:
Thank you for the suggestion. We have generated 100 trajectories of CNS data which are 10$\times$ longer than those in the main text. We found that although noise injection reduces performance for shorter rollouts, it plays a key factor in maintaining long-term stability. We discuss this in greater depth and present results in Appendix F.5, including the averaged rollout error, as well as a plot showing the evolution of the rollout error across time.

Reviewer comment (Experiments):
In Figure 3, the error seems to keep decreasing with the number of waves. Have you tested with more waves? This would be interesting to see when the error arrives at a plateau and this could help improve the performances of the model if the error keeps decreasing with the number of waves.

Author response:
Thank you for the suggestion. On all three datasets, we have added results for SineNet-10,12,14, and 16. Results are included in Figure 3, and have also been reported numerically in Appendix F.8.

Reviewer comment (Experiments):
Have you tried using noise to improve the rollout stability? (as has been done for instance in [4])

Author response:
Thank you for the suggestion. We have generated 100 trajectories of CNS data which are 10$\times$ longer than those in the main text. We found that although noise injection reduces performance for shorter rollouts, it plays a key factor in maintaining long-term stability. We discuss this in greater depth and present results in Appendix F.5, including the averaged rollout error, as well as a plot showing the evolution of the rollout error across time.

Reviewer comment (Experiments):
In the experiments, null Dirichlet and Neumann boundary conditions were chosen. What happens if different values for these conditions are chosen? Will zero padding still be effective?

Author response:
This is a good point. We have added a discussion on challenges presented by mixed boundary conditions and potential approaches for dealing with non-null boundary conditions in Appendix E on limitations. Although it would appear to be relatively simple to encode both via padding, we leave this to future work.

Reviewer comment (Experiments):
Some ablation study is missing. For example, what about if we replace each $V_{k}$ with standard U-Net with the equations (2) and (3)? Can the authors also provide visualization on this ablation version of the model in a way similar to Figure 5?

Author response:
## **W2. Replacing $V_k$ with standard U-Net**

> Some ablation study is missing. For example, what about if we replace each
 $V_k$ with standard U-Net with the equations (2) and (3)? Can the authors also provide visualization on this ablation version of the model in a way similar to Figure 5?  


There are two main difference between $V_k$ and the standard U-Net (equations (2) and (3)):

1. The dual multi-resolution processing which we ablate with SineNet-8-Entangled in Table 2.
2. We maintain a high-dimensional feature map between waves instead of decoding to the dimension of the solution space and encoding back to the latent space in the following wave.

In Appendix F.7, we ablate (2) with SineNet-8-BottleNeck by decoding and re-encoding the latent solution between waves such that each $V_k$ becomes a standard U-Net. There is a loss of information by linearly projecting from high dimensions to low dimensions and then back again (the bottleneck). We therefore observe that the 1-step loss for SineNet-8-BottleNeck is slightly higher than than of SineNet-8. This has a regularizing effect on the rollout loss, where we observe a slight improvement in performance. Moreover, we visualize the decoded feature maps between each wave in Figure 15 of the revised draft.

Reviewer comment (Experiments):
Why does Dil-ResNet which you benchmark against only have 4.2M parameters (compared to 35M in your model)? This is a fairly critical point, since neither the inference time nor the required memory in Table 4 seem to justify such a small parameter count for Dil-ResNet. Indeed, I was curious to see how vanilla ResNet would perform and like that you have included Dil-ResNet. This actually forced me to only rate your submission as "marginally above threshold" instead of voting for a clear accept.

Author response:
- **Dil-ResNet:** We use the DilResNet from the PDEArena benchmark [1] (DilResNet128-norm in their paper). Although it has fewer parameters, it is one of the strongest model in the PDEArena benchmark -- see Table 8 of [1].
For a more fair comparison with SineNet, we trained a larger Dil-ResNet by increasing the hidden channels from 128 to 256 (dubbed Dil-ResNet-256), which has 16.7M parameters. Although it still has fewer parameters that SineNet, it is especially computationally expensive. We show this in Table 4 of Appendix C, where we report the computational cost of all models, which has been revised to include training stats. We can see that Dil-ResNet is slower in inference and is comparable with SineNet in training speed. However, it almost doubles the GPU memory usage of SineNets in both training and testing. This is because all convolutions and non-linearities take place on the full-resolution, as opposed to reducing the resolution as done in U-Nets and SineNet. The high computational cost prohibits us from using use more parameters for Dil-ResNet. Nontheless, the increase in channels substantially improves performance for Dil-ResNet, particularly on the INS data, where it is now extremely competive with SineNet-8, as can be seen in Table 1.


- **ResNet:** According to results presented in PDEArena [1], ResNet underperforms other methods, including U-Nets, by more than one order of magnitude. Please see Table 8 in [1] for INS results and Table 2 for SWE results; note that their reported metric is MSE, as opposed to the scaled L2 loss as we use here, however, we expect similar results. The results are likely due to the limited receptive field in ResNets due to lack of downsampling or other multi-scale processing mechanisms such as dilation, which limit its ability to model long-range dependencies. Moreover, similar to Dil-ResNet, the computation is also much more expensive than U-Nets due to the lack of downsampling (see Table 1 in [1] for the comparison of computational costs).

Reviewer comment (Experiments):
Have you tried to set $\delta_k$ as learnable parameter, or is the tuning of the time-subscale subject to the practitioner and model designer?

Author response:
This is an interesting question. Recall that $\delta_k$ is the amount of time that the $k$-th wave in SineNet advances the solution in the latent space. In training SineNet end-to-end to advance time from $u_t\to u_{t+1}$, $\delta_k$ is in fact learnable and likely varies between different examples, referred to as **adaptable temporal resolution**, which we further discuss in the second-to-last paragraph of Section 3.3.

However, the $\delta_k$ are latent and therefore cannot be observed or even explicitly set without obtaining the ground truth solution at times $t+\Delta_t$ and applying deep supervision on each wave during training, where $\Delta_t\in(0,1)$ are chosen dependent on the desired $\delta_k$. Instead, one may parameterize a U-Net as the time derivative $\frac{dh}{d\tau}$ of $h(\tau)$, where $h_\tau$ is the latent of $u_{t+\tau}$. Then, to map from $u_t\to u_{t+1}$, one can integrate the derivative using a numerical integrator. Using this approach, the practitioner maintains the flexibility to determine the $\delta_k$ in selecting the time discretization for the numerical integration.

Alternatively, one may choose to employ an adaptive time-stepping integrator, which selects the $\delta_k$ automatically so as to induce stability in the integration. This is the approach taken by neural ODE [2], and we have added an experiment using their framework to Appendix F.2. We find that although this approach leads to expensive training, it is able to achieve competitive results with SineNet-8. Moreover, as suggested by Reviewer **BgSq**, we also included experiments on the conditional INS task in Appendix F.1 where the models are trained to generalize over variable-sized time steps so that the user can control the prediction time interval, which implicitly modifies $\delta_k$.