GridMix: Exploring Spatial Modulation for Neural Fields in PDE Modeling

Thank you for the valuable comments. Our responses to the comments are as follows.

Reply to W1:

We appreciate the insightful comment and thank you for pointing out these relevant works. We agree that spatial modulation has been explored in various contexts. We have carefully considered the suggested references and incorporated a more detailed comparison in the appendix D. To address your concern regarding novelty, we would like to emphasize that the key distinction of our approach lies in its unique focus on the domain generalization of spatial modulations. This aspect is under-explored in previous research and is particularly crucial for PDE modeling, where adaptability to diverse domains is essential.

Reply to W2 and Q2:

Thank you for raising this important consideration. We have provided a detailed analysis of computational complexity in Appendix F. Specifically, we unroll a trajectory up to T = 40 with a batch size of 1 and report runtime and memory usage in Figure 7. All experiments were conducted on an NVIDIA A100 GPU.

Time Complexity: Our method shows a moderate increase in runtime compared to CORAL. For example, at a resolution of $128 \times 128$, our approach is approximately 1.6× slower than CORAL. This additional overhead arises from steps like spatial modulation and interpolation in our framework. Nonetheless, our method still maintains substantial efficiency compared to traditional numerical methods. The pseudo-spectral method, the numerical baseline used to generate the dataset, requires approximately an order of magnitude more time than our method at the same resolution.

Memory Usage: Our method incurs about 6.5× the memory overhead compared to CORAL, primarily due to additional parameters and intermediate computations required by our enhanced representation. However, the memory cost scales predictably and remains well within the capabilities of modern GPUs for the tested resolutions, ensuring feasibility for larger datasets and higher resolutions in practical applications. As discussed in Section 5, this memory overhead could potentially be mitigated through techniques such as vector quantization[3] and tensor factorization[4], offering a promising direction for future research.

Reply to Q1

Thank you for your insightful question. In our work, we chose to optimize the decoder and Neural ODE models separately, following the practices established in prior work[1]. However, exploring end-to-end training is an intriguing direction, and we acknowledge its potential benefits, as demonstrated in [2].

We attempted end-to-end training in our case but encountered challenges. A key issue is that Neural ODE training can be unstable during the early stages due to its auto-regressive nature, which provides less meaningful supervision for the decoder. Additionally, our current training strategies and network architectures were selected specifically for separate training, and adapting them to an end-to-end framework would require careful consideration and redesign.

Since this was not the primary focus of our work, we did not pursue it further but consider it a promising avenue for future exploration.

Reference

[1] Louis Serrano, Lise Le Boudec, Armand Kassa¨ı Koupa¨ı, Thomas X Wang, Yuan Yin, Jean-No¨el Vittaut, and Patrick Gallinari. Operator learning with neural fields: Tackling pdes on general geometries. Advances in Neural Information Processing Systems, 36, 2023.

[2] David M Knigge, David R Wessels, Riccardo Valperga, Samuele Papa, Jan-Jakob Sonke, Efstratios Gavves, and Erik J Bekkers. Space-time continuous pde forecasting using equivariant neural fields. arXiv preprint arXiv:2406.06660, 2024.

[3] Towaki Takikawa, Alex Evans, Jonathan Tremblay, Thomas M¨uller, Morgan McGuire, Alec Ja- cobson, and Sanja Fidler. Variable bitrate neural fields. In ACM SIGGRAPH 2022 Conference Proceedings, pp. 1–9, 2022.

[4] Anpei Chen, Zexiang Xu, Andreas Geiger, Jingyi Yu, and Hao Su. Tensorf: Tensorial radiance fields. In European conference on computer vision, pp. 333–350. Springer, 2022.

Reply to Q3:

Thank you for the comment. We follow the practice in prior work, CORAL[1], where the same entries are reported as "N.A." The "N.A." values in Table 2 indicate that the corresponding results are "Not Available." Specifically, for the $*Shallow-Water*$ dataset, the data is represented in a 2D spherical grid (latitude and longitude), which is not directly compatible with the linear interpolation method we used, resulting in the "N.A." values.

We hope this clarifies the reason behind the missing entries. Please let us know if further details are needed.

Reply to Q4:

