# Conditional Injective Flows for Bayesian Imaging

AmirEhsan Khorashadizadeh, Konik Kothari, Leonardo Salsi, Ali Aghababaei
Harandi, Maarten de Hoop and Ivan Dokmanić  Manuscript received April 18,
2022; revised September 5, 2022 and January 1, 2023; accepted February 9,
2023.
AmirEhsan Khorashadizadeh and Ivan Dokmanić were supported by the European
Research Council Starting Grant 852821–SWING. Maarten de Hoop gratefully
acknowledges support from the Department of Energy under grant DE-SC0020345,
the Simons Foundation under the MATH + X program, and the corporate members of
the Geo-Mathematical Imaging Group at Rice University. (Corresponding author:
AmirEhsan Khorashadizadeh.) AmirEhsan Khorashadizadeh and Leonardo Salsi are
with the Department of Mathematics and Computer Science of the University of
Basel, 4001 Basel, Switzerland (e-mail<EMAIL_ADDRESS>leonardo.salsi@outlook.com). Konik Kothari is with the Department of
Electrical and Computer Engineering, the University of Illinois at Urbana-
Champaign, Urbana, IL 61801 USA (e-mail: kkothar3@illinois.edu). Ali
Aghababaei Harandi is with the Department of Electrical Engineering, Sharif
University of Technology, Tehran 14399-57131, Iran (e-mail:
aghababaei.ali94@gmail.com). Maarten V. de Hoop is with the Department of
Computational and Applied Mathematics, Rice University, Houston, TX 77005 USA
(e-mail: mdehoop@rice.edu). Ivan Dokmanić is with the Department of
Mathematics and Computer Science of the University of Basel, 4001 Basel,
Switzerland, and also with the Department of Electrical, Computer Engineering,
the University of Illinois at Urbana-Champaign, Urbana, IL 61801 USA (e-mail:
ivan.dokmanic@unibas.ch).
Our implementation is available at https://github.com/swing-
research/conditional-trumpets. 0000-0003-0660-6823 0000-0002-6333-0379
0000-0001-7132-5214

###### Abstract

Most deep learning models for computational imaging regress a single
reconstructed image. In practice, however, ill-posedness, nonlinearity, model
mismatch, and noise often conspire to make such point estimates misleading or
insufficient. The Bayesian approach models images and (noisy) measurements as
jointly distributed random vectors and aims to approximate the posterior
distribution of unknowns. Recent variational inference methods based on
conditional normalizing flows are a promising alternative to traditional MCMC
methods, but they come with drawbacks: excessive memory and compute demands
for moderate to high resolution images and underwhelming performance on hard
nonlinear problems. In this work, we propose C-Trumpets—conditional injective
flows specifically designed for imaging problems, which greatly diminish these
challenges. Injectivity reduces memory footprint and training time while low-
dimensional latent space together with architectural innovations like fixed-
volume-change layers and skip-connection revnet layers, C-Trumpets outperform
regular conditional flow models on a variety of imaging and image restoration
tasks, including limited-view CT and nonlinear inverse scattering, with a
lower compute and memory budget. C-Trumpets enable fast approximation of point
estimates like MMSE or MAP as well as physically-meaningful uncertainty
quantification.

###### Index Terms:

Bayesian Imaging, Amortized Variational Inference, Deep Generative Models,
Computed Tomography, Inverse Scattering, Normalizing Flows

##  I Introduction

In Bayesian modeling of computational imaging problems, we assume that the
(unknown) object of interest $x$ and the observed measurements $y$ are
realizations of random vectors $X\in\mathcal{X}$ and $Y\in\mathcal{Y}$ with a
joint distribution $p_{X,Y}$. This general model includes the common setting
of a deterministic forward operator and additive Gaussian noise,

$y=A(x)+\epsilon,$ (1)

where

$Y|X\sim\mathcal{N}(A(X),\sigma^{2}I)$ (2)

as well as other relevant models like $Y=\mathrm{Poisson}(A(X),\lambda)$ or
even random or uncertain forward operators $A$.

Common machine-learning approaches to ill-posed inverse problems yield point
estimates, that is, they output a single reconstruction. For example, training
a deep neural network $f_{\theta}$ with the mean-squared error (MSE) loss
$\mathbb{E}~{}\|X-f_{\theta}(Y)\|^{2}$ approximates the minimum-mean-squared-
error (MMSE) estimator111Often called the regression function in machine
learning and statistics. $\mathbb{E}[X|Y]$ (the posterior mean) [1].

In many situations, however, a single point estimate can be misleading or
incomplete. For example, in radio interferometric imaging which aims to
reconstruct astronomical images from radio telescope measurements, there can
be multiple solutions that fit the observed measurements; a now-famous example
is the imaging of a black hole [2]. This can happen for a variety of reasons,
all stemming from the ill-posedness of the imaging problem. The posterior may
be multimodal, in which case the MMSE estimator blends the modes together and
maximum a posteriori estimate (MAP) (or posterior mode)
$\operatorname*{arg\,max}_{x}p_{X|Y}(x|y)$, returns only one of the many
modes. Even when the posterior is unimodal, providing a point estimate when
measurements have low signal-to-noise ratio does not convey the amount of
uncertainty in the estimate, thus requiring cautious interpretation. As a
remedy, uncertainty quantification (UQ) on top of a reconstructed image can
greatly help medical professionals make more informed decisions or order
additional measurements in uncertain regions [3].

Figure 1: Travel-time tomography with boundary sensors; we measure the wave
travel time between sensors in Figure 6(c). The column $x$ shows the ground
truth, the column $y$ the pseudo-inverse of the forward operator applied to
the measurements. Conditional injective flows provide MMSE, surrogate MAP
(Section III-C) and physically meaningful uncertainty quantification (UQ).
Given no sensors in the top regions of the image, the trained C-Trumpet
assigns higher uncertainty to the top half of the domain. While both standard
point estimates look close to “9”, many posterior samples look like “4”s which
indicates the significance of posterior sampling for ill-posed inverse
problems.

In Figure 1 we use a toy problem to illustrate access to the posterior can
help in the presence of multiple modes. Both standard point estimates look
close to a “9”, but many posterior samples look like “4”s. While innocuous on
MNIST [4], such failure modes come with risks in medical imaging.

Another approach is then to characterize the posterior $p_{X|Y}$ upon
observing $y$. Given the joint distribution $p_{X,Y}$, computing the posterior
$p_{X|Y}$ generally involves intractable high-dimensional integrations. A
standard way to circumvent this is by sampling methods such as Markov chain
Monte Carlo (MCMC) [5, 6]. These methods work well in low-dimensional problems
but become computationally intensive when used in high-dimensional imaging
tasks due to the need to compute the forward operator many times [7]. An
alternative is to perform variational inference. This requires defining a
tractable, parameterized family of distributions $\mathcal{Q}$ and choosing a
$q\in\mathcal{Q}$ that is “close” to $p_{X|Y}$.

In this paper, we define a new class of approximate posteriors $\mathcal{Q}$
using injective deep neural networks, suitable for imaging problems. We build
on the recently-proposed injective flows [8, 9, 10] and conditional coupling
layers [11]. Injective flows map a low-dimensional latent space to a high-
dimensional image space using a sequence of injective functions whose inverses
(on the range) can be computed efficiently and exactly. They combine the
favorable aspects of GANs [12] and invertible normalizing flows [13, 14, 15]
in a way especially well-suited for imaging problems. Indeed, modern GANs with
the various innovations in architectures and training protocols generate
subjectively high-quality high-resolution images [16, 17] and can be trained
relatively straightforwardly with a reasonable compute budget but lack an
exact, fast inverse and access to (approximate) likelihoods. Normalizing flows
have fast inverses and enable fast and exact likelihood computation (although
this depends on the particular architecture), but they lack a low-dimensional
latent space thereby requiring large memory and compute budget when training
at higher image resolutions [18]. Injective flows have a low-dimensional
latent space with a fast and exact inverse on the range while maintaining a
low compute and memory budget.

_Conditional_ normalizing flows [11] inherit the favorable properties of their
non-conditional counterparts—easy access to the likelihoods and inverses of
generated samples. They were applied to inverse problems [19, 20] where they
enable posterior estimation, efficient sampling and even uncertainty
quantification. In this work, we propose conditional injective flows termed
C-Trumpets. While injective flows model image datasets supported on low-
dimensional manifolds, the range of a C-Trumpet is a (potentially) different
low-dimensional manifold of posterior samples _for each_ measurement. As we
show in Section IV, this makes them an effective model for data distributions
supported on _fiber bundles_ with the base space corresponding to the space of
measurements $\mathcal{Y}$, and in particular effective models of posterior
distributions in imaging inverse problems. Moreover, thanks to a low-
dimensional latent space, they can be trained for high-dimensional data
($256\times 256$) using a single GPU while training conditional bijective
flows at this resolution requires significantly more memory.

While generating approximate posterior samples is important, many applications
still call for Bayesian point estimates such as the MAP or the MMSE estimator,
and it is convenient if those can be computed with the same model. While it is
clear (at least conceptually) how to do it for the MMSE estimator—generate
many samples and average them—the MAP estimator is more elusive. There are not
many deep learning approaches to imaging which compute (or approximate) the
MAP estimator since the corresponding training loss (purely formally) is the
highly irregular $\delta(x-x^{\prime})$ as opposed to the “nice”
$\|x-x^{\prime}\|^{2}$ [21]. (A notable exception is amortized MAP for image
super-resolution [22], although it is limited to noiseless linear low-rank
projections.) We could in principle obtain a MAP estimate from C-Trumpets via
iterative maximization; however, that is slow and it is not guaranteed to
converge. We thus propose a modification of coupling layers which results in a
fixed volume change with respect to the input. We use this newly designed
layer to efficiently approximate MAP estimators without the need to run an
iterative process or evaluate the forward operator.

Our main contributions can be summarized as follows:

  * •

We define a class of deep learning models for amortized variational inference
called C-Trumpets, with a smaller memory and compute footprint compared to
bijective flows; C-Trumpets can be trained on $256\times 256$ images on a
single V100 GPU in a day.

  * •

The new flows include bespoke architectural innovations: fixed-volume-change
layers provide efficient MAP estimates without iterative optimization, while
skip connections improve the quality of the generated samples and uncertainty
quantification.

  * •

We show that C-Trumpets outperform conditional bijective flows in solving
computational imaging problems including nonlinear electromagnetic scattering,
limited-view CT and seismic travel-time tomography; C-Trumpets produce much
better posterior samples and uncertainty estimates that are consistent with
the physics of the forward operators in various inverse problems.

  * •

While standard injective flows parameterize manifolds and are thus a natural
(regularizing) choice when we believe the manifold assumption holds, we show
that conditional injective flows can parameterize fiber bundles [23].

##  II Variational Bayesian Inference

In this section we work formally and assume that all probability measures have
a density; this allows us to simply present the main ideas. A training
strategy suitable for distributions supported on low-dimensional manifolds is
detailed in Section III.

We consider random vectors $X\in\mathcal{X}$ (representing the unknown object)
and $Y\in\mathcal{Y}$ (representing the observed noisy measurements). The
posterior distribution $p_{X|Y}$ can be expressed as

$p_{X|Y}(x|y)=\dfrac{p_{Y|X}(y|x)p_{X}(x)}{\int_{x}p_{X,Y}(x,y)dx}.$ (3)

In high-dimensional imaging applications, calculating $\int_{x}p_{X,Y}(x,y)dx$
is intractable. Moreover, the prior distribution $p_{X}$ is unknown and needs
to be estimated [24, 25]. MCMC-based sampling methods do not require access to
$\int_{x}p_{X,Y}(x,y)dx$ but they are slow if one desires likely posterior
samples [26, 7]. An alternative to MCMC is to perform variational inference
[27, 28]: we define a parameterized class of distributions

