# Semilinear Feynman-Kac Formulae for $B$-Continuous Viscosity Solutions

Lukas Wessels

(March 17, 2023)

###### Abstract

We prove the existence of a $B$-continuous viscosity solution for a class of
infinite dimensional semilinear partial differential equations (PDEs) using
probabilistic methods. Our approach also yields a stochastic representation
formula for the solution in terms of a scalar-valued backward stochastic
differential equation. The uniqueness is proved under additional assumptions
using a comparison theorem for viscosity solutions. Our results constitute the
first nonlinear Feynman-Kac formula using the notion of $B$-continuous
viscosity solutions and thus introduces a framework allowing for
generalizations to the case of fully nonlinear PDEs.

Mathematics Subject Classification (2020) — 60H30, 35R15, 60H15

Keywords and phrases — backward stochastic differential equations, partial
differential equations in infinite dimensional spaces, $B$-contnuous viscosity
solutions

Email: •<EMAIL_ADDRESS>
##  1 Introduction

The classical Feynman-Kac formula going back to Richard Feynman [13] and Mark
Kac [22] gives a stochastic representation for the solution of a linear
partial differential equation (PDE) in terms of a path integral over a
diffusion process. This result and generalizations thereof have since found
many applications in various fields such as quantitative finance and
stochastic optimal control, see e.g. [32]. Recent years have shown a rising
interest in control problems associated with stochastic partial differential
equations (SPDEs), see e.g. [11]. Therefore, generalizations of the Feynman-
Kac formula to infinite dimensions are of increasing importance. Another
emerging field that relies on nonlinear generalizations of the Feynman-Kac
formula is numerical analysis for high-dimensional PDEs. The stochastic
representation for solutions of PDEs can be utilized to develop numerical
schemes based on machine learning methods, see e.g. [19, 3, 21]. In this
paper, we prove a generalization of the classical Feynman-Kac formula to
semilinear PDEs in infinite dimensional spaces using the theory of
$B$-continuous viscosity solutions and backward stochastic differential
equations (BSDEs).

For finite initial and terminal times $0\leq t<T<\infty$, respectively, and a
separable Hilbert space $H$, we consider the SPDE

$\begin{cases}\mathrm{d}X^{t,x}_{s}=[AX^{t,x}_{s}+b(s,X^{t,x}_{s})]\mathrm{d}s+\sigma(s,X^{t,x}_{s})\mathrm{d}W_{s},\quad
s\in[t,T]\\\ X^{t,x}_{t}=x\in H,\end{cases}$ (1)

where $A:\mathcal{D}(A)\subset H\to H$ is an unbounded linear operator, and
$b:[0,T]\times H\to H$ and $\sigma:[0,T]\times H\to L_{2}(\Xi,H)$ are the
drift and noise coefficient, respectively. Moreover, $(W_{s})_{s\in[t,T]}$ is
a cylindrical Wiener process on some separable Hilbert space $\Xi$ and defined
on some probability space
$(\Omega,\mathcal{F},(\mathcal{F}_{s})_{s\in[t,T]},\mathbb{P})$ with
$(\mathcal{F}_{s})$ being its natural filtration augmented by all
$\mathbb{P}$-null sets. Furthermore, we consider the infinite dimensional PDE

$\begin{cases}v_{t}(t,x)+\langle Ax+b(t,x),Dv(t,x)\rangle_{H}\\\
\quad+\frac{1}{2}\text{tr}(\sigma^{\ast}(t,x)D^{2}v(t,x)\sigma(t,x))-f(t,x,v(t,x),\sigma^{\ast}Dv(t,x))=0,\quad(t,x)\in[0,T]\times
H\\\ v(T,x)=g(x),\quad x\in H,\end{cases}$ (2)

where $f:[0,T]\times H\times\mathbb{R}\times\Xi\to\mathbb{R}$ and
$g:H\to\mathbb{R}$ are given. In order to introduce the main ideas, let us
assume that equation (2) admits a smooth solution $v$ with
$Dv(t,x)\in\mathcal{D}(A^{\ast})$, where $A^{\ast}$ denotes the adjoint of
$A$. Then, applying Itô’s formula and plugging in equation (2) immediately
yields that

$\begin{cases}Y^{t,x}_{s}=v(s,X^{t,x}_{s})\\\
Z^{t,x}_{s}=\sigma^{\ast}Dv(s,X^{t,x}_{s})\end{cases}$ (3)

solves the BSDE

$\begin{cases}\mathrm{d}Y^{t,x}_{s}=f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})\mathrm{d}s+\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi},\quad s\in[t,T]\\\
Y^{t,x}_{T}=g(X^{t,x}_{T}).\end{cases}$ (4)

In particular, for $s=t$, equations (3) and (4) yield the following stochastic
representation of $v$:

$\begin{split}v(t,x)&=Y^{t,x}_{t}\\\
&=g(X^{t,x}_{T})-\int_{t}^{T}f(s,X^{t,x}_{s},v(s,X^{t,x}_{s}),\sigma^{\ast}Dv(s,X^{t,x}_{s}))\mathrm{d}s\\\
&\quad-\int_{t}^{T}\langle\sigma^{\ast}Dv(s,X^{t,x}_{s}),\mathrm{d}W_{s}\rangle_{\Xi}.\end{split}$
(5)

The preceding discussion is only formal since the PDE (2) in general, does not
admit a smooth solution and therefore, the solution of the BSDE (4) can not be
given in terms of the derivative of $v$ as in (3). However, the stochastic
representation (5) remains valid even if the solution of equation (2) is not
differentiable. The objective of this paper is to prove that the stochastic
representation

$u(t,x):=Y^{t,x}_{t}$ (6)

yields a $B$-continuous viscosity solution for the PDE (2), see Theorem 4.
Furthermore, we prove under additional assumptions that $u$ as given in (6) is
the unique solution of the PDE (2).

######  Remark 1.1.

The representation (6) can be viewed as a nonlinear generalization of the
classical Feynman-Kac formula: If $f(s,x,y,z)=k(s,x)y+l(s,x)$ for some
functions $k,l:[t,T]\times H\to\mathbb{R}$, then the BSDE (4) admits an
explicit solution formula for $Y^{t,x}$ in terms of $X^{t,x}$ and $Z^{t,x}$.
The $Z^{t,x}$-dependency drops out when taking the expectation in (5),
yielding the classical Feynman-Kac formula, see e.g. [23, Chapter 5, Theorem
7.6].

In the finite dimensional case, nonlinear Feynman-Kac formulae using BSDEs go
back to Pardoux and Peng [28]. Following their seminal work, there were many
generalizations of this relationship between PDEs and BSDEs in various
directions such as relaxing the assumptions on the coefficients [1, 24], and
deriving a similar relationship for integral-partial differential equations
and BSDEs involving a foward jump-diffusion process [2], for nonlocal PDEs and
mean-field BSDEs [5], for quasilinear PDEs and forward-backward SDEs [12, 30,
31], for reflected BSDEs and related obstacle problems for PDEs [10] and, more
recently, even for fully nonlinear PDEs and second order BSDEs [7, 33, 34].

In the infinite dimensional case, the first nonlinear Feynman-Kac formula
using BSDEs was obtained by Fuhrman and Tessitore in [15], again followed by
generalizations relaxing the assumptions on the coefficients [4, 17, 25, 27],
and deriving similar relationships for elliptic PDEs and infinite horizon
BSDEs [16, 20, 18], for PDEs on a space of continuous functions and stochastic
delay differential equations [14, 26], and for quasilinear PDEs and foward-
backward SDEs [6].

All of the results in the infinite dimensional case study the PDE (2) in the
framework of mild solutions. In this work, our objective is to prove such a
relationship between BSDEs and PDEs in infinite dimensional spaces usind the
framework of $B$-continuous viscosity solutions. As a byproduct, our
probabilistic approach yields a new method to prove existence of
$B$-continuous viscosity solutions for infinite dimensional PDEs involving
unbounded terms. There are various methods in the existing literature for
proving the existence of solutions, see [11, Chapter 3] and the references
therein. These existing approaches construct solutions even for fully
nonlinear PDEs exploiting the connection with infinite dimensional stochastic
control problems, using finite dimensional approximations, or by extending
Perron’s method to the infinite dimensional case. However, in our setting, we
do not impose any structural assumptions on the nonlinearity $f$ and
therefore, in general one cannot introduce an associated stochastic control
problem to construct a solution to the PDE (2). Furthermore, the method using
finite dimensional approximations imposes the additional assumption that the
unbounded operator $A$ satisfies the strong $B$-condition for a compact
operator $B$ (see Assumption (A2) for a definition of the strong
$B$-condition). Finally, Perron’s method requires an additional coercivity
assumption on $A$.

Finally, let us mention that our work introduces a framework for the extension
of the classical Feynman-Kac formula to the case of infinite dimensional fully
nonlinear PDEs using second order BSDEs in the spirit of the corresponding
results in finite dimensions, see [7, 33, 34]. This problem will be
investigated in future work.

The remainder of the paper is organized as follows: In Section 2, we introduce
the notation and the assumptions we are working with in the following
sections. In Section 3, we prove some results for scalar-valued BSDEs driven
by an infinite dimensional Wiener process. Section 4 contains our main result,
the existence of a $B$-continuous viscosity solution of equation (2) and its
stochastic representation (5). Furthermore, the uniqueness is also discussed
in Section 4.

##  2 Assumptions

We impose the following assumptions on the unbounded operator $A$.

######  Assumption 2.1.

  1. (A1)

Let $A:\mathcal{D}(A)\subset H\to H$ be a linear, densely defined, maximal
dissipative operator.

  2. (A2)

Let $B\in L(H)$ be a strictly positive, self-adjoint operator that satisfies
the strong $B$-condition for $A$, i.e., $A^{\ast}B\in L(H)$ and

$-A^{\ast}B+c_{0}B\geq I$ (7)

for some $c_{0}\geq 0$.

Under Assumption (A1), the operator $A$ is the generator of a
$C_{0}$-semigroup which we denote in the following by $(S(s))_{s\in[t,T]}$.
Furthermore, using the operator $B$, for $\alpha>0$, we define the space
$H_{-\alpha}$ as the completion of the space $H$ with respect to the norm

$\|x\|_{H_{-1}}:=\langle Bx,x\rangle_{H}.$ (8)

We impose the following assumptions on $b$ and $\sigma$.

######  Assumption 2.2.

  1. (B1)

Let $b:[0,T]\times H\to H$ be
$\mathcal{B}([0,T])\otimes\mathcal{B}(H)/\mathcal{B}(H)$-measurable and let
there be a constant $C>0$ such that

$\displaystyle\|b(s,x)-b(s,x^{\prime})\|_{H}$ $\displaystyle\leq
C\|x-x^{\prime}\|_{H}$ (9) $\displaystyle\|b(s,x)\|_{H}$ $\displaystyle\leq
C(1+\|x\|_{H})$ (10)

for all $s\in[0,T]$, $x,x^{\prime}\in H$.

  2. (B2)

Let $\sigma:[0,T]\times H\to L_{2}(\Xi;H)$ be
$\mathcal{B}([0,T])\otimes\mathcal{B}(H)/\mathcal{B}(L_{2}(\Xi,H))$-measurable
and let there be a constant $C>0$ such that

$\displaystyle\|\sigma(s,x)-\sigma(s,x^{\prime})\|_{L_{2}(\Xi,H)}$
$\displaystyle\leq C\|x-x^{\prime}\|_{H_{-1}}$ (11)
$\displaystyle\|\sigma(s,x)\|_{L_{2}(\Xi,H)}$ $\displaystyle\leq
C(1+\|x\|_{H}),$ (12)

for all $s\in[0,T]$, $x,x^{\prime}\in H$.

Finally, we impose the following assumptions on the coefficients of the BSDE.

######  Assumption 2.3.

  1. (C1)

Let $f:[0,T]\times H\times\mathbb{R}\times\Xi\to\mathbb{R}$ be uniformly
continuous on bounded subsets of $(0,T)\times H\times\mathbb{R}\times\Xi$ and
satisfy the following conditions:

     * •

There exists a constant $C>0$ such that

$|f(s,x,y,z)-f(s,x^{\prime},y^{\prime},z^{\prime})|\leq
C(\|x-x^{\prime}\|_{H}+|y-y^{\prime}|+\|z-z^{\prime}\|_{\Xi})$ (13)

for all $s\in[0,T]$, $x,x^{\prime}\in H$, $y,y^{\prime}\in\mathbb{R}$,
$z,z^{\prime}\in\Xi$.

     * •

It holds

$\int_{0}^{T}|f(s,0,0,0)|^{2}\mathrm{d}s<\infty.$ (14)

  2. (C2)

For $g:H\to\mathbb{R}$, let there be a constant $C>0$ such that

$|g(x)-g(x^{\prime})|\leq C\|x-x^{\prime}\|_{H}$ (15)

for all $x,x^{\prime}\in H$.

######  Remark 2.4.

Note that the SPDE (1) as well as the BSDE (4) are well-posed under these
assumptions, see e.g. [9, Theorem 7.2] and [11, Proposition 6.20],
respectively.

##  3 BSDEs

In this section, we prove some basic results for scalar-valued BSDEs with
infinite dimensional driving noise. For the corresponding theory when the
driving noise is finite dimensional, see e.g. [10, 29, 35].

###  3.1 Explicit Solution of Linear BSDE

First, we need the following explicit solution formula for linear BSDEs.
Consider the equation

$\begin{cases}\mathrm{d}Y_{s}=-[a_{s}Y_{s}+b_{s}+\langle
c_{s},Z_{s}\rangle_{\Xi}]\mathrm{d}s+\langle
Z_{s},\mathrm{d}W_{s}\rangle_{\Xi}.\quad s\in[t,T]\\\ Y_{T}=\eta\in
L^{2}(\Omega),\end{cases}$ (16)

where $a,b:[t,T]\times\Omega\to\mathbb{R}$ and $c:[t,T]\times\Omega\to\Xi$ are
progressively measurable processes and $\eta$ is $\mathcal{F}_{T}$-measurable.
Let

$\Gamma_{s}=\exp\left[\int_{t}^{s}a_{r}-\frac{1}{2}\|c_{r}\|_{\Xi}^{2}\mathrm{d}r+\int_{t}^{s}\langle
c_{r},\mathrm{d}W_{r}\rangle_{\Xi}\right].$ (17)

