Transform Once Efficient Operator Learning in Frequency Domain
Abstract
Spectrum analysis provides one of the most effective paradigms for information–preserving dimensionality reduction in data: often, a simple description of naturally occurring signals can be obtained via few terms of periodic basis functions.
Neural operators designed for frequency domain learning – frequency domain models (FDMs) – are based on complex–valued transforms i.e. Fourier Transforms (FT), and layers that perform computation on the spectrum and input data separately.
This design introduces considerable computational overhead: for each layer, a forward and inverse FT.
Instead, this work introduces a blueprint for frequency domain learning through a single transform: transform once ( T1 ).
Our results noticeably streamline the design process of FDMs, pruning redundant transforms, and leading to speedups of 3 x to 10 x that increase with data resolution and model size.
We perform extensive experiments on learning to solve partial differential equations, including incompressible
Navier–Stokes, turbulent flows around airfoils and high-resolution video of smoke dynamics.
T1 models improve on the test performance of
SOTA FDMs while requiring significantly less computation, with over 20% reduction in predictive error across tasks.
1. Introduction
Naturally occurring signals are often sparse when projected on periodic basis functions (Strang, 1999)
Central to recently–introduced instances of frequency–domain neural operators (Li et al., 2020; Tran et al., 2021), which we refer to as frequency–domain models (FDMs), is the idea of learning to modify specific frequency components of inputs to obtain a desired output in data space.
With a hierarchical structure that blends learned transformations on frequency domain coefficients with regular convolutions, FDMs are able to effectively approximate global, long–range dependencies in higher resolution signals without requiring prohibitively deep architectures.
Yet, state–of–the–art variants of these models suffer from several drawbacks:
Poor parameter scaling: FDMs employ complex– valued transforms, followed by parameter– inefficient complex–valued layers 1 . 2.
Slow inference: every layer of an FDM performs a forward and inverse frequency domain transform, introducing a considerable computational overhead.
Suboptimal initialization: parameter initialization schemes and layers devised to learn directly in data space can be highly suboptimal when introduced without modifications to FDMs.
Although some attempts to improve performance (Gupta et al., 2021) or alleviate the above–mentioned limitations (Tran et al., 2021), scaling FDMs to larger data resolutions and model sizes remains fundamentally challenging.
In example, instantiating a medium–sized FDM (tens of layers with a non–trivial number of channels) requires a parameter count in the hundred millions.
In this work, we start by posing the question:
To reap the benefits of learning on frequency domain representations, is it necessary to construct hierarchical deep models that perform forward and inverse frequency transforms at each layer?
We provide the answer in Transform Once ( T1 ), a model that builds representations directly in frequency domain, after a single forward transform.
Each aspect of T1 addresses a specific limitation of existing FDMs:
Favourable scaling: T1 employs a single real– valued transform ( Discrete Cosine Transform – DCT), which reduces parameter scaling overhead w.r.t complex–valued layers.
Fast: by performing a single forward transform and optimizing directly on DCT coefficients of target data, T1 iterations are at least 3 x to 10 x faster.
When scaling to larger models and higher resolutions, the relative speedups increase as the overhead of each transform grows.
Enhanced compatibility: removing redundant transforms streamlines the design space for T1 architectures compared to existing FDMs, allowing direct introduction of optimized layers developed for other applications
e.g. UNets
(Ronneberger et al., 2015)
In §2.1 we provide a (short) history on frequency domain approaches in deep learning, followed by background on
FDMs in §2.3.
In Section §3, we describe how to train
T1 directly in frequency domain, in §3.1 we motivate the choice of DCT, in §3.2 we derive a loss decomposition for reduced–order FDMs and in §3.3 we introduce a simple variance–preserving weight initialization scheme for all
FDMs.
Finally, in §4 we evaluate T1 on a suite of benchmarks related to learning solution operators for a variety of dynamics: incompressible Navier–Stokes, flow around different airfoil geometries, and high–resolution videos of turbulent smoke plumes (Eckert et al., 2019).
Across tasks, T1 is 3 × to 10 × faster and reduces predictive errors by 20% on average.
Training T1 models on high resolution videos ( 600 x 1062 ) of turbulent dynamics is significantly faster, requiring 5 hours instead of 32 hours (FNOs) for the same number of iterations.
2. Related Work and Background 2.1. Learning and Frequency Domain: A Short History
Links between frequency–domain signal processing and neural network architectures have been explored for decades, starting with the original CNN designs (Fukushima and Miyake, 1982).
(Mathieu et al., 2013;
Rippel et al., 2015) proposed replacing convolutions in pixel space with element–wise multiplications in Fourier domain.
In the context of learning to solve partial differential equations (PDEs), Fourier Neural Operators (FNOs) (Li et al., 2020) popularized the state–of–the–art
FDM layer structure: forward transform → learned layer → inverse transform.
Similar architectures had been previously proposed for generic image classification tasks in (Pratt et al., 2017; Chi et al., 2020).
Modifications to the basic FNO recipe are provided in (Tran et al., 2021;
Guibas et al., 2021; Wen et al., 2022).
A frequency domain representation of convolutional weights has also been used for model compression (Chen et al., 2016).
Fourier features of input domains and periodic activation functions play important roles in deep implicit representations (Sitzmann et al., 2020; Dupont et al., 2021;
Poli et al., 2022) and general–purpose P ERCEIVERS (Jaegle et al., 2021).
2.2. Learning to Solve Differential Equations
A variety of deep learning approaches have been developed to solve differential equations: neural operators and physics–informed networks (Long et al., 2018;
Raissi et al., 2019;
Lu et al., 2019; Karniadakis et al., 2021), specialized architectures (Wang et al., 2020; Lienen and Günnemann, 2022), hybrid neural–numerical methods (Poli et al., 2020;
Kochkov et al., 2021; Mathiesen et al., 2022), and FDMs (Li et al., 2020; Tran et al., 2021), the focus of this work.
2.3. Frequency–Domain Models
Let D n ( n -space) to be the set of real–valued discrete signals 2 of resolution N .
Our objective is to develop efficient neural networks to process discrete signals x ∈ D n , x 0 , x 1 , . . .
, x N − 1 , x n ∈ R .
We define a layer of FDMs mapping x to an output signal ˆ y ∈ D n as the structured operator:
[Equation]
Forward Transform [Equation] Learned Map
[Equation]
Inverse Transform
[Equation]
where T is an orthogonal (possibly complex ) linear operator.
We denote the T –transformed n -space with D k ( kspace) so that T : D n → D k .
Typically, we assume T to be a Fourier–type transform 3 (Oppenheim, 1999, Chapter 8) so that the k -space corresponds to the frequency domain and its elements form the spectrum of the input signal x .
The learned parametric map f θ : D k
→ D k is the stem of a FDM layer: it maps the k -space into itself and is typ- ically chosen to be rank–deficient in the linear case, e.g. f θ ( X ) =
S ⊤ m A ( θ ) S m X, A ( θ ) ∈
C m × m ( m ≤ N ) .
The matrix S m ∈ R n × m selects m desired elements of X , setting the rest to zero.
In the case of frequency domain transforms, this allows (1) to preserve or modify only specific frequencies of the input signal x .
Residual connections or residual convolutions g (Li et al., 2020; Wen et al., 2022) are optionally added to reintroduce frequency components filtered by the S m .
An FDM mixes global transformations applied to coefficients of the chosen transform to local transformations g i.e. convolutions with finite kernel sizes.
To ensure that such models can approximate generic nonlinear functions, nonlinear activations are introduced after each inverse transform.
Fourier Neural Operators Layers of the form (1) appear in recent FDMs such as Fourier Neural Operators (FNOs) (Li et al., 2020) and variants (Tran et al., 2021;
Guibas et al., 2021; Wen et al., 2022).
In example, an FNO is recovered from (1) by letting T be a Discrete Fourier Transform (DFT)
[Equation]
where W ∈ C N × N is the standard N -dimensional DFT matrix and W ∗ its conjugate transpose.
The Discrete
Fourier Transforms (DFTs) is a natural choice of T as it can be computed in O ( N log N ) via Fast Fourier Transform (FFT) algorithms (Oppenheim, 1999, Chapter 9.2).
We identify two major limitations of FDMs in the form (1); each layer performs T and T − 1 , and DFTs are complex– valued, requiring parameter–inefficient complex–valued layers.
With T1 , we aim to develop an FDM that does not require more than a single T , while preserving or improving on predictive accuracy.
Ideally, the transform in T1 should be (1) real–valued, to avoid introducing parameter scaling overhead, (2) universal, to allow the representation of target signals, and (3) approximately sparse or structured, to allow dimensionality reduction.
[Equation]
Commutative diagrams for FDM layers (1) and linear FNOs (frequency domain part).
3. Transform Once: The T1 Recipe
With T1 , we introduce major modifications to the way
FDMs are designed and optimized.
In particular, T1 is defined, inferred and trained directly in the frequency domain with only a single direct transform required to process data.
Hence follows the name: transform once ( T1 ).
Direct learning in the frequency domain Consider two signals x ∈ D n , y ∈ D n and suppose there exists a function φ : D n → D n mapping x to y , i.e. y = φ ( x ) .
Then, there must also exist another function ψ : D k → D k that relates the spectra of the two signals, i.e. Y = ψ ( X ) being X = T ( x ) and Y
= T ( y ) .
In particular,
[Equation]
It follows that, from a learning perspective, we can aim to approximate ψ directly in the k -space rather than φ in the n -space.
To do so, we define a learnable parametric function f θ : D k
→ D k and train it to minimize the approximation error J θ of the output signal spectrum Y in the k –space.
Given a distribution p ( x ) of input signals, T is characterized by the following nonlinear program
[Equation]
subject to ˆ
[Equation]
If T is a DFT, the above turns out to be a close approximation (or equivalent, depending on the function class of f θ ) to the minimization of ∥ y − ˆ y ∥ in n -space by the ParsevalPlancherel identity .
Theorem 3.1 (Parseval-Plancherel Identity (Stein and
Shakarchi, 2011, pp. 223)) .
Let T be the normalized DFT.
Given a signal v ∈ D n and its transform V = T ( v ) , it holds ∥ v ∥ =
∥ V ∥ .
This result also applies to any other norm–preserving transform T , e.g. a normalized type–II DCT (Oppenheim, 1999, pp. 679).
For the linear transforms considered in this work, T ( x ) =
Wx, W ∈ C N × N , condition for Th.
3.1 to hold is W to be orthonormal, i.e. W ∗ W
= I .
Note that T1 retains, in principle, the same universal approximation properties of FNOs (Kovachki et al., 2021) as f θ is allowed to operate on the entirety of the input spectrum.
Given enough capacity, f θ can arbitrarily approximate ψ , implicitly reconstructing φ via T − 1 ◦ f θ ◦ T .
Speedup measurements We provide a concrete example of the effect of pruning redundant transforms on computational costs.
We measure wall–clock inference time speedups of depth d T
[Equation]
over an equivalent depth d FNO with layers (1).
The only difference concerns the application of transforms between layers.
Figure 3.1 provides the speedups on two–dimensional signals: on the left, we fix model depth d 6 and investigate the scaling in signal width ( .e. number of channels) and signal resolution.
the right, we fix signal width to be 32 and visualize the interaction of model depth and signal resolution.
For common experimental settings e.g. resolutions of 64 or 128 , 6 layers and width 32 , T1 is at least 10 x faster than other FDMs.
It will later be shown (§4) that T also preserves or improves on predictive accuracy of other
FDMs across tasks.
When T1 is not preceded by online preprocessing steps for inputs x , such as other neural networks or randomized data augmentations, the transform on T ( x ) can be done once on the dataset, amortizing the cost over training epochs, and increasing the speed of T1 further.
3.1. Choosing the right transform
The transform T in T1 is chosen to be in the class of
Discrete Cosine Transforms (DCTs) (Ahmed et al., 1974;
Strang, 1999), in particular the normalized DCT–II,
[Equation]
DCT–II transforms can be computed in O ( N log N ) via an extended FFT (Makhoul, 1980), are real–valued, and can be shown to provide effective representations of smooth 4
Signal Width
S i g n a l R e s o l u t
i o n
Model Depth
S i g n a l R e s o l u t
i o n
T1 Speedup on Two-Dimensional Data continuous signals leveraging classical results on function approximation with Chebyshev polynomials (Trefethen, 2019).
Moreover, DCT–II transforms are the backbone of modern compression codecs for digital images i.e. JPEG, providing further evidence of approximate sparsity induced in D k by this transform.
A real–valued transform allows T1 to at least halve the required number of parameters by f θ , avoiding the overhead of complex–valued parameters.
Although our initial design of T1 included training deep f θ in FFT k -spaces, we observe poor performance and optimization instabilities, also noted by (Li et al., 2020).
Along with parameter inefficiency, these issues are resolved by letting T be a DCT–II.
3.2. Reduced–Order T1 Model and Irreducible Loss Bound
We seek to leverage structure induced in D k by T .
To this end we allow T1 , similarly to (1), to modify specific elements of X and consequently affect only certain frequency components of x .
The reduced–order T1 model is designed to operate only on m < N elements (selected by S m ∈ R N × m ) of the input k -space, i.e. on a reduced k -space D
m ≡ R m of lower dimension.
Thus, we can employ a smaller neural network γ θ : D m → D m for mapping S m X to the corresponding m elements S m Y of the output k -space.
Thus, training involves a truncated objective that compares predictions with elements in the output signal spectrum also selected by S m : [Equation] subject to
[Equation]
Irreducible losses Without any loss of generality, let us assume to select with S
m the first m elements of the k -
[Figure]
space, i.e.
[Equation]
The truncated training objective J θ ( X, Y ) reads as
[Equation]
However, J θ does not represent a good metric to measure the real performance of the reduced–order T1 model.
For that we need to evaluate the complete loss L θ of the approximation task, including the N − m elements of the output k -space discarded by our model, i.e.
[Equation]
It follows that the overall loss L θ is higher than T1 ’s training objective J θ , i.e. L θ =
J θ + R o > J θ , whilst R o represents the irreducible residual loss due to truncation of the predictions ˆ Y k .
The choice of elements and thus magnitude of irreducible loss is a strong inductive bias for reduced–order T1 .
Nonetheless, given a dataset of input–output signals it is possible to perform an a priori analysis on R o to inform the choice of m and S m .
Often, we empirically observe the irreducible error R o for reduced–order T1 to be smaller than for non–reduced–order FDMs i.e R o < (cid:80) K − 1 k
= m ∥
Y k −
T k
(ˆ y ) ∥ with layers of type (1) 4 .
We note further that the reachable component J θ of the objective cannot always be minimized to zero regardless of the approximation power of γ θ .
For each k < m , S m discards N − m frequency components of the input signal which, if different than zero, likely contain the necessary information to approximate ψ k ( X ) exactly.
Specifically, the irreducible lower bound on J θ should depend on “how much” the output’s m frequency components depend on the discarded N − m input’s elements.
A rough quantification of such bound can be obtained by inspecting the mismatch between the gradients of ψ k − γ θ,k ◦ S m with respect to X .
In particular, it holds
[Equation]
+
N − 1 (cid:88)
j = m (cid:12)(cid:12)(cid:12)(cid:12) ∂ψ k ( X ) ∂X j (cid:12)(cid:12)(cid:12)(cid:12)
Unless ∂ X j ψ k ( X ) = 0 holds for all j = m, . .
. , N −and all k i.e. no dependency of the ground truth map in kspace on the truncated elements, there will be an irreducible overall gradient mismatch and thus a nonzero J θ .
3.3. Weight Initialization for Reduced–Order FDMs
FDMs (Li et al., 2020; Tran et al., 2021; Wen et al., 2022) opt for a standard Xavier–like (Glorot and Bengio, 2010) initialization distribution that takes into account the input channels c to a layer i.e. N (0 , 1 c ) .
However, well–known variance preserving properties of Xavier schemes do not hold for FDM layers truncating N − m elements of the kspace.
Notably, Xavier schemes do not scale the variance of the weight initialization distribution based on the number of elements m kept after truncation of the spectrum performed by f θ , leading to the collapse of the outputs to zero.
To avoid this issue in T1 and other FDMs, we develop a simple variance–preserving (vp) that introduces a variance scaling factor based on m and the class of transform.
[Variance Preserving ( vp ) Initialization] Let ˆ x =
W ∗ S ⊤ m AS m Wx be a k -space reduced–order layer and W is a normalized DCT–II transform.
If x ∈ R N is a random vector with [Equation]
Then,
[Equation]
Corollary 3.1 ( vp initialization for DFTs) .
Under the assumptions of Theorem 3.3, if W is a normalized DFT ma- trix we have Re( A ij ) , Im( A ij ) ∼ N (0 , N 2 m 2 ) ⇒ V
[ˆ x ] = V [ x ] .
We report informal proofs of the above results in Appendix
A.
The collapse phenomenon is empirically shown in Figure 3.2 for m = 24 , comparing a single layer of FNO and
FFNO (with Xavier initialization) with FNO equipped with the proposed vp scheme.
Under the assumptions of Corollary ?? , we sample A and compute empirical variances of ˆ x = W ∗ S ⊤ m A ( θ ) S m Wx for several finite batches of input signals x .
We repeat the experiment for signals of different lengths N .
The vp scheme preserves unitary variances whereas the other layers concentrate output variances towards zero at a rate that grows with N − m .
When the learned frequency–domain transformation f θ is obtained, instead of the single low–rank linear layer f θ =
A ( θ ) S m X , as the composition of several layers, preserving variances can be achieved by applying the vp scheme only to the first layer.
For some variants of FDMs e.g. FNO that truncate the spectrum at each layer, vp initialization should instead be applied to all.
4. Experiments
We validate T1 on learning to approximate solution operators of dynamical systems from images.
• In Section 4.1, we apply T1 on the standard task of learning solution operators for incompressible Navier– Stokes, comparing against other FDMs.
In 4.1 we perform a series of ablation experiments on each ingredient of the T1 recipe, including weight initialization and architecture.
In 4.1 we provide scaling laws.
• Section 4.2 we deal with fluid–solid interaction dynamics in the form of higher resolution images ( 128 ).
We consider turbulent flows around varying airfoil geometries, benchmarking against current SOTA (Thuerey et al., 2020).
• In Section 4.3 we show how the computational efficiency of T1 allows learning on unwieldy data without downsampling or building low–resolution meshes.
We consider learning on high–resolution video ( 600 × 1062 ) capturing the turbulent dynamics of smoke (Eckert et al., 2019).
Configuration and model details are reported in the supplementary material.
4.1. Incompressible Navier–Stokes
We show that T1 matches or outperforms SOTA FDMs with less computation on the standard incompressible Navier–
Stokes benchmark.
Losses are reported in n -space (signal space) for comparison.
Setup We consider two–dimensional Navier–Stokes equations for incompressible fluid in vorticity form as described in (Li et al., 2020).
Given a dataset of initial conditions, we train all models to approximate the solution operator at time 50 seconds for high viscosity ( ν = 1 e − 3 ) and at time 15 for lower viscosity ( ν = 1 e − 4 ).
As a metric, we report normalized mean squared error (N–MSE).
Both initial condition as well as solution are provided as images of resolution 64 .
We include as baseline established FDMs, such as Fourier
Neural Operators (FNOs) (Li et al., 2020) and Factorized
Fourier Neural Operators (FFNOs) (Tran et al., 2021).
We indicate with the suffix vp models that employ the proposed variance preserving initialization scheme.
All models truncate to m = 24 , except FFNOs with m = 32 .
Results We perform 20 training runs for each model and report mean and standard deviation in Table 4.1.
T1 reduces solution error w.r.t FNOs by over 20% and FFNOs by over 40% .
A single forward pass of T1 models is on average 2 x faster than FNO and 10 x than FFNOs.
We note that
FFNOs are designed to share parameters between layers, and thus require deeper architectures – and slower, due to more transforms.
In particular, training time ( 500 epochs) for T1 is cut to 20 minutes down from 40 of FNOs, matching the model speedup.
Finally, we report an improvement in performance for FNOs with parameters initialized following our proposed scheme (FNOvp).
Figure 4.1 provides sample predictions in n -space (left) to contextualize the task, in addition to prediction errors in frequency domain (right).
Despite being a reduced order model with m = 24 , T1 +vp produces smaller errors on truncated kspace elements ( k > m ) compared to FNOvp and FFNO.
Ablations on weight scheme and architecture We repeat the previous experiment and report prediction errors for four variants of T1 : same architecture and weight initialization scheme as FNOs ( T1 ), T1 with our proposed vp scheme ( T1 vp), a reduced–order variant with k -space model f θ defined as a UNet architecture ( T1 +), and T1 + with variance preserving scheme ( T1 +vp).
The results in
Table 4.2 provide empirical evidence in support of the vp scheme and its synergistic effect with the proposed architecture.
In particular, combining vp scheme and UNet structure in frequency domain reduces error by half compared to the naive T1 approach.
Scaling laws We verify whether the reduction in predictive error of T1 over neural operator baselines is preserved as the size of training dataset grows.
We perform 10 training runs on the Navier–Stokes ν = 1 e − 4 experiment, each time with a larger dataset size, and report the scaling laws in Figure 4.2. With additional data, the gaps in test errors [Table] [Table]
[Table] [Table] T e s t N – M
S E
Data Scaling
[Figure]
narrow slightly, with noticeable improvements obtained by applying the vp scheme to both FNO and T1+ .
4.2. Flow Around Airfoils
We investigate the performance of T1 in predicting steady– state solutions of flow around airfoils.
Setup We use data introduced in (Thuerey et al., 2020) in the form of 10000 training pairs of initial conditions, specifying freestream velocities and the airfoil mask, with the target steady–state velocity and pressure fields.
This task introduces additional complexity in the form of higher resolution input images ( 128 ) and a full k -space due to the discontinuity in the field produced by the mask.
We compare a SOTA UNet architecture (DFPNet) introduced by (Thuerey et al., 2020) to FNOs and T1 with vp initialization schemes.
We perform a search on the most representative hyperparameters (detailed in the Appendix). Averages for 5 runs are reported in Table 4.3.
Method
N–MSE
Time (hrs)
DFPNet
[Table]
[Table]
FNO
[Table]
T1 +
[Table]
Results All models are able to accurately predict steady– state solutions for different airfoils.
FDMs and T1 offer a slight improvement over DFPNets (Thuerey et al., 2020), with FNOs introducing a significant computational overhead.
Training of T1 is as fast as DFPNets and as accurate as FNOs.
We note standard deviation of the results σ = 0 .
001 ; differences between FNO are explainable by natural variation T1 ( ≤ 1 σ ) whereas improvements over
DFPNet are ≥ 6 σ .
[Figure]
4.3. Turbulent Smoke
We investigate the performance of T1 in predicting iterative rollouts from high–resolution video of real rising smoke plumes.
Setup We use the ScalarFlow dataset introduced in (Eckert et al., 2019) consisting of 104 sequences of 150 frames each collected from video recordings of rising hot smoke plumes.
The dataset consists of raw video data at high– resolution ( 600 × 1062 ) collected at 60 fps.
This task scales up complexity by involving real–world high–definition data, capturing highly–turbulent dynamics.
We perform rollouts iteratively based on previous predictions: all models are trained on 3 –step rollouts and evaluated over 10 – steps extrapolation to test their generalization in time.
We compare FNOs against T1 and T1 + of similar model sizes, after performing a search on most representative hyperparameters (Appendix B).
Method
[Table]
Results Figure 4.3 provides a sample rollout of different model predictions in k -space (DCT–II).
T1 accumulates smaller errors over the rollout and is less prone to generation of non–physical artifacts by performing prediction only on a subset of the k -space (Table 4.4).
Notably, T1 and
T
+ are 4 × to 7 × faster, providing a reduction in training time from 32 . 4 hours to 4 . 7 .
In Appendix B, we include additional visualization, including prediction errors on kspace.
5. Conclusion
We present a streamlined class of frequency domain models (FDM): Transform Once ( T1 ).
T1 models are optimized directly in frequency domain, after a single transform, and achieve similar or improved predictive performance at a fraction of the computational cost ( 3 x to 10 x speedups across tasks).
Further, a simple truncation–aware weight initialization scheme is introduced and shown to improve performance of T1 and existing FDMs.
A. Weight Initialization
Let us assume the input sequence x to be real and distributed according to an unknown distribution p , i [Equation]
e. x ∼ p ( x ) .
We seek a principled initialization scheme for A that allows the first central moments of p X be preserved by a reduced order
FDM layer.
Moments under a DFT:
Let x be a real–valued input distributed according to
[Equation]
Consider a single element of X
[Equation]
with
[Equation]
For clarity, we will treat the real part Re( X k ) first.
[Equation]
and
[Equation]
where we have used the fact that
[Equation]
Thus,
[Equation]
We observe that (a) the first central moment is preserved and (b) while the variance term can be simplified as
[Equation]
We follow a similar procedure for Im( X k ) , arriving at
[Equation]
where the variance again simplifies to
[Equation]
Since X k = Re( X k )
+ j
Im( X k ) , its statistics are given by
[Equation]
A similar argument can be developed using basic properties of circular–symmetry of complex Normals.
It is critical that the normalization factor 1 √ N be included in W in order to preserve the variance of V
[ X ] .
Indeed, normalization factors used in different conventions lead to different results
[Equation]
As N can easily be in the order of hundreds or thousands for generic signals, explosion of variance can be an issue if the orthogonalization factor 1 N is not applied to W .
A.1. Variance response of transformations in spectral layers
Having determined the preservation of probability under Fourier Transforms, we now analyze of the first two central moments of the distribution of a spectral layer output ˆ x are affected by a learned transformation on X .
Gaussian Initialization with Dense Linear Layer Case:
Once again, we consider a single element ˆ X k of ˆ X
[Equation]
In this case, p ˆ X k is a sum of product distributions involving independent random variables A ki and X t .
The first central moment is readily obtained
[Equation]
if at least one of the conditions
[Equation]
or ∀ k, i : E
[ A ki ] = 0 holds.
The variance
[Equation]
Since
[Equation]
Solving for σ A , we arrive at a sufficient condition for variance preservation i.e.
[Equation]
The result for DCT can be derived by following the same steps.
Since the k–space for DCTs is real, the variance sums to
[Equation]
B. Additional Details
Broader impact FDMs are widely used in the context of learning to predict the evolution of dynamical systems.
The model class presented in this work, T1 , provides an accessible way to train and evaluate large–scale FDMs, reducing memory overhead and overall training times.
When predicting the solution of e.g. a partial differential equation (PDE), care should be taken especially when the prediction is used to inform downstream decision making, as many systems are optimally predictable only for a certain time scale (Strogatz, 2018, pp. 366).
We anticipate a potential positive environmental impact from the adoption of T1 as a replacement for the largest FDMs currently in use.
Experimental setup Experiments have been performed on an NVIDIA© DGX workstation equipped with a 128 threads
AMD© E PYC 7742 CPU, 512GB of RAM and four NVIDIA© A100 GPUs.
The main software implementation has been done within the PyTorch (Paszke et al., 2017) ecosystem building upon the pytorch - lightning (Falcon et al., 2019) framework.
Common experimental settings
B.1. Incompressible Navier–Stokes
Dataset We use data generated in (Li et al., 2020) in the form of pairs of initial conditions and solutions of the incompressible Navier–Stokes equations in vorticity form solved with a pseudospectral method.
The dataset 5 is comprised rollouts of solutions as images of resolution 64 .
Models and training The training configuration is shared by all models: 1 datamodule : ntrain : ntest : batch_size : history_size :6 train : optimizer : type : AdamW learning_rate : 1eweight_decay : 1escheduler : type : Step step_size : gamma : 0.
scheduler_interval : epoch {+
+} loss_fn : RelativeL2Loss
For the high viscosity ( 1 e − 3 ) setting, the models are trained to predict the solution at time T = 50 seconds directly, without producing rollouts and supervising the model with solutions at times between 0 and 50 .
Crucially, this ensures that the task is much more challenging than that of (Li et al., 2020), where for a single training sample the entire rollout is used as supervision.
For the low viscosity setting ( 1 e − 4 ), target times are T = 15 seconds.
Model configurations are given below:
FNO 1 modes {+:2 nlayers :
+}
width :
T  1 modes :2 nlayers {+:
+}
width :
FFNO 1 modes {+:2 nlayers :
+}
width :where each layer in a model shares the same structure.
In FNOs and FFNOs, we employ a regular FDM layer following (Li et al., 2020; Tran et al., 2021) with k–space convolutions and residual connections given by n–space layers (pointwise convolutions for FNOs, dense for FFNOs).
T1 uses a similar layer without n–space residual paths.
The differences in number of layers and width have been introduced to keep parameter counts comparable.
At a given channel width, FNOs require the largest number of parameters due to k–space convolutions on complex numbers given by the DFT coefficients.
Although FFNOs (Tran et al., 2021) are most parameter efficient due to parameter sharing, we found them unable to tackle the task and produce high–quality predictions.
T
+ employs a UNet on the patch constructed by the elements of the k–space kept, and shares its structure with T otherwise.
The vp parameter initialization scheme in T1 is applied only to the first layer performing the truncation in k–space, not to the following layers which use standard Kaiming initialization
(He et al., 2015). In FNOvp the scheme is applied to all layers.
Hyperparameter tuning We start with the basic model structure of FNOs as detailed (Li et al., 2020) and perform a basic hyperparameter search on a small slice of the training set, with the goal of ensuring proper convergence of a model.
We did not find the number of layers to have a significant impact on convergence.
Width plays an important role and is best kept above 24 .
64 DCT modes kept 10 −10 −
M
[Figure] 64 DCT modes kept [Figure] I n i t i a l
V o r t
i c i t y
T a r g e t
T1+vp
F N O v p
FF N O
Scaling laws We use the same settings as the main experiment, repeating separate training runs for the low viscosity setting.
In particular, we increase the dataset size for each set of runs by a factor of 2 : 1024 , 2048 , 4096 , 8192 .
The total number of epochs is kept fixed, so that more iterations are performed for larger datasets.
The same test set of size 200 is used in all cases.
Further comments Additional predictions are provided in Figure B.2.
Figure B.1 shows the approximation error on the
Navier-Stokes solutions due to truncation at different number of k–space elements m .
B.2. Flow Around Airfoils
Dataset We use a slice of the dataset introduced by (Thuerey et al., 2020) in the form of 11000 training pairs of initial conditions and solutions.
The solutions are obtained via OpenFOAM (Jasak et al., 2007)
SIMPLE, a steady–state solver for incompressible and turbulent flows.
In particular, the initial conditions are specified as freestream velocities over the domain (two–directional components), in addition to a specification of the airfoil in point cloud format.
Delaunay triangulation is used for mesh generation.
After simulation, data is provided as initial condition and steady–state solution pairs.
The initial condition is a three channel 128 × 128 image: two channels for freestream velocities and one for the airfoil mask.
The solution is a three channel 128 × 128 image: a velocity field and a scalar pressure field.
All data is normalized using training set statistics.
Models and training Training configuration is given as 1 datamodule : ntrain : nval : ntest : batch_size :6 train : optimizer : type : AdamW learning_rate : 1eweight_decay : 1escheduler : type : Step step_size : gamma : 0.
scheduler_interval : epoch {+
+}
loss_fn : RelativeL2Loss
The baseline UNet matches the architecture of (Thuerey et al., 2020)
(DFPNet).
The FNO architecture is comprised of a standard stack of FDM layer as discussed in B.1.
The k–space UNet in T1 + has the same structure as a DFPNet.
FNO 1 modes {+:2 nlayers :
 width+} :
