remarkRemark hypothesisHypothesis claimClaim Efficient OUU using Derivative-
Informed Neural OperatorsD. Luo, T. O’Leary-Roseberry, P. Chen, O. Ghattas

#  Efficient PDE-constrained optimization under high-dimensional uncertainty
using Derivative-Informed Neural Operators ††thanks: Submitted to the editors
DATE. This research partially supported by DOE grants DE-SC0019303 and DE-
SC0023171; NSF DMS grant 2012453; and DOD MURI FA9550-21-1-0084.

Dingcheng Luo Oden Institute for Computational Engineering and Sciences, The
University of Texas at Austin, Austin, TX, USA<EMAIL_ADDRESS>Thomas
O’Leary-Roseberry22footnotemark: 2 Peng Chen  School of Computational Science
and Engineering, Georgia Institute of Technology, Atlanta, GA, USA  Omar
Ghattas22footnotemark: 2 Walker Department of Mechanical Engineering, The
University of Texas at Austin, Austin, TX, USA

###### Abstract

We propose a novel machine learning framework for solving optimization
problems governed by large-scale partial differential equations (PDEs) with
high-dimensional random parameters. Such optimization under uncertainty (OUU)
problems may be computational prohibitive using classical methods,
particularly when a large number of samples is needed to evaluate risk
measures at every iteration of an optimization algorithm, where each sample
requires the solution of an expensive-to-solve PDE. To address this challenge,
we propose a new neural operator approximation of the PDE solution operator
that has the combined merits of (1) accurate approximation of not only the map
from the joint inputs of random parameters and optimization variables to the
PDE state, but also its derivative with respect to the optimization variables,
(2) efficient construction of the neural network using reduced basis
architectures that are scalable to high-dimensional OUU problems, and (3)
requiring only a limited number of training data to achieve high accuracy for
both the PDE solution and the OUU solution. We refer to such neural operators
as multi-input reduced basis derivative informed neural operators (MR-DINOs).
We demonstrate the accuracy and efficiency our approach through several
numerical experiments, i.e. the risk-averse control of a semilinear elliptic
PDE and the steady state Navier–Stokes equations in two and three spatial
dimensions, each involving random field inputs. Across the examples, MR-DINOs
offer $10^{3}$—$10^{7}\times$ reductions in execution time, and are able to
produce OUU solutions of comparable accuracies to those from standard PDE
based solutions while being over $10\times$ more cost-efficient after
factoring in the cost of construction.

###### keywords:

PDE-constrained optimization, optimization under uncertainty, neural operator,
operator learning, scientific machine learning, adjoint methods, reduced
basis, dimension reduction

49M41, 65C20, 65D15, 68T07, 75D55, 90C15, 90C90, 93E20

##  1 Introduction

PDE-constrained optimization problems arise in many computational science and
engineering fields. Canonical examples of such problems include optimal
design, where the goal is to find the best system configuration given
constraints, and optimal control, where the aim is to determine the optimal
operation of a system while adhering to constraints. In real-world
applications, uncertainties are inevitable and arise from many sources in PDE
models, e.g., PDE coefficients that parametrize the system properties, initial
and boundary conditions, source terms, and computational geometries. In order
to achieve robustness of optimal solutions, it is crucial to account for these
uncertainties in solving PDE-constrained optimization problems. In such PDE-
constrained optimization under uncertainty (OUU) problems, the uncertainty is
modeled by a probability distribution, and the optimization objective is
formulated using risk measures of a performance function, which can often be
written as the integral of a scalar quantity over the probability
distribution.

Solution of PDE-constrained OUU problems is challenging for the following
reasons. (1) Each optimization iteration requires solving numerous PDEs to
estimate the risk measure objective, e.g. by sample average approximation
(SAA). Accurately estimating optimization risk measures may require a large
number of samples, especially for risk measures that focus on tail
probabilities and rare events. (2) To scale up to high-dimensional
optimization variables, methods that compute or approximate the Hessian are
required, leading to the need for additional linearized (adjoint) PDEs to be
solved at each sample. (3) The space of the random parameters may be high-
dimensional, or even infinite-dimensional, which precludes the use of
deterministic quadrature methods, which suffer from the curse of
dimensionality. As a result of these computational challenges, solving PDE-
constrained OUU problems for complex systems, such as large-scaled,
multiphysics, or multiscale systems, is often computationally prohibitive
using traditional PDE solver-based methods. Recent developments to partially
address these challenges include methods such as multigrid [8], multifidelity
[51], multilevel [3, 14], stochastic Galerkin [35], stochastic collocation
[32, 62], Taylor approximation [2, 11, 12, 20], model reduction [4, 13, 36,
67], and neural networks [23, 25]. Despite the progress, the computational
challenges remain formidable, particularly for OUU problems constrained by
large-scale nonlinear PDEs under high-dimensional uncertainties.

In this work, we investigate the feasibility of overcoming the above
computational challenges via a new development of neural operators. Neural
operators have gained significant interest in recent years because of their
remarkable ability to efficiently approximate high-dimensional mappings such
as those arising in parametric PDE problems. We propose a novel framework that
utilizes neural operators in the solution of PDE-constrained OUU problems. In
this framework, the neural operator learns a mapping over a product measure of
the random parameter distribution and an auxiliary distribution for the
optimization variables that covers the feasibility sets of the optimization
problems for which the neural operator is constructed. This presents a
challenging task, since the mapping is over a product measure that is formally
infinite-dimensional in at least one of the inputs. Mesh-dependent neural
network strategies often suffer from deteriorating performance as the
dimension of the problem increases. To avoid this difficulty, we use reduced
bases to encode or compress the high-dimensional uncertain parameter field and
the PDE state, which has gained traction as a popular architectural strategy
for neural operators [6, 24, 27, 44, 54, 56]. Moreover, a key contribution is
to train the neural operator not only on the input-output solution map but
also on its derivative with respect to the optimization variables. Building on
recent work, the derivative training data can be efficiently computed and
imposed in the training process by using reduced basis neural operators [53].
This derivative-informed neural operator (DINO) strategy allows us to obtain
significantly improved approximations not only of the joint parametric map but
also of derivative quantities such as gradients, which are essential for
efficient optimization methods in high dimensions.

We demonstrate the efficiency and accuracy of our proposed method on three
challenging OUU problems subject to random field uncertainties; an optimal
source control of a semilinear elliptic PDE, and an optimal boundary control
of flow around a bluff body governed by Navier–Stokes equations in both two
and three space dimensions. We consider a challenging risk measure, the
conditional value at risk (CVaR), which typically requires a large number of
samples for accurate optimal solutions. We show that our neural operators
offer reductions in execution time by factors of $10^{3}-10^{7}$ depending on
the specific PDEs. Moreover, for the same quality of the OUU solution, the
neural operators are over $10\times$ more cost-efficient than a traditional
PDE-based optimization method over a single optimization run, even after
factoring in the construction cost for the neural operator. Derivative
training proves to be critical for accurate approximation of the solution
operator and more importantly its derivatives, yielding OUU solutions of much
higher quality than their counterparts without the derivative training. Once
trained, the neural operators can be reused to solve a family of OUU problems,
e.g. with alternative optimization objectives and risk measures, at virtually
no additional cost.

###  1.1 Related work

In recent years there has been significant work on developing neural operators
for approximating high-dimensional, complex parametric maps arising in PDE
problems [24, 29, 34, 38, 39, 50, 54, 56, 57, 66]. Additionally there has been
interest in deploying neural operators to solve “outer-loop” problems such as
Bayesian inverse problems [9, 39], Bayesian optimal experimental design [65],
optimal design [22, 54]. In particular, neural operators have been considered
as surrogates for a variety of deterministic PDE-constrained optimization
problems in [28, 30, 45, 61, 68], in addition to two recent works [23, 25]
considering the use of neural networks for PDE-constrained OUU problems.

Specifically, in [23] the authors consider the use of neural networks for
topology optimization under uncertain material parameters and loading
conditions governed by linear elasticity using a phase-field based
formulation. The neural networks output gradients with respect to the
optimization variable given the states and gradients at previous iterations.
This replaces the gradient computation step, which involves solving the
adjoint PDE and an additional gradient PDE that arises as part of the phase-
field formulation, thereby accelerating the optimization process. We note that
the approach of [23] differs from our proposed approach in that the state PDE
is still solved during the optimization, as the states are used as inputs to
the neural network.

Neural networks have also been considered for PDE-constrained optimization
problems using an “all-at-once” or “one-shot” approach. Here, neural networks
are used to represent the state and optimization variables, and potentially
additional model parameters. The training of the neural network attempts to
optimize the weights of the neural network using a loss function defined in
terms of the optimization objective and the PDE residual, thereby trying to
simultaneously achieve optimality (objective minimization) and feasibility
(satisfying the PDE). Examples of this approach include [26, 46, 64] for
deterministic optimization problems and [25] in the OUU setting. This approach
aims to solve single instances of the optimization problem, and the trained
neural networks are not intended to be reused. Instead, a new neural network
needs to be trained for different choices of optimization performance
objective, control cost, and constraints. Moreover, the training of the neural
network may still be computationally expensive due to the ill-conditioning of
the optimization problem involving both the optimization objective and the PDE
residual.

The remainder of this paper is organized as follows. In Section 2, we
introduce the formulation for PDE-constrained OUU problems. This is followed
by a presentation of the proposed neural operator architecture in Section 3.
Numerical results are then shown in Section 4 before concluding with some
remarks in Section 5.

##  2 PDE-constrained optimization under uncertainty

We consider systems that consist of a state $u\in\mathcal{U}$, random
parameters $m\in\mathcal{M}$, and optimization/control variables
$z\in\mathcal{Z}$, where $\mathcal{U},\mathcal{M},\mathcal{Z}$ are the
respective Hilbert spaces to which they belong. The system is governed by a
PDE written abstractly as

(1) $R(u,m,z)=0,$

where
$R:\mathcal{U}\times\mathcal{M}\times\mathcal{Z}\rightarrow\mathcal{U}^{\prime}$
is a differential operator and $\mathcal{U}^{\prime}$ is the dual of
$\mathcal{U}$. We assume that the PDE (1) admits a unique solution
$u=u(m,z):\mathcal{M}\times\mathcal{Z}\rightarrow\mathcal{U}$ for each given
$m$ and $z$, i.e. the PDE problem is well-posed in the spaces of $m$ and $z$.

The uncertainty in $m$ is described by its distribution $\nu_{m}$, which is a
measure over the Borel sigma algebra $\mathcal{B}(\mathcal{M})$. In this work,
we consider $m$ to be a random field, in particular a Gaussian random field,
$\nu_{m}=\mathcal{N}(\bar{m},\mathcal{C})$, with a mean
$\bar{m}\in\mathcal{M}$ and a covariance operator of Matérn class
$\mathcal{C}=\mathcal{A}^{-\alpha}$ [41], where $\mathcal{A}$ is an elliptic
differential operator, e.g., $\mathcal{A}=-\gamma\Delta+\delta I$ with
Laplacian $\Delta$, identity $I$, and homogeneous Neumann boundary condition.
The parameters $\alpha,\gamma,\delta>0$ control the smoothness, variance, and
correlation of the random field. We consider a performance function
$Q:\mathcal{U}\rightarrow\mathbb{R}$ as a function of the state variable
$u(m,z)$ that measures the performance of the system at a given $m$ and $z$.
Due to the stochasticity of $m$, the quantity $Q(u(m,z))$ is a random
variable. Thus, the PDE-constrained optimization problem is typically
formulated in terms of a risk or statistical measure $\rho(Q)$ of $Q$. We can
therefore write the PDE-constrained OUU problem as

(2) $\min_{z\in\mathcal{Z}_{ad}}\mathcal{J}(z):=\rho(Q)(z)+\mathcal{P}(z),$

