FNP: Fourier Neural Processes for Arbitrary-Resolution Data Assimilation

Dear reviewer,

Thank you very much for your thorough review, highly constructive comments, and feedback! We sincerely appreciate your recognition of the effectiveness and significance of our method. Below, we will address each of your questions and concerns in sequence.

The ERA5 dataset is global meteorological data with limited high-frequency information. Validation on datasets with richer high-frequency information (e.g., HRRR) could better demonstrate the effectiveness of the Neural Fourier layer and provide more confidence.

Excellent suggestions! Conducting experiments on the HRRR dataset can better demonstrate the effectiveness of NFL and enhance the experimental content of this paper. Unfortunately, we regret that we are unable to download the HRRR dataset and train models on it within the limited time frame, resulting in the unavailability of experimental results here. We will strive to include experiments on the HRRR dataset in our future work.

Nevertheless, we have conducted additional experiments to showcase the effectiveness of NFL from a different perspective. Table A from the PDF in global response provides ablation study on the generalization of different modules at various resolutions. As the resolution increases, the impact of NFL on performance becomes more pronounced. Similarly, in terms of visual perception, FNP's ability to capture high-frequency information strengthens with higher resolutions. This indirectly reflects the correlation between NFL and the high-frequency information gain brought by our method.

The role of FNP in weather forecast tasks is not thoroughly discussed or experimentally validated.

Your feedback is highly professional! Data assimilation aims to improve weather forecasting results by reducing initial errors. Therefore, it is crucial to explore the impact of different methods on forecast error reduction. Figure A from the PDF in global response provides results on the forecast RMSE improvement of z500 variable over the next ten days through data assimilation, where lead time 0 corresponds to the reduction of initial errors. Darker colors indicate stronger improvements, meaning a greater reduction in forecast errors compared to not using data assimilation.

Similar to the results of data assimilation, FNP consistently achieves state-of-the-art results in most cases, with its advantage becoming more pronounced as the resolution increases. Moreover, FNP is the only model that strictly enhances forecast improvement with increasing resolution and observational information. Additionally, apart from the accuracy of initial values affecting forecast errors, the physical characteristics of the initial states (e.g., physical balance) also influence the rate of forecast error growth. FNP demonstrates greater improvements in forecast errors at all lead times compared to improvements in initial errors, indicating that FNP not only reduces forecast errors but also slows down the growth rate of forecast errors. Other models do not exhibit the same trend, further highlighting the superior characteristics of the initial states produced by FNP. We appreciate your insights once again and will include this part in the appendix of the final version.

We sincerely hope that we have addressed your concerns and questions and look forward to your reading and response. If you have any further questions, please feel free to let us know and we will do our best to answer them.

Dear reviewer,

Thank you very much for your detailed review, thoughtful comments, and feedback! We sincerely appreciate your recognition of the effectiveness of our method and its relevance to the research field. At the same time, we deeply regret and apologize for any confusion our writing may have caused you. Below, we will address each of your concerns in sequence.

I found the paper rather hard to follow. Concepts like "spatial variable functional representation" and "dynamic alignment and merge" are not well explained. As a result, even having previous experience of Fourier neural operators it is difficult to get a sense of what the network actually does. The authors refer the reader to previous papers on the subject, but I feel the current paper would need to better explain these ideas to a reader from outside the specific niche of research to be accessible for the NeurIPS audience.

Spatial-variable functional representation: All neural process methods share a common core idea of encoding contextual information to model a representation $R$, which is then used for decoding the target values. In CNP, $R$ is a static global representation, while in ANP, $R$ is dynamic and related to the absolute position. In ConvCNP, the encoder's mapping is translation equivariant, resulting in $R$ that is dependent on relative position, which is termed as functional representation (FR). FNP adopts the ConvCNP concept, utilizing SetConv to encode context coordinates and values separately to capture density and signal, which are then concatenated and fed into deep feature extraction module composed of NFL blocks to model FR. Tailored to meteorological data characteristics, we model a distinct spatial FR for each variable and a variable FR for all variables instead of a mixed FR. This is called spatial-variable functional representation, and the modeling process for each FR follows the aforementioned procedure.

