A composable autoencoder-based algorithm for accelerating numerical simulations

Abstract: Numerical simulations for engineering applications solve partial differential equations (PDE) to model various physical processes. Traditional PDE solvers are very accurate but computationally costly. On the other hand, Machine Learning (ML) methods offer a significant computational speedup but face challenges with accuracy and generalization to different PDE conditions, such as geometry, boundary conditions, initial conditions and PDE source terms. In this work, we propose a novel ML-based approach, CoAE-MLSim, which is an unsupervised, lower-dimensional, local method, that is motivated from key ideas used in commercial PDE solvers. This allows our approach to learn better with relatively fewer samples of PDE solutions. The proposed ML-approach is compared against commercial solvers for better benchmarks as well as latest ML-approaches for solving PDEs. It is tested for a variety of complex engineering cases to demonstrate its computational speed, accuracy, scalability, and generalization across different PDE conditions. The results show that our approach captures physics accurately across all metrics of comparison (including measures such as results on section cuts and lines).

Reviews and author responses:

 We were wondering if you had any other questions or comments about the recent revision we submitted. Thanks.

The paper proposed CoAE-MLSim to learn with relatively fewer samples of PDE solutions and solve PDEs. CoAE-MLSim uses the idea of domain decomposition: first learn the solution on local subdomains using autoencoder, and then couple these sub-solutions by an iterative algorithm. Numerical experiments are performed to test the efficiency of the method. Strengths:
- The proposed method uses the idea of domain decomposition, and thus it has the potential to only need a small dataset to train the networks.

Weaknesses:
- The method requires 100-1000 PDE solutions for training, which is not really a very small dataset as they claimed. For a comparison, in the paper arXiv:2010.08895, it also uses 1000 solutions for the NS equation, and in this paper, it uses 500 in Section 4.2.
- The method part is not well written and hard to read. After reading Section 3.2/3.3, I only capture the main idea, but the detail of the algorithm is very vague and not well presented, e.g., how the neural networks are used here, what the network input and output are, how the flux conservation is applied, etc. The authors should rearrange Section 3 so that it is easy to read.
- During the inference stage, because the method needs to couple all the solutions in subdomains together by using an iterative algorithm. The inference time is significantly slower than other ML methods, which only requires one forward pass of the network.
- Also, it is not fair to compare the speed with Ansys solver, because Ansys is much more accuracy than the proposed method. If we need to compare with Ansys solver, then the authors should use a low resolution mesh of Ansys solver, such that it has a low accuracy similar to CoAE-MLSim, and then we can compare the speed.
- The accuracy is not very good. In Section 4.1, the MSE of the temperature is 5.3e-2, so without the square, the error is > 0.2, which is really large.
- For the result in Section 4.2, the timestep is dt = 1e-4, and here the authors only predict for 50 timesteps, i.e., t = 50 * dt = 0.005, which is too small. Even for the first few timesteps, the L2 error in Fig. 8 is almost 0.1, which is large.
- The authors claims that the method has a good generalization, but there is no evidence to support this. The prediction of future time in Fig. 8 is not really extrapolation. For example, we would not say time series prediction is extrapolation, and it is still interpolation in the space of PDE condition/solution.

Other comments:
- What does “CoAE-MLSim” stands for? Composable Autoencoder - Machine Learning Simulation? In the whole paper, “Composable” is only mentioned in the title, but not any where in the text, and thus it is not clear what it means exactly.
- In the abstract, the authors only mentioned the name of the method. After reading the abstract, I have no idea what “CoAE-MLSim” is.
- In the abstract, “unsupervised” is conflict with “relatively fewer samples of PDE solutions”. The method is not unsupervised.
- The Point 4 “The iterative inferencing algorithm is completely unsupervised and allows for coupling with traditional PDE solvers.” The iterative algorithm is part of inference, and the inference of networks of course doesn’t need the data. It is strange to say the inference stage is “unsupervised”. Also in Section 3.8.
- In Algorithm 1 Line 5, there is a “?”.
- There is empty space in the top right corner of Page 6.
- Is the SDF for the local domain or the global domain?
- In Fig. 4, the SDF looks strange. SDF is a continuous function, but in the figure, it changes from red color to blue color sharply and then keep the same blue color.
- The fonts in many figures are too small to read.
 The main issue is that their numerical experiments are not convincing enough to support their claims on small data, accuracy, and extrapolation. 

 The paper has been improved, especially on the writing. But there are still concerns about the numerical results.

- The main strength of the method is that it needs fewer training data, but the factor is not significant enough in the current results. 
- The iterative algorithm has its flexibility, but it is much more expensive than only one forward pass of other ML methods.

I raised my score, but I would like to see that the results will be further improved in terms of either the required dataset size or inference speed.

 Just checking if you had a chance to review responses to your questions as well as other posts updated on this thread. Please let us know if you have additional questions.

 Just checking if you had a chance to review responses to your questions as well as other posts updated on this thread. Please let us know if you have additional questions. 

 We thank the reviewer for their constructive comments and suggestions. We have addressed them below and made required changes to the revised manuscript. Additionally, we have also revised our manuscript to address most of the minor comments from the previous discussion. 

The reviewer comments are numbered and marked in bold.

**1) Since FCNN is such a strong competitor, it shouldn't be only in the appendix.**

As mentioned in our previous comments, FCNN performs better than UNet and FNO only because of the representation of boundary conditions between these methods, which is only important for the Laplace equations. Nonetheless, we are in the process of running FCNN for all the other use cases and add it to the final version of the manuscript if accepted.

**2) The authors didn't add train/test comparisons and don't think this is relevant.**

We apologize for the misunderstanding. We do think that this experiment is relevant and we provided a table in the previous comments comparing the average training error to average testing error. We mentioned in the previous comments that “If required, this analysis can be added to the revised paper.”. We have added this discussion to the appendix section A.2.3 of the revised paper.

**3) I still don't understand the authors' comments about "scaling." I don't know what definition of the word they are using.**

We have removed this statement from the paper to avoid any confusion. 

**4) The authors agree with me that they cannot back up all of their claims, but still left some of them in.**

