Maximal Update Parametrization and Zero-Shot Hyperparameter Transfer for Fourier Neural Operators

The authors apply a framework for transferring hyperparameters from small to large FNO models.

I have not verified the details of the theoretical proofs.

The datasets used in this study appear overly simplistic. The necessity for FNO models with billions of parameters in these tasks is questionable. More challenging datasets, such as those involving multiple mixed equations for training PDE foundation models [1,2], would provide more meaningful scenarios for investigating hyperparameter tuning in large-scale models. [1] Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers [2] PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations

This work may contribute to hyperparameter tuning research.

N/A

Yes, the methods and evaluation criteria make sense for the problem of hyperparameter transfer in FNOs for PDE solving.

• µP as a Method: The use of µP as a foundation for hyperparameter transfer is well-motivated.The core idea of maintaining consistent training dynamics across model sizes is directly relevant to the goal of transferring hyperparameters. The theoretical derivation provides a principled basis for the proposed scaling factors.

• Algorithm 1: The µTransfer-FNO algorithm is a straightforward and logical application of the µP theory. It clearly outlines the steps involved in transferring hyperparameters from a small proxy model to a larger target model.

• Choice of PDEs: The selected PDEs (Burgers' Equation, Darcy Flow, and Navier-Stokes Equation) represent a reasonable range of complexity and dimensionality. They are commonly used as benchmarks in the FNO literature, allowing for comparison with existing work. While a broader set of PDEs (PDEBench) would be even better, the chosen set provides a good starting point.

• Evaluation Metrics: Using the L2 relative test error is a standard and appropriate metric for evaluating the performance of PDE solvers. It directly measures the accuracy of the solution obtained by the FNO.

• Comparison with Baseline: Comparing µTransfer-FNO against directly tuning the large model provides a relevant baseline.This allows for assessing the computational cost savings and potential performance gains of the proposed method. However, as mentioned in the previous response, a stronger baseline would involve using the same expanded hyperparameter search space for both methods.

I've reviewed the provided proof of Theorem 3.5, which establishes the µP for FNOs. While the overall structure of the proof seems reasonable, there are some specific points and potential issues:

• Spectral Norm Calculation: The proof relies heavily on calculating the spectral norm of the kernel integral operator $\mathcal{K}$. The argument that this norm is equivalent to the maximum absolute value of the parameters r seems plausible given the diagonal structure of $\mathcal{R}$ after the Fourier transform. However, the interaction between the multidimensional Fourier transform, the truncation operator $\mathcal{T}_K$, and the parameter tensor $\mathcal{R}$ could be complex, and a more detailed justification of this step would be beneficial.

• High-Probability Bound: The proof invokes a standard result from high-dimensional probability to bound the maximum of $K^d$ sub-Gaussian random variables. While this is a common technique, the specific constants and assumptions underlying this result need careful checking against the properties of the parameters $r$. The paper assumes these are sub-Gaussian, which needs to be verified in practice.

• Simplification to $m=1$: The proof simplifies the analysis by assuming a hidden dimension $m=1$. While extending to general $m$ is mentioned as straightforward, explicitly showing this extension, or at least providing a sketch, would strengthen the proof. The interaction between $m$ and $K$ in the scaling factors could be non-trivial.

• Discretization Effects: The proof works with the discretized version of the FNO. While this is necessary for practical implementation, the impact of discretization on the theoretical results is not explicitly discussed. Ideally, the proof should connect back to the continuous formulation of FNOs.

I've reviewed the experimental designs and analyses presented in the paper. While they provide support for the claims, there are some areas where the soundness and validity could be improved:

• Hyperparameter Search Space: The paper doesn't explicitly define the hyperparameter search space $\Xi$ used in Algorithm 1. Knowing the range and granularity of the search for each hyperparameter (learning rate, batch size, $\beta_2$) is crucial for interpreting the results. A larger search space could lead to better results regardless of µTransfer-FNO, so specifying the search space is essential for a fair comparison.

• Baseline Comparison: As mentioned previously, the comparison with the baseline of directly tuning the large model is not entirely fair. µTransfer-FNO benefits from exploring a potentially larger search space with the smaller proxy model. A stronger baseline would involve using the same expanded search space for tuning the large model directly, even if it's computationally more expensive. This would isolate the benefit of the transfer method itself.

• Limited Number of Runs: The paper presents results for single runs of each experiment. Given the stochastic nature of neural network training, reporting results averaged over multiple runs (e.g., $10-20$) with standard deviations or error bars would provide a more robust evaluation and account for potential variability.

• Computational Cost Analysis: The analysis of computational cost is somewhat limited. Table 1 only provides relative training costs ($1\times$, $1.38\times$, $0.30\times$). Reporting absolute training times or FLOPs would be more informative. Furthermore, the cost of the hyperparameter search itself is not factored into the comparison. A more comprehensive analysis should include the total cost of both the search and the final training.

• Lack of Ablation Study: An ablation study investigating the individual contributions of the initialization variance scaling and the learning rate scaling would provide further insights into the effectiveness of the proposed µP. This would help determine the relative importance of each component.

• Details of Data Generation: While the paper describes the PDEs used, more details about the data generation process would improve reproducibility.

• Code Availability: Providing the code used for the experiments would significantly enhance reproducibility and allow for independent verification of the results.

This paper's main contributions relate to various areas within the broader scientific literature:

• Fourier Neural Operators (FNOs) and Operator Learning: The paper directly builds upon the existing literature on FNOs for solving PDEs (Li et al., 2021).

• Physics-Informed Neural Networks (PINN): The paper extends its method to Physics-Informed Neural Operators (PINOs) (Li et al., 2024), demonstrating its applicability beyond standard supervised learning for PDEs.

• Maximal Update Parametrization (µP) and µTransfer: The core theoretical contribution relies heavily on the µP framework introduced by Yang & Hu (2021) and the concept of µTransfer (Yang et al., 2022).

• Hyperparameter Transfer Learning: The overall goal of the paper is to enable efficient hyperparameter tuning for large models by transferring knowledge from smaller models.

The paper covers the most relevant literature regarding FNOs, PINNs, and µP/µTransfer. However, a few potential areas could benefit from additional discussion or referencing:

• Theoretical Analysis of Hyperparameter Transfer: While the paper provides a theoretical justification for µP in FNOs, it could benefit from discussing any existing theoretical work on hyperparameter transfer in general. Are there any theoretical guarantees or bounds on the effectiveness of transferring hyperparameters across different model sizes or architectures? Connecting to this broader theoretical literature would strengthen the paper's contribution.

• Alternative Parametrizations: The paper focuses on µP, but other parametrization schemes might exist for FNOs. Discussing potential alternatives and their potential advantages or disadvantages compared to µP would provide a more complete picture. For example, are there parametrizations specifically designed for different optimization algorithms or different types of PDEs?

• It makes sense to try to increase the number of Fourier modes and the proposed transfer HPO method seems effective. To do so, the proposed method scales the kernel integral operator parameters as the number of modes $K$ is increased so that the optimal hyperparameters stay approximately the same across model sizes. This motivates HPO on the small model and then directly transferring to larger FNOs.
• The result cover a good range of PDEs, i.e., Burgers, Darcy Flow and the challenging Navier-Stokes equations. I would like to see a larger number of equations tested. See the comprehensive Neural Operator benchmark in Saad et al., "Guiding continuous operator learning through Physics-based boundary constraints", ICLR 2023.
• The authors show that the proposed transfer learning preserves the same optimal learning rate, training batch size and optimizer configurations across increasing model sizes.
• With the proposed approach, the method is able to scale to large FNOs with nearly 1B parameters with better accuracy and only 0.3x training compute.
• In Figure 2, even with the proposed approach the optimal learning rates are not exactly the same and seem slightly shifted but less than the base especially in Figure 2a-b. Is there a quantitative metric to better measure this difference?

• The main theoretical result is presented in Theorem 3.5 on calculating the scaling rate for kernel up the FNO kernel integral operator. The theorem shows the abc-parametrization is a $\mu$P of FNO with the Adam optimizer and scales according to 1 / sqrt(d), where d is the dimensionality. The proof is detailed and provided in the appendix with a proof outline of the main concepts in the main body.
• The authors highlight a difference in their proof than in the standard $\mu$P proofs that the past proofs rely on the CLT to measure the average of random variables, whereas there is a technically difference in the structure of FNO and the kernel integral operator that leads to the analysis of the maximum of the random variables. It is good that the authors clarify this difference.
• It is also good that the proposed scaling function in the proof is not resolution dependent so that the desired resolution invariance property of NOs is not broken with this methods - do the authors have proof of this, i.e.,that resolution invariance still holds?

• Nice that experiment section is organized according to various questions to show that aforementioned theory also holds empirically.
• Please clarify that Burgers' in 4.1 is actually viscous Burgers' for $\nu \ne 0$, which is simpler to solve than the sharp shock solution case where $\nu = 0$ and there is no artificial diffusion from the viscosity term.
• Good that the authors consider 1D and 2D problems and also vary mapping the initial condition to solution in Burgers' and PDE coefficient to the solution in Darcy Flow.
• Since one of the main motivations of the proposed method is for large-scale problems, I think 3D spatial cases should also be used because that when $d=3$ is when the FNO becomes very computational expensive. It is good that 3D Navier-Stokes is tested. Is this FNO-3D because it is the standard 2D space + time test case? I think 3D space should also be considered. An example practical, real-world 3D Car Surface Pressure Prediction test case is provided in Ma et al., "Calibrated uncertainty quantification for operator learning via conformal prediction", 2024 and Li et al., " Geometry-informed neural operator for large-scale 3d pdes", 2023. Both of these works should be cited as well.

• FNOs have shown large impact on solving PDE problems with ML. This paper focuses on more challenging PDEs where larger parameter FNOs are needed with more Fourier modes. The authors in the introduction go directly into discussing the specifics of FNO and the effect of its modes. I think the authors should first motivate solving PDE problems, why ML methods have been developed to solve these problems, mention the three main method classes, e.g., PINNs, Neural Operators and MeshGraphNets and then state that this paper will focus on FNOs. The authors can tell clearly identify the more challenging PDEs with regular size parameters FNOs struggle to motivate these larger-parameter FNO models which require HPO.
• The authors mention the $\mu$ P (Yang & Hu (2021) and $\mu$Transfer methods (Yang et al. (2022)). Please clearly differentiate the novelty of the proposed approach. Does $\mu$Transfer-FNO just directly apply these methods to FNO? If so, please explain the technical challenges of applying them to the FNO architecture.
• My main concern is that the proposed method is too similar to the past Yang & Hu (2021) and Yang et al. (2022) in literature and hence the novelty needs to be clarified.
• The authors state that the number of hidden dimensions $m$ and model depth $L$ are fixed in this work and have been studied in Yang & Hu (2021) and Yang et al. (2022). Please briefly summarize their effect here or in an appendix so that this work can be fully understood end-to-end without relying on these prior works for key concepts.

• It is good that the authors include the known boundary conditions in the problem definition in Eqn. 1. There is missing reference to Saad et al., "Guiding continuous operator learning through Physics-based boundary constraints", ICLR 2023 that enforces boundary conditions as an exact constraint.
• The authors mention that FNO is most commonly NO. They should state the reason is because the FFT is very computational efficient to compute the kernel vector $\mathcal{K}x$ products. Having said that, the paper is missing references to other neural operators with different bases, e.g., Gupta et al., “Multiwavelet-based Operator Learning for Differential Equations”. In: Advances in Neural Information Processing Systems. Vol. 34. PMLR, pp. 24048–24062, 2021, Li, Zongyi, Nikola Kovachki, et al. (2020b). “Multipole Graph Neural Operator for Parametric Partial Differential Equations”. In: arXiv preprint arXiv:2006.09535.
• It is good that the authors also compare to PINO. I think the authors should be a bit careful referring to the PINNs loss as "more advanced training techniques". Both the PINO paper and Saad et al., ICLR 2023 show that in several cases the unconstrained FNO outperforms PINO. Also there have been several instabilities with PINNs training that have been reported in the literature that are not cited, e.g., Krishnapriyan et al., "Characterizing possible failure modes in physics-informed neural networks", Advances in neural information processing systems 34, 26548-26560, 2021. This looks like it is also observed by the authors own training of PINO in Figure 2(c) with the spike and oscillations in the loss function.
• References to transfer learning methods other than $\mu$ P (Yang & Hu (2021) and $\mu$Transfer methods (Yang et al. (2022)) are missing.

Yes. The experiments are mainly based on the Burgers’ Equation (1D), the Darcy Flow Equation (2D), and the Navier-Stokes Equations (3D), which are standard examples used in the original FNO paper. The authors studies the performance of $\mu$Transfer-FNO on the finetuning of multiple parameters like the learning rates and batch sizes with respect to different number of Fourier modes. Experiments on other variants of FNO like PINO (Physics-Informed Neural Operator) are also included to justify the efficacy of the proposed methodology.

The only theoretical claim in the paper is Theorem 1.1 (informal version of Theorem 3.5). Its proof in the supplement has been verified to be correct.

Yes, please refer to the "Methods And Evaluation Criteria" section above.

This work is based on the Maximal Update Parametrization ($\mu$P) framework proposed in the series of work on tensor programs, which is originally used in the studies of MLPs and Transformers. However, this paper should be mainly posited within the literature (FNOs, DeepOnets and their variants) on operator learning and neural PDE solvers.

Below is one preprint that probably needs to be discussed in the paper is [1]. Specifically, the main result in [1] discussed the connection between neural operators and transformers (the attention architecture), which is also partially addressed in the cited paper [2]. The authors should discuss how the $\mu$Transfer-FNO methodology proposed in this paper differs from the $\mu$Transfer framework originally proposed for transformers, since neural operators can be linked with transformers under certain circumstances.

References:

[1] Calvello, Edoardo, Nikola B. Kovachki, Matthew E. Levine, and Andrew M. Stuart. "Continuum attention for neural operators." arXiv preprint arXiv:2406.06486 (2024).

[2] Kovachki, Nikola, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. "Neural operator: Learning maps between function spaces with applications to pdes." Journal of Machine Learning Research 24, no. 89 (2023): 1-97.