Scalable Bayesian uncertainty quantification with data-driven priors for radio interferometric imaging

The interferometric measurement equation for a radio telescope (Thompson et al., 2017) in the monochromatic setting relates our observations represented by the visibility function $\mathcal{Y}$ to the sky brightness $\mathcal{X}$, which we want to reconstruct,

where $\mathbf{u}=(u,v,w)$ are the interferometer baseline coordinates with units depending on the observation wavelength, $\mathbf{l}=(l,m,n)$ are cosine sky coordinates restricted to the unit sphere, and $\mathcal{A}$ includes direction-dependent effects (DDEs) like the primary beam of the dishes. The previous general model allows us to consider different DDEs through $\mathcal{A}$ and non-coplanar effects through the exponential term in $w$. These effects become considerable when considering wide fields of view and long baselines. There exists a rich body of literature incorporating such effects, e.g., Smirnov, O. M. (2011a, b, c, d); Thompson et al. (2017), and there are scalable algorithms that take them into account, e.g., Pratley et al. (2019b).

In this article, for the sake of simplicity but without loss of generality, we assume the coplanar setting, where the antennas are located in the same $w$ plane. We also assume that we observe a small field of view such that $1-l^{2}-m^{2}\approx 1$. Consequently, we have $\exp\left[-2\pi\text{i}w\left(\sqrt{1-l^{2}-m^{2}}-1\right)\right]\approx 1$, and Equation 1 reduces to

From the previous equation, we can notice the remarkable result of $\mathcal{Y}(u,v)=\mathcal{F}(\mathcal{A}\mathcal{X})(u,v)$ where $\mathcal{F}$ is the two-dimensional Fourier transform.

To further simplify the problem, we will avoid using the continuous $\mathcal{Y}$ and $\mathcal{X}$ and work with their discrete counterparts, $\bm{x}$ and $\bm{y}$, respectively. The observational model we study for our RI imaging problem writes

where $\bm{y}\in\mathbb{C}^{M}$ are the $M$ observed complex visibilities, $\bm{x}\in\mathbb{R}^{N}$ is the discrete sky brightness sampled on a $N$ point grid, and $\mathbf{\Phi}\in\mathbb{C}^{M\times N}$ is the linear measurement operator that models the acquisition process. Without loss of generality, the observational and instrumental noise $\bm{n}\in\mathbb{C}^{M}$ is assumed to be independent and identically distributed (iid) white Gaussian noise with zero mean and standard deviation $\sigma$. If the noise is not white, we can incorporate a noise whitening matrix in the $\mathbf{\Phi}$ operator such that the previous white noise assumption holds.

Each pair of antennas provides us with one visibility, which is a noisy Fourier component of the intensity image. Using an array of $n$ radio antennas allows us to sample $\binom{n}{2}=(n^{2}-n)/2$ points in the $uv$-plane (or Fourier plane). The distribution of these points depends on the configuration of the radio antenna array. If different time intervals are considered, the Earth’s rotation can be exploited to increase the number of $uv$ points. The $uv$ coverage is incomplete in all practical cases, and the measurements are noisy. Therefore, the linear operator $\mathbf{\Phi}$ is ill-posed. If we also consider a large number of measurements, recovering $\bm{x}$ from $\bm{y}$ becomes a challenging inverse problem.

The most basic reconstruction of $\bm{x}$ is often referred to as the naturally weighted dirty image that we will denote from now on as $\hat{\bm{x}}_{\text{dirty}}$ and call by dirty image as calibration weigths are not currently being taken into account in this work. This estimation is obtained by applying the adjoint of $\mathbf{\Phi}$ to the visibilities $\bm{y}$. To obtain a higher fidelity solution to the RI imaging inverse problem, we must regularise the problem by incorporating some prior information about the desired solutions $\bm{x}$. A broad range of methods can be characterized by what type of prior information is used to regularise the inverse problem and which algorithm is used to compute the reconstructed image $\hat{\bm{x}}$.

The last two decades have brought us a great number of RI imaging methods based on sparse representations. The prior information exploited is that the solution $\bm{x}$ is known to be sparsely represented in some bases or dictionaries. The bases are often built using multi-scale wavelets, or a dictionary is constructed with a collection of wavelets (Mallat, 2008). We can represent our image $\bm{x}$ in a dictionary $\mathbf{\Psi}\in\mathbb{C}^{N\times L}$,

