Geometry Aware Operator Transformer as an efficient and accurate neural surrogate for PDEs on arbitrary domains

We start by thanking the reviewer for their careful reading of our paper, their appreciation of its merits and their suggestions for improving it. We provide a detailed response to their points of concern below.

[1.] Regarding novelty of our model, while the reviewer is correct that the ViT processor in GAOT is not novel nor is the encode-process-decode paradigm, we reiterate from l41-55 (Main Text) and Fig. 1 that there is an accuracy-efficiency trade-off for Neural PDE surrogates and accurate GNN type models (RIGNO is state of the art in this family) are not efficient while efficient models (GINO) are not as accurate. So, our main rationale for this paper was to provide a model that could be efficient and accurate at the same time. While combining available elements, we did need to come up with 3 key novel features to do this: i) multiscale attentional GNO for encoding/decoding ii) geometric embeddings based on statistical descriptors to enhance geometric inputs to the model and iii) several tricks allowing for an optimized implementation (see pt. [3] below). To the best of our knowledge, these are original contributions and they clearly enhance model performance. See Tab. F.5 and reply pt [8] below, for results on how multiscale information in the encoder/decoder significantly improves upon a single-scale version, Similarly, see Tab F.4. to observe how statistical geometric embeddings improve accuracy over just inputting coordinates and SDFs. Finally, the role of efficient implementation is described below in pt [3].

[2.] We acknowledge the reviewer's point that related work is rather succinct, primarily due to page limits for the main text. Following the reviewer's suggestion, we are happy to elaborate on the related work section in a camera-ready version (CRV) esp. SM, if accepted. In particular, we thank the reviewer for bringing Refs. [1,3] to our attention and we will gladly add them to related work. Regarding Transolver ++, we will certainly add it with the context described in our reply pt [6.] to reviewer diAB above.

[3.] To answer the reviewer's pertinent question on why is GAOT faster than GINO ?, we ran GINO (6.07 M params) for a single sample of the NACA0012 dataset on a NVIDIA-4090 GPU to obtain the following breakdown of compute time in ms for GINO: Encoder (Graph building: 1.984, GNO: 2.793, Total: 4.777), Processor (5-Layer FNO): 5.89, Decoder (Graph Build: 1.984, GNO: 1.140, Total: 3.124). Total GINO run time per sample = 13.791 ms. As it happens, GINO does not support batch processing when the underlying graphs change in a batch (as in general for PDEs on arbitrary domains). Thus to process a batch of 32, GINO has to loop over the entire batch, leading to a runtime of 441.32 ms ! Thus, we observe 2 computational bottlenecks here: i) Graph building in the encoder/decoder is as (or more) expensive than the GNO module and ii) lack of batch processing severely impedes computational throughput. Our challenge was to overcome these bottlenecks.

To this end, we came up with the following tricks: a) we take graph building out of the GAOT model evaluation loop to reduce computational wastage by precomputing the graphs and storing them as cache files and loading with efficient data loaders and b) batch processing the transformer processor while looping over the encoder/decoder sequentially. Please see our reply pt [4] to reviewer diAB above for further details on these tricks.

With these implementation tricks, we have the following compute breakdown for a 5.62M GAOT (same setting as in GINO) for a batch size of 1: Encoder (AGNO: 2.296, GeoEmb: 3.412, Total: 5.708), 5-Layer ViT processor with patch size 2: 10.1, Decoder (AGNO:1.356, GeoEmb:1.865, Total: 3.221). Total GAOT run time per sample=19.029 ms. Hence, GAOT is slower than GINO for a batch size of 1 as the multi-scale GNO and GeoEmb embeddings add cost although Graph building has been taken out of the loop. However, the results for a batch size of 32 are: Encoder (AGNO+GeoEmb) = 126 ms, Transformer processor = 30.1 ms, Decoder (AGNO+GeoEmb) = 107 ms, leading to an overall batch compute time of 263.1 ms, which is 178 ms faster than GINO !!. The payoff happens in the batch processing of the transformer as this module only increased from 10.1 ms (batch size 1) to 30.1 ms (batch size = 32). It was precisely this novel trick that allows GAOT to be more efficient than GINO, while being much more accurate. As a final trick for very large (particularly 3-D) graphs, we also introduce an edge masking procedure to further reduce memory bottlenecks in the encoder/decoder (l214). We will add this quantitative computational breakdown in the CRV to better inform the readers on how our model achieves computational efficiency.

[4.] As total number of training samples in each dataset is exactly the same for all models, training time per epoch = (Total Number of Samples)/Training throughput, scales exactly as the training throughput for each dataset. Hence, GOAT is 1.6 times faster to train than GINO from Tab E.1. We can also provide the exact training times for each model in a CRV, if accepted.

[5.] Regarding the reviewer's question about irregular queries in time, our model provides the current time t and lead time $\tau$ as inputs (l191), allowing it to be evaluated at any query point in time and at any query point in space (neural field property). To demonstrate this feature, we consider the NS-PwC dataset, where we had used training data at time stamps $t=0.0,0.1,0.2,0.3,0.4.0.5,0.6,0.7$. At test time, we output the velocity fields at times $t=0.35$ and $t=0.45$, which are not included in the training set to obtain relative % errors of 1.4 and 1.91, respectively. This interpolates very nicely with errors of 1.2 ($t=0.3$), 1.65 ($t=0.4$) and 2.2 ($t=0.5$). We thank this reviewer for this suggestion and will add the results in a CRV.