Then

$\displaystyle\mathrm{d}\Gamma_{s}$
$\displaystyle=\Gamma_{s}a_{s}\mathrm{d}s+\Gamma_{s}\langle
c_{s},\mathrm{d}W_{s}\rangle_{\Xi}$ (18)
$\displaystyle\mathrm{d}\Gamma_{s}^{-1}$
$\displaystyle=\Gamma_{s}^{-1}(-a_{s}+\|c_{s}\|_{\Xi}^{2})\mathrm{d}s-\Gamma_{s}^{-1}\langle
c_{s},\mathrm{d}W_{s}\rangle_{\Xi}.$ (19)

By [11, Proposition 6.18], there exists an adapted process $V\in
L^{2}([t,T]\times\Omega;\Xi)$ such that

$\Gamma_{T}\eta+\int_{t}^{T}\Gamma_{s}b_{s}\mathrm{d}s=\mathbb{E}\left[\Gamma_{T}\eta+\int_{t}^{T}\Gamma_{s}b_{s}\mathrm{d}s\right]+\int_{t}^{T}\langle
V_{s},\mathrm{d}W_{s}\rangle_{\Xi}.$ (20)

######  Proposition 3.1.

The solution of the linear BSDE (16) is given by

$\begin{cases}Y_{s}=\Gamma_{s}^{-1}\mathbb{E}\left[\Gamma_{T}\eta+\int_{s}^{T}\Gamma_{r}b_{r}\mathrm{d}r\middle|\mathcal{F}_{s}\right]\\\
Z_{s}=\Gamma_{s}^{-1}V_{s}-c_{s}Y_{s}.\end{cases}$ (21)

###### Proof.

Using (20), we derive from (21) that

$\displaystyle Y_{s}$
$\displaystyle=\Gamma_{s}^{-1}\left(\mathbb{E}\left[\Gamma_{T}\eta+\int_{t}^{T}\Gamma_{r}b_{r}\mathrm{d}r\right]-\int_{t}^{s}\Gamma_{r}b_{r}\mathrm{d}r+\int_{t}^{s}\langle
V_{r},\mathrm{d}W_{r}\rangle_{\Xi}\right)$ (22)
$\displaystyle=\Gamma_{s}^{-1}\left(\mathbb{E}\left[\Gamma_{T}\eta+\int_{t}^{T}\Gamma_{r}b_{r}\mathrm{d}r\right]-\int_{t}^{s}\Gamma_{r}b_{r}\mathrm{d}r+\int_{t}^{s}\langle\Gamma_{r}Z_{r}+\Gamma_{r}c_{r}Y_{r},\mathrm{d}W_{r}\rangle_{\Xi}\right).$
(23)

Applying Itô’s formula yields

$\displaystyle\mathrm{d}Y_{s}$
$\displaystyle=\left[\Gamma_{s}^{-1}\left(-a_{s}\mathrm{d}s+\|c_{s}\|_{\Xi}^{2}\mathrm{d}s\right)-\Gamma_{s}^{-1}\langle
c_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right]\Gamma_{s}Y_{s}$ (24)
$\displaystyle\qquad+\Gamma_{s}^{-1}[\langle\Gamma_{s}Y_{s}c_{s}+\Gamma_{s}Z_{s},\mathrm{d}W_{s}\rangle_{\Xi}-\Gamma_{s}b_{s}\mathrm{d}s]-\Gamma_{s}^{-1}\langle
c_{s},c_{s}\Gamma_{s}Y_{s}+\Gamma_{s}Z_{s}\rangle_{\Xi}\mathrm{d}s$ (25)
$\displaystyle=-[a_{s}Y_{s}+b_{s}+\langle
c_{s},Z_{s}\rangle_{\Xi}]\mathrm{d}s+\langle
Z_{s},\mathrm{d}W_{s}\rangle_{\Xi}.$ (26)

Since $Y_{T}=\eta$, this proves that $(Y,Z)$ as given in (21) solves the BSDE
(16). ∎

###  3.2 Comparison for BSDEs

In the remainder of this section, let
$F:\Omega\times[t,T]\times\mathbb{R}\times\Xi\to\mathbb{R}$ be
$\mathcal{P}\times\mathcal{B}([t,T]\times\mathbb{R}\times\Xi)/\mathcal{B}(\mathbb{R})$-measurable,
where $\mathcal{P}$ denotes the predictable $\sigma$-field on
$\Omega\times[s,T]$ and $\mathcal{B}(\Lambda)$ denotes the Borel
$\sigma$-field of any topological space $\Lambda$. Furthermore, assume that
there exists a constant $C>0$ such that

$|F(s,y,z)-F(s,y^{\prime},z^{\prime})|\leq
C(|y-y^{\prime}|+\|z-z^{\prime}\|_{\Xi})$ (27)

$\mathbb{P}$-almost surely and for every $s\in[t,T]$,
$y,y^{\prime}\in\mathbb{R}$ and $z,z^{\prime}\in\Xi$. Finally, assume that

$\mathbb{E}\left[\int_{t}^{T}|F(s,0,0)|^{2}\mathrm{d}s\right]<\infty.$ (28)

We consider the BSDE

$\begin{cases}\mathrm{d}Y_{s}=F(s,Y_{s},Z_{s})\mathrm{d}s+\langle
Z_{s},\mathrm{d}W_{s}\rangle_{\Xi},\quad s\in[t,T]\\\ Y_{T}=\eta\in
L^{2}(\Omega),\end{cases}$ (29)

where $\eta$ is $\mathcal{F}_{T}$-measurable.

######  Definition 3.2.

A pair $(Y,Z)$ of adapted processes $Y\in L^{2}(\Omega;C([t,T]))$ and $Z\in
L^{2}([t,T]\times\Omega;\Xi)$ is a supersolution of the BSDE (29) if for every
$t\leq s<r\leq T$ it holds

$\begin{cases}Y_{r}\leq
Y_{s}+\int_{s}^{r}F({s^{\prime}},Y_{s^{\prime}},Z_{s^{\prime}})\mathrm{d}{s^{\prime}}+\int_{s}^{r}\langle
Z_{s^{\prime}},\mathrm{d}W_{s^{\prime}}\rangle_{\Xi}\\\
Y_{T}=\eta.\end{cases}$ (30)

A pair $(Y,Z)$ of adapted processes $Y\in L^{2}(\Omega;C([t,T]))$ and $Z\in
L^{2}([t,T]\times\Omega;\Xi)$ is a subsolution of the BSDE (29) if for every
$t\leq s<r\leq T$ it holds

$\begin{cases}Y_{r}\geq
Y_{s}+\int_{s}^{r}F({s^{\prime}},Y_{s^{\prime}},Z_{s^{\prime}})\mathrm{d}{s^{\prime}}+\int_{s}^{r}\langle
Z_{s^{\prime}},\mathrm{d}W_{s^{\prime}}\rangle_{\Xi}\\\
Y_{T}=\eta.\end{cases}$ (31)

Note that this definition is equivalent with the following definition.

######  Definition 3.3.

A triple $(Y,Z,I)$ of adapted processes $Y\in L^{2}(\Omega;C([t,T]))$, $Z\in
L^{2}([t,T]\times\Omega;\Xi)$ and $I\in L^{2}(\Omega;C([t,T]))$ is a
supersolution of the BSDE (29) if $I(t)=0$ $\mathbb{P}$-almost surely, $I$ is
increasing, and for every $s\in[t,T]$ it holds

$Y_{s}=\eta-\int_{s}^{T}F(r,Y_{r},Z_{r})\mathrm{d}r+I_{T}-I_{s}-\int_{s}^{T}\langle
Z_{r},\mathrm{d}W_{r}\rangle_{\Xi}.$ (32)

A triple $(Y,Z,D)$ of adapted processes $Y\in L^{2}(\Omega;C([t,T]))$, $Z\in
L^{2}([t,T]\times\Omega;\Xi)$ and $D\in L^{2}(\Omega;C([t,T]))$ is a
subsolution of the BSDE (29) if $D(t)=0$ $\mathbb{P}$-almost surely, $D$ is
decreasing, and for every $s\in[t,T]$ it holds

$Y_{s}=\eta-\int_{s}^{T}F(r,Y_{r},Z_{r})\mathrm{d}r+D_{T}-D_{s}-\int_{s}^{T}\langle
Z_{r},\mathrm{d}W_{r}\rangle_{\Xi}.$ (33)

######  Remark 3.4.

The fact that (30) and (31) follow from (32) and (33), respectively, is
immediate. For the opposite direction, consider

$I_{s}:=Y_{t}+\int_{t}^{s}F(r,Y_{r},Z_{r})\mathrm{d}r+\int_{t}^{s}\langle
Z_{r},\mathrm{d}W_{r}\rangle_{\Xi}-Y_{s}.$ (34)

In the case of linear BSDEs, we have the following comparison result.

######  Proposition 3.5.

Let $(Y,Z)$ be a supersolution of the linear BSDE (16). Then it holds

$Y_{s}\geq\Gamma_{s}^{-1}\mathbb{E}\left[\Gamma_{T}\eta+\int_{s}^{T}\Gamma_{r}b_{r}\mathrm{d}r\middle|\mathcal{F}_{s}\right],$
(35)

i.e., a supersolution to the linear BSDE (16) dominates the solution to the
linear BSDE.

###### Proof.

Let $(Y,Z,I)$ be a supersolution of the linear BSDE (16). Then an application
of Itô’s formula yields that

$\Gamma_{s}Y_{s}+\int_{t}^{s}\Gamma_{r}b_{r}\mathrm{d}r$ (36)

is a local supermartingale. From this fact, one easily deduces (35). ∎

Next, let us prove a comparison result for sub- and supersolutions.

######  Proposition 3.6.

Let $(F^{i},\eta^{i})$, $i=1,2$, satisfy the same assumptions as $(F,\eta)$
above, and let $(Y^{1},Z^{1})$ and $(Y^{2},Z^{2})$ be a sub- and supersolution
of the BSDE (29) associated with $(F^{1},\eta^{1})$ and $(F^{2},\eta^{2})$,
respectively. Assume

  1. (i)

$\eta^{2}\geq\eta^{1}$

  2. (ii)

$F^{2}(s,Y^{2}_{s},Z^{2}_{s})\leq F^{1}(s,Y^{2}_{s},Z^{2}_{s})$
$\mathrm{d}s\otimes\mathbb{P}$-almost surely.

Then for every $s\in[t,T]$, it holds $Y^{2}_{s}\geq Y^{1}_{s}$
$\mathbb{P}$-almost surely. Moreover, if there exists an $s\in[t,T]$ such that

$\eta^{2}-\eta^{1}+\int_{s}^{T}F^{2}(r,Y^{2}_{r},Z^{2}_{r})-F^{1}(r,Y^{2}_{r},Z^{2}_{r})\mathrm{d}r>0$
(37)

$\mathbb{P}$-almost surely, then $Y^{2}_{s}>Y^{1}_{s}$ $\mathbb{P}$-almost
surely.

###### Proof.

Note that

$Y^{2}_{r}-Y^{1}_{r}\leq
Y^{2}_{s}-Y^{1}_{s}-\int_{s}^{r}a_{s^{\prime}}(Y^{2}_{s^{\prime}}-Y^{1}_{s^{\prime}})+b_{s^{\prime}}+\langle
c_{s^{\prime}},Z^{2}_{s^{\prime}}-Z^{1}_{s^{\prime}}\rangle_{\Xi}\mathrm{d}{s^{\prime}}+\int_{s}^{r}\langle
Z^{2}_{s^{\prime}}-Z^{1}_{s^{\prime}},\mathrm{d}W_{s^{\prime}}\rangle_{\Xi},$
(38)

where

$a_{s^{\prime}}=\frac{1}{Y^{2}_{s^{\prime}}-Y^{1}_{s^{\prime}}}(F^{1}({s^{\prime}},Y^{1}_{s^{\prime}},Z^{1}_{s^{\prime}})-F^{1}({s^{\prime}},Y^{2}_{s^{\prime}},Z^{1}_{s^{\prime}}))\mathbf{1}_{\\{Y^{2}_{s^{\prime}}\neq
Y^{1}_{s^{\prime}}\\}}$ (39)

and

$b_{s^{\prime}}=F^{1}({s^{\prime}},Y^{2}_{s^{\prime}},Z^{2}_{s^{\prime}})-F^{2}({s^{\prime}},Y^{2}_{s^{\prime}},Z^{2}_{s^{\prime}})$
(40)

and

$c_{s^{\prime}}=\frac{Z^{2}_{s^{\prime}}-Z^{1}_{s^{\prime}}}{\|Z^{2}_{s^{\prime}}-Z^{1}_{s^{\prime}}\|_{\Xi}^{2}}(F^{1}({s^{\prime}},Y^{2}_{s^{\prime}},Z^{1}_{s^{\prime}})-F^{1}({s^{\prime}},Y^{2}_{s^{\prime}},Z^{2}_{s^{\prime}}))\mathbf{1}_{\\{Z^{2}_{s^{\prime}}\neq
Z^{1}_{s^{\prime}}\\}}.$ (41)

Therefore, we obtain from Proposition 3.5 and assumptions $(i)$ and $(ii)$

$Y^{2}_{s}-Y^{1}_{s}\geq\Gamma_{s}^{-1}\mathbb{E}\left[\Gamma_{T}(\eta^{2}-\eta^{1})+\int_{s}^{T}\Gamma_{r}b_{r}\mathrm{d}r\middle|\mathcal{F}_{s}\right]\geq
0,$ (42)

where $\Gamma$ is given by equation (17). The strict inequality follows from
(37). ∎

###  3.3 A Piori Estimates for BSDEs

######  Lemma 3.7.

Let $(Y,Z)$ be the solution of the BSDE (29). Then it holds

$\mathbb{E}\left[\sup_{s\in[t,T]}|Y_{s}|^{2}+\int_{t}^{T}\|Z_{s}\|_{\Xi}^{2}\mathrm{d}s\right]\leq
C\mathbb{E}\left[|\eta|^{2}+\left(\int_{t}^{T}|F(s,0,0)|\mathrm{d}s\right)^{2}\right].$
(43)

The proof of this result can be found in [15, Proposition 4.3].

######  Lemma 3.8.