where $\bm{a}\in\mathbb{C}^{L}$ is a vector of coefficients of $\bm{x}$ weighting the corresponding dictionary atoms of $\mathbf{\Psi}$. The assumption made to regularise the inverse problem is that $\bm{a}$ is sparse or compressible, meaning that most of the coefficients are zero-valued or near zero, respectively. An array $\bm{a}$ is called $k$-sparse if it has only $k$ non-zero elements, which can be written as $\|\bm{a}\|_{0}=k$, where $\|\cdot\|_{0}$ denotes the $\ell_{0}$ pseudonorm.

Sparsity should be ideally enforced through the $\ell_{0}$ pseudonorm, which is non-convex. Consequently, a convex relaxation to the $\ell_{1}$ norm is used, which is a sparsity-promoting norm. The optimisation problem is formulated such that its solution coincides with the inverse problem solution. Therefore, the inverse problem can be tackled with an optimisation algorithm. The optimisation objective comprises two competing terms: (i) a data-fidelity term $f(\cdot)$ that promotes consistency with the observed visibilities and depends on the statistics of the noise $\bm{n}$; and (ii) a regularisation term $r(\cdot)$ that encodes our prior knowledge of $\bm{x}$. The optimisation problem reads

where we are using a sum of regularisation terms $r_{k}$, each with its corresponding regularisation strength parameter $\lambda_{k}$. Substituting the RI data fidelity and sparsity-enforcing regularisation in an overdetermined wavelet dictionary $\Psi$ terms into Equation 3 we obtain

where $\sigma$ is the noise standard deviation.

The previous formulation is referred to as unconstrained. Other works consider the constrained formulation, which minimises the $\ell_{1}$ term with respect to a hard $\ell_{2}$-ball constraint over $\bm{y}$ with a radius of $\epsilon$, which is related to the noise’s $\sigma$ (Carrillo et al., 2012; Pratley et al., 2017). In this article, we will focus on the unconstrained formulation as it has a natural Bayesian interpretation. Obtaining the solution from Equation 6 involves solving a convex optimisation problem, where we have the sum of a differentiable and a non-differentiable term. Proximal algorithms (Parikh & Boyd, 2014) are well suited to tackle such optimisation problems. Recent developments brought us a wide collection of proximal optimisation algorithms, such as the forward-backward (FB) algorithm (Combettes & Pesquet, 2009), the FISTA algorithm (Beck & Teboulle, 2009), the alternating direction method of multipliers (ADMM, Boyd et al., 2011), and the primal-dual forward-backward algorithm (Chambolle & Pock, 2011; Condat, 2013), to mention a few. A rich literature exists exploiting the aforementioned concepts to tackle the RI imaging problem (Wiaux et al., 2009; McEwen & Wiaux, 2011; Carrillo et al., 2012, 2014; Onose et al., 2016; Pratley et al., 2017, 2019b, 2019a; Pratley & McEwen, 2019; Cai et al., 2018b). For example, the Sparsity Averaging Reweighted Analysis (SARA) family of methods (Carrillo et al., 2012) use an over-complete dictionary composed of a concatenation of the Dirac basis and the first eight Daubechies wavelets (Daubechies, 1992) and has shown good performance in RI imaging.

Precursor of RI image reconstructions, the CLEAN algorithm (Högbom, 1974) is a highly successful RI imaging method and it is still being used (Event Horizon Telescope Collaboration et al., 2019a, b) despite various negative characteristics. The CLEAN algorithm is a non-linear iterative method that assumes a sparse sky model. CLEAN iteratively removes the contribution of the brightest source convolved with the instrument’s point spread function or dirty beam. This method can be interpreted as a matching pursuit algorithm (Wiaux et al., 2009), or an $\ell_{0}$ regularisation with a basis composed of a sum of Dirac spikes. Several extensions of CLEAN have been developed over time (Bhatnagar & Cornwell, 2004; Cornwell, 2008; Stewart et al., 2011; Offringa et al., 2014; Offringa & Smirnov, 2017) achieving better reconstruction performance. See Rau et al. (2009) for a review of CLEAN-based algorithms.

