Model Inversion for Spatio-temporal Processes using
the Fourier Neural Operator

Dan MacKinlay
CSIRO Data61
dan.mackinlay@data61.csiro.au

Dan Pagendam
CSIRO Data61
dan.pagendam@data61.csiro.au

Petra M. Kuhnert
CSIRO Data61
petra.kuhnert@data61.csiro.au

Tao Cui
Office of Groundwater Impact Assessment
Queensland Government
taocuisunny@gmail.com

David Robertson
CSIRO Land and Water
david.robertson@csiro.au

Sreekanth Janardhanan
CSIRO Land and Water
sreekanth.janardhanan@csiro.au

Abstract
We explore black-box model inversion using the Fourier Neural Operator (FNO)of Li et al [4]. The approach learns an emulator of a partial differential equation
forward operator from simulated realisations and then infers unobserved systemparameters by minimising emulator predictive loss with respect to the observationsof the system outputs. Our results suggest that this underdetermined inverse
problem is significantly harder than the forward problems or initial condition
inference of [4], but by careful regularisation we are able to improve our inference
substantially.

1

Introduction

Many environmental problems (e.g. predicting and forecasting the dynamics of ocean currents,
weather events, and dynamics of water and solute movement described by surface and groundwaterhydrological models) are challenging to model due to complex non-linearities and dynamical processes that characterise the system. Models of these systems can be described by partial differential
equations (PDEs). In many industrial application, standard solvers are black-box systems which
do not provide gradient information. For such systems, inference of high-dimensional unobserved
parameters is challenging because gradients must be estimated by finite difference, which rapidly
becomes prohibitively expensive.
The challenges have sparked interest in neural network approximation to the PDE solution operators [5,6, 3]. Classical solution methods for a PDE might involve approximating some continuous domain
process into a discrete approximation with complex dependency structure, e.g. stylized in Figure 1a.
Naive inference over the graphical model [1] in such systems is forbiddingly complex. A PDE
emulator can conceptually collapse the complex dependency graph into a smaller more tractable one,e.g. Figure 1b at the cost of requiring us to accept entire functions as inputs and output of the network,
i.e. to learn operators on functions. Recent work by Li et al. [4] shows the potential of the FNO for
precisely this kind of purpose. The FNO is lauded for a number of features; notably, it is rapid totrain, achieves high accuracy with little tuning, and may be interpreted as a resolution-independent
mapping between functions on continuous domains without presuming a fixed quantisation of the
Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

(1)

𝜃

a)

(2)

𝜃

(3)

𝜃

(4)

(5)

𝜃

𝜃

(6)

𝜃

vt(1)

vt(2)

vt(3)

vt(4)

vt(5)

vt(6)

vt+ (1)vt+ (2)vt+ (3)vt+ (4)vt+ (5)vt+ (6)

𝜀

𝜀

𝜀

𝜀

𝜃

𝜃

𝜀

𝜀

– parameters

vt – solution values

vt

b)

vt+

𝜀

Figure 1: The conditional independence graph of a PDE model with with one spatially-varying
parameter, represented as a) as the solution of a discretized PDE solver, evaluated at 6 spatial
locations, and b) in terms of continuous operators.

domain. Furthermore, the FNO has recently been shown to satisfy a universal approximation property
for such operators [2]. Once learned, the emulator can be used to predict a solution surface withsignificant speed up over direct PDE solvers, which is differentiable with respect to all the input. Li
et al. [4] emphasise the potential application of such operators in downstream tasks such as Bayesianmodel inversion (i.e., estimating physical system parameters from observations). A further usefulproperty of the FNO to this end is that differentiation of the operator is trivial in all inputs using
off-the-shelf backpropagation.
We also follow this thread of inquiry, learning a FNO for the forward operator of the PDE and thencomputing derivatives with respect to the latent parameters within an optimisation routine to performinversion, i.e. estimation of the value of an unobserved parameter.

2

Methods

2.1

The Physical System