This decoupled modeling offers several advantages. Firstly, the individual spatial FRs diversify and specialize the model's spatial relationship modeling, as different variables may exhibit distinct spatial patterns. Secondly, the variable FR can capture the interrelationships among different variables. Modeling the correlations in these two dimensions is also crucial in traditional data assimilation techniques (i.e., the role of error covariance). Thirdly, it clarifies the objectives of each FR and reduces complexity, thereby easing model training, accelerating convergence, and enhancing performance. Finally, this decoupling allows for lower data embedding dimensions for each FR (i.e., number of channels), which can significantly reduce computational resource consumption. In the ablation study in Table 3, the FLOPs with and without SVD are 67.872G and 167.932G, respectively, corresponding the data embedding dimensions of 128 and 256. The SVD makes the model achieve better performance with lower FLOPs, thus enhancing computational efficiency.

Dynamic alignment and merge: The DAM aligns the FR of the background and observations in the spatial dimension and then merges them. Alignment is achieved through interpolation operations, ensuring that the spatial dimensions of observations are consistent with the background in all circumstances for subsequent information fusion. The merge is implemented by selecting based on similarity to shared features, with the calculation for similarity and selection rules provided in the formula in Section 3.3.

Figure 4 and section 4.2 refer to fine-tuning the model and present better results for the fine-tuned version. What does the fine-tuning involve? I don't get a good picture of this from the paper.

In our experiments in Sections 4.2 and 4.3, we present the performance of models without fine-tuning and with fine-tuning. "Without fine-tuning" refers to training in one setting and directly testing in another, showcasing the model's generalization capability. "Fine-tuning" involves further training for a certain number of epochs in the testing setting, demonstrating the model's optimal performance in that specific context.

Specifically, in Section 4.2, "without fine-tuning" denotes assimilating observations all at 1.40625° during training, while assimilating observations at 0.703125° and 0.25° in testing. The performance after fine-tuning refers to the testing results after the model is further trained at the corresponding resolutions.

In Section 4.3, "without fine-tuning" means using the model weights trained from the data assimilation task directly for reconstructing the observational information. We infer the FR of the observations using only the encoding weights of the observation branch, and then infer the output based on the FR using the decoding weights. The performance after fine-tuning refers to the testing results after continuing to train the weights of these two parts in observational information reconstrution task.

We sincerely hope that we have addressed your concerns and questions and look forward to your reading and response. If you have any further questions, please feel free to let us know and we will do our best to answer them.

Dear reviewer,

Thank you for your question. We believe we have adequately addressed all of the issues you raised in your review. We would like to emphasize that data assimilation is a very important area that has taken AI-based weather forecasting one step further toward practical operational applications. Compared to existing work, our approach has greater potential and a wider range of applications to advance the meteorology field and benefit society. If your question has been satisfactorily addressed, we hope you will review and possibly update your score. We are willing to address any additional questions or concerns you may have.

Thank you.

Dear reviewer,

Thank you very much for your thorough review, highly constructive comments, and feedback! We greatly appreciate your positive reception of our work and recognition of its practical value. At the same time, we deeply regret any confusion or difficulties you may have encountered regarding our methodological expressions and terminologies. We will address each of your points in the following responses.

It would be beneficial if the authors could provide the precise mathematical formalisms for key modules such as the “Spatial-variable decoupled representation,” the Neural Fourier layer, and DAM. Although one can how these modules work, a definitive mathematical expression would greatly alleviate any confusion.

Excellent suggestion! Following the description in Figure 1, we will endeavor to provide mathematical formulations for the key modules in FNP.

• Functional representation (FR): $FR(x^c,y^c)=Concate(SetConv(x^c),SetConv(y^c))$

• Spatial-Variable decoupled representation (SVD): As described in Section 3.2, we model a separate spatial FR for each variable and a variable FR for all variables instead of a mixed FR, with each FR following the formula above.

• Neural Fourier layer (NFL): $NFL(·)=\mathcal{F}^{-1}(Linear(\mathcal{F}(·)))+Conv(·)+Identity(·)$

