Efficient Fourier Neural Operators by Group Convolution and Channel Shuffling

Authors: myungjoon  kim
junhyung  park
jonghwa  shin
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 13:00:05

Abstract: Fourier neural operators (FNOs) have emerged as data-driven alternatives to conventional numerical simulators for solving partial differential equations (PDEs). However, these models typically require a substantial number of learnable parameters. In this study, we explore parameter-efficient FNO architectures through modifications in their width, depth, and applications of group convolutiopn and channel shuffling. We perform benchmark on different problems on learning the operator of Maxwell's equations and Darcy flow equations. Our approach leads to significant improvement in prediction accuracy for both small and large FNO models. The proposed methods are widely adaptable across various problem types and neural operator architectures, aiming to boost prediction accuracy.

Main paper: 
 introduction: 
 Science and engineering problems are often characterized by partial differential equations (PDEs). To solve these complex PDEs, numerical simulations, such as the finite difference method (FDM) or the finite element method (FEM), are employed. These traditional approaches solve PDEs through particular algorithms that discretize space and time, often requiring substantial computational resources. Such a computationally demanding part is a bottleneck, hindering the rapid calculation and analysis of the system.Recently, the advancement of deep neural networks has made data-driven models a promising approach for solving PDEs. While convolutional neural networks (CNNs) have been applied to these challenges, they often fall short for PDE problems due to their inherent design for mapping between Euclidean spaces. A new class of neural PDE solvers, known as neural operators, has emerged to address these limitations. Among various methods, Fourier neural operators (FNOs) stand out as a solution for solving complex PDEs, offering high accuracy and rapid prediction capabilities. Numerous research aimed at exploring diverse applications and enhancing functionalities of FNOs.In this work, we implement an architectural modification to FNOs aimed at increasing both efficiency and accuracy. Through modifications in width, depth, and applying group convolution and channel shuffling, inspired by the traditional CNN architectures, we can find the optimal architecture with best performance. The main contribution of our works are summarized as follows:• We present an adjustment to the architecture of the Fourier layer, modifying its width and depth, incorporate group convolution along with channel shuffling to improve performance, and show our approach is memory and computation efficient way to improve large FNO models.• Thorough benchmarking and detailed analysis show that modifications to the architecture significantly boost model performance. The benchmark results highlight the importance of selecting the appropriate model for accurately predicting the solution field across different PDE problems.Parameter efficient architectures in convolutional neural networks Since the introduction of the ResNet architecture (He et al., 2016) , numerous modifications have been proposed to enhance its representational capacity while maintaining the complexity of neural networks. Zagoruyko & Komodakis (2016) suggested WideResNet and explored how increasing the width of the network could contribute to higher accuracy in image classification tasks. The ResNeXt architecture employs group convolution to enhance the cardinality of its layers, contributing to its effectiveness and efficiency (Xie et al., 2017) . ShuffleNet extends the application of group convolution to include pointwise operations and incorporates channel shuffling to optimize performance and efficiency (Zhang et al., 2018) .Fourier neural operators FNOs are advanced computational models that efficiently solve PDEs by leveraging Fast Fourier transforms to parameterize and predict complex dynamics. FNOs learn mapping between infinite-dimensional function spaces (Li et al., 2020) , thus they are well-suited as surrogate models for conventional PDE simulators. Recently, numerous variants of the FNO have been proposed to enhance data and memory efficiency (Rahman et al., 2022; Li et al., 2022; 2021; Helwig et al., 2023) . Notably, Factorized Fourier Neural Operator (F-FNO) introduces a dimension-wise factorization approach (Tran et al., 2022) . This was achieved by representing 2D Fourier domains through the summation of two orthogonal 1D Fourier transforms.Surrogate models for numerical simulations Numerous neural network-based methods have been developed to solve PDEs (Kovachki et al., 2023) . In the field of optics, researchers have tackled the problem of field distribution prediction using CNNs (Wiecha & Muskens, 2019; Chen et al., 2022; Lim & Psaltis, 2022) . More recently, specialized variants of FNOs designed to solve Maxwell's equations have shown superior performance over traditional CNNs (Gu et al., 2022; Augenstein et al., 2023) . 