On top of a very early introduction, the success of CLEAN resides in its scalability. However, CLEAN has been shown to produce artefacts when point sources do not well describe the underlying sky model limiting CLEAN’s image quality and justifying the need for more advanced techniques, e.g. based on Section 2.2. CLEAN often requires manual intervention, making its use less practical. Furthermore, CLEAN and its extensions do not provide meaningful uncertainty quantification of its reconstruction.

The advent of deep learning models has affected many imaging applications, and RI imaging is no exception. Handcrafted models and priors are limited in the information they can capture or represent with respect to recent, more expressive neural networks. Learned or data-driven methods can encode complex information existing in the data, e.g., astrophysical simulations, used in their training. In general, these approaches produce reconstructions with improved quality, a computational speed-up, or both. These reasons make learned approaches very relevant to the RI imaging reconstruction problem. However, there are issues regarding the robustness of learned methods to data distribution shifts (Hendrycks et al., 2020) and scalable methods for uncertainty quantification to the reconstruction.

Allam (2016) proposed a learned method for RI imaging based on convolutional neural networks (Dong et al., 2016) originally considered for super-resolution. The approach consists of learning to post-process dirty images with variants for both, known and unknown PSF. More recently, Gheller & Vazza (2021) proposed to use a convolutional denoising autoencoder to learn to post-process radio images, e.g., the dirty image or CLEAN’s output. Connor et al. (2022) proposed a residual deep neural network (DNN) coined POLISH that works as a learned post-processing and super-resolution network. The DNN is based on the architecture proposed in Yu et al. (2018) and takes as input dirty images at different wavelengths and resolutions. POLISH outputs a clean image at a higher resolution for each wavelength and shows a better reconstruction quality than CLEAN. The proposed method has been applied to simulations from the upcoming Deep Synoptic Array-2000 (Hallinan et al., 2019) and real data from the Very Large Array (VLA, Perley et al., 2011).

The Plug-and-Play (PnP) framework (Venkatakrishnan et al., 2013) provides a way to incorporate a deep learning model into a modern optimisation algorithm. The central idea is to replace a proximal regularisation term with a denoising deep neural network. Ryu et al. (2019) studied conditions for the convergence of PnP algorithms. Pesquet et al. (2021) proposed a new term for the denoiser’s training loss that enforces the firm nonexpansiveness of the denoiser, which is usually deep learning-based. This training procedure allows the denoiser to suit a PnP framework with theoretical convergence conditions. The PnP framework with the nonexpansiveness enforced to the deep learning-based denoiser has been applied to the RI imaging problem in Terris et al. (2022), where the approach has been called AIRI for Artificial Intelligence for Regularisation in RI imaging. The approach achieved similar or better performance than competing prior-based approaches whilst providing a significant acceleration potential. The AIRI method was later validated on observations from the Australian Square Kilometre Array Pathfinder (ASKAP, Wilber et al., 2023).

Two learned approaches for an interferometric-based imager named Segmented Planar Imaging Detector for Electro-Optical Reconnaissance (SPIDER) were proposed by Mars et al. (2023). The first approach consists of a learned post-processing step from the dirty reconstruction based on a convolutional U-Net architecture (Ronneberger et al., 2015). The second approach consists of a learned multiscale iterative method coined GU-Net, which incorporates the measurement operator to include measurement information at the different steps and scales of the method. GU-Net is more efficient than standard unrolling methods due to its multi-scale nature. The numerical results show an improved reconstruction quality and a faster convergence than proximal optimisation-based methods. In the following work (Mars et al., 2024), the GU-Net was applied to the RI imaging problem. The variations of the $uv$-coverage are handled by training the neural network on a broad distribution of simulated $uv$-coverages and subsequently fine-tuning the network for a specific sampling distribution.

