CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs

Thank you for your effort in reviewing our work and your valuable feedback. We appreciate that you recognize the method, featuring the hierarchical query point design and the limited neighborhood, as important to effectively solve PDEs. We are happy to address the mentioned weakness and questions below.

(W1) Except for the experiment on the 2D Navier-Stokes problem (ν = 1e−4), the accuracy and efficiency of CALM-PDE are generally comparable to those of FNO [1] or Geo-FNO[2], while FNO appears to have a simpler architectural design. In this regard, the advantages of CALM-PDE merit further investigation.

We would like to emphasize that we use a curriculum strategy to train the autoregressive models, including FNO and Geo-FNO, which improves the error compared to a fully autoregressive training (see Appendix I, Table 6, line 799). For instance, the curriculum training of FNO improves the error by 39% on the 2D Navier-Stokes $\nu=1e^{-4}$ dataset. Thus, we use strong baselines and incorporate additional techniques to improve them further. In general, CALM-PDE is a viable and competitive alternative to commonly used FNOs and Transformer architectures for solving PDEs. The model demonstrates the effectiveness of convolutional models on irregularly sampled spatial domains and the benefit of sampling the spatial domain via adaptive query points.

(Q1) How is the number of query points in each layer of the encoder and decoder determined? Is there a quantitative method to calculate how many query points are needed?

We determine the number of query points based on the principle that the number of query points, receptive field, and channel dimension must cover the entire spatial domain. Similar to conventional CNNs, a small receptive field with a large stride (e.g., only a few query points) misses information, while a large receptive field with a small stride could have an averaging effect. With this principle, we obtain the number of query points (1024, 256, 16 for the encoder layers), which we use for all 2D experiments. Thus, a selected number of query points works across different problems and does not need individual hyperparameter search.

(Q2) The authors mention that their method does not perform well on PDE problems with many fine details. Did they actually conduct experiments on the Kolmogorov Flow problem? If so, what were the results?

Yes, we have experimented with the Kolmogorov Flow problem from [1] by training only the encoder-decoder model with a self-reconstruction loss (autoencoder) and obtained a relative error of 0.0928 for a simple reconstruction, which is already worse compared to the relative error of 0.8156 for a rollout reported for DilResNet in [1]. Thus, we did not continue the experiment.

References
[1] Z. Li, D. Shu, and A. B. Farimani, “Scalable Transformer for PDE Surrogate Modeling,” arXiv [cs.LG], 2023.

Thank you for your feedback and your questions. Your constructive feedback helped us to improve our submission. We appreciate that you consider our method interesting and helpful. Below, we address your concerns and questions in detail.

(W1) The contribution of the method seems to be the encoding/decoding since the process is a Neural ODE. I think a deeper study of encoding/decoding performances/effects would help the reader clearly identify the benefits of the model.

Yes, the main contribution of the method is the novel discretization-agnostic encoder and decoder based on continuous and adaptive convolution. The latent time stepping can be interpreted as solving a neural ODE. However, there are different ways of parametrizing the neural ODE, and previous work relied on point-wise MLPs [1, 2], while we utilize a Transformer with a combination of vanilla and position-attention [3] to parametrize the neural ODE. We added additional experiments that investigate the scaling of the latent query points and the output query points to provide more insights into the encoder-decoder architecture. Additionally, we add an explanation of the computational costs for the encoder. Please see (Q6) and (W3).

(W2) I would lower the first claim "a new model that solves ... PDEs", because the major contribution relies in the encoding/decoding architectures.

As mentioned in the response for (W1), we agree that the main contribution is the novel encoder-decoder architecture and not the paradigm itself. Thus, we’ve updated the paper and changed it from “a new model” to “a model with a new encoder-decoder mechanism” to better highlight the main contribution.

(W3) One of the arguments for the model being the computational cost, I think a deeper comparison on this aspect with other model would strengthen the argumentation. [...] How performs CALM-PDE in the encoding tasks wrt to other encoding architectures?

