Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

We thank you for your thoughtful review. We appreciate that you find our work has high quality theoretical and empirical results, and that our work is novel. Furthermore, we appreciate that you found our work to be clear and detailed. We reply to your questions below.

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. Our experiments include architectures such as GINO, SFNO that are different from regular FNO and apply to climate modeling and 3D car geometries. Recently, new research such as FourCastNet 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.

We thank the reviewer for taking the time to review the manuscript and provide valuable feedback. We are glad to see that you found our experimental results to be impressive. We address your questions below.

W1: Unclear methodology. Thank you for pointing this out. We are happy to clarify that our methodology does run the entire FNO block in half-precision, with the forward FFT, the tensor contraction, and the inverse FFT all in half precision. We have now updated Section 4.2 and Figure 2 to clarify this further. Based on your question, we now also added a more extensive ablation study on which FNO block operations to keep in half-precision, which we have added to Appendix B.7. Forward FFT, tensor contraction, and inverse FFT can each be set to each full or half precision, creating a total of 8 settings. Our study shows the half-precision FNO block (our method) achieving better runtime, memory usage, and training error than all 7 other settings (this is consistent with our original Darcy flow results, in which half-precision outperforms full-precision).

The full results are as follows: Performance of each setting on Darcy Flow. The fully half-precision setting (our method) is advantageous across all the metrics.

W2: On novelty. We respectfully note that our novelty “lies in leveraging the analysis of precision error in Fourier transform, which successfully reduced the computational cost and maintained the accuracy of FNO” (as pointed out by reviewers XXPr, 5SPZ). Second, the goal of our paper is to establish a method that improves the throughput and memory footprint within the domain of scientific machine learning and particularly neural operators. In the literature prior to our work, it was not clear if mixed-precision techniques could be used in this domain, especially since classical numerical solvers traditionally suffer from stability issues and sometimes need higher than standard precision. However, our paper shows that numerical stability is quite promising (with up to 50% reduction in memory and throughput) and our work generally lays a theoretical and empirical foundation in this area.

Q1: discretization convergent. Thank you for this question. Yes, this paper gives a formal proof of discretization convergence for neural operators, including FNO (Section 9.2, Theorem 8). We have now added this reference in the paper. This shows that the approximation error is approximately constant as the resolution increases. As you mentioned, we also empirically verify this property via our zero-shot super-resolution experiments (Table 1).

We would also like to highlight a few other updates we made to our paper

• Improved theoretical bounds: while we originally proved bounds using the Fourier basis, we have now proven matching results for arbitrary basis functions (Appendix A.2).
• 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 very much, once again, for your generally positive view of our work and your excellent suggestions. If you find our responses satisfying, we respectfully ask that you consider increasing your score. On the other hand, we would be very happy to answer any follow-up or additional questions you have.

Thank you for your positive feedback. We appreciate that you find our work is well-written and covers an application that is significant. We reply to your points below.

W1: Illustration of empirical characterization. Thank you for the suggestion; this is a great idea! Following your suggestion, we added new experiments to Appendix A.2 that directly compare our theoretical resolution and precision upper and lower bounds, to the empirical resolution and precision error. We find that the discretization error is higher than the precision error, as expected by our theory, which is additional confirmation of the correctness of our theory and the applicability of our results. These experiments also make it easy to see qualitative features of our theoretical results, such as how the discretization error decays as the grid size increases, while the precision error stays constant.

W2: Mixed-precision with FP8. Thank you, this is another interesting suggestion. Despite there being no native support for FP8 in pytorch, we are able to simulate FP8 training through two approaches:

