LOGLO-FNO: Local Global Fourier NO

Core Building Blocks of Base FNO vs. LOGLO-FNO

The core idea of LOGLO-FNO for achieving local spectral convolution in FNO is to first partition or decompose the domain D into M non-overlapping hypercubes called patches, such that $$\bigcup _{m=1}^M{P}_m=D,$$ and then performing (learnable) spectral convolutions on these sub-domains without any truncation of Fourier modes.

Fourier layer of FNO:

$$\text{Let}\mathbf{Z}\in {\mathbb{R}}^{\cdots \times N_c\times N_x\times N_y[\times N_z]}\text{be the output of the (initial) lifting layer. Then,}$$ $$\mathbf{{\rm Y}}:=\sigma [\mathcal{K}(\mathbf{Z})+{\mathbf{W}}_f\ast \mathbf{Z}+{\mathbf{b}}_f]$$ $$\mathcal{L}(\mathbf{Z}):={\mathbf{W}}_{c2}\ast (\sigma [{\mathbf{W}}_{c1}\ast \mathbf{{\rm Y}}+{\mathbf{b}}_{c1}])+{\mathbf{b}}_{c2}+[{\mathbf{W}}_c\odot \mathbf{Z}+{\mathbf{b}}_c]$$ $$\text{where}{\mathbf{W}}_{◻}\text{and}{\mathbf{b}}_{◻}\text{are learnable parameters and}\mathcal{K}(\cdot )\text{is global kernel integral.}$$

Fourier layer of LOGLO-FNO:

$${\mathbf{{\rm Y}}}_g:=\sigma [{\mathcal{K}}_g(\mathbf{Z})+{\mathbf{W}}_g\ast \mathbf{Z}+{\mathbf{b}}_g],{\mathbf{{\rm Y}}}_l:=\sigma [{\mathcal{K}}_l(\hat{\mathbf{Z}})+{\mathbf{W}}_l\ast \hat{\mathbf{Z}}+{\mathbf{b}}_l],$$ $$\mathcal{L}(\mathbf{Z},\hat{\mathbf{Z}},{\mathbf{Z}}^{{}^{\prime }}):={\mathbf{W}}_{gc2}\ast (\sigma [{\mathbf{W}}_{gc1}\ast {\mathbf{{\rm Y}}}_g+{\mathbf{b}}_{gc1}])+{\mathbf{b}}_{gc2}+[{\mathbf{W}}_{gc}\odot \mathbf{Z}+{\mathbf{b}}_{gc}]+$$ $$[{\mathbf{W}}_{lc}\odot \hat{\mathbf{Z}}+{\mathbf{b}}_{lc}]+{\mathbf{W}}_{lc2}\ast (\sigma [{\mathbf{W}}_{lc1}\ast {\mathbf{{\rm Y}}}_l+{\mathbf{b}}_{lc1}])+{\mathbf{b}}_{lc2}+$$ $${\mathbf{W}}_{hfc2}\ast (\sigma [{\mathbf{W}}_{hfc1}\ast {\mathbf{Z}}^{{}^{\prime }}+{\mathbf{b}}_{hfc1}])+{\mathbf{b}}_{hfc2},$$

Complexity Analysis of Local and Global Spectral Convolutions in Baseline FNO and LOGLO-FNO:

(i) Global Branch - 2D Spectral Convolutions on 2D Spatial Data

Let $N_b$ be the batch size, $N_x$ and $N_y$ the resolutions of the spatial dimensions, $N_c$ be the width or the number of hidden channels, and $c_{\text{in}}$ and $c_{\text{out}}$ are typically set to be of same dimension (e.g., 128). $${\text{FLOPs}}_{\text{FFT}}=5\cdot N_b\cdot C_{\text{in}}\cdot N_x\cdot N_y\cdot {\log }_2(N_x\cdot N_y)$$ $${\text{FLOPs}}_{\text{IFFT}}=5\cdot N_b\cdot C_{\text{out}}\cdot N_x\cdot N_y\cdot {\log }_2(N_x\cdot N_y)$$ Therefore, FFT computation on the global branch with the full 2D spatial resolution has a per channel complexity of $$\mathcal{O}(N_x\cdot N_y\cdot {\log }_2(N_x\cdot N_y)),$$ making it expensive for large spatial resolutions such as $2048\times 2048$ or higher.