As mentioned in our comments earlier, there are a lot of small components in our system that can be validated. Although, we have not validated all the claims, we wanted the reader to have the knowledge of our experience and intuition of developing this method. As a result, many claims in the paper were modified to “We recommend …”. The comment highlighted by the reviewer got missed out and we thank the reviewer for bringing it to our attention. We have removed this claim from the appendix of the paper. 

**5) I looked at Ranade et al., 2021a, Ranade et al., 2021b, and Maleki, et al., 2021 again, and I still think that they are not cited in a way that fairly conveys to the reader which parts of this paper are new.**

We have expanded on this in the revised paper. The following discussion has been added:

*The use of autoencoders for compressed solution and condition representation in tandem with iterative inferencing solution algorithm is inspired from Ranade et al. (2021a, 2021b) but the main difference in this work is that the solution procedure is carried out on local subdomains as opposed to entire computational domains to align with principles used in traditional PDE solvers.*

**Additional comments from previous review:**

All other comments have been addressed in the revised paper. Comments that require additional explanation are presented below.

**1) Why is Err(T_max) negative sometimes? Maybe I don't understand the definition, but I would expect to use absolute value?**

Err (T_max) is the difference in the maximum temperature predicted and ground truth. As mentioned in the paper, this seems to be an important metric for the chip cooling problem as it indicates whether the hottest spot is modeled correctly or not. Therefore, it has a sign to indicate whether the prediction exceeds or falls short of the ground truth.

**2) Sequence of autoencoder training**

We thought that this order may be clear from the solution algorithm, but we have added a comment to clarify this in the revised paper. The following statement has been added:

*Since, flux conservation and time integration networks are trained using the compressed latent vectors of PDE solutions and conditions, the PDE solution and condition autoencoders are required to be trained first.*

**3) Boundary condition encoding confusion**

We have changed the description in the paper to reflect that the boundary conditions have been manually encoded and not learnt with an autoencoder in this work.

Please let us know if there are any additional questions.


 Thank you to the authors - the paper is much improved. 

I still have some significant concerns before the paper is ready to be published. 
- Since FCNN is such a strong competitor, it shouldn't be only in the appendix.
- The authors didn't add train/test comparisons and don't think this is relevant. "As a result, the algorithm can perform the same on the training PDE conditions as it would on the testing conditions. This is different from traditional ML methods, which work better in training than in testing." As we can see from the train/test pair added to the response, this is definitely not true - it's possible for this method to overfit just like other deep learning methods. I understand that the iterative algorithm has nice features, but the autoencoders can still overfit. This is important to check in any deep learning paper, but especially when a major selling point of the method is supposed to be improved generalization. Without further information, it's still not clear to me if this method *generalizes* better than the other methods. 
- I still don't understand the authors' comments about "scaling." I don't know what definition of the word they are using.
- The authors agree with me that they cannot back up all of their claims, but still left some of them in. For example, "Additionally, the CoAE-MLSim apprach solves the PDE solution in the latent space, and hence, the idea of conditioning at the bottleneck layer improves solution predictions near geometry and boundaries, especially when the solution latent vector prediction has minor deviations." This is something that could be checked by running a variant of their method without the conditioning. 
- I looked at Ranade et al., 2021a, Ranade et al., 2021b, and Maleki, et al., 2021 again, and I still think that they are not cited in a way that fairly conveys to the reader which parts of this paper are new. In particular, some parts of this paper that I thought were interesting were also in Ranade et al., 2021a, and some of the figures are even quite similar. Of course I agree that there is enough new in this paper that it should be published, but I also think that it needs to be more clear to the reader which parts are new and which parts already appeared in Ranade et al., 2021a.

The paper propsed a composable AE-based PDE solver that combines traditional PDE solver with ML techniques. It split the big domain into sub-domains and compress the high dimensional representation into lower dimension using autoencoder. Overall the clarity of the paper needs some seriour improvement. 

Questions/Weaknesses:
1. The related work sections covers a wide range of works that related to pde solving using neural networks. Although the paper provided a great list of related works, the relationship between this work and those mentioned is unclear to me. This is a critical piece that I am missing. If the related work section is delivered with more emphasis on the relationship instead of a mere list, it would be much improved.

2. The clarity of method description needs some significant improvement. Possibly rising from the size of the figures being too small, I had a hard time understanding the general flow and what to expect for the method proposed. A lot of questions like what is the dimension of each latent space, what is the size of your encoder network and do all the subdomains share the same encoder (I believe so)? How does the solution conditional AE take the source and geometry as conditions? Are they just appended to the input of both encoder and decoder?  

3. The autoencoders usually does not generalize well for the unseen data and the decoded result from unseen latent space is usually not valid in the original space unless being regularized properly (like Variational autoencoder etc.). How does this method do when facing out-of-distribution test cases?

Couple of minor points:
1. I believe the Figure 1 is not mentioned in the main text?
2. The figures are really too smallto see clearly. It heavily undermines the clarity of the paper
3. Page 4 alg 1 line 5 there is a error in latex? I believe it is <?
4. What is \THETA in page 4 alg 1 line 8? 
5. In Section 3.4 line 6 what is the citation [Ranade et al., 2021a] for? CoAE-MLSim is the proposed algorithm right?
6. For result part, again the size of the plot makes it impossible to read
7. The final table, there is no caption or table label and the unit is unclear The paper is lacking in clarity as mentioned above. Although technically might be novel and can benefit the scientific computing community, this is far from publishable from a perspective of readability. 

Update on Nov. 20th during rebuttal:
I think during the rebuttal period, I have gain much more understanding of the paper and willing to raise the score. My remaining concerns are clarity of the paper (including the fact that I still need to use 200% to see whats in the figures, but understandable given the page limit) and out-of-distribution generalizability.