Let $(F^{i},\eta^{i})$, $i=1,2$, satisfy the same assumptions as $(F,\eta)$
above and let $(Y^{1},Z^{1})$ and $(Y^{2},Z^{2})$ be the solution of the BSDE
(29) associated with $(F^{1},\eta^{1})$ and $(F^{2},\eta^{2})$, respectively.
Then it holds

$\begin{split}&\mathbb{E}\left[\sup_{s\in[t,T]}\left|Y^{1}_{s}-Y^{2}_{s}\right|^{2}+\int_{t}^{T}\left\|Z^{1}_{s}-Z^{2}_{s}\right\|_{\Xi}^{2}\mathrm{d}s\right]\\\
&\leq
C\mathbb{E}\left[\left|\eta^{1}-\eta^{2}\right|^{2}+\left(\int_{t}^{T}|F^{1}(r,Y^{1}_{r},Z^{1}_{r})-F^{2}(r,Y^{1}_{r},Z^{1}_{r})|\mathrm{d}r\right)^{2}\right].\end{split}$
(44)

###### Proof.

Note that for all $s\in[t,T]$ it holds

$\begin{split}Y^{1}_{s}-Y^{2}_{s}&=\eta^{1}-\eta^{2}-\int_{s}^{T}F^{1}(r,Y^{1}_{r},Z^{1}_{r})-F^{2}(r,Y^{2}_{r},Z^{2}_{r})\mathrm{d}r-\int_{s}^{T}\langle
Z^{1}_{r}-Z^{2}_{r},\mathrm{d}W_{r}\rangle_{\Xi}\\\
&=\eta^{1}-\eta^{2}-\int_{s}^{T}F^{1}(r,Y^{1}_{r},Z^{1}_{r})-F^{2}(r,Y^{1}_{r},Z^{1}_{r})\mathrm{d}r\\\
&\qquad-\int_{s}^{T}a_{r}(Y^{1}_{r}-Y^{2}_{r})+b_{r}(Z^{1}_{r}-Z^{2}_{r})\mathrm{d}r-\int_{s}^{T}\langle
Z^{1}_{r}-Z^{2}_{r},\mathrm{d}W_{r}\rangle_{\Xi},\end{split}$ (45)

where

$a_{r}:=\frac{F^{2}(r,Y^{1}_{r},Z^{1}_{r})-F^{2}(r,Y^{2}_{r},Z^{1}_{r})}{Y^{1}_{r}-Y^{2}_{r}}\mathbf{1}_{\\{Y^{1}_{r}\neq
Y^{2}_{r}\\}}$ (46)

and

$b_{r}:=\frac{F^{2}(r,Y^{2}_{r},Z^{1}_{r})-F^{2}(r,Y^{2}_{r},Z^{2}_{r})}{Z^{1}_{r}-Z^{2}_{r}}\mathbf{1}_{\\{Z^{1}_{r}\neq
Z^{2}_{r}\\}}.$ (47)

Now, Lemma 3.7 yields the result. ∎

##  4 Feynman-Kac Formulae for $B$-Continuous Viscosity Solutions

In this section, we prove that the stochastic representation formula (6)
yields a $B$-continuous viscosity solution of the PDE (2). Under additional
assumptions, we establish uniqueness.

###  4.1 $B$-Continuous Viscosity Solutions

First, we recall the notion of $B$-continuous viscosity solutions. Let us
begin by introducing the class of test functions we are using, see [11,
Definition 3.32].

######  Definition 4.1.

A function $\psi$ is a test function if $\psi=\varphi+h(t,\|x\|_{H})$, where:

  1. (i)

$\varphi\in C^{1,2}((0,T)\times H)$ is locally bounded, and is such that
$\varphi$ is $B$-lower semicontinuous, and $\varphi_{t}$, $A^{\ast}D\varphi$,
$D\varphi$, $D^{2}\varphi$ are uniformly continuous on $(0,T)\times H$.

  2. (ii)

$h\in C^{1,2}((0,T)\times\mathbb{R})$ and is such that for every $t\in(0,T)$,
$h(t,\cdot)$ is even and $h(t,\cdot)$ is non-decreasing on $[0,+\infty)$.

Now, we define the notion of $B$-continuous viscosity solution, see [11,
Definition 3.35].

######  Definition 4.2.

A locally bounded and upper semicontinuous function $u$ on $(0,T]\times H$
which is $B$-upper semicontinuous on $(0,T)\times H$ is a viscosity
subsolution of (2) if the following holds: whenever $u-\psi$ has a local
maximum at a point $(t,x)\in(0,T)\times H$ for a test function $\psi$ as in
Definition 4.1 then

$\psi_{t}(t,x)+\langle x,A^{\ast}D\varphi(t,x)\rangle_{H}+\langle
b(t,x),D\psi(t,x)\rangle_{H}\\\
+\frac{1}{2}\text{tr}(\sigma^{\ast}(t,x)D^{2}\psi(t,x)\sigma(t,x))-f(t,x,u(t,x),\sigma^{\ast}(t,x)D\psi(t,x))\geq
0.$ (48)

A locally bounded and lower semicontinuous function $u$ on $(0,T]\times H$
which is $B$-lower semicontinuous on $(0,T)\times H$ is a viscosity
supersolution of (2) if the following holds: whenever $u+\psi$ has a local
minimum at a point $(t,x)\in(0,T)\times H$ for a test function $\psi$ as in
Definition 4.1 then

$-\psi_{t}(t,x)-\langle x,A^{\ast}D\varphi(t,x)\rangle_{H}-\langle
b(t,x),D\psi(t,x)\rangle_{H}\\\
-\frac{1}{2}\text{tr}(\sigma^{\ast}(t,x)D^{2}\psi(t,x)\sigma(t,x))-f(t,x,u(t,x),-\sigma^{\ast}(t,x)D\psi(t,x))\leq
0.$ (49)

A viscosity solution of (2) is a function which is both a viscosity
subsolution and a viscosity supersolution of (2).

###  4.2 Existence and Stochastic Representation

First, we recall the following a priori estimates for the forward SPDE (1),
see [11, Lemma 3.23] and [15, Proposition 3.3], respectively.

######  Lemma 4.3.

Let Assumptions 2.1 and 2.2 be satisfied. Then, there exists a constant $C>0$
independent of $t\in[0,T]$ such that

$\mathbb{E}\left[\|X^{t,x}_{T}-X^{t,x^{\prime}}_{T}\|_{H_{-1}}^{2}+\int_{t}^{T}\|X^{t,x}_{s}-X^{t,x^{\prime}}_{s}\|_{H}^{2}\mathrm{d}s\right]\leq
C\|x-x^{\prime}\|_{H_{-1}}^{2}$ (50)

for all $x,x^{\prime}\in H$. Furthermore, for every $t\in[0,T)$, there exists
a constant $C>0$ such that

$\mathbb{E}\left[\|X^{t,x}_{T}-X^{t,x^{\prime}}_{T}\|_{H}^{2}\right]\leq\frac{C}{T-t}\|x-x^{\prime}\|_{H_{-1}}^{2}$
(51)

for all $x,x^{\prime}\in H$.

######  Lemma 4.4.

Let Assumptions 2.1 and 2.2 be satisfied. Then, for every $x\in H$, it holds

$\mathbb{E}\left[\sup_{s\in[t^{\prime},T]}\|X^{t^{\prime},x}_{s}-X^{t,x}_{s}\|_{H}^{2}\right]\to
0$ (52)

as $t^{\prime}\downarrow t$.

######  Remark 4.5.

In [15, Proposition 3.3], the authors actually prove a stronger result than
this. However, we only need the stated result.

Now, let us state and prove our main result.

######  Theorem 4.6.

Let $(Y,Z)$ be the solution of the BSDE (4), and define $u(t,x):=Y^{t,x}_{t}$.
Then $u$ is a $B$-continuous viscosity solution of the PDE (2) in the sense of
Definition 4.2.

###### Proof.

In the first step, we are going to prove the $B$-continuity of $u$, and in the
second step, we are going to prove the viscosity property.

_Step 1_ : Let $(t,x)\in(0,T)\times H$. Let $t_{n}\to t$,
$x_{n}\rightharpoonup x$ and $Bx_{n}\to Bx$ and fix $N\in\mathbb{N}$ such that
$|t_{n}-t|<(T-t)/2$ for all $n\geq N$. We have

$|u(t_{n},x_{n})-u(t,x)|=\left|Y^{t_{n},x_{n}}_{t_{n}}-Y^{t,x}_{t}\right|\leq\left|Y^{t_{n},x_{n}}_{t_{n}}-Y^{t_{n},x}_{t_{n}}\right|+\left|Y^{t_{n},x}_{t_{n}}-Y^{t,x}_{t}\right|.$
(53)

For the first term, choosing

$\begin{cases}F^{1}(s,y,z):=f(s,X^{t_{n},x_{n}}_{s},y,z)\\\
\eta^{1}:=g(X^{t_{n},x_{n}}_{T})\end{cases}$ (54)

and

$\begin{cases}F^{2}(s,y,z):=f(s,X^{t_{n},x}_{s},y,z)\\\
\eta^{2}:=g(X^{t_{n},x}_{T})\end{cases}$ (55)

and applying Lemma 3.8 yields

$\begin{split}&\left|Y^{t_{n},x_{n}}_{t_{n}}-Y^{t_{n},x}_{t_{n}}\right|^{2}\\\
&\leq
C\mathbb{E}\Bigg{[}\left|g(X^{t_{n},x_{n}}_{T})-g(X^{t_{n},x}_{T})\right|^{2}\\\
&\qquad\qquad+\left(\int_{t_{n}}^{T}|f(s,X^{t_{n},x_{n}}_{s},Y^{t_{n},x_{n}}_{s},Z^{t_{n},x_{n}}_{s})-f(s,X^{t_{n},x}_{s},Y^{t_{n},x_{n}}_{s},Z^{t_{n},x_{n}}_{s})|\mathrm{d}s\right)^{2}\Bigg{]}.\end{split}$
(56)

Using Assumption 2.3, we obtain

$\left|Y^{t_{n},x_{n}}_{t_{n}}-Y^{t_{n},x}_{t_{n}}\right|^{2}\leq
C\mathbb{E}\left[\|X^{t_{n},x_{n}}_{T}-X^{t_{n},x}_{T}\|_{H}^{2}+\int_{t_{n}}^{T}\|X^{t_{n},x_{n}}_{s}-X^{t_{n},x}_{s}\|_{H}^{2}\mathrm{d}s\right].$
(57)

Applying Lemma 4.3 yields

$\left|Y^{t_{n},x_{n}}_{t_{n}}-Y^{t_{n},x}_{t_{n}}\right|^{2}\leq
C\|x_{n}-x\|_{H_{-1}}^{2}$ (58)

for some constant that is uniform in $n\geq N$.

Now let us turn to the second term in (53). Assume without loss of generality
that $t_{n}\downarrow t$. Then we have

$\begin{split}&\left|Y^{t_{n},x}_{t_{n}}-Y^{t,x}_{t}\right|^{2}\\\ &\leq
C\mathbb{E}\Bigg{[}\left|g(X^{t_{n},x}_{T})-g(X^{t,x}_{T})\right|^{2}+\int_{t_{n}}^{T}|f(s,X^{t_{n},x}_{s},Y^{t_{n},x}_{s},Z^{t_{n},x}_{s})-f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s\\\
&\quad+\left|\int_{t_{n}}^{T}\langle
Z^{t_{n},x}_{s}-Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}+\int_{t}^{t_{n}}|f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s+\left|\int_{t}^{t_{n}}\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}\Bigg{]}.\end{split}$ (59)

The fourth and the fifth term tend to zero as $t_{n}\downarrow t$ due to the
continuity of the integral. For the first term, we have

$\left|g(X^{t_{n},x}_{T})-g(X^{t,x}_{T})\right|^{2}\leq
C\|X^{t_{n},x}_{T}-X^{t,x}_{T}\|^{2}_{H}.$ (60)

For the second term we have by Assumption (C1)

$\begin{split}&\int_{t_{n}}^{T}|f(s,X^{t_{n},x}_{s},Y^{t_{n},x}_{s},Z^{t_{n},x}_{s})-f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s\\\
&\leq
C\int_{t_{n}}^{T}\|X^{t_{n},x}_{s}-X^{t,x}_{s}\|^{2}_{H}+|Y^{t_{n},x}_{s}-Y^{t,x}_{s}|^{2}+\|Z^{t_{n},x}_{s}-Z^{t,x}_{s}\|^{2}_{\Xi}\mathrm{d}s.\end{split}$
(61)

Therefore, applying Lemma 3.8 and using (60), we obtain

$\begin{split}&\mathbb{E}\left[\int_{t_{n}}^{T}|f(s,X^{t_{n},x}_{s},Y^{t_{n},x}_{s},Z^{t_{n},x}_{s})-f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s\right]\\\
&\leq
C\mathbb{E}\left[\sup_{s\in[t_{n},T]}\|X^{t_{n},x}_{s}-X^{t,x}_{s}\|_{H}^{2}\right].\end{split}$
(62)

For the third term, we have by Lemma 3.8

$\begin{split}\mathbb{E}\left[\left|\int_{t_{n}}^{T}\langle
Z^{t_{n},x}_{s}-Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}\right]&=\mathbb{E}\left[\int_{t_{n}}^{T}\|Z^{t_{n},x}_{s}-Z^{t,x}_{s}\|_{\Xi}^{2}\mathrm{d}s\right]\\\
&\leq
C\mathbb{E}\left[\sup_{s\in[t_{n},T]}\|X^{t_{n},x}_{s}-X^{t,x}_{s}\|_{H}^{2}\right].\end{split}$
(63)

Applying Lemma 4.4 yields that the right-hand side of (59) tends to zero as
$t_{n}\downarrow t$. Together with (58), this concludes the proof of the
$B$-continuity.

_Step 2_ : First, we show that $u$ is a viscosity subsolution. Let
$\psi=\varphi+h$ be a test function as in Definition 4.1 such that

$0=(u-\psi)(t,x)=\max_{\begin{subarray}{c}s\in[t,T]\\\ x^{\prime}\in
H\end{subarray}}(u-\psi)(s,x^{\prime})$ (64)

We can assume without loss of generaly that the maximum is strict, see [11,
Lemma 3.37]. Assume, for sake of contradiction, that