We also compare our model against OFormer (linear attention) and LNO (vanilla attention), which are encoder-decoder architectures that utilize attention for encoding and decoding. As shown in Figure 4 and Table 3 on page 9, CALM-PDE achieves a lower inference time and memory consumption compared to these models. Furthermore, we added an explanation to the paper to clarify the computational cost of the models. Let N be an arbitrary number of input points. For simplicity, we only consider one CALM encoder layer. The output of the layer are features for N’ output or query points, where N’ is a constant number. We also have to consider the percentile p, which determines the size of the receptive field and the feature or channel dimension d. This yields that $\mathcal{O}(p \cdot N \cdot N’ \cdot d)$ computations are needed. Models such as AROMA or LNO do not consider a receptive field, which requires $\mathcal{O}(N \cdot N’ \cdot d)$ computations for N’ query tokens. Since p is usually small (we use p=0.01), the complexity is reduced by two orders of magnitude compared to models without a receptive field.

(Q1) What are the differences between LNS/LTS and RO discussed in the related work?

We use the term Reduced Order Modeling (RO) to generally refer to models that reduce the model complexity in some way. We further introduce the terms Latent Neural Surrogates (LNS) and Latent Time Stepping (LTS) as concrete ways of reducing the model complexity. For instance, LNS reduces the model complexity by reducing the spatial resolution, and LTS reduces the model complexity by computing the PDE trajectory completely in the latent space. LNS and LTS can be combined (e.g., AROMA), but they can also be used separately (only LNS for LNO and only LTS for OFormer).

(Q2) Can the number of point N differs between samples or must it be fixed ?

Yes, the number of input points N can be different between samples, which is the case in the cylinder dataset. Due to continuous convolution, each query point can take an arbitrary and changing number of input points.

(Q3) What happens if there are no points in a considered neighboorhood ?

This problem can only happen if the epsilon is fixed, which is not the case for our model. We adopt the mechanism from PiT [3], which computes the epsilon dynamically based on the p-th percentile of the Euclidean distances. This simply means that each query point has at least N*p points as input. If the input points are far away, the epsilon will be increased, and vice versa. Please see Appendix E.1, line 574 for details.

(Q4) Lines 40-42, « mainly utilises CNNs ... or transformers » but no example are given here, I think this sentence should be illustrated. You could add some references [...], this point is only for clarity when reading the paper

Thanks for pointing that out. We changed the corresponding lines and added LE-PDE [4] as an example for a model using CNNs as well as AROMA and UPT for transformer-based models.

(Q5) what makes LNO 3x heavier in terms of memory consumption as FNO?

The key difference that makes LNO heavier in terms of memory consumption and time is the use of attention compared to the use of the fast Fourier Transform (FFT). The encoder and decoder of LNO are based on cross-attention, which makes it expensive. In particular, cross-attention between a fixed number of query tokens (queries in attention) and all input points (values, keys in attention) is computed to compress the input into a smaller latent representation. In the 2D Navier-Stokes experiments, LNO uses 128 query tokens and the input has 4096 input points, which requires $\mathcal{O}(128\cdot4096)$ computations. For decoding, this process is reversed by using the latent tokens as key, value in attention, and the query positions as queries. The latent representation is processed with a Transformer that computes self-attention, which requires $\mathcal{O}(128\cdot128)$ computations. In contrast, FNO makes use of the FFT in each layer with a complexity of $\mathcal{O}(N \cdot log(N))$, which yields $\mathcal{O}(4096 \cdot log(4069))$ for the 2D Navier-Stokes dataset, which is more efficient.

(Q6) How do evolve the performances wrt to query grid size ? in term of computational cost/accuracy?

We’ve conducted an experiment on the 2D Navier-Stokes $\nu=1e^{-4}$ dataset (max. resolution of ground truth is 64x64) where we increase the output query grid size during inference. The numbers are obtained on an A100 40GB GPU with a batch size of 32 and predicting one timestep.