• Dynamic Alighment and Merge (DAM): The DAM module aligns the functional representations of the background and observations in the spatial dimension and merges them. The alignment is achieved through interpolation operations, while the merge is implemented by selecting based on similarity to shared features, with the calculation for similarity and selection rules provided in the formula in Section 3.3.

Conditional points/conditional domain. Do you mean "context" as supposed to "target"?

Conditional points and conditional domain refer to the context points and context domain, corresponding to the input of the model, whlie target points and target domain correspond to the output of the model.

Variable dimension. Do you mean the dimension across different meteorological variables at the same spatial location?

Yes, it is also the channel dimension, because the data of different variables are concatenated on the channel dimension.

What do you mean by "dynamics" at line 175.

It means the ability to support inputs of varying sizes, i.e., assimilating observations with different resolutions.

Line 228, the term RMSE is a statistical metric used way beyond geospatial analysis and atmospheric science. So you might want to use something else to refer to the special "latitude-weighted" version.

You are right, in atmospheric science, we usually use RMSE to refer to the latitude-weighted version rather than the standard version. We appreciate and respect your rigor, and we will use the term "WRMSE" to replace "RMSE" in final version.

What do you think is the main reason that FNP outperforms ConvCNPs?

The main difference between FNP and ConvCNP lies in our tailored design for data assimilation task. The SVD reduces the complexity of model training and computational resource consumption while achieving better performance. Unified coordinate transformation and dynamic alignment ensure that the model comprehends the spatial correspondences between data of different resolutions, while dynamic selection and merge mechanism enhances the effectiveness of information fusion. The visualization in Figure 2 demonstrates the enhancement in the model's ability to extract deep features and capture high-frequency information through the NFL. In our supplementary experiments (see Table A from the PDF in global response), as the observational resolution increases, the impact of the NFL in ablation study on generalization becomes increasingly evident, further validating its effectiveness.