Thank you for your question. In our work, we adhered to the experimental setup and utilized the original datasets provided by Serrano et al.[1] to maintain consistency in our comparisons. We reproduced their experiments by following the training recipes for CORAL and other baselines as detailed in their Appendix B. Our reproduction closely aligned with the reported results in their paper. Consequently, we directly used their reported results as baselines for our comparison. No additional data was generated for these experiments. For clarity and reproducibility, we include detailed dataset preparation steps in Appendix A.1, so our approach remains fully self-contained.

Reply to Q5:

Thank you for highlighting this point. We acknowledge that in certain settings, such as Shallow Water ($\pi=5\%$) and Pipe, MARBLE’s performance was slightly below some baselines. Here is our analysis:

1. Shallow Water ($\pi=5\%$): MARBLE underperforms DINo in the In-t setting but outperforms it in the Out-t setting. This difference primarily stems from their distinct approaches to embedding optimization in the auto-decoder. In DINo, each sample’s embedding is initialized once and optimized throughout training without reset, resulting in modulation space that are finely tuned to the training data. While this strategy can yield highly optimized embeddings, it also increases the risk of overfitting, particularly when the input data is sparse, limiting temporal extrapolation. In contrast, MARBLE follows a meta-learning strategy as outlined in Section 3.2, where each inner-loop re-initializes embeddings from scratch and optimizes them for only three steps. This method, combined with an outer-loop that optimizes the INR weights conditioned on the inner-loop embeddings, promotes smoothness in the modulation space and enhances temporal extrapolation, leading to stronger performance in the Out-t setting.
2. Pipe Dataset: MARBLE performs below GEO-FNO and Factorized-FNO on the Pipe dataset, likely due to limitations of the SIREN backbone used. The Pipe dataset exhibits high frequencies predominantly along the vertical dimension. While SIREN is effective for isotropic frequency distributions, it may struggle with data featuring strong directional anisotropy due to its sensitivity to frequency-related hyperparameters. To address this, future work could explore improved INR designs[2] capable of capturing directional frequency variations.

We have incorporated these analyses into the experiments and discussion sections to provide further clarity. Thank you again for your valuable feedback.

Reply to Q6:

Thank you for your insightful question. In this specific experiment, our goal is to predict the future states (output function) of a physical quantity, given its current state (input function). The data we utilize consists of various initial conditions and time instances, all sampled from a similar underlying distribution. As a result, we can assume that the input and output functions come from the same function space.

Given the relatively homogeneous nature of this dataset, we found that a single modulated INR was sufficient to capture the underlying variations. This simplified architecture allows for efficient training, while still achieving strong performance.

In contrast, the general setup outlined in our method section, involving two modulated INRs, is designed for more complex scenarios where the input and output functions may exhibit distinct characteristics. This approach provides greater flexibility and can be applied to a wider range of tasks.

Reference

[1] Louis Serrano, Lise Le Boudec, Armand Kassa¨ı Koupa¨ı, Thomas X Wang, Yuan Yin, Jean-No¨el Vittaut, and Patrick Gallinari. Operator learning with neural fields: Tackling pdes on general geometries. Advances in Neural Information Processing Systems, 36, 2023.

[2] Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park. Separable physics-informed neural networks. Advances in Neural Information Processing Sys- tems, 36, 2024.

Thank you for the valuable comments. Our responses to the comments are as follows.

Reply to W1:

Thank you for the comment. We appreciate the opportunity to clarify these points.

• To derive the modulation parameter $\phi(x)$ for a given spatial location $x$, we first identify the neighboring grid points surrounding this location (in 2D, this involves four points). We then compute $\phi(x)$ by interpolating the values at these neighboring points, using bilinear interpolation in the 2D case.
• As outlined in Equation 8 of the original submission, we use the latent embedding $z$ to derive the mixture coefficients $c$ through a hypernetwork, which in turn determines the value of $\phi$. For clarity, we have combined Equations 7 and 8 into a single equation.

Reply to W2

Thank you for raising this point. The baseline models differ between the two classes of experiments because each experimental setup evaluates distinct aspects of model performance. We followed prior work[1] to select appropriate baselines for each task.

1. In the dynamics modeling setting, the focus is on assessing temporal forecasting capabilities. For this, we compare our method against:

• Two classical neural operators (DeepONet and FNO),
• A well-designed auto-regressive graph-based model (MP-PDE),
• Two coordinate-based approaches (DINo and CORAL).

2. In the geometry-aware inference setting, the focus shifts to evaluating models' abilities to generalize across varying geometries. Here, we compare with:

• Two regular-grid baselines (FNO and UNet)
• Two improved FNO variants, Geo-FNO and Factorized-FNO, which are specifically designed for enhanced performance on geometry-aware tasks.
• A coordinate-based approach (CORAL)

The selection of baselines is tailored to align with the primary goal of each experiment, ensuring fair and meaningful comparisons. Based on your feedback, we have incorporated the following updates:

1. Added Factorized-FNO as a baseline in the dynamics modeling setting (see Reply to Q1).
2. Provided an analysis explaining why Geo-FNO is unsuitable for irregular-grid settings in dynamics modeling (see Reply to Q2).

We hope these clarifications address your concerns and enhance the understanding of our baseline choices.

Reply to Q1:

Thank you for the suggestion. We have included a comparison with Factorized-FNO (FFNO) in the dynamics modeling setting, as shown in the following table. Specifically, we use a 12-layer Factorized-FNO with mode=16 and width=64. As observed, Factorized-FNO demonstrates improved performance compared to FNO across all settings. However, our method, MARBLE, consistently achieves superior performance. We have included these results in the Appendix E.

Reply to Q2:

Thank you for your suggestion. “FNO + lin. int.” refers to preprocessing irregular grid data by linearly interpolating it onto a regular grid, which is then used as input for the standard FNO. To ensure clarity, we have included this explanation in the table caption.

Regarding Geo-FNO, while it is designed to handle irregular meshes, it is not suitable for this setting due to the following reasons:

1. In our setup, the irregular points are sampled from a regular grid. Transforming these irregular points back into a regular grid, as required by Geo-FNO’s deformation mechanism, disrupts their original structure.
2. In our setup, the same irregular grid is used across all training samples, while a different irregular grid is employed for testing samples. Geo-FNO relies on learnable grid deformation, which is prone to overfitting to the training grid due to the limited diversity of training grids in our case. Consequently, it performs poorly when generalizing to unseen testing grids.

Given these limitations, we sought a more robust alternative. To this end, we have incorporated Factorized-FNO as an enhanced baseline as shown in Reply to Q1. Linear interpolation is used to preprocess the irregular grid for compatibility with this approach. This approach ensures a fair and meaningful comparison while maintaining compatibility with our irregular grid setting

I thank the authors for the detailed responses to my questions and for updating the manuscript. I have raised my score from a 6 to an 8.

Thank you for the valuable comments. Our responses to the comments are as follows.

Reply to W1:

Thank you for your comment. We have included a detailed introduction to classical grid-based methods, including finite difference, finite element and spectral methods, in Appendix D for further context and elaboration. Please let us know if there are any specific aspects you would like us to expand upon further.

Reply to W2 and Q2:

Thank you for highlighting these important points. We have provided a detailed analysis of computational complexity in Appendix F. Specifically, we evaluated the runtime of our method by unrolling a trajectory up to T = 40 with a batch size of 1, as shown in Figure 7. All experiments were conducted on an NVIDIA A100 GPU.

Our method exhibits a moderate increase in runtime compared to CORAL. For instance, at a resolution of $128 \times 128$, it is approximately 1.6× slower than CORAL. This additional computational cost stems from steps such as spatial modulation and interpolation within our framework. Despite this overhead, our method still maintains substantial efficiency compared to traditional numerical methods. For example, the pseudo-spectral method, the numerical baseline used to generate the dataset, requires roughly an order of magnitude more time than our method at the same resolution.

Reply to Q1:

Thank you for your question regarding the grid settings.

1. Grid Shape and Density: The grid we employ is a regularly spaced, uniform grid, consistent across all experiments. The density of this grid refers to the level of resolution (H*W for the 2D case), which determines how finely we sample and represent the spatial field. The choice of uniform density aims to balance computational efficiency with the desired detail level.
2. Discretization Method: Since our approach does not involve traditional PDE solving, we do not use explicit discretization techniques such as finite difference or finite element methods. Instead, the grid is uniformly dense to provide consistent feature sampling, which helps in representing spatial variations for modulated INRs. Each cell in this grid contains a learnable feature, allowing us to implicitly encode spatial information that the model can leverage for scene reconstruction.
3. Handling Complex Geometries: Traditional simulation methods often employ adaptive grids to refine mesh elements around complex geometries and capture fine-grained details. In contrast, in the Geometry-Aware Inference section, we treat each geometry as a deformation function $X_i$, which describes how the geometry is obtained by transforming a shared, uniform reference domain $X$. We model the family of geometry deformation functions using a modulated INR, where $X_i(x) = f_{\theta, \phi_i}(x)$ with $x \in X$. By learning the deformation function $f_{\theta, \phi_i}$ on the uniform grid of the reference domain, we can implicitly represent complex geometric variations. Each geometry $X_i$ is then associated with a grid-based modulation parameter $\phi_i$, which is defined in the reference domain and takes the form of a uniform grid. This approach enables us to efficiently handle diverse geometries without the need for explicit grid adaptation.