Additionally, we’ve conducted an experiment where we change the latent query grid size (number of latent points). We measure the relative L2 error on the test set and total training time on the 2D Navier-Stokes dataset with $\nu=1e^{-4}$. The bold row shows the configuration used in the paper. The results show that we can expect a saturation effect in the error.

(Q7) Have you considered other architecture for the process ?

Besides a Transformer, we’ve considered a pointwise MLP as the latent processor. However, we decided to use a Transformer instead of an MLP to allow the processor to have a global view, which would not be the case for a pointwise MLP, and to allow latent tokens to dynamically interact with each other with attention.

(Q8) Why not considering AROMA or UPT as baselines?

We've added AROMA as an additional baseline. Please note that we use a different setup (e.g., different rel. L2 error computed on non-normalized data, no in-t/ out-t split, etc.) compared to the AROMA paper.

We've also measured the inference times on the 2D Navier-Stokes $\nu=1e^{-4}$ dataset, and AROMA needs 5.18s, in comparison to CALM-PDE with 0.37s. Reducing the refinement steps to 1 reduces the inference time to 3.81s.

could CALM-PDE handle parametric PDEs?

We did not experiment with parametric PDEs with varying PDE parameters, but CALM-PDE can be easily adapted to handle such cases. For instance, the PDE parameter could be added as an additional input channel, or a more sophisticated conditioning mechanism could be used to incorporate the PDE parameter [5, 6].

References
[1] Z. Li, K. Meidani, and A. B. Farimani, “Transformer for partial differential equations’ operator learning,” arXiv [cs.LG], 2022.
[2] L. Serrano et al., “Operator learning with neural fields: Tackling PDEs on general geometries,” arXiv [cs.LG], 2023.
[3] J. Chen and K. Wu, “Positional knowledge is all you need: Position-induced Transformer (PiT) for operator learning,” arXiv [cs.LG], 2024.
[4] T. Wu, T. Maruyama, and J. Leskovec, “Learning to accelerate Partial Differential Equations via latent global evolution,” arXiv [cs.LG], 2022.
[5] M. Takamoto, F. Alesiani, and M. Niepert, “Learning neural PDE solvers with parameter-guided Channel Attention,” arXiv [cs.LG], 2023.
[6] J. Hagnberger, M. Kalimuthu, D. Musekamp, and M. Niepert, “Vectorized Conditional Neural Fields: A framework for solving time-dependent parametric Partial Differential Equations,” arXiv [cs.LG], 2024.

Thank you for carefully reviewing our work and asking insightful questions, which helped us to improve our submission. We are grateful that you recognize our approach as a strategy with the potential for wide applicability across various neural operator frameworks. We address your concerns and questions separately as follows.

(W1) The results on the 2D Navier--Stokes equation with high Reynolds number seem not strong enough compared to FNO [...]

We developed the architecture for irregularly sampled spatial domains, which are important for many real-world applications, as you already pointed out in strength 3. We included the benchmark on regular grids to demonstrate that our approach is also directly applicable to such settings and yields competitive errors. Thus, CALM-PDE is not specially optimized for regular grids. Since the grid is uniformly structured, the inference time, training time, and memory consumption could be improved by exploiting the uniform structure, which we consider as interesting future work. However, we specifically target PDEs with irregularly sampled spatial domains with our method.

(W2) There could be more hyperparameters [...] to tune compared to other baseline models. And it is unclear how robust the model is to these hyperparameters' change.

The model introduces the softmax temperature, receptive field radius, channel dimension, and number of query points as hyperparameters. The receptive field radius is equivalent to the kernel size in conventional CNNs and the number of query points to downsampling (e.g., stride and pooling) in CNNs. The channel dimension is equivalent to the channel dimension in CNNs. Alternatively, the number of query points can be compared to the query points used in other models, such as AROMA [1] or LNO [2]. Thus, our architecture does not introduce more hyperparameters.

Selecting good hyperparameters for our model does not require an extensive hyperparameter search. As shown in Table 11 in Appendix J.4 (line 856), we use the same temperatures, receptive field sizes (percentiles), and number of query points for all 2D experiments, which demonstrates that the hyperparameters are robust and work across different problems.

