Spectral Shaping for Neural PDE Surrogates

Thanks for your comments and requests for a more in-depth theoretical treatment of how nonlinearities introduce spectral junk. We can certainly strengthen our exposition with a detailed mathematical description of the mechanism whereby it is introduced. We can explain the phenomenon through amplitude distortion. Mathematically what happens is as follows:

Consider a scalar-valued function $u(x)$ and a pointwise nonlinearity $f: \mathbb{R} \to \mathbb{R}$, such that $v(x) = f(u(x))$. We will study how the Fourier transform $V(k)$ of $v(x)$ for wavenumber $k$ is related to the Fourier transform $U(k)$ of $u(x)$.

To do this we approximate $f$ with its Taylor expansion about zero (we assume f is smooth) as

$v(x) = f(u(x)) = f(0) + u(x) \cdot f’(0) + u(x)^2 \cdot f’’(0) / 2! + u(x)^3 \cdot f’’’(0) / 3! + \text{h.o.t.}$

Then note how the nonlinear terms are monomials. The Fourier transform of $u(x)^2$ is the autoconvolution $U * U$, where $*$ denotes convolution. For the nth power, we have the convolution $U$ with $n-1$ copies of itself. This serves to broaden out the spectrum, mixing high and low frequencies of the spectrum and makes the spectrum more Gaussian in shape (indeed this is the same reason why sums of random variables undergo Gaussianization in central limit theorems, since sums of random variables are convolutions of their PDFs). For smooth signals $u(x)$ most of the energy is in the low frequencies. Passing $u(x)$ through a nonlinearity $f$ moves some of the energy through to the higher frequencies through this autoconvolutional mechanism. This is why we see an increase in power in the higher frequencies of the spectrum.

Note that we do not require there to be any discontinuities for this effect to be seen and in the above treatment all the operations were smooth. Indeed in our submission all PDEs are smooth and all architectures are composed of smooth operations. Regarding our comment about spectral junk being “noise”, we can tighten up our phrasing here and reword this as interference of the signal with itself, which is exactly what this is. Indeed this phenomenon is different to aliasing, being induced by a different mechanism.

We hope this clears up any confusion concerning the origin of spectral junk. We shall include this in an updated draft of the paper.

Minor points

• In acronyms are used before they are defined, e.g. the use of KS in the caption of Figure 3 before the use of Kuramoto-Shivashinsky in the text. Thanks for spotting this! This is being corrected.

• In some cases the language could be slightly more formal in an academic paper. For example the use of ``sanity check" and "probably" in Section 5.1. This is certainly personal preference, but we shall strive to update this.

Thanks for the review. Please see our rebuttal and let us know if anything is missing. We have tried to address all your concerns. If you believe we have improved the exposition, please do upgrade your rating.

The presented study combines multiple known aspects of the rollout problem but has limited novelty on its own. Besides the eigen-spectrum analysis, most of the insights from the presented study have already been presented in other works. The inability to model high-frequency noise is a known problem and has been described in previous publications. (Lippe 2024, McCabe 2023). In McCabe et al. (2023), the noise problem was explicitly attributed to the non-linearity, and they proposed to learn to filter this noise out (although for FNO). In Raonić et al. (2024), formal solutions are introduced to avoid aliasing produced by non-linearities in convolutional models.

Thanks for pointing out those connections. The papers by Lippe, McCabe, and Raonić are all great works.

Lippe et al. concerns itself with spectral shape, and we compare against their paper PDE-Refiner below. We find while they achieve good decorrelation time, so can we, with only a single pass through the model. We also achieve good spectral shape (low MELR), which implies better statistics in the far-term after we have rolled out far into the future.

Raonić et al. concerns itself with aliasing. We do not. We are concerned with the introduction of spectral junk due to amplitude distortion from nonlinearities. Indeed the two processes are related and this is a well-known phenomenon, which we explain in detail to reviewer TbPC, which you may read at your leisure. The technique in the CNO paper by Raonić et al. is to oversample, but this only works for sufficiently bandlimited signals, which cannot be guaranteed with general nonlinearities. We accept this and opt for low pass filtering, instead.

The proposed method is not compared against any other architectural methods that tackle the noise problem, e.g., PDERefiner, in the experiments. While the presented method is computationally cheaper, the question remains whether it can increase the decorrelation time as much as the other methods tackling the high-frequency noise problem.