[6.] Perhaps l203 was not clearly articulating our point, which was to say that we would like the compute of the processor to scale very favorably with increasing number of layers as we clearly show in the argument presented in pt [3] above. Clearly the encoder-decoder are necessary to compress information but we want them to be as lightweight computationally as possible as they don't scale as well as the processors. We will adjust this sentence accordingly in a CRV, if accepted.

[7.] The reviewer's aim to understand the precise role of encoder/decoder vs. processor in our model is very much appreciated. We had infact presented results on a Cartesian grid in SM Tab E.2, where we compared GOAT to standard Cartesian grids models such FNO, CNO and the SWIN transformer based scOT, to observe that GAOT was clearly more (or as) accurate as the dedicated cartesian grid architectures on problems on Cartesian grids. Following the reviewer's excellent suggestion, we ablated GAOT further by just running the ViT processor (without the encoder/decoder) on these Cartesian datasets to obtain the following relative % errors: NS-Gauss (ViT (3.16) vs GAOT (2.29)), NS-PwC (ViT (3.89) vs. GAOT (1.23)), NS-SL (ViT (0.73) vs. GAOT (0.98)), NS-SVS (ViT (0.4) vs. GAOT (0.46)), CE-Gauss (ViT (6.81) vs. GAOT (5.28)), Wave-Layer (ViT (5.48) vs. GAOT (5.40)), Poisson Gauss (ViT (2.47) vs. GAOT (0.79)). From these results, it is clear that in some cases, ViT is comparable to or slightly better than GAOT, indicating that the power of GAOT in these cases stemmed from a good processor. However in many more cases, GAOT is significantly superior to ViT also showing that the encoder/decoder contribute significantly to expressivity and together, this combination achieves SOTA performance. This interesting experiment clearly delineates the relative contributions of encoder/decoder vs. processor and we thank the reviewer again for asking us to perform it. Needless to say, we will include it in a CRV, if accepted.

[8.] To answer the reviewer's interesting question about if the GAOT encoder-decoder can work as autoencoders (AE), we consider the output distribution of the NACA0012 compressible flow dataset. As shown in SM Fig. G.4, second column from left, this is a challenging distribution to learn as it contains sharp shocks separated by smooth regions in the flow and there is considerable variation when the increasing Mach number transitions the flow from sub-sonic to supersonic. Thus, it is worthwhile to see how the MAGNO + geometric embeddings based encoder-decoder of GAOT fairs as an AE. As a baseline, we chose GINO with its GNO encoder/decoder as the only point of difference now is the multi-scale attentional features in MAGNO compared to the single-scale GNO and the geometric embeddings. The resulting relative L1 reconstruction errors are 0.08 (GAOT) vs 0,11 (GINO) showing that the MAGNO encoder-decoder+ geometric embeddings are clearly more powerful than a vanilla GNO. We had already shown the advantage of MAGNO in SM Tab F.5 and geometric embeddings in SM Tab F.4 and this experiments further reinforces these advantages. We thank the reviewer for this interesting suggestion and will be happy to add this point to a CRV.

[9.] We don't see any empirical limitations with GAOT given its very broad scope, scalability and excellent performance. However,, as we wrote on l364, we have not shown any rigorous theoretical results such as an universal approximation property etc. This is a limitation that we hope to address in future work.

We take this opportunity to thank the reviewer again and hope that we have satisfactorily addressed their concerns, particularly regarding the computational efficiency of GAOT and the contribution of individual elements (encoder/decoder vs. processor). If so, we request the reviewer to kindly upgrade their assessment accordingly, while remaining at your disposal to answer any further questions about our paper.

I thank the authors for their answer.

• Regarding [1.] : Thanks for the updated results. These should be consistently reported in the manuscript to avoid any misleading for the reader. In my opinion, inference time should be reported in an end-to-end manner, reflecting the realistic scenario of deploying the model on new samples. The detailed inference time is also interesting as it illustrates some of the claims of the paper.
• Regarding limitations, thanks for your answer. I think the section l356-365 reads more like a discussion of future work rather than a clear articulation of the method’s limitations, which is why I raised the question.

To conclude, with the suggestions and clarifications provided in the rebuttal, I will raise my score to 4. I am particularly refering to points 1, 2, 3, 9, which contribute to an honest and rigorous comparison of the proposed method with existing ones (in the related work as well as in the experiments). Additionally, the new experiments presented (e.g., points 5 and 7) help illustrate and support some of the properties and claims of the model.

We start by thanking the reviewer for their careful reading of our paper, their appreciation of its merits and their suggestions for improving it. We provide a detailed response to their points of concern below.

[1.] Regarding the reviewer's point about clarity of exposition, we would be happy to significantly elaborate on the overview section (l102-109) in the main text in a camera-ready version (CRV), if accepted. This will allow us to better inform the reader on how our model is designed.