Only the temperature is a new hyperparameter that controls the distance weighting within the receptive field. We also use the same temperatures across all experiments, demonstrating its robustness. However, we’ve conducted a hyperparameter study for the temperature, which shows that it is also robust to changes ($T \in [0.1, 5]$ yields small changes in error). We ran the experiments on the 2D Navier-Stokes $\nu=1e^{-4}$ dataset.

The table shows that T=1 yields the lowest errors and that a slightly smaller (e.g., T=0.1) or larger (e.g., T=5) temperature results only in small changes in error. This demonstrates the robustness of the temperature parameter.

(W3) There has been a lot of follow-up work on reducing the memory consumption of the transformer-based methods. [...]

In contrast to (efficient) transformer-based methods, CALM-PDE introduces a fundamentally different and efficient architecture grounded in continuous convolution. The model demonstrates that convolution can be applied to irregularly sampled spatial domains, where usually transformer-based methods dominate [1, 2, 3], and that restricting the receptive field inherently reduces the computational and memory complexity. Thus, CALM-PDE is a competitive alternative to (efficient) transformer-based methods. To address concerns around efficiency, we compare our method with OFormer, which leverages linear attention to improve memory and runtime of the attention mechanism (see Figure 4 and Table 3 on page 9). Importantly, the implementation of CALM-PDE is not heavily optimized for GPUs, similar to the implementations of the baselines, which we consider a fair comparison. However, with optimized implementations, the efficiency of CALM-PDE can be improved further, which we aim to investigate as interesting future work.

(Q1) What values of positions $P$ of latent tokens in Equation (8) can take? How are they initialized and learned?

The positions $P$ of the latent tokens are the learned query positions of the last encoder layer. We did not constrain the positions that the latent tokens can take. Thus, they could potentially take an arbitrary position. For periodic domains, this is not an issue because the positions are corrected and will be periodically mapped to $[0,1]^n$. For non-periodic domains, we noticed that the positions do not move too far away and stay close to the spatial domain. We opt for this design to keep the method simple. However, additional regularization to force the positions to stay in $[0,1]^n$ could be tested for future work. The plots with the title “Latent Position” in Appendix K (pages 27 and 28) show the learned positions $P$.

We have two ways of initializing the query positions. a) They are uniformly sampled from $[0,1]^n$ to achieve a uniform cover of the spatial domain. b) They are sampled from a training mesh to have a better initialization and to capture the underlying mesh structure. We use b) only for the airfoil experiment, where we have one irregular mesh for all samples, and a) for the remaining experiments.

The query points are learned similarly to the weights and biases of a neural network. The query points influence how the input features are weighted. Thus, we can compute the gradient of the loss function with respect to the query positions. The gradient is used to optimize the query positions (i.e., move the query points) with a variant of stochastic gradient descent (we use AdamW [4]).

Details about the initialization and learning of the query positions are explained in the paragraph “Learnable Query Points" in line 132 in the paper.

(Q2) How to determine the appropriate dimensionality of the latent variables?

The dimensionality of the latent variables is defined by the number of latent variables or query points and their channel dimension. We determine the dimensionality of the latent variables based on the principle that the number of query points, receptive field, and channel dimension must cover the entire spatial domain. Similar to conventional CNNs, a small receptive field with a large stride (e.g., only a few query points) misses information, while a large receptive field with a small stride could have an averaging effect. With this principle, we obtain the number of query points (1024, 256, 16 for the encoder layers), which we use for all 2D experiments. Thus, a fixed number of query points works across different problems.

Will we observe a saturation effect or diminishing returns in prediction accuracy as the latent dimension increases?

We’ve conducted an experiment where we change the latent dimensionality by changing the number of latent points, and the results show that we will observe a saturation effect. We measure the relative L2 error on the test set and total training time on the 2D Navier-Stokes dataset with $\nu=1e^{-4}$. The bold row shows the configuration used in the paper.