$\mathcal{Q}=\\{q_{\theta}\in\mathcal{P}(\mathcal{X})\ :\ \theta\in\Theta\\},$

where $\mathcal{P}(\mathcal{X})$ is the space of probability distributions
over $\mathcal{X}$, and we search for $\theta$ such that $q_{\theta}$ is close
to $p_{X|Y}$. Examples of $\mathcal{Q}$ include Gaussian mixtures and
densities induced by generative neural networks.

The remaining ingredients are a measure of “closeness” and a fitting
algorithm. In standard variational inference it is common to use the
Kullback–Leibler (KL) divergence as a measure of fit,

$\displaystyle\text{KL}(q\|p)$
$\displaystyle=\int_{\mathcal{X}}q(x)\log\left(\dfrac{q(x)}{p(x)}\right)dx$
$\displaystyle=\mathbb{E}_{X\sim q}[\log q(X)-\log p(X)].$

Given a measurement $y$, we choose $q_{\theta}$ by solving either

$\theta_{\text{rev}}^{\ast}(y)=\operatorname*{arg\,min}_{\theta\in\Theta}~{}\text{KL}(q_{\theta}\|p_{X|Y}(\cdot\,|\,y)),$

or

$\theta_{\text{fwd}}^{\ast}(y)=\operatorname*{arg\,min}_{\theta\in\Theta}~{}\text{KL}(p_{X|Y}(\cdot\,|\,y)\|q_{\theta}),$

respectively called reverse and forward KL minimization [29, 6]. The two
estimates are in general different due to the asymmetry of the KL divergence.

Minimizing the reverse KL requires computing the expectation
$\mathbb{E}_{X\sim q_{\theta}}\log p_{Y|X}(y|X)p_{X}(X)$. While $p_{Y|X}$ is
usually asssumed known in imaging problems with a known forward operator—in
(2) it corresponds to additive Gaussian noise—the prior distribution $p_{X}$
is unknown. Recently, normalizing flows were used to estimate $p_{X}$ to then
allow for downstream minimization of the reverse KL divergence [2, 30, 10]. On
the other hand, minimizing the forward KL does not require access to the prior
distribution, the noise model or the forward operator[31, 32].

Computing $\theta^{*}(y)$ in both forward and reverse KL formulation involves
solving a separate optimization problem for every new measurement $y$. The
reason is that $q_{\theta}\in\mathcal{Q}$ is a function of $x$ alone, not $x$
and $y$, and we hope that for each $y$,

$q_{\theta^{*}(y)}(x)\approx p_{X|Y}(x|y).$

This separate optimization for every $y$ may be inefficient if implemented via
standard iterative solvers.

We could, however, work directly with a family of conditional distributions
$q_{\theta}(x|y)$ which depend on both $x$ and $y$,

$\mathcal{Q}_{\text{cond}}=\\{(x,y)\mapsto q_{\theta}(x|y)\ :\
\theta\in\Theta_{\text{cond}}\\},$

where for each $y$, $q_{\theta}(\cdot|y)\in\mathcal{P}(\mathcal{X})$ is a
probability distribution over $\mathcal{X}$.

We can now compute a conditional variational approximator $q_{\theta}(x|y)$ by
minimizing the average KL divergence over all measurements $y$: this procedure
is known as amortized inference. It leads to the following optimization
problem:

$\displaystyle\theta^{*}$
$\displaystyle=\operatorname*{arg\,min}_{\theta}~{}\mathbb{E}_{Y\sim
p_{Y}}\text{KL}(p_{X|Y}(\,\cdot\,|Y)\|q_{\theta}(\,\cdot\,|Y))$ (4)
$\displaystyle=\operatorname*{arg\,max}_{\theta}~{}\mathbb{E}_{X,Y\sim
p_{X,Y}}\log q_{\theta}(X|Y).$

The key observation is that the population expectation over $p_{X,Y}$ can now
be approximated by an empirical expectation over the training data
$\\{(x^{(i)},y^{(i)})\\}_{i=1}^{N}$.

The question that remains is: how to parameterize $q_{\theta}(x|y)$ so that we
can 1) learn $\theta^{*}$ efficiently from data, 2) easily obtain conditional
samples from $q_{\theta^{*}}(x|y)\approx p_{X|Y}(x|y)$ for a given $y$, and 3)
efficiently compute standard point estimators such as the MAP estimator? We
answer this question in the remainder of the paper by describing conditional
injective flows called C-Trumpets.

In a nutshell, we will define a family of neural networks $f_{\theta}(Z;y)$
where $y$ is the conditioning input. The first argument, $Z$, will be a
standard Gaussian random vector over a low-dimensional latent space, and the
parameter $\theta$ adjusted so that for each $y$, the random vector
$f_{\theta}(Z;y)$ is close in distribution to $X|Y\\!\\!=\\!\\!y$. In other
words, denoting the standard Gaussian distribution by $p_{Z}$, we will obtain
$q_{\theta}(\cdot~{}|~{}y)$ as a pushforward of $p_{Z}$ via
$f_{\theta}(z;y)$,222For simplicity we lightly abuse notation by identifying a
probability measure and its density.

$q_{\theta}(\cdot~{}|~{}y)=\left[f_{\theta}(\,\cdot\,;y)\right]_{\\#}p_{Z}.$
(5)

The approximate posterior samples can then be obtained in a standard way as
$f_{\theta}(Z;y)$.

##  III C-Trumpets: Conditional Injective Flows

Figure 2: Conditional injective flows. Different measurements  $y_{1}$,
$y_{2}$ and  $y_{3}$ give different manifolds.

For many natural and medical image classes, the posterior distribution
$p_{X|Y}$ is low-dimensional as its support is a subset of the support of the
prior distribution $p_{X}$ which is assumed to concentrate close to a low-
dimensional manifold [33, 34]. C-Trumpets (Figure 2) are conditional injective
normalizing flows that map a low-dimensional latent space to a high
dimensional data space using an injective transformation for each conditioning
sample. Injectivity guarantees that for each conditioning sample the range of
the network is a manifold. Moreover, the efficiently-invertible layers
facilitate training and inference. As shown in Figure 2, the proposed model
has two subnetworks: an injective generator $g_{\gamma}$ that maps a low-
dimensional space $\mathcal{Z}=\mathbb{R}^{d}$ to the data space
$\mathcal{X}=\mathbb{R}^{D}$, $d\ll D$, and a bijective mapping
$h_{\eta}:\mathcal{Z}\to\mathcal{Z}$ maintaining dimensionality. The end-to-
end mapping is then given as,

$f_{\theta}(z;y)=g_{\gamma}(h_{\eta}(z;y);y)$

with learnable parameters $\theta=(\gamma,\eta)$.

C-Trumpets are inspired by non-conditional injective flows [10] called
Trumpets. They comprise a sequence of injective [10] and bijective [15] revnet
blocks. Each revnet block consists of three components: activation
normalization, (injective or bijective) $1\times 1$ convolution and affine
coupling layer. It is worth noting that all these components are non-
conditional, and we will later use conditional affine coupling layers to make
the generation process conditional.

  1. 1.

activation normalization,

$\begin{array}[]{ll}\text{{Forward:}}&~{}x=\dfrac{z-\mu}{\sigma}\\\
\text{{Inverse:}}&~{}z=\sigma x+\mu\end{array}$ (6)

  2. 2.

$1\times 1$ convolution with a kernel $w$,

    1. (a)

Bijective version:

$\begin{array}[]{ll}\text{{Forward:}}&~{}x=w\ast z\\\
\text{{Inverse:}}&~{}z=w^{-1}\ast x\end{array}$ (7)

    2. (b)

Injective version:

$\begin{array}[]{ll}\text{{Forward:}}&~{}x=w\ast z\\\
\text{{Inverse:}}&~{}z=w^{\dagger}\ast x\end{array}$ (8)

$w$ is a $1\times 1$ convolutional filter, which is simply a matrix
multiplication along the channel dimension where $w\in\mathbb{R}^{c_{in}\times
c_{out}}$ and $w^{\dagger}$ is the pseudo-inverse of $w$ (a non-square matrix
in the injective dimension-expanding case). We note that $w$ should obey the
appropriate constraints to guarantee injectivity [35]. We use LU decomposition
in matrix inversion, which significantly reduces training time of the
injective part of C-Trumpets (see Section S-I for more details).

  3. 3.

affine coupling layer

$\begin{array}[]{lll}\text{{Forward:}}&x_{1}=z_{1},\quad x_{2}=s(z_{1})\circ
z_{2}+b(z_{2})\\\ \text{{Inverse:}}&z_{1}=x_{1},\quad
z_{2}=s(x_{1})^{-1}\circ(x_{2}-b(x_{1})),\end{array}$

where $z=[z_{1},z_{2}]^{T}$ and $x=[x_{1},x_{2}]^{2}$. The mappings $s$ and
$b$ are respectively the scale and the shift networks.

In order to make the generative process conditional, we adapt the conditional
affine coupling layers proposed in [11]. Conditional affine coupling layers
keep the advantages of a regular flow model—fast inverses and tractable
Jacobian computations, while benefiting from the conditioning framework. Since
scale $s(\,\cdot\,)$ and shift $b(\,\cdot\,)$ networks are never inverted, we
can concatenate the features of the conditioning sample $y$ to their input
without losing invertibility and tractable $\log\det$ Jacobian computation.
Accordingly, $s(\,\cdot\,)$ and $b(\,\cdot\,)$ are replaced by
$s(\,\cdot,c_{\varphi}(y))$ and $b(\,\cdot,c_{\varphi}(y))$,

$\begin{array}[]{lll}\text{{Forward:}}&x_{1}=z_{1}\\\
&x_{2}=s(z_{1},c_{\varphi}(y))\circ z_{2}+b(z_{2},c_{\varphi}(y))\\\
\text{{Inverse:}}&z_{1}=x_{1}\\\
&z_{2}=s(x_{1},c_{\varphi}(y))^{-1}\circ(x_{2}-b(x_{1},c_{\varphi}(y))),\end{array}$

where $z=[z_{1},z_{2}]^{T}$, $x=[x_{1},x_{2}]^{T}$ and
$c_{\varphi}(\,\cdot\,)$ represents the conditioning network that extracts
appropriate features from $y$. We deploy conditional affine coupling layers in
both injective and bijective subnetworks of C-Trumpets.

###  III-A Conditioning network

The role of the conditioning network $c_{\varphi}(\,\cdot\,)$ is to extract
features from conditioning samples $y$ to be used by the affine coupling
layers. The architecture of the conditioning network depends on the nature of
the conditioning samples. When $y$ is structured as an image, we use
convolutional layers; when it is an unstructured 1D vector, we use fully
connected layers, as for example in class-based image generation (see Section
S-III-A in the supplemental materials) where the conditioning data are one-hot
class encodings. The output dimension of the conditioning network is set to
match the input dimension of the scale and shift modules of the coupling
layers. The weights of the conditioning networks are trained jointly with the
remaining parameters in C-Trumpets via back-propagation using paired training
data. Further details about conditioning networks in all numerical experiments
are given in Sections S-II-B and S-III-A in the supplemental materials.

###  III-B Skip connections

The fact that we use expanding revnet layers allows us to augment the
architecture with elements that are not compatible with the standard injective
layers. The conditioning samples often have a pixel-wise dependency on target
signals. For example, in image inpaiting with a fixed mask location, the out-
of-mask pixels should be simply forwarded to the output. In order to not re-
learn these features, we introduce skip connections after every revnet block
between the measurements, $y$ and the different resolution levels of the
injective subnetwork of C-Trumpets,

$\begin{array}[]{ll}\text{{Forward:}}&~{}x=({\mathbf{1}}-S)\circ\text{rev}(z;y)+S\circ\text{resize}(y)\\\\[5.69054pt]
\text{{Inverse:}}&~{}z=\text{rev}^{-1}\left(\dfrac{x-S\circ\text{resize}(y)}{{\mathbf{1}}-S};y\right),\end{array}$
(9)

where $\text{rev}(\,\cdot\,)$ is the conditional revnet block, $S$ is a
learnable matrix with entries between $0$ and $1$; $S$ adjusts the amount of
the direct mixing of the measurements in the generated posterior sample. We
empirically find that skip connections help the injective part of C-Trumpets
converge faster and they yield better reconstruction in several imaging
problems.

###  III-C Fast MAP estimation via fixed volume-change layers

MAP estimation traditionally requires an iterative solution of a maximization
problem. Deep neural networks for image reconstruction are traditionally
trained with an $\ell^{p}$ loss ($p=2$ giving the ubiquitous MSE loss and
ultimately an approximation of the MMSE estimator $\mathbb{E}[X|Y]$). There
are, however, few attempts at using deep neural networks to approximate the
MAP estimate in imaging, possibly because the associated “loss” would be a
tempered distribution $\delta(x-x^{\prime})$ [21]. While the MAP and the MMSE
estimates coincide when $X$ and $Y$ are jointly Gaussian, they are in general
different. For posteriors supported on a low-dimensional manifold the MMSE
estimate will generally not lie on the manifold. (We show a qualitative
manifestation of this effect in Section V-D)

A notable deep-learning approach to amortized MAP inference for image super-
resolution has been proposed by Sønderby et al. [22]. However, their method is
only applicable for super-resolution in the ideal noise-free scenario. In this
section, we propose a new variant of affine coupling layers, which enables us
to obtain the MAP estimate instantaneously with a single forward pass of the
trained network. This method is exact when used with bijective flows and
approximate for injective flows where it computes the MAP estimate of the pre-
image samples $z^{\prime}$ in Figure 2.

Consider a trained conditional normalizing flow model $x=f(z;y)$. The MAP
estimate can be obtained by solving

$\begin{array}[]{rcl}x_{\text{MAP}}&=&\displaystyle\operatorname*{arg\,max}_{x}~{}\log(p_{X|Y}(x|y))\\\
&=&\displaystyle\operatorname*{arg\,max}_{x}~{}\log(p_{Z}(z))-\sum_{k=1}^{K}\log|\det
J_{f^{k}}|,\end{array}$ (10)

where $z=f^{-1}(x;y)$. In principle, we can run an iterative maximization
process to compute $x_{\text{MAP}}$. In general, this may be slow and it is
not guaranteed to converge, even with multiple random restarts.

Let us analyze (10) more closely. While the first term, $\log(p_{Z}(z))$, has
the highest value at $z=\mu_{z}$ (the mean of the Gaussian), the second term
has three components: activation normalization, $1\times 1$ convolution, and
the conditional affine coupling layer. The first two components are linear
layers with data-independent Jacobians; their $\log\det$s are thus constant
with respect to $x$ and can be omitted from (10). The $\log\det$ of the
Jacobian of conditional affine coupling layers is

$\log(|\det J(z)|)=\sum_{i=1}^{l}\log(s^{i}(z_{1},c_{\varphi}(y))),$ (11)

where $s^{i}(\,\cdot\,)$ is the $i$th element of the output of the scale
network. This term is in general data-dependent. In order to make it data-
independent, we propose to use the following activation for the scale network,

$s_{\text{FVC}}(z,c_{\varphi}(y))=\exp(\text{softmax}(m(z,c_{\varphi}(y)))),$
(12)

where $m(\,\cdot\,)$ is an arbitrary neural network and $\text{softmax}(x)$ is
defined as

$\text{softmax}(x)_{i}=\dfrac{e^{x_{i}}}{\sum_{j=1}^{l}e^{x_{j}}}\quad\text{for}\quad
i=1,2,...,l.$

Then we have,

$\log(|\det
J_{\text{FVC}}(z)|)=\sum_{i=1}^{l}\log(s_{\text{FVC}}^{i}(z_{1},c_{\varphi}(y)))=1.$
(13)

The newly proposed layer has a data-independent $\log\det$ Jacobian, without
losing expressivity, as verified empirically in Sections V-B and V-C.

Now all terms in (10) are independent of $x$, so that

$\begin{array}[]{rcl}x_{\text{MAP}}&=&\displaystyle\operatorname*{arg\,max}_{x}~{}\log(p_{X|Y}(x|y))\\\
&=&\displaystyle\operatorname*{arg\,max}_{x}~{}\log(p_{Z}(z))\\\
&=&f(\mu_{z};y);y),\end{array}$ (14)