$\psi_{t}(t,x)+\langle x,A^{\ast}D\varphi(t,x)\rangle_{H}+\langle
b(t,x),D\psi(t,x)\rangle_{H}\\\
+\text{tr}(\sigma^{\ast}(t,x)D^{2}\psi(t,x)\sigma(t,x))-f(t,x,u(t,x),\sigma^{\ast}D\psi(t,x))<0.$
(65)

Fix $\varepsilon>0$ such that for all $t\leq s\leq t+\varepsilon$ and
$\|x^{\prime}-x\|_{H}\leq\varepsilon$, it holds

$\psi_{t}(s,x^{\prime})+\langle
x^{\prime},A^{\ast}D\varphi(s,x^{\prime})\rangle_{H}+\langle
b(s,x^{\prime}),D\psi(s,x^{\prime})\rangle_{H}\\\
+\text{tr}(\sigma^{\ast}(s,x^{\prime})D^{2}\psi(s,x^{\prime})\sigma(s,x^{\prime}))-f(s,x^{\prime},u(s,x^{\prime}),\sigma^{\ast}D\psi(s,x^{\prime}))<0.$
(66)

Now, we define the stopping time

$\tau:=\inf\\{s\geq
t|\|X^{t,x}_{s}-x\|_{H}\geq\varepsilon\\}\wedge(t+\varepsilon).$ (67)

Let $(Y^{t,x},Z^{t,x})$ be the solution of the BSDE (4), and define for $t\leq
s\leq t+\varepsilon$

$(Y^{1}_{s},Z^{1}_{s})=(Y^{t,x}_{s\wedge\tau},\mathbf{1}_{[0,\tau]}(s)Z^{t,x}_{s}).$
(68)

Then we have

$u(t+h,X^{t,x}_{t+h})=Y^{t+h,X^{t,x}_{t+h}}_{t+h}=Y^{t,x}_{t+h}$ (69)

due to the uniqueness of the solution of the SPDE (1) and the uniqueness of
the solution of the BSDE (4). Therefore, for every $r\in[t,t+\varepsilon]$

$Y^{t,x}_{r\wedge\tau}=u(r\wedge\tau,X^{t,x}_{r\wedge\tau}),$ (70)

which shows that $(Y^{1},Z^{1})$ solves the BSDE

$\begin{cases}\mathrm{d}Y^{1}_{s}=\mathbf{1}_{[0,\tau]}(s)f(s,X^{t,x}_{s},u(s,X^{t,x}_{s}),Z^{1}_{s})\mathrm{d}s+\langle
Z^{1}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\\\
Y^{1}_{\tau}=u(\tau,X^{t,x}_{\tau}).\end{cases}$ (71)

On the other hand, we obtain from Itô’s formula for test functions (see [11,
Proposition 1.166]) that for every $r,s\in[t,t+\varepsilon]$, $r\geq s$,

$\begin{split}&\psi(r\wedge\tau,X^{t,x}_{r\wedge\tau})\\\
&\leq\psi(s\wedge\tau,X^{t,x}_{s\wedge\tau})+\int_{s}^{r}\mathbf{1}_{[0,\tau]}(s^{\prime})\left(\psi_{t}(s^{\prime},X^{t,x}_{s^{\prime}})+\langle
X^{t,x}_{s^{\prime}},A^{\ast}D\varphi(s^{\prime},X^{t,x}_{s^{\prime}})\rangle_{H}\right)\mathrm{d}s^{\prime}\\\
&\quad+\int_{s}^{r}\mathbf{1}_{[0,\tau]}(s^{\prime})\langle
b(s^{\prime},X^{t,x}_{s^{\prime}}),D\psi(s^{\prime},X^{t,x}_{s^{\prime}})\rangle_{H}\mathrm{d}s^{\prime}\\\
&\quad+\frac{1}{2}\int_{s}^{r}\mathbf{1}_{[0,\tau]}(s^{\prime})\text{tr}\left(\sigma^{\ast}(s^{\prime},X^{t,x}_{s^{\prime}})D^{2}\psi(s^{\prime},X^{t,x}_{s^{\prime}})\sigma(s^{\prime},X^{t,x}_{s^{\prime}})\right)\mathrm{d}s^{\prime}\\\
&\quad+\int_{s}^{r}\langle\mathbf{1}_{[0,\tau]}(s^{\prime})\sigma^{\ast}(s^{\prime},X^{t,x}_{s^{\prime}})D\psi(s^{\prime},X^{t,x}_{s^{\prime}}),\mathrm{d}W_{s^{\prime}}\rangle_{\Xi}.\end{split}$
(72)

Therefore,

$(Y^{2}_{s},Z^{2}_{s}):=(\psi(s\wedge\tau,X^{t,x}_{s\wedge\tau}),\mathbf{1}_{[0,\tau]}(s)\sigma^{\ast}D\psi(s,X^{t,x}_{s}))$
(73)

is a supersolution of the BSDE

$\begin{cases}\mathrm{d}Y_{s}=\mathbf{1}_{[0,\tau]}(s)\Bigl{[}\psi_{t}(s,X^{t,x}_{s})+\langle
X^{t,x}_{s},A^{\ast}D\varphi(s,X^{t,x}_{s})\rangle_{H}+\langle
b(s,X^{t,x}_{s}),D\psi(s,X^{t,x}_{s})\rangle_{H}\\\
\qquad\qquad\qquad\qquad\qquad+\frac{1}{2}\text{tr}(\sigma^{\ast}(s,X^{t,x}_{s})D^{2}\psi(s,X^{t,x}_{s})\sigma(s,X^{t,x}_{s}))\Bigr{]}\mathrm{d}s+\langle
Z_{s},\mathrm{d}W_{s}\rangle_{\Xi}\\\
Y_{\tau}=\psi(\tau,X^{t,x}_{\tau}).\end{cases}$ (74)

Due to inequality (66) and the definition of $\tau$, we can apply Proposition
3.6 to obtain

$\psi(t,x)=Y^{2}_{t}>Y^{1}_{t}=Y^{t,x}_{t}=u(t,x)$ (75)

which contradicts (64) and therefore concludes the proof that $u$ is a
$B$-continuous viscosity subsolution. The proof that $u$ is a $B$-continuous
viscosity supersolution is similar. ∎

###  4.3 Uniqueness

In order to prove uniqueness, we rely on a comparison theorem for the PDE (2)
which is proved using analytic methods. For this result to hold, we have to
impose additional assumptions on the coefficients of the PDE.

######  Assumption 4.7.

  1. (C1)

Let $\sigma$ be uniformly continuous on bounded subsets of $(0,T)\times H$.

  2. (C2)

For $y^{\prime}\geq y$, let

$f(s,x,y^{\prime},z)-f(s,x,y,z)\geq 0$ (76)

for all $(s,x,z)\in(0,T)\times H\times\Xi$.

######  Theorem 4.8.

Let Assumptions 2.1, 2.2, 2.3, 4.7 be satisfied. Then $u(t,x):=Y^{t,x}_{t}$ is
the unique $B$-continuous viscosity solutions satisfying

$\lim_{t\to T}|u(t,x)-g(S(T-t)x)|=0$ (77)

uniformly on bounded subsets of $H$ and

$|u(t,x)|\leq
C_{1}\exp\left(C_{2}\left(\ln\left(1+\|x\|_{H}\right)\right)^{2}\right)$ (78)

for some constants $C_{1},C_{2}\geq 0$ and all $(t,x)\in(0,T)\times H$.

###### Proof.

Once we know that $u$ satisfies (77) and (78), the result follows from [11,
Theorem 3.54]. Let us first prove (77). Note that

$\begin{split}&|u(t,x)-g(S(T-t)x)|^{2}\\\
&=\mathbb{E}\left[|Y^{t,x}_{t}-g(S(T-t)x)|^{2}\right]\\\
&=\mathbb{E}\left[\left|g(X^{t,x}_{T})-\int_{t}^{T}f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})\mathrm{d}s-\int_{t}^{T}\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}-g(S(T-t)x)\right|^{2}\right]\\\
&\leq\mathbb{E}\left[\int_{t}^{T}|f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s+\left|\int_{t}^{T}\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}\right]\\\
&\quad+\mathbb{E}\left[|g(X^{t,x}_{T})-g(S(T-t)x)|^{2}\right].\end{split}$
(79)

For the first term, using Itô’s isometry and Assumption 2.3, we have

$\begin{split}&\mathbb{E}\left[\int_{t}^{T}|f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s+\left|\int_{t}^{T}\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}\right]\\\ &\leq
C\int_{t}^{T}|f(s,0,0,0)|^{2}+\mathbb{E}\left[\|X^{t,x}_{s}\|_{H}^{2}+|Y^{t,x}_{s}|^{2}+\|Z^{t,x}_{s}\|_{\Xi}^{2}\right]\mathrm{d}s.\end{split}$
(80)

Applying Lemma 3.7, we obtain

$\begin{split}&\mathbb{E}\left[\int_{t}^{T}|f(s,X^{t,x}_{s},Y^{t,x}_{s},Z^{t,x}_{s})|^{2}\mathrm{d}s+\left|\int_{t}^{T}\langle
Z^{t,x}_{s},\mathrm{d}W_{s}\rangle_{\Xi}\right|^{2}\right]\\\ &\leq
C\int_{t}^{T}|f(s,0,0,0)|^{2}+1+\mathbb{E}\left[\|X^{t,x}_{s}\|_{H}^{2}\right]\mathrm{d}s.\end{split}$
(81)

For the last term in (79), using Assumption (C2), we have

$\begin{split}&\mathbb{E}\left[|g(X^{t,x}_{T})-g(S(T-t)x)|^{2}\right]\\\ &\leq
C\mathbb{E}\left[\|X^{t,x}_{T}-S(T-t)x\|_{H}^{2}\right]\\\
&=C\mathbb{E}\left[\left\|\int_{t}^{T}S(T-s)b(s,X^{t,x}_{s})\mathrm{d}s+\int_{t}^{T}S(T-s)\sigma(s,X^{t,x}_{s})\mathrm{d}W_{s}\right\|_{H}^{2}\right]\\\
&\leq
C\mathbb{E}\left[\int_{t}^{T}\|S(T-s)b(s,X^{t,x}_{s})\|_{H}^{2}\mathrm{d}s+\left\|\int_{t}^{T}S(T-s)\sigma(s,X^{t,x}_{s})\mathrm{d}W_{s}\right\|_{H}^{2}\right].\end{split}$
(82)

Using again Itô’s isometry and the linear growth assumption on $b$ and
$\sigma$, we obtain

$\mathbb{E}\left[|g(X^{t,x}_{T})-g(S(T-t)x)|^{2}\right]\leq
C\int_{t}^{T}1+\mathbb{E}\left[\|X^{t,x}_{s}\|_{H}^{2}\right]\mathrm{d}s.$
(83)

Therefore, together with (81), we have

$|u(t,x)-g(S(T-t)x)|^{2}\leq
C\int_{t}^{T}|f(s,0,0,0)|^{2}+1+\mathbb{E}\left[\|X^{t,x}_{s}\|_{H}^{2}\right]\mathrm{d}s.$
(84)

Due to Assumption (C1) and well-known a priori estimates for the solution of
the SPDE (1) (see e.g. [9, Theorem 7.2]), the right-hand side converges to
zero as $t\to T$ uniformly on bounded subsets of $H$. The growth condition
(78) follows from Lemma 3.7 and the same a priori estimates for the solution
of the SPDE (1). ∎

## Acknowledgments

The author would like to thank Andrzej Święch for his constructive criticism
of the manuscript. During the preparation of this work, the author was
supported by a fellowship of the German Academic Exchange Service (DAAD).