[2.] We will add a reference to CVIT paper in a CRV as it is indeed relevant. We agree with the reviewer that the multiscale encoder-decoder and statistical geometric embeddings are our main novel contributions at the level of model design.

[3.] On the reviewer's question about different choices of the latent grid, we would like to point to SM F.2 (l1344-1358), in particular Tab. F.2, where we have precisely explored this question. From this table, we observe that reducing the latent grid size does affect the overall accuracy to a minor extent, but not in any significant manner, particularly for the correct patch sizes. This is consistent with our heuristic that the main point of the latent grid should be to provide adequate coverage of the physical domain. As long as it holds, even coarser latent grids are admissible. Summarizing, GAOT seems to be very robust to varying latent grid size, radii of aggregation and patch size. We will highlight this point further in a CRV, if accepted.

[4.] The reviewer is absolutely spot on questioning the computational costs of GINO. To further quantify your concerns, we ran GINO (6.07 M params) for the NACA0012 dataset on a NVIDIA-4090 GPU to obtain the following breakdown of compute time in ms for GINO for a single sample: Encoder (Graph building: 1.984, GNO: 2.793, Total: 4.777), Processor (5-Layer FNO): 5.89, Decoder (Graph Build: 1.984, GNO: 1.140, Total: 3.124). Total GINO run time per sample = 13.791 ms. However, as the reviewer very astutely pointed out, GINO does not support batch processing when the underlying graphs change in a batch (as in general for PDEs on arbitrary domains). Thus to process a batch of 32, GINO has to loop over the entire batch, leading to a runtime of 441.32 ms ! Thus, we observe 2 computational bottlenecks here: i) Graph building in the encoder/decoder is as (or more) expensive than the GNO module and ii) lack of batch processing severely impedes computational throughput. Our challenge was to overcome these bottlenecks.

To this end, we came up with the following tricks: First, we take graph building out of the GAOT model evaluation loop to reduce computational wastage by precomputing the graphs and storing them as cache files and loading with efficient data loaders. Next to enable batch processing, initially we explored vectorizing the encoder and decoder operations. However, this approach proved impractical for two key reasons. First, the grid sizes within a batch may exhibit significant variations, complicating efficient vectorization. Second, for relatively large graphs, vectorization introduces an operation that triggers a transient spike in memory consumption, which can cause the overflow of GPU memory. Consequently, straightforward vectorization is still neither memory-efficient nor scalable. To address this, we adopted a sequential processing strategy for the encoder and decoder while batch processing the transformer. Importantly, this only introduces a lightweight loop whose overhead is negligible compared to the transformer processor, ensuring that it does not become a computational bottleneck when the processor is scaled.

With these implementation tricks, we have the following compute breakdown for a 5.62M GAOT (same setting as in GINO) for a batch size of 1: Encoder (AGNO: 2.296, GeoEmb: 3.412, Total: 5.708), 5-Layer ViT processor with patch size 2: 10.1, Decoder (AGNO:1.356, GeoEmb:1.865, Total: 3.221). Total GAOT run time per sample=19.029 ms. Thus, as the reviewer anticipated GAOT is slower than GINO for a batch size of 1 as the multi-scale GNO and GeoEmb embeddings add cost although Graph building has been taken out of the loop. However, the results for a batch size of 32 are: Encoder (AGNO+GeoEmb) = 126 ms, Transformer processor = 30.1 ms, Decoder (AGNO+GeoEmb) = 107 ms, leading to an overall batch compute time of 263.1 ms, which is 178 ms faster than GINO !!. The payoff happens in the batch processing of the transformer as this module only increased from 10.1 ms (batch size 1) to 30.1 ms (batch size = 32). It was precisely this novel trick that allows GAOT to be more efficient than GINO, while being much more accurate.

Moreover, this set of implementation tricks become even more impactful when the input size and model size is scaled up. As this happens, the memory bottleneck in the encoder/decoder and the compute usage in the processor grows proportionately. Consequently, moving graph building outside the encoder/decoder and batch processing the transformer processor makes GAOT scale significantly better than GINO in both input size and model size as we observe from Figs. 3a and 3b.

As a final trick for very large (particularly 3-D) graphs, we also introduce an edge masking procedure to further reduce memory bottlenecks in the encoder/decoder (l214).

Following the reviewer's excellent suggestion, we will elaborate on this point in a CRV and provide the above detailed computational breakdown such that readers can better appreciate the effort we put into optimizing our implementation for efficiency and scalability.