which is to say that the MAP estimate is obtained simply by feeding the mean
of the Gaussian base distribution into the (bijective) flow.

While the proposed technique is exact for bijective conditional flows, however
the $\log\det$ term for conditional _injective_ flows is given as

$\log p_{X|Y}(x|y)=\log p_{Z}(f_{\theta}^{\dagger}(x;y))\\\
-\dfrac{1}{2}\log|\det[J_{f_{\theta}}(f_{\theta}^{\dagger}(x;y))^{T}J_{f_{\theta}}(f_{\theta}^{\dagger}(x;y))]|.$
(15)

and equation (15) cannot be decomposed into a sum of the $\log\det$s of its
constituent components. We can nevertheless use this technique to obtain a MAP
estimate in the intermediate $z^{\prime}$-space as in Figure 2. Therefore, we
obtain a surrogate of the end-to-end MAP estimate and call it surrogate MAP
and denote it as $g^{\dagger}$-MAP.

###  III-D Training strategy

Thanks to simple, tractable inverses and $\log\det$s of layer Jacobians, the
parameters of normalizing flows can be fitted via maximum likelihood (ML).
However, as C-Trumpets have a low-dimensional latent space, the likelihood is
only defined in the range of the injective generator. Since the weights of the
networks are initialized randomly, prior to training this range does not
contain the targets in the training set and their end-to-end likelihoods are
not defined. We thus split the training of C-Trumpets,
$f_{\theta}(z;y)=g_{\gamma}(h_{\eta}(z;y);y)$, into two phases, following the
method of Brehmer and Cranmer [9] for non-conditional injective flows: (1) The
MSE training phase where we only fit the trainable parameters $\gamma$ of the
injective network with the goal of adjusting the range of $f_{\theta}$ to
contain the training data, and (2) The ML training phase where we optimize the
trainable parameters $\eta$ of the low-dimensional bijective part, maximizing
the likelihood of the pre-image (through $g_{\gamma}$) of the training
data.333For simplicity we denote all trainable parameters of the injective
network, _including the weights of the conditioning networks and the skip
connections_ by $\gamma$, and all trainable parameters of the bijective
network, _including the weights of the conditioning networks_ by $\eta$.

In the first phase, we optimize $\gamma$ by minimizing

$\mathcal{L}_{\text{MSE}}(\gamma)=\\\
\dfrac{1}{N}\sum_{i=1}^{N}\|x^{(i)}-g_{\gamma}(g_{\gamma}^{\dagger}(x^{(i)};y^{(i)});y^{(i)})\|_{2}^{2},$
(16)

where $\\{(x^{(i)},y^{(i)})\\}_{i=1}^{N}$ are the training samples and
$g^{\dagger}$ is the layer-wise inverse of the injective subnetwork;
therefore, the projection operator on the range of $g_{\gamma}(\,\cdot\,;y)$
is given as $P_{g_{\gamma}}(x;y):=g_{\gamma}(g_{\gamma}^{\dagger}(x;y);y)$.
After training the injective network for a fixed number of epochs, the range
of the network approximately contains the training data.

We now switch to the second phase: maximizing the likelihood of the projected
training samples in the intermediate space ($z^{\prime}$ in Figure 2),
$\\{g_{\gamma}^{\dagger}(x^{(i)})\\}_{i=1}^{N}$, by minimizing the following
KL divergence over $\eta$ (cf. Section II):

$\mathcal{L}_{\text{ML}}(\eta)=\dfrac{1}{N}\sum_{i=1}^{N}\left(-\log
p_{Z}(z^{(i)})+\log|\det J_{h_{\eta}}|\right),$ (17)

where $z^{(i)}=h_{\eta}^{-1}(g_{\gamma}^{\dagger}(x^{(i)};y^{(i)});y^{(i)})$
and $p_{Z}(z)$ is a standard Gaussian distribution in $\mathbb{R}^{d}$. In
summary, this training strategy first ensures that the range of the injective
generator “interpolates” the training data and then maximizes the intermediate
space likelihoods as proxies to the image-space likelihoods. After training,
sampling an approximate posterior sample for a given $y$ is performed by
sampling a $z$ from a normal distribution and using the forward pass:
$x_{\text{gen}}=g_{\gamma}(h_{\eta}(z;y);y)$.

##  IV The C-Trumpets Signal Model

In this section we briefly discuss the geometric and topological aspects of
C-Trumpets. Readers who care mostly about the practical aspects and numerical
results may safely skip ahead to Section V.

The primary motivation for introducing (non-conditional) injective flows is to
model data supported on low-dimensional manifolds [35, 9, 10, 36]. It was
empirically shown that for common structured datasets these models are indeed
simpler, faster to train and generate higher quality samples than globally
invertible normalizing flows which maintain the same dimension across all
layers.

On the other hand, we introduce C-Trumpets from the point of view of
uncertainty quantification and posterior sampling rather than geometrical and
topological considerations about the class of signals it models. It is
nevertheless important to make this class explicit, since this understanding
will guide design choices and circumscribe the range of problems in which
C-Trumpets are the right tool of choice.

For every (fixed) conditioning sample $y$, a C-Trumpet becomes a conditional
flow modeling the corresponding conditional distribution. Thus for every
conditioning sample $y$ the range is a low-dimensional manifold

$\mathcal{M}_{y}=\\{f_{\theta}(z;y)\ :\ z\in\mathcal{Z}\\}$

with topology induced by the topology of the latent space (the support of
$p_{Z}$). The range of a C-Trumpet as a function of $z$ and $y$ then
corresponds to the union of the family of these manifolds indexed by $y$,

$\mathcal{R}_{\theta}=\\{f_{\theta}(z;y)\ :\
z\in\mathcal{Z},y\in\mathcal{Y}\\}=\bigcup_{y\in\mathcal{Y}}\mathcal{M}_{y}.$

For a given $(y,z)$ let

$\beta(y)=\\{x\ :\ x=f_{\theta}(z;y),\ z\in\mathcal{Z}\\},$

denote the set of all possible signals $x$ that could have caused the
observation $y$, or in other words the support of the distribution
$p_{X|Y=y}$.

Let us now assume a slightly stronger condition than injectivity of
$f_{\theta}$ in the first argument: that the Jacobian (with respect to both
$z$ and $y$) of $f_{\theta}$ has full rank
$\dim(\mathcal{Z})+\dim(\mathcal{Y})$. Note that this does not guarantee
global injectivity of the map $(z,y)\mapsto f_{\theta}(z;y)$, but it does
imply that there is a (finite) collection of open sets $\\{U_{i}\\}$ (thought
of as neighborhoods in $\mathcal{Y}$) which cover $\mathcal{Y}$,
$\bigcup_{i}U_{i}=\mathcal{Y}$, such that $f_{\theta}$ is injective on
$\mathcal{Z}\times U_{i}$ for all $i$ and in fact a homeomorphism between
$\mathcal{Z}\times U_{i}$ and $\beta(U_{i})$. The signal model we just
described is called a fiber bundle [23, 37].