where $\mathcal{Z}_{ad}\subset\mathcal{Z}$ is an admissible set of the
optimization variable $z$ and $\mathcal{P}(z)$ is a penalization or
regularization term that controls the cost or regularity of $z$.

###  2.1 Risk measures

In the OUU problem, the risk measure $\rho$ quantifies the uncertainty in the
performance function $Q$ due to the random parameter $m$. This effectively
specifies a statistical goal we have in optimizing the PDE system over the
distribution $\nu_{m}$. For example, one may be interested in optimizing the
average performance of the system. This can be achieved using the expectation

(3) $\rho_{\text{Mean}}(Q)(z)=\mathbb{E}_{\nu_{m}}[Q(u(\cdot,z))]$

as a risk-neutral measure. In engineering applications, large values of $Q$
often correspond to undesirable or even failure states of the system. Although
occurrence of such events may be rare, they can have catastrophic
consequences. In such cases, it is insufficient to consider the expectation
alone. Instead one can use risk-averse measures that account for the risk of
large deviations from the mean.

In this work, we will consider the superquantile, or conditional value-at-risk
(CVaR) [59]. Originally developed for financial risk management, the CVaR has
become of a risk measure of interest in engineering applications and PDE-
constrained OUU [10, 31, 33, 37]. For a value $\beta\in[0,1]$, the
$\beta$-quantile of $Q$, or the $\beta$-value at risk, is defined as

(4) $\mathrm{VaR}_{\beta}[Q](z):=F^{-1}_{Q(u(\cdot,z))}(\beta),$

where $F_{Q((\cdot,z))}$ is the cumulative distribution function of $Q$. The
CVaR is then defined as the conditional expectation of $Q$ given that it
exceeds the $\beta$-quantile, i.e.,

(5)
$\rho_{\text{CVaR},\beta}(Q)(z):=\mathrm{CVaR}_{\beta}[Q]=\frac{1}{1-\beta}\mathbb{E}_{\nu_{m}}[Q(u(\cdot,z))1_{Q(u(\cdot,z))>\mathrm{VaR}_{\beta}}[Q](z)].$

The value of $\beta$ specifies the level of risk-aversion. For $\beta=0$, the
superquantile is simply the expectation with the weakest risk aversion (risk
neutral), while for $\beta=1$, it is the essential supremum with the strongest
risk aversion (worst case scenario).

The OUU problem (2) with the CVaR risk measure $\rho_{\text{CVaR},\beta}(Q)$
in (5) can be equivalently formulated as [59]

(6) $\min_{\begin{subarray}{c}z\in\mathcal{Z}_{ad}\\\
t\in\mathbb{R}\end{subarray}}\mathcal{J}_{\text{CVaR},\beta}(z,t):=t+\frac{1}{1-\beta}\mathbb{E}_{\nu_{m}}[(Q(u(\cdot,z))-t)^{+}]+\mathcal{P}(z)$

by introducing an additional optimization variable $t\in\mathbb{R}$, where
$(\cdot)^{+}=\max(\cdot,0)$. The cost functional (6) is non-differentiable
because of the maximum function and smooth approximations of the maximum
function are often used instead, e.g., [33]

(7) $(x)^{+}_{\epsilon}=\begin{cases}0&\text{if }x<0,\\\
\left({x^{3}}/{\epsilon^{2}}-{x^{4}}/{2\epsilon^{3}}\right)&\text{if
}0<x<\epsilon,\\\ x-\epsilon/2&\text{if }x\geq\epsilon,\\\ \end{cases}$

with $\epsilon\ll 1$. This approximation has continuous second order
derivatives, making the optimization amenable to gradient-based optimizers.
For simplicity, we will assume that $\epsilon$ is a fixed value (e.g.
$\epsilon=10^{-4})$ such that the CVaR approximation is sufficiently accurate.
In practice, one may need to solve the OUU problem with successively
decreasing values of $\epsilon$ to obtain a more accurate solution.

In this work, we focus on the CVaR risk measure, but note that our framework
applies to a wide class of risk measures that can be formulated in terms of
expectations of functions of $Q$. This includes moments of $Q$, probability of
failure, buffered probability of failure [58], and so we refer to [16, 60] for
a more extensive exposition of risk measures for optimization under
uncertainty.

###  2.2 Sample average approximation

We compute the risk measures via sample average approximation (SAA), i.e.
approximating the risk measure $\rho$ by a Monte Carlo estimator
$\widehat{\rho}$. For concreteness, in the case of the mean, the estimator is
simply

(8) $\widehat{\rho}_{\text{Mean}}(z):=\frac{1}{N}\sum_{i=1}^{N}Q(u(m_{i},z)),$

where $m_{i}\sim\nu_{m}$ are i.i.d. samples. The SAA optimization problem then
becomes

(9)
$\min_{z\in\mathcal{Z}_{ad}}\widehat{\mathcal{J}}(z):=\widehat{\rho}_{\text{Mean}}(z)+\mathcal{P}(z).$

Similarly, for the CVaR risk measure, we use the sample average approximation
for $\mathcal{J}_{\text{CVaR},\beta}$ along with the smoothed maximum function
to obtain

(10) $\min_{\begin{subarray}{c}z\in\mathcal{Z}_{ad}\\\
t\in\mathbb{R}\end{subarray}}\widehat{\mathcal{J}}_{\text{CVaR},\beta}(z):=t+\frac{1}{N}\sum_{i=1}^{N}\frac{1}{1-\beta}(Q(u(m_{i},z))-t)_{\epsilon}^{+}+\mathcal{P}(z).$

Note that when the samples $\\{m_{i}\\}_{i=1}^{N}$ are fixed during the
solution of the optimization problem, (9) and (10) become deterministic
optimization problems, and can be solved using conventional algorithms for
deterministic PDE-constrained optimization.

###  2.3 Gradient-based optimization for OUU

It is well known that for high-dimensional optimization problems, derivative-
free methods suffer from very slow convergence, requiring numerous function
evaluations to obtain a solution. On the other hand, gradient-based methods
such as Newton and quasi-Newton methods can typically attain asymptotic
superlinear convergence rates that are independent of the dimension of the
discretized optimization variable. Since each optimization iteration involves
solving the state PDE for each of the $N$ samples $\\{m_{i}\\}_{i=1}^{N}$, it
is imperative that the number of optimization iterations is kept small,
necessitating the use of derivative-based optimization methods when the $z$ is
high-dimensional.

Under the assumption that the risk measures can be formulated as expectations
of functions of $Q$, i.e., $\rho(Q)=\mathbb{E}_{\nu_{m}}[f(Q)]$, and that they
are differentiable with respect to $z$, the gradients can be computed by the
chain rule

(11)
$D_{z}\rho(Q)=D_{z}\mathbb{E}_{\nu_{m}}[f(Q)]=\mathbb{E}_{\nu_{m}}\left[D_{Q}f\;\partial_{z}Q\right],$

where the expectation is replaced by a sum in the case of SAA. For the
computation of $\partial_{z}Q$, we can write this (interpreted formally in the
infinite dimensional setting) as

$\displaystyle g^{T}_{z}(m,z)$
$\displaystyle:=\partial_{z}Q(u(m,z))=\partial_{u}Q(u(m,z))\partial_{z}u(m,z)$
$\displaystyle=-\partial_{u}Q(u(m,z))\left[\partial_{u}R(u,m,z)\right]^{-1}\partial_{z}R(u,m,z)$
(12) $\displaystyle=p^{T}\partial_{z}R(u,m,z),$

where we have used $p$ to denote the adjoint variable, which is given by
solving the adjoint system
$p=-\left[\partial_{u}R(u,m,z)\right]^{-T}\partial_{u}Q(u(m,z))^{T}.$ Thus,
for a given $m$, $z$, and $u(m,z)$, the dominant cost of computing the
gradient $g_{z}(m,z)$ is in solving an additional linear PDE involving adjoint
operator $\partial_{u}R(u,m,z)^{T}$. The gradients can then be used in
gradient-based optimization methods to minimize the SAA of the cost
functional.

In summary, to estimate the risk measure via SAA, one needs to solve the state
PDE for each sample $m_{i}$, amounting to $N$ state PDE solves. Moreover,
computing the gradient of the risk measure with respect to the optimization
variable requires an additional $N$ adjoint PDE solves, one for each sample
$m_{i}$. Assuming a fixed sample size $N$, the overall cost of the OUU problem
then involves $N_{\text{opt}}\times N$ state and $N_{\text{opt}}\times N$
adjoint PDE solves, where $N_{\text{opt}}$ is the number of optimization
iterations, assuming a quasi-Newton method and ignoring the solves needed in
the line search.

Solving OUU problems thus becomes computationally prohibitive for large $N$
and $N_{\text{opt}}$ when each PDE solve is expensive, which motivates the
approximation of the map from the product space of the random parameter field
and optimization variables to the PDE state, $(m,z)\mapsto u(m,z)$, by
surrogates that are accurate for not only the PDE state but also its
derivative with respect to the optimization variables $z$. This leads to the
development of derivative-informed neural operators in the following section.

##  3 Derivative-informed neural operators

Neural operators are neural network approximations of operators as mappings
between input and output function spaces. Let $x\in\mathcal{X}$ denote an
input function in the function space $\mathcal{X}$ equipped with a probability
measure $\nu_{x}$, and $y\in\mathcal{Y}$ denote an output function in the
function space $\mathcal{Y}$. For a given operator mapping
$T:\mathcal{X}\rightarrow\mathcal{Y}$, the goal of neural operator learning is
to construct an approximation $T_{w}$ parametrized by neural network
parameters or weights $w$ that is optimal in a parametric Bochner space, e.g.
$L^{2}_{\nu_{x}}=L^{2}(\mathcal{X},\nu_{x};\mathcal{Y})$. That is, one seeks a
solution to the expected risk minimization problem,

(13)
$\min_{w}\|T_{w}-T\|^{2}_{L^{2}_{\nu_{x}}}:=\int_{\mathcal{X}}\|T_{w}(x)-T(x)\|_{\mathcal{Y}}^{2}\,d\nu_{x}(x).$

Due to the intractability of directly integrating with respect to the measure
$\nu_{x}$, one typically approximates the risk minimization problem (13) using
finitely many samples of input-output pairs,
$\\{(x_{i},T(x_{i}))\\}_{i=1}^{N_{s}}$, leading to empirical risk
minimization. We refer to this as the data-driven approach. Alternatively,
equivalent objective functions such as norms of the PDE residual or other
physics-based objective functions can be used as loss functions in the so-
called “physics-informed” machine learning approach [57, 66]. Hybrid
approaches combining the data-driven approach with additional physics-informed
losses have also been considered [40]. However, for simplicity and without
loss of generality, we consider only the data-driven approach.

Neural operator construction typically consists of the following challenges:
(1) For solution operators of PDEs, generating input-output pairs requires
solving the PDE, which is typically computationally expensive. Thus, one may
be able to afford only a limited number of samples, leading to sampling errors
in the empirical risk minimization problem. (2) The neural network training
problem, i.e., the empirical risk minimization problem, is non-convex, and is
solved using a nonlinear stochastic optimization method. Global optimization
is typically NP-hard and one can only settle for local minimizers. Moreover,
the neural operator training problem can be very sensitive to the optimization
procedure, initial guesses, and additional factors such as scaling of the
data. (3) Appropriate neural operator architectures are required in order to
accommodate a sufficiently rich functional representation.