(ii) Local Branch - 2D Spectral Convolutions on 2D Spatial Data

Let $N_p$ be the number of patches obtained by $\frac{N_x\cdot N_y}{P_s^2}$, $N_x$ and $N_y$ the resolutions of the spatial dimensions, $P_s\times P_s$ the patch size (e.g., $16\times 16$ ), $N_c$ be the width or the number of hidden channels, and $c_{\text{in}}$ and $c_{\text{out}}$ are typically set to be of same dimension (e.g., 128). $$\begin{array}{rl}{\text{FLOPs}}_{\text{FFT}}& =5\cdot N_b\cdot C_{\text{in}}\cdot N_p\cdot P_s\cdot P_s\cdot {\log }_2(P_s\cdot P_s)\\ & =5\cdot N_b\cdot C_{\text{in}}\cdot N_x\cdot N_y\cdot {\log }_2(P_s\cdot P_s)\end{array}$$ $$\begin{array}{rl}{\text{FLOPs}}_{\text{IFFT}}& =5\cdot N_b\cdot C_{\text{out}}\cdot N_p\cdot P_s\cdot P_s\cdot {\log }_2(P_s\cdot P_s)\\ & =5\cdot N_b\cdot C_{\text{out}}\cdot N_x\cdot N_y\cdot {\log }_2(P_s\cdot P_s)\end{array}$$ Thus, the FFT and IFFT computations on the local branch, which is operating on the patches has a per channel computational complexity of $$\mathcal{O}(N_x\cdot N_y\cdot {\log }_2(P_s\cdot P_s)).$$ Since, in practice, $$\mathcal{O}(N_x\cdot N_y\cdot {\log }_2(P_s\cdot P_s))\ll \mathcal{O}(N_x\cdot N_y\cdot {\log }_2(N_x\cdot N_y)),$$ the FFT computations are significantly cheaper in the local branch. Moreover, it is highly parallelizable when computing FFTs since each patch can be processed independently, leveraging modern accelerators such as GPUs, TPUs, and other forms of processing units.

(iii) Global Branch - 3D Spectral Convolutions on 3D Spatial Data

Let $N_b$ be the batch size, $N_x$, $N_y$, and $N_z$ the resolutions of the spatial dimensions, $N_c$ be the width or the number of hidden channels, and $c_{\text{in}}$ and $c_{\text{out}}$ are typically set to be of same dimension (e.g., 128). $${\text{FLOPs}}_{\text{FFT}}=5\cdot N_b\cdot C_{\text{in}}\cdot N_x\cdot N_y\cdot N_z\cdot {\log }_2(N_x\cdot N_y\cdot N_z)$$ $${\text{FLOPs}}_{\text{IFFT}}=5\cdot N_b\cdot C_{\text{out}}\cdot N_x\cdot N_y\cdot N_z\cdot {\log }_2(N_x\cdot N_y\cdot N_z)$$ Therefore, FFT computation on the global branch with the full 3D spatial resolution has a per channel complexity of $$\mathcal{O}(N_x\cdot N_y\cdot N_z\cdot {\log }_2(N_x\cdot N_y\cdot N_z)),$$ making it highly expensive for large values of $N_x$, $N_y$, and $N_z$ such as $512\times 512\times 512$ or further higher spatial resolutions.

(iv) Local Branch - 3D Spectral Convolutions on 3D Spatial Data