Formally, a generic fiber bundle comprises sets
$(\mathcal{E},\mathcal{B},\mathcal{Z})$ called the total space
($\mathcal{E}$), the base space $(\mathcal{B})$, and the (typical) fiber
$(\mathcal{Z})$, together with a surjective map
$\pi:\mathcal{E}\rightarrow\mathcal{B}$, with the property that for each
$b\in\mathcal{B}$, there exists an open neighborhood $b\in
U_{i}\subseteq\mathcal{B}$ on which $\phi:\mathcal{Z}\times
U_{i}\rightarrow\pi^{-1}(U_{i})$ is a homeomorphism. The definition can be
illustrated by a commutative diagram [23],

${\pi^{-1}(U_{i})}$${\mathcal{Z}\times
U_{i}}$${U_{i}}$$\scriptstyle{\pi}$$\scriptstyle{\phi^{-1}}$$\scriptstyle{\phi}$proj
(18)

where $\text{proj}:\mathcal{Z}\times U_{i}\to U_{i}$. We have the following
correspondence between the standard fiber bundle notation and the notation
specific to C-Trumpets used in this paper:

$\begin{array}[]{rcl}\textbf{C-Trumpets}&&\textbf{Fiber
bundles}\\\\[5.69054pt]
\mathcal{R}_{\theta}\phantom{ttt}&\longleftrightarrow&\phantom{ttttt}\mathcal{E}\\\
\mathcal{Z}\phantom{ttt}&\longleftrightarrow&\phantom{ttttt}\mathcal{Z}\\\
\mathcal{Y}\phantom{ttt}&\longleftrightarrow&\phantom{ttttt}\mathcal{B}\\\
\beta\phantom{ttt}&\longleftrightarrow&\phantom{ttttt}\pi^{-1}\\\
f_{\theta}\phantom{ttt}&\longleftrightarrow&\phantom{ttttt}\phi.\end{array}$

Fiber bundles model a variety of situations in imaging, scientific computing,
and physics. In CryoEM we can see the 3D rotation group as a bundle of circles
over the 2-sphere, which leads to denoising algorithms [38, 39]. Symmetries
for the forward operator (for example in phase retrieval) naturally split the
data space into equivalence classes that can be modeled as fiber bundles. More
generally, any many-to-one machine learning task is well-modeled by a fiber
bundle [23]. We also mention natural connections with gauge equivariance and
geometric deep learning [40].

That said, our aim here is to identify a natural assumption we make when
modeling data with a conditional generative model, analogously to how
unconditional generative models are naturally compatible with data that lives
close to a low-dimensional manifold. Continuing the analogy, just like
manifolds are used to regularize ill-posed inverse problems, we imagine fiber
bundle models can be used as regularizers with additional structure even when
the data only approximately lives on a fiber bundle.

Fiber bundles model spaces that are locally product spaces but may have non-
trivial topology globally. An example of a space that is globally a product
space is the cylinder $\mathcal{E}=S^{1}\times[0,1]$ with the base space being
the circle $\mathcal{B}=S^{1}$ and all fibers being translates of the line
segment $\mathcal{Z}=[0,1]$. The projection $\pi$ associates to each point
$x\in\mathcal{E}$ its position along the base circle; $\pi^{-1}$ takes a
position along the base circle and returns the corresponding line segment. The
same $\mathcal{B}$ and $\mathcal{Z}$ can generate a rather different space—the
Möbius band—which is globally not a product space.444To explain this
mathematically we would need to introduce the transition maps and the
fundamental group, which is beyond the scope of our sketch.

Thus (under appropriate conditions) the range of a C-Trumpet is locally
homeomorphic to a product space (it is a product space up to a stretch), even
though it globally need not be depending on the topologies of $\mathcal{Z}$
and $\mathcal{Y}$. The requirement that $\phi$ be a homeomorphism on $U_{i}$
implies that (locally) the fibers do not intersect. This can of course only
hold if $\dim(\mathcal{Y})+\dim(\mathcal{Z})\leq\dim(\mathcal{X})$ and the
appropriate conditions on the Jacobian are satisfied, but in any case it gives
a useful intuition for the kind of problems and datasets that admit modeling
by C-Trumpets.

We show that C-Trumpets can indeed model simple low-dimensional fiber bundles:
an embedded solid555By a “solid” torus we refer to a bundle whose fibers are
disks. Modeling a standard torus is challenging with coupling layers which can
only be defined starting in dimension 2. Even in this case their expressivity
is limited so this low-dimensional example is only meant as illustration.
torus in Figure 3(a) and a solid Möbius band with elliptic cross-sections
(fibers) in Figure 3(b). As shown in Figure 3, the conditioning samples for
both the torus and the elliptic Möbius are taken to live on the base circle
($y\in[0,2\pi)$). For each angle $y$, C-Trumpets generate samples from a
distribution on the disk for the torus and an elliptical disk for the elliptic
Möbius band. Accordingly, we train C-Trumpets with latent dimension two.
Figure 4 demonstrates the generated samples by C-Trumpets; we sample the
resulting models in increments of $6^{\circ}$ to make the fiber bundle
structure clear.

(a) Torus

(b) Elliptic Möbius

Figure 3: Two examples of fiber bundles: Torus and elliptic Möbius as fiber
bundles over the base circle.

(a) Torus

(b) Elliptic Möbius

Figure 4: Samples on fiber bundles generated by a C-Trumpet; in both (a) and
(b) the base manifold $S^{1}$ is sampled every $6^{\circ}$.

##  V Experimental Results

We start by showcasing how C-Trumpets provide MMSE, MAP and uncertainty
estimates in a variety of imaging inverse problems. The MMSE estimate,
$\operatorname{\mathbb{E}}[X|Y=y]$, is approximated by averaging a fixed
number of posterior samples from a C-Trumpet fitted to the training data. The
MAP estimate can be efficiently approximated using fixed volume-change layers
in Section III-C, without iterative maximization over $x$. Finally, we compute
the uncertainty estimate (UQ) through a simple pixel-wise standard deviation
as

$\begin{array}[]{rcl}\widehat{\text{MMSE}}&=&\displaystyle\dfrac{1}{K}\sum_{k=1}^{K}f_{\theta}(z_{k};y),\\\
\text{UQ}&=&\displaystyle\sqrt{\dfrac{1}{K}\sum_{k=1}^{K}|f_{\theta}(z_{k};y)-\widehat{\text{MMSE}}|^{2}},\end{array}$
(19)

where $|\cdot|^{2}$ is applied to each pixel. In all experiments, we use
$K=25$ posterior samples to approximate the MMSE and UQ estimates (we find
that the quality of MMSE and UQ saturate for a higher number of posterior
samples).

We experiment with limited-view computed tomography (CT), nonlinear
electromagnetic inverse scattering [41], and (linearized) seismic travel-time
tomography [42]. Additional evaluations on popular compressed-sensing-style
benchmarks—denoising, inpainting, super resolution and random masking—are
given in Section V-C.

We compare C-Trumpets with two different types of conditional bijective flows,
C-INN [11] and C-Glow [43]. C-INN used conditional coupling layers for the
first time, while C-Glow has a rather different architecture by using two
parallel bijective flows for simultaneously modeling the target and
conditioning samples. In order to emphasize the importance of an expansive
injective model over a bijective one, we build a comparison baseline model,
C-Rev, which consists of the bijective portion of C-Trumpets with latent-space
dimension equaling that of the image data, i.e., $h_{\eta}(z,y)$. Finally, in
Section V-D, we visually compare the MMSE and MAP estimates for several ill-
posed inverse problems.

###  V-A Gaussian Random Fields

We choose realizations of Gaussian random fields (GRF) as instances of
$p_{X}(x)$ where we have access to exact posterior distributions. We compare
the quality of the posterior approximated by C-Trumpets with the true
posterior obtained by analytical solution. The closed-form posterior
distribution $p_{X|Y}(x|y)$ where $y=Ax+n$,
$x\sim\mathcal{N}(\mu_{x},\Sigma_{x})$ and $n\sim\mathcal{N}(0,\lambda^{2}I)$
as follows,

$\begin{array}[]{cc}&p_{X|Y}(x|y)=\mathcal{N}(\mu_{x|y},\Sigma_{x|y}),\\\
&\mu_{x|y}=\mu_{x}+\Sigma_{x}A^{T}(A\Sigma_{x}A^{T}+\lambda^{2}I)^{-1}(y-A\mu_{x}),\\\
&\Sigma_{x|y}=\Sigma_{x}-\Sigma_{x}A^{T}(A\Sigma_{x}A^{T}+\lambda^{2}I)^{-1}A\Sigma_{x}.\end{array}$
(20)

We consider GRFs in resolution $64\times 64$, $\lambda=5\times 10^{-3}$, and
let $A$ be the mask operator that replaces a $32\times 32$ patch at the center
of the image with zeros. We train C-Trumpets over $60000$ training samples. In
Figure 5 (second to fifth columns), we show the MMSE estimate followed by
three posterior samples, and finally the UQ obtained from both the analytical
solution (20) and C-Trumpets. MMSE and UQ are computed over 500 random
posterior samples. We compare the MMSE estimates obtained from C-Trumpet and
the analytical solution (20) against the ground truth. In Table I we find that
C-Trumpet performs similar to the analytical solution in both SNR and SSIM
metrics. Since the analytical posterior distribution is Gaussian, MAP and MMSE
estimators are equal. We find that the MAP and MMSE estimates obtained from
C-Trumpet are quite close: an SNR of 44.43dB hinting at the efficacy of our
way of computing MAP estimates.

Figure 5: Performance comparison between C-Trumpet and analytical solution
(true posterior) (20) in mask problem where the target signals come from a
Gaussian distribution; C-Trumpets present MMSE and UQ close to the analytical
solution. TABLE I: Performance comparison of MMSE estimates between C-Trumpets
and the analytical solution obtained from 20; C-Trumpets presents MMSE
estimates close to the analytical solution, which shows the effectivity of the
proposed method.

| SNR (dB) | SSIM
---|---|---
C-Trumpets | 19.71 | 0.89
Analytical solution | 20.22 | 0.91

###  V-B Computational imaging

(a) CT vertical

(b) CT horizontal

(c) Travel time

(d) Scattering

Figure 6: Sensing geometry for the various imaging problems explored in
Section V. The red regions show the locations of the sensors. In the case of
CT ((a) and (b)), the line segments indicate that “sensors” measure entire
projection images as opposed to complex scalars in nonlinear scattering (d).

#### Limited-view CT

CT is a major medical imaging modality. The corresponding inverse problem is
to recover an image from its integrals along straight lines, arranged in a so-
called sinogram. In limited-view CT (cf. cryo-electron tomography [44] and
dental CT [45]), a contiguous cone of angles is missing from the acquisition;
Figures 6(a) and 6(b), illustrate vertical and horizontal missing cones of
$60^{\circ}$. As there are no measurements in a vertical (horizontal) missing
cone, we should expect higher uncertainty in horizontal (vertical) components.
We use the filtered back projection (FBP) reconstruction as our measurements
$y$, and train on 40000 $256\times 256$ samples from the LoDoPaB-CT [46]
dataset. The measurement SNR is set to 40 dB. In Figure 7, we show posterior
samples, $g^{\dagger}$-MAP, MMSE, and uncertainty quantification (UQ)
estimates from C-Trumpets. The real-space UQ estimate shows higher uncertainty
in the vertical (horizontal) components where we have horizontal (vertical)
missing cone. This is consistent with our expectations from the physics of CT.