The paper raised an interesting idea, using changing the simulation to the latent space by compression of autoencoders. However as the overall goal is still to simulate new designs, the generalizability of it due to autoencoders might still bottleneck its practicle impact.



 Yes, the understanding is correct that our approach is designed to operate like a numerical solver but on the level of subdomains (group of computational elements). The speed-up over traditional PDE solvers comes from this as well as the fact that the solutions/PDE conditions on subdomains are compressed further. We use PDE solution examples on subdomains to train the weights of the autoencoders. Similarly, groups of neighboring subdomains to learn the weights of flux conservation autoencoders and time integration networks. The latent vectors are bounded between -1 and 1 because of tanh activation. The variance divided by mean of the latent vectors over all the subdomains in the vortex decay case is -6.33607 and standard deviation divided by mean is -40.667. 

Our approach will be slower than direct/static inferencing approaches such as FNO, UNet etc. as we use an iterative procedure during inference, where the flux conservation autoencoder is evaluated multiple times. However, the computational cost is comparable and moreover, we are still faster than traditional PDE solvers by 40-50x in steady-state settings and about 100x in transient settings.  We understand that there is a cost paid in terms of computational time with dynamic inferencing approaches, however the benefits in terms of accuracy and generalizability as compared to static inferencing approaches outweigh them. 

Point 3 in response to Reviewer WDpc is also along similar lines and we request the reviewer to read that as well for more information. Additionally, in response to reviewer hE8v, we presented a table comparing the computational cost at inference for different flux conservation autoencoer layer sizes. This table provides a ball park of the computational time required for the solution of the Darcy's equations.

 Thank authors for the detailed response.

If I am understanding correctly, CoAE-MLSim is encoding all PDE condition and solutions into low dimensional space and operate just like a numerical solver. Is the 100-1000 samples required to learn the sub-domain autoencoder weights number of sub-domain examples or total examples? Also I wonder how diverse is the sub-domain solution latent variables (maybe measured by the variance of that dimension divided by mean?)

Another question regarding to the effectiveness of this method is the running/evaluation time difference between benchmarks. How fast is this comparing to FNO? 

 great, thanks!

 This is referring to the full PDE solutions. In the above experiment, we train several flux conservation autoencoders with different bottleneck layer sizes and perform an inference on 50 test cases for each layer size. The table reports the $L_2$ norm error ratio, convergence iterations and computational time averaged over these different test cases.

 The notation $Y_{true}$ and $Y_{pred}$ seems new. Is this error referencing the full PDE solutions, or are these errors specific to the flux autoencoder aspect?

 **5. The authors claims that the method has a good generalization, but there is no evidence to support this.**

We claim good generalization and support it with the variety of numerical experiments presented in the paper in Section 4 of the main paper and Appendix section A.2. Here we again provide a brief summary of all the experiments and explain how they demonstrate generalization. 

1)	**Electronic chip cooling:** Here we demonstrate generalization over source term conditions. The source term is sampled from a high-dimensional space of 4096 using a Gaussian mixture model that can contain up to 25 Gaussians each with their own means and variances. Learning PDE solutions corresponding to this space with just 300 solutions on a 128^3 resolution computational domain for random power maps sampled from the mixture model is an extremely challenging task. The comparison with UNet baseline in Section 4.3 demonstrates this. Yet, our method is able learn from this data set and accurately generalize to 300 other unseen source terms sampled from this high-dimensional space.
2)	**Vortex decay:** Here we show generalization across different initial conditions as well as extrapolation capabilities in time of our method. Again, the dataset is based on only 500 initial conditions with 25 timesteps per condition (first 10 are fixed and not used, so only 15 timesteps are used for training the autoencoders). The results show that we can accurately generalize to 500 other initial conditions and moreover, predict with controlled error accumulation up to 50 timesteps, which is twice as long as the training regime.
3)	**Laplace equations:** Here we show generalization across different boundary conditions. Again, our approach is trained on limited data (300 solutions) generated for arbitrary choices of boundary conditions and demonstrate that we can accurately predict for other choices of boundary conditions. We also show that other baselines don’t predict as well even though we train them with a significantly larger number of samples (6400).
4)	**Navier-Stokes equations for external flow:** Here we demonstrate generalization over different geometry shapes that are immersed in fluid flow on a computational domain resolution of 320x160x80. The geometry representation is high-dimensional, with up to 4096 dimensions. Again, we train the autoencoders on 300 solutions and show accurate predictions for any arbitrary sampled geometry from the remaining 300 test solutions.
5)	**Darcy’s equation:** Finally, to address the comments of Reviewer 1 we added another use case, and our approach continues to accurately generalize for this use case as well. 

More details about each of these numerical experiments may be found in the paper (Section 4) as well as Appendix (Section A.2). 

Can the reviewer please clarify why they think we have not provided enough evidence to validate our claims?

**6. The prediction of future time in Fig. 8 is not really extrapolation. For example, we would not say time series prediction is extrapolation, and it is still interpolation in the space of PDE condition/solution.**

This argument is unclear and possibly incorrect. Can the reviewer please clarify why time series cannot be extrapolation? If you consider the physics represented in this problem, we start with an initial vortex in the domain and as time goes on this vortex breaks down and diffuses within the bounds of the computational domain. This flow dynamics for the first 25 timesteps will be completely different to what can be observed from 25 to 50 timesteps. Considering, none of the solutions used to train the autoencoders have seen solutions beyond 25 timesteps, why would you not consider this an extrapolation in time dimension? 

**Other comments**

Most of the minor comments are incorporated in the revised paper. Here are responses to a few.

**1. Is the SDF for the local domain or the global domain? In Fig. 4, the SDF looks strange. SDF is a continuous function, but in the figure, it changes from red color to blue color sharply and then keep the same blue color.**

The SDF is for local subdomains because most of the global domain does not have representation of geometry. The SDF is continuous within the local subdomain. There is no need for an SDF in subdomains where the geometry is absent.

**2. Reviewer comment on unsupervised learning**

The autoencoders are trained with PDE solutions but the solution algorithm used at iterative inferencing is unsupervised and is never taught the procedure of solving PDEs. Examples of PDE solutions are only used in the autoencoder to teach them solution representation, the rules of flux conservation etc. If not for our solution algorithms, the autoencoders cannot deduce the solutions by themselves. So, we think the solution algorithm may be considered as unsupervised.
 

 We thank the reviewer for spending the time to review our manuscript. The concerns raised are addressed below.