[5.] As the reviewer correctly claims, UPT is a scalable model. Our problem with UPT is that it lacked stability while training, making convergence very difficult, if not impossible. Despite considerable efforts, we could not make UPT (particularly for larger model sizes) converge for many for our time-dependent datasets. Hence, we did not present these results in the main text. However, we obtained the following relative % errors for the time-dependent datasets of Tab. 1 : NS-Gauss (did not converge (UPT) vs. 2.91 (GAOT)), NS-PwC (did not converge (UPT) vs. 1.5 (GAOT)), NS-SL (51.5 (UPT) vs. 1.21 (GAOT)), NS-SVS (4.20 (UPT) vs. 0.46 (GAOT)), CE-Gauss (64.2 (UPT) vs. 6.4 (GAOT)), CE-RP (26.8 (UPT) vs. 5.97 (GAOT)), Wave-Layer (19.6 (UPT) vs. 5.78 (GAOT)), Wave-C-Sines (12.7 (UPT) vs. 4.65 (GAOT)). Together with the time-independent results from Tab. 1, we see that GOAT is often one to two orders of magnitude more accurate than UPT. We are happy to add these UPT results in a CRV, if accepted.

[6.] Regarding Transolver, we found this model to be very slow computationally. In our understanding, this is because the Transolver model, although it also adopts an encoder–processor–decoder architecture, this architecture is applied at the level of a single Transolver block and there is no sequential processing of information in the latent space. As a result, when stacking many Transolver blocks, the model repeatedly executes both the encoder and the decoder. Since the encoder and decoder each time need to establish mappings between the input grid size and the latent tokens, this leads to a substantial amount of computation (as input grids are typically large) and significant memory overhead. Moreover, the computations in the encoder and decoder cannot fully exploit the hardware’s throughput capabilities, and the large memory footprint forces us to use very small batch sizes as soon as the input grid size becomes slightly larger. Consequently, the model struggles to achieve high throughput as we observe from Figs 3a and 3b. Using the original setup described in the Transolver paper, we found it nearly infeasible to train time-dependent datasets (each with around 20,000 samples) in an all-to-all manner on our hardware.

Transolver++ was published concurrently with the submission of our article, preventing us from using it as a baseline. Although Transolver++ also benchmarks on the Drivaernet++ dataset, their model is trained with only 200 of the total of 8000 samples, making it impossible for us to compare with their results as more data will surely lead to more accuracy. Instead, we have compared with models that sit on top of the Drivaernet++ leaderboard to emerge as the best performing model for this entire dataset. Nevertheless, we do agree with the reviewer that Transolver++ is a contribution to neural PDE surrogate efficiency and we will refer it in a CRV, if accepted.

We take this opportunity to thank the reviewer again and hope that we have satisfactorily addressed their concerns, particularly regarding the computational efficiency of GAOT and its comparison with UPT and Transolver. If so, we request the reviewer to kindly upgrade their assessment accordingly, while remaining at your disposal to answer any further questions about our paper.

[4.] Given the above clarification about novelty of implementation, we would like to take this opportunity to reiterate the key elements of novelty in our paper.

In our opinion, we did not posit the novelty in implementation as our key original contribution in the main text but as one of the many novel contributions (l68-69). We still claim that the originality in our model design, concretized in terms of multiscale encoders/decoders and the use of statistical geometric embeddings contributes significantly to the performance of GAOT and have not been used in the literature to the best of our knowledge.

The role of these design elements has been clearly demonstrated with ablation studies in SM Table F.5 (for multiscale encoders/decoders) and Table F.4 (for statistical embeddings). Following the suggestion of another reviewer, we ablated GAOT further by just running the ViT processor (without the encoder/decoder of GAOT) on the Regular grid datasets from Table E.2 where ViT can be run without any need of encoding/decoding, to obtain the following relative % errors: NS-Gauss (ViT (3.16) vs GAOT (2.29)), NS-PwC (ViT (3.89) vs. GAOT (1.23)), NS-SL (ViT (0.73) vs. GAOT (0.98)), NS-SVS (ViT (0.4) vs. GAOT (0.46)), CE-Gauss (ViT (6.81) vs. GAOT (5.28)), Wave-Layer (ViT (5.48) vs. GAOT (5.40)), Poisson Gauss (ViT (2.47) vs. GAOT (0.79)). From these results, it is clear that in some cases, ViT is comparable to or slightly better than GAOT, indicating that the power of GAOT in these cases stemmed from a good processor. However in many more cases, GAOT is significantly superior to ViT also showing that our novel design of the encoder/decoder contribute significantly to expressivity and together, this combination achieves SOTA performance, even on regular grids, let alone on arbitrary domains.

Summarizing, we believe that the design novelties in GAOT are more important as contributions than its efficient implementation and these underpin its wide-ranging SOTA performance, which we have demonstrated over a wide range of PDE solution operators on different types of domains, including 3D industrial use cases.

We sincerely hope that this (admittedly long) discussion has further clarified the reviewer's concern about the relative importance of contributions of our paper and will play a role in their final assessment of our paper. We thank the reviewer again for providing us with a very detailed feedback and with multiple suggestions that will help us improve the quality and clarity of a CRV, if accepted.

We start by thanking the reviewer for your comments and sincerely appreciate the possibility to engage in further discussion. We thoroughly respect the reviewer's opinion and score but take this opportunity to reply to your comment as it pertains to the very core of our contributions, particularly in terms of novelty. We take the liberty to divide this response into two parts, the first one focussed on computational efficiency and the second on novelty of design. We apologize in advance for the long response and request the reviewer's patience in reading it.