We further show the UQ estimate in the Fourier domain computed by averaging
over the individual DFT bins. Quite pleasingly, this estimate aligns perfectly
with the theoretical prediction from the Fourier slice theorem; higher
uncertainty (bright regions) are indeed inside the missing cone while it is
worth emphasizing that C-Trumpets are not specifically designed for the
forward operator of the CT problem. Moreover, as expected, there is higher
uncertainty in higher frequency components (see Figure S3 in the supplemental
materials for more samples). We emphasize that C-Trumpets are the only
conditional generative architecture that gives such physically meaningful
posterior samples.

It is also worth mentioning that C-Trumpets model this $256\times 256$
resolution dataset with only 10M trainable parameters in less than 24 hours of
training time on a single NVIDIA V100 GPU. Due to memory constraints, we could
simply not train the baseline bijective models C-INN, C-Glow, and C-Rev on
$256\times 256$ images. The large latent space dimension of these models leads
to very large memory footprints.

We can still compare our uncertainty estimates with these models at a lower
resolution of $64\times 64$. Figure 8 shows such an experiment at the SNR of
25dB. In the first panel (second to fifth row), we show MMSE estimates of
different flow models and three random posterior samples. We further provide
UQ estimates in real and Fourier space. As we can see, not only do C-Trumpets
outperform the conditional bijective flows in terms of reconstruction quality
(MMSE estimate) (see also Table II) but they also give a more meaningful
uncertainty estimate even in high noise. Bijective models generate significant
uncertainty outside the missing cone, indicating that they do not learn a
reconstruction map consistent with the physics of CT.

(a) Vertical missing cone ($60^{\circ}$ to $120^{\circ}$)

(b) Horizontal missing cone ($-30^{\circ}$ to $30^{\circ}$)

Figure 7: Limited-view CT in resolution $256\times 256$. The frequency-domain
uncertainty estimate shows the log-scale standard deviation of DFT-bins.
Alignment with the theoretical prediction from the Fourier slice theorem
signifies physically-meaningful uncertainty quantification; higher uncertainty
(bright regions) inside the missing cone.

(a) Vertical missing cone ($60^{\circ}$ to $120^{\circ}$)

(b) Horizontal missing cone ($-30^{\circ}$ to $30^{\circ}$)

Figure 8: Performance comparison in limited-view CT with resolution $64\times
64$. The frequency-domain uncertainty estimate shows the log-scale standard
deviation of DFT-bins. The estimated $\log(\text{fft}(\text{UQ}))$ by
C-Trumpets aligns with the theoretical prediction from the Fourier slice
theorem; higher uncertainty (bright regions) inside the missing cone. Figure
9: Performance of C-Trumpets on nonlinear inverse scattering with
$\epsilon_{r}=6$ with a full-view sensing geometry. Conditioning inputs are
scattered fields.

#### Electromagnetic inverse scattering

We consider the non-linear electromagnetic inverse scattering problem as
described in [41]. We consider reconstruction of the finite number of
parameters from the scattered fields. Although inverse scattering becomes
well-posed and Lipschitz stable with continuous measurements [47]666This paper
considers 2D inverse scattering problem, 3D case is addressed in [48], it is
an ill-posed inverse problem with finite number of measurements, which means
that noise in the measurements may translate to exponential errors in the
reconstruction [41]. Moreover, the problem becomes increasingly non-linear as
the relative permittivity of the objects being imaged, $\epsilon_{r}$,
increases. In the first experiment, we use 36 incident plane waves and 36
receivers, distributed uniformly around the object with maximum permittivity
of $\epsilon_{r}=6$ and dimension 20 cm $\times$ 20 cm. We work at the
frequency of 3 GHz. and simulate the measurements by solving the Helmholtz
equation explained in Section S-II-C in the supplemental materials. We add
noise to the measurements for a target measurement SNR of 30 dB. We build a
dataset of $64\times 64$ images of overlapping ellipses with 60000 samples.

Figure 9 shows the performance of C-Trumpets where the scattered fields are
used as conditioning samples. This experiments clearly shows that C-Trumpets
can generate meaningful posterior samples, even for a highly non-linear
($\epsilon_{r}=6$) and ill-posed problem. Additional experiments with
different contrasts $\epsilon_{r}$ and conditioning schemes (scattered fields
vs back projections) are illustrated in Figures S4, S5, S6 and S7 in the
supplemental materials.

In order to assess the performance of different conditional flows in
uncertainty quantification, we design another experiment with 180 incident
plane waves and 180 receivers only on the top side of the object (6(d)). In
this setup, we expect higher uncertainty in the lower part of the object.

Figure 10 illustrates the performance of C-Glow, C-Rev and C-Trumpets for
$\epsilon_{r}=6$. We use a simple back projection (BP)[49] as conditioning
measurements $y$. In the first panel (from the second row to the fifth row),
we show MMSE estimates of different flow models: C-Trumpets perform better
than bijective flows (see also Table II). As expected, however, the
reconstructions are not as good as in the previous experiment since we only
have measurements on one side of the domain. We see that C-Trumpets again
provide meaningful uncertainty estimates, with more uncertainty in the lower
part of the object (red regions in the UQ panel correspond to higher
uncertainty).

#### Seismic travel-time tomography (NS)

We work with the linearized seismic tomography operator as described in [42].
Here we are given travel times of a seismic wave between each sensor pair in a
network of NS sensors placed on the ground. The travel times are assumed to
linearly depend on the “slowness” map which is taken to be an MNIST image [4].
We use NS$=10$ sensors on the boundary of the lower part of the domain which
yields 33 measurements as shown in Figure 6(c). We use the pseudo-inverse of
the measurements $y$ as the conditioning samples and work at the SNR of 40dB.
Figure 11 compares the performance of C-Trumpets and C-Rev. In the first
panel, the second and the third rows show MMSE estimates of C-Rev and
C-Trumpets. C-Trumpets outperform C-Rev in both MMSE estimate and posterior
sampling (see also Table II). Furthermore, given no sensors (Figure 6(c)) in
the top regions of the image (or slowness map), we would expect higher
uncertainty there. The UQ column shows that estimates from C-Trumpets assign
higher uncertainty to the top half of the domain compared to C-Rev. Additional
results are shown in Figure S8 in the supplemental materials.

Figure 10: Performance comparison in electromagnetic inverse scattering
($\epsilon_{r}=6$, top-view, BP conditioned); C-Trumpets demonstrate more
meaningful UQ by assigning higher uncertainty in the bottom of the object (red
regions), which lacks measurements. Figure 11: Performance comparison in
seismic travel-time tomography with 10 sensors. C-Trumpets demonstrate better
posterior samples and MMSE estimate, both models assign higher uncertainty in
top of the image (black regions) which lacks measurements.

###  V-C Image restoration

In this section, we compare C-Trumpets with models on various image
restoration tasks. Similarly as in computational imaging problems, each task
requires training a different model but once trained, the model can be used
instantaneously for arbitrary measurements $y$. We consider four standard
restoration tasks: (i) Image denoising: We train a model to generate plausible
clean images given a noisy image at an SNR of $-1$ dB; (ii) Image super
resolution ($f$): Generate high-resolution images given an image downsampled
by a factor of $f$ along each axis; (iii) Random mask ($p$): The degradation
process replaces every pixel with zero with probability $p$; and (iv) Mask
($s$): The degradation process replaces an $s\times s$ patch of the image with
zeros.

We used the 8-bit RGB CelebA [50] dataset with 80000 $64\times 64$ training
samples. Figure S2 in the supplemental materials compares the performance of
flow models on different image restoration tasks. Table II further gives the
SNR and SSIM of the MMSE estimate. C-Trumpets consistently outperform other
conditional bijective flows. The difference in the performance of C-Rev and
C-Trumpets suggests that the low-dimensional latent space of C-Trumpets acts
as an effective regularizer in the restoration mapping, $f_{\theta}$.
C-Trumpets also provide a meaningful uncertainty quantification. For example,
although the forward operator is random in the random mask problem, C-Trumpets
still capture a meaningful uncertainty estimate by assigning higher
uncertainty inside the masked region (see Figure 2(d)).

In order to assess the memory requirements of the different models, we compare
the number of trainable parameters used for training over $64\times 64$ RGB
images: C-INN: 13M, C-Glow: 14M, C-Rev: 22M and C-Trumpets: 4M. We did not
have sufficient resources to train bijective models over $256\times 256$
images but it is clear that the differences in the memory footprint at that
resolution would be further exacerbated.

TABLE II: Performance of MMSE estimate (computed over 25 posterior samples) of
different models on solving inverse problems averaged over 5 test images

(a) SNR (dB)

| C-INN | C-Glow | C-Rev | C-Trumpets
---|---|---|---|---
Denoising | 15.98 | 15.67 | 15.99 | 16.86
Super resolution ($\times 4$) | 18.57 | 19.24 | 19.22 | 20.70
RandMask ($p=0.2$) | 20.84 | 19.70 | 13.31 | 20.11
Mask ($s=32$) | 18.91 | 17.86 | 18.54 | 21.01
Limited-view CT | - | 11.63 | 13.13 | 13.58
Scattering ($\epsilon_{r}=6$) | - | 1.45 | 1.19 | 4.04
Travel-time (NS = 10) | - | - | 14.83 | 18.19

(b) SSIM

| C-INN | C-Glow | C-Rev | C-Trumpets
---|---|---|---|---
Denoising | 0.63 | 0.60 | 0.66 | 0.70
Super Resolution ($\times 4$) | 0.79 | 0.82 | 0.79 | 0.84
RandMask ($p=0.2$) | 0.86 | 0.82 | 0.57 | 0.87
Mask ($s=32$) | 0.83 | 0.80 | 0.84 | 0.88
Limited-view CT | - | 0.59 | 0.71 | 0.74
Scattering ($\epsilon_{r}=6$) | - | 0.67 | 0.65 | 0.73
Travel-time (NS = 10) | - | - | 0.59 | 0.64

###  V-D MAP vs MMSE for ill-posed inverse problems

Figure 12 demonstrates the MMSE, $g^{\dagger}$-MAP and a random posterior
sample for four types of ill-posed inverse problems. The MMSE estimate is
obtained by averaging over 25 posterior samples. Although the MMSE estimate is
the optimal reconstruction in terms of the $\ell^{2}$-error, we see from
Figure 12 that it is often blurry, especially when the true posterior is
multi-modal; $g^{\dagger}$-MAP estimates are sharper. Moreover, as the MMSE
estimate is obtained by averaging over the posterior samples, it is not
generally on the manifold, while the $g^{\dagger}$-MAP estimate is always on
the manifold.

Figure 12: MMSE and $g^{\dagger}$-MAP estimations in different inverse
problems; the proposed $g^{\dagger}$-MAP estimate gives much sharper
reconstruction than MMSE.

##  VI Related Work

There is by now a very large body of work on solving inverse imaging problems
using deep neural networks. On the supervised regression end of the spectrum
arguably the most important architecture is the U-Net [51]. It has been
applied with great success to a variety of imaging problems including CT [1],
magnetic resonance imaging (MRI) [52] and electromagnetic inverse scattering
[53]. Its success may be attributed to the particular multiscale structure
[54, 55] which matches both the physical description of the imaging problems
and the representation of the involved image classes. Many alternatives have
been proposed for specific problems where the assumptions that make the U-Net
a natural choice do not hold, for example for wave-based problems [56, 57].