The reviewer comments are numbered and marked in bold.

**1. The method requires 100-1000 PDE solutions for training, which is not really a very small dataset as they claimed. For a comparison, in the paper arXiv:2010.08895, it also uses 1000 solutions for the NS equation, and in this paper, it uses 500 in Section 4.2.**

As the reviewer points out, we use solutions from 500 initial conditions to train our autoencoders. This cuts our data set requirement into half as compared to the paper arXiv:2010.08895. Additionally, we only retain solutions corresponding to the first 25 timesteps for each initial condition as opposed to 40-50 timesteps used in the paper arXiv:2010.08895. Since the first 10 timesteps are fixed in both approaches, the data set requirements in our approach are a quarter of the data set used in arXiv:2010.08895. This story is consistent across all the numerical experiments presented in the paper, where we can accurately capture PDE solutions across a high-dimensional space of PDE conditions by learning with a very small set of samples. Moreover, the comparisons with baseline in Section 4.2, show that our approach outperforms arXiv:2010.08895, when only a quarter of the dataset is used. 

**2. The method part is not well written and hard to read.**

We have improved the figures and the corresponding explanation.

**3. During the inference stage, because the method needs to couple all the solutions in subdomains together by using an iterative algorithm. The inference time is significantly slower than other ML methods, which only requires one forward pass of the network. Also, it is not fair to compare the speed with Ansys solver, because Ansys is much more accuracy than the proposed method. If we need to compare with Ansys solver, then the authors should use a low resolution mesh of Ansys solver, such that it has a low accuracy similar to CoAE-MLSim, and then we can compare the speed.**

The reviewer misses the point about why we have proposed an iterative inferencing algorithm. Dynamic iterative inferencing is in fact a strength of our approach, as pointed out by reviewer 2. Dynamic inferencing offers flexibility and more control of the solution trajectory. Moreover, it enables seamless and easy coupling with traditional PDE solvers. Additionally, it provides us with the flexibility of preconditioning our algorithm with an initial solution to result in faster convergence. Dynamic inferencing also provides us with the possibility of handling boundary conditions and geometry in a specialized manner. For example, periodic or Dirichlet boundary conditions can be explicitly enforced on subdomains adjacent to the boundary. Similarly, for subdomains across geometries we could explicitly provide wall models to handle turbulence. Finally, the main goal of our work is to develop a fast machine learning based approach that closely resembles the solution algorithms in traditional PDE solvers to achieve good stability and robustness. Moreover, this work provides us a successful towards building true, generalizable ML-based PDE solver in the future. We understand that with dynamic inferencing we pay a cost in terms of computational speed as compared to static inferencing approaches but the benefits this strategy provides outweighs those costs. The reason we compare the speed with Ansys solver is to simply demonstrate that even though our method is slower than static inferencing approaches (which are 1000x faster than Ansys solvers) it is still much faster than an Ansys solver and is a viable path for development of ML-Solvers. 

**4. The accuracy is not very good. In Section 4.1, the MSE of the temperature is 5.3e-2, so without the square, the error is > 0.2, which is really large.**

We don’t agree that the error is very large. If the reviewer looks closely, the temperature scale ranges from 300-600 K. A mean squared error of 5.3e-2 on that range can be considered as very good accuracy.

**5. For the result in Section 4.2, the timestep is dt = 1e-4, and here the authors only predict for 50 timesteps, i.e., t = 50 * dt = 0.005, which is too small. Even for the first few timesteps, the L2 error in Fig. 8 is almost 0.1, which is large.**

The dataset used in this experiment is taken from the paper arXiv:2010.08895. The data is available up to 50 timesteps. But we wanted to show that with the same dataset we outperform the baselines and show better extrapolation capability. Moreover, we learn the dynamics from only 15 timesteps for 500 initial conditions and yet achieve results with reasonable accuracy and controlled error accumulation. As more data become available, our model can be significantly improved.

**Continued in the comment below due to character limit** 

 We thank the reviewer for their efforts and time in reviewing our manuscript. The concerns raised are addressed below. 

The reviewer comments are numbered and marked in bold.

**1. The related work sections covers a wide range of works that related to pde solving using neural networks. Although the paper provided a great list of related works, the relationship between this work and those mentioned is unclear to me. This is a critical piece that I am missing. If the related work section is delivered with more emphasis on the relationship instead of a mere list, it would be much improved.**

Most of the work referenced in the literature share a common goal of solving PDEs with data-driven machine learning methods. Our approach is different in that we use unsupervised solution algorithm to extend the bounds of generalization and achieve accurate results (in comparison to Ansys Fluent, throughout the domain not just improving average error norms) with low number of samples. We have reworked this section in the revised paper to provide some more details. We are happy to modify this further if there are any specific issues.

**2. A lot of questions like what is the dimension of each latent space, what is the size of your encoder network and do all the subdomains share the same encoder (I believe so)?**

Even though the network architectures of the different autoencoders remains pretty much the same across all applications, the size of the latent vector (a hyper-parameter) may vary for different applications, depending on the complexity of physics, complexity of PDE conditions, size of subdomains etc. As a result, we chose to provide this information in the appendix where we discuss each of the numerical experiments in more detail. In the paper, we have only discussed the general setup and network architecture since that remains similar for different applications. Yes, all subdomains pass through the same autoencoder. 

**3. How does the solution conditional AE take the source and geometry as conditions? Are they just appended to the input of both encoder and decoder?**

Yes, encodings of PDE conditions obtained from pretrained PDE condition encoders are simply concatenated with the solution encodings while training the solution autoencoder. This information is provided in in Section 3.4 of the paper.

**4. The clarity of method description needs some significant improvement. Possibly rising from the size of the figures being too small, I had a hard time understanding the general flow and what to expect for the method proposed.**

We have improved the figures and the corresponding explanations in the revised paper. Hope it helps.

**5. The autoencoders usually does not generalize well for the unseen data and the decoded result from unseen latent space is usually not valid in the original space unless being regularized properly (like Variational autoencoder etc.).**