Thank you for the valuable comments. Our responses to the comments are as follows.

W1: Some quantitative results, especially those comparing single-resolution grids with global modulation in Table 1, appear counterintuitive. The single-channel GridMix approach yields lower performance in some cases compared to simpler global modulation, which contradicts the initial motivation that spatial features should perform better.

Q1: In Table 1, why single resolution feature grid (Single-channel GridMix) gives worse results than Global Modulation? As the spatial feature representation should give already good results given the initial motivation

Reply to W1 and Q1:

Thank you for this question. Single-channel GridMix can be viewed as a regularized version of vanilla spatial modulation, where the modulation space is deliberately constrained to enhance domain generalization. This limited modulation space, however, imposes constraints on its performance. As shown in Table 1, the vanilla spatial modulation indeed achieves better reconstruction in the training domain $\mathcal{X}_{tr}$ , with an error of 3.95e-5 compared to 1.32e-4 for Global Modulation. This supports the original motivation for spatial modulation’s use in capturing finer details.

However, vanilla spatial modulation struggles to generalize to unseen spatial domains. To mitigate this, Single-channel GridMix constructs spatial feature representations as a linear combination of shared basis functions across all samples. This design implicitly regularizes the modulation space, improving generalization to new spatial domains at the expense of reduced representational capacity. This trade-off ultimately leads to a slightly lower performance on the original domain but increases robustness to the change of spatial domain. To further enhance reconstruction performance while preserving this robustness, we propose Multi-channel GridMix, which demonstrates superior results compared to Global Modulation.

W2: Missing mentioning Factor Fields [4] paper that also works in similar framework and solves similar problems but not in PDE domain.

Reply to W2:

Thank you for pointing out the Factor Fields [4] paper. We appreciate the opportunity to address this. While Factor Fields decomposes a signal into components such as coefficient fields and basis functions, our work performs a similar decomposition within the modulation space rather than the signal space. Additionally, Factor Fields primarily addresses tasks outside the PDE domain, whereas our approach focuses on PDE-related challenges, with an emphasis on the domain generalization of spatial modulations.

We have included a comparison with Factor Fields in the related work section to highlight the differences in application focus and the unique contributions of our method. Please feel free to let us know if further clarifications or expanded analyses would be helpful.

Q2: I dont fully understand the explanation of why more basis functions lead to worse results at some point? As in case there are too many basis, they can be zeroed out with corresponding weights.

Reply to Q2:

Thank you for your question. A possible reason why adding more basis functions leads to worse results at some point is the increased complexity of the latent trajectories. While it’s theoretically possible to zero out redundant basis functions through corresponding weights, in practice, the model often utilizes this additional capacity to refine the reconstruction, resulting in more complex latent trajectories. These trajectories can become overly intricate, potentially making them difficult for the NeuralODE to learn effectively. As shown in Appendix G, the trajectories generated with 64 basis functions fluctuate more than those with 32, which may hinder forecasting performance.

Q3: Figure 3 is hard to read – need legend on the Y axis. It is not clear to me from the footnote what’s going on.

Reply to Q3:

Thank you for your feedback. We’ve refined Figure 3 to improve its readability and have added a legend on the Y-axis to clarify the data presented. Additionally, we revised the footnote for greater clarity on the figure’s context and content.

To summarize, Figure 3 illustrates key findings from Table 1:

1. Spatial modulation significantly reduces reconstruction error on $\mathcal{X}\_{tr}$ but struggles to capture global information and generalize to $\mathcal{X}_{te}$.
2. GridMix mitigates the overfitting issue associated with spatial modulation while achieving superior performance compared to global modulation.

Please let us know if there are any further improvements or clarifications we can provide.