We have added some more baselines. But first an explanation why we did not initially include them. Models such as FNO, Galerkin Transformer, and other neural operator models are learned maps directly from initial conditions to solution values at time t, namely, $u(., 0) -> u(., t)$. Meanwhile the focus of this paper is on autoregressive methods. That said, nothing stops us from training a Neural Operator model in an autoregressive fashion and we do so for the rebuttals, but this is really not the arena where such models shine. Regarding Galerkin Transformer, notions of spatial geometry are lost in the attention layer and so it is not obvious how to apply spectral shaping in that space. PDE-Refiner, a diffusion-type model, is stochastic and was not the focus of this paper, but we include it, since it is state-of-the-art. It cannot be combined with spectral shaping.

We provide these for the KS equation, which is all we had time for in the rebuttal period. We report MELR (Mean energy log ratio, a measure of spectral shape) at 128 time steps (lower is better) and decorrelation time at 0.8 Pearson’s correlation coefficient (higher is better).

(Bold indicates best in each row)

Decorrelation time at 0.8 (first timestep when correlation drops below 0.8)

The findings we had in the paper that spectral shaping improves MELR holds for all baselines apart from the Galerkin Transformer, where, as we stated, it is not directly obvious how to apply spectral shaping. Furthermore, note how most models with spectral shaping outperform PDE-Refiner on MELR. Note that PDE-Refiner was originally developed and motivated with improving spectral shape in mind.

In terms of decorrelation time, spectral shaping always improved this.

The method is only evaluated for a single model (U-Net). This leads to the question of whether the good experimental results are due to using a model that may be especially prone to producing high-frequency noise. For example, the paper could be improved by evaluating whether the filter helps FNO. This would also allow a comparison to the method of McCabe et al. (2023).

Please see above

Thanks for the review. We have tried to address all critique and suggestions to the best we can in the given time window. If you too believe we have addressed some of your concerns please do revise you score to reflect this.

First, a general comment is that the paper is hard to read and assumes the reader is well-read with the literature. Further explanation of concepts, such as defining how the PDE solutions' low, medium, and high-frequency bands are identified and explaining the eigenvalue spectrum of the linearized flow map, would help the reader understand the challenges and motivate them to read the work.

Thanks for this feedback. We have tried to strike a balance between explaining everything from scratch and keeping things short, to cater to a broad readership and indeed some reviewers had the opposite experience of the paper. We will update the explanations in the areas you mention to improve readability. Specifically the different frequency bands are more qualitative regions than anything very specific, read, left, middle and right of spectra, but we can tighten this up by referring to the “energy containing range”, “inertial range”, and “dissipative range” of classical kinetic turbulent energy spectra. Regarding the eigen-spectrum, we can provide an exact description in the Appendix.

The chosen baselines are vanilla, and more advanced variants of the architectures could have been considered as the baseline, for e.g. advanced variants of U-Net have been applied for temporal conditioning and long rollouts.

The original idea to focus on simple baselines was motivated by the fact that phenomena are easier to analyze. That said, we understand in such a deeply empirical field demonstrating an idea works is more convincing on a wider range of models. We accept this criticism and present a number of extra baselines. We provide these for the KS equation, which is all we had time for in the rebuttal period

We have added some more baselines. But first an explanation why we did not initially include them. Models such as FNO, Galerkin Transformer, and other neural operator models are learned maps directly from initial conditions to solution values at time t, namely, $u(., 0) -> u(., t)$. Meanwhile the focus of this paper is on autoregressive methods. That said, nothing stops us from training a Neural Operator model in an autoregressive fashion and we do so for the rebuttals, but this is really not the arena where such models shine. Regarding Galerkin Transformer, notions of spatial geometry are lost in the attention layer and so it is not obvious how to apply spectral shaping in that space. PDE-Refiner, a diffusion-type model, is stochastic and was not the focus of this paper, but we include it, since it is state-of-the-art. It cannot be combined with spectral shaping.

We provide these for the KS equation, which is all we had time for in the rebuttal period. We report MELR (Mean energy log ratio, a measure of spectral shape) at 128 time steps (lower is better) and decorrelation time at 0.8 Pearson’s correlation coefficient (higher is better).

(Bold indicates best in each row)

Decorrelation time at 0.8 (first timestep when correlation drops below 0.8)

The findings we had in the paper that spectral shaping improves MELR holds for all baselines apart from the Galerkin Transformer, where, as we stated, it is not directly obvious how to apply spectral shaping. Furthermore, note how most models with spectral shaping outperform PDE-Refiner on MELR. Note that PDE-Refiner was originally developed and motivated with improving spectral shape in mind.

In terms of decorrelation time, spectral shaping always improved this.

The challenge of the long rollout is even persistent for 3D stiff problems; how well the method scales to a higher-dimensional problem still needs to be checked.