Aghabiglou et al. (2023) recently proposed a series of DNNs that combines notions of PnP algorithms and unrolled optimisation methods (Adler & Öktem, 2018; Monga et al., 2019). Each DNN is trained to transform a back-projected residual into an image residual, thus ideally improving the reconstruction of the previous iteration. The results show a significant speed-up with respect to AIRI or SARA-based methods while maintaining a similar reconstruction quality. Other recent approaches based on deep neural networks include: Wang et al. (2023) who proposed a denoising diffusion probabilistic model conditioned on the visibilities and the dirty reconstruction; and Schmidt et al. (2022) who proposed a convolutional neural network based on residual blocks that intend to inpaint the measurements, or recover the entire $uv$ plane from an incomplete coverage.

Bayesian inference provides a principled statistical framework to solve the inverse problem in Equation 3 with statistical guarantees. This framework builds upon Bayes’ famous theorem,

Bayes’ theorem relates the posterior distribution to the likelihood and prior terms that are the main constituents of a Bayesian model. The likelihood is associated with the data-fidelity term depending on the observational model and the noise statistics. The prior models expected properties of the solution $\bm{x}$, for example, smoothness, and piecewise regularity. This prior knowledge regularises the estimation problem.

The term in the denominator, commonly known as the Bayesian evidence, does not depend on $\bm{x}$ as we are marginalising over that variable, and it describes the likelihood of the observed data based on the modelling assumptions. The Bayesian evidence is crucial for making Bayesian model comparison (Robert, 2007), which provides us with a consistent way to compare models. Such high dimensional integrals can be effectively estimated by, for example, nested sampling techniques (Skilling, 2006; Ashton et al., 2022), or the recently introduced learned harmonic mean estimator (McEwen et al., 2021; Spurio Mancini et al., 2022; Polanska et al., 2023). Recent developments have focused on nested sampling to compute the model evidence in high-dimensional imaging problems with sparsity-based handcrafted priors (Cai et al., 2022) and deep learning-based priors (McEwen et al., 2023). Carrying out model selection is out of the scope of this work.

Under the Bayesian framework, we have the posterior distribution $p(\bm{x}|\bm{y})$ which assigns a probability to each possible solution $\bm{x}$ given some observations $\bm{y}$ and a model $\mathcal{M}$ consisting of the likelihood and prior terms. In imaging settings, explaining the information contained in the posterior distribution is not trivial due to its high-dimensional nature. The posterior distribution is generally characterised by samples computed by MCMC sampling. Efficiently sampling from high-dimensional posterior distributions is a current research topic, see e.g. Klatzer et al. (2023). Once $p(\bm{x}|\bm{y})$ is defined, we can say that the reconstruction method will be a point estimator of the posterior that will provide us with $\hat{\bm{x}}$. There are several choices for point estimators (Robert, 2007; Arridge et al., 2019), each with advantages and drawbacks. Some examples are a sample from the posterior, $\hat{\bm{x}}\sim p(\bm{x}|\bm{y})$, the maximum-a-posteriori estimator, $\hat{\bm{x}}=\operatorname*{argmax}_{\bm{x}}p(\bm{x}|\bm{y})$, or the posterior mean, $\hat{\bm{x}}=\mathbb{E}[\bm{x}|\bm{y}]$.

The posterior also provides us with consistent ways of quantifying the uncertainty of the chosen point estimate or reconstruction (Robert, 2007). For example, one way to represent uncertainty is to compute the posterior standard deviation. The pixels with a higher standard deviation are less constrained by the data and the prior allowing for more significant fluctuations.

One of the most significant drawbacks of Bayesian imaging methods is that they are known to be computationally expensive, even if there is a continuous effort targetting the scalability of these methods (Pereyra, 2016; Durmus et al., 2018; Pereyra et al., 2020, 2022; Klatzer et al., 2023).

As stated before, we need an expressive regulariser that is convex and has an explicit potential. Modern regularisers based on deep neural networks, like convolutional neural networks, used in RI imaging reconstruction methods satisfy neither of the two constraints. This last constraint, i.e., with an explicit potential required by the UQ approach, excludes a range of denoisers whose potentials are defined implicitly. PnP approaches (Terris et al., 2022) only require the denoising of the image without explicitly computing the regularisation potential. For example, a typical iteration from a PnP algorithm writes