[1.] We agree that applying our implementation tricks to GINO is going to increase its speed and possibly make it faster than GAOT although it is not the case with current implementations. Moreover, we did not find these tricks elsewhere in the literature.

[2.] However, what is undeniable is the striking gain in accuracy of GAOT over GINO. This is clearly demonstrated in Table 1 where we have compared GAOT and GINO over 15 PDE learning tasks to find that the relative error with GAOT is, on an average, 3.27x lower than GINO. Furthermore, GAOT is more accurate than GINO in every single benchmark that we tested on. These include not only datasets generated by us (compressible flows past objects and Poisson-C-Sines), but also publicly available datasets released in other papers, including the elasticity benchmark which was proposed in Ref. [26] by many of the authors of the GINO paper. On this dataset, GAOT's error is 3.26x times lower than GINO's, which is exactly the average gain in accuracy of GAOT over GINO. Even if someone makes GINO faster than GAOT, it is very unlikely that the significant gains in accuracy with GAOT can be compensated by a proportionate gain in computational efficiency of GINO over GAOT.

[3.] Moreover, we still have degrees of freedom to further increase the computational efficiency of GAOT. One of them is to increase the patch size of the VIT processor. The current default patch size is 2 (SM Table B.2) and all the quantitative metrics that we provided about GAOT (GAOT-2 hereon) are with this patch size. However, we have also considered a version of GAOT with a patch size of 8, calling it GAOT-8. The corresponding scaling plots comparing training throughput and peak memory to input grid size and model size are provided in Figure E.1. To provide some concrete numbers, in the training time example that we detailed earlier, while the encoder and decoder runtimes with GAOT-8 are identical as that of GAOT-2, the processor is faster. For a batch size of 1, it takes 8.2 ms (compared to 10.1 ms for GAOT-2), but for a batch size of 32, the ViT processor of GAOT-8 takes only 9.8 ms to run (compared to 30.1 ms for GAOT-2), so total runtime for GAOT-8 is 20ms faster than GAOT-2 in this case.

However, as you have pointed out, it is imperative to fairly compare GAOT and GINO for computational efficiency. As GOAT-8 is our fastest model, we do this comparison in terms of the training throughout for a dataset (Poisson-Gauss) where the underlying graphs do not change in a batch and GINO supports batch processing in this case. We have also implemented storing graphs on Cache for GINO and taking graph building out of the runtime loop. So, this is a totally fair apples to apples comparison. We observe that GAOT-8 has a training throughput of 957 samples/sec compared to GINO's throughput of 753 samples/sec in this setting.

Hence, we observe that GAOT-8 is clearly more computationally efficient than GINO in this metric, without even resorting to our implementation tricks. Of course, the key question remains: is GAOT-8 accurate enough ? We have answered this question in the ablation study of SM Table F.2 where we report the following relative errors for the Elasticity benchmark: GAOT-2 (1.80), GAOT-8 (1.60) and GINO (4.38) and for the Poisson-Gauss benchmark (no change in underlying graphs): GAOT-2 (1.05), GAOT-8 (1.65) and GINO (7.57). Note that, as this was an ablation study, multiscale radii in the MAGNO encoder/decoder were not used here. These numbers clearly show that the larger patch size in GAOT still leads to a model significantly more accurate than GINO, while being more efficient, without resorting to any novelties in implementation. Viewed in this light, our implementation tricks were aimed at making our most expressive model competitive with the fastest models around, while still retaining the superior accuracy of GAOT.

(See the second part of the response below)

We start by thanking the reviewer for their careful reading, appreciation of the merits of our paper and their comments and suggestions for improving it. Below we address their points of concern.

[1.] Regarding the reviewer's question about neural operators, our claim is based on the fact that GOAT is endowed with a neural field property (l173-175) that allows it to process inputs on arbitrary point clouds and provide outputs on any point in the underlying domain. This enables GAOT to possess a key attribute for a neural operator, i.e., the ability to process inputs and outputs at arbitrary underlying grid resolutions. Moreover, works such as Refs. [22,4] argue that this attribute does not suffice for a neural operator and it should possess some (approximate) resolution invariance for genuine operator learning. We test this attribute for GAOT in SM E.5. In this context, Fig E.5 clearly demonstrates how GAOT can generalize seamlessly to test resolutions, different from the training resolution. Thus, our labelling of GAOT as a neural operator is justified in this sense.

[2.] We will gladly increase the size of the figures in a camera-ready version (CRV), if accepted. This will be aided by the additional page that the CRV is allowed to have.

We take this opportunity to thank the reviewer again for their constructive comments and suggestions.

We start by thanking the reviewer for their comments and suggestions. Below, we provide a detailed response to all of the reviewer's points of criticism.