This is certainly true. For papers targeting machine learning conferences, we believed that 3 equations in 1D and 2D would be enough. 3D problems present an entirely different beast, that many in the field do not come on to. We would like to keep it that way for this paper

The training done in fp64 for 1D problem and fp32 for 2D problem.

We chose the floating point format to match the format in which the data is generated. It did not appear to be much of an issue for us. Importantly, however, running the 1D problem in fp64 highlights how spectral shaping is able to perform with a much higher dynamic range compared to other method, such as PDE-Refiner.

The authors wrote "the choice of sigma is important", but never reported the comparison.

Please revisit Figure 8. You will notice the box plots show the range of results after sweeping over sigma. We will update the paper to highlight this fact.

Why not making sigma learnable in the filter?

We did try this and it did not work for us. We are not sure why, but the most likely answer is that sigma controls the ability of the model to overfit over a single step. Therefore, making sigma learnable will not work unless we add a regularizer on the sigma to encourage it towards larger values.

What is a "high-frequency regularizer"? In “Learning to minimize spectral gap” we added the MELR loss to the L2 loss. This is in effect a regularizer on the high frequency content to match the ground truth. Alone this is not enough to encourage the model to match spectra, low pass filtering is needed. That is what we meant. We shall update the paper accordingly.

Thanks for your review. We are generally satisfied and believe that you have ingested the central ideas of the paper. Below we address your comments and critique. We hope this improves the paper. If you also believe this, please update your rating to reflect this.

No modern baselines, such as FNO or FNO variants, Transformer-based NOs (Galerkin, OFormer, GNOT, and many others), or DeepONets.

We have added some more baselines. But first an explanation why we did not initially include them. Models such as FNO, Galerkin Transformer, and other neural operator models are learned maps directly from initial conditions to solution values at time t, namely, $u(., 0) -> u(., t)$. Meanwhile the focus of this paper is on autoregressive methods. That said, nothing stops us from training a Neural Operator model in an autoregressive fashion and we do so for the rebuttals, but this is really not the arena where such models shine. Regarding Galerkin Transformer, notions of spatial geometry are lost in the attention layer and so it is not obvious how to apply spectral shaping in that space. PDE-Refiner, a diffusion-type model, is stochastic and was not the focus of this paper, but we include it, since it is state-of-the-art. It cannot be combined with spectral shaping.

We provide these for the KS equation, which is all we had time for in the rebuttal period. We report MELR (Mean energy log ratio, a measure of spectral shape) at 128 time steps (lower is better) and decorrelation time at 0.8 Pearson’s correlation coefficient (higher is better).

(Bold indicates best in each row)

Decorrelation time at 0.8 (first timestep when correlation drops below 0.8)

The findings we had in the paper that spectral shaping improves MELR holds for all baselines apart from the Galerkin Transformer, where, as we stated, it is not directly obvious how to apply spectral shaping. Furthermore, note how most models with spectral shaping outperform PDE-Refiner on MELR. Note that PDE-Refiner was originally developed and motivated with improving spectral shape in mind.

In terms of decorrelation time, spectral shaping always improved this.

There should be an ablation of different filters (at least some filters with sharper cut-offs than Gaussian, Butterworth for example).

We found between Gaussian, Chebyshev, and Butterworth filters and found them all to perform similarly. The Gaussian filter was simple enough, which is why we just focussed on that. We will add the ablations to the appendix, but you will find that the effects are small. Essentially, any filter which suppresses high frequencies works.

The linearized flow map is unclear, is the linearization at the current time step?

We assume you are referring to Figure 1a. Indeed, the linearization is performed about the current timestep for a randomly sampled trajectory in the test set. This is mainly intended as an exemplar of a phenomenon we have observed across many trajectories and timesteps.

There are some minor notational errors, such as the PDE for the Kolmogorov flow is wrong, and the divergence-free condition in RBC is written as something else.

Thanks for catching the error in the forcing for Kolmogorov flow! Regarding RBC, we have used the problem description as implemented in https://dedalus-project.readthedocs.io/en/latest/pages/examples/ivp_2d_rayleigh_benard.html

I would argue differently from Figure 8 on the usage of only filtering the last layer vs every layer. RBC is a different beast than others, as one of the major difficulties for neural surrogates is the boundary conditions (not so big for periodic boxes) after repeated application.

This is an interesting point. We do agree that for RBC the boundary conditions are important. However, what Figure 8 highlights is that the model is not able to achieve a lengthened decorrelation window without some level of filtering (or equivalent such as denoising, pushforward, etc.). Indeed, as you say, solving the boundary conditions is an important component in solving RBC, but it is complementary to spectral shaping, not in competition.