In our approach we rely on learning local PDE solutions on subdomains with a small number of voxels/computational elements. The reason for this is to limit the space of PDE solutions that exist on the small subdomain and subsequently reduce the amount of variation in PDE solution seen by the subdomain. This local learning aspect enables us to learn generalizable representations with vanilla autoencoders. If the same autoencoder were to be used to represent solution on the entire computational domain, then we agree that the generalizability could be limited. We agree with the reviewer that other ideas such as variational autoencoders, physics-based regularization etc. will help improve the robustness of the autoencoder and further improve our results but in this work, we show that even with vanilla autoencoders we are able to get good accuracy and generalizability across arbitrary, high-dimensional PDE conditions for a variety of PDEs and applications because they are used in context of local learning where patterns repeat.

**6. How does this method do when facing out-of-distribution test cases?**

This is a great question for the entire ML community. However, data-driven methods are always limited in their ability to perform out-of-distribution generalization. Even though our method is different, there is always an extent of generalization and we have discussed that in the paper in Section 4.3. We agree that our method, like all other ML approaches, is not application agnostic yet. We don’t expect the autoencoders trained on one application will work for a different application as a traditional solver would or for that matter for a substantially different distribution of PDE conditions. However, our method shows a great generalization ability in the same bounds defined by other ML approaches. In future, we will be working on extending in this direction by using physics-based regularization and learning from truly random solution fields. 

**Minor points**

All the minor points identified by the reviewer are addressed in the revision.


 **7. "It may be observed that the FCNN performs better than both UNet and FNO and this points to an important aspect about representation of PDE conditions and its impact on accuracy."**

All of these networks were optimized to get the best performance out of them. Unfortunately, we did not record these experiments. However, the main difference between them was the input representation. The encoded form of input representation enabled FCNN to predict better. However, UNet and FNO cannot take a similar representation as an input but requires the boundary conditions represented as an image, which is sparse and high-dimensional. 

All in all, we appreciate that the reviewer has gone through our paper in great detail and identified these points. The approach we have developed has a lot of moving parts and it is very difficult to provide results to verify all of these claims in this paper. Some of the claims made are based on our observations and learnings that we gained while developing this approach. However, we will continue to work on generating results in verifying all of those and add them to our paper as we have more results.

**Relationship to literature:**

There is not much overlap between Ranade et al., 2021a, Ranade et al., 2021b, and Maleki, et al., 2021 and our approach. The only thing we have adopted from these papers is the autoencoder network architectures. These papers have some commonality with our paper in that the goal is accelerate numerical simulations and solve PDEs. We agree with the reviewer that in the quoted statement, the reference is misplaced. We have changed this in the revised version of the paper and expanded in the related works on explaining the overlaps with these papers.

**1. I thought that otherwise the related work section was okay but was largely just summarizing some papers without giving context for how they relate to this paper.**

We have modified the related works section to provide more context.

**We again thank the reviewer for their detailed feedback and look forward to continuing the discussion. The feedback on minor details be addressed in a new version that we will upload in a couple of days.**

 Additionally, in the table below we provide the average errors for the 3 approaches in the interpolation and extrapolation regions separately. For the interpolation regime the averages are computed over 50 testing samples and until 25 timesteps, while for extrapolation regime they are computed over 50 testing samples and for timesteps between 25 and 50. The error metric is same as the one used in the paper.

| || | | | Average error in interpolation regime | | | Average error in extrapolation regime | |
| ---| ---| ---| ---| ---| ---| ---| ---| ---| ---|
|CoAE-MLSim |||||0.012 |||0.058|
|FNO |||||0.009 |||0.186 |
|Unet |||||0.017 |||0.152 |

It may be observed that the CoAE-MLSim performs as well as FNO and UNet in the interpolation region but outperforms both of them for long time flow dynamics. This table is also added to the revised version of the paper.

We thank the reviewer for suggesting this experiment as it adds more clarity about the capabilities of our approach and further validates it.

**Backing up claims**

**1. The abstract says that the method is tested for a variety of cases to demonstrate a list of things, including "scalability." The list of "significant contributions" also includes "This enables scaling to arbitrary PDE conditions..." I might have missed/forgotten something, but I think this wasn't tested?**

We meant that our approach scale to arbitrary PDE conditions from a certain distribution. For example, in the case of electronic cooling, the source term in sampled from a Gaussian mixture model with a maximum of 25 Gaussians each having their own mean, variance and amplitude. This is a high dimensional distribution, and we require about 300 samples to learn the flow dynamics through our system of autoencoders. But, once the autoencoders are pretrained it can scale to all different choices of source terms from this wide distribution. We have changed the text in the paper to reflect this. 

**2. "Hence, the choice of subdomain size depends on the trade-off between speed and accuracy." This isn't clear to me from the results. It seems like 32^3 is the fastest and most accurate?**

The choice of the subdomain size can be application specific. It mainly depends on the spatial scales of flow that can be observed in a certain application. In scenarios, where the flow scales are very fine, we would generally use smaller subdomains to reduce the variation of flow within a subdomain but for coarser scales we can get away with larger sizes. For the application we considered in the ablation study in Section 4.3, $32^3$ turns out to be the most accurate and fastest but that may not always be true for all applications. Through our experience we have found out that 16^3 provides the best trade-off between capturing different scales of flow as well as speed for a wide range of applications. We see how this comment can cause this confusion. We have removed this statement from the revised version of our paper.

**3. "Physics constrained optimization at inference time can be used to improve convergence robustness and fidelity with physics."**

We have modified this to: “"Physics constrained optimization at inference time might be used to improve convergence robustness and fidelity with physics.””

**4. "The decoupling allows for better modeling of long range time dynamics and results in improved stability and generalizability."**

The reviewer is correct that this claim is not verified in the paper. We make this comment based on our observations from a simpler use case that decoupling of time and space works better as opposed to a fully coupled system. The use case is that of a 1-D Burger’s equation with time. Here, we perform 2 experiments, one in which the space and time are decoupled and another where they are coupled. We have added those results to Appendix section A.2.7 to verify this claim. 

