Paper: Guaranteed Approximation Bounds for Mixed-Precision Neural Operators (ICLR-2024)

Reviewer comment (Novelty):
The major doubt for this paper is regarding the significance. While it is dedicated to improving FNO, an important architecture of neural operators, FNO is not the only one and arguably the best one.

Author response:
**Significance.**
Thank you for motivating this discussion! We give our answer in two parts; first, by motivating the large class of high-performing models that use FNO as a base, and second, by presenting new theoretical results that generalize our original results beyond FNO.

First, our method and findings hold true for any model that is based on FNO, including a large class of state-of-the-art models. Although initially solving classical PDEs, the versatility of FNOs for real-world applications is [gaining recognition](https://www.technologyreview.com/2020/10/30/1011435/ai-fourier-neural-network-cracks-navier-stokes-and-partial-differential-equations/). Our experiments include architectures such as GINO, SFNO that are different from regular FNO and apply to [climate modeling](https://arxiv.org/abs/2306.03838) and [3D car geometries](https://arxiv.org/abs/2309.00583). Recently, new research such as [FourCastNet](https://arxiv.org/abs/2202.11214) have been unveiled that leverages FNO variants and targets high-impact applications, including weather and climate modeling. 

Second, we still agree that our work could be even more impactful if we generalize our results. Therefore, we have now added **new theoretical results**, that generalize our original bounds using the Fourier basis, to arbitrary basis functions (Theorems A.1 and A.2 in Appendix A). This extends our results beyond Fourier neural operators to a larger class of integral operators, that can include e.g. Wavelet bases or graph neural operators. We also have initial simulation results which confirm our new theoretical results, in Appendix A.3.

Finally, we would also like to highlight a few other updates we made to our paper
- We give new experiments that directly compare our theoretical resolution and precision upper and lower bounds, to the empirical discretization and precision error (Appendix A.2). We find that the errors behave as expected by our theory, which is additional confirmation of the correctness of our theory and the applicability of our results.
- We added another experiment setting, the Ahmed body CFD dataset. On this high-resolution 3D task, we achieve 38.4% memory reduction while the error stays within 0.05% of the full-precision model.

Thank you once again, for your excellent point. Please let us know if you have any follow-up questions or further suggestions. We would be happy to continue the discussion.

Reviewer comment (Novelty):
If FNO is yet to be a dominant model for large scale applications of neural operators, the significance of this work is limited.

Author response:
**Significance.**
Thank you for motivating this discussion! We give our answer in two parts; first, by motivating the large class of high-performing models that use FNO as a base, and second, by presenting new theoretical results that generalize our original results beyond FNO.

First, our method and findings hold true for any model that is based on FNO, including a large class of state-of-the-art models. Although initially solving classical PDEs, the versatility of FNOs for real-world applications is [gaining recognition](https://www.technologyreview.com/2020/10/30/1011435/ai-fourier-neural-network-cracks-navier-stokes-and-partial-differential-equations/). Our experiments include architectures such as GINO, SFNO that are different from regular FNO and apply to [climate modeling](https://arxiv.org/abs/2306.03838) and [3D car geometries](https://arxiv.org/abs/2309.00583). Recently, new research such as [FourCastNet](https://arxiv.org/abs/2202.11214) have been unveiled that leverages FNO variants and targets high-impact applications, including weather and climate modeling. 

Second, we still agree that our work could be even more impactful if we generalize our results. Therefore, we have now added **new theoretical results**, that generalize our original bounds using the Fourier basis, to arbitrary basis functions (Theorems A.1 and A.2 in Appendix A). This extends our results beyond Fourier neural operators to a larger class of integral operators, that can include e.g. Wavelet bases or graph neural operators. We also have initial simulation results which confirm our new theoretical results, in Appendix A.3.

Finally, we would also like to highlight a few other updates we made to our paper
- We give new experiments that directly compare our theoretical resolution and precision upper and lower bounds, to the empirical discretization and precision error (Appendix A.2). We find that the errors behave as expected by our theory, which is additional confirmation of the correctness of our theory and the applicability of our results.
- We added another experiment setting, the Ahmed body CFD dataset. On this high-resolution 3D task, we achieve 38.4% memory reduction while the error stays within 0.05% of the full-precision model.

Thank you once again, for your excellent point. Please let us know if you have any follow-up questions or further suggestions. We would be happy to continue the discussion.

Reviewer comment (Theory):
Similar theoretical justification should have been provided for the learning rate schedule.

Author response:
**W2: Justification for learning rate schedule.**
First, we would like to clarify that you meant the precision schedule (our contribution), instead of the learning rate schedule (well-known in the literature)? The precision schedule follows the following intuition: in the early stages of training, it is okay for gradient updates to be more “coarse”, since the gradient updates are larger overall, so full precision is less important. However, in the later stages of training, the average gradient updates are much smaller, so full precision is important. In fact, we find that at the end stages of training, the magnitudes of some gradients in the FNO block fall below $10^{-4}$, which is below the precision constant of float16. This aligns with the empirical results of our precision schedule experiments: full precision performs well at the last stages of training.  We have now updated our manuscript with this explanation.

Reviewer comment (Experiments):
Why is it that the standard mixed precision training used for training ConvNets and ViTs is ineffective for training Fourier Neural Operators?

Author response:
**Q1: Standard techniques ineffective for FNO.**
This is a great question; we have now updated the paper to clarify this. The standard techniques are designed for common operations such as convolution and transformers, however, the most memory-intensive operations of FNO are in the spectral domain, which are not supported by AMP and require a different procedure that takes into account both casting to half precision as well as several vectorization operations (moving to and from complex-valued representation). Since complex-valued operations are not supported in AMP, it simply skips the most important operations inside the FNO block, and therefore only results in a small improvement for FNO. We introduce the first mixed-precision training method for the FNO block by devising a simple and lightweight greedy strategy to find the optimal contraction order, taking into account the vectorization operations for each intermediate tensor. Furthermore, numerical stability methods also must be handled differently for neural operators, since the typical methods fail, which we discuss below in Q2.