where D is the denoiser, $f$ is the data-fidelity term, $\gamma$ is the step size, and $k$ is the iteration number. The algorithm’s convergence can be assured if D and the stepsize satisfies some conditions (Pesquet et al., 2021; Ryu et al., 2019). Even if the denoiser D is convex, we cannot use it for our approach as we must evaluate the potential.

Mukherjee et al. (2020) proposed a learned convex regulariser parametrised by the architecture of a deep input-convex neural network (ICNN, Amos et al., 2016), which is convex by construction. The training of the regulariser is done with an adversarial framework introduced by Lunz et al. (2018).

Very recently, a learnable convex-ridge regulariser neural network (CRR-NN²²2https://github.com/axgoujon/convex_ridge_regularizers, Goujon et al., 2023b) has been proposed, which comes with the required properties of being convex and having an explicit potential. In addition, the model focuses on being reliable and interpretable while still being expressive enough to provide excellent reconstruction quality. The CRR-NN regulariser, $R_{\bm{\theta}}$, has the form

where $\bm{h}_{n}$ are learnable 2D convolution kernels, $(\bm{h}_{i}\ast\bm{x})[k]$ denotes the $k$-th pixel of the resulting convolution, $N_{\text{C}}$ is the number of channels or kernels, $\psi_{i}:\mathbb{R}\mapsto\mathbb{R}$ are learnable non-linear convex profile functions with a Lipschitz continuous derivative, i.e., $\psi_{i}\in C^{1,1}(\mathbb{R})$, and $\bm{\theta}$ in $R_{\bm{\theta}}$ represents all learnable parameters. The convexity constraint on the learnable activation functions, $\psi_{i}$, is enforced by making the pointwise $\sigma_{i}:\mathbb{R}\to\mathbb{R}$ monotonically increasing, with $\psi_{i}^{\prime}=\sigma_{i}$, where $\sigma_{i}\in C^{0,1}_{\uparrow}(\mathbb{R})$, and $C^{0,1}_{\uparrow}(\mathbb{R})$ is the set of scalar Lipschitz continuous and increasing functions on $\mathbb{R}$. The $\sigma_{i}$ functions are chosen as learnable linear splines, which have been shown to be more expressive than ReLU functions in Propositions 3.3 and 3.5 from Goujon et al. (2023b, §III.B). The main difference between a prior based on the CRR-NN and a wavelet dictionary is that the kernels (or filters) and the activation (or thresholding) functions are learnt in the first one. In the second one, they are fixed or handcrafted. We refer the reader to Goujon et al. (2023b, Figs. 5 and 6) for examples of learned kernels $\bm{h}_{n}$ and activation functions for two trainings with two different noise levels. See Goujon et al. (2023b); Bohra et al. (2020) for more information on learnable splines.

In the spirit of PnP approaches, the CRR-NN training is based on the denoising problem that reads

where $\bm{y}$ is a noisy version of $\bm{x}$, and $\lambda$ is a parameter controlling the regularisation strength. The denoising problem is addressed through the fixed point of the problem, which given the convexity assumptions, is unique. A gradient step of Equation 19 reads

where $\alpha$ is the stepsize. Convergence can be guaranteed if the stepsize satisfies $\alpha\in(0,2/(1+\lambda\,\text{Lip}(\nabla R_{\bm{\theta}})))$, where $\text{Lip}(\cdot)$ denotes the Lipschitz constant. By composing $t$ gradient descent updates of Equation 20, i.e., a $t$-fold composition, we obtain a multi-gradient step denoiser that we denote $T_{R_{\bm{\theta}},\lambda,\alpha}^{t}$ following the notation of Goujon et al. (2023b).

The denoising problem in Equation 19 can be formulated as a fix point problem for the $t$-step denoiser $T_{R_{\bm{\theta}},\lambda,\alpha}^{t}$ as follows,

We build the CRR-NN training by penalising the residual of the fix point problem in Equation 21 with a loss function $\mathcal{L}$, for a training set of pairs of noiseless and noisy images $\{\bm{x}^{(m)},\bm{y}^{(m)}\}_{m=1}^{M}$, and reads

After having trained the denoiser, we define our prior potential as

where we have dropped the $\bm{\theta}_{t}^{\ast},\lambda_{t}^{\ast}$ notation for $\bm{\theta},\lambda$ and added a scaling parameter, $\mu$, to boost performance following Goujon et al. (2023b). For the optimisation algorithm, we need the Lipschitz constant of the gradient of the potential in Equation 23, which can be expressed as

which is calculated in Goujon et al. (2023b, Prop. IV.1), and $\Sigma_{\infty}=\text{diag}(\|{\sigma}_{1}^{\prime}\|_{\infty},\ldots,\|\sigma_{N_{C}}^{\prime}\|_{\infty})$, and $\mathbf{W}=[\bm{w}_{1}\cdots\bm{w}_{N_{C}}]^{T}$ where $\bm{w}_{i}$ corresponds to the filter $\bm{h}_{i}$: $\bm{h}_{i}\ast\bm{x}=\bm{w}_{i}^{T}\bm{x}$.

In our case, computing the MAP reduces to solving a convex optimisation problem. Following Equation 9, the optimisation problem we address is the following one,

where in addition we include $\iota_{\mathbb{R}^{N}}$, an indicator function enforcing the reconstructed image to be real. The proximal operator of the indicator function to a convex set is known and it amounts to a projection onto that convex set. In the the last term of Equation 25 the proximal operator of the reality constraint is the projection of the vector to the real number, which is written as $\text{Re}(\cdot)$. We have assumed a (white) Gaussian likelihood and the prior term is based on a previously trained CRR-NN. The CRR-NN is smooth with Lipschitz continuous gradients. However, the non-smoothness of the reality enforcing constraint forces us to rely on proximal algorithms (Parikh & Boyd, 2014) instead of an accelerated gradient descent method (Nesterov, 2018). In this case, we use the FISTA algorithm (Beck & Teboulle, 2009).

For the optimisation, we need the gradient of the likelihood and prior terms

where, in our case, $(\cdot)^{\dagger}$ is the complex conjugate transpose.

To ensure the algorithm’s convergence we use the stepsize $\tau=1/L$, where $L=\text{Lip}(\nabla_{\bm{x}}f(\bm{x},\bm{y})+\nabla g(\bm{x}))$. We can estimate a simple bound for the Lipschitz constant as follows

where we have exploited the result from Equation 24, and $\|\Phi^{\dagger}\Phi\|$ denotes the spectral norm, which in the case of a linear operator coincides with its maximum singular value. In the simplified problem we are considering in Section 2.1 with gridded visibilities, we have that $\|\Phi^{\dagger}\Phi\|=1$. If a more realistic linear operator should be considered, the maximum singular value could be computed iteratively via the power method (Golub & van Loan, 2013).

We initialise the optimisation with the naturally weighted dirty image, $\bm{x}_{(0)}=\text{Re}(\Phi^{\dagger}\bm{y})$. The optimisation procedure is summarised in Algorithm 1. We optimise for a fixed number of iterations $N_{\text{max}}$, or until a tolerance criterion of $\xi$ is reached. The stepsize is computed using the bound from Equation 28.

Let us define a posterior credible region with a credible level of $100(1-\alpha)\%$ as a set $C_{\alpha}\in\mathbb{R}^{N}$ satisfying

with $\mathbbm{1}_{C_{\alpha}}$ being being unity if $\bm{x}\in C_{\alpha}$ and zero otherwise. There are many regions satisfying the previous equation. We will focus on the highest posterior density region (HPD), which is defined as

where $f$ and $g$ are the potentials of our likelihood and prior terms, and $\gamma_{\alpha}$ is a constant that defines a level-set of the log-posterior such that Equation 29 holds. The HPD region has the property of having minimum volume (Robert, 2007, §5.5).

Our posterior $p(\bm{x}|\bm{y})=\exp[-f(\bm{x})-g(\bm{x})]/Z$ is log-concave on $\mathbb{R}^{N}$, where $Z$ is the Bayesian evidence. Then, following Pereyra (2017, Theorem 3.1), for any $\alpha\in(4\exp[(-N/3)],1)$, the HPD region $C_{\alpha}$ from Equation 30 is contained in

where

with a positive constant $\tau_{\alpha}=\sqrt{16\log(3/\alpha)}$ independent of $p(\bm{x}|\bm{y})$.