In recent years many neural operators have been developed that construct
effective representations of parametric PDE maps by encoding a priori known
mathematical structure of the maps into the architecture. For maps that admit
compressible representations in linear bases of the inputs and outputs
$\mathcal{X}$ and $\mathcal{Y}$, reduced basis neural operators have been
developed [6, 24, 27, 44, 54, 56]. Other architectures have exploited compact
nonlinear representations, such as a Fourier representation [39]. These
various neural operators often come with universal approximation theorems [6,
34, 43, 54], asserting that mappings in a parametric Bochner space (e.g.
$L^{2}_{\nu_{x}}$) can be approximated arbitrarily well by finite dimensional
neural network representations. Finding suitable representations in practice
remains a challenge, and is highly problem-dependent.

###  3.1 Use of neural operators for OUU

In this work, we consider the use of neural operators to approximate the
solution map $u_{w}(m,z)\approx u(m,z)$ in the OUU problem (i.e.
$\mathcal{Y}=\mathcal{U}$). Specifically, we require the neural operator to be
sufficiently accurate over the input spaces
$\mathcal{X}=\mathcal{M}\times\mathcal{Z}$ equipped with a joint probability
measure $\nu_{x}=\nu_{m}\otimes\nu_{z}$, where we have introduced an auxiliary
distribution $\nu_{z}$ from which training data for the control variables $z$
are generated. We note that although in certain contexts, it may be sufficient
to learn directly the scalar performance function $Q(u(m,z))$, learning the
full solution map is a more general approach. Neural operators can be
constructed independent of the optimization objective/performance function,
allowing for flexibility in its deployment and amortization of construction
costs to solve families of optimization problems involving different choices
of performance functions.

The selection of $\nu_{z}$ is problem dependent, but should contain in its
support the admissible set $\mathcal{Z}_{ad}$. For example, when $z$ has bound
constraints, a natural choice for $\nu_{z}$ is the uniform distribution over
the bounds. The choice of $\nu_{z}$ can additionally encapsulate prior
information about regions of $\mathcal{Z}$ over which we need the neural
operator to be most accurate. For simplicity, we will use uniform or Gaussian
distributions supported over the admissible set, and defer a study on
strategically selecting $\nu_{z}$ to future work.

The neural operator then replaces the PDE solution in the sample average
approximation of the risk measures. In the case of the expectation, we take

(14) $\mathbb{E}[Q](z)\approx\frac{1}{N}\sum_{i=1}^{N}Q(u_{w}(m_{i},z)),$

where $m_{i}\sim\nu_{m}$. Analogous to (2.3), the gradient of the neural
operator based approximation for the performance function can be computed by
chain rule, i.e.,

(15)
$\partial_{z}Q(u_{w}(m,z))=\partial_{u}Q(u_{w}(m,z))\;\partial_{z}u_{w}(m,z),$

where the derivative $\partial_{z}u_{w}(m,z)$ can be efficiently computed by
automatic differentiation. The evaluation of the neural operator $u_{w}(m,z)$
and its derivative $\partial_{z}u_{w}(m,z)$ can typically be orders of
magnitdue faster than solving the corresponding PDEs, allowing us to use a
large sample size to solve the OUU problem with SAA and gradient-based
optimization methods. This can effectively eliminate sampling error in SAA,
albeit at the cost of introducing a bias due to the approximation error of the
neural operator. Thus, the effectiveness of using the neural operator to solve
OUU problems depends on the trade-off between this sampling error and the
approximation error, which will be investigated in detail in the numerical
experiments in Section 4.

Additionally, the gradient of the performance function $\partial_{z}Q$ with
respect to the optimization variable $z$ plays an important role in the OUU
problem. In particular, the optimal solution $z^{*}$ is characterized by the
first-order necessary condition $D_{z}\mathcal{J}(z^{*})=0$. Here, the
gradient of the cost functional $\mathcal{J}$ involves the gradient of the
performance function, $\partial_{z}Q(u(m,z))$, as illustrated in (11). Poor
approximation of the gradient by the neural operator
$\partial_{z}Q(u_{w}(m,z))$ leads to spurious local minimizers of
$\mathcal{J}(z)$, resulting in inaccurate optimal solutions. The gradient
error due to the neural operator approximation can be bounded in terms of the
approximation error of the neural operator and its derivative as follows.

######  Proposition 3.1.

Assume that the operator $u(m,z)$ is differentiable and the performance
$Q:\mathcal{U}\rightarrow\mathbb{R}$ is Lipschitz continuously differentiable
with constant $L_{Q}^{1}$. Then, at a given $(m,z)$,

(16) $\|\partial_{z}Q(u)-\partial_{z}Q(u_{w})\|_{\mathcal{Z}^{\prime}}\leq
L_{Q}^{1}\left\|\partial_{z}u\right\|_{\mathcal{L}(\mathcal{Z},\mathcal{U})}\|u-u_{w}\|_{\mathcal{U}}+\|\partial_{u}Q(u_{w})\|_{\mathcal{U}^{\prime}}\|\partial_{z}u-\partial_{z}u_{w}\|_{\mathcal{L}(\mathcal{Z},\mathcal{U})}$

where $\|\cdot\|_{\mathcal{L}(\mathcal{Z},\mathcal{U})}$ denotes the operator
norm of bounded linear operators from $\mathcal{Z}$ to $\mathcal{U}$.

######  Proof 3.2.

The bound follows from a triangle inequality,

$\displaystyle\|\partial_{z}Q(u)$
$\displaystyle-\partial_{z}Q(u_{w})\|_{\mathcal{Z}^{\prime}}=\left\|\partial_{u}Q(u)\partial_{z}u-\partial_{u}Q(u_{w})\partial_{z}u_{w}\right\|_{\mathcal{Z}^{\prime}}$
$\displaystyle\leq\left\|(\partial_{u}Q(u)-\partial_{u}Q(u_{w}))\partial_{z}u\right\|_{\mathcal{Z}^{\prime}}+\left\|\partial_{u}Q(u_{w})(\partial_{z}u-\partial_{z}u_{w})\right\|_{\mathcal{Z}^{\prime}}$
$\displaystyle\leq
L_{Q}^{1}\left\|\partial_{z}u\right\|_{\mathcal{L}(\mathcal{Z},\mathcal{U})}\|u-u_{w}\|_{\mathcal{U}}+\|\partial_{u}Q(u_{w})\|_{\mathcal{U}^{\prime}}\left\|\partial_{z}u-\partial_{z}u_{w}\right\|_{\mathcal{L}(\mathcal{Z},\mathcal{U})}$

This suggests that accuracy of both the solution operator $u(m,z)$ and its
Jacobian $\partial_{z}u(m,z)$ are needed in order to guarantee accuracy of the
gradient $\partial_{z}Q$.

In light of the aforementioned points, we focus on the following interrelated
challenges in deploying neural operators for the task of solving OUU problems
subject to high-dimensional uncertainty. First, due to the large computational
costs of forward simulations, as well as the high-dimensionality of the random
parameter field $m$ and state $u$, one is faced with the task of learning
complex high-dimensional operators from limited samples. Second, as the neural
operator is to be deployed in the solution of an optimization problem, we are
concerned not just with the operator approximation accuracy, but also the
derivatives of the operator with respect to the optimization variable $z$. It
remains to establish whether or not sufficiently accurate neural operators can
be constructed for OUU in a cost effective manner when compared to solving the
OUU problem directly with the PDE.

###  3.2 Derivative-informed neural operators

In order to accurately and efficiently solve OUU problems it is important that
neural operator errors do not lead to inaccurate approximations of the risk
measure and its gradient. To this end, we propose to train the neural network
approximation on not only evaluations of the solution operator, but also its
derivative with respect to the optimization variable $z$. Until recently,
parametric derivative training was not addressed in neural operator learning.
The work of [53] investigates the feasibility of constructing derivative-
informed neural operators (DINOs), where neural operators are trained on both
the operator value and its derivative with respect to the input variable. That
is, the operator regression task is performed in $H^{1}_{\nu_{x}}$ instead of
$L^{2}_{\nu_{x}}$ as in (13), i.e.,

(17)
$\min_{w}\|T_{w}-T\|^{2}_{H^{1}_{\nu_{x}}}:=\int_{\mathcal{X}}\|T_{w}(x)-T(x)\|_{\mathcal{Y}}^{2}+\underbrace{\|D_{x}T_{w}(x)-D_{x}T(x)\|_{\text{HS}(\mathcal{X},\mathcal{Y})}^{2}}_{\text{Jacobian
error term}}d\nu_{x}(x),$

where $\|A\|_{\mathrm{HS}(\mathcal{X},\mathcal{Y})}$ denotes the
Hilbert–Schmidt norm of a linear operator
$A\in\mathcal{L}(\mathcal{X},\mathcal{Y})$ defined using an orthonormal basis
$e_{i}$ of $\mathcal{X}$ as
$\|A\|_{\mathrm{HS}(\mathcal{X},\mathcal{Y})}^{2}:=\sum_{i}\langle
Ae_{i},Ae_{i}\rangle_{\mathcal{Y}}^{2}.$

In the OUU setting, we consider the following derivative-informed neural
operator training motivated by the error expression in (16),

(18)
$\displaystyle\min_{w}\|u_{w}-u\|^{2}_{H^{(0,1)}_{\nu_{m}\otimes\nu_{z}}}:=$
$\displaystyle\int_{\mathcal{M}\times\mathcal{Z}}\|u_{w}(m,z)-u(m,z)\|_{\mathcal{U}}^{2}$
$\displaystyle+$
$\displaystyle\underbrace{\|\partial_{z}u_{w}(m,z)-\partial_{z}u(m,z)\|_{\text{HS}(\mathcal{Z},\mathcal{U})}^{2}}_{\text{control
Jacobian error term}}d\nu_{m}\otimes\nu_{z}(m,z),$

where we refer to the derivative of the state with respect to the control
variables, $\partial_{z}u$, as the control Jacobian. This objective is
ostensibly intractable in comparison to the $L^{p}_{\nu_{x}}$ learning problem
(13), due to the computational costs of evaluating $\partial_{z}u$ for
training data, and the large online memory and arithmetic costs associated
with computing and differentiating through the Jacobian error term. However,
we show that the data generation and training costs may be significantly
reduced using reduced basis architectures, which we introduce in Section 3.3.

Besides improving accuracy of the Jacobian, it is numerically shown in [53]
that the Jacobian training also improves the generalization accuracy of the
neural operator output itself. Since the optimal solution $z^{*}$ is unlikely
to coincide with any of the training samples, generalization accuracy of the
neural operator within the admissible space $\mathcal{Z}_{\mathrm{ad}}$ is
important to obtaining accurate OUU solutions. In this regard, the Jacobian
$\partial_{z}u(m,z)$ contains additional information about the dependence of
$u$ on $z$, which can be particularly valuable in preventing overfitting of
the neural network when the training sample size is small. Often, this
information can be obtained at little additional cost, as discussed in Section
3.4.

###  3.3 Reduced basis architectures

In order to address the high dimensionality of the input and output spaces, we
propose the use of reduced basis neural operator architectures that exploit
the intrinsic low dimensionality of the map $m,z\mapsto u(m,z)$. In this
framework, neural networks are used to approximate the mapping between reduced
basis representations of the input and output spaces. Since the construction
of the neural network depends only on the intrinsic dimensionality of the
solution mapping, reduced basis architectures have the capability to learn
complex high-dimensional PDE-based maps in an efficient and dimension
independent manner when intrinsic dimensionality is low.

For the OUU problem, we assume that the spaces $\mathcal{M}$ and $\mathcal{U}$
have dimensions $d_{M}$ and $d_{U}$, arising from the discretization of
infinite dimensional function spaces. In particular, $d_{M}$ and $d_{U}$ may
be arbitrarily large as the mesh resolution increases. On the other hand, we
consider the space $\mathcal{Z}$ to be inherently finite dimensional, with
dimension $d_{Z}\ll d_{U},d_{M}$. This is typical of many optimization
problems in engineering, since design/control choices are often finite-
dimensional by nature. Therefore, we consider reduced basis representations of
$\mathcal{M}$ and $\mathcal{U}$, but do not employ dimension reduction on
$\mathcal{Z}$. Despite this, our method extends to the case where
$\mathcal{Z}$ is infinite dimensional by applying similar dimension reduction
techniques to $\mathcal{Z}$.

