SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lorenzo Baldassari
lorenzo.baldassari@unibas.ch
Ali Siahkoohi
alisk@rice.edu
Josselin Garnier
Ecole Polytechnique, IP Paris
josselin.garnier@polytechnique.edu
Knut Sølha
ksolna@math.uci.edu
Maarten V. de Hoop
mvd2@rice.edu
Since their initial introduction, score-based diffusion models (SDMs) have been successfully applied to solve a variety of linear inverse problems in finite-dimensional vector spaces due to their ability to efficiently approximate the posterior distribution. However, using SDMs for inverse problems in infinite-dimensional function spaces has only been addressed recently, primarily through methods that learn the unconditional score. While this approach is advantageous for some inverse problems, it is mostly heuristic and involves numerous computationally costly forward operator evaluations during posterior sampling. To address these limitations, we propose a theoretically grounded method for sampling from the posterior of infinite-dimensional Bayesian linear inverse problems based on amortized conditional SDMs. In particular, we prove that one of the most successful approaches for estimating the conditional score in finite dimensions—the conditional denoising estimator—can also be applied in infinite dimensions. A significant part of our analysis is dedicated to demonstrating that extending infinite-dimensional SDMs to the conditional setting requires careful consideration, as the conditional score typically blows up for small times, contrarily to the unconditional score. We conclude by presenting stylized and large-scale numerical examples that validate our approach, offer additional insights, and demonstrate that our method enables large-scale, discretization-invariant Bayesian inference.
Inverse problems seek to estimate unknown parameters using noisy observations or measurements. One of the main challenges is that they are often ill-posed. A problem is ill-posed if there are no solutions, or there are many (two or more) solutions, or the solution is unstable in relation to small errors in the observations [1]. A common approach to transform the original ill-posed problem into a well-posed one is to formulate it as a least-squares optimization problem that minimizes the difference between observed and predicted data. However, minimization of the data misfit alone negatively impacts the quality of the obtained solution due to the presence of noise in the data and the inherent nullspace of the forward operator [2, 3]. Casting the inverse problem into a Bayesian probabilistic framework allows, instead, for a full characterization of all the possible solutions [4-6]. The Bayesian approach consists of putting a prior probability distribution describing uncertainty in the parameters of interest, and finding the posterior distribution over these parameters [7]. The prior must be chosen appropriately in order to mitigate the ill-posedness of the problem and facilitate computation of the posterior. By adopting the Bayesian formulation, rather than finding one single solution to the inverse
problem (e.g., the maximum a posteriori estimator [8]), a distribution of solutions—the posterior—is finally obtained, whose samples are consistent with the observed data. The posterior distribution can then be sampled to extract statistical information that allows for uncertainty quantification [9].
Over the past few years, deep learning-based methods have been successfully applied to analyze linear inverse problems in a Bayesian fashion. In particular, recently introduced score-based diffusion models (SDMs) [10] have become increasingly popular, due to their ability of producing approximating samples from the posterior distribution [11, 12]. An SDM consists of a diffusion process, which gradually perturbs the data distribution toward a tractable distribution according to a prescribed stochastic differential equation (SDE) by progressively injecting Gaussian noise, and a generative model, which entails a denoising process defined by approximating the time-reversal of the diffusion. Crucially, the denoising stage is also a diffusion process [13] whose drift depends on the logarithmic gradients of the noised data densities—the scores—which are estimated by Song et al. [10] using a neural network. Among the advantages of SDMs over other deep generative models is that they produce high-quality samples, matching the performance of generative adversarial networks [14], without suffering from training instabilities and mode-collapse [10, 15]. Additionally, SDMs are not restricted to invertible architectures like normalizing flows [16], which often limits the complexity of the distributions that can be learned. Finally, and most importantly to the scope of this work, SDMs have demonstrated superior performance in a variety of inverse problems, such as image inpainting [10, 17], image colorization [10], compressing sensing, and medical imaging [12, 18].
In the aforementioned cases, SDMs have been applied by assuming that the data distribution of interest is supported on a finite-dimensional vector space. However, in many inverse problems, especially those governed by partial differential equations (PDEs), the unknown parameters to be estimated are functions (e.g., coefficient functions, boundary and initial conditions, or source functions) that exist in a suitable function space, typically an infinite-dimensional Hilbert space. The inverse heat equation or the elliptic inverse source problem presented in [19] are typical examples of ill-posed inverse problems that are naturally formulated in infinite-dimensional Hilbert spaces. In addition to these PDE-based examples, other interesting cases that are not PDE-based include geometric inverse problems (e.g., determining the Riemann metric from geodesic information or the background velocity map from travel time information in geophysics [20]) and inverse problems involving singular integral operators [21]. A potential solution for all of these problems could be to discretize the input and output functions into finite-dimensional vectors and apply SDMs to sample from the posterior. However, theoretical studies of current diffusion models suggest that performance guarantees do not generalize well on increasing dimensions [22-24]. This is precisely why Andrew Stuart's guiding principle to study a Bayesian inverse problem for functions—"avoid discretization until the last possible moment" [5]—is critical to the use of SDMs.
Motivated by Stuart's principle, in this work we define a conditional score in the infinite-dimensional setting, a critical step for studying Bayesian inverse problems directly in function spaces through SDMs. In particular, we show that using this newly defined score as a reverse drift of the diffusion process yields a generative stage that samples, under specified conditions, from the correct target conditional distribution. We carry out the analysis by focusing on two cases: the case of a Gaussian prior measure and the case of a general class of priors given as a density with respect to a Gaussian measure. Studying the model for a Gaussian prior measure provides illuminating insight, not only because it yields an analytic formula of the score, but also because it gives a full characterization of SDMs in the infinite-dimensional setting, showing under which conditions we are sampling from the correct target conditional distribution and how fast the reverse SDE converges to it. It also serves as a guide for the analysis in the case of a general class of prior measures. Finally, we conclude this work by presenting, in Section 6, stylized and large-scale numerical examples that demonstrate the applicability of our SDM. Specifically, we show that our SDM model (i) is able to approximate non-Gaussian multi-modal distributions, a challenging task that poses difficulties for many generative models [25]; (ii) is discretization-invariant, a property that is a consequence of our theoretical and computational framework being built on the infinite-dimensional formulation proposed by Stuart [5]; and (iii) is applicable to solve large-scale Bayesian inverse problems, which we demonstrate by applying it to a large-scale problem in geophysics, i.e., the linearized wave-equation-based imaging via the Born approximation that involves estimating a  $256 \times 256$ -dimensional unknown parameter.
Related works Our work is primarily motivated by Andrew Stuart's comprehensive mathematical theory for studying PDE-governed inverse problems in a Bayesian fashion [5]. In particular, we are
interested in the infinite-dimensional analysis [7, 26], which emphasizes the importance of analyzing PDE-governed inverse problems directly in function space before discretization.
Our paper builds upon a rich and ever expanding body of theoretical and applied works dedicated to SDMs. Song et al. [10] defined SDMs integrating both score-based (Hyvärinen [27]; Song and Ermon [17]) and diffusion (Sohl-Dickstein et al. [28]; Ho et al. [29]) models into a single continuous-time framework based on stochastic differential equations. The generative stage in SDMs is based on a result from Anderson [13] proving that the denoising process is also a diffusion process whose drift depends on the scores. This result holds only in vector spaces, which explains the difficulties to extend SDMs to more general function spaces. Initially, there have been attempts to project the input functions into a finite-dimensional feature space and then apply SDMs (Dupont et al. [30]; Phillips et al. [31]). However, these approaches are not discretization-invariant. It is only very recently that SDMs have been directly studied in function spaces, specifically infinite-dimensional Hilbert spaces. Kerrigan et al. [32] generalized diffusion models to operate directly in function spaces, but they did not consider the time-continuous limit based on SDEs (Song et al. [10]). Dutordoir et al. [33] proposed a denoising diffusion generative model for performing Bayesian inference of functions. Lim et al. [34] generalized score matching for trace-class noise corruptions that live in the Hilbert space of the data. However, as Kerrigan et al. [32] and Dutordoir et al. [33], they did not investigate the connection to the forward and backward SDEs as Song et al. [10] did in finite dimensions. Three recent works, Pidstrigach et al. [24], Franzese et al. [35] and Lim et al. [36], finally established such connection for the unconditional setting. In particular, Franzese et al. [35] used results from infinite-dimensional SDEs theory (Föllmer and Wakolbinger [37]; Millet et al. [38]) close to Anderson [13].
Among the mentioned works, Pidstrigach et al. [24] is the closest to ours. We adopt their formalism to establish theoretical guarantees for sampling from the conditional distribution. Another crucial contribution comes from Batzolis et al. [39], as we build upon their proof to show that the score can be estimated by using a denoising score matching objective conditioned on the observed data [17, 40]. A key element in Pidstrigach et al. [24], emphasized also in our analysis, is obtaining an estimate on the expected square norm of the score that needs to be uniform in time. We explicitly compute the expected square norm of the conditional score in the case of a Gaussian prior measure, which shows that a uniform in time estimate is not always possible in the conditional setting. This is not surprising, given that the singularity in the conditional score as noise vanishes is a well-known phenomenon in finite dimensions and has been investigated in many works, both from a theoretical and a practical standpoint [41, 42]. In our paper, we provide a set of concrete conditions to be satisfied to ensure a uniform estimate in time for a general class of prior measures in infinite dimensions.
Pidstrigach et al. [24] have also proposed a method for performing conditional sampling, building upon the approach introduced by Song et al. [12] in a finite-dimensional setting. Like our approach, their method can be viewed as a contribution to the literature on likelihood-free, simulation-based inference [43, 44]. Specifically, the algorithm proposed by Pidstrigach et al. [24] relies on a projection-type approach that incorporates the observed data into the unconditional sampling process via a proximal optimization step to generate intermediate samples consistent with the measuring acquisition process. This allows Pidstrigach et al. [24] to avoid defining the conditional score<sup>1</sup>. While their method has been shown to work well with specific inverse problems, such as medical imaging [12], it is primarily heuristic, and its computational efficiency varies depending on the specific inverse problem at hand. Notably, their algorithm may require numerous computationally costly forward operator evaluations during posterior sampling. Furthermore, their implementation does not fully exploit the discretization-invariance property achieved by studying the problem in infinite dimensions since they employ a UNet to parametrize their score, limiting the evaluation of their score function to the training interval. The novelty of our work is then twofold. First, we provide theoretically grounded guarantees for an approach that is not heuristic and can be implemented such that it is not constrained to the grid on which we trained our network. As a result, we show that we effectively take advantage of the discretization-invariance property achieved by adopting the infinite-dimensional
<sup>1</sup>We note that, in a newer version of their paper submitted to arXiv on October 3, 2023 (four months after our submission to arXiv and NeurIPS 2023), Pidstrigach et al. [24] abandoned the projection-type approach. Instead, they invoke the conditional score function to perform posterior sampling and solve an inverse problem in a similar fashion to ours (we refer to Section 8.3 of their paper for details). However, they still do not address the well-posedness of the forward-reverse conditional SDE and the singularity of the conditional score, and their implementation is based on UNets and, thus, is not discretization-invariant.
formulation proposed by Stuart [5]. Second, we perform discretization-invariant Bayesian inference by learning an amortized version of the conditional score. This is done by making the score function depending on the observations. As a result, provided that we have access to high-quality training data, during sampling we can input any new observation that we wish to condition on directly during simulation of the reverse SDE. In this sense, our method is data-driven, as the information about the forward model is implicitly encoded in the data pairs used to learn the conditional score. This addresses a critical gap in the existing literature, as the other approach using infinite-dimensional SDM resorts to projections onto the measurement subspace for sampling from the posterior—a method that not only lacks theoretical interpretation but may also yield unsatisfactory performance due to costly forward operator computations. There are well-documented instances in the literature where amortized methods can be a preferred option in Bayesian inverse problems [45-50], as they reduce inference computational costs by incurring an offline initial training cost for a deep neural network that is capable of approximating the posterior for unseen observed data, provided that one has access to a set of data pairs that adequately represent the underlying joint distribution.
Main contributions The main contribution of this work is the analysis of conditional SDMs in infinite-dimensional Hilbert spaces. More specifically,
- We introduce the conditional score in an infinite-dimensional setting (Section 3).
- We provide a comprehensive analysis of the forward-reverse conditional SDE framework in the case of a Gaussian prior measure. We explicitly compute the expected square norm of the conditional score, which shows that a uniform in time estimate is not always possible for the conditional score. We prove that as long as we start from the invariant distribution of the diffusion process, the reverse SDE converges to the target distribution exponentially fast (Section 4).
- We provide a set of conditions to be satisfied to ensure a uniform in time estimate for a general class of prior measures that are given as a density with respect to a Gaussian measure. Under these conditions, the conditional score—used as a reverse drift of the diffusion process in SDMs—yields a generative stage that samples from the target conditional distribution (Section 5).
- We prove that the conditional score can be estimated via a conditional denoising score matching objective in infinite dimensions (Section 5).
- We present examples that validate our approach, offer additional insights, and demonstrate that our method enables large-scale, discretization-invariant Bayesian inference (Section 6).
Here, we review the definition of unconditional score-based diffusion models (SDMs) in infinite-dimensional Hilbert spaces proposed by Pidstrigach et al. [24], as we will adopt the same formalism to define SDMs for conditional settings. We refer to Appendix A for a brief introduction to key tools of probability theory in function spaces.
Let  $\mu_{\mathrm{data}}$  be the target measure, supported on a separable Hilbert space  $(H,\langle \cdot ,\cdot \rangle)$ . Consider a forward infinite-dimensional diffusion process  $(X_{t})_{t\in [0,T]}$  for continuous time variable  $t\in [0,T]$ , where  $X_0$  is the starting variable and  $X_{t}$  its perturbation at time  $t$ . The diffusion process is defined by the following SDE:
$$
d X _ {t} = - \frac {1}{2} X _ {t} d t + \sqrt {C} d W _ {t}, \tag {1}
$$
where  $C:H\to H$  is a fixed trace class, positive-definite, symmetric covariance operator and  $W_{t}$  is a Wiener process on  $H$ . Here and throughout the paper, the initial conditions and the driving Wiener processes in (1) are assumed independent.
The forward SDE evolves  $X_0 \sim \mu_0$  towards the Gaussian measure  $\mathcal{N}(0, C)$  as  $t \to \infty$ . The goal of score-based diffusion models is to convert the SDE in (1) to a generative model by first sampling  $X_T \sim \mathcal{N}(0, C)$ , and then running the correspondent reverse-time SDE. In the finite-dimensional case, Song et al. [10] show that the reverse-time SDE requires the knowledge of the score function  $\nabla \log p_t(X_t)$ , where  $p_t(X_t)$  is the density of the marginal distribution of  $X_t$  (from now on denoted  $\mathbb{P}_t$ ) with respect to the Lebesgue measure. In infinite-dimensional Hilbert spaces, there is no natural analogue of the Lebesgue measure (for additional details, see [51]) and the density is thus no longer well defined. However, Pidstrigach et al. [24, Lemma 1] notice that, in the finite-dimensional setting
where  $H = \mathbb{R}^D$ , the score can be expressed as follows:
$$
C \nabla_ {x} \log p _ {t} (x) = - (1 - e ^ {- t}) ^ {- 1} \left(x - e ^ {- t / 2} \mathbb {E} [ X _ {0} | X _ {t} = x ]\right), \tag {2}
$$
for  $t > 0$ . Since the right-hand side of the expression above is also well-defined in infinite dimensions, Pidstrigach et al. [24] formally define the score as follows:
Definition 1. In the infinite-dimensional setting, the score or reverse drift is defined by
$$
S (t, x) := - \left(1 - e ^ {- t}\right) ^ {- 1} \left(x - e ^ {- t / 2} \mathbb {E} [ X _ {0} | X _ {t} = x ]\right). \tag {3}
$$
Assuming that the expected square norm of the score is uniformly bounded in time, Pidstrigach et al. [24, Theorem 1] shows that the following SDE
$$
d Z _ {t} = \frac {1}{2} Z _ {t} d t + S (T - t, Z _ {t}) \mathrm {d} t + \sqrt {C} d W _ {t}, \quad Z _ {0} \sim \mathbb {P} _ {T}, \tag {4}
$$
is the time-reversal of (1) and the distribution of  $Z_{T}$  is thus equal to  $\mu_0$ , proving that the forward-reverse SDE framework of Song et al. [10] generalizes to the infinite-dimensional setting. The reverse SDE requires the knowledge of this newly defined score, and one approach for estimating it is, similarly to [10], by using the denoising score matching loss [40]
$$
\mathbb {E} \left[ \left\| \widetilde {S} (t, X _ {t}) + (1 - e ^ {- t}) ^ {- 1} \left(X _ {t} - e ^ {- t / 2} X _ {0}\right) \right\| ^ {2} \right], \tag {5}
$$
where  $\widetilde{S}(t, X_t)$  is typically approximated by training a neural network.
3 The conditional score in infinite dimensions
Analogous to the score function relative to the unconditional SDM in infinite dimensions, we now define the score corresponding to the reverse drift of an SDE when conditioned on observations. We consider a setting where  $X_0$  is an  $H$ -valued random variable and  $H$  is an infinite-dimensional Hilbert space. Denote by
$$
Y = A X _ {0} + B, \tag {6}
$$
a noisy observation given by  $n$  linear measurements, where the measurement acquisition process is represented by a linear operator  $A: H \to \mathbb{R}^n$ , and  $B \sim \mathcal{N}(0, C_B)$  represents the noise, with  $C_B$  a  $n \times n$  nonnegative matrix. Within a Bayesian probabilistic framework, solving (6) amounts to putting an appropriately chosen prior probability distribution  $\mu_0$  on  $X_0$ , and sampling from the conditional distribution of  $X_0$  given  $Y = y$ .
To the best of our knowledge, the only existing algorithm which performs conditional sampling using infinite-dimensional diffusion models on Hilbert spaces is based on the work of Song et al. [12]. The idea, adapted to infinite dimensions by Pidstrigach et al. [24], is to incorporate the observations into the unconditional sampling process of the SDM via a proximal optimization step to generate intermediate samples that are consistent with the measuring acquisition process. Our method relies instead on utilizing the score of infinite-dimensional SDMs conditioned on observed data, which we introduce in this work. We begin by defining the conditional score, by first noticing that, in finite dimensions, we have the following lemma:
Lemma 1. In the finite-dimensional setting where  $H = \mathbb{R}^D$ , we can express the conditional score function for  $t > 0$  as
$$
C \nabla_ {x} \log p _ {t} (x | y) = - (1 - e ^ {- t}) ^ {- 1} \left(x - e ^ {- t / 2} \mathbb {E} \left[ X _ {0} | Y = y, X _ {t} = x \right]\right). \tag {7}
$$
Since the right-hand side of (7) is well-defined in infinite dimensions, by following the same line of thought of Pidstrigach et al. [24] we formally define the score as follows:
Definition 2. In the infinite-dimensional setting, the conditional score is defined by
$$
S (t, x, y) := - \left(1 - e ^ {- t}\right) ^ {- 1} \left(x - e ^ {- t / 2} \mathbb {E} [ X _ {0} | Y = y, X _ {t} = x ]\right). \tag {8}
$$
Remark 1. It is possible to define the conditional score in infinite-dimensional Hilbert spaces by resorting to the results of [37, 38], see Appendix C.1.
For Definition 2 to make sense, we need to show that if we use (8) as the drift of the time-reversal of the SDE in (1) conditioned on  $y$ , then it will sample the correct conditional distribution of  $X_0$  given  $Y = y$  in infinite dimensions. In the next sections, we will carry out the analysis by focusing on two cases: the case of a Gaussian prior measure  $\mathcal{N}(0,C_{\mu})$ , and the case where the prior of  $X_0$  is given as a density with respect to a Gaussian measure, i.e.,
$$
X _ {0} \sim \mu_ {0}, \quad \frac {d \mu_ {0}}{d \mu} \left(x _ {0}\right) = \frac {e ^ {- \Phi \left(x _ {0}\right)}}{\mathbb {E} _ {\mu} \left[ e ^ {- \Phi \left(X _ {0}\right)} \right]}, \quad \mu = \mathcal {N} \left(0, C _ {\mu}\right), \tag {9}
$$
where  $C_\mu$  is positive and trace class and  $\Phi$  is bounded with  $\mathbb{E}_{\mu}[\| C_{\mu}\nabla_{H}\Phi (X_{0})\|^{2}] < + \infty$
4 Forward-reverse conditional SDE framework for a Gaussian prior measure