## References

  * [1] V. Bally, E. Pardoux, L. Stoica,  _Backward stochastic differential equations associated to a symmetric Markov process_ ,  Potential Anal., 22 (2005), 17–60.
  * [2] G. Barles, R. Buckdahn, E. Pardoux,  _Backward stochastic differential equations and integral-partial differential equations_ ,  Stochastics Stochastics Rep., 60 (1997), 57–83.
  * [3] C. Beck, W. E, A. Jentzen,  _Machine learning approximation algorithms for high-dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations_ ,  J. Nonlinear Sci., 29 (2019), 1563–1619.
  * [4] P. Briand, F. Confortola,  _BSDEs with stochastic Lipschitz condition and quadratic PDEs in Hilbert spaces_ ,  Stochastic Process. Appl., 118 (2008), 818–838.
  * [5] R. Buckdahn, J. Li, S. Peng,  _Mean-field backward stochastic differential equations and related partial differential equations_ ,  Stochastic Process. Appl., 119 (2009), 3133–3154.
  * [6] S. Cerrai, G. Guatteri, G. Tessitore,  _Nonlinear random perturbations of PDEs and quasi-linear equations in Hilbert spaces depending on a small parameter_ ,  preprint, https://arxiv.org/abs/2208.12867.
  * [7] P. Cheridito, H. M. Soner, N. Touzi, N. Victoir,  _Second-order backward stochastic differential equations and fully nonlinear parabolic PDEs_ ,  Comm. Pure Appl. Math., 60 (2007), 1081–1110.
  * [8] N. El Karoui, S. Peng, M. C. Quenez,  _Backward Stochastic Differential Equations in Finance_ ,  _Math. Finance_ , 7 (1997), 1–71.
  * [9] G. Da Prato, J. Zabczyk,  _Stochastic equations in infinite dimensions_ ,  second ed., Cambridge University Press, Cambridge, 2014.
  * [10] N. El Karoui, C. Kapoudjian, E. Pardoux, S. Peng, M. C. Quenez,  _Reflected solutions of backward SDE’s, and related obstacle problems for PDE’s_ ,  Ann. Probab., 25 (1997), 702–737.
  * [11] G. Fabbri, F. Gozzi, A. Święch,  _Stochastic optimal control in infinite dimension_ ,  Springer, Cham, 2017.
  * [12] C. Feng, X. Wang, H. Zhao,  _Quasi-linear PDEs and forward–backward stochastic differential equations: weak solutions_ ,  J. Differential Equations, 264 (2018), 959–1018.
  * [13] R. P. Feynman,  _Space-time approach to non-relativistic quantum mechanics_ ,  Rev. Modern Physics, 20 (1948), 367–387.
  * [14] M. Fuhrman, F. Masiero, G. Tessitore,  _Stochastic equations with delay: optimal control via BSDEs and regular solutions of Hamilton-Jacobi-Bellman equations_ ,  SIAM J. Control Optim., 48 (2010), 4624–4651.
  * [15] M. Fuhrman, G. Tessitore,  _Nonlinear Kolmogorov equations in infinite dimensional spaces: the backward stochastic differential equations approach and applications to optimal control_ ,  Ann. Probab., 30 (2002), 1397–1465.
  * [16] M. Fuhrman, G. Tessitore,  _Infinite horizon backward stochastic differential equations and elliptic equations in Hilbert spaces_ ,  Ann. Probab., 32 (2004), 607–660.
  * [17] M. Fuhrman, G. Tessitore,  _Generalized directional gradients, backward stochastic differential equations and mild solutions of semilinear parabolic equations_ ,  App. Math. Optim., 51 (2005), 279–332.
  * [18] G. Guatteri, G. Tessitore,  _Ergodic BSDEs with multiplicative and degenerate noise_ ,  SIAM J. Control Optim., 58 (2020), 2050–2077.
  * [19] J. Han, A. Jentzen, W. E,  _Solving high-dimensional partial differential equations using deep learning_ ,  Proc. Natl. Acad. Sci. USA, 115 (2018), 8505–8510.
  * [20] Y. Hu, P.-Y. Madec, A. Richou,  _A probabilistic approach to large time behavior of mild solutions of HJB equations in infinite dimension_ ,  SIAM J. Control Optim., 53 (2015), 378–398.
  * [21] C. Huré, H. Pham, X. Warin,  _Deep backward schemes for high-dimensional nonlinear PDEs_ ,  Math. Comp., 89 (2020), 1547–1579.
  * [22] M. Kac,  _On distributions of certain Wiener functionals_ ,  Trans. Amer. Math. Soc., 65 (1949), 1–13.
  * [23] I. Karatzas, S. E. Shreve,  _Brownian motion and stochastic calculus_ ,  second ed., Springer, New York, 1991.
  * [24] M. Kobylanski,  _Backward stochastic differential equations and partial differential equations with quadratic growth_ ,  Ann. Probab., 28 (2000), 558–602.
  * [25] F. Masiero,  _Stochastic optimal control problems and parabolic equations in Banach spaces_ ,  SIAM J. Control Optim., 47 (2008), 251–300.
  * [26] F. Masiero, C. Orrieri, G. Tessitore, G. Zanco,  _Semilinear Kolmogorov equations on the space of continuous functions via BSDEs_ ,  Stochastic Process. Appl., 136 (2021), 1–56.
  * [27] F. Masiero, A. Richou,  _HJB equations in infinite dimensions with locally Lipschitz Hamiltonian and unbounded terminal condition_ ,  J. Differential Equations, 257 (2014), 1989–2034.
  * [28] E. Pardoux, S. Peng,  _Backward stochastic differential equations and quasilinear parabolic partial differential equations_ ,  in: Stochastic partial differential equations and their applications (Charlotte, NC, 1991), Springer, Berlin, 1992, pp. 200–217.
  * [29] E. Pardoux, A. Răşcanu,  _Stochastic differential equations, backward SDEs, partial differential equations_ ,  Springer, Cham, 2014.
  * [30] E. Pardoux, S. Tang,  _Forward-backward stochastic differential equations and quasilinear parabolic PDEs_ ,  Probab. Theory Related Fields, 114 (1999), 123–150.
  * [31] S. Peng, Z. Wu,  _Fully coupled forward-backward stochastic differential equations and applications to optimal control_ ,  SIAM J. Control Optim., 37 (1999), 825–843.
  * [32] H. Pham,  _Continuous-time stochastic control and optimization with financial applications_ ,  Springer, Berlin, 2009.
  * [33] H. M. Soner, N. Touzi, J. Zhang,  _Wellposedness of second order backward SDEs_ ,  Probab. Theory Related Fields, 153 (2012), 149–190.
  * [34] H. M. Soner, N. Touzi, J. Zhang,  _Dual formulation of second order target problems_ ,  Ann. Appl. Probab., 23 (2013), 308–347.
  * [35] J. Zhang,  _Backward stochastic differential equations_ ,  Springer, New York, 2017.


# Flow Annealed Kalman Inversion for Gradient-Free Inference in Bayesian
Inverse Problems

Richard D.P. Grumitt 1 1 Department of Astronomy, Tsinghua University, Beijing
100084, China
2 Berkeley Center for Cosmological Physics and Department of Physics,
University of California, Berkeley, CA 94720
3 Physics Department, Lawrence Berkeley National Laboratory, Cyclotron Rd,
Berkeley, CA 94720
<EMAIL_ADDRESS>Minas Karamanis 2,3 1 Department of Astronomy,
Tsinghua University, Beijing 100084, China
2 Berkeley Center for Cosmological Physics and Department of Physics,
University of California, Berkeley, CA 94720
3 Physics Department, Lawrence Berkeley National Laboratory, Cyclotron Rd,
Berkeley, CA 94720
<EMAIL_ADDRESS>Uroš Seljak 2,3 1 Department of Astronomy,
Tsinghua University, Beijing 100084, China
2 Berkeley Center for Cosmological Physics and Department of Physics,
University of California, Berkeley, CA 94720
3 Physics Department, Lawrence Berkeley National Laboratory, Cyclotron Rd,
Berkeley, CA 94720
<EMAIL_ADDRESS>
###### Abstract

For many scientific inverse problems we are required to evaluate an expensive
forward model. Moreover, the model is often given in such a form that it is
unrealistic to access its gradients. In such a scenario, standard Markov Chain
Monte Carlo algorithms quickly become impractical, requiring a large number of
serial model evaluations to converge on the target distribution. In this paper
we introduce Flow Annealed Kalman Inversion (FAKI). This is a generalization
of Ensemble Kalman Inversion (EKI), where we embed the Kalman filter updates
in a temperature annealing scheme, and use normalizing flows (NF) to map the
intermediate measures corresponding to each temperature level to the standard
Gaussian. In doing so, we relax the Gaussian ansatz for the intermediate
measures used in standard EKI, allowing us to achieve higher fidelity
approximations to non-Gaussian targets. We demonstrate the performance of FAKI
on two numerical benchmarks, showing dramatic improvements over standard EKI
in terms of accuracy whilst accelerating its already rapid convergence
properties (typically in $\mathcal{O}(10)$ steps).

##  1 Introduction

Many scientific inference tasks are concerned with inverse problems of the
form

$y=\mathcal{G}(x)+\eta,$ (1)

where $y\in\mathbb{R}^{d_{y}}$ are the data, $x\in\mathbb{R}^{d}$ are the
model parameters, $\mathcal{G}$ is the forward map, and $\eta$ is the
observation noise. Throughout this work we will assume that we do not have
access to gradients of $\mathcal{G}$ with respect to the parameters, and that
$\eta\sim\mathcal{N}(0,\Gamma)$ where $\Gamma$ is a fixed noise covariance.
The assumption of additive Gaussian noise is the standard setting for Ensemble
Kalman Inversion (EKI) [1, 2, 3, 4, 5, 6, 7, 8], and whilst we are restricted
to problems with Gaussian likelihoods, this covers a large family of
scientific inverse problems. The goal of the Bayesian inverse problem is then
to recover the posterior distribution over the model parameters given our
observations, $p(x|y)$.

Typical gradient-free inference methods often involve some variant on Markov
Chain Monte Carlo (MCMC) algorithms e.g., random walk Metropolis [9, 10, 11],
or Sequential Monte Carlo (SMC) [12]. However, these methods typically require
$\gtrsim 10^{3}$ serial model evaluations to achieve convergence, making them
intractable for problems with expensive forward models. EKI by contrast
utilizes embarrassingly parallel model evaluations to update parameter
estimates, typically converging to an approximate solution in
$\mathcal{O}(10)$ iterations [2, 6, 7, 8].

EKI leverages ideas originally developed in the context of Ensemble Kalman
Filtering (EKF) for data assimilation [13]. Since its development, EKI has
seen applications across a range of disciplines, including studies of fluid
flow [14], climate models [15] and machine learning tasks [16]. EKI can be
understood in the context of annealing, where seek to move from the prior to
the posterior through a sequence of intermediate measures. In standard EKI,
this involves constructing a sequence of Gaussian approximations to the
intermediate measures. In the regime where we have a Gaussian prior
$\pi_{0}(x)=\mathcal{N}(m_{0},C_{0})$ and a linear forward model
$\mathcal{G}(x)=Gx$, the particle distribution obtained via EKI converges to
the true posterior in the limit where the ensemble size $J\rightarrow\infty$.
However, outside this linear, Gaussian regime EKI is an uncontrolled
approximation to the posterior that is constructed on the basis of matching
first and second moments of the target distribution. Nonetheless, EKI has been
shown to perform well on problems with nonlinear forward models and slightly
non-Gaussian targets [1, 2, 6].

In this work we propose the application of normalizing flows (NF) [17, 18, 19,
20] to relax the Gaussian ansatz made by standard EKI for the intermediate
measures. Instead of assuming a Gaussian particle distribution at each
iteration, the NF is used to fit for the empirical particle distribution and
map to a Gaussian latent space, where the EKI updates are performed. In doing
so, we are better able to capture non-Gaussian target geometries. The
structure of this paper is as follows: in Section 2 we describe the Flow
Annealed Kalman Inversion (FAKI) algorithm, in Section 3 we demonstrate the
performance of the method on two Bayesian inference tasks with non-Gaussian
target geometries and we summarize our work in Section 4.

##  2 Methods

###  2.1 Regularized Ensemble Kalman Inversion

A number of versions of EKI have been proposed in the literature. Of interest
here is the regularized, perturbed observation form of EKI [6]. Starting with
an ensemble of particles drawn from the prior, $\\{x_{0}^{j}\\}_{j=1}^{J}$,
the particles are updated at each iteration according to

$x_{n+1}^{j}=x_{n}^{j}+C_{n}^{x\mathcal{G}}(C_{n}^{\mathcal{G}\mathcal{G}}+\alpha_{n}\Gamma)^{-1}(y-\mathcal{G}(x_{n}^{j})+\sqrt{\alpha_{n}}\xi_{n}^{j}).$
(2)

The empirical covariances $C_{n}^{x\mathcal{G}}$ and
$C_{n}^{\mathcal{G}\mathcal{G}}$ are given by

$\displaystyle C_{n}^{x\mathcal{G}}$
$\displaystyle=\frac{1}{J-1}\sum_{j=1}^{J}(x_{n}^{j}-\langle
x_{n}\rangle)\otimes(\mathcal{G}(x_{n}^{j})-\langle\mathcal{G}_{n}\rangle),$
(3) $\displaystyle C_{n}^{\mathcal{G}\mathcal{G}}$
$\displaystyle=\frac{1}{J-1}\sum_{j=1}^{J}(\mathcal{G}(x_{n}^{j})-\langle\mathcal{G}_{n}\rangle)\otimes(\mathcal{G}(x_{n}^{j})-\langle\mathcal{G}_{n}\rangle).$
(4)

At each iteration we perturb the forward model evaluations with the Gaussian
observation noise $\xi^{j}_{n}\sim\mathcal{N}(0,\Gamma)$. The parameter
$\alpha_{n}$ is a Tikhonov regularisation parameter, which can be viewed as an
inverse step size in the Bayesian annealing context. In particular, given a
set of annealing parameters $\beta_{0}\equiv
0<\beta_{1}<\ldots<\beta_{N}<\beta_{N+1}\equiv 1$, we have the corresponding
set of target distributions

$\pi_{n}(x)\propto\pi_{0}(x)\exp\left(-\frac{\beta_{n}}{2}\norm{\Gamma^{-1/2}(y-\mathcal{G}(x))}^{2}\right),$
(5)

with

$\alpha_{n}=\beta_{n+1}-\beta_{n}.$ (6)

EKI proceeds by constructing a sequence of ensemble approximations to Gaussian
distributions that approximate the intermediate targets.

The choice of the regularization parameter, $\alpha_{n}$ controls the
transition from the prior to the posterior. Previous proposals for an adaptive
choice have taken inspiration from SMC by using a threshold on the effective
sample size (ESS) of the particles at each temeperature level [21, 22]. In
this work we adopt the same approach, calculating pseudo-importance weights at
each temperature given by

$w_{n}^{j}=\exp\left(-\frac{1}{2}(\beta_{n+1}-\beta_{n})\norm{\Gamma^{-1/2}(y-\mathcal{G}(x_{n}^{j}))}^{2}\right).$
(7)

The next temperature level can then be selected by solving

$\left(\sum_{j=1}^{J}w_{n}^{j}(\beta_{n+1})^{2}\right)^{-1}\left(\sum_{j=1}^{J}w_{n}^{j}(\beta_{n+1})\right)^{2}=\tau
J,$ (8)

using the bisection method, where $0<\tau<1$ is the target fractional ESS
threshold. Throughout our work we set $\tau=0.5$. Full pseudocode for EKI is
given in Algorithm 1.

Algorithm 1 Ensemble Kalman Inversion

1: Input: $J$ prior samples $\\{x_{0}^{j}\sim\pi_{0}(x)\\}_{j=1}^{J}$, data
$y$, observation error covariance $\Gamma$ and fractional ESS target threshold
$\tau$

2: Initialize inverse temperature $\beta_{0}=0$, iteration counter $n=0$

3: while $\beta<1$ do

4: Evaluate $\mathcal{G}_{j}=\mathcal{G}(x_{n}^{j})$, $j\in\\{1,\ldots,J\\}$

5: Solve for $\beta_{n+1}$ using the bisection method with Equation 8

6: $\alpha_{n}\leftarrow\beta_{n+1}-\beta_{n}$

7: Update particles using Equation 2 to obtain $\\{x_{n+1}^{j}\\}_{j=1}^{J}$

8: $n\leftarrow n+1$

9: end while

10: Output: Converged particle ensemble $\\{x_{N}^{j}\\}_{j=1}^{J}$

###  2.2 Normalizing Flows

As discussed above, standard EKI proceeds by constructing a sequence of
ensemble approximations to Gaussian distributions. The procedure works well in
the situation where the target and all the intermediate measures are close to
Gaussian. However, when any of these measures are far from Gaussian, EKI can
dramatically fail to capture the final target geometry.