We can write the proposed reduced basis neural operator as

(19) $u_{w}(m,z)=\Phi_{r_{U}}\varphi_{w}(m_{r},z)+b,\qquad
m_{r}=\Psi_{r_{M}}^{T}m,$

where $\Psi_{r_{M}}\in\mathbb{R}^{d_{M}\times r_{M}}$ and
$\Phi_{r_{U}}\in\mathbb{R}^{d_{U}\times r_{U}}$ are reduced bases for the
parameter and state spaces $\mathcal{M}$ and $\mathcal{U}$, respectively. The
mapping $\varphi_{w}:\mathbb{R}^{r_{M}\times
d_{Z}}\rightarrow\mathbb{R}^{r_{U}}$ is a reduced basis neural network
parametrized by weights $w\in\mathbb{R}^{d_{W}}$, and $b\in\mathbb{R}^{r_{U}}$
is a bias term. We will refer to this architecture as the multi-input reduced
basis neural operator (MR-NO) and use the name MR-DINO when the architecture
is trained with the derivative-informed loss (18). A schematic for MR-DINO is
shown in Figure 1.

Uncertain parameter reduced basis Eigenvectors of
$\mathcal{C}$$\mathbb{R}^{d_{M}}\rightarrow\mathbb{R}^{r_{M}}$Optimizationvariable$z\in\mathbb{R}^{d_{Z}}$Neural
Network:$\mathbb{R}^{r_{M}}\times\mathbb{R}^{d_{Z}}\rightarrow\mathbb{R}^{r_{U}}$State
reduced
basis$\mathbb{E}_{\nu_{m}\otimes\nu_{z}}[uu^{T}]$$\mathbb{R}^{r_{U}}\rightarrow\mathbb{R}^{d_{U}}$Derivative-
Informed Neural
Operator$\min_{w}\mathbb{E}_{\nu_{m}\otimes\nu_{z}}[\|u_{w}-u\|_{2}^{2}+\|\partial_{z}u_{w}-\partial_{z}u\|_{\mathrm{HS}}^{2}]$
Figure 1: Schematic for the MR-DINO in the solution of OUU problems.

We note that for certain classes of performance functions, such as linear or
quadratic forms, its evaluation at the neural operator output, i.e.
$Q(u_{w})$, can be further accelerated using the reduced basis representation
of $u_{w}$. For example, in the quadratic case with $Q(u)=u^{T}Wu$,
$W\in\mathbb{R}^{d_{U}\times d_{U}}$, and $b=0$ for simplicity, we have

(20)
$Q(u_{w})=u_{w}^{T}Wu_{w}=\varphi_{w}^{T}\Phi_{r_{U}}^{T}W\Phi_{r_{U}}\varphi_{w}=\varphi_{w}^{T}W_{r_{U}}\varphi_{w},$

where the reduced matrix
$W_{r_{U}}:=\Phi_{r_{U}}^{T}W\Phi_{r_{U}}\in\mathbb{R}^{r_{U}\times r_{U}}$
can be precomputed so that the costs of evaluating $Q(u_{w})$ scale only with
the rank $r_{U}$. This applies to many commonly used optimization objectives,
such as $L^{2}(\Omega)$ norms and data misfits.

In this work, we obtain reduced bases by proper orthogonal decomposition (POD)
for the state, and principal component analysis (PCA) for the uncertain
parameter. This architecture is an extension of the PCANet [6] to multiple
inputs. The POD basis is computed by solving the eigenvalue problem

(21)
$\mathbb{E}_{\nu_{x}}[(u-\bar{u})(u-\bar{u})^{T}]\phi_{i}=\lambda_{u}^{(i)}\phi_{i},$

where $\bar{u}=\mathbb{E}_{\nu_{x}}[u]$. The reduced basis is then taken as
the $r_{U}$ eigenvectors $\Phi_{r_{U}}=[\phi_{1},\dots\phi_{r_{U}}]$
corresponding to the $r_{U}$ largest eigenvalues, and the bias is taken as the
mean $b=\bar{u}$. In practice, the eigenvectors are computed by SVD of data
matrix $U=[u_{1},\dots u_{N}]$, where $\\{u_{i}\\}_{i=1}^{N}$ are snapshots of
the solution from training data. Moreover, when $u\in\mathbb{R}^{d_{U}}$
represents the coefficients of a finite element discretization, the inner
product of $\mathcal{U}$ becomes a weighted inner product $\langle
u_{1},u_{2}\rangle_{M_{u}}:=u_{1}^{T}M_{u}u_{2}$, where the symmetric positive
definite matrix $M_{u}$ arises from the discretization of the underlying
function space inner product. For $L^{2}(\Omega)$, $M_{u}$ is simply the mass
matrix. In this case, to maintain consistency with the infinite dimensional
setting, we consider a weighted POD in which the SVD is carried out on
$M_{u}^{1/2}U$, such that the resulting basis is orthonormal in the $\langle
u_{1},u_{2}\rangle_{M_{u}}$ inner product.

Reduced basis networks based on POD have become very popular in neural
operator learning [6, 24, 44, 54, 56]. Approximation errors based on the POD
truncation can be proven via the Hilbert–Schmidt Theorem or Fan’s Theorem [6,
48, 54], with an upper bound by the sum of the trailing eigenvalues
$\\{\lambda_{u}^{(i)}\\}_{i>r_{M}}$. More general a-priori reduced basis error
analyses can be found in [7] based on their comparison to Kolmogorov $n$-width
for optimal linear approximations and for specific parametric problems with
explicit exponential rates in [13, 47] and algebraic rates in [17, 19], see
review in [15, 18, 21, 52].

The basis for the random parameter field is analogously computed using the
PCA. This uses the dominant $r_{M}$ eigenvectors of the covariance of
$\nu_{m}$ as the reduced basis,
$\mathcal{C}\psi_{i}=\lambda_{m}^{(i)}\psi_{i},$ giving rise to the rank
$r_{M}$ basis $\Psi_{r_{M}}=[\psi_{1},\dots,\psi_{r_{M}}]$. In the case of a
Gaussian random field, this simply corresponds to the Karhunen–Loève expansion
(KLE) of the random field. The KLE basis works well in cases where the most
sensitive modes of the solution operator align well with the dominant modes of
the covariance operator. When this is not the case, low-dimensional bases that
directly capture sensitivity of the solution operator $u$ with respect to the
random parameter field $m$ can be identified using an active subspace approach
[56]. For simplicity, we will use only the KLE basis and note that our
proposed approach works analogously with the active subspace basis, which may
be an appealing option for problems where this $m$-sensitivity is informative.

###  3.4 Efficient Jacobian training for reduced basis architectures

To efficiently perform Jacobian training, we use the fact that the range of
the control Jacobian of the reduced basis neural operators,
$\partial_{z}u_{w}$, is precisely the span of the output reduced basis.
Therefore, for the reduced basis neural operator presented in (19), the
minimization of Jacobian error can be re-written in the reduced output space
as

(22) $\displaystyle\min_{w}$
$\displaystyle\|u_{w}(m,z)-u(m,z)\|_{\mathcal{U}}^{2}+\|\partial_{z}u_{w}(m,z)-\partial_{z}u(m,z)\|_{\text{HS}(\mathcal{Z},\mathcal{U})}^{2}$
$\displaystyle\Leftrightarrow$
$\displaystyle\min_{w}\|\varphi_{w}(m_{r},z)+\Phi_{r_{U}}^{T}M_{u}(b-u(m,z))\|_{\ell^{2}}^{2}+\|\partial_{z}\varphi_{w}(m,z)-\Phi_{r_{U}}^{T}M_{u}\partial_{z}u(m,u)\|_{F}^{2},$

due to Theorem 1 in [53], where $\|\cdot\|_{F}^{2}$ is the Frobenius norm. We
refer to the term
$\Phi_{r_{U}}^{T}M_{u}\partial_{z}u\in\mathbb{R}^{r_{U}\times d_{Z}}$ as the
reduced control Jacobian of $u$. Thus, the loss term in (22) allows the online
memory and arithmetic costs for training our reduced basis neural operators to
scale only with $d_{Z}\times r_{U}$ instead of $d_{Z}\times d_{U}$. Morever,
the reduced control Jacobian $\Phi_{r_{U}}^{T}M_{u}\partial_{z}u$ can be
efficiently computed after computing the PDE solution $u(m,z)$ at relatively
little additional cost. Recall that by the implicit function theorem, we have

(23)
$\Phi_{r_{U}}^{T}M_{u}\partial_{z}u(m,z)=-\Phi_{r_{U}}^{T}M_{u}\left[\partial_{u}R(u,m,z)\right]^{-1}\partial_{z}R(u,m,z),$

where $\partial_{z}R$ is of size $d_{U}\times d_{Z}$. When $d_{Z}<r_{U}$, we
solve the linearized state PDE,
$\left[\partial_{u}R(u,m,z)\right]^{-1}\partial_{z}R(u,m,z)$ for each of the
$d_{Z}$ columns of $\partial_{z}R$, On the other hand, if $d_{Z}\geq r_{U}$,
we instead solve the adjoint PDE,
$[\partial_{u}R(u,m,z)]^{-T}M_{u}\Phi_{r_{U}}$ for each of the $r_{U}$ basis
vectors $\Phi_{r_{U}}$. In any case, given the state $u(m,z)$, the additional
computational costs of computing the Jacobian training data are associated
with the $\min(d_{Z},r_{U})$ linearized PDE solves, all with the same
linearized operator, $\partial_{u}R(u,m,z)$, or its adjoint.

For steady-state PDEs, which is the focus of this work, the costs to compute
the Jacobian training data can be mitigated as follows. When the state PDE is
linear, the linearized PDE has the same linear operator as the state PDE.
Hence, when one uses a direct solver, the linearized PDE solves for the
Jacobian computation can reuse the same triangular factors of the state PDE
solution, which require only the back substitution step of the solve and
become negligible in cost relative to the state PDE solve. For linear problems
requiring iterative methods, the costs of the preconditioner construction can
be amortized across the linearized PDE solves, since they use the same
preconditioner as the linear state PDE solve. When the state PDE is nonlinear,
multiple linearized state PDEs need to be solved in order to arrive at a
solution, e.g. via Picard or Newton iterations. The linearized PDE solves
involved in the Jacobian computation therefore cost only a fraction of the
state PDE solve. We will demonstrate the cost of the Jacobian computation
relative to that of the state PDE solve in the numerical experiments in
Section 4.

##  4 Numerical experiments

In this section, we demonstrate the accuracy and efficiency of our method
using three PDE-constrained OUU problems; the optimal source control of a
semilinear elliptic PDE in 2D with an uncertain diffusion coefficient field,
and the optimal boundary control of fluid flow governed by Navier–Stokes
equations in both 2D and 3D under an uncertain inflow velocity field. We
present a detailed comparison of the accuracy and cost of our method against
the ground truth solutions for the 2D cases which are still affordable, and
demonstrate the power of our method to the 3D flow control problem where the
ground truth solution is prohibitively expensive to compute.

For the numerical results, we use FEniCS [42] to implement the finite element
discretizations with direct solvers from PETSc [5]. Data generation is
conducted with the use of hIPPYlib [63] and hIPPYflow [55]. Neural network
approximations are implemented using TensorFlow [1]. Unless otherwise
specified, timings are carried out on a single compute node with an Intel Xeon
Gold 6248R processor and NVIDIA A100 40GB GPU for the neural network
computations.