On the other hand, trained generative models have been shown to be effective
priors [58, 59] in ill-posed inverse problems that can be trained in an
unsupervised manner. Normalizing flows in particular were used to approximate
MAP estimates using iterative optimization [59, 60]. However, normalizing
flows are bijective, requiring large memory and compute budget even at
moderate resolution. Moreover, they lack a low-dimensional latent space which
has been shown to effectively regularize the inversion. Brehmer and Cranmer
[9] proposed the first injective model for densities supported on low-
dimensional manifolds. Kothari et al. [10] proposed injective flows with
significantly optimized compute and memory requirements which were shown to
outperform earlier variants of injective and bijective flows in computational
imaging problems.

A variety of methods have been proposed for Bayesian imaging. The latter aims
to approximate posterior distribution and / or the various point estimators
related to the posterior. The authors of [61] consider a convex log-likelihood
function for posterior distribution around the MAP estimate, which is suitable
for modeling unimodal posteriors. Normalizing flows have been proposed as
variational approximators to the posterior distribution for a given
measurement [2]. The authors of [30, 10] propose to train a flow model to
approximate the posterior corresponding to a prior which is also modeled using
a flow model. More recently, pre-trained GANs were used in conjunction with
MCMC to generate posterior samples in non-linear inverse problems [25]. The
authors of [62, 63] use a style-GAN generator [16] to regularize ill-posed
inverse problems and generate posterior samples. All these methods train a new
generative model or run an iterative process for every measurement, which
makes them slow when applied to multiple reconstructions. Further, training
for each conditioning sample requires many calls to the forward operator. To
tackle these issues, one may consider amortized inference to make the
generative models conditioned based on the measurements.

Conditional versions of generative adversarial networks (GANs) [12, 64] and
variational autoencoders (VAEs) [65, 66] rely on injecting conditioning data
into the different layers of the generator model. However, the lack of access
to the posterior distribution make conditional GANs difficult to be used in
inference tasks. On the other hand, VAEs provide lower bounds on likelihoods
of generated samples. While these bounds can be made tighter by importance
weighting [67], C-Trumpets allow one to compute exact end-to-end likelihoods
via stochastic estimation, as well as to obtain fast exact values of
likelihoods before the high-dimensional expansion. These generative models
also suffer from mode collapse and training instabilities. Conditional
normalizing flows were introduced in [11] to estimate the posterior by
modifying the scale and shift terms of the coupling layers. The authors of
[19] additionally make the mean and covariance of the base Gaussian
distribution depend on the measurements. More recently, [68] proposed to
append the measurements to all layers of the Glow network [15] in order to
enable greater information flow from conditioning data to the generated
samples. A different approach to conditioning the flow models has been
developed by [43, 69], benefitting from two parallel flow models for
simultaneously modeling of target and conditioning samples. All these
conditional normalizing flows are bijective mappings from latent space to the
target domain for each measurement. It is worth mentioning that all these
models can be used in the bijective part of C-Trumpets to exploit their
advantages in posterior modeling. On the side of theory, Puthawala et al. [36]
establish universality of density and manifold approximation of injective
flows such as those in [9, 10].

##  VII Limitations and Conclusions

We proposed C-Trumpets, a conditional injective flow model that enables
amortized inference with approximate posteriors that live on low-dimensional
manifolds. Our proposed model is considerably cheaper to train in terms of
memory and compute costs compared to the regular conditional flows. The
experiments we performed indicate that C-Trumpets generate better posterior
samples and more accurate uncertainty estimates over a variety of ill-posed
inverse problems. The proposed fixed-volume-change coupling layers enable us
to approximate the sharp MAP estimates instantaneously after training. High
computational demands of training bijective flows at high resolution have thus
far impeded their wider adoption in computational imaging workflows. The
comparably lightweight memory footprint of C-Trumpets together with
physically-consistent UQ makes them an attractive architecture for imaging
problems where characterizing uncertainty is paramount.

#### Limitations

C-Trumpets have several limitations that warrant discussion. The latent space
dimension in C-Trumpets is chosen arbitrarily and it may be quite different
from the true dimension of the posterior support. Recent work [70] proposes an
injective flow architecture that estimates the dimension of the data manifold.
Similar ideas may extend to C-Trumpets but for the moment we rely on rules of
thumb rather than principled choices. Another limitation is that it is not
straightforward to estimate the likelihoods of samples generated by C-Trumpets
(cf. Section III-C), the reason being that the Jacobian determinant of
compositions of maps between spaces of different dimension cannot be written
as a product of Jacobian determinants of the constituent maps. Likelihood
estimates can still be obtained by sampling but that is considerably slower
than what is possible with bijective flows. Recently, Ross et al. [71]
proposed an injective generator which provides access to the exact likelihood
of the generated samples, but the constraints they impose on the architecture
in order to enable this feat seem to severely limit expressivity. The design
of an injective model that is at once expressive, lightweight, and gives fast
exact likelihoods remains an open problem. On the theoretical side, further
studies are needed to characterize the types of posterior distributions that
can be modeled by C-Trumpets, especially with fixed-volume-change layers. The
important open question is that of universality of C-Trumpets as models of
conditional distributions. Finally, Siahkoohi et al. [72, 73] fully exploit
the depth-independent memory complexity of normalizing flows to handle high-
dimensional data. This strategy can also be used in C-Trumpets to further
improve memory efficiency and apply the model to super high-dimensional
imaging problems.