DFPNET 1 channel_exponent :
T1 + 1 modes :2 channel_exponent :
Hyperparameter tuning This is an example of a dataset where the k–space is full due to discontinuity in the solution given by the airfoil mask.
We use the training and validation sets to inspect the k–space and set m to 100 for the irreducible loss term to be sufficiently small as shown in Figure B.5.
We swept over m for FNOs and found larger than 24 to perform worse, likely due to k–space convolution being sufficient to capture higher frequency components.
We observe DFPNets with larger channel exponents perform worse due to overfitting.
Further comments A sample of predictions is given in Fig. B.3. Figure B.4 shows the n–space and corresponding DCT k–space of a data point. As can be observed, the k–space is structured but full due to the discontinuity caused by the airfoil mask.
Figure B.5 shows the approximation error on solution fields due to truncation in k–space at different m .
In this task, the DCT is more efficient, given a budget of modes to keep, as it yields lower errors.
This error provides a theoretical lower bound for the predictive error achievable by a T1 model with a given budget, reachable only if the T1 predicts the first m modes perfectly.
The vertical line indicates the budget used for the main text T1 experiments ( m = 100 ), and the horizontal one the test N–
MSE achieved.
Various segments of the vertical line indicate reducible and irreducible components of the loss as discussed in Section 3.2.
The theoretical limit at m = 100 is well below what has been empirically achieved by T1 and other models.
Indeed, the irreducible loss is an order of magnitude smaller than what the best model (including non–reduced–order variants) achieves on the task.
T a r g e t
T1vp
F N O v p
D F P
N e t
[Figure]
”Flow around airfoils” n-space and k-space
[Figure]
N-space k- modes kept B.3. Turbulent Smoke Dataset We employ for this experiment the ScalarFlow dataset introduced in (Eckert et al., 2019) which is available online under the Creative Commons license CC-BY-NC-SA 4.0 6 .
(Eckert et al., 2019) created an environment for controlling the release of smoke plumes: a fog machine generated fog inside of a container; the fog was then heated up by a heating cable and a valve controlled its release.
Data was captured via multiple calibrated cameras in high resolution at 60 fps (frames per second) for 150 frames
[Figure] DCT modes. The dataset contains 3D reconstructions of the smoke plumes and 2D input and rendered images: input images are used by (Eckert et al., 2019) to solve an optimization problem in which the goal is to generate a 3D reconstruction that minimizes the difference between input and rendered images.
2D input images are obtained directly from raw data on which only post–processing is applied by (Eckert et al., 2019) in the form of gray scaling and denoising: these are saved in compressed numpy (Harris et al., 2020) arrays named imgsTarget _ 000xxx .
npz .
Each resulting frame comprises 5 different camera views 600 × 1062 in size.
Since we want to use T1 on high–resolution experimental data, we directly utilize the central camera view of these input images in our learning task without any further downsampling or data processing.
Similarly to (Lienen and Günnemann, 2022), we divide the 104 recordings into the first 64 for training and use the remaining 20 for validation and 20 for testing.
Data is normalized to the [0 , 1] range based on training dataset statistics.
Hyperparameter selection and tuning We performed a search on the most representative hyperparameters.
One of the most important hyperparameters to choose from is the number of DCT modes to keep, i [Figure] e. first m elements in k –space.
We note that for simplicity as well as for compatibility with the UNet inside of T1 +, we consider a square mode pruning, i [Figure] e. we keep the same number of frequencies on both height and width of the image and refer to the modes kept in both dimensions as m .
Figure B.7 and Figure B.6 show trends of DCT modes in terms of errors and visual quality: while the first modes m contribute the most to the quality of the representation in n –space, the last elements contribute only to high–frequency details whose effect is minor on the overall reconstruction.
Thus, we set T1 + to m = 224 and consequently
T1 to m = 512 to have comparable model sizes.
We set m = 48 for FNO due to memory and model size limitations, noting that its residual connections effectively enlarge the training spectrum to all possible frequencies as shown in Figure 4.3.
Similarly to other experiments (B2) [Figure] we observe raising m in FNO to not significantly improve predictive error, even when the additional k–space elements would include a larger portion of the dataset [Figure]
We also experiment with different iterative rollout update strategies as in (Pfaff et al., 2020).
We consider the time step ∆ t to be unitary, i.e. ∆ t = 1 , given that the training frames are sampled consistently at 60 fps.
We call 0 –order integration an update of the type: x t +1 =
h θ ( x t ; x t − 1 , . . .
, x t − H ) in which h θ denotes a learned model which takes as inputs the current state x t and optionally a history of size H of past states x t − 1 , . . .
, x t − H and directly predicts the next state x t +1 .
A 1 –order integrator performs the following update: x t +1 = x t
+ h θ ( x t ; · ) , in which the model predicts the state update, i.e. the velocity , similarly to an Euler step.
A 2 –order integrator, also known as basic Störmer—Verlet (Verlet, 1967) can be written as following: x t +1 = 2 x t − x t −
+ h θ ( x t ; · ) ; the model h
θ predicts the acceleration of the system.
We empirically found the zero–order integration to be more prone to generating artifacts with slower convergence, which may be because the model has to directly predict the next step with no "help" from the current step information.
We found models trained with first–order integrators to have lower predictive errors than those trained with second–order ones, and we thus use it in all the experiments.
As for the history size, we selected H = 1 since it provided noticeable benefits compared to H = 0 , in which the model has no way of knowing previous states and thus inferring velocities.
Larger history sizes did not seem to provide any improvements and only made the models larger as also noted in (Pfaff et al., 2020).
Models and training All models share the following configuration for training: 1 datamodule : ntrain : nval : ntest :
batch_size : history_size : target_steps_train : target_steps_val_test :9 train : optimizer : type : AdamW learning_rate : 1eweight_decay : 1escheduler : type : CosineAnnealingWarmRestarts
T_0 :
step_size : scheduler_interval : step 19 loss_fn : RelativeL2Loss
Where we used the implementation in PyTorch of the cosine annealing schedule with warm restarts 7 .
The FNO architecture comprises a standard stack of FDM layers as discussed in B.1.
The k–space UNet in T1+ has the same structure as a
DFPNet.
[Figure] nlayers [Figure] width :
T  1 modes {+:2 nlayers :
 width+} :
T1 + 1 modes {+:2 nlayers :
 width :4 channel_exponent+} :where we note that all models employ GeLU (Hendrycks and Gimpel, 2016) activation functions between inner layers.
[Figure]
Analysis of results Table B.1 provides a larger version of the table in the main text, including 1 –step mean absolute errors (MAE).
We note that while FNO produces smaller errors in one–step predictions, it quickly accumulates larger errors in extrapolation.
Figure B.8 shows mean errors in k –space of FNO vs T1 and T1 +.
T1 demonstrate smaller overall errors and lower maxima compared to the FNO.
[Table]