Consider a PDE, with parameters θ, that describes the evolution of some vector-valued field,v(s, t|θ, φ) ∈ V ⊆ Rdv over two-dimensional space (s = (x, y)> , s ∈ S ⊆ R2 ) and time(t ∈ T ⊆ R). Such a PDE might be used to model an environmental process (e.g. groundwaterpressure head across an aquifer (dv = 1)). Typically, the PDE will exhibit dynamics that are governedby: (i) function spaces of parameters, for example a spatial field θ(s), θ : S → G where G ⊆ Rdp
encodes physical properties of the spatial environment; and (ii) boundary conditions and/or initialconditions φ that define constraints on the state of the system at particular locations and/or times. For
our purposes, we may think of the boundary conditions as implicit constraints on the solutions which
are encoded within the PDE simulations (i.e. the data) used in this study. Both θ and φ may varyover time and/or space.
We handle the observations of the PDE in discrete time. For each time t, we assume the instantaneous“slices” of solutions to the PDE are functions vt belonging to the Banach space F ∗ , with vt : D∗ →R∗ where D∗ = S × Θ is also a compact set of positive Lebesgue measure and R∗ = V .
A fundamental property of a PDE is the forward operator M : F ∗ × Θ → F ∗ which produces theentire solution surface at some future time t + , given the current state and boundary conditions andparameters: (see also Figure 1 b.)
vt+ |θ, φ, vt = M [vt |θ, φ](1)

The forward operator of a PDE is dependent not only on the current state, but also upon temporal
derivatives of the state field. We augment vt (s) with additional components in Rdv +d∂ , where d∂ isthe number of temporal derivatives sufficient to define the PDE. In practice, these derivatives could beapproximately encoded by augmenting the state vector to be (vt (s)> , vt− (s)> , . . . , vt−m (s)> )>for sufficiently small values of  and m sufficiently large. Herafter we will hold the boundaries
contraints φ fixed, and they are suppressed.
22.2

The Fourier Neural Operator as a Model Emulator

t=0

t=1
0.25
0.00
0.25

Figure 2: Samples of the vorticity field vt from Navier-stokes simulation. Viscosity ν = 10−4 .
We use a simple 2d model of flow, the well-known Navier-Stokes equations as a case study,1 andinvestigate the behaviour of optimisation-based inference for this problem. In two dimensions
the Navier-Stokes system is described by a set of equations that characterise the flow vorticity
field (Figure 2). We closely follow the methods of Li et al. [4] in fitting a discrete-time FNO toapproximate the forward solution of a 2d Navier-Stokes equation on a toroidal domain, by training it
to minimise prediction error data generated by a classic PDE solver. However, unlike the original
paper, we generate simulations subject to a time-invariant, spatially-varying random forcing parameter
θ(s) simulated from a Gaussian random field. Our goal is to infer the value of θ for new data bynumerically solving 3 by gradient descent. Using a classic PDE solver, we generate simulations fromthis system and use those to train the FNO network. For details of that network we refer the reader
to [3].
2.3

Inverse Inference of Parameter Fields by Forward Prediction Error Minimisationb to minimises discrepancy d between
A pragmatic form of inference for a θ is to choose an estimate θobservations vt+ε and predictions of those same observations from vt through M ,
b := argmin
θ(2)θ∈Θ d M [vt |θ], vt+ |θ .
As a computationally intensive, infinite dimensional problem, this is typically infeasible in a naive
implementation, but by emulation we may efficiently approach some approximation of this if wec in place of the true PDE operator. A realisable inference procedure involves
use the emulation M
additional approximations; we represent θ by some finite parametrisation, θ̄ κ : Rdκ → Θ. We use
an approximate mean-squared discrepancy d ≈ d¯ between predicted solution surfaces at a finite
set of sites {sj }. If we have chosen a flexible parametrisation for our candidate θ̄, this problem ispotentially under-specified, so we allow for a regularization term p(θ̄ κ ) ≥ 0 and a penalty weightλ ≥ 0 which allow us to bias the solution towards parameter ranges.
We have replaced the original problem with one which, for appropriately chosen θ̄ κ and p, isdifferentiable in its (finite-dimensional) arguments, allowing incremental updating of solutions usingbκ := θ̄ κ̂ via
the loss gradient. Putting this together, we defines estimate θc [vt |θ̄ κ ], vt+ |θ̄ κ + λp(κ).
b := argminκ d¯ Mκ(3)Nothing in this set up guarantees that the solution we find this way is unique, or that the optimum
matches the true value.