###  4.1 Cost-accuracy comparison of the neural operator

We train neural operators using varying training data sizes, both with (MR-
DINO) and without Jacobian loss (MR-NO). The neural operators are then used to
solve the OUU problem by SAA of the cost functional using a large sample size.
We denote optimal solutions computed by the neural operator by
$z^{*}_{\mathrm{NN}}$. For comparison, we solve the OUU problems using the PDE
by SAA with different sample sizes, representing PDE-constrained OUU under
different computational budgets. We denote the optimal solutions computed
using the PDE by $z^{*}_{\mathrm{PDE}}$.

The accuracies of the approximations are evaluated with respect to the
reference solutions, which are obtained by solving the PDE-based OUU problem
with a large sample size for SAA. This reference optimal solution is denoted
as $z^{*}_{\mathrm{ref}}$. The performance of the approximate optimal
solutions $z_{\mathrm{approx}}^{*}$ are compared to $z^{*}_{\mathrm{ref}}$
using

(24) $\textrm{relative optimal cost error
}:=\frac{|\mathcal{J}(z_{\mathrm{approx}}^{*})-\mathcal{J}(z^{*}_{\mathrm{ref}})|}{|J(z^{*}_{\mathrm{ref}})|}.$

We compare the computational cost in terms of the total number of state PDE
solves used to obtain the optimal solution, which is the dominant cost of the
optimization process. For the neural operator-based optimization, this
corresponds to the total number of training samples. For the PDE-based
optimization, it is given by the number of optimization iterations times the
sample size for SAA.

Additionally in Section 4.4, we present the computational costs of the state
PDE solves in comparison to the costs of the linearized PDE solves required to
compute Jacobian (vector products) used in both the PDE-based gradient
computation and neural operator Jacobian training data generation using [55].
These costs are also compared to the neural operator training and evaluation
costs.

###  4.2 Optimal source control of a semilinear elliptic PDE

We first consider the source control of a semilinear elliptic PDE over a
square domain, $\Omega=(0,1)^{2}$, with a log-normal parameter field. The PDE
is given by

(25) $-\nabla\cdot(e^{m}\nabla u)+ru^{3}=\sum_{i=1}^{49}z_{i}f_{i},$

with homogeneous Dirichlet boundary conditions. The coefficient $r=0.1$ is a
constant reaction coefficient. The log permeability
$m\sim\nu_{m}=\mathcal{N}(\bar{m},\mathcal{C})$ is a Gaussian random field for
which we specify $\bar{m}=-1$ and $\mathcal{C}=(-\gamma\Delta+\delta I)^{-2}$
with $\gamma=0.1,\delta=5.0$ for the distribution of $m$. The control
variables $z=\\{z_{i}\\}_{i=1}^{49}$, for which we consider box bounds
$\mathcal{Z}_{\mathrm{ad}}=[-4,4]^{49}$, define the strengths of invidiual
sources $\\{f_{i}\\}_{i=1}^{49}$ in a $7\times 7$ grid of localized Gaussian
sources, with $f_{i}$ being given by

(26)
$f_{i}(x)=\frac{1}{\sigma\sqrt{2\pi}}\exp\left(-\frac{|x-x_{i}|^{2}}{2\sigma^{2}}\right),$

where $x_{i}$ is the position of the source and $\sigma=0.08$ is a width
parameter. The PDE is discretized in a $64\times 64$ uniform triangular mesh
using piecewise linear elements for both the random parameter and state
spaces, leading to $d_{U}=d_{M}=4225$. The performance function $Q$ is of the
following tracking type,

(27) $Q(u):=\int_{\Omega}(u-u_{\text{target}})^{2}dx,$

where $u_{\text{target}}$ is the target state for the system. This allows us
to define an optimal control problem that minimizes the CVaR of the tracking
objective, i.e.,

(28)
$\min_{z\in\mathcal{Z}_{\mathrm{ad}}}\mathrm{CVaR}_{\alpha}[Q](z)\qquad\text{subject
to }\text{\eqref{eq:elliptic_pde}}.$

Note that we solve the CVaR optimization problem using the reformulation (6).
As examples, we consider a sinusoidal target state
$u_{\mathrm{target}}=\sin(2\pi x_{1})\sin(2\pi x_{2})$, and a quadratic target
state $u_{\mathrm{target}}=4x_{2}(1-x_{2})$.

####  4.2.1 Solution by neural operator

We generate the training data for the neural operator by sampling from the
input distribution $\nu_{m}\otimes\nu_{z}$, where we use the auxiliary
distribution $\nu_{z}=\text{Uniform}(-4,4)^{49}$ for the control variables.
For the reduced basis generation, we use $512$ samples to compute the POD
basis for the output and use the KLE basis for the input parameter, which are
given as the discrete eigenfunctions of the covariance operator of $\nu_{m}$.

We train neural networks of two different architectural sizes. One is a small
network with input rank $r_{M}=50$ for the input basis and output rank
$r_{U}=100$, connected by a dense neural network with two hidden layers of
width $200$ that approximates the mapping between the reduced basis
coefficients. We also consider a larger network with $r_{M}=100$ and
$r_{U}=300$, and two hidden layers of width $400$. All networks use softmax
activation functions. We train neural networks using training data sets of
size 512, 1,024, 2,048, and 4,096, both with and without Jacobian training.
The neural networks are trained using Adam for 1,600 epochs with an initial
learning rate of $10^{-3}$ that is reduced to $2.5\times 10^{-4}$ after 800
epochs. In Figure 2, we plot the mean relative $L^{2}(\Omega)$ errors of the
neural operator and the relative errors of its Jacobian measured in the
Hilbert–Schmidt norm. The errors are computed on a test set of 1,024 samples
from the input distribution, and averaged over 10 different runs of the
training process with different initializations. The incorporation of Jacobian
training consistently leads to smaller errors in both the solution and its
Jacobian. These improved approximations give rise to smaller bounds for the
control gradient errors as discussed in Proposition 3.1. Interestingly, we
observe that without Jacobian training, the larger architecture performs worse
than the smaller architecture as a consequence of overfitting. With Jacobian
training, the small architecture saturates in accuracy, while the larger
architecture yields lower errors and continues to improve. This suggests that
the additional Jacobian training data allows one to use larger and more
expressive architectures while mitigating the tendency of overfitting.

Figure 2: State ($L^{2}(\Omega)$) (left) and Jacobian
($\mathrm{HS}(\mathcal{Z},\mathcal{U})$) (right) testing errors of elliptic
PDE neural operators trained with (MR-DINO) and without Jacobian training (MR-
NO).

Figure 3: Top-left: a sinusoidal target $u_{\text{target}}$. Top-middle: a
random parameter sample $m$. Right: optimal control $z_{\mathrm{NN}}^{*}$
using the neural operator. Bottom-left and bottom-middle: neural surrogate
$u_{w}(m,z_{\mathrm{NN}}^{*})$ and PDE solution $u(m,z_{\mathrm{NN}}^{*})$ at
$m$ and $z_{\mathrm{NN}}^{*}$. The SAA optimization problem with MR-DINO is
solved in 52 seconds.

The OUU problem is solved using the neural operator by SAA with a sample size
of $N=2,048$, using L-BFGS-B to obtain the optimal control
$z_{\mathrm{NN}}^{*}$. For clarity of presentation, we only consider the
larger network architecture with $r_{M}=100$ and $r_{U}=300$ for the remainder
of this section. Figure 3 shows an example of the optimization solution for
the sinusoidal target state with MR-DINO trained on 2,048 samples. In
particular, we present the target state, the computed optimal control
$z_{\mathrm{NN}}^{*}$, a sample of the random coefficient field $m$, and the
state corresponding to the optimal control for the sampled $m$ as predicted by
both the neural operator approximation and the true PDE. Comparing the neural
network prediction of the state to the true PDE, we observe that visually, the
neural operator can accurately approximate the PDE solution operator.
Moreover, the computed minimizer aims to match the state to the target while
being robust to the different possible realizations of $m$.

####  4.2.2 Cost-accuracy comparison with PDE solutions using SAA

To quantify the accuracy of the neural operator OUU solution relative to its
costs, we compare the CVaR at the optimal controls $z_{\mathrm{NN}}^{*}$ and
$z_{\mathrm{PDE}}^{*}$. For the PDE, we solve the optimal control problem for
$z_{\mathrm{PDE}}^{*}$ using SAA with sample sizes of $N=$16, 32, 64, 128 and
256. An accurate reference solution of the OUU problem $z_{\mathrm{ref}}^{*}$
is also computed using SAA with $N=4,096$.

Figure 4: Optimal cost values for the sinusoidal (left) and quadratic (right)
target states at optimal solutions versus the number of state PDE solves
required.

The CVaR values at optimal controls are plotted in Figure 4 as a function of
the number of state solves required to obtain the optimal control. The CVaR
values are evaluated by a Monte Carlo estimator with 8,192 samples, and are
averaged across 10 different runs of the OUU problem with different sets of
parameter samples. In the case of the neural operators, the training
initializations are also randomized across the 10 runs. Notice that due to
sampling errors, the CVaR values for the PDE-based optimal controls are
suboptimal relative to the reference value, and tend to vary significantly
between runs. The optimality improves as the sample sizes increase. On the
other hand, the CVaR values from MR-DINO are much closer to the reference
value, and are much more consistent across runs. This suggests that the bias
committed by the neural operators is less problematic than the variance
(sampling error) arising from the PDE. Moreover, for the same number of
training samples, the MR-DINO with Jacobian training greatly improves the
optimal solutions.

Figure 5: Relative error of the cost functional versus the number of state PDE
solves required for the optimization with different target states and $\beta$
values.

The accuracy is more clearly compared in Figure 5, where we plot the relative
errors of the CVaR values at the optimal controls with respect to the
reference PDE solution. We additionally include OUU runs where the cost
functional is CVaR with quantiles $\beta=0.9$ and $\beta=0.99$. Here we
observe that across all four cases, the neural operators with Jacobian
training (MR-DINO) are able to obtain OUU solutions with lower cost values
than the PDE-based solutions, despite using over $10\times$ fewer state PDE
solves. On the other hand, neural operators without Jacobian training are only
more cost-effective than the PDE-based solutions across a single OUU run in
the sinusoidal cases. Evidently, Jacobian training adds a valuable source of
information at a fraction of the cost of the state PDE solve, and is highly
beneficial to the accuracy of the neural operator OUU solutions, especially in
the low-data regime. It is important to note that in these results, the same
neural operators are used to solve the OUU problems across all four cases. The
training cost can easily be amortized across different choices of the
performance functions and risk measures, which in this example, correspond to
different target states and $\beta$ values. Moreover, we observe that for both
the neural operator and PDE results, the optimality of the OUU solutions
degrade as the quantile value $\beta$ increases. Larger $\beta$ values
represent increased weighting of the tail of the distribution, which require
more samples for both neural operator training and accurate estimation of the
CVaR. Nevertheless, for the $\beta=0.99$ case considered, the MR-DINO
solutions are still more than an order of magnitude more cost-effective than
the PDE solutions.

###  4.3 Optimal boundary control of flow around a bluff body

Next, we consider the boundary control of flow around a bluff body with
uncertain inlet conditions governed by the steady state Navier–Stokes
equation. We consider the two dimensional domain shown in Figure 6, where the
setup is analogous to that considered in [49]. We write the flow equations as