## References

  * [1] K. H. Jin, M. T. McCann, E. Froustey, and M. Unser, “Deep convolutional neural network for inverse problems in imaging,” _IEEE Transactions on Image Processing_ , vol. 26, no. 9, pp. 4509–4522, 2017.
  * [2] H. Sun and K. L. Bouman, “Deep probabilistic imaging: Uncertainty quantification and multi-modal solution characterization for computational imaging,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 35, no. 3, 2021, pp. 2628–2637.
  * [3] F. Vasconcelos, B. He, N. Singh, and Y. W. Teh, “UncertaINR: Uncertainty Quantification of End-to-End Implicit Neural Representations for Computed Tomography,” _arXiv preprint arXiv:2202.10847_ , 2022.
  * [4] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner, “Gradient-based learning applied to document recognition,” _Proceedings of the IEEE_ , vol. 86, no. 11, pp. 2278–2324, 1998.
  * [5] M.-H. Chen, Q.-M. Shao, and J. G. Ibrahim, _Monte Carlo methods in Bayesian computation_. Springer Science & Business Media, 2012.
  * [6] C. M. Bishop and N. M. Nasrabadi, _Pattern recognition and machine learning_. Springer, 2006, vol. 4, no. 4.
  * [7] J. Marino, Y. Yue, and S. Mandt, “Iterative Amortized Inference,” in _Proceedings of the 35th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, J. Dy and A. Krause, Eds., vol. 80. PMLR, 10–15 Jul 2018, pp. 3403–3412.
  * [8] A. Kumar, B. Poole, and K. Murphy, “Regularized autoencoders via relaxed injective probability flow,” in _International Conference on Artificial Intelligence and Statistics_. PMLR, 2020, pp. 4292–4301.
  * [9] J. Brehmer and K. Cranmer, “Flows for simultaneous manifold learning and density estimation,” in _Advances in Neural Information Processing Systems_ , vol. 33. Curran Associates, Inc., 2020, pp. 442–453.
  * [10] K. Kothari, A. Khorashadizadeh, M. de Hoop, and I. Dokmanić, “Trumpets: Injective flows for inference and inverse problems,” in _Uncertainty in Artificial Intelligence_. PMLR, 2021, pp. 1269–1278.
  * [11] L. Ardizzone, J. Kruse, C. Lüth, N. Bracher, C. Rother, and U. Köthe, “Conditional invertible neural networks for diverse image-to-image translation,” in _Pattern Recognition: 42nd DAGM German Conference, DAGM GCPR 2020, Tübingen, Germany, September 28–October 1, 2020, Proceedings 42_. Springer, 2021, pp. 373–387.
  * [12] M. Mirza and S. Osindero, “Conditional generative adversarial nets,” _arXiv preprint arXiv:1411.1784_ , 2014.
  * [13] L. Dinh, D. Krueger, and Y. Bengio, “Nice: Non-linear independent components estimation,” in _International Conference on Learning Representations, ICLR_ , 2015.
  * [14] L. Dinh, J. Sohl-Dickstein, and S. Bengio, “Density estimation using real nvp,” in _International Conference on Learning Representations, ICLR_ , 2017\.
  * [15] D. P. Kingma and P. Dhariwal, “Glow: Generative flow with invertible 1x1 convolutions,” _Advances in neural information processing systems_ , vol. 31, 2018.
  * [16] T. Karras, S. Laine, and T. Aila, “A style-based generator architecture for generative adversarial networks,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 4401–4410.
  * [17] T. Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila, “Analyzing and improving the image quality of stylegan,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2020, pp. 8110–8119.
  * [18] J. J. Yu, K. G. Derpanis, and M. A. Brubaker, “Wavelet flow: Fast training of high resolution normalizing flows,” _Advances in Neural Information Processing Systems_ , vol. 33, pp. 6184–6196, 2020.
  * [19] C. Winkler, D. Worrall, E. Hoogeboom, and M. Welling, “Learning likelihoods with conditional normalizing flows,” _arXiv preprint arXiv:1912.00042_ , 2019\.
  * [20] G. A. Padmanabha and N. Zabaras, “Solving inverse problems using conditional invertible neural networks,” _Journal of Computational Physics_ , vol. 433, p. 110194, 2021.
  * [21] R. Gribonval, “Should penalized least squares regression be interpreted as maximum a posteriori estimation?” _IEEE Transactions on Signal Processing_ , vol. 59, no. 5, pp. 2405–2410, 2011.
  * [22] C. K. Sønderby, J. Caballero, L. Theis, W. Shi, and F. Huszár, “Amortised map inference for image super-resolution,” _International Conference on Learning Representations, ICLR_ , 2017.
  * [23] N. Courts and H. Kvinge, “Bundle networks: Fiber bundles, local trivializations, and a generative approach to exploring many-to-one maps,” _International Conference on Learning Representations, ICLR_ , 2022.
  * [24] A. Jalal, M. Arvinte, G. Daras, E. Price, A. G. Dimakis, and J. Tamir, “Robust compressed sensing mri with deep generative priors,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 14 938–14 954, 2021.
  * [25] P. Bohra, T.-a. Pham, J. Dong, and M. Unser, “Bayesian Inversion for Nonlinear Imaging Models using Deep Generative Priors,” _arXiv preprint arXiv:2203.10078_ , 2022.
  * [26] K. Wang, T. Bui-Thanh, and O. Ghattas, “A randomized maximum a posteriori method for posterior sampling of high dimensional nonlinear Bayesian inverse problems,” _SIAM Journal on Scientific Computing_ , vol. 40, no. 1, pp. A142–A171, 2018.
  * [27] G. E. Hinton and D. Van Camp, “Keeping the neural networks simple by minimizing the description length of the weights,” in _Proceedings of the sixth annual conference on Computational learning theory_ , 1993, pp. 5–13.
  * [28] A. Graves, “Practical variational inference for neural networks,” _Advances in neural information processing systems_ , vol. 24, pp. 2348–2356, 2011.
  * [29] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, “Normalizing flows for probabilistic modeling and inference,” _Journal of Machine Learning Research_ , vol. 22, no. 57, pp. 1–64, 2021.
  * [30] J. Whang, E. Lindgren, and A. Dimakis, “Composing normalizing flows for inverse problems,” in _International Conference on Machine Learning_. PMLR, 2021, pp. 11 158–11 169.
  * [31] T. P. Minka, “Expectation propagation for approximate bayesian inference,” in _Uncertainty in Artificial Intelligence_. PMLR, 2001, p. 362––369.
  * [32] ——, “A family of algorithms for approximate Bayesian inference,” Ph.D. dissertation, Massachusetts Institute of Technology, 2001.
  * [33] Y. Bengio, A. Courville, and P. Vincent, “Representation learning: A review and new perspectives,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 35, no. 8, pp. 1798–1828, 2013.
  * [34] B. C. Brown, A. L. Caterini, B. L. Ross, J. C. Cresswell, and G. Loaiza-Ganem, “Verifying the union of manifolds hypothesis for image data,” in _The Eleventh International Conference on Learning Representations_ , 2023.
  * [35] M. Puthawala, K. Kothari, M. Lassas, I. Dokmanic, and M. de Hoop, “Globally injective relu networks,” _Journal of Machine Learning Research_ , vol. 23, no. 105, pp. 1–55, 2022.
  * [36] M. Puthawala, M. Lassas, I. Dokmanic, and M. De Hoop, “Universal joint approximation of manifolds and densities by simple injective flows,” in _International Conference on Machine Learning_. PMLR, 2022, pp. 17 959–17 983.
  * [37] Y. Choquet-Bruhat and C. DeWitt-Morette, _Analysis, Manifolds and Physics Revised Edition_. Gulf Professional Publishing, 1982.
  * [38] Y. Fan and Z. Zhao, “Cryo-electron microscopy image denoising using multi-frequency vector diffusion maps,” in _2021 IEEE International Conference on Image Processing (ICIP)_. IEEE, 2021, pp. 3463–3467.
  * [39] V. Debarnot, V. Kishore, C. Shi, and I. Dokmanić, “Manifold rewiring for unlabeled imaging,” _arXiv preprint arXiv:2209.05168_ , 2022.
  * [40] M. M. Bronstein, J. Bruna, T. Cohen, and P. Veličković, “Geometric deep learning: Grids, groups, graphs, geodesics, and gauges,” _arXiv preprint arXiv:2104.13478_ , 2021.
  * [41] X. Chen, _Computational methods for electromagnetic inverse scattering_. John Wiley & Sons, 2018.
  * [42] S. Gupta, K. Kothari, M. V. de Hoop, and I. Dokmanić, “Random mesh projectors for inverse problems,” in _7th International Conference on Learning Representations, ICLR 2019_ , 2019.
  * [43] M. Sorkhei, G. E. Henter, and H. Kjellström, “Full-glow: Fully conditional glow for more realistic image generation,” in _Pattern Recognition: 43rd DAGM German Conference, DAGM GCPR 2021, Bonn, Germany, September 28–October 1, 2021, Proceedings_. Springer, 2022, pp. 697–711.
  * [44] A. Doerr, “Cryo-electron tomography,” _Nature Methods_ , vol. 14, no. 1, pp. 34–34, 2017.
  * [45] E. Venkatesh and S. V. Elluru, “Cone beam computed tomography: basics and applications in dentistry,” _Journal of istanbul University faculty of Dentistry_ , vol. 51, no. 3 Suppl 1, p. S102, 2017.
  * [46] J. Leuschner, M. Schmidt, D. O. Baguer, and P. Maass, “LoDoPaB-CT, a benchmark dataset for low-dose computed tomography reconstruction,” _Scientific Data_ , vol. 8, 2021.
  * [47] A. I. Nachman, “Global uniqueness for a two-dimensional inverse boundary value problem,” _Annals of Mathematics_ , pp. 71–96, 1996.
  * [48] J. Sylvester and G. Uhlmann, “A global uniqueness theorem for an inverse boundary value problem,” _Annals of mathematics_ , pp. 153–169, 1987.
  * [49] K. Belkebir, P. C. Chaumet, and A. Sentenac, “Superresolution in total internal reflection tomography,” _JOSA A_ , vol. 22, no. 9, pp. 1889–1897, 2005.
  * [50] Z. Liu, P. Luo, X. Wang, and X. Tang, “Deep learning face attributes in the wild,” in _Proceedings of the IEEE international conference on computer vision_ , 2015, pp. 3730–3738.
  * [51] O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” in _International Conference on Medical image computing and computer-assisted intervention_. Springer, 2015, pp. 234–241.
  * [52] C. M. Hyun, H. P. Kim, S. M. Lee, S. Lee, and J. K. Seo, “Deep learning for undersampled MRI reconstruction,” _Physics in Medicine & Biology_, vol. 63, no. 13, p. 135007, 2018.
  * [53] Z. Wei and X. Chen, “Deep-learning schemes for full-wave nonlinear inverse scattering problems,” _IEEE Transactions on Geoscience and Remote Sensing_ , vol. 57, no. 4, pp. 1849–1860, 2018.
  * [54] D. Lee, J. Yoo, and J. C. Ye, “Deep residual learning for compressed sensing MRI,” in _2017 IEEE 14th International Symposium on Biomedical Imaging (ISBI 2017)_ , 2017, pp. 15–18.
  * [55] T. Liu, A. Chaman, D. Belius, and I. Dokmanić, “Learning multiscale convolutional dictionaries for image reconstruction,” _IEEE Transactions on Computational Imaging_ , vol. 8, pp. 425–437, 2022.
  * [56] Y. Fan and L. Ying, “Solving inverse wave scattering with deep learning,” _Annals of Mathematical Sciences and Applications_ , pp. 23–48, 2022.
  * [57] K. Kothari, M. de Hoop, and I. Dokmanić, “Learning the geometry of wave-based imaging,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 8318–8329.
  * [58] A. Bora, A. Jalal, E. Price, and A. G. Dimakis, “Compressed sensing using generative models,” in _International Conference on Machine Learning_. PMLR, 2017, pp. 537–546.
  * [59] M. Asim, M. Daniels, O. Leong, A. Ahmed, and P. Hand, “Invertible generative models for inverse problems: mitigating representation error and dataset bias,” in _Proceedings of the 37th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, vol. 119. PMLR, 13–18 Jul 2020, pp. 399–409.
  * [60] J. Whang, Q. Lei, and A. Dimakis, “Compressed sensing with invertible generative models and dependent noise,” in _NeurIPS 2020 Workshop on Deep Learning and Inverse Problems_ , 2020.
  * [61] A. Repetti, M. Pereyra, and Y. Wiaux, “Scalable Bayesian uncertainty quantification in imaging inverse problems via convex optimization,” _SIAM Journal on Imaging Sciences_ , vol. 12, no. 1, pp. 87–118, 2019.
  * [62] S. Bhadra, U. Villa, and M. A. Anastasio, “Mining the manifolds of deep generative models for multiple data-consistent solutions of ill-posed tomographic imaging problems,” _arXiv preprint arXiv:2202.05311_ , 2022.
  * [63] R. V. Marinescu, D. Moyer, and P. Golland, “Bayesian image reconstruction using deep generative models,” _arXiv preprint arXiv:2012.04567_ , 2020\.
  * [64] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, “Generative adversarial nets,” _Advances in neural information processing systems_ , vol. 27, 2014.
  * [65] K. Sohn, H. Lee, and X. Yan, “Learning structured output representation using deep conditional generative models,” _Advances in neural information processing systems_ , vol. 28, pp. 3483–3491, 2015.
  * [66] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” _International Conference on Learning Representations, ICLR_ , 2014.
  * [67] Y. Burda, R. Grosse, and R. Salakhutdinov, “Importance weighted autoencoders,” _International Conference on Learning Representations_ , 2015\.
  * [68] Y. Lu and B. Huang, “Structured output learning with conditional generative flows,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 34, no. 04, 2020, pp. 5005–5012.
  * [69] A. Pumarola, S. Popov, F. Moreno-Noguer, and V. Ferrari, “C-flow: Conditional generative flow models for images and 3d point clouds,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 7949–7958.
  * [70] M. Zhang, Y. Sun, S. McDonagh, and C. Zhang, “Flow based models for manifold data,” _arXiv preprint arXiv:2109.14216_ , 2021.
  * [71] B. Ross and J. Cresswell, “Tractable density estimation on learned manifolds with conformal embedding flows,” _Advances in Neural Information Processing Systems_ , vol. 34, 2021.
  * [72] A. Siahkoohi, G. Rizzuti, M. Louboutin, P. A. Witte, and F. J. Herrmann, “Preconditioned training of normalizing flows for variational inference in inverse problems,” _arXiv preprint arXiv:2101.03709_ , 2021.
  * [73] A. Siahkoohi, G. Rizzuti, R. Orozco, and F. J. Herrmann, “Reliable amortized variational inference with physics-based latent distribution correction,” _Geophysics_ , vol. 88, no. 3, pp. 1–137, 2023.
  * [74] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _International Conference on Learning Representations, ICLR_ , 2015.
  * [75] A. Nagrani, J. S. Chung, and A. Zisserman, “Voxceleb: A large-scale speaker identification dataset,” in _Interspeech_ , 2017.

Supplemental Materials: Conditional Injective Flows for Bayesian Imaging

##  S-I Fast Inverses

Inverting the revnet coupling layers requires the matrix inverse of the kernel
of the $1\times 1$ convolution layer. As this may be slow, Kingma and Dhariwal
[15] proposed to use the LU decomposition to reduce the computational
complexity,

$w=PL(U+\text{diag}(s)),$ (S1)

where $P$ is a permutation matrix, $L$ is a lower-triangular matrix with ones
on diagonal and $s$ is a vector. Computing the $\log\det$ of the layer
Jacobian then simplifies to

$\log|\det(w)|=\sum_{i=1}^{d}\log(|s_{i}|).$ (S2)

The LU decomposition thus reduces the complexity of computing the $\log\det$
of the Jacobian of the $1\times 1$ convolution layer from $\mathcal{O}(c^{3})$
to $\mathcal{O}(c)$ (and one-time factorization cost). This trick, however,
was only used in inference (when sampling from the model), but not for
training the bijective flows as matrix inversion is not critical in that case.
Unlike bijective flows, C-Trumpets do require matrix inversion during the MSE
training phase (cf. (16)). In order to reduce the computational cost of
inversion, we leverage the LU decomposition as follows

$w^{-1}=(U+\text{diag}(s))^{-1}L^{-1}P^{-1}.$ (S3)

Since $L$ and $U$ are triangular matrices and $P$ is a fixed rotation,
computing the inverses costs $\mathcal{O}(2\times c)$ instead of
$\mathcal{O}(c^{3})$. This significantly reduces the training time of the
injective part of C-Trumpets, especially in high-dimensional problems.

##  S-II Additional Experimental Details

###  S-II-A Fiber bundles

_Datasets:_ We analyzed the performance of C-Trumpets on two fiber bundle
datasets:

  * •

Torus

$\begin{array}[]{l}x=\cos(t)[R+r\cos(s)]\\\ y=\sin(t)[R+r\cos(s)]\\\
z=r\sin(s)\end{array}$ (S4)

where $t\in[0,2\pi)$ and $s\in[0,2\pi)$. We set $R=1$ and $r=0.25$ and
generate $18000$ training samples.

  * •

Elliptic Möbius

$\begin{array}[]{l}x=\cos(t)[R-b\sin(t/2)\sin(s)+a\cos(t/2)cos(s)]\\\
y=\sin(t)[R-b\sin(t/2)\sin(s)+a\cos(t/2)cos(s)]\\\
z=b\cos(t/2)\sin(s)+a\sin(t/2)\cos(s)\end{array}$ (S5)

where $t\in[0,2\pi)$ and $s\in[0,2\pi)$. We set $R=1$, $a=0.4$ and $b=0.1$ and
generate 18000 training samples.