parameter-efficient fourier neural operators: 
 The Fourier integral operator in the original Fourier neural operator (Li et al., 2020) is defined asEQUATIONwhere K (l) represents the learnable kernel integral operator using the Fourier transform F and inverse Fourier transform F -1 , defined asEQUATIONThe parameter numbers in FNO is O LH 2 M D , where L, H, and M represents the number of Fourier layer, the channel, the Fourier modes being kept, respectively. D denotes the dimension of the problem.In our proposal, we apply modifications to FNO models to enhance its parameter efficiency. Initially, we tune the depth (the number of Fourier layers) and the width (the number of channels) of the model. Specifically, we incorporate pointwise group convolution and channel shuffling into the architecture. The design of our modified version of the architecture is depicted in Figure 1 . Within each Fourier layer, we implement group convolution in both the frequency and spatial domains. This technique reduces the number of parameters toO LH 2 M D g, where g represents the group number. By doing so, we decrease the model's complexity by a factor of g, enhancing its parameter efficiency. Moreover, we utilize channel shuffle operations at every Fourier layer to boost the representation capability by enabling effective crosstalk among features. These operations involve reshaping the feature dimension into (g, m), transposing it, and then flattening it back to the input format for the subsequent layer.  
experiment 4.1 dataset: 
 Periodic metagrating In nanophotonics, periodic nanostructures with the size comparable to wavelength can be used for diffraction control devices. These have been implemented to steer light beams (Sell et al., 2017) in specific directions, and phase masks for photolithographic process (Nam et al., 2022) . The optical responses are governed by Maxwell's equations. We use a metasurface dataset originally provided by Chen et al. (2022) , available in the online database (Jiang et al., 2020) . We sampled 12,000 and 1,500 data for training and validation, respectively. Figure 2 illustrates an example of the metagrating structure along with the fields computed by the FDM solver. The system comprises periodic boundaries along the x direction and perfectly matched layers, which are artificial layers designed to absorb all radiated fields, along the y direction.These structures are converted to [1, 64, 256] image format, where each pixel corresponds to a mesh size of 6.25nm × 6.25nm. The incident light has a wavelength of 1050 nm and is polarized in the transverse magnetic (TM) mode. The relative permittivity of the materials ranges from 6.12 to 11.90, and there is a substrate with a permittivity of 2.10.We demonstrate an additional examples using parameter-efficient FNO to solve different type of PDE. Darcy flow equations describe the movement of fluids through porous media, relating fluid velocity to the pressure gradient, permeability of the medium, and fluid viscosity. We use standard two-dimensional Darcy Flow equation on the unit box. The dataset consists of 1000 training samples and 200 validation samples, all arranged in a uniform grid of 241 × 241.Training We use F-FNO for all experiments. For training, we use AdamW optimizer (Loshchilov & Hutter, 2017) with an initial learning rate of 0.001 and a batch size of 32. For all models, we use exponential learning rate scheduling and the total number of epochs is set to 200. GeLU serves as the activation function, and we maintain [12, 12] Fourier modes for all experiments. For the loss function and the prediction accuracy metric, we use the normalized mean absolute error (NMAE) defined as:NMAE = 1 N N i=1 ∥ŷ -y∥ 1 ∥y∥ 1 (3)where y and ŷ are the predicted field values and ground truth, respectively, and ∥y∥ 1 is the L 1 norm of the field distribution.Depth vs. Width In our initial experiment, we evaluated the impact of depth and width while maintaining a fixed number of parameter of small (0.56M) and large (2.2M) models. We tested configurations with 5, 10, 15, and 20 Fourier layers, adjusting the number of channels for each layer accordingly. The results showed that a setup with 15 layers and 27 channels for the small model, and 10 layers with 64 channels for the large model, exhibits the best performance. From these findings, we conclude that there are specific optimal values for the number of layers and channels that produce the highest efficiency within a given parameter budget. We explore the impact of group convolution and channel shuffling on model performance.As in the earlier experiment, we performed benchmarks with both small and large models. We kept the number of parameters constant, increasing the number of groups and channels simultaneously (Table 2 ). In both cases, the models achieved their best performance when the group number is set to 8. As the number of groups increases, it enhances the feature representation capability, which in turn improves the overall performance of the model. Overall, it is evident that neural operators perform best when their width, depth, and group numbers are set optimally within the given parameter budget constraints. Prediction accuracy increases by 33% and 31% through the appropriate selection of width/depth ratios, and further improvements of 16% and 21% are achieved by incorporating group convolution and channel shuffling in small (0.56M) and large (2.2M) models, respectively. These results highlight the importance of the selection of architectures to maximize model performance.We also evaluated the impact of group convolution on FNO models while keeping the number of channels (Table 3 ). In small FNO models, we observed a performance decline as the number of parameters was reduced by increasing the number of groups. On the other hand, in large models, performance enhanced with the addition of group convolutions These suggest that for large models, sparse connection may enhance generalization capability. We anticipate that by integrating memoryefficient models (Rahman et al., 2022) , we can develop more efficient FNO models. Generalization with FNO We examined the effect of pointwise group convolution and channel shuffling on the original FNO architectures. Within the constraints of parameter number budgets or with the same number of channels for large models, pointwise group convolution and channel shuffling enhance FNO models. Therefore, our approach appears to consistently boost the performance of various FNO models.CONSLUSIONIn this study, we introduce parameter-efficient FNOs, modified neural operator achitectures capable of solving PDE problems. Leveraging previous development in conventional CNNs, we investigated the effectiveness of changing width, depth, and group convolutions and channel shuffling in FNOs. Our proposed methods achieve superior representational capabilities by increasing the number of channels while maintaining the overall parameter count. We conducted two benchmark problems, Maxwell's equations and Darcy flow equation. The architectural modifications significantly improve prediction accuracy. Moving forward, we aim to tackle a wider array of PDE systems, including diffusion processes, and Navier-Stokes equations. 
 
 Appendix:  