What do you think is the most important innovation of FNP for the data assimilation problem, given that ConvNPs (https://proceedings.nips.cc/paper/2020/file/5df0385cba256a135be596dbe28fa7aa-Paper.pdf) already has applications for Environmental Data.

FNP has the capability to assimilate observational data with arbitrary resolution without the need for prior interpolation of observations that typically have higher resolutions. It avoids information loss, thereby enhancing the performance of data assimilation. The types and resolutions of observational data in practical applications are highly diverse and complex. FNP can directly assimilate such data without pre-processing, offering greater practical value. Additionally, the challenge of dealing with high-dimensional data is a common issue faced by traditional methods and AI approaches. As our experiments have shown, the outstanding generalization ability of FNP enables training at low resolutions and direct application to high resolutions. This not only significantly reduces computational resource consumption but also provides a pathway for data assimilation at higher resolutions.

Why the neural Fourier layer (NFL) is more computational efficient than standard conv layers, if NFL contains three branches, one one branch is Conv operations?

In the ablation study presented in Table 3, FNP utilized a structure encoded with 4 NFL blocks, amounting to 67.872G FLOPs. In contrast, the FNP without NFL opted for a replacement comprising 12 convolutional blocks (consistent with the official implementation in [1]), totaling 167.932G FLOPs. FNP's ability to achieve better performance with lower computational complexity demonstrates its higher computational efficiency.

• [1] Yann Dubois, Jonathan Gordon, and Andrew YK Foong. Neural process family. http://yanndubs.github.io/Neural-Process-Family/, 2020.

We sincerely hope that we have addressed your concerns and questions and look forward to your reading and response. If you have any further questions, please feel free to let us know and we will do our best to answer them.

Dear reviewer,

Thank you very much for your detailed review, thoughtful comments, and feedback! We appreciate your recognition of the effectiveness of our method, while deeply regretting and apologizing for any confusion or concerns we may have caused you regarding the motivation and design of our approach. Below, we address each of your questions in turn.

Since the authors aim to address arbitrary-resolution data assimilation with good generalization ability, it would be beneficial to clearly highlight the key observations or components for this problem upfront.

Weather forecasting is of paramount importance to both science and society. In recent years, AI-based medium-range weather forecasts have advanced rapidly, yet they still rely on traditional data assimilation techniques from conventional numerical weather prediction (NWP) systems to provide initial states. Consequently, data-driven data assimilation methods have garnered increasing attention as a crucial component in constructing end-to-end weather forecasting systems based on AI. These AI-based approaches not only significantly reduce resource consumption but also offer new possibilities for overcoming bottlenecks in traditional numerical methods.

Existing work is primarily designed and trained based on specific resolutions (often matching the background resolution), resulting in models that can only be utilized in the same settings. In this paper, we propose the Fourier Neural Processes (FNP) for arbitrary-resolution data assimilation, whose necessity and importance lie mainly in the following points. Firstly, FNP eliminates the need to interpolate observational data, typically with higher resolution, which can avoid the information loss and improve the performance of data assimilation. Secondly, observational data in practical applications vary greatly in type and resolution; FNP can directly assimilate the diverse data without preprocessing, enhancing its practical application value. Lastly, the challenge of dealing with high-dimensional data is a common issue faced by traditional and AI methods; As our experiments demonstrate, the outstanding generalization ability of FNP enables training at low resolutions and direct application at high resolutions, thereby significantly reducing computational resource consumption.

The authors quickly delve into the pipeline description without first presenting the key and unique idea of their method for this problem. It takes readers some time to figure out, 'How does this method solve the issue of arbitrary resolution in data assimilation?’

From an implementation perspective, the ability to support inputs of arbitrary resolutions lies in the fact that fundamental operations (such as convolution, MLP, etc.) are independent of input size. Neural processes operate on paired (context or target) coordinate-value units, making them highly suitable for data assimilation tasks where observational data may take irregular forms. Therefore, we introduced a convolutional version of neural processes (i.e., ConvCNP) and incorporated custom modules independent of input size to achieve data assimilation with arbitrary resolutions. However, as indicated by our experiments, the performance of ConvCNP is unsatisfactory. Therefore, from a performance standpoint, the key to effectively addressing the issue of arbitrary resolutions lies in the efficacy of our module design (e.g., unified coordinate transformation and dynamic alignment to ensure the model comprehends spatial correspondences between data of different resolutions, and the neural Fourier layer providing significant high-frequency information gain when assimilating high-resolution observations).

Could the authors clarify the unique contributions of their method in addressing the issue of arbitrary resolution from a design perspective?

Based on the aforementioned ideas, we can categorize the core contributions of our approach in design into two aspects. Firstly, we astutely recognized the unique advantages of neural processes in addressing data assimilation challenges, introducing neural processes into data assimilation tasks to enable data assimilation with arbitrary resolutions. Secondly, through efficient module design, we achieved a significant improvement in performance and generalization, making FNP the only model whose performance demonstrates consistency with theoretical understanding (with the increase of resolution and the increase of observational information, the model's performance continues to improve significantly).

It seems that the core of your approach relies on interpolation in the feature space. If the generalization ability for arbitrary resolutions is primarily due to this interpolation, does this mean that integrating similar interpolation techniques into any deep learning method could achieve similar improvements in generalization?

Based on the above summary, interpolation in feature space is only a small part of our method. In fact, in the results we presented, ConvCNP also employed interpolation and simple fusion instead of DAM, yet its generalization was poor. Here, we further provide ablation study on generalization of different modules, as shown in Table A from the PDF in global response. The models in the table are all trained at 1.40625° resolution and directly tested at other resolutions to validate their generalization. It can be observed that all modules contribute to enhancing generalization to varying degrees. In addition, not all deep learning methods can support inputs of arbitrary sizes, thereby combining interpolation to achieve data assimilation with arbitrary resolutions. Even if they can, simple interpolation does not necessarily lead to improved generalization.

We sincerely hope that we have addressed your concerns and questions and look forward to your reading and response. If you have any further questions, please feel free to let us know and we will do our best to answer them.

Thank you for the detailed responses. I’ve reviewed the comments from other reviewers as well as your replies. I’m curious about the performance if both NFL and SVD are removed simultaneously. Additionally, if SVD alone is removed, does the generalization capability (arbitrary resolution) of the proposed method primarily stem from the unified coordinate transformation?