Should we use the same hyperparameters for different latent dimensions?

The number of query points, receptive field, and channel dimension are related to each other. Thus, it is recommended to adapt them jointly. When decreasing the number of latent points, the receptive field and channel dimension have to be increased and vice versa.

(Q3) [...] That is, can it generalize across varying domain shapes or boundary configurations as part of the input space?

Yes, the model can also handle varying domain shapes and boundary conditions. The cylinder dataset covers exactly this case. Each sample in the train and test set has a different domain shape and boundary condition (i.e., different diameter and position of the cylinder). The results show that CALM-PDE can also generalize across different domain shapes and boundary conditions.

We have also conducted an additional experiment on the time-independent airfoil and elasticity datasets from Geo-FNO [5]. The task is to map an input geometry (e.g., changing airfoil shape) to a physical quantity. We follow the setup from the LNO [2] paper and take the errors for the baselines from their paper. The results also demonstrate that CALM-PDE can generalize across complex, varying geometries.

The results show that CALM-PDE outperforms 6 out of 8 baselines and that it ranks among the top models alongside LNO [2] and Transolver [6].

References
[1] AROMA: Preserving spatial structure for latent PDE modeling with local neural fields
[2] Latent Neural Operator for solving forward and inverse PDE problems
[3] Transformer for partial differential equations’ operator learning
[4] Decoupled weight decay regularization
[5] Fourier neural operator with learned deformations for PDEs on general geometries
[6] Transolver: A fast transformer solver for PDEs on general geometries

Thank you for your insightful feedback and thoughtful questions, which helped us to enhance our submission. We appreciate acknowledging the versatility of our approach in handling both regular and irregular grids effectively. We've carefully reviewed each of your points and provide our responses below.

(W1) The accuracy improvements are not universal. [...] Other papers in this field compare to various other irregularly spaced datasets such as elasticity, plasticity etc. See the LNO paper for example.

As you already pointed out with strength 1, CALM-PDE works effectively for both regular and irregular grids. We would like to emphasize that we do not claim CALM-PDE to be the best model across all problems, as no single model can perform best on all PDE tasks. Rather, our goal is to demonstrate that CALM-PDE is a strong and competitive alternative, particularly in challenging settings such as irregular spatial domains. Indeed, CALM-PDE outperforms all baselines in 3 out of 6 benchmarks and achieves second-best errors in 2 out of 6 benchmarks. Even on the remaining benchmark, CALM-PDE outperforms two transformer-based baselines, reducing the error by 58%. Overall, these results provide evidence that CALM-PDE is a viable and competitive alternative to commonly used FNOs and Transformer architectures for solving PDEs. Notably, the experiments demonstrate the effectiveness of convolutional models on irregularly sampled spatial domains and the benefit of sampling the spatial domain via adaptive query points.

Moreover, we use strong baselines and incorporate additional techniques to further improve them. For instance, we train the autoregressive models (e.g., FNO and Geo-FNO) with a curriculum strategy, which significantly reduces their errors by 39% compared to full autoregressive training (see Table 6 in Appendix I, line 799). Taken together, these results position CALM-PDE as a novel and competitive approach for neural PDE solving.

Following your suggestion, we’ve incorporated the airfoil and elasticity datasets into our analysis. While these datasets are time-independent, our primary focus remains on time-dependent PDEs. We follow the benchmark setup of the LNO paper and provide the baseline errors reported in that work (see table below).

The results show that CALM-PDE outperforms 6 out of 8 baselines and ranks among the top models alongside LNO and Transolver, demonstrating strong performance on both time-dependent and time-independent problems. While Geo-FNO underperforms in this setting, it remains among the top models on time-dependent tasks, highlighting that model performance can vary with problem type. We do not claim universal superiority, but CALM-PDE offers a competitive, convolution-based alternative to Fourier and transformer models.

(W2) Some important baselines such as factorized FNO (FFNO) and transolver, which are quite competitive in this domain are missing.