[1.] Regarding the point about reasons behind formulations: we start by reiterating the main motivation for our paper: as clearly stated in the introduction (l41-50 and Fig.1), we observed an accuracy-efficiency tradeoff for all existing ML models for PDEs on arbitrary domains. In particular, accurate models such as GNNs (RIGNO) were not efficient and efficient models such as GINO were not as accurate. This motivates the design of a model that can be both efficient and accurate at the same time -- this is main premise of our paper. To do so, we observe that end-to-end GNNs such as RIGNO involve repeated sparse computations and memory accesses (l204-205) and will not necessarily be efficient and scalable on modern GPU hardware. Hence, we would like to minimize sparse computations. To this end, we focus on a paradigm, where the inputs are encoded to a latent grid, where they can be batch processed efficiently with transformers and then decoded to the desired output. Although our processor is a standard ViT, the encoder and decoder are designed to enhance accuracy by adding 2 novel features: multiscale attentional GNO and geometric embeddings based on statistical descriptors.

The reason behind a multiscale GNO encoder-decoder lies in the fact that solutions of PDEs contain a whole range of spatial scales (SM l764). Almost all the benchmarks that we consider in Tab. 1 contain multiscale solutions (see SM Eq D.5 for an explicit example in the Poisson-C-Sines dataset). Hence, adding multiscale features to the encoder-decoder equips the model with the ability to sense different perceptive fields, making it a reasonable choice in this regard. This is also justified by the ablation study presented in SM F.5 (Tab F.5), where we show that using a multiscale encoder-decoder is significantly more accurate than a single-scale version.

We have already discussed the reasons behind our novel geometry embeddings briefly in l134-140 and more extensively in SM l675-692. Repeating the main points, existing models either provide coordinates of input points or SDFs to encode geometric information. We argue that this does not suffice for problems where domain geometry itself plays an important role in the solution operator (e.g flow past airfoils) and hence additional geometric information needs to be provided to the model as done with the statistical embeddings based on neighborhood information (l140-149 and SM l693-719). The addition of these embeddings is also justified with the ablation study presented in SM F.4 (Tab F.4) where adding statistical embedding significantly increases accuracy over just inputting coordinates and SDFs. We would be happy to elaborate on the above points further in a camera ready version (CRV) if accepted.

[2.] Regarding the reviewer's concern about lack of clarity regarding efficient implementation, we begin by pointing out that our implementation of GAOT is the key to achieving efficiency, which is the second axis in our rationale for this paper. Following the reviewer's suggestion, we can readily include theoretical complexity estimates for GAOT and baselines in a CRV. Given E edges, C channels, T Latent tokens, p Patch size and L Layers, compute for GAOT scales as $O((E+L(T/p)^2)C)$, GINO as $O((E+LT\log T)C)$ and GNN-based RIGNO as $O(LEC)$. Clearly, GAOT and GINO will outperform RIGNO when $E>>T$ (which can be ensured by design) as we see from Fig 3 (a). This also suggests GAOT will be competitive wrt GINO for large patch sizes. Yet, we show in Fig. 3(a) that GAOT is even more efficient than GINO. To achieve this was the key objective of our efficient implementation and we outline the main steps in l209-216. Following a detailed analysis of the compute cost of each individual component of GINO (pl. see the exact numbers in pt. [4] in the reply to Reviewer diAB below), we observed that for GINO: i) graph building is almost as expensive as encoding/decoding and ii) there is no batch processing due to the heterogeneity in the underlying graphs for PDEs on arbitrary domains, implying that the realized computational cost of GINO is not as low as the theoretical scaling would suggest. To alleviate these issues, as described in l209-216, we were able to i) move graph construction for GAOT outside the model evaluation and ii) implement batch processing in the GAOT transformer processor while sequentially processing inputs for the encoder/decoder. It is precisely these implementation tricks (which can only be described qualitatively and realized in the submitted code) which allow GAOT to be more efficient than GINO (pl. see exact numbers in pt [4] in the reply to Reviewer diAB below) and provide the context for the scaling results in Fig 3a and 3b. We are happy to add theoretical complexity estimates as well as a detailed compute breakdown in a CRV, if accepted and thank the reviewer for suggesting it.

[3.] RIGNO (Ref. [39]) is a state of the art end-to-end GNN that greatly improves upon MeshGraphNets (MGN). In fact, MGN can be realized as a special case of RIGNO. As we considered RIGNO as a baseline, we did not include results with MGN. However, following the suggestions of the reviewer, we ran MGN on several of our benchmarks to obtain the following % relative errors: NS-Gauss (76.4 (MGN) vs 2.91 (GAOT)), NS-PwC (33.1 (MGN) vs 1.5 (GAOT)), NS-SL (18.8 (MGN) vs 1.21 (GAOT)), NS-SVS (76.4 (MGN) vs 6.4 (GAOT)), CE-RP (13.9 (MGN) vs 5.97 (GAOT)), Wave-Layer (84.8 (MGN) vs 5.78 (GAOT)), Wave-C-Sines (17.3 (MGN) vs 4.65 (GAOT)), Poisson-Gauss (30.9 (MGN) vs. 0.83 (GAOT)), Elasticity (11.9 (MGN) vs. 1.34 (GAOT)). All the experiments were run with 2M parameter MGN. Given the limited time of the rebuttal, we could only benchmark MGN on 10 of the 15 tasks of Tab. 1. But, the trend is crystal clear. GOAT is an order (sometime two orders) of magnitude more accurate than MGN. Moreover, one can see from Tab. 1 that RIGNO is also much more accurate than MGN. This justifies our choice of benchmarking against a state of the art GNN such as RIGNO rather than a more traditional MGN baseline. Nevertheless, we will include MGN results in a CRV as an additional baseline. MP-PDE has mostly been run on simple 1-D PDEs and was already shown to be significantly inferior to UPT, which have already benchmarked against (Tab. 1).