_Network architecture and training details:_ The injective part of C-Trumpets
maps $\mathbb{R}^{2}\to\mathbb{R}^{3}$ and consists of an expansion layer (a
matrix of size $\mathbb{R}^{3\times 2}$) followed by 24 revent blocks without
activation normalization layer. The bijective part which maps
$\mathbb{R}^{2}\to\mathbb{R}^{2}$ also consists of 32 revent blocks. We use
coupling layers proposed in [15] with 3 fully-connected layers for scale and
bias. We train for 100 epochs using the Adam optimizer [74] with a learning
rate of $10^{-3}$ for both the injective and the bijective parts of
C-Trumpets. We note that this minimal input dimension of the coupling layers
has yields poor expressivity, but we only use this example for illustration.

###  S-II-B Limited-view CT

We consider the 2D parallel-beam CT problem with a missing cone of sensors.

_Network architecture:_ We describe the architecture of C-Trumpets for
$256\times 256$ images. The injective part consists of 6 injective revnet
blocks, each increasing the dimension by a factor of 2. Between the injective
layers, we intersperse 36 bijective revnet blocks. We choose a latent space of
size 2048. The bijective part consists of 48 bijective revnet blocks. We use 3
convolutional layers for the scale and bias networks of fixed-volume-change
coupling layers. The conditioning networks have 3 conv layers along with a
‘squeeze’ layer that helps match the dimension of the conditioning sample to
the input of the scale network.

The architecture of C-Trumpets for $64\times 64$ images is similar to the
$256\times 256$ variant except that we use 18 (resp. 24) revnet blocks in
injective (resp. bijective) subnetworks and choose a 64-dimensional latent
space. The C-Rev model used for comparison has 24 revnet blocks, all at the
highest resolution (i.e $64\times 64$).

We initialize the weights of the revnet blocks as in [15]. All elements of the
skip connection matrix $S$ are initialized to $0.5$. All models are trained
for 300 epochs (150 per phase) using the Adam optimizer [74] with a learning
rate of $10^{-4}$.

###  S-II-C Electromagnetic inverse scattering

We consider the 2D transverse magnetic inverse scattering problem. We
illuminate the domain of investigation $D_{inv}$—a $D{\times}D$ square—by
$N_{i}$ plane waves; $N_{r}$ antennas measure the scattered field for each
incident wave.

The forward scattering problem can be described by two coupled equations [41].
The first equation, called the state equation, relates the total electric
field ($E^{t}$) in an unknown domain to its contrast current density ($J$) as

$E^{t}(r)=E^{i}(r)+k_{0}^{2}\int\displaylimits_{D_{inv}}g(r,r^{\prime})J(r^{\prime})dr^{\prime},$
(S6)

where $g(r,r^{\prime})=(1/4i)H^{(2)}_{0}(k_{0}|r-r^{\prime}|)$, $H^{(2)}_{0}$
(a Hankel function of the second kind) is the 2-D free space Green’s function,
and $J(r^{\prime})$ can be calculated using the equivalence theorem as
$J(r^{\prime})=\chi(r^{\prime})E^{t}(r^{\prime})$ with the contrast
$\chi(r^{\prime})=\epsilon_{r}(r^{\prime})-1$.

The second equation, called the data equation, maps the contrast current
source to the scattered electric field ($E^{s}$) at antenna locations. For
$r\in S$,

$E^{s}(r)=k_{0}^{2}\int\displaylimits_{D_{inv}}g(r,r^{\prime})J(r^{\prime})dr^{\prime}.$
(S7)

The task of inverse scattering is to reconstruct the contrast $\chi$ from the
scattered field $E^{s}$.

_Network architecture:_ We use the same architecture and training parameters
as in the $64\times 64$ 2D-CT problem.

###  S-II-D Linearized travel-time tomography

Linearized travel-time tomography is inspired by seismic travel-time
tomography that aims to reconstruct the wavespeed variation inside a planet.
Sensors are placed on the ground and we measure arrival times of surface
waves. We assume that the receivers and transmitting sensors are co-located.
We further assume that the pixel intensities represent the “slowness” (inverse
wave speed). We can therefore measure the travel times between a transmitting
sensor, $s_{i}$ and a receiving sensor $s_{j}$ as

$t(s_{i},s_{j})=\int_{0}^{1}f(s_{i}+\lambda(s_{j}-s_{i}))d\lambda,$ (S8)

where $f$ represents the image. Note that $t(s_{i},s_{j})=t(s_{j},s_{i})$.

_Network architecture:_ The injective subnetwork consists of 4 injective
revnet blocks, each increasing the dimension by a factor of 2. We use 12
bijective revnet blocks interspersed in between the injective layers. The
bijective part takes 64-dimensional latent vectors and consists of 12
bijective revnet blocks. C-Rev used in comparison has 16 revnet blocks all at
the $32\times 32$ resolution. The training hyperparameters are same as the CT
problem.

###  S-II-E Image restoration tasks

We use the same architecture and training hyperparameters as that of the CT
problem except that we work with $64\times 64\times 3$ resolution images. The
latent dimension of C-Trumpets is chosen to be 192 ($64\times 3$).

##  S-III Additional Experiments

###  S-III-A Class-based image generation

We perform class-based image generation over MNIST digits [4] and a subset of
10 people from the voxceleb [75] face dataset with 5000 $64\times 64$ training
samples. We use the one-hot class labels as conditioning vectors to generate
samples from the given class and use two fully-connected layers followed by a
reshaping module for the conditioning network of C-Trumpets. Figure S1 shows
the class-based generated samples by C-Trumpets. This experiment indicates
that our proposed model can generate good quality class-conditioned samples.

(a) MNIST

(b) Voxceleb

Figure S1: Class-based image generation

###  S-III-B Image restoration and inverse problems

Figure S2 compares the performance of conditional normalizing flows in image
restoration tasks. Figures S3 to S12 demonstrate further results on different
ill-posed inverse problems. We can observe that C-Trumpets have a significant
edge over the baselines.

(a) Super resolution ($f=4$)

(b) Random mask ($p=0.2$)

(c) Denoising

(d) Mask ($s=32$)

Figure S2: Performance comparison over image restoration problems.

(a) Vertical missing cone ($60^{\circ}$ to $120^{\circ}$)

(b) Horizontal missing cone($-30^{\circ}$ to $30^{\circ}$)

Figure S3: Limited-view CT in resolution $256\times 256$ Figure S4: Inverse
scattering ($\epsilon_{r}=1.5$, full-view) Figure S5: Inverse scattering
($\epsilon_{r}=2$, full-view) Figure S6: Inverse scattering ($\epsilon_{r}=6$,
full-view) Figure S7: Inverse scattering ($\epsilon_{r}=6$, top-view) Figure
S8: Travel-time tomography ($NS=10$) Figure S9: Denoising Figure S10: Mask
($s=32$) Figure S11: Super resolution ($f=4$) Figure S12: Random mask
($p=0.2$)


* (8) J. Maldacena, D. Stanford and Z. Yang, _Diving into traversable wormholes_ , _Fortsch. Phys._ 65 (2017) 1700034, [1704.05333].
  * (9) J. Maldacena and X.-L. Qi, _Eternal traversable wormhole_ , 1804.00491.
  * (10) J. Maldacena, A. Milekhin and F. Popov, _Traversable wormholes in four dimensions_ , 1807.04726.
  * (11) E. Witten, _Anti-de Sitter space, thermal phase transition, and confinement in gauge theories_ , _Adv. Theor. Math. Phys._ 2 (1998) 505–532, [hep-th/9803131].
  * (12) T. Sakai and S. Sugimoto, _Low energy hadron physics in holographic QCD_ , _Prog. Theor. Phys._ 113 (2005) 843–882, [hep-th/0412141].
  * (13) B. M. Levitan and M. G. Gasymov, _Determination of a Differential Equation by Two of its Spectra_ , _Russian Mathematical Surveys_ 19 (Apr., 1964) R01.
  * (14) O. H. Hald, _The inverse Sturm-Liouville problem with symmetric potentials_ , _Acta Mathematica_ 141 (1978) 263 – 291.
  * (15) M. Kac, _Can one hear the shape of a drum?_ , _The american mathematical monthly_ 73 (1966) 1–23.
  * (16) T. Faulkner, M. Li and H. Wang, _A modular toolkit for bulk reconstruction_ , _JHEP_ 04 (2019) 119, [1806.10560].
  * (17) S. Ryu and T. Takayanagi, _Holographic derivation of entanglement entropy from the anti–de sitter space/conformal field theory correspondence_ , _Physical review letters_ 96 (2006) 181602.
  * (18) S. Cooper, M. Rozali, B. Swingle, M. Van Raamsdonk, C. Waddell and D. Wakeham, _Black hole microstate cosmology_ , _Journal of High Energy Physics_ 2019 (2019) 1–70.
  * (19) S. Antonini and B. Swingle, _Cosmology at the end of the world_ , _Nature Phys._ 16 (2020) 881–886, [1907.06667].
  * (20) S. Antonini and B. Swingle, _Holographic boundary states and dimensionally reduced braneworld spacetimes_ , _Phys. Rev. D_ 104 (2021) 046023, [2105.02912].
  * (21) A. Almheiri, A. Mousatov and M. Shyani, _Escaping the Interiors of Pure Boundary-State Black Holes_ , 1803.04434.
  * (22) S. Fallows and S. F. Ross, _Constraints on cosmologies inside black holes_ , 2203.02523.
  * (23) C. Waddell, _Bottom-up holographic models for cosmology_ , 2203.03096.
  * (24) I. S. Gradshteyn and I. M. Ryzhik, _Table of integrals, series, and products_.  Academic press, 2014.
  * (25) O. Aharony, D. Marolf and M. Rangamani, _Conformal field theories in anti-de Sitter space_ , _JHEP_ 02 (2011) 041, [1011.6144].
  * (26) E. Hijano and D. Neuenfeld, _Soft photon theorems from CFT Ward identites in the flat limit of AdS/CFT_ , _JHEP_ 11 (2020) 009, [2005.03667].
  * (27) E. Witten, _SL(2,Z) action on three-dimensional conformal field theories with Abelian symmetry_ , in _From Fields to Strings: Circumnavigating Theoretical Physics: A Conference in Tribute to Ian Kogan_ , pp. 1173–1200, 7, 2003.  hep-th/0307041.
  * (28) H.-U. Yee, _A Note on AdS / CFT dual of SL(2,Z) action on 3-D conformal field theories with U(1) symmetry_ , _Phys. Lett. B_ 598 (2004) 139–148, [hep-th/0402115].
  * (29) D. Marolf and S. F. Ross, _Boundary Conditions and New Dualities: Vector Fields in AdS/CFT_ , _JHEP_ 11 (2006) 085, [hep-th/0606113].
  * (30) P. Breitenlohner and D. Z. Freedman, _Stability in gauged extended supergravity_ , _Annals of physics_ 144 (1982) 249–281.
  * (31) P. Breitenlohner and D. Z. Freedman, _Positive energy in anti-de sitter backgrounds and gauged extended supergravity_ , _Physics letters B_ 115 (1982) 197–201.
  * (32) S. Bilson, _Extracting Spacetimes using the AdS/CFT Conjecture: Part II_ , _JHEP_ 02 (2011) 050, [1012.1812].
  * (33) T. Hartman and J. Maldacena, _Time evolution of entanglement entropy from black hole interiors_ , _Journal of High Energy Physics_ 2013 (2013) , [1303.1080].
  * (34) D. Colton, _The confluent hypergeometric function. by herbert buchholz. springer-verlag, new york (1969). 238 pp._ , _Canadian Mathematical Bulletin_ 13 (1970) 164–164.