To address this shortcoming we propose the use of NFs to approximate each
intermediate target, instead of using the Gaussian ansatz of standard EKI. NFs
are powerful generative models that can be used for flexible density
estimation and sampling [17, 18, 19, 20]. An NF model maps from the original
space $x\in\mathbb{R}^{d}$ to a latent space $z\in\mathbb{R}^{d}$, through a
sequence of invertible transformations $f=f_{1}\circ f_{2}\circ\ldots\circ
f_{L}$, such that we have a bijective mapping $z=f(x)$. The mapping is such
that the latent variables are mapped to some simple base distribution,
typically chosen to be the standard Normal distribution, giving $z\sim
p_{z}(z)=\mathcal{N}(0,I)$.

The NF density can be evaluated through the change of variables formula,

$q(x)=p_{z}(f(x))\left|\mathrm{det}\,Df(x)\right|=p_{z}(f(x))\prod_{l=1}^{L}\left|\mathrm{det}\,Df_{l}(x)\right|,$
(9)

where $Df(x)=\partial f(x)/\partial x$ denotes the Jacobian of $f$. Efficient
evaluation of this density requires the Jacobian of the transformation to be
easy to evaluate, and efficient sampling requires the inverse of the mapping
$f$ to be easy to calculate. In this work we use Masked Autoregressive Flows
(MAF) [18], which have previously been found to perform well in the context of
preconditioned MCMC sampling within SMC without the need for expensive hyper-
parameter searches during sampling [23].

###  2.3 Flow Annealed Kalman Inversion

Given particles distributed as $\pi_{n}(x)$, the subsequent target can be
written as

$\pi_{n+1}(x)\propto\pi_{n}(x)\exp\left(-\frac{1}{2\alpha_{n}}\norm{\Gamma^{-1/2}(y-\mathcal{G}(x))}^{2}\right).$
(10)

We may therefore view $\pi_{n}(x)$ i.e., the posterior at the temperature
level $\beta_{n}$, as an effective prior for $\pi_{n+1}(x)$, with a data
likelihood annealed by $\alpha_{n}^{-1}$. By fitting an NF to the particles
$\\{x_{n}^{j}\\}_{j=1}^{J}$, we obtain an approximate map from the
intermediate target $\pi_{n}(x)$ to $\mathcal{N}(z|0,I)$. The latent space
target is then given by the change of variables formula as

$\pi_{n+1}(z)=\pi_{n+1}(x=f_{n}^{-1}(z))\left|\mathrm{det}\,Df_{n}^{-1}(z)\right|.$
(11)

By controlling the choice of $\alpha_{n}$, we control the distance between the
Gaussianized effective prior and this latent space target density. For FAKI,
we therefore perform the EKI updates in the NF latent space at each
temperature level, allowing us to relax the Gaussian ansatz of standard EKI by
constructing an approximate map from each $\pi_{n}(x)$ to a Gaussian latent
space. It is worth noting that, whilst this method relaxes the Gaussianity
assumptions of standard EKI, it does not address the linearity assumptions
used in deriving EKI.

The FAKI update for the latent space particle locations is given by

$z_{n+1}^{j}=z_{n}^{j}+\mathcal{C}_{n}^{z\mathcal{G}}(\mathcal{C}_{n}^{\mathcal{G}\mathcal{G}}+\alpha_{n}\Gamma)^{-1}(y-\mathcal{G}(f_{n}^{-1}(z_{n}^{j}))+\sqrt{\alpha_{n}}\xi_{n}^{j}),$
(12)

where the latent space empirical covariances are given by

$\displaystyle\mathcal{C}_{n}^{z\mathcal{G}}$
$\displaystyle=\frac{1}{J-1}\sum_{j=1}^{J}(z_{n}^{j}-\langle
z_{n}\rangle)\otimes(\mathcal{G}(f_{n}^{-1}(z_{n}^{j}))-\langle\mathcal{G}_{n}\rangle),$
(13) $\displaystyle\mathcal{C}_{n}^{\mathcal{G}\mathcal{G}}$
$\displaystyle=\frac{1}{J-1}\sum_{j=1}^{J}(\mathcal{G}(f_{n}^{-1}(z_{n}^{j}))-\langle\mathcal{G}_{n}\rangle)\otimes(\mathcal{G}(f_{n}^{-1}(z_{n}^{j}))-\langle\mathcal{G}_{n}\rangle).$
(14)

Full pseudocode for FAKI is given in Algorithm 2.

Algorithm 2 Flow Annealed Kalman Inversion

1: Input: $J$ prior samples $\\{x_{0}^{j}\sim\pi_{0}(x)\\}_{j=1}^{J}$, data
$y$, observation error covariance $\Gamma$ and fractional ESS target threshold
$\tau$

2: Initialize inverse temperature $\beta_{0}=0$, iteration counter $n=0$

3: while $\beta<1$ do

4: Evaluate $\mathcal{G}_{j}=\mathcal{G}(x_{n}^{j})$, $j\in\\{1,\ldots,J\\}$

5: Solve for $\beta_{n+1}$ using the bisection method with Equation 8

6: $\alpha_{n}\leftarrow\beta_{n+1}-\beta_{n}$

7: Fit NF map $f_{n}$ to current samples $\\{x_{n}^{j}\\}_{j=1}^{J}$

8: Map particles to latent space $z_{n}^{j}=f_{n}(x_{n}^{j})$,
$j\in\\{1,\ldots,J\\}$

9: Update particles using Equation 12 to obtain $\\{z_{n+1}^{j}\\}_{j=1}^{J}$

10: Map back to the data space $x_{n+1}^{j}=f_{n}^{-1}(z_{n+1}^{j})$,
$j\in\\{1,\ldots,J\\}$

11: $n\leftarrow n+1$

12: end while

13: Output: Converged particle ensemble $\\{x_{N}^{j}\\}_{j=1}^{J}$

##  3 Results

In this section we demonstrate the performance of FAKI compared to standard
EKI on two numerical benchmarks, a two dimensional Rosenbrock distribution and
a stochastic Lorenz system [24, 25]. Both models display significant non-
Gaussianity at some point during the transition from prior to posterior,
severely frustrating the performance of EKI. This manifests in both reduced
fidelity of the final ensemble approximations to the posterior, and in a
larger number of iterations being required for convergence following the ESS-
based annealing scheme described in Section 2.1.

In Table 1 we provide statistics summarizing the performance of EKI and FAKI
on our numerical benchmarks. We measure the quality of the posterior
approximations by computing the 1-Wasserstein distance, $W_{1}$ [26, 27]
between the samples obtained through FAKI and EKI, against reference posterior
samples obtained via long runs of Hamiltonian Monte Carlo (HMC) [28, 29].
These reference samples are thinned to be approximately independent when
computing the 1-Wasserstein distances111We use the Python Wasserstein library:
https://github.com/pkomiske/Wasserstein/.. The 1-Wasserstein distance may be
interpreted as the cost involved in rearranging one probability measure to
look like another, with lower values indicating the two probability measures
are closer to one another. In addition to this assessment of the approximation
quality, we report the number of iterations, $N_{\mathrm{iter}}$ required by
FAKI and EKI for convergence. For both quantities we report the median and
median absolute deviation (MAD), estimated over 10 independent runs using
different random seeds.

Table 1: Median and MAD values for $N_{\mathrm{iter}}$ and 1-Wasserstein distances for each model and algorithm combination, calculated over 10 independent runs using different random seeds. For both the numerical benchmarks we see that FAKI results in a reduced number of iterations for convergence, and a lower value of the 1-Wasserstein distance between the converged samples and the ground truth. Model | Algorithm | $\mathrm{Median}[N_{\mathrm{iter}}]$ | $\mathrm{MAD}[N_{\mathrm{iter}}]$ | $\mathrm{Median}[W_{1}]$ | $\mathrm{MAD}[W_{1}]$
---|---|---|---|---|---
Rosenbrock | EKI | 100 | 7.0 | 0.72 | 0.05
Rosenbrock | FAKI | 34.0 | 7.0 | 0.43 | 0.14
Lorenz | EKI | 10.0 | 0.0 | 69.8 | 1.08
Lorenz | FAKI | 8.0 | 0.0 | 5.65 | 0.86

###  3.1 $d=2$ Rosenbrock

In our first numerical experiment we consider the two dimensional Rosenbrock
distribution. This toy model allows us to clearly see the impact of non-
Gaussianity on the performance of EKI, and how FAKI is able to alleviate these
issues. For the Rosenbrock model we assume a Gaussian prior over the
parameters $x\in\mathbb{R}^{2}$,

$x\sim\mathcal{N}(0,10^{2}I).$ (15)

The data, $y\in\mathbb{R}^{2}$ are distributed according to the likelihood,

$y\sim\mathcal{N}(\mathcal{G}(x)=(x_{1}-x_{0}^{2},x_{0})^{\intercal},\Gamma=\mathrm{diag}(0.01^{2},1^{2})).$
(16)

To generate simulated data we evaluate $y=G((1,1)^{\intercal})+\eta$, where
$\eta\sim\mathcal{N}(0,\mathrm{diag}(0.01^{2},1^{2}))$. The large difference
in noise scales results in a highly non-Gaussian posterior geometry that poses
a significant challenge for EKI. For each run of EKI and FAKI we use 100
particles.

In Figure 1 we show pair-plots comparing the final particle distributions
obtained with EKI and FAKI against samples obtained through a long run of HMC.
The NF mapping means that the ensemble approximation obtained by FAKI is able
to capture the highly nonlinear target geometry. In comparison, EKI struggles
to fill the tails of the Rosenbrock target. Moreover, whilst FAKI converges
within $\sim 34$ iterations, EKI required a median number of $\sim 100$
iterations to converge using the ESS-based annealing scheme.

(a) EKI-MCMC comparison

(b) FAKI-MCMC comparison

Figure 1: Pair-plots for the Rosenbrock target. Panel (a): Pair-plot
comparison of samples from EKI and a long HMC run. Panel (b): Pair-plot
comparison of samples from FAKI and a long HMC run. Samples from FAKI are able
to correctly capture the highly nonlinear target geometry. Standard EKI
struggles to fill the tails of the target, and requires $\sim 100$ iterations
to converge, compared to $\sim 34$ iterations for FAKI.

###  3.2 Stochastic Lorenz System

The Lorenz equations are a set of coupled differential equations used as a
simple model of atmospheric convection. Notably, for certain parameter values
the Lorenz equations are known to exhibit chaotic behaviour [24]. In this work
we follow [25] and consider the stochastic Lorenz system,

$\displaystyle\mathrm{d}X_{t}$
$\displaystyle=10(Y_{t}-X_{t})\mathrm{d}t+\mathrm{d}W^{x}_{t},$ (17)
$\displaystyle\mathrm{d}Y_{t}$
$\displaystyle=X_{t}(28-Z_{t})\mathrm{d}t-Y_{t}\mathrm{d}t+\mathrm{d}W^{y}_{t},$
(18) $\displaystyle\mathrm{d}Z_{t}$
$\displaystyle=X_{t}Y_{t}\mathrm{d}t-\frac{8}{3}Z_{t}\mathrm{d}t+\mathrm{d}W^{z}_{t},$
(19)

where $W^{x}_{t}$, $W^{y}_{t}$ and $W^{z}_{t}$ are Gaussian white noise
processes with standard deviation $\sigma_{0}=0.1$. To generate simulated data
we integrated these equations using an Euler–Maruyama scheme with
$\mathrm{d}t=0.02$ for 30 steps, with initial conditions
$X_{0},Y_{0},Z_{0}\sim\mathcal{N}(0,1^{2})$. The observations are then taken
to be the $X_{t}$ values over these 30 time steps, with Gaussian observational
noise $\eta_{t}\sim\mathcal{N}(0,\sigma^{2}=1^{2})$.

(a) EKI-MCMC $\mathbb{E}[X]$ comparison

(b) FAKI-MCMC $\mathbb{E}[X]$ comparison

(c) EKI-MCMC $\mathrm{Var}[X]^{1/2}$ comparison

(d) FAKI-MCMC $\mathrm{Var}[X]^{1/2}$ comparison

Figure 2: Comparison of first and second moment estimates along each dimension
for the stochastic Lorenz system. Panel (a): Comparison between the mean
estimates from EKI and a long HMC run. Panel (b): Comparison between the mean
estimates from FAKI and a long HMC run. Panel (c): Comparison between the
standard deviation estimates from EKI and a long HMC run. Panel (d):
Comparison between the standard deviation estimates from FAKI and a long HMC
run. Blue bars indicate the moment estimates obtained via HMC along each
dimension, with the adjacent orange bars showing the estimates obtained
through EKI/FAKI. EKI is unable to obtain accurate mean estimates for much of
the $Z_{t}$ trajectory, whilst FAKI is able to obtain accurate mean estimates
for each dimension. FAKI outperforms EKI in its estimates of the marginal
standard deviations, with EKI drastically overestimating the standard
deviations along many dimensions.

The goal of our inference here is to recover the initial conditions, the
trajectories $(X_{t},Y_{t},Z_{t})$ and the innovation noise scale
$\sigma_{0}$, giving a parameter space of $d=94$ dimensions. We assign priors
over these parameters as,