[4.] The neural field property is a desirable feature for any ML surrogate for PDEs as it enables evaluation of the input and output fields at any point in the underlying spatial domain. In particular, it allows the model to be trained and tested at different resolutions, a highly desirable feature in PDE operator learning (Ref [22]). For example, with a neural field model, one can train at a coarser resolution (for training efficiency) but test at a higher resolution (for potentially more accuracy). It is this feature of GAOT that we wanted to demonstrate in Tab. 3 for the industrial scale 3-D Drivaernet++ dataset. However, as described in l296-301, having the neural field property does not imply that the model is more accurate. We can observe from Tab. 3 that the version of GAOT trained with only 10% of the total input points but tested at the full dataset resolution of 500K surface points using the neural field property is nowhere as accurate as the version of GAOT trained and tested at full resolution, although it is still better than some of the baselines. Thus, comparing with a neural field model on this dataset will not lead to more accurate results. Nevertheless, to further explore the consequences of model endowed with the neural field property, we observe that the UPT baseline provides a neural field. As testing UPT on an industrial scale dataset in the limited time of rebuttal is not possible, we considered the 2-D NACA0012 airfoil dataset with 2 different settings: i) training and testing GAOT and UPT at full resolution leading to relative % errors of 6.81 (GOAT) vs. 16.1 (UPT) and ii) training on 1/4th of the resolution and testing at full resolution leading to errors of 9.54 (GOAT) vs. 16.3 (UPT). This experiment reinforces that a) training at full resolution is simply more accurate than training at coarser resolutions b) GOAT is significantly more accurate than UPT in both settings and iii) UPT does not improve much with full resolution training, indicating lack of convergence with increasing resolution. We will gladly add these results together with a more detailed discussion of the neural field property in a CRV.

[5.] We will address the pertinent minor points raised by the reviewer in a CRV, if accepted.

Finally, in addition to the very low overall score, the reviewer has given our article the lowest possible score of 1 for originality and clarity. We take this opportunity to bring to your attention the fact that i) multiscale attentional encoder/decoder ii) statistical geometric embeddings and iii) multiple tricks for efficient implementation as described in pt [2.] of our reply above are completely novel to the best of our knowledge and are acknowledged as such by the other reviewers. These original contributions have enabled us to propose a model that is highly accurate, very efficient and scalable. We have demonstrated its state of the art performance on 24 different datasets, corresponding to a wide spectrum of challenging PDEs, including industrial scale 3-D benchmarks. To the best of our knowledge, no existing ML PDE surrogate model has the scalability, range and performance of GAOT. We hope that this rebuttal has allowed us to clarify our contributions accordingly and to reiterate the novelty, scope and performance of our model. Therefore, we kindly request the reviewer to upgrade their assessment accordingly and we are at your disposal to answer any further questions about our article.

We start by thanking the reviewer for your response and for stating that your concerns are partially addressed.

Regarding your question about the runtime of classical solvers, we sincerely apologize for not having responded in our original rebuttal due to space limitations. The comparison of runtime of classical solvers and that of ML surrogates is the very premise of ML for PDEs and is certainly an important issue. We did not discuss it in detail in our paper as a lot of the datasets that we compare GAOT and the baselines on are not generated by us. Rather, they are publicly available datasets.

To be more specific, for the results in Table 1, Poisson-Gauss, NS-Gauss, NS-PwC, NS-SL, NS-SVS, CE-Gauss, CE-RP, and Wave-Layer datasets are taken from Ref. [20], with the specific random point cloud inputs from Ref. [39], the elasticity dataset is from Ref. [26] and Wave-C-Sines is from Ref. [39]. In Page 72 of [20], the authors of [20] state that runtime of a highly optimized GPU spectral viscosity solver that they use for the Navier-Stokes problems (NS-Gauss, NS-PwC, NS-SL, NS-SVS) at the resolution of $128^2$ is 0.1 sec and that of the FEM solver for Poisson-Gauss is 10 sec.,The authors of [20] also state that CE-Gauss, CE-RP are generated with highly optimized GPU-based finite volume solvers for the compressible Euler equations but do not provide the runtimes. As these solvers are more expensive than spectral solvers and the compressible Euler equations have more variables (4), we expect that the runtime of these solvers for the CE-RP and CE-Gauss datasets are between 0.1 sec and 1 sec. Unfortunately, the runtimes of Wave-Layer, elasticity and Wave-C-Sines are not available in the above references. But given that either finite difference or finite element solvers were used, we would expect runtimes between 1 sec and 10 secs for these 2D problems.

We have generated the data for the NACA0012, NACA2412, RAE2822 and Bluff-body datasets for compressible flow past objects as well as the Poisson-C-Sines dataset. The compressible flow datasets took between 283-419 seconds (ca. 5 to 7 min) per sample to generate as the underlying code was CPU based and highly-adaptive grids with suitable iterative solvers were needed to obtain accurate results, please SM D.2 for details. The FEM simulation for Poisson-C-Sines took 10 secs to simulate.