(29a) $\displaystyle(\mathbf{u}\cdot\nabla)\mathbf{u}+\nabla
p-\nu\Delta\mathbf{u}$ $\displaystyle=0\qquad x\in\Omega,$ (29b)
$\displaystyle\nabla\cdot\mathbf{u}$ $\displaystyle=0\qquad x\in\Omega,$ (29c)
$\displaystyle\mathbf{u}-e^{m}\mathbf{e}_{1}$ $\displaystyle=0\qquad
x\in\Gamma_{I},$ (29d) $\displaystyle\mathbf{T}(\mathbf{u},p)\mathbf{n}$
$\displaystyle=0\qquad x\in\Gamma_{O},$ (29e)
$\displaystyle\mathbf{u}\cdot\mathbf{n}=0,\;\mathbf{T}(\mathbf{u},p)\mathbf{n}\cdot\mathbf{t}$
$\displaystyle=0\qquad x\in\Gamma_{W},$ (29f) $\displaystyle\mathbf{u}$
$\displaystyle=0\qquad x\in\Gamma_{B},$ (29g)
$\displaystyle\mathbf{u}-\phi(z)\mathbf{n}$ $\displaystyle=0\qquad
x\in\Gamma_{C}.$

In the above equations, the state $u=(\mathbf{u},p)$ consists of the velocity
and pressure fields, $\nu=0.005$ is the viscosity,
$\mathbf{T}=-p\mathbf{I}+2\nu\;\mathrm{symm}(\nabla\mathbf{u})$ is the stress
tensor, where $\mathrm{symm}(\mathbf{A}):=(\mathbf{A}+\mathbf{A}^{T})/2$. The
boundaries $\Gamma_{I}$, $\Gamma_{O}$, $\Gamma_{W}$, $\Gamma_{B}$, and
$\Gamma_{C}$ are as labeled in Figure 6, $\mathbf{n}$ and $\mathbf{t}$ are the
unit normal and tangent vectors along the boundaries respectively. The inflow
velocity $\mathbf{u}|_{\Gamma_{I}}$ is given by the trace of a 2D lognormal
random field, $e^{m}|_{\Gamma_{I}}\mathbf{e}_{1}$, where
$m\sim\mathcal{N}(0,\mathcal{C})$ and $\mathbf{e}_{1}=(1,0)$. The control
variables define the normal flow velocity along the sides of the bluff body
$\Gamma_{B}$ using a cubic B-spline representation,
$\phi(z)(x)=\sum_{i=1}^{18}z_{i}\phi_{i}(x),$ where $z_{i},\phi_{i}$ are the
weights and basis functions for the upper ($i=1,...,9$) and lower
($i=10,...,18$) sides respectively.

Figure 6: Flow domain $\Omega=(2,0)\times(1,0)$ with labelled boundaries.

In this problem, we consider two different performance functions. The first is
the viscous dissipation rate of the flow, defined as

(30)
$Q_{\text{Dissipation}}(\mathbf{u}):=2\nu\int_{\Omega}\mathrm{symm}(\nabla\mathbf{u}):\mathrm{symm}(\nabla\mathbf{u})dx.$

This corresponds to the drag force on the bluff body. We also consider a
tracking type objective, where we seek to minimize the difference between the
velocity field and a target velocity field $\mathbf{u}_{\text{target}}$ behind
the bluff body,

(31)
$Q_{\text{Tracking}}(\mathbf{u}):=\int_{\Omega_{o}}|\mathbf{u}-\mathbf{u}_{\text{target}}|^{2}dx,$

taking $\Omega_{o}=(0.6,2)\times(1,0)$ and $\mathbf{u}_{\text{target}}=(1,0)$.
Additionally, we adopt an $L^{2}$ penalization term on the velocity profile
induced by the control,

(32) $\mathcal{P}(z)=\alpha\int_{\Gamma_{C}}|\phi(z)|^{2}ds,$

where $\alpha$ is a weighting parameter on the penalization term. We consider
the problem of minimizing of the CVaR with $L^{2}$ penalization subject to the
PDE constraint, i.e.,

(33)
$\min_{z\in\mathcal{Z}_{\mathrm{ad}}}\mathrm{CVaR}_{\beta}[Q](z)+\mathcal{P}(z)\qquad\text{s.t.
}\eqref{eq:ns_start}-\eqref{eq:ns_end}.$

The PDEs are discretized on a triangular mesh using Taylor–Hood elements for
the state $u=(\mathbf{u},p)$ with quadratic elements for the velocity field
$\mathbf{u}$ and linear elements for the pressure field $p$. This leads to the
state dimension $d_{U}=42,649$. Quadratic elements are also used for the
random parameter field. This discretization has a nominal dimension of
$d_{M}=18,921$ since it is sampled on the entire domain $\Omega$, but its
trace on the left boundary has 101 degrees of freedom. The state PDE is solved
using a backtracking Newton method with Galerkin–Least Squares (GLS)
stabilization.

####  4.3.1 Solution by neural operator

We generate the training data for the 2D control problem using the input
distribution $\nu_{m}\otimes\nu_{z}$, with a Gaussian distribution
$\nu_{z}=\mathcal{N}(0,I)$ as the auxiliary control distribution. For the
reduced bases, we consider a POD basis with rank $r_{U}=200$ for the state,
computed from 256 samples from the training data. Though dimension reduction
of the random inflow parameter is not necessary for this discretization, we
still consider its representation in the reduced basis, since this
representation is amenable to further mesh refinement. Here, we adopt a rank
of $r_{M}=100$ for the KLE basis. We use dense neural networks with $2$ hidden
layers of width $400$ to approximate the mapping from the reduced input space
to the reduced output space. We train the neural networks using training data
sets of size 256, 512, 1,024, and 2,048, both with and without Jacobian
training. Testing errors of the trained neural operators and their Jacobians
are shown in Figure 7. Similar to the previous example, we see that Jacobian
training improves both the state and Jacobian accuracy of the neural
operators.

Figure 7: State ($L^{2}(\Omega)$) (left) and Jacobian
($\mathrm{HS}(\mathcal{Z},\mathcal{U})$) (right) testing errors of the 2D
Navier–Stokes neural operators trained with (MR-DINO) and without Jacobian
training (MR-NO).

We solve the OUU problem using the trained neural operator by SAA with a
sample size of $N=2,048$, and using the L-BFGS algorithm to obtain the optimal
controls $z_{\mathrm{NN}}^{*}$. As an example, we consider the CVaR of the
viscous dissipation rate objective with $\beta=0.95$ and an $L^{2}$
penalization with $\alpha=10^{-2}$. Figure 8 compares the uncontrolled flow to
the controlled flow for a random inflow sample from $\nu_{m}$, where the
controlled flow uses the optimal control computed from a MR-DINO with only 256
training samples. The controlled flow exhibits significantly smaller
recirculation regions behind the bluff body, which effectively reduces the
drag force.

Figure 8: Samples of flow fields using no control (left) and the optimal
control (right) computed using MR-DINO with 256 training samples. The SAA
optimization problem with MR-DINO is solved in 34 seconds.

####  4.3.2 Comparison with PDE solutions using SAA

To quantify the accuracy of the neural operator, we compare the neural
operator solutions against PDE solutions obtained by SAA with 16, 32, 64, and
128 samples. A reference solution is obtained by the PDE using 4,096 samples
for the SAA. In addition to the viscous dissipation objective, we also
consider the CVaR of the tracking objective with $\beta=0.95$ and the penalty
parameter $\alpha=1$. We present in Figure 9 the relative error in the cost
functional values at the optimal controls with respect to the reference PDE-
based optimal solution for both the viscous dissipation and tracking
objectives. As in the semilinear elliptic PDE examples, the errors are
averaged across 10 runs with different random inflow samples and neural
network initializations.

Figure 9: Relative error of the cost functional versus the number of state PDE
solves required for the optimization with different objectives and penalty
parameters.

For the viscous dissipation objective, we observe that the PDE-based approach
achieves near optimal OUU solutions with small sample sizes. This suggests
that the SAA cost functional is well-correlated with the true CVaR cost
functional, and that the OUU solution is not very sensitive to the random
inflow parameter. Nevertheless, we observe that the neural operators with
Jacobian training are able to achieve similar accuracy in optimal cost as the
PDE-based approach using approximately 10 times fewer state PDE solves. In
contrast, without the Jacobian training the neural operators are not as cost-
effective than the PDE-based OUU solve if used only for a single optimization.
On the other hand, the PDE-based optimal solutions for the tracking objective
exhibit much larger errors compared to the reference. This is likely because
the optimal control for the tracking objective is much more sensitive to the
inflow profile. In this example, the neural operators with Jacobian training
attain an order of magnitude smaller errors than the PDE solutions, while only
using 100$\times$ fewer state PDE solves. Even without Jacobian training, the
neural operators demonstrate a 10$\times$ improvement in cost-effectiveness
compared to the PDE solutions.

We make an important observation here on the relationship between the
generalization accuracy of the neural operators and their performance in
solving OUU problems. Figure 7 shows that the average generalization error for
MR-DINO trained on 256 samples is larger than that of the MR-NO trained on
2,048 samples (without Jacobian training). However, for the OUU problems as
shown in Figure 9, we see that the MR-DINO trained on 256 samples achieves
lower cost values than the MR-NO trained on 2,048 samples in both the
dissipation rate and tracking objectives. This suggests that a high $L^{2}$
generalization accuracy cannot guarantee the neural operator’s performance in
solving OUU problems, which can be corrupted by the Jacobian errors that are
propagated to the gradients as shown in Proposition 3.1.

###  4.4 Comparison of timings

First of all, to demonstrate the small incremental computational cost for
generating the Jacobian data, we report the computation time for the state PDE
solve, the first linearized PDE solve in Jacobian computation for which an LU
factorization is computed and stored for the linear operator, and the
subsequent LU solve for which only a back substitution is needed. The results
are presented in Table 1 for the semilinear elliptic and 2D Navier–Stokes
examples.

Time (in seconds) | Semilinear Elliptic | 2D Navier–Stokes
---|---|---
State PDE solve | 0.869 | 12.81
Jacobian (LU) | 0.366 | 2.95
Jacobian (Back sub.) | 0.004 | 0.02
Table 1: Time (in seconds) for the state PDE solve, the first linearized PDE
solve in Jacobian computation for which an LU factorization is computed and
stored for the linear operator, and the subsequent LU solve for which only a
back substitution is needed. The reported timings are averaged over 100 random
samples of $(m,z)$.

We then present in Table 2 the average training time for the neural operators
with and without Jacobian training in the two examples. The training time is
broken down into three parts for (1) the training data generation (state and
Jacobian training pairs), (2) preprocessing of data, and (3) neural network
training. We use the networks with two hidden layers of widths $(400,400)$,
and use the reduced dimension $(r_{M},r_{U})=(100,300)$ and $(100,200)$ for
the semilinear elliptic PDE and 2D Navier–Stokes examples respectively, and
use 1,024 training samples.

Time (in seconds) | Semilinear Elliptic | 2D Navier–Stokes
---|---|---
|  $L^{2}$ Only | DINO |  $L^{2}$ Only | DINO
Data Generation | 889.9 | 1,461.3 | 13,117.4 | 16,486.4
Pre-processing | 2.6 | 22.8 | 29.9 | 122.5
Training | 140.7 | 650.9 | 142.0 | 486.7
Total | 1,033.2 | 2,135.0 | 13,289.3 | 17,095.6
Table 2: Time (in seconds) for data generation, pre-processing (reduced basis
construction), and training in neural operator construction with 1,024
training samples.

