Complex Space-Frequency Learning Module

• Complex Space-Frequency Learning Module (CSFLM) is a neural architecture that leverages both complex frequency and spatial domain representations to efficiently encode long-range dependencies and reduce data dimensionality.
• It employs a 'Transform Once' paradigm with a global Fourier transform and spectral operator layers that ensure variance-preserving initialization and significant computational savings.
• CSFLMs demonstrate practical benefits with reduced training times, lower memory usage, and improved accuracy in operator learning tasks, as shown in fluid dynamics and inverse problem applications.

A Complex Space-Frequency Learning Module (CSFLM) refers to a class of neural architectures that explicitly leverage representations and learning in both the complex frequency domain (typically via Fourier or related transforms) and the spatial (or time) domain. These modules are designed to efficiently encode long-range dependencies, exploit spectral sparsity, reduce data dimensionality, and facilitate operator learning or robust feature extraction by direct, often parameterized, manipulation in frequency space. The approach is foundational to frequency-domain models (FDMs), state-of-the-art operator networks, hybrid spatial–spectral backbones, and advanced inverse-problem solvers.

1. Motivation and Principles

The rationale for complex space–frequency learning arises from two key properties of natural and physical signals:

• Long-Range Correlations in Frequency Domain: Many solution operators for PDEs and physical processes exhibit structures (such as translation equivariance or conservation laws) that are best captured by basis functions in the frequency domain, where convolution becomes elementwise multiplication. Spectral analysis enables dimensionality reduction because a significant portion of relevant information is encoded in a small number of low-frequency components (Poli et al., 2022).
• Computational Efficiency and Robustness: Working in the frequency domain allows for information-preserving dimensionality reduction and efficient global filtering. However, standard FDMs that switch back and forth between domains for each layer incur prohibitive computational cost. CSFL modules streamline this by operating predominantly (or exclusively) in frequency space between a single input transform and output inverse transform, thus achieving superior efficiency (Poli et al., 2022).

2. Core Architecture: Transform-Once Blueprint

The "Transform Once" (T1) paradigm formalizes the core operational workflow:

1. Global Complex Transform: Apply a unitary Discrete Fourier Transform (DFT) or Fast Fourier Transform (FFT) to the input $x \in \mathbb{R}^N$ at the network entry:

$$\hat x[k] = \frac{1}{\sqrt{N}} \sum_{n=0}^{N-1} x[n] e^{-2\pi i k n / N}$$

The forward and inverse transforms are normalized to ensure variance preservation (unitary), as established in Lemma A2 (Poli et al., 2022).

2. Spectral Operator Layers: Within frequency space, a series of $L$ parameterized operator layers perform channelwise (and potentially cross-channel) complex multiplications. For each mode $k$ and channel $c$:

$$\hat Y_c[k] = \sum_{c'=1}^{C} \hat W_{c, c'}[k] \hat X_{c'}[k]$$

The design may use dense, block-diagonal, or diagonal complex kernels, always acting directly in k-space. Only one inverse DFT is performed after all spectral layers to return to the spatial domain.

3. Variance-Preserving Initialization: For stability, complex weight matrices are initialized so that output activations at each layer preserve the input variance:

$\sigma_A^2 = \frac{N}{m^2}$

for general $m \times m$ kernels, or $N/m$ for diagonal kernels (Poli et al., 2022). This initialization is critical for avoiding vanishing/exploding signals during deep frequency-domain propagation.

4. Frequency Selection and Reduced-Order Modelling: Physical signals typically possess rapidly decaying spectra, allowing truncation to $m \ll N$ lowest-frequency modes with empirically negligible reconstruction error. Boolean frequency masks $M \in \{0,1\}^N$ select the retained set, and per-layer computations are limited accordingly, yielding substantial computational savings while maintaining accuracy.

3. Practical Implementation and Pseudocode

Efficient CSFLM implementation leverages hardware-friendly batched FFTs and memory layouts, minimizing per-layer domain conversions. Below is a representative PyTorch-style module architecture (Poli et al., 2022):

class T1Module(nn.Module): def __init__(self, C, m, L): super().__init__() self.mask = make_lowpass_mask(N, m) self.A = nn.ParameterList([ complex_weight_matrix(C, C, m, init='vp') for _ in range(L) ]) self.activation = GeLU() def forward(self, x): Xk = fft(x, norm='ortho') Xk_trunc = Xk[..., self.mask] for A in self.A: Yk = complex_batch_matmul(A, Xk_trunc) Xk_trunc = self.activation(Yk) Xk_full = torch.zeros_like(Xk) Xk_full[..., self.mask] = Xk_trunc y = ifft(Xk_full, norm='ortho').real return y

Key efficiency details:

• Use of torch.fft.rfft/rfft2 for real inputs (exploiting conjugate symmetry).
• All FFTs batched once at the beginning, one IFFT at the end; no per-layer transforms.
• Complex parameters stored as native complex types or in separate real/imag parts.
• Fused spectral multiplications to minimize kernel-launch overhead.

4. Frequency Selection and Dimensionality Reduction

Empirical analysis reveals that retaining approximately 25–50% of the frequency modes typically suffices for less than 1% error in signal reconstruction (see Figs. B.1, B.9, B.12 in (Poli et al., 2022)). This enables aggressive pruning of irrelevant spectral bands via masks and selection matrices, which drastically lowers the computational cost from $O(C^2N)$ to $O(C^2m)$ per layer in multi-channel settings.

In experimental settings for operator learning tasks (e.g., 2D incompressible Navier–Stokes, airfoil flow, turbulent smoke):

• T1 models achieve 3×–10× wall-clock speedups versus standard FNO baselines, with up to 2×–5× reduction in memory footprint.
• Reduced channel count and absence of residual spatial convolutions yield parameter efficiency comparable or superior to FNO, FFNO, and U-Net-based designs (Poli et al., 2022).

5. Empirical Results and Performance Benchmarks

T1 CSFLMs have demonstrated robust performance across a suite of high-resolution spatio-temporal operator learning tasks:

• Speed: Wall-clock training/inference times reduced from 32 hours (FNO) to 5 hours (T1) in large-scale experiments.
• Accuracy: Average predictive error reduced by over 20% across tasks.
• Resource Utilization: Model memory footprint is halved due to reduced mode storage and elimination of per-layer FFT buffers.
• Generalization: On incompressible NSE at $T = 50$ s, T1 models reduce relative-L2 error by ∼20% versus FNO and ∼40% versus FFNO; in ScalarFlow, 10-step rollout N-MSE drops from $2.32 \times 10^{-1}$ to $2.28 \times 10^{-1}$ for T1+variance-preserving init.

6. Module Generalization and Theoretical Insights

The Transform Once CSFLM directly exploits properties of the DFT as a unitary, information-preserving operator over $\mathbb{C}^N$. The approach is theoretically justified by the preservation of total variance throughout the frequency-domain pathway, permitting deep stacking of spectral operator layers without destabilization. Its efficacy is rooted in the mathematical structure of physical fields and natural signals whose dominant correlations concentrate in low-frequency bands, supporting aggressive model-order reduction (Poli et al., 2022).

Furthermore, CSFLMs naturally extend to multidimensional settings (e.g., images, videos, physical simulation grids), where frequency truncation becomes even more impactful due to increased spectral redundancy.

7. Summary Table: Core Design Steps

Stage	Description	Key Formula/Parameter
Global Transform	Apply DFT to input; unitary, variance-preserving	$\hat x[k] = \frac{1}{\sqrt{N}} \sum x[n]\cdots$
Spectral Operator Layer	Complex matrix multiplies in k-space	$$\hat{Y}_c[k] = \sum \hat{W}_{c,c'}[k] \hat{X}_{c'}[k]$$
Variance-Preserving Initialization	Scale random complex weights to match input/output variance	$\sigma_A^2 = N/m^2$ (full), $N/m$ (diagonal)
Frequency Mode Selection	Retain mask of low-$m$ modes per empirical decay	$m = 0.25N$ to $0.5N$ for $<1\%$ error
Efficient GPU Implementation	Batched FFTs, complex dtypes, fused multiplies	Use torch.fft.rfft, channel-last layout
Empirical Performance	3–10× speedup, $\leq$50% memory, $>$20% error reduction	Benchmarks: NSE, airfoil, ScalarFlow

This organizational pattern generalizes to operator learning, high-dimensional regression, and inverse problems, establishing CSFLM (esp. the T1 blueprint) as an efficient and theoretically sound alternative to iterative per-layer frequency–spatial transforms in FDMs (Poli et al., 2022).