Finally, the industrial scale 3D Drivaernet++ dataset has been taken from Ref. [12]. From Page 9 of [12], we observe that on an average, generating each sample of Drivaernet++ took 375 hours of CPU compute !!. This is not unexpected for this challenging 3-D problem with a RANS model on a dense volumetric grid.

Summarizing, the runtime of classical solvers ranged from 0.1 secs to 7 min in 2D and a massive 375 hours in 3D.

Next, we find that inference runtime of GAOT per sample for the 2D datasets ranged from 8.95 to 10.14 milli-seconds (we will provide the exact numbers for each dataset in a CRV, if accepted). On the other hand, for 3D Drivaernet++ dataset, the GAOT inference runtime was ca. 365 milliseconds per sample.

Thus, for the 2D datasets, the runtime gain for GAOT over a classical solver ranged from 1 to 5 orders of magnitude, depending on the dataset considered. On the other hand for the industrial scale 3D dataset, the gain was a massive factor of 37M (7 orders of magnitude), while still being the SOTA model on this benchmark.

One can argue that these gains come at cost of errors with respect to the numerical solvers. However, these errors are small. For instance, in the Poisson-Gauss, the error with GAOT is less than 1% (with a runtime gain of a factor of 10000) and for Bluff body, the error is ca 2% (with a runtime gain of 29000). For the 3D benchmark, the gain is very significant. We also refer the reviewer to an earlier study on the accuracy-runtime efficiency performed in Ref. [45] (Page 58) for the NS-SL problem where the authors of Ref. [45] found that their model (CNO) was able to provide accuracy comparable to a spectral solver on a $100^2$ grid. We use exactly the same setup as [45] and find that GAOT has even lower errors than CNO on this benchmark (SM Table E.2). Hence, for this NS-SL dataset, we can show that GAOT has comparable errors to the spectral GPU solver, while still being a factor of 10 faster to run.

We are happy to add a discussion on runtime (and accuracy) comparisons of GAOT and classical solvers in the appendix of a CRV, if accepted and thank the reviewer for pointing out this avenue to further highlight the merits of our model, especially on industrial scale 3D simulations.

We thank the reviewer again for your questions and hope that we have addressed your remaining concerns satisfactorily. If so, we request the reviewer to kindly upgrade your assessment accordingly, while remaining at your disposal for further questions.

We sincerely thank the reviewer for raising your score to borderline accept and for your comments and suggestions which can help us improve the quality and clarity of our paper.

We are happy to discuss the cost-accuracy tradeoff of numerical solvers for PDEs vs. ML surrogates in a CRV, if accepted. To this end, we will mention and contextualize the study performed in Ref. [45], where in page 58, the authors of [45] precisely present such a study for the NS-SL benchmark on regular grids and convincingly show that the accuracy of their model (CNO) was comparable to (error of 2%) with accuracy of a spectral viscosity solver at a resolution of $100^2$ (also error of 2%), when compared to a high-resolution spectral solver for this problem (Ground Truth). So for this dataset, [45] has already shown that the error of $100^2$ grid with a spectral-viscosity solver is 2%, with a run-time of 0.1 sec. As it happens, we had run this dataset with GAOT and presented the results in SM Table E.2 to obtain an error of less than 1%, with an inference time of 10 ms. Thus, our ML surrogate is 10 times faster than a state of the art GPU optimized numerical solver, even for the same accuracy. Nevertheless, we cannot perform this analysis of every dataset that we test on as we have not generated the underlying data in most cases. We will acknowledge this fact as a a limitation, not just of our ML model but most ML models for PDEs as this cost-benefit accuracy with respect to classical solvers is rarely discussed in the literature.

We thank the reviewer again for your comments.

Motivated by the reviewer's comments about the cost-accuracy comparison of classical numerical solvers and ML surrogates, we wanted to explore if the cost-accuracy results of Ref. [45] (as extensively discussed above), which were obtained on regular grids, can be replicated for the case of arbitrary point clouds that we have considered as the main use case of GAOT.

To this end, we focussed on one of the NACA airfoils dataset, where the input and output grids are unstructured. Using a highly refined grid (similar to the one shown in SM Figure D.2) to generate the ground truth, we coarsened this mesh by 2 levels of refinement. Preliminary results indicate that the runtime came down to 1 min (from 7 min for the full mesh) and the resulting relative error was approximately 5%. As reported in Table. 1, the GOAT relative error is approximately 6.5%. Thus, GAOT has comparable errors to a lower resolution numerical solver. Yet, its runtime of 10 ms is 6K times faster. This experiment clearly justifies the use of an accurate and efficient ML surrogate like GOAT, especially on problems on unstructured grids and 3D, even more than regular grids as considered in [45].

We will include these results in a CRV, if accepted and thank the reviewer for pressing us to provide such a quantitative comparison. We hope that these results have further clarified the reviewer's remaining concern and play a role in their final assessment of our paper.