$\displaystyle\log\sigma_{0}$ $\displaystyle\sim\mathcal{N}(-1,1^{2}),$ (20)
$\displaystyle X_{0},Y_{0},Z_{0}$ $\displaystyle\sim\mathcal{N}(0,1^{2}),$
(21) $\displaystyle X_{t}$
$\displaystyle\sim\mathcal{N}(X_{t-1}+f_{X}(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,\sigma_{0}^{2}\mathrm{d}t),t\in\\{1,\ldots,30\\},$
(22) $\displaystyle Y_{t}$
$\displaystyle\sim\mathcal{N}(Y_{t-1}+f_{Y}(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,\sigma_{0}^{2}\mathrm{d}t),t\in\\{1,\ldots,30\\},$
(23) $\displaystyle Z_{t}$
$\displaystyle\sim\mathcal{N}(Z_{t-1}+f_{Z}(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,\sigma_{0}^{2}\mathrm{d}t),t\in\\{1,\ldots,30\\},$
(24)

where $f_{X}$, $f_{Y}$ and $f_{Z}$ are the transition functions corresponding
to Equations 17-19 respectively. The Gaussian likelihood has the form

$\hat{X}_{t}\sim\mathcal{N}(X_{t},\sigma^{2}),t\in\\{1,\ldots,30\\},$ (25)

where $\hat{X}_{t}$ are the observations of the $X_{t}$ trajectory. The
chaotic dynamics of the Lorenz system results in a highly non-Gaussian prior
distribution, with the inversion having to proceed through a sequence of
highly non-Gaussian intermediate measures towards the posterior. This severely
frustrates the performance of EKI, with the Gaussian ansatz failing to
describe the geometry of the intermediate measures. For each run of EKI and
FAKI we use 940 particles.

In Figure 2 we show the ensemble estimates for the mean and standard deviation
along each dimension obtained by EKI and FAKI, compared to reference estimates
obtained through long runs of HMC. FAKI is able to obtain accurate mean
estimates along each dimension, whereas EKI is unable to obtain the correct
means for much of the $Z_{t}$ trajectory. EKI severely overestimates the
marginal standard deviations along many dimensions. This situation is
alleviated by the NF mappings learned by FAKI. The greater fidelity of the
FAKI posterior approximations are reflected in the median estimates for the
1-Wasserstein distances, with a value of $5.65$ for FAKI and $69.8$ for EKI.

##  4 Conclusions

In this work we have introduced Flow Annealed Kalman Inversion (FAKI), a
gradient-free inference algorithm for Bayesian inverse problems with expensive
forward models. This is a generalization of Ensemble Kalman Inversion (EKI),
where we utilize Normalizing Flows (NF) to replace the Gaussian ansatz made in
EKI. Instead of constructing a sequence of ensemble approximations to Gaussian
measures that approximate a sequence of intermediate measures, as we move from
the prior to the posterior, we learn an NF mapping at each iteration to a
Gaussian latent space. Provided the transition between temperature levels is
controlled, we can perform Kalman inversion updates in the NF latent space. In
the NF latent space, the Gaussianity assumptions of EKI are more closely
satisfied, resulting in a more stable inversion at each temperature level.

We demonstrate the performance of FAKI on two numerical benchmarks, a $d=2$
Rosenbrock distribution and a $d=94$ stochastic Lorenz system. Both examples
exhibit significant non-Gaussianity in the transition from prior to posterior
that frustrate standard EKI. In the presence of strong non-Gaussianity, we
find FAKI produces higher fidelity posterior approximations compared to EKI,
as measured by the 1-Wasserstein distance between FAKI/EKI samples and
reference HMC samples. In addition to the improved fidelity of the posterior
approximations, we find FAKI tends to reduce the number of iterations required
for convergence.

Whilst the application of NFs is able to relax the Gaussian ansatz of EKI, it
does not address the linearity assumptions used in deriving EKI. As such, FAKI
is still not exact for general forward models. In future work, it will be
interesting to explore methods to address this, for example the combination of
FAKI with unbiased MCMC or importance sampling methods. It would also be
interesting to consider generalizations of FAKI that are able to accommodate
non-Gaussian likelihoods and/or parameter-dependent noise covariances. The use
of NFs means that we typically require ensemble sizes $J\gtrsim 10d$ to learn
accurate NF maps with the MAF architecture employed in this work. It would be
useful to explore alternative NF architectures and regularization schemes that
are able to learn accurate NF maps with smaller ensemble sizes, in order to
enable FAKI to scale to higher dimensions. In this work, we have found that
the MAF architecture is able to capture a wide range of target geometries
without the need for expensive NF hyper-parameter searches. However, it may be
possible to exploit NF architectures with inductive biases that are
particularly suited to common target geometries e.g., the nonlinear
correlations that often appear in hierarchical models.

## Acknowledgments

This research was funded by NSFC (grant No. 12250410240) and the U.S.
Department of Energy, Office of Science, Office of Advanced Scientific
Computing Research under Contract No. DE-AC02-05CH11231 at Lawrence Berkeley
National Laboratory to enable research for Data-intensive Machine Learning and
Analysis. RDPG was supported by a Tsinghua Shui Mu Fellowship.

The authors thank Qijia Jiang and David Nabergoj for helpful discussions.

## References

  * [1] Marco A Iglesias, Kody JH Law, and Andrew M Stuart.  Ensemble kalman methods for inverse problems.  Inverse Problems, 29(4):045001, 2013.
  * [2] Marco A Iglesias.  A regularizing iterative ensemble kalman method for pde-constrained inverse problems.  Inverse Problems, 32(2):025002, 2016.
  * [3] Claudia Schillings and Andrew M Stuart.  Analysis of the ensemble kalman filter for inverse problems.  SIAM Journal on Numerical Analysis, 55(3):1264–1290, 2017.
  * [4] Neil K Chada, Marco A Iglesias, Lassi Roininen, and Andrew M Stuart.  Parameterizations for ensemble kalman inversion.  Inverse Problems, 34(5):055009, 2018.
  * [5] Claudia Schillings and Andrew M Stuart.  Convergence analysis of ensemble kalman inversion: the linear, noisy case.  Applicable Analysis, 97(1):107–123, 2018.
  * [6] Marco Iglesias and Yuchen Yang.  Adaptive regularisation for ensemble kalman inversion.  Inverse Problems, 37(2):025008, 2021.
  * [7] Daniel Zhengyu Huang, Tapio Schneider, and Andrew M Stuart.  Iterated kalman methodology for inverse problems.  Journal of Computational Physics, 463:111262, 2022.
  * [8] Daniel Zhengyu Huang, Jiaoyang Huang, Sebastian Reich, and Andrew M Stuart.  Efficient derivative-free bayesian inference for large-scale inverse problems.  Inverse Problems, 38(12):125006, 2022.
  * [9] Charles J. Geyer.  Practical markov chain monte carlo.  Statistical Science, 7(4):473–483, 1992.
  * [10] A. Gelman, W. R. Gilks, and G. O. Roberts.  Weak convergence and optimal scaling of random walk Metropolis algorithms.  The Annals of Applied Probability, 7(1):110 – 120, 1997.
  * [11] S. L. Cotter, G. O. Roberts, A. M. Stuart, and D. White.  MCMC Methods for Functions: Modifying Old Algorithms to Make Them Faster.  Statistical Science, 28(3):424 – 446, 2013.
  * [12] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra.  Sequential Monte Carlo Samplers.  Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411–436, 05 2006.
  * [13] Geir Evensen.  Sequential data assimilation with a nonlinear quasi-geostrophic model using monte carlo methods to forecast error statistics.  Journal of Geophysical Research: Oceans, 99(C5):10143–10162, 1994.
  * [14] Heng Xiao, J-L Wu, J-X Wang, Rui Sun, and CJ3540338 Roy.  Quantifying and reducing model-form uncertainties in reynolds-averaged navier–stokes simulations: A data-driven, physics-informed bayesian approach.  Journal of Computational Physics, 324:115–136, 2016.
  * [15] Tapio Schneider, Shiwei Lan, Andrew Stuart, and João Teixeira.  Earth system modeling 2.0: A blueprint for models that learn from observations and targeted high-resolution simulations.  Geophysical Research Letters, 44(24):12–396, 2017.
  * [16] Nikola B Kovachki and Andrew M Stuart.  Ensemble kalman inversion: a derivative-free technique for machine learning tasks.  Inverse Problems, 35(9):095005, 2019.
  * [17] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real NVP.  In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
  * [18] George Papamakarios, Iain Murray, and Theo Pavlakou.  Masked autoregressive flow for density estimation.  In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 2338–2347, 2017.
  * [19] Diederik P. Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 10236–10245, 2018.
  * [20] Biwei Dai and Uros Seljak.  Sliced iterative normalizing flows.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 2352–2364. PMLR, 2021.
  * [21] Lia De Simon, Marco Iglesias, Benjamin Jones, and Christopher Wood.  Quantifying uncertainty in thermophysical properties of walls by means of bayesian inversion.  Energy and Buildings, 177:220–245, 2018.
  * [22] Marco Iglesias, Minho Park, and MV Tretyakov.  Bayesian inversion in resin transfer molding.  Inverse Problems, 34(10):105002, 2018.
  * [23] Minas Karamanis, Florian Beutler, John A Peacock, David Nabergoj, and Uroš Seljak.  Accelerating astronomical and cosmological inference with preconditioned monte carlo.  Monthly Notices of the Royal Astronomical Society, 516(2):1644–1653, 2022.
  * [24] Edward N. Lorenz.  Deterministic Nonperiodic Flow.  Journal of Atmospheric Sciences, 20(2):130–148, March 1963.
  * [25] Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and Marcel van Gerven.  Automatic structured variational inference.  In International Conference on Artificial Intelligence and Statistics, pages 676–684. PMLR, 2021.
  * [26] Cédric Villani.  Optimal transport – Old and new, volume 338, pages xxii+973.  01 2008.
  * [27] Lu Zhang, Bob Carpenter, Andrew Gelman, and Aki Vehtari.  Pathfinder: Parallel quasi-newton variational inference.  Journal of Machine Learning Research, 23(306):1–49, 2022.
  * [28] Radford M Neal et al.  Mcmc using hamiltonian dynamics.  Handbook of Markov Chain Monte Carlo, 2(11):2, 2011.
  * [29] Matthew D Hoffman, Andrew Gelman, et al.  The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo.  Journal of Machine Learning Research, 15(1):1593–1623, 2014.


# Independent Modular Networks
††thanks: † School of Computer Science, University of Adelaide, Australia
∗ College of Engineering and Computing, University of South Carolina,
Columbia, South Carolina, USA.

Hamed Damirchi†<EMAIL_ADDRESS>Forest Agostinelli∗
<EMAIL_ADDRESS>Pooyan Jamshidi∗<EMAIL_ADDRESS>
###### Abstract

Monolithic neural networks that make use of a single set of weights to learn
useful representations for downstream tasks explicitly dismiss the
compositional nature of data generation processes. This characteristic exists
in data where every instance can be regarded as the combination of an identity
concept, such as the shape of an object, combined with modifying concepts,
such as orientation, color, and size. The dismissal of compositionality is
especially detrimental in robotics, where state estimation relies heavily on
the compositional nature of physical mechanisms (e.g., rotations and
transformations) to model interactions. To accommodate this data
characteristic, modular networks have been proposed. However, a lack of
structure in each module’s role, and modular network-specific issues such as
module collapse have restricted their usability. We propose a modular network
architecture that accommodates the mentioned decompositional concept by
proposing a unique structure that splits the modules into predetermined roles.
Additionally, we provide regularizations that improve the resiliency of the
modular network to the problem of module collapse while improving the
decomposition accuracy of the model.

###### Index Terms:

modular networks, representation learning, world models

##  I Introduction

Representation learning using monolithic models (models that use a single set
of weights for every input) typically involves passing all data samples
through the model and extracting features after training on pretext tasks
deemed helpful for downstream tasks. However, this approach ignores the
compositional characteristics of data generation processes. As an example of
this common characteristic for images of objects, every object can be
described as a composition of the object’s shape, color, size, orientation,
texture, etc. Instead of learning these concepts separately while considering
the compositional nature of each possible combination of the mentioned
concepts, monolithic models attempt to directly extract high-level features
without any constraints that would impose such a structure explicitly.

Modular neural networks [1] was proposed as a potential solution that
considers this compositional characteristic, where a set of modules instead of
a singular neural network is used to extract features from any given input.
Generally, the feature extraction process for these models proceeds using a
scoring method to determine which module or set of modules from the available
modules should process the input. Then, the input is passed through the chosen
modules. In the case where multiple modules were chosen, a combination method,
such as adding features, is chosen to combine the features extracted by each
module.

Regardless of the module selection approach, while modular networks take a
step towards solutions considering the mentioned compositional characteristics
of data, one particular structural element is still missing from the currently
available works. To have a compositional structure, there is a need for an
identity state to be altered through combination with compositional concepts.
We define this identity state as one of the true generative factors of the
dataset that does not require another factor to be defined. Using the previous
example of images of objects, one can consider the shape of the object (a
physical concept) as the identity state that can be modified by compositional
concepts (non-physical concepts) such as rotations, color, and scaling. In
this example, while the object’s shape can be defined separately from the
other factors, the remaining factors require a shape to be defined in a
hierarchical manner beforehand. Therefore, we propose structural changes to
modular networks that allow the learning of two separate sets of modules,
where one set of modules automatically learns a notion of identity, and the
other set only encodes the compositional concepts. In particular, this is done
by only allowing the compositional modules to observe the input while the
identity modules do not. Instead, the identity modules will use a set of
learnable parameters to output a static state that is modified using
compositional modules. The reason for this change is that the identity modules
will now have to learn a specific, unchangeable concept that remains static
among large portions of the dataset.

A problem that modular networks commonly deal with is when the model routes
all the inputs through the same module or set of modules. Different works
propose varying solutions, such as regularizing the router to choose diverse
modules or modifying the way modules are combined during feature extraction.
In this work, we propose to solve this problem indirectly by imposing
independence constraints for the features extracted by each module and among
the compositional and identity modules. By penalizing correlations between the
concepts embedded by each module, this constraint would force modules to learn
to embed information not considered by other modules, which leads to the
decomposition of the concepts present in the dataset into separate modules.

  * •

Propose a new architecture for modular networks that promotes learning
features that are in line with the compositional nature of data generation
processes

  * •

Propose an independence-based solution that indirectly solves module collapse.

  * •

Provide experiments showcasing the automatic decomposition capabilities of the
proposed approach in the extraction of the identity states.

##  II Related Work

Algorithmically, the previous works on modular networks share a few structural
traits, such as the usage of a router and the recursiveness of the application
of modules to the input. Meanwhile, other details such as the training method,
learnability of different components of the methods, and the application
differ between the works in the literature. In the following, a summary of the
literature alongside the distinction between each work and ours is delineated.

In [1], a router and a set of modules are trained using a generalized
expectation-maximization (EM) algorithm. For each input, the router selects a
set of modules from a library of modules to process the input. The outputs
from selected modules are then combined (either through concatenation or
addition) to compute the output of the layer. While a modular approach, this
method does not impose any structural constraints on the modules to decompose
features into identity and compositional states.

While making use of reinforcement learning algorithms instead of EM, [2]
allows the routing module to decide when to stop processing the original input
rather than using a predetermined number of recursions. While the
decomposition of the input into multiple atomic representations using multiple
modules is similar to our work, we do not constrain our modules to encode
specific information such as rotation, and the learning is done automatically
during the training.

Each of the mentioned works deals with a problem called module collapse, where
during training, the routing algorithm routes all the inputs through only a
few of the modules. [1] uses a different number of M-step compared to the
E-step in the EM algorithm to mitigate this issue. On the other hand, [3]
finds that using multiple agents per task mitigates the module collapse
problem. Even though previous works have, in some cases, shown that module
collapse is prevented using various regularization techniques, without an
independency constraint, diversity in usage of modules may not mean that
module collapse does not occur since multiple modules might still encode the
same concept. Conventional module collapse detection methods cannot detect
this. In contrast, our proposed independence-based approach solves this
problem indirectly by preventing the same concept from being encoded by
multiple modules.

##  III Proposed approach

The general architecture of the proposed approach is shown in Fig. 1. The pool
of available modules can be split into two categories, where one group of
modules is only used for learning compositional concepts, and modules from
this group are able to observe the input. The other modules are not able to
observe the input and output of a learnable set of parameters in the shape of
a transformation matrix to be combined with compositional modules. Each
compositional module’s architecture is similar to that of the encoder of Lie
group Variational Autoencoders (LVAE) [4]. Modules designated $T^{I}$
represent the learnable identity matrices we name identity modules, while
$m^{c}$ designates compositional modules. While the proposed approach can be
used with any set of data with compositional characteristics, we use an image
reconstruction application in this paper. Therefore, a decoder is required to
reconstruct the image using the features extracted by the modules. This module
is designated by $d$ in Fig. 1. To train the model, we first output a
reconstruction of the input image using every combination of the available
modules. Then, a winning combination is chosen by choosing the combination
that provides an output with the lowest loss. In this section, we give a brief
background of LVAEs. Then, feature extraction and the reconstruction processes
are delineated. Finally, the regularizations and loss calculation methods are
provided.

Figure 1: The general structure of the proposed approach

###  III-A Lie-group VAE

LVAE [4] is a variational autoencoder that extracts a vector of disentangled
representations for the input image, while the same vector is an element of a
Lie group represented using the tangent space of the group. This way, smooth
equivariant representations are learned where the encoder predicts values
along the axes of the Lie algebra vector that are converted to their matrix
representations through exponential mapping. In feature extraction, the input
is first passed through an encoder typically chosen as a CNN for feature
extraction from images and a distribution is inferred over the position of the
latent variables. After sampling from this distribution, the latent variables
are used alongside the learnable Lie algebra bases to infer the element of the
group representing the input through the exponential operator. This process is
formulated as follows [4].

$\begin{split}\mu,\sigma&=g(x),\;z=\mu+\sigma\epsilon,\\\
T&=exp(zA),\;\hat{x}=d(T)\\\ \end{split}$ (1)

Where $g$ and $d$ are the encoders and the decoder of this network, $z$ is the
latent vector, and $A$ is the learnable Lie algebra bases. Additionally, $T\in
R^{u\times x}$ represents the extracted features as a transformation matrix,
i.e., an element of the group representing the input. To recover disentangled
representations, regularizations are applied to the network. [4] showed that
by imposing $A_{i}A_{j}=0,\forall i\neq j$, the following would be satisfied.

$H_{ij}=\frac{\partial^{2}g(t)}{\partial z_{i}\partial z_{j}}=0$ (2)

where $A_{i}$ is the $i^{th}$ Lie algebra basis and $H$ represents the Hessian
matrix. If (2) is satisfied for every $i$ and $j$ when $i\neq j$, then the
changes in the output of the network with respect to a latent variable are not
dependent on the changes of other latent variables. This ensures that
independent concepts are encoded into each latent variable during training.

###  III-B Modular feature extraction

To extract the features from an image, every compositional module first
processes the input image as follows:

$T^{c}_{i}=m^{c}_{i}(x)$ (3)

where $T^{c}_{i}$ represents the transformation features extracted by the
$i^{th}$ compositional module. Thereafter, every combination of the
compositional features is calculated, denoted as:

$\mathbf{T}^{c}=\\{T^{c}_{1},\hbox
to10.00002pt{.\hss.\hss.},T^{c}_{m},T^{c}_{1}T^{c}_{2},\hbox
to10.00002pt{.\hss.\hss.},T^{c}_{m-1}T^{c}_{m},\hbox
to10.00002pt{.\hss.\hss.},T^{c}_{1}\hbox
to10.00002pt{.\hss.\hss.}T^{c}_{m-1}T^{c}_{m}\\}$ (4)

where $\mathbf{T}^{c}$ represents the set of all combinations of the
compositional features, and $m$ is the number of compositional modules set
manually at the start of training. Now, to obtain the combined representation
of the compositional features with the identity transformations, every element
of the set $\mathbf{T}^{c}$ is multiplied by the learnable transformation
matrix representing each identity module as follows:

$\mathbf{T}=\\{T_{c_{i}}T_{I_{j}}\>|\;T_{c_{i}}\in\mathbf{T}^{c}\;,\;T_{I_{j}}\in\mathbf{T}^{I}\\},$
(5)

where $\mathbf{T}^{I}=\\{T^{I}_{1},\cdots,T^{I}_{n}\\}$ represents the set of
learnable transformations from every identity module where the pool of
available modules consists of $n$ identity modules. Every element of the set
$\mathbf{T}$ will be considered the candidate representation for input $x$
until the combination that produces the best output is chosen based on the
scoring criteria depicted in Section III-C1. But first, we need to produce a
reconstruction based on every combined feature in $\mathbf{T}$ so that every
element of this set can be evaluated in the image space. In this work, we use
a single decoder network for every feature to get a set of reconstructions as
follows:

$\hat{\mathbf{X}}=\\{d(T_{i})\;|\;T_{i}\in\mathbf{T}\\},$ (6)

where $d$ represents the decoder network.

###  III-C Scoring, regularization, and loss function

In order to update the weights of the modules, two steps remain. The winning
combination needs to be chosen first. Then, based on the output from the
winning combination, the loss will be calculated, and after the computation of
the gradients, the weights of the modules responsible for the generation of
the winning output will be updated.

####  III-C1 Scoring the output from each combination

Three criteria are used to score the reconstruction and the features extracted
using each combination of modules. First, the reconstruction is evaluated
using image mean-squared-error loss

$\mathcal{L}^{img}_{i}=\|x-\hat{x}_{i}\|,$ (7)

where $\hat{x}_{i}$ is the $i^{th}$ element of $\hat{\mathbf{X}}$. The second
criterion quantifies the level of independence between the features extracted
by each module, alongside the independency of the features between the
different modules. By promoting intra-module independency, the modules that
output higher-quality latent variables are prioritized. This is while inter-
module independency constraints incentivize combinations of modules that embed
the least amount of correlated information with respect to other modules in
the combination. To quantify intra-module independence (which will also be
used in the loss function later), we use the same approach as [4] where the
changes of the gradient of the output with respect to a single element of the
latent vector are computed with respect to every other element of the same
vector. We extend this formulation to inter-module independence as well. The
following shows the formulation for both independence quantification methods.

$\mathcal{L}^{ind}=\sum_{i,j}\frac{\partial^{2}\hat{x}}{z_{i}z_{j}}+\sum_{i,j}\frac{\partial^{2}\hat{x}}{z^{\prime}_{i}z^{\prime}_{j}}$
(8)

where $z_{i},z_{j}\in Z_{k}^{I}$ and $z^{\prime}_{i}\in
Z_{k}^{I},z^{\prime}_{j}\in Z_{k}^{I}{{}^{\prime}}$. Additionally, $Z_{k}^{I}$
represents the latent variables extracted by the $k^{th}$ module in the
combination under evaluation. Therefore, the first term on the right-hand side
of (8) quantifies how correlated the variables in the latent vector of one
module are while the second term quantifies how correlated the variables of
one module are compared to that of another module in the combination under
evaluation. The final term of the scoring function is the Kullback-Leibler
Divergence (KLD) between the predicted and the true posterior distribution
over the latent used as a part of the originally proposed VAE [5] loss
function. Putting this scoring formula together, we have the following
equation.

$S_{i}=-(\mathcal{L}^{img}_{i}+\mathcal{L}^{ind}_{i}+\sum_{j}\mathcal{L}^{KL}_{ij})$
(9)

where $S_{i}$ determines the score for the $i^{th}$ combination from (5) and
$j$ represents the KL divergence for the latent variable of the $j^{th}$
module in combination $i$. To choose the winning combination, the one with the
maximum score is selected and the loss is calculated for this winning
combination.

####  III-C2 Calculating the loss for the winning combination

With the winning combination selected, the only step left is to calculate the
loss and update the weights of the modules and the learnable parameters
representing the learned identity of the winning combination. To this end, the
following is used.

$\mathcal{L}_{i}=\mathcal{L}^{img}_{i}+\mathcal{L}^{ind}_{i}+\sum_{j}\mathcal{L}^{KL}_{ij}+\mathcal{L}^{extra}_{i}$
(10)

where the same image, independence, and KLD loss are re-used by adding an
extra LVAE-specific loss term that is not included in the scoring process and
is only used to ensure training stability in [4]. So far, the loss function
does not include any terms that would give the modules information about what
the encoded identity matrices should represent. To explore the effects of the
introduction of such an element, we use an additional loss and compare the
results to the case where (10) is used. This term aims to guide the modules
toward learning a specific ground truth factor as identity representations. To
do this, we introduce a classifier with a fully-connected architecture trained
to predict the shape of the input based on the identity transformation of the
winning combination. Moreover, to ensure the imposed identity remains
unchanged by the compositional modules, we add another loss that lowers the
$L^{2}$ distance between the classifier features of the identity and the
combined representations that we name ID classifier loss.

(a) Example of the 3dshapes dataset

(b) Learned identity reconstruction on 3Dshapes

(c) Learned identity reconstruction on 3Dshapes with ID classifier

Figure 2: Reconstructions of identity representations

##  IV Experiments

###  IV-A Datasets

We used the 3Dshapes [6] dataset, a synthetically generated dataset of 3D
objects in the middle of a room. The variable factors used to generate this
dataset consist of the camera’s orientation with respect to the object, shape,
and scale, and the color of the walls, object, and floor. The objects in this
dataset are capsules, spheres, cylinders, and cubes. A few example images of
this dataset are shown in Fig. 2a. A pool of 5 compositional modules and 5
learnable identity representations is available for every training session.

###  IV-B What are the learned identities?

To visualize what the identity representations learn, we use the decoder
without any input images to reconstruct the image represented by each learned
identity. Fig. 2b shows one image per identity reconstruction. Note that the
loss used for this training session is based on (10), and the ID classifier
mentioned in Section III-C2 is not used. Based on this figure, the 4 shapes
available in the dataset are split among the 5 modules, where one of the
shapes (cube) is learned by 2 modules. Therefore, it is evident that the
proposed identity-preserving approach to modular networks is able to decompose
the shapes into different learnable representations automatically without the
need for labels for the different shapes. Additionally, despite the presence
of 6 different ground truth generative factors, the learned representations
are the shape of the object and not any of the other 5 factors. Fig. 2c
visualizes the reconstructions when an ID classifier is used. Similar to the
previous reconstructions, one shape is learned twice, and the decomposition is
based on shapes.

ID Cls. | GT Shape | $m_{1}(\%)$ | $m_{2}(\%)$ | $m_{3}(\%)$ | $m_{4}(\%)$ | $m_{5}(\%)$
---|---|---|---|---|---|---
✗ | Cube | 28.03 | 28.80 | 0 | 43.17 | 0
Cylinder | 100.0 | 0 | 0 | 0 | 0
Sphere | 0 | 0 | 100.0 | 0 | 0
Capsule | 40.32 | 0 | 0 | 0 | 59.68
✓ | Cube | 93.62 | 0 | 0 | 0 | 6.377
Cylinder | 0 | 0 | 0 | 0 | 100.0
Sphere | 0 | 0 | 12.33 | 87.66 | 0
Capsule | 0 | 77.80 | 13.34 | 7.992 | 0

TABLE I: Shape decomposition evaluation

###  IV-C What modules are responsible for what shapes?

The previous experiment confirms that the learned identities represent shapes
when reconstructed. However, the same experiment needs to be performed when
inputs are present to study the effects of compositional modules on these
identities. To this end, we separate the images in the test split with respect
to their ground truth shapes. These images are then passed through the
dataset, and the winning combination for each image is stored. Using this
data, Table I is created where each entry in the table quantifies the fraction
of images of that row’s ground truth shape that had that column’s learned
identity in the winning combination for that image. This table is split row-
wise for the two experiments based on the usage of an ID classifier. Without
an ID classifier, all images containing a cylinder are passed through $m_{1}$,
and all images of spheres are passed through $m_{3}$, which shows perfect
decomposition for the two shapes. However, due to a lack of regularization on
the compositional modules, these modules modify the shape of the learned
identity alongside other variable factors, causing an imperfect decomposition
for shapes cube and capsule. With the introduction of the ID classifier, the
decomposition is improved significantly, where a large majority of the images
for each shape are passed through one module only, leaving $m_{3}$ with nearly
no images.

##  V Conclusion

In this work, we introduced a modular network-based model that is able to
decompose the generative factors of a synthetically generated dataset into
multiple compositional modules and static transformation matrices representing
a physical factor present in the ground truth factors of the dataset. This is
done by modifying the structure of the model and preventing a set of modules
from observing the input while imposing independence constraints for each
module separately and between different modules. Each identity is embedded in
a single transformation matrix modifiable by the learned compositional
modules.

## References

  * [1] L. Kirsch, J. Kunze, and D. Barber, “Modular networks: Learning to decompose neural computation,” in _NeurIPS_ , 2018.
  * [2] M. Chang, A. Gupta, S. Levine, and T. L. Griffiths, “Automatically composing representation transformations as a means for generalization,” in _International Conference on Learning Representations_ , 2019. [Online]. Available: https://openreview.net/forum?id=B1ffQnRcKX
  * [3] C. Rosenbaum, T. Klinger, and M. Riemer, “Routing networks: Adaptive selection of non-linear functions for multi-task learning,” in _International Conference on Learning Representations_ , 2018. [Online]. Available: https://openreview.net/forum?id=ry8dvM-R-
  * [4] X. Zhu, C. Xu, and D. Tao, “Commutative lie group vae for disentanglement learning,” 2021. [Online]. Available: https://arxiv.org/abs/2106.03375
  * [5] D. P. Kingma and M. Welling, _An Introduction to Variational Autoencoders_ , 2019.
  * [6] C. Burgess and H. Kim, “3d shapes dataset,” https://github.com/deepmind/3dshapes-dataset/, 2018.