The semilinear elliptic PDE is mildly nonlinear, requiring 2–3 Newton
iterations on average. The data generation for DINO (both state and Jacobian)
takes approximately $1.5\times$ the time of generating the state data alone.
The Navier–Stokes problem is more nonlinear, requiring many more Newton
iterations. The DINO data generation takes approximately $1.25\times$ the time
of generating the state data alone, meaning that the Jacobian data generation
comes with only a small increase in cost. Moreover, as the semilinear elliptic
PDE solve is not expensive for the given discretization dimension, the neural
network training takes a large portion of the overall time. On the other hand,
since the neural network training is conducted in the reduced basis spaces,
the Navier–Stokes example shows similar training times, despite having a state
discretization that is about $10\times$ larger than the semilinear elliptic
PDE example. The increased nonlinearity and state dimension for the
Navier–Stokes problem means the training costs are much lower relative to the
data generation costs, being only 1–3% of the overall time. In both examples,
the small additional computation time for training the MR-DINOs results in
large performance improvements in the OUU problem as shown in the comparisons
of Sections 4.2 and 4.3.

We also present the time for the evaluation of the state, the performance
function $Q$, and the gradient of the performance function with respect to the
control variables $z$ in Table 3. Here, we report the time for an evaluation
using a single random parameter sample and that using a batch of 2,048
samples. Since there are no architectural differences between the MR-NO and
MR-DINO, we only report the timing for the MR-DINO.

Time (in seconds) | Semilinear Elliptic | 2D Navier–Stokes
---|---|---
| Single | Batched (2,048) | Single | Batched (2,048)
State evaluation | 0.0023 | 0.10 | 0.0036 | 0.28
$Q$ evaluation | 0.0034 | 0.13 | 0.0039 | 0.27
$Q$ $z$-gradient | 0.0041 | 0.16 | 0.0045 | 0.29
Table 3: Neural operator evaluation time for the state, the performance
function, and its gradient, reported for a single sample and batched
evaluation over 2,048 samples.

Comparing the time of the PDE solve in Table 1 and of the neural operator
evaluation in Table 3, we see that once trained, the neural operator offers on
average a $380\times$ speed up for a single solve of the semilinear elliptic
PDE and a $3,600\times$ for a single solve of the 2D Navier–Stokes PDE. The
speed ups are much more significant in the batched case, where the time taken
to evaluate the neural operator for 2,048 different samples is a fraction of
that for a single PDE solve in both examples, and $18,000\times$ and
$94,000\times$ for the total samples. Comparisons for the performance function
evaluation and gradient are similar, noting that the time required to evaluate
the gradient of $Q$ using the PDE is comparable to that of a Jacobian action
with LU factorization reported in Table 1.

###  4.5 3D Navier–Stokes example

Finally, we demonstrate the scalability of our approach by considering a 3D
Navier–Stokes example for the boundary control of the flow around a bluff body
in the presence of uncertain inflow conditions, with the setup analogous to
that in 2D. This problem is not amenable to traditional PDE-based OUU methods
as a single state PDE solve takes 30 minutes with parallel computation using
48 cores in one CPU node of Frontera at TACC. We consider a domain of
dimensions $2\times 1\times 1$, with a bluff body that is defined similar to
the 2D case, at an angle-of-attack of 30 degrees. The inflow condition at
$x=0$ is given by $e^{m}\mathbf{e}_{1}$ with
$m\sim\mathcal{N}(0,(-\gamma\Delta+\delta I)^{-2})$, where
$\mathbf{e}_{1}=(1,0,0)$. We choose $\gamma=1.5$ and $\delta=7.5$ such that
the pointwise variance and correlation lengths are similar to the 2D case. The
control variables define the normal flow velocity along the sides of the bluff
body using a tensor product of quadratic B-splines. Figure 10 illustrates the
obstacle geometry and the controlled region. Using a tetrahedral mesh with
Taylor–Hood elements for the state and quadratic elements for the random
parameter, the discretization yields $d_{U}=1,035,243$, $d_{M}=4,369$
(boundary degrees of freedom), and $d_{Z}=50$.

We consider the CVaR optimization problem with $\beta=0.95$, using the viscous
dissipation objective with the $L^{2}(\Gamma_{C})$ penalization term on the
control. To solve the OUU problem, we train a MR-DINO with 448 training
samples, using input rank $r_{M}=100$ and output rank $r_{U}=200$. With
Jacobian training, the neural operator is able to achieve a mean relative
$L^{2}(\Omega)$ generalization error of $1.8\%$. The neural operator is then
deployed to solve the OUU problem by SAA with a sample size of $N=1,024$,
using the L-BFGS algorithm to obtain the optimal controls
$z_{\mathrm{NN}}^{*}$. The neural operator takes 0.13 seconds to evaluate the
performance function for all 1,024 samples, which is over $10^{7}\times$
faster than using the PDE solver.

A comparison of the flow field at a random control and the optimal control
$z_{\text{NN}}^{*}$ is shown in Figure 10 for a sample inflow profile. In the
controlled flow field, the recirculation region is shifted towards the rear of
the bluff body due to the boundary control along the bluff body, thereby
reducing the overall viscous dissipation rate and hence drag. We remark that
in the 3D problem, the boundary control is only defined in the central region
of the top and bottom faces, as illustrated in Figure 10. Unlike the 2D case,
the flow field exhibits a more complex 3D structure, as the flow also wraps
around the bluff body in the $x_{3}$ direction, for which the boundary control
is not able to completely eliminate the recirculation region behind the bluff
body.

Figure 10: Top: the obstacle geometry for the 3D Navier–Stokes control
problem. The control variables prescribe the normal velocity on the top and
bottom faces of the obstacle over a $0.2\times 0.4$ plane. Tangential velocity
is set to zero. Bottom: Streamlines for a sample of the flow field with a
random control (left) and an optimal control (right) computed using the MR-
DINO with 448 training samples. The SAA optimization problem with MR-DINO is
solved in 53 seconds.

##  5 Conclusions

In this work, we have presented a novel framework for solving PDE-constrained
OUU problems using neural operators to approximate the mapping from the joint
input spaces of the uncertain parameters and optimization variables to the
solution of the underlying PDE. The key contribution is the DINO training of
neural operators on the derivatives of the solution map with respect to the
optimization variable, along with the use of reduced basis architectures that
enable scalable and efficient data generation and training.

Through our numerical experiments, we have demonstrated that reduced basis
neural operators can be constructed to efficiently solve a range of PDE-
constrained OUU problems. In particular, we consistently showed that Jacobian
training was extremely effective in improving the function approximation, the
gradients, and critically, the quality of the optimization solution. We also
observed that the MR-DINOs were more cost-effective for OUU than standard SAA-
based PDE solutions, with over 10$\times$ fewer state PDE solves for the same
accuracy. Moreover, once trained, the online evaluation cost with the neural
operator approximation was reduced by several orders of magnitude, giving rise
to the potential for real-time solution of PDE-constrained OUU problems. We
remark that when further accuracy is required, MR-DINO can be employed in a
multifidelity Monte Carlo framework [51] to guarantee convergence to the exact
OUU solution. In this setting, the low construction cost of MR-DINO relative
to its accuracy makes it an appealing control variate.

Our demonstrations focused on steady-state control problems with finite
dimensional optimization variables. In future work, we will apply our method
to the solution of time-dependent OUU problems with function-valued
optimization variables and more general risk measures and probability
constraints. Furthermore, strategies for selecting the training distribution
of the control variable $\nu_{z}$ can also be explored in future work.