**5. "Each solution variable can be trained using a different autoencoder to improve accuracy."**

This claim has not been verified in the paper. But is meant to be a general direction based on our experience for someone who wants to use our approach. We have modified this to say:
“We recommend the training of each solution variable using a different autoencoder to improve accuracy.”

**6. "Since, the PDE solutions are dependent and unique to PDE conditions, establishing this explicit dependency in the autoencoder improves robustness." "Additionally, the CoAE-MLSim apprach solves the PDE solution in the latent space, and hence, the idea of conditioning at the bottleneck layer improves solution predictions near geometry and boundaries, especially when the solution latent vector prediction has minor deviations."**

Again, this more of a general comment based on our experience with this approach. We have currently not provided any experiments in the paper to verify this.

**Continued in the next comment due to character limit** 




 We thank the reviewer for reviewing our manuscript and providing useful comments and suggestions. We have addressed those comments with additional experiments and explanations provided below. We also request the reviewer to review the revised manuscript.

The reviewer comments are numbered and marked in bold.

**Comparison to baselines:**

**1. FCNN is by far the strongest competitor for the Laplace example in the appendix. Why is this left off of the baseline comparison table in the main paper? Further, is there any reason that FCNN couldn't have been used for the other examples?**

The boundary condition representation is sparse for Laplace equations and this is not true for the other problems considered in the paper. Hence, FCNN is expected to underperform for other cases. But, in the case Laplace equation, FCNN preforms better than UNet and FNO because it uses the compressed encoding of boundary conditions as inputs to predict the solutions. Another motivation for demonstrating this is to show the importance of how better representation of inputs can improve accuracy even for simpler architectures. But this is an unfair comparison for UNet and FNO because they cannot use compressed representations of boundary conditions and have to use the sparse, high-dimensional representations of boundary conditions in the form of images. UNet and FNO are still superior to FCNN for the same input representation and hence we only chose them as baselines. This discussion can be found in Section A.2.1 in the Appendix.

**2. Why is FNO not applied to the Chip cooling (Temperature) example?**

We tried running FNO for the chip cooling problem, but their code was giving an error which we were not able to resolve. This error in the code was observed only with 3D-Steady state cases.

**3. In other words, is CoAE-MLSim's error lower than UNet's error because the approach fit the training data better, or is it because it generalized better? Further, in some cases, it's not obvious to me if the test errors are impressive, so maybe it is having a hard time generalizing. It would be helpful to see train vs. test errors, and ideally I like to see train vs. val. vs. test.**

Our model performs same on the training set as it would perform on the testing set. The training portion of our model involves training of a bunch autoencoders. At inference time, the pretrained encoders cannot directly provide the solution directly for a particular setting of PDE conditions. We have to use the iterative algorithm to compute an equilibrium solution for a specified set of PDE conditions. The iterative algorithm must figure out the solution and is never taught anything about the trajectory to get there. As a result, the algorithm can perform the same on the training PDE conditions as it would on the testing conditions. This is different from traditional ML methods, which work better in training than in testing. However, to prove this point we carried out an experiment to compare the accuracy of the inference algorithm on both training and testing PDE conditions. This experiment is carried out on the electronic chip cooling problem. In this problem, we use 300 solutions for training the autoencoders and another 300 for testing. The PDE conditions with respect to both the training and testing set are used in the iterative solution algorithm to generate a solution. In the table below, we compare the average results over the training set and the testing set for the infinity norm and maximum value of temperature in the computational domain. We also provide the best error that we obtain from all the training samples. If required, this analysis can be added to the revised paper. 

| || | | | Temperature $L_{\infty}(Y_{pred},Y_{true})$ | | | $Err. (T_{max})$ | |
| ---| ---| ---| ---| ---| ---| ---| ---| ---| ---|
|Best training err |||||3.45 |||-0.3 |
|Avg. training err |||||13.6 |||2.78 |
|Avg testing err |||||15.1 |||3.93 |

Moreover, we would like to clarify that the PDE conditions are randomly sampled from a distribution on a high dimensional space and because of this they can be significantly different in the training and testing sets.

**3. Do the results in the table for vortex decay back up the claim that CoAE-MLSim is generalizing over initial conditions better than FNO, or is it just better at extrapolation in time?**

We agree that this will be a great experiment to show the capability of our approach in interpolation as well as extrapolation. As per the reviewer suggestions, in this experiment, we have averaged over 50 unseen testing samples from different initial conditions at each of the 50 timesteps. The figure corresponding to this experiment is attached in the Appendix section, A.2.4. It may be observed from the figure that our approach performs as good as FNO and UNet in the interpolation regime but outperforms both methods in the extrapolation regime by a large margin.

**Continued in the next comment due to character limit**



 We thank the reviewer for their efforts in reviewing our paper and providing valuable suggestions. Below, we have addressed the concerns raised. We also request the reviewer to review the revised manuscript.

The reviewer comments are numbered and marked in bold.

**1.	Discussions on how or why this method can be faster than traditional solvers by the ratio given by the authors (e.g., 40-50x faster).**

The CoAE-MLSim method has a very similar solution algorithm as traditional solvers and we plan to evolve this work as ML based solver as we are able to address issues pointed out in the “Future work” section in the paper. The main difference being that the solutions procedure is carried out in a latent space that represents the solutions on local subdomains, where each subdomain may be considered as a group of computational elements. The number of elements in each computational domain as well as the compression achieved in determining the solution latent vector on each subdomain determines the computational speed-up. In a scenario where, each subdomain consists of a single computational element, our method will have a lesser computational performance as a traditional solver. This is because our model represents the flux conservation with non-linear functions as opposed to traditional solvers, which linearize the fluxes. As we increase the subdomain resolution, the computational performance improves, and this may be observed in Figure 9. It may be observed that the speed-up is not linear with the increase in subdomain sizes, and this is because the latent vector compression ratio decreases as subdomain sizes increase because the space of PDE solutions on a bigger subdomain is larger.

