Fourier Feature Pyramids for Physics-Informed Neural Networks

TL;DR

beignet replaces the random Fourier feature embedding used by current physics-informed neural networks (PINNs) with a learnable multi-resolution Fourier feature pyramid over equispaced grids. Querying a continuous coordinate uses Fourier interpolation at each pyramid level; spatial derivatives for the PDE residual are computed spectrally via FFT and chained through a small MLP decoder with automatic differentiation. The architecture beats JAX-PI and PirateNet by 6-17× relative L2 error on four JAX-PI benchmarks (Allen-Cahn, KdV, Ginzburg-Landau, Gray-Scott) at a quarter to half the parameter count [Table 1], and reaches near machine-precision residuals on a 1D inviscid Burgers self-similar blowup problem using only Adam — an accuracy regime previously only reached with expensive second-order optimizers like BFGS [Table 2].

Key claims

Random Fourier feature embeddings underperform a learnable multi-resolution Fourier pyramid on PINN benchmarks at matched parameter count, because the pyramid directly controls spatial bandwidth per level and lets the model allocate feature capacity across resolutions instead of relying on a single Gaussian bandwidth hyperparameter [§4.1].
Spatial derivatives of the feature field can be computed in Fourier space by multiplying DFT coefficients by wavenumbers and inverse-FFTing to the collocation grid, then chained through the MLP decoder with Faà di Bruno’s formula via nested forward-mode autodiff — avoiding the quadratic-in-parameter cost of autodiff through a large MLP that dominates existing PINN compute [§4.3].
On four JAX-PI benchmarks (Allen-Cahn, KdV, Ginzburg-Landau, Gray-Scott) beignet reduces relative L2 error by 6-17× over PirateNet at 0.25-0.5× the parameter count; the gap is largest on KdV (third-order dispersive derivative amplifies high-frequency content) and Gray-Scott (coupled 2D reaction-diffusion with fine-scale pattern formation) [Table 1].
On a 1D inviscid Burgers self-similar blowup problem, a 12-level beignet pyramid trained with Adam reaches PDE MSE ~1e-14, matching coordinate-MLP + BFGS second-order optimization at a residual regime two to three orders of magnitude lower than coordinate-MLP + Adam [Table 2].
Different PDEs favor different spectral allocations under a fixed parameter budget: Allen-Cahn (sharp interfaces, diffusive damping) benefits from more pyramid levels and finer grids, while KdV (third-derivative wavenumber-cubed amplification) prefers compact bandwidth with capacity concentrated at low/mid frequencies [§5.3, Figure 2].
Stochastic grid shift — randomly offsetting the collocation grid in Fourier space via a phase ramp — is needed to prevent the model from satisfying the residual at grid points while remaining inaccurate between them, an aliasing failure mode specific to fixed-grid Fourier feature architectures [§4.4].

Method

The core primitive is a Fourier feature pyramid: a stack of learnable periodic feature grids at coarse-to-fine spatial resolutions (typically dyadic, e.g. 8, 16, …, 256). To evaluate at a continuous coordinate x, the model performs Fourier interpolation on each grid — multiplying DFT coefficients by complex exponentials at x and summing over modes — to obtain a per-level feature vector; these are concatenated and decoded by a small MLP trunk. For time-dependent PDEs, time is split into bins and each bin stores its own pyramid; evaluating at time t linearly interpolates between the two bracketing pyramids before Fourier interpolation.

The derivative pipeline is the key engineering lever. Rather than back-propagating through the whole MLP for every residual term, beignet differentiates the feature field spectrally — the kth spatial derivative is DFT coefficients multiplied by (2πik)^k — then composes with MLP autodiff via the chain rule (Faà di Bruno for higher orders). Because the feature grid is equispaced, feature values and derivatives at the structured collocation grid are evaluated in a single inverse FFT; the MLP side is lightweight because the backbone can stay small when the feature capacity scales independently. Stochastic grid shift during training breaks the alignment between the collocation grid and the Fourier basis, preventing high-frequency aliasing between grid points.

Results

JAX-PI benchmark suite, relative L2 vs PirateNet: Allen-Cahn 1.1e-5 vs 4.6e-5 (4.2×), KdV 2.3e-4 vs 4.0e-3 (17×), Ginzburg-Landau 7.7e-4 vs 2.1e-3 (2.7×), Gray-Scott 1.6e-3 (field 1) and 1.9e-3 (field 2) vs 2.4e-2 and 3.1e-2 (15-16×) [Table 1]. beignet uses ~412K parameters vs PirateNet’s ~1.9M on 1D problems.
1D Burgers self-similar blowup, PDE MSE on fixed diagnostic grid: coordinate MLP + Adam 3.1e-9, coordinate MLP + BFGS 1.2e-13, beignet 12-level pyramid + Adam 1.1e-14 [Table 2]. Peak pointwise profile residual drops from ~2e-4 (coordinate MLP + Adam) to ~2e-7 (beignet + Adam).
Spectral allocation ablation [§5.3, Figure 2]: holding parameter count fixed near 400K, Allen-Cahn error monotonically improves as max grid resolution increases up to the reference-grid resolution; KdV error is non-monotone and prefers a mid-range max resolution with more capacity concentrated at lower levels — direct evidence that one-knob RFF bandwidth is strictly worse than per-level allocation.
Stochastic grid shift ablation [Appendix B.1]: without the shift, training-grid residual drops to machine precision while the off-grid test residual plateaus orders of magnitude higher — confirms the aliasing failure mode is real and that the shift is load-bearing, not cosmetic.

Why it’s interesting

The architectural idea — replacing a scalar bandwidth hyperparameter with a learnable per-level feature grid, then differentiating spectrally via FFT and chaining through autodiff — is a cleanly transferable primitive beyond PINNs. In the wiki’s existing coverage, PixNerd: Pixel Neural Field Diffusion already pushes neural fields into pixel-space generation by having each DiT token predict the weights of a small MLP queried per-pixel; beignet is a complementary design that keeps the MLP shared and tiny but puts the capacity into a fixed-structure, bandlimited feature pyramid whose derivatives are computed in closed form. Both bets fight the same enemy (autodiff through a large MLP dominates compute), from opposite directions.

The second connection runs through the Navier-Stokes thread the wiki filed in September: the paper explicitly motivates the Burgers blowup experiment as a stepping stone toward 3D Navier-Stokes singularity analysis, and the author’s thread cites arxiv 2509.14185 on unforced Navier-Stokes. This sits beside On the Navier–Stokes Millennium Prize Problem‘s agent-swarm claim and Statement on OpenAI Navier–Stokes priority dispute (Buckmaster, 2026-09-08)‘s priority statement — but represents a very different bet: a tighter neural representation that lets a first-order optimizer reach residuals where second-order methods currently monopolize, enabling direct computer-assisted blowup studies without the LLM-agent-swarm machinery.