We’ve added F-FNO as an additional baseline and benchmarked it on the 2D Navier-Stokes datasets as well as the irregularly sampled cylinder dataset. Please see the following table. For comparison with Transolver and a further comparison with F-FNO, we added the airfoil and elasticity datasets. Please see the table in our response to (W1).

(W3) By compressing the solution into a smaller latent space, the method inevitably discards some high-frequency or fine-grained information.

Compressing the solution into a smaller latent space inevitably discards some high-frequency details, a trade-off we accept for improved efficiency. We acknowledge this limitation explicitly in the paper (see line 321 in the Limitations section), and it is a common challenge in reduced-order models.

(Q1) Can the authors show a comparison of performance and efficiency to Factorized FNO? [...] Given its relevance and efficiency, particularly on similar problems, it seems like an important missing baseline.

We added F-FNO as an additional baseline. Please see our response for (W1) and (W2).

(Q2) I also think another recent architecture like Transolver should be discussed and compared to. The best comparisons would be on the irregular grid problems like airfoil, elasticity, plasticity or pipe problems commonly used as baselines for these problems.

We benchmarked our model on the mentioned airfoil and elasticity datasets and included the errors for Transolver. Please see our response for (W1).

[...] I do believe that broader benchmarking [...] is needed to more convincingly demonstrate the method’s advantages in terms of both accuracy and efficiency.

We already compare our method against many recent architectures. For instance, we compare our model against recent state-of-the-art methods, including PiT (ICML 2024) [7], LNO (NeurIPS 2024), and AROMA (NeurIPS 2024) [3]. We also include models such as FNO/Geo-FNO (ICLR 2021 and JMLR 2023) and OFormer (TMLR 2023), since they are commonly used and still achieve state-of-the-art performance. With the new experiments, we include a comparison against F-FNO (ICLR 2023) on time-dependent and time-independent problems as well as a comparison with Transolver (ICML 2024) on time-independent problems.

Furthermore, we compare our method on 8 different datasets and problems, including time-dependent and time-independent problems. The used datasets are commonly used in the community to demonstrate the effectiveness of new models (e.g., see AROMA paper).

Thus, we have benchmarked CALM-PDE on a broad set of benchmark problems and compared it against a broad set of state-of-the-art models. The experiments show that the model outperforms or achieves competitive errors on these versatile benchmark problems, which demonstrates the effectiveness of our new approach.

References
[1] Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger, D., & Niepert, M. (2022). PDEBENCH: An extensive benchmark for Scientific machine learning. In arXiv [cs.LG]. http://arxiv.org/abs/2210.07182
[2] Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). Fourier neural operator for parametric partial differential equations. In arXiv [cs.LG]. http://arxiv.org/abs/2010.08895
[3] Wu, T., Maruyama, T., & Leskovec, J. (2022). Learning to accelerate Partial Differential Equations via latent global evolution. In arXiv [cs.LG]. http://arxiv.org/abs/2206.07681
[4] Serrano, L., Wang, T. X., Naour, E. L., Vittaut, J.-N., & Gallinari, P. (2024). AROMA: Preserving spatial structure for latent PDE modeling with local neural fields. In arXiv [cs.LG]. http://arxiv.org/abs/2406.02176
[5] Lippe, P., Veeling, B. S., Perdikaris, P., Turner, R. E., & Brandstetter, J. (2023). PDE-Refiner: Achieving accurate long rollouts with neural PDE solvers. In arXiv [cs.LG]. http://arxiv.org/abs/2308.05732
[6] Oommen, V., Bora, A., Zhang, Z., & Karniadakis, G. E. (2024). Integrating neural operators with diffusion models improves spectral representation in turbulence modeling. In arXiv [cs.LG]. http://arxiv.org/abs/2409.08477
[7] Chen, J., & Wu, K. (2024). Positional knowledge is all you need: Position-induced Transformer (PiT) for operator learning. In arXiv [cs.LG]. http://arxiv.org/abs/2405.09285