**2. Besides, it would be more solid if the author compare between using different sizes of layers (number of parameters) in the autoencoders. How is the balance between performance, accuracy, as well as generalization ability, regarding to the size of the autoencoders? How to choose the layer size regarding to different problems or some performance requirement? Summary of review: I would like to see the acceptance of this paper if the authors can provide more comparisons and discussion on the performance / accuracy regarding to the size of layers.**

The reviewer has raised a very good point here. While we have provided other ablation studies in the paper in section 4.3, here we provide results from the experiment suggested by the reviewer. We present results from the experiments to analyze the effect of flux-conservation autoencoder bottleneck layer size on the generalizability, accuracy, and performance of our approach. We choose the flux-conservation autoencoder for this analysis because it is one of the primary workhorse of our approach and is evaluated repeatedly during the iterative inferencing. We perform 6 experiments with bottleneck layer sizes, 8, 16, 32, 64, 128, 256 for an input size of 340. This experiment is carried out on a completely new use case, Steady-State Darcy’s equation on a 2-D domain. The reason for choosing this new use case is two-fold: a) the PDE in Darcy flow is relative simpler and there is a public dataset associated with it (Li. et. al, 2020) and b) it serves as another use case to validate our approach. The problem setup as well as the data to train the autoencoders is taken from Li et al., 2020. We use about 400 samples for training the flux conservation autoencoders, and 50 randomly picked samples for testing from the remaining data. In the tables below, we have provided the testing accuracy and the solve time for each encoding. The iterative algorithm convergence tolerance is stringent and set to L_2< 1e-7. Additional details and analysis related to this case are added to Appendix Section A.2.2 of the revised version of the paper.

**Table 1: Bottleneck size(Compression ratio) vs average testing error, average iterations to convergence and average solution time**

|  Bottleneck size (compression ratio)  | Avg. $L_{2}(Y_{true}-Y_{pred})/L_2(Y_{true})$  | Avg. Convergence Iterations  | Avg. Solution Time  |
|---|---|---|---|
|  8(42.5) | 0.0241  | 15  | 0.23  |
|  16(21.25) |  0.0107 |  37 |  0.57 |
|  32(10.62) |  0.0067 |  45|  0.706 |
|  64(5.32) |  0.0081 |  55 |  0.85 |
|  128(2.65) |  0.133 |  140 |  2.18 |
|  256(1.32) |  0.198 |  168 |  2.73 |

All the autoencoder networks used in this work have a similar architecture across all the numerical experiments presented here. The only hyperparameter that might vary is the size of the input layer and the size of latent vector in the bottleneck layer. These hyperparameters depend on the complexity of physics, complexity of PDE conditions, subdomains size etc. But, in all the numerical experiments the compression ratio between the input and latent vector size consistently varies between 10 and 30. 



This paper used autoencoders to approximate different phases of a PDE solver. The method proposed in this paper is novel and validated. Since this work is about accelerating the PDE algorithms, I would expect to read more discussions on how or why this method can be faster than traditional solvers by the ratio given by the authors (e.g., 40-50x faster). Besides, it would be more solid if the author compare between using different sizes of layers (number of parameters) in the autoencoders. How is the balance between performance, accuracy, as well as generalization ability, regarding to the size of the autoencoders? How to choose the layer size regarding to different problems or some performance requirement? I would like to see the acceptance of this paper if the authors can provide more comparisons and discussion on the performance / accuracy regarding to the size of layers.

This paper proposes a new ML approach called CoAE-MLSim that is a faster alternative to PDE solvers. Compared to previous ML work on this problem, it aims to be more accurate and generalize better across PDE conditions. They also aim to require fewer PDE solutions to train the model.  *Strengths:* 
- I think this is an important application area with growing interest. 
- The approach seems smart to me. In particular, I could imagine this parallelizing well, it makes sense to me that applying the same network to each subdomain would allow for simpler networks, and it makes sense to me that learning low-dimensional representations with autoencoders could help. This is especially interesting to me in the context of learning a representation for the PDE conditions that you wish to generalize over. 
- My favorite part might be the dynamic inferencing with the iterative algorithm to establish consistency, enabling the networks to focus on subdomains. 
- Being able to initialize the ML solution with coarse-grained solutions from the PDE solver to speed up convergence is also intriguing. 
- I appreciate that the paper included results on quite a few problems, covering a variety of scenarios (steady-state vs. transient and generalizing across boundary conditions, sources, geometry, and initial conditions). 

*Weaknesses:*
 
I think that:
- The comparison to baselines could be improved. 
- Some of the claims are not carefully backed up. 
- The explanation of the relationship to the existing literature could be improved. 

*More details on the above weaknesses:*

*Comparison to baselines:*

"We did not find good benchmarks to compare our unsupervised, iterative inferencing algorithm against" I think this is a slightly unfair comment. The unsupervised and iterative inferencing aspects are only positives if they have the claimed benefits, as compared to other ML methods (more accurate and better generalization). There is a lot of recent work addressing the same ML task (as mentioned in the related work section.) This paper contains some comparisons to previous work, but as I detail below, there seem to be some holes.

FCNN is by far the strongest competitor for the Laplace example in the appendix. Why is this left off of the baseline comparison table in the main paper? Further, is there any reason that FCNN couldn't have been used for the other examples? 

Why is FNO not applied to the Chip cooling (Temperature) example?

A major point in this paper is improved generalization across PDE conditions. However, I think that's hard to check when only looking at the test errors for each method. In other words, is CoAE-MLSim's error lower than UNet's error because the approach fit the training data better, or is it because it generalized better? Further, in some cases, it's not obvious to me if the test errors are impressive, so maybe it is having a hard time generalizing. It would be helpful to see train vs. test errors, and ideally I like to see train vs. val. vs. test. 

For the second main example (vortex decay over time), looking at Figures 8 and 33 (four of the fifty test conditions), CoAE-MLSim has much lower error than the baselines in the extrapolation phase but noticeably higher in the interpolation phase. In some cases, it's hard to tell how close the FNO line is to zero - it could be that CoAE-MLSim even has orders of magnitude more error. Since we can see that there's a big difference between interpolation and extrapolation, it would be helpful to see the test error averaged over the 50 test cases but not averaged over the 50 time steps. When averaged over all 50 time steps for the table on page 9, it could be that CoAE-MLSim looks better than FNO just because of the extrapolation regime. In practice, someone might pick FNO over CoAE-MLSim if they aren't interested in extrapolating in time. Do the results in the table for vortex decay back up the claim that CoAE-MLSim is generalizing over initial conditions better than FNO, or is it just better at extrapolation in time?