Decision: Thanks to reviewers for their feedback. The paper proposes the advantage of group convolution and channel shuffling in neural operator based models to improve model computational efficiency as well as accuracy. I encourage the authors to carefully review the comments of reviewers and clearly address them. Due to workshop limited page number, consider adding more details to appendices. I vote accept to this paper.

Reviews: review: title: title: This paper presents a development of architecture in Factorized - Fourier Neural Operator (F-FNO) to reduce the computational overheads and so prediction performance in a case study. It introduces group convolution to the F-FNO framework to decrease computational complexity by a factor of the number of groups. The authors test the performance evaluation for Maxwell PDEs as a preliminary benchmark investigation.
main_review: main_review: The study conducted benchmarking to demonstrate progress in both computational complexity and prediction accuracy by combining group convolution with hyperparameter optimizations, which is a strength of this paper. However, the authors' claim of a reduction of N-MSE from 0.054 to 0.031 should be clarified. It appears that the benefit of reduction by introducing the group convolution should be calculated as the reduction from 0.043 to 0.031, resulting in a 27.9% reduction, especially concerning the same number of Fourier layers (i.e., block) used.

While the theoretical aspects are clear, the paper lacks clarity regarding the practical memory efficiency gained during actual computations. However, the dataset and problem are well-defined, aiding understanding.

The authors justify its originality in combining group convolution with the F-FNO architecture. This approach contributes to the field by proposing a novel method for addressing computational efficiency and prediction accuracy simultaneously.

The use of the proposed method can enhance prediction accuracy and computational efficiency, as demonstrated in the study. The performance improvements observed in a case study involving Maxwell's equations in the nanophotonic structure domain underscore the importance of accelerating FNO to other scientific domains governed by partial differential equations (PDEs). However, it remains unclear whether the proposed architecture can achieve similar improvements in prediction accuracy across other domains, which warrants further investigation.
title: title: Efficient fourier neural operators by group convolution and channel shuffling
main_review: main_review: ## General comments
The paper introduces enhancements to Fourier Neural Operators, incorporating
group convolution and channel shuffling, aiming for improved efficiency and
accuracy in solving partial differential equations. The approach is evaluated
through benchmarks, showing significant performance enhancements. The
methodology combines theoretical foundations with empirical validation. The
document is structured to clearly outline the contributions, facilitating
understanding among readers with interests in scientific computing and data-
driven modeling.
The work is ready for publication, offering valuable insights and advancements
without the need for further modifications.