Let $N_p$ be the number of patches obtained by $\frac{N_x\cdot N_y\cdot N_z}{P_s^3}$, $N_x$, $N_y$, and $N_z$ the resolutions of the spatial dimensions, $P_s\times P_s\times P_s$ the patch size (e.g., $16\times 16\times 16$, $32\times 32\times 32$, etc.), $N_c$ be the width or the number of hidden channels, and $c_{\text{in}}$ and $c_{\text{out}}$ are typically set to be of same dimension (e.g., 128). $$\begin{array}{rl}{\text{FLOPs}}_{\text{FFT}}& =5\cdot N_b\cdot C_{\text{in}}\cdot N_p\cdot P_s\cdot P_s\cdot P_s\cdot {\log }_2(P_s\cdot P_s\cdot P_s)\\ & =5\cdot N_b\cdot C_{\text{in}}\cdot N_x\cdot N_y\cdot N_z\cdot {\log }_2(P_s\cdot P_s\cdot P_s)\end{array}$$ $$\begin{array}{rl}{\text{FLOPs}}_{\text{IFFT}}& =5\cdot N_b\cdot C_{\text{out}}\cdot N_p\cdot P_s\cdot P_s\cdot P_s\cdot {\log }_2(P_s\cdot P_s\cdot P_s)\\ & =5\cdot N_b\cdot C_{\text{out}}\cdot N_x\cdot N_y\cdot N_z\cdot {\log }_2(P_s\cdot P_s\cdot P_s)\end{array}$$ Thus, the FFT and IFFT computations on the local branch operating on 3D patches has a per channel computational complexity of $$\mathcal{O}(N_x\cdot N_y\cdot N_z\cdot {\log }_2(P_s\cdot P_s\cdot P_s)).$$ Since, in practice, $$\mathcal{O}(N_x\cdot N_y\cdot N_z\cdot {\log }_2(P_s\cdot P_s\cdot P_s))\ll \mathcal{O}(N_x\cdot N_y\cdot N_z\cdot {\log }_2(N_x\cdot N_y\cdot N_z)),$$ the FFT computations are significantly cheaper in the local branch. Furthermore, it is highly parallelizable when computing FFTs since each patch can be processed independently, leveraging modern accelerators such as GPUs, TPUs, and other forms of PUs.

🌈 Radially Binned Spectral Energy of Errors

In addition, we propose a spectral loss term based on the radial binning of spectral energy of errors, which is as follows:

🖥 Pseudocode - Radially Binned Spectral Energy of Errors

            
                def RadialBinnedSpectralLoss(preds, target):
                    # input data shape and params
                    nb, nc, nx, ny, nt = target.size()
                    iLow, iHigh = 4, 12
                    Lx, Ly = 1.0, 1.0

                    # Compute error in Fourier space
                    err_phys = preds - target
                    err_fft = torch.fft.fftn(err_phys, dim=[2, 3])
                    err_fft_sq = torch.abs(err_fft)**2
                    err_fft_sq_h = err_fft_sq[Ellipsis, :nx//2, :ny//2, :]

                    # Create radial indices
                    x = torch.arange(nx//2)
                    y = torch.arange(ny//2)
                    X, Y = torch.meshgrid(x, y, indexing="ij")
                    radii = torch.sqrt(X**2 + Y**2).floor().to(torch.int) # Radial dist.
                    max_radius = int(torch.max(radii))

                    # flatten radii for binary mask
                    radii_flat = radii.flatten() # (nx//2 * ny//2)
                    
                    # Spatially flatten Fourier space error; (nb, nc, nx//2 * ny//2, nt)
                    err_fft_sq_flat = err_fft_sq_h.contiguous().reshape(nb, nc, -1, nt)
                    
                    # initialize output tensor to hold the Fourier error
                    # for each radial bin at distance r from the origin
                    err_F_vect_full = torch.zeros(nb, nc, max_radius + 1, nt)
                    
                    # Apply ‘index_add_‘ for all radii and accumulate the errors
                    valid_r = radii_flat <= max_radius # binary mask to find valid radii
                    
                    # Sum for all valid radial indices
                    err_F_vect_full.index_add_(2,
                                               radii_flat[valid_r],
                                               err_fft_sq_flat[:, :, valid_r]
                                              )

                    # Normalize & compute mean over batch; (nc, min(nx//2, ny//2), nt)
                    nrm = (nx * ny) * Lx * Ly
                    _err_F = torch.sqrt(torch.mean(err_F_vect_full, dim=0)) / nrm

                    # Classify Fourier space error into three bands
                    err_F = torch.zeros([nc, 3, nt])
                    err_F[:, 0] += torch.mean(_err_F[:, :iLow], dim=1) # low freqs
                    err_F[:, 1] += torch.mean(_err_F[:, iLow:iHigh], dim=1) # mid freqs
                    err_F[:, 2] += torch.mean(_err_F[:, iHigh:], dim=1) # high freqs

                    # mean or sum over the channels and time dimensions
                    if reduction == "mean":
                        freq_loss = torch.mean(err_F, dim=[0, -1])
                    elif reduction == "sum":
                        freq_loss = torch.sum(err_F, dim=[0, -1])

                    return freq_loss
            
        

🏹 Training Objective

The 1-step training loss for N trajectories, each comprising T timesteps, is given by,

$${\theta }^{\ast }=\arg \underset{\theta }{min}\sum _{n=1}^N\sum _{t=1}^{T-1}\mathcal{C}({\mathcal{N}}_{\theta }(u^t),u^{t+1}),\mathcal{C}={\mathcal{C}}_{\mathrm{MSE}}+\lambda \cdot {\mathcal{C}}_{\mathrm{freq}},0\le \lambda \le 1$$

High-Frequency Adaptive Gaussian Noise

Let ${\mathbf{X}}^{hf}\in {\mathbb{R}}^{N_b\times N_x\times N_y\times (N_t\cdot N_c)}$ denote the input tensor of HF features, where $N_b$ is the batch size, $N_x$ and $N_y$ are the resolutions of the spatial dimensions, and $N_t\cdot N_c$ represents the combined temporal and channel dimensions. The high-frequency feature adaptive Gaussian noise ${\mathbf{N}}_{dynamic}$ is then computed as follows:

1. Compute (per sample) Mean ${\mu }_b$ and Standard Deviation ${\sigma }_b$ of High-Frequency Features $${\mu }_b=\frac{1}{N_x\cdot N_y\cdot N_t\cdot N_c}\sum _{i=1}^{N_x}\sum _{j=1}^{N_y}\sum _{k=1}^{N_t{N}_c}{\mathbf{X}}_{b,i,j,k}^{\mathrm{hf}}$$ $${\sigma }_b=\sqrt{\frac{1}{N_x\cdot N_y\cdot N_t\cdot N_c}\sum _{i=1}^{N_x}\sum _{j=1}^{N_y}\sum _{k=1}^{N_t{N}_c}({\mathbf{X}}_{b,i,j,k}^{\mathrm{hf}}-\mu {)}^2}+ϵ$$ where $ϵ$ is a small constant added for numerical stability, and $\mu$ and $\sigma$ are obtained by stacking the per sample statistics along the batch dimension.

2. Generate Standard Gaussian Noise $$\mathbf{N}\sim \mathcal{N}(0,1)$$ 3. Scale Noise Dynamically $${\mathbf{N}}_{dynamic}=\mu +\alpha \cdot \sigma \cdot \mathbf{N}$$ where $\alpha$ is a small value such as 0.025 and ${\mathbf{N}}_{dynamic}$ has the same shape as the input ${\mathbf{X}}^{hf}$.

${\mathbf{N}}_{dynamic}$ can now be added to the batch of inputs to the global and local branches during the training phase of LOGLO-FNO.

🌀 Visualizing the Radially Binned Spectral Loss

In the below slideshow, we visualize the radially binned spectral energy errors of predictions of the considered neural operators and LOGLO-FNO on the Kolmogorov Flow 2D PDE. Note that we show only alternate radial bins prioritizing uncluttered representation over completeness.

@inproceedings{loglo-fno-kalimuthu:2025,
    title={{LOGLO}-{FNO}: Efficient Learning of Local and Global Features in Fourier Neural Operators},
    author={Marimuthu Kalimuthu and David Holzm{\"u}ller and Mathias Niepert},
    booktitle={ICLR 2025 Workshop on Machine Learning Multiscale Processes},
    year={2025},
    url={https://openreview.net/forum?id=OCM7OkVg9C}
}