## References

  * [1] M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Jozefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Vasudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng, TensorFlow: Large-scale machine learning on heterogeneous systems, 2015, https://www.tensorflow.org/.  Software available from tensorflow.org.
  * [2] A. Alexanderian, N. Petra, G. Stadler, and O. Ghattas, Mean-variance risk-averse optimal control of systems governed by PDEs with random parameter fields using quadratic approximations, SIAM/ASA Journal on Uncertainty Quantification, 5 (2017), pp. 1166–1192, https://doi.org/10.1137/16M106306X.
  * [3] A. A. Ali, E. Ullmann, and M. Hinze, Multilevel Monte Carlo analysis for optimal control of elliptic PDEs with random coefficients, SIAM/ASA Journal on Uncertainty Quantification, 5 (2017), pp. 466–492.
  * [4] A. Alla, M. Hinze, P. Kolvenbach, O. Lass, and S. Ulbrich, A certified model reduction approach for robust parameter optimization with pde constraints, Advances in Computational Mathematics, 45 (2019), pp. 1221–1250.
  * [5] S. Balay, S. Abhyankar, M. F. Adams, S. Benson, J. Brown, P. Brune, K. Buschelman, E. M. Constantinescu, L. Dalcin, A. Dener, V. Eijkhout, J. Faibussowitsch, W. D. Gropp, V. Hapla, T. Isaac, P. Jolivet, D. Karpeev, D. Kaushik, M. G. Knepley, F. Kong, S. Kruger, D. A. May, L. C. McInnes, R. T. Mills, L. Mitchell, T. Munson, J. E. Roman, K. Rupp, P. Sanan, J. Sarich, B. F. Smith, S. Zampini, H. Zhang, H. Zhang, and J. Zhang, PETSc Web page.  https://petsc.org/, 2023, https://petsc.org/.
  * [6] K. Bhattacharya, B. Hosseini, N. B. Kovachki, and A. M. Stuart, Model reduction and neural networks for parametric PDEs, SMAI Journal of Computational Mathematics, Volume 7, (2021).
  * [7] P. Binev, A. Cohen, W. Dahmen, R. DeVore, G. Petrova, and P. Wojtaszczyk, Convergence rates for greedy algorithms in reduced basis methods, SIAM journal on mathematical analysis, 43 (2011), pp. 1457–1472.
  * [8] A. Borzì, Multigrid and sparse-grid schemes for elliptic control problems with random coefficients, Computing and Visualization in Science, 13 (2010), pp. 153–160.
  * [9] L. Cao, T. O’Leary-Roseberry, P. K. Jha, J. T. Oden, and O. Ghattas, Residual-based error correction for neural operator accelerated infinite-dimensional Bayesian inverse problems, Journal of Computational Physics, (2023), p. 112104.
  * [10] A. Chaudhuri, B. Kramer, M. Norton, J. O. Royset, and K. Willcox, Certifiable risk-based engineering design optimization, AIAA Journal, 60 (2022), pp. 551–565, https://doi.org/10.2514/1.j060539.
  * [11] P. Chen and O. Ghattas, Taylor approximation for chance constrained optimization problems governed by partial differential equations with high-dimensional random parameters, SIAM/ASA Journal on Uncertainty Quantification, 9 (2021), pp. 1381–1410.
  * [12] P. Chen, M. Haberman, and O. Ghattas, Optimal design of acoustic metamaterial cloaks under uncertainty, Journal of Computational Physics, 431 (2021), p. 110114.
  * [13] P. Chen and A. Quarteroni, Weighted reduced basis method for stochastic optimal control problems with elliptic PDE constraints, SIAM/ASA J. Uncertainty Quantification, 2 (2014), pp. 364–396.
  * [14] P. Chen, A. Quarteroni, and G. Rozza, Multilevel and weighted reduced basis method for stochastic optimal control problems constrained by Stokes equations, Numerische Mathematik, 133 (2016), pp. 67–102.
  * [15] P. Chen, A. Quarteroni, and G. Rozza, Reduced basis methods for uncertainty quantification, SIAM/ASA Journal on Uncertainty Quantification, 5 (2017), pp. 813–869.
  * [16] P. Chen and J. O. Royset, Performance bounds for PDE-constrained optimization under uncertainty, arXiv:2110.10269, accepted in SIAM Journal on Optimization, (2023).
  * [17] P. Chen and C. Schwab, Sparse-grid, reduced-basis Bayesian inversion, Computer Methods in Applied Mechanics and Engineering, 297 (2015), pp. 84 – 115.
  * [18] P. Chen and C. Schwab, Model order reduction methods in computational uncertainty quantification, Handbook of Uncertainty Quantification, Springer, (2016).
  * [19] P. Chen and C. Schwab, Sparse-grid, reduced-basis Bayesian inversion: Nonaffine-parametric nonlinear equations, Journal of Computational Physics, 316 (2016), pp. 470–503.
  * [20] P. Chen, U. Villa, and O. Ghattas, Taylor approximation and variance reduction for PDE-constrained optimal control under uncertainty, Journal of Computational Physics, 385 (2019), pp. 163–186, https://arxiv.org/abs/1804.04301.
  * [21] A. Cohen and R. DeVore, Approximation of high-dimensional parametric PDEs, Acta Numerica, 24 (2015), pp. 1–159.
  * [22] X. Du, J. R. Martins, T. O’Leary-Roseberry, A. Chaudhuri, O. Ghattas, and K. E. Willcox, Learning Optimal Aerodynamic Designs through Multi-Fidelity Reduced-Dimensional Neural Networks, in AIAA SCITECH 2023 Forum, 2023, p. 0334.
  * [23] M. Eigel, M. Haase, and J. Neumann, Topology optimisation under uncertainties with neural networks, Algorithms, 15 (2022), p. 241, https://doi.org/10.3390/a15070241.
  * [24] S. Fresca and A. Manzoni, POD-DL-ROM: enhancing deep learning-based reduced order models for nonlinear parametrized PDEs by proper orthogonal decomposition, Computer Methods in Applied Mechanics and Engineering, 388 (2022), p. 114181.
  * [25] P. A. Guth, C. Schillings, and S. Weissmann, A general framework for machine learning based optimization under uncertainty, 2021, https://doi.org/10.48550/arXiv.2112.11126.
  * [26] Z. Hao, C. Ying, H. Su, J. Zhu, J. Song, and Z. Cheng, Bi-level physics-informed neural networks for PDE constrained optimization using broyden’s hypergradients, arXiv preprint arXiv:2209.07075, (2022).
  * [27] J. S. Hesthaven and S. Ubbiali, Non-intrusive reduced order modeling of nonlinear problems using neural networks, Journal of Computational Physics, 363 (2018), pp. 55–78, https://doi.org/10.1016/j.jcp.2018.02.037.
  * [28] R. Hwang, J. Y. Lee, J. Y. Shin, and H. J. Hwang, Solving PDE-constrained control problems using operator learning, in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, 2022, pp. 4504–4512.
  * [29] P. Jin, S. Meng, and L. Lu, MIONet: Learning multiple-input operators via tensor product, SIAM Journal on Scientific Computing, 44 (2022), pp. A3490–A3514.
  * [30] T. Keil, H. Kleikamp, R. J. Lorentzen, M. B. Oguntola, and M. Ohlberger, Adaptive machine learning-based surrogate modeling to accelerate PDE-constrained optimization in enhanced oil recovery, Advances in Computational Mathematics, 48 (2022), p. 73.
  * [31] A. Kodakkal, B. Keith, U. Khristenko, A. Apostolatos, K.-U. Bletzinger, B. Wohlmuth, and R. Wüchner, Risk-averse design of tall buildings for uncertain wind conditions, Computer Methods in Applied Mechanics and Engineering, 402 (2022), p. 115371, https://doi.org/10.1016/j.cma.2022.115371, https://www.sciencedirect.com/science/article/pii/S0045782522004443.  A Special Issue in Honor of the Lifetime Achievements of J. Tinsley Oden.
  * [32] D. Kouri, D. Heinkenschloos, M. Ridzal, and B. Van Bloemen Waanders, A trust-region algorithm with adaptive stochastic collocation for PDE optimization under uncertainty, SIAM Journal on Scientific Computing, 35 (2012), pp. 1847–1879.
  * [33] D. P. Kouri and T. M. Surowiec, Risk-averse PDE-constrained optimization using the conditional value-at-risk, SIAM Journal on Optimization, 26 (2016), pp. 365–396, https://doi.org/10.1137/140954556.
  * [34] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar, Neural operator: Learning maps between function spaces, arXiv preprint arXiv:2108.08481, (2021).
  * [35] A. Kunoth and C. Schwab, Sparse adaptive tensor Galerkin approximations of stochastic PDE-constrained control problems, SIAM/ASA Journal on Uncertainty Quantification, 4 (2016), pp. 1034–1059.
  * [36] O. Lass and S. Ulbrich, Model order reduction techniques with a posteriori error control for nonlinear robust optimization governed by partial differential equations, SIAM Journal on Scientific Computing, 39 (2017), pp. S112–S139.
  * [37] D. Lee and B. Kramer, Bi-fidelity conditional value-at-risk estimation by dimensionally decomposed generalized polynomial chaos expansion, Structural and Multidisciplinary Optimization, 66 (2023), p. 33, https://doi.org/10.1007/s00158-022-03477-6, https://doi.org/10.1007/s00158-022-03477-6.
  * [38] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, Multipole graph neural operator for parametric partial differential equations, Neural Information Processing Systems, (2020).
  * [39] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, Fourier neural operator for parametric partial differential equations, International Conference on Learning Representations, (2021).
  * [40] Z. Li, H. Zheng, N. Kovachki, D. Jin, H. Chen, B. Liu, K. Azizzadenesheli, and A. Anandkumar, Physics-informed neural operator for learning partial differential equations, arXiv preprint arXiv:2111.03794, (2021).
  * [41] F. Lindgren, H. Rue, and J. Lindström, An explicit link between Gaussian fields and Gaussian Markov random fields: the stochastic partial differential equation approach, Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73 (2011), pp. 423–498, https://doi.org/10.1111/j.1467-9868.2011.00777.x, http://dx.doi.org/10.1111/j.1467-9868.2011.00777.x.
  * [42] A. Logg, K.-A. Mardal, and G. Wells, Automated Solution of Differential Equations by the Finite Element Method: The FEniCS book, vol. 84, Springer Science & Business Media, 2012.
  * [43] L. Lu, P. Jin, G. Pang, and G. E. Karniadakis, DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators, Nature Machine Intelligence, (2021).
  * [44] L. Lu, X. Meng, S. Cai, Z. Mao, S. Goswami, Z. Zhang, and G. E. Karniadakis, A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data, Computer Methods in Applied Mechanics and Engineering, 393 (2022), p. 114778.
  * [45] L. Lu, R. Pestourie, S. G. Johnson, and G. Romano, Multifidelity deep neural operators for efficient learning of partial differential equations with application to fast inverse design of nanoscale heat transport, Physical Review Research, 4 (2022), p. 023210, https://doi.org/10.1103/physrevresearch.4.023210.
  * [46] L. Lu, R. Pestourie, W. Yao, Z. Wang, F. Verdugo, and S. G. Johnson, Physics-informed neural networks with hard constraints for inverse design, SIAM Journal on Scientific Computing, 43 (2021), pp. B1105–B1132, https://doi.org/10.1137/21m1397908.
  * [47] Y. Maday, A. T. Patera, and G. Turinici, A priori convergence theory for reduced-basis approximations of single-parameter elliptic partial differential equations, Journal of Scientific Computing, 17 (2002), pp. 437–446.
  * [48] A. Manzoni, F. Negri, and A. Quarteroni, Dimensionality reduction of parameter-dependent problems through proper orthogonal decomposition, Annals of Mathematical Sciences and Applications, 1 (2016), pp. 341–377, https://doi.org/10.4310/AMSA.2016.v1.n2.a4.
  * [49] A. Manzoni, A. Quarteroni, and S. Salsa, Optimal Control of Partial Differential Equations, Springer International Publishing, 2021, https://doi.org/10.1007/978-3-030-77226-0.
  * [50] N. H. Nelsen and A. M. Stuart, The random feature model for input-output maps between banach spaces, SIAM Journal on Scientific Computing 43 (5), A3212-A3243, (2021).
  * [51] L. Ng and K. Willcox, Multifidelity approaches for optimization under uncertainty, International Journal for Numerical Methods in Engineering, 100 (2014), pp. 746–772, https://doi.org/10.1002/nme.4761.
  * [52] M. Ohlberger and S. Rave, Reduced basis methods: Success, limitations and future challenges, arXiv preprint arXiv:1511.02021, (2015).
  * [53] T. O’Leary-Roseberry, P. Chen, U. Villa, and O. Ghattas, Derivate informed neural operator: An efficient framework for high-dimensional parametric derivative learning, arXiv:2206.10745, (2022), http://arxiv.org/abs/2206.10745.
  * [54] T. O’Leary-Roseberry, X. Du, A. Chaudhuri, J. Martins, K. Willcox, and O. Ghattas, Learning high-dimensional parametric maps via reduced basis adaptive residual networks, Computer Methods in Applied Mechanics and Engineering, 402 (2022), p. 115730.
  * [55] T. O’Leary-Roseberry and U. Villa, hIPPYflow: Dimension reduced surrogate construction for parametric PDE maps in Python, 2021, https://doi.org/10.5281/zenodo.4608729, https://github.com/hippylib/hippyflow.
  * [56] T. O’Leary-Roseberry, U. Villa, P. Chen, and O. Ghattas, Derivative-informed projected neural networks for high-dimensional parametric maps governed by PDEs, Computer Methods in Applied Mechanics and Engineering, 388 (2022), p. 114199.
  * [57] M. Raissi, P. Perdikaris, and G. E. 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 (2019), pp. 686–707.
  * [58] R. Rockafellar and J. Royset, On buffered failure probability in design and optimization of structures, Reliability Engineering & System Safety, 95 (2010), pp. 499–510, https://doi.org/10.1016/j.ress.2010.01.001.
  * [59] R. T. Rockafellar and S. Uryasev, Optimization of conditional value-at-risk, Journal of risk, 2 (2000), pp. 21–42.
  * [60] A. Shapiro, D. Dentcheva, and A. Ruszczynski, Lectures on Stochastic Programming: Modeling and Theory, Society for Industrial and Applied Mathematics, third ed., July 2021, https://doi.org/10.1137/1.9781611976595.
  * [61] K. Shukla, V. Oommen, A. Peyvan, M. Penwarden, L. Bravo, A. Ghoshal, R. M. Kirby, and G. E. Karniadakis, Deep neural operators can serve as accurate surrogates for shape optimization: A case study for airfoils, 2023, https://doi.org/10.48550/ARXIV.2302.00807.
  * [62] H. Tiesler, R. M. Kirby, D. Xiu, and T. Preusser, Stochastic collocation for optimal control problems with stochastic PDE constraints, SIAM Journal on Control and Optimization, 50 (2012), pp. 2659–2682.
  * [63] U. Villa, N. Petra, and O. Ghattas, HIPPYlib: An Extensible Software Framework for Large-Scale Inverse Problems Governed by PDEs: Part I: Deterministic Inversion and Linearized Bayesian Inference, ACM Trans. Math. Softw., 47 (2021), https://doi.org/10.1145/3428447, https://doi.org/10.1145/3428447.
  * [64] S. Wang, M. A. Bhouri, and P. Perdikaris, Fast PDE-constrained optimization via self-supervised operator learning, arXiv preprint arXiv:2110.13297, (2021).
  * [65] K. Wu, T. O’Leary-Roseberry, P. Chen, and O. Ghattas, Large-Scale Bayesian Optimal Experimental Design with Derivative-Informed Projected Neural Network, Journal of Scientific Computing, 95 (2023), p. 30.
  * [66] J. Yu, L. Lu, X. Meng, and G. E. Karniadakis, Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems, Computer Methods in Applied Mechanics and Engineering, 393 (2022), p. 114823.
  * [67] M. J. Zahr, K. T. Carlberg, and D. P. Kouri, An efficient, globally convergent method for optimization under uncertainty using adaptive model reduction and sparse grids, SIAM/ASA Journal on Uncertainty Quantification, 7 (2019), pp. 877–912.
  * [68] Q. Zhao, D. B. Lindell, and G. Wetzstein, Learning to solve PDE-constrained inverse problems with graph networks, 2022.