3

Results

For brevity, we defer description of the basic FNO to Li et al.[4] and note only where we diverge. All
simulation and model parameters are included in released code2 . All execution times are on a Tesla
P100 GPU. We differ in that we have simulated a new dataset of 2d Navier-Stokes simulation, withboth initial conditions and latent forcing parameter generated from the same Gaussian Random field,
1

The Navier-Stokes model is sufficiently common that non-black-box solvers are available; we do not exploitthat here.
2
https://github.com/csiro-mlai/fno_inversion_ml4ps2021

3

TargetEstimate

Error
0.050
0.025
0.000
0.025
0.050

Targetbκ ) ≈ 0.30 (45s).
(a) Unregularised, `(θ

Estimate

Error
0.050
0.025
0.000
0.025
0.050

bκ ) ≈ 0.14 (1m59s).
(b) With λ = 0.3, `(θbκ at convergence.
Figure 3: Estimated latents θ1
in the latter scaled by a factor of 100. There are 1000 training and 200 of each validation and testingtime series on a (256 × 256) grid, the generation of which takes ∼ 15 hours. Example realisations
are shown in Figure 2. Forward model fit terminates by early stopping after 2h 28m.

relative error

In the inversion stage we hold the weights of theforward model fixed. For the prediction error
100
minimisation we choose a simple parametrisation for θ̄ κ , specifying its value on a discrete
6 × 10 1
lattice of K ×K points {kδ, jδ} spanning D, i.e.
θ̄ κ (kδ, jδ) ≡ κjk . For this problem K = 256,
4 × 10 1
and we choose the same lattice upon which the
3 × 10 1
training PDE simulations are evaluated. Writing k · k2 for the empirical 2-norm evaluated on
2 × 10 1
the lattice, we define the target predictive divergence to be the same used in the NN training,¯ 1 , f2 ) := kf1 − f2 k2 , and we measure the
d(f
10 1
2
0
20
40
60
80
100
120
quality in the inverse problem by relative error
Gradient update stepbκ ) = kθbκ − θk2 /kθk2 scaled to the magni`(θtude of the estimand θ, where Ekθk2 = 0.0112. Figure 4: Convergence of 50 different inference
In this toy example we know the latent functions problems to different local optima.
are smooth by construction, so we impose anempirical roughness penalty in p, specifically,PK
the mean squared empirical first order finite differences, p(θ̄ κ ) = K12 j,k=1 (κj,k − κj,k−1 )2 +(κj,k − κj−1,k )2 . We draw initial conditions Ki,j ∼ N (0, 0.01).
Optimisation proceed by first-order gradient descent via a stock Adam optimiser with learning rate
0.0025. For this naïve parameterisation, there are 2562 parameters so higher order optimisation is not
tenable. However, individual optimisation steps here are cheap. In batches of 50 examples, estimand
gradient update steps take 2.23 seconds on a Tesla P100 GPU, and convergence is attained withat most 200 Adam iterations. However, the estimation procedure is challenging, requiring careful
tuning to achieve convergence. Without regularisation of the latent field, λ = 0, the results are poor
even at optima, with large and systematic errors (see Figure 3a). By contrast, when regularisation is
performed, gradient descent may find plausible solutions. Selecting optimisation parameters by gridsearch to minimise median reconstruction error over a random validation set sample (figure 5), we
obtain a improved prediction of the latent field, with estimated relative loss 0.2 ± 0.05 at λ = 0.3.
4

4Discussion

In practice, inverse modelling of the latent field using the FNO emulator presents a number of challenges. The FNO by design does not observe physical constraints, such as conservation of matter or energy. Moreover, the inversion problem is underdetermined and weneed to be aware of instability in the optimisation process. Careful regularisation of the inferred field may be needed to avoid physically implausible solutions to the model inversion.
When comparing the unregularised optimisation
results (Figure 3) to those obtained in the experiments of Li et al. [4], our results suffer from thepresence of high frequency artefacts whereas
theirs do not. Notwithstanding the estimands
are slightly different (latent parameters versus
initial conditions), this is an interesting finding
and we speculate that the inferred field in Li etal. [4] avoids these artifacts by taking spatiallypointwise means over many field realisations
which are individually noisy. One avenue forfurther enquiry is whether the average of an ensemble of estimates removes the need for regularisation of the inferred field.
Solutions are for the high regularization values are visually plausible, but may be far fromground truth. That is, an optimum of low predictive error may nonetheless differ from thegenerating parameters (Figure 4). This is compatible with the hypothesis of finding alternative solutions to an underdetermined problem. Choosing
parametrisations, regularisation and hyperparameters to improve the quality of inference in thisdomain is a topic of ongoing research.
bκ ) onFigure 5: λ vs estimated relative error `(θheld-out set (2h30). Error bars denote 90% range.

A further practical concern is in real-world inference process we may not know the generating processof the true θ as we do here and in Li et al. Further research to attain more plausible predictions
could involve eliciting domain-expert priors, or selecting a nonparametric functional with parameterschosen by a hierarchical Bayes method.

Broader Impact
The modeling of physical systems of hydrological processes, and processes in general, can directly
benefit from the ideas presented in this paper. Groundwater models and models of flood inundation for
example represent complex systems with partial knowledge and noisy, observational data. Inference
about these systems are usually employs hierarchical model to infer parameters or infer intervention
effects. However, such inference is non-trivial as some variables are defined in terms of PDE
solution steps, which have complex implementation and demanding compute requirement, and whose
solutions are provided by black-box solvers. The FNO provides one approach for conducting an
invertible approximate forward simulation of a PDE cheaply. The inversion method presented here
has the potential to solve more general inference problems, where the focus is on learning theconditional inverse mapping of outputs to latent features. This generates a wider range of inference
methods, scenario simulations and research questions within computational reach of the hydrologicalcommunity.

Acknowledgments and Disclosure of Funding
We thank Alasdair Tran for helpful feedback on the methods and code.

References[1] Daphne Koller and Nir Friedman. Probabilistic Graphical Models : Principles and Techniques.
MIT Press, Cambridge, MA, 2009.
5

[2] Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation anderror bounds for Fourier Neural Operators. arXiv:2107.07562 [cs, math], July 2021.
[3] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya,
Andrew Stuart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function
Spaces. arXiv:2108.08481 [cs, math], September 2021.
[4] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential
Equations. arXiv:2010.08895 [cs, math], October 2020.
[5] Lu Lu, Pengzhan Jin, and George Em Karniadakis. DeepONet: Learning nonlinear operatorsfor identifying differential equations based on the universal approximation theorem of operators.
arXiv:1910.03193 [cs, stat], April 2020.
[6] Maziar Raissi, P. Perdikaris, and George Em Karniadakis. Physics-informed neural networks: A
deep learning framework for solving forward and inverse problems involving nonlinear partial
differential equations. Journal of Computational Physics, 378:686–707, February 2019.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes](c) Did you discuss any potential negative societal impacts of your work? [N/A](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes](b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes](c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes](b) Did you mention the license of the assets? [Yes](c) Did you include any new assets either in the supplemental material or as a URL? [Yes](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]6

A

Appendix: Architecture of the FNO

The architecture of a FNO has three basic steps: (i) projection of the input data into a higherdimensional space through a shallow, feed-forward neural network; (ii) a series of K layers, each ofwhich simultaneously performs both a non-local integral operation and local linear operation thensums these two tensors and passes the result through a non-linear activation function, σ; and (iii) a
projection to the outputs using a final feed-forward neural network. The structure of the model isoutlined in Figure 6.
h1
Boundary(φ)
Conditions
Parameters (θ)

State
Field (v)

{

d2
d1

{

{

φ

.φ B

φ.

h 1,v

2

hK

h2
h 2,v

1

h k,v

2

k

y1

1

..

yd

.v

..θA

θ
θ 2

INPUTS

(x)

h 1,2

P

Legendmatrix

W

x

h k,1

R

+

W

x
FFT

FFT

x

IFFT

Q

h k,2

h 2,2
h 2,1

h 1,1

W

h k,3

h 2,3

h 1,3

1

d
.. v

R

+

x

OUTPUTS
FFT

x

IFFT

R

(y)
x

+

operation
FFT

Fast Fourier Transform

IFFT

Inverse Fast Fourier Transform

Figure 6: The FNO approximation to operator M .

7

IFFT