1. Calibrating network activations from FP16 to FP8 using two format options, E5M2 and E4M3. Unfortunately, for both formats, the training instantly fails in the first iteration due to FNO producing inf values. While prior work has used FP8 for GEMM and convolution operations (as in this paper, to the best of our knowledge, we are not aware of any work that computes Fourier transforms in FP8.
2. Clipping out-of-range activations to the upper and lower limits of FP8. To mitigate training divergence, we constantly check for out-of-range values and clip them to the extrema of FP8 formats. Although training could proceed as normal, the FP8-FNO is unable to converge as shown in the Darcy Flow training curve comparing FP8 and FP16 training. Training curves are added in Appendix B.11.

Furthermore, the worse result of FP8 is actually predicted by our theory. In our paper, Theorem 3.2 showed that for FP16, precision error is comparable to discretization error when plugging in $\epsilon=10^{-4}$, the multiplicative precision of FP16. However, for FP8 $\epsilon>10^{-2}$, and we can no longer prove that the precision error is lower than the discretization error. We have added this discussion to Appendix B.11.

Finally, we would also like to highlight a few other updates we made to our paper

• Improved theoretical bounds: while we originally proved bounds using the Fourier basis, we have now proven matching results for arbitrary basis functions (Appendix A.2).
• 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 favorable assessment of our work and your excellent suggestions. We would be very happy to answer follow-up questions or follow any additional suggestions.

Q2. Common numerical stability techniques. Thank you for the question! While the main source of instability in mixed-precision neural nets is the gradient updates, we find that for neural operators, there is even more numerical instability within the forward FFT operations. Therefore, the standard techniques such as loss scaling, gradient clipping, and delayed updates, are focused on reducing the dynamic range of the gradients, but for neural operators, we are interested in addressing the numerical instability of the FFT. We now further verify these claims by significantly extending our original ablation study, which we have also added to Appendix B.5.

• Loss scaling: we now confirm in Appendix B.5 that the typical loss scaling (GradScaling) is ineffective, due to the reason above. Instead, we tried scaling just before each FFT. However, this fails due to the wide range of values. For example, if we typically have a few outliers of $10^6$, while most of the numbers are within the range $[-1,1]$, then we would expect scaling to perform much worse than tanh, since tanh negates the outliers while keeping the same range in the rest of the values. We confirm this in a new ablation study in Appendix B.5.
• Gradient clipping: this also does not work due to the above reason. In our new ablation study, we clip gradients to a default value of 5 during training, yet it does not resolve the NaN issue. Instead, we tried two different forms of clipping the activations just before the FFT (from our original submission), in Appendix B.5.
• Normalization: note that we already do have normalization in our models. In our new ablation study, we added another normalization layer before the FFT, which delays the NaN occurrence only for 2 epochs of training. Additionally, this incurs a 30% GPU memory overhead due to intermediate values stored for the backward pass. Our tanh approach is more lightweight and effective.
• Delaying updates: for gradient accumulation, we tried updating weights every 3 batches, which did not resolve the NaN problem.

Finally, we would like to highlight a few other updates we made to our paper

• Improved theoretical bounds: while we originally proved bounds using the Fourier basis, we have now proven matching results for arbitrary basis functions (Appendix A.2).
• 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 very much, once again, for your positive view of our work and your excellent suggestions. We would be happy to answer any additional questions. (2/2)

Thank you for your favorable and detailed review. We appreciate that you find our work very interesting and provide a strong contribution to the research community. Further, we are glad that you found that our work has the potential to advance the efficiency and scalability of neural operators in downstream applications. We address each of your questions below:

W1: Justification for pre-activation before each FFT. This is a great question. In fact, pre-activation is predicted by our theory, and we have now updated our paper to point this out. Given a function $v(x)$ defined on $[0,1]$, as long as $|\max_x v(x)-\min_x v(x) |>1$, then $\text{tanh}(v(x))$ has both a lower $L_\infty$ norm and a lower Lipschitz constant compared to $v(x)$, both of which are terms that show up in the upper bounds of Theorem 3.1 and Theorem A.1. In other words, tanh decreases the magnitude and gradient of the functions, which decrease the theoretically-predicted discretization and precision errors. Finally, as you say, we also have an ablation study that empirically motivates tanh in Appendix B.8.

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.

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.

(1/2)

Dear Reviewers,
Thank you once again for your suggestions. In addition to the above updates, we have also just added one more section (Appendix B.12) that presents an even more detailed ablation study for our method. This ablation study shows the benefit of (1) decomposition of einsum into sub-equations, (2) einsum path caching, (3) our greedy tensor contraction procedure, and (4) both weights and inputs in half-precision.

Please let us know if you have any questions or follow-up comments to our responses, and we would be very happy to reply before the deadline tonight; thanks again!