*Backing up claims:*

The abstract says that the method is tested for a variety of cases to demonstrate a list of things, including "scalability." The list of "significant contributions" also includes "This enables scaling to arbitrary PDE conditions..." I might have missed/forgotten something, but I think this wasn't tested?

"Hence, the choice of subdomain size depends on the trade-off between speed and accuracy." This isn't clear to me from the results. It seems like 32^3 is the fastest and most accurate? 

I noticed some other claims that I think are speculations, not backed up with reported experiments. If I didn't miss something, this could be fixed by adding words like "might." 
- "Physics constrained optimization at inference time can be used to improve convergence robustness and fidelity with physics."
- "The decoupling allows for better modeling of long range time dynamics and results in improved stability and generalizability."
- "Each solution variable can be trained using a different autoencoder to improve accuracy."
- "Since, the PDE solutions are dependent and unique to PDE conditions, establishing this explicit dependency in the autoencoder improves robustness."
- "Additionally, the CoAE-MLSim apprach solves the PDE solution in the latent space, and hence, the idea of conditioning at the bottleneck layer improves solution predictions near geometry and boundaries, especially when the solution latent vector prediction has minor deviations."
- "It may be observed that the FCNN performs better than both UNet and FNO and this points to an important aspect about representation of PDE conditions and its impact on accuracy." The representation of the PDE conditions could be why, but it's hard to say without careful ablation studies. There's a lot different about the networks. 
- Similarly: "Furthermore, compressed representations of sparse, high-dimensional PDE conditions improves generalizability."


*Relationship to literature:*

The citation in this sentence is abrupt and confusing because it sounds like CoAE-MLSim is a method from that paper instead of the new method: "Figure 4 shows a schematic of the autoencoder setup used in the CoAE-MLSim (Ranade et al., 2021a)." More broadly, Ranade et al., 2021a, Ranade et al., 2021b, and Maleki, et al., 2021 are all cited and all quite related to this paper. It should be more clear how the authors are building on those papers (what exactly they are citing them for), and which parts of CoAE-MLSim are new. (The Maleki part is clearer in the appendix, but the reader shouldn't have to check the appendix to know what is new in a paper.) 

I thought that otherwise the related work section was okay but was largely just summarizing some papers without giving context for how they relate to this paper. 


~~~~



Additional feedback (minor details, could fix in a later version, but no need to discuss in the discussion phase):
- The abstract could be clearer about what the machine learning task is that CoAE-MLSim addresses. 
- The text in the figures is often too small. 
- "using pre-trained decoders (g)" - probably meant g_u? 
- Many of the figures would be more clear if they said pre-trained solution encoders & solution decoders, since there are multiple types of autoencoders.
- The notation is inconsistent, especially with nu. For example, the notation in Figures 2 & 3 doesn't seem to match the notation in Alg 1. Then on Page 4 & Figure 4, the notation changes again.
- Why is the error table not ordered 8^3, 16^3, 32^3 like Figure 9? The order makes it harder for the reader to reason about the tradeoff. 
- Why is Err(T_max) negative sometimes? Maybe I don't understand the definition, but I would expect to use absolute value?
- I don't think the study about different subdomain sizes is an "ablation" study since they aren't removing a component of the method. 
- Figure 11: I'm guessing that the y-axis is log error, but this isn't labeled as such. I didn't really understand the the legend or the figure in general until I got to the appendix, since there's little discussion of it in the main paper. 
- "Figure 30 shows comparisons of CoAE-MLSim with Ansys Fluent for 4 unseen objects in addition to the example shown in the main paper." - probably from previous draft. Now this whole example is in the appendix, unless I missed something. 
- My understanding is that each type of autoencoder is trained separately and that there's an ordering that makes sense to do this in, so you can use one trained autoencoder for the next one (i.e. train the PDE condition AEs, then the PDE solution AE, then the flux conservation AE, then the time integration AE). This took me a while to understand though, so maybe this could be mentioned in the body of the paper. (Or perhaps I missed that!)
- It seems that the time integration autoencoder isn't actually an autoencoder if it's outputting the solution at the next time step, not reconstructing the input. 
- Either I don't understand Figure 5 or the labels are wrong. 
- It's implied in the paper (like in Algorithm 1) that the boundary conditions are encoded like the other PDE conditions. In the Appendix (A.1), it's stated that "The training portion of the CoAE-MLSim approach proposed in this work corresponds to training of several autoencoders to learn the representations of PDE solutions, conditions, such as geometry, boundary conditions and PDE source terms as well as flux conservation and time integration." But then later in the appendix (A.1.3), it's stated that boundary conditions could be learned with autoencoders but are actually manually encoded for this paper. That seems misleading.  I think that CoAE-MLSim seems like a good idea and could be a great method. However, the results are not yet strongly backing up the claims, especially that it is more accurate and generalizes better than the previous ML methods for this task. I think that a revision during the discussion phase could raise my score to 8/accept.

Meta-review: This paper proposes a new approach called CoAE-MLSim as a faster alternative to PDE solvers.  According to the authors, the main strength of the method is that it needs fewer training data, however, as pointed out by the reviewers, the factor is not significant enough in the current results. Another concern raised by the reviewers is that the iterative algorithm is much more expensive than only one forward pass of other ML methods. Furthermore, the reviewers criticized that the comparisons with baselines are not sufficient, and some claims in the paper were not backed up. The authors provided their rebuttals and had a long discussion with most of the reviewers. Some clarifications have been made, and some reviewers increased their scores accordingly. However, overall speaking, the major problems with the paper still remain, and the rebuttal has not successfully convinced the reviewers to turn to the positive side. Therefore, we cannot give a green light to this paper yet.

Decision: Reject