# Denoising diffusion delensing: reconstructing the non-Gaussian CMB lensing potential with diffusion models

Thomas Flöss, William R. Coulton, Adriaan J. Duivenvoorden, Francisco Villaescusa-Navarro, Benjamin D. Wandelt

► To cite this version:

Thomas Flöss, William R. Coulton, Adriaan J. Duivenvoorden, Francisco Villaescusa-Navarro, Benjamin D. Wandelt. Denoising diffusion delensing: reconstructing the non-Gaussian CMB lensing potential with diffusion models. Monthly Notices of the Royal Astronomical Society, 2024, 533, pp.423-432. 10.1093/mnras/stae1818. insu-04851529

HAL Id: insu-04851529

https://insu.hal.science/insu-04851529v1

Submitted on 20 Dec 2024

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Denoising diffusion delensing: reconstructing the non-Gaussian CMB lensing potential with diffusion models

Thomas Flöss<sup>1,2*</sup> William R. Coulton<sup>3,4</sup> Adriaan J. Duivenvoorden,<sup>5</sup>

Francisco Villaescusa-Navarro<sup>5,6</sup> and Benjamin D. Wandelt<sup>5,7</sup>

Accepted 2024 July 22. Received 2024 July 22; in original form 2024 June 7

## ABSTRACT

Optimal extraction of cosmological information from observations of the cosmic microwave background (CMB) critically relies on our ability to accurately undo the distortions caused by weak gravitational lensing. In this work, we demonstrate the use of denoising diffusion models in performing Bayesian lensing reconstruction. We show that score-based generative models can produce accurate, uncorrelated samples from the CMB lensing convergence map posterior, given noisy CMB observations. To validate our approach, we compare the samples of our model to those obtained using established Hamiltonian Monte Carlo methods, which assume a Gaussian lensing potential. We then go beyond this assumption of Gaussianity, and train and validate our model on non-Gaussian lensing data, obtained by ray-tracing *N*-body simulations. We demonstrate that in this case, samples from our model have accurate non-Gaussian statistics beyond the power spectrum. The method provides an avenue towards more efficient and accurate lensing reconstruction, which does not rely on an approximate analytical description of the posterior probability. The reconstructed lensing maps can be used as an unbiased tracer of the matter distribution, and to improve delensing of the CMB, resulting in more precise cosmological parameter inference.

**Key words:** cosmic background radiation – large-scale structure of Universe.

## 1 INTRODUCTION

The deflection of photons of the cosmic microwave background (CMB) by the large-scale structure between today and the last-scattering surface changes the CMB statistics, leaving a characteristic imprint (Cole & Efstathiou 1989; Linder 1990; Seljak 1996; Met-calf & Silk 1997; Lewis & Challinor 2006). This effect, known as weak gravitational lensing, affects our ability to infer properties of the early universe such as primordial gravitational waves through B-mode polarization (Seljak & Hirata 2004) or the presence of primordial non-Gaussianity (Babich & Zaldarriaga 2004; Coulton et al. 2020; Coulton, Miranthis & Challinor 2023), both of which would constitute strong evidence for the inflationary paradigm and are thus important targets for experimental endeavours (Achúcarro et al. 2022). Furthermore, the gravitational potential that lenses the CMB is an unbiased tracer (albeit projected) of the total matter distribution of the universe and thus contains important cosmological information by itself. Accurately reconstructing the CMB lensing potential is therefore of considerable interest (Green, Meyers &

van Engelen 2017; Hotinli et al. 2022; Ange & Meyers 2023; Trendafilova, Hotinli & Meyers 2023), for example, to constrain the mass of neutrinos (Abazajian et al. 2015), the late-time growth of structure (i.e. $\sigma_8$), and local primordial non-Gaussianity, especially in cross-correlation with other tracers, through cosmic-variance cancellation (Seljak 2009; Schmittfull & Seljak 2018).

Several methods have been developed to reconstruct the lensing potential from observed CMB data, starting with the quadratic estimator by Seljak and Zaldarriaga (Zaldarriaga & Seljak 1999), and Hu and Okamoto (Hu & Okamoto 2002; Okamoto & Hu 2003). These estimators have been used in virtually all CMB lensing analyses to date (Das et al. 2011; van Engelen et al. 2012; Ade et al. 2014a, b, 2016; Story et al. 2015), for which they proved to be sufficiently optimal. The high resolution and low noise levels of current [South Pole Telescope 3G (SPT-3G); Benson et al. 2014; Prabhu et al. 2024] and next-generation [Simons Observatory (Ade et al. 2019) and CMB-S4 (Abazajian et al. 2016)] CMB surveys allow sub-per cent level lensing reconstruction and an unprecedented sensitivity to primordial B modes. At this point, however, the quadratic estimator ceases to be nearly optimal, and more sophisticated techniques are required to achieve optimal lensing reconstruction, such as iterative maximum a posteriori (MAP) estimators (Hirata & Seljak 2003a,

<footnote n="1"><sup>1</sup>Van Swinderen Institute for Particle Physics and Gravity, University of Groningen, Nijenborgh 3, NL-9747 AG Groningen, the Netherlands</footnote>

<footnote n="2"><sup>2</sup>Kapteyn Astronomical Institute, University of Groningen, NL-9700 AV Groningen, the Netherlands</footnote>

<footnote n="3"><sup>3</sup>Kavli Institute for Cosmology, Madingley Road, Cambridge CB3 0HA, UK</footnote>

<footnote n="4"><sup>4</sup>DAMTP, Centre for Mathematical Sciences, Wilberforce Road, Cambridge CB3 0WA, UK</footnote>

<footnote n="5"><sup>5</sup>Center for Computational Astrophysics, 162 5th Avenue, New York, NY 10010, USA</footnote>

<footnote n="6"><sup>6</sup>Department of Astrophysical Sciences, Princeton University, 4 Ivy Lane, Princeton, NJ 08544, USA</footnote>

<footnote n="7"><sup>7</sup>Sorbonne Université, CNRS, UMR 7095, Institut d'Astrophysique de Paris, 98 bis bd Arago, F-75014 Paris, France</footnote>

<footnote>*E-mail: tsfloss@gmail.com</footnote>
b; Carron & Lewis 2017; Millea, Anderes & Wandelt 2019) and gradient inversion techniques (Hadzhiyska et al. 2019). Besides these analytical approaches, the use of machine learning methods in CMB lensing problems has been an active topic of research in recent years as well (Caldeira et al. 2019; Guzman & Meyers 2021; Li et al. 2022; Yan et al. 2023). Although the quality of the reconstructed lensing potential with iterative or machine-learned estimators may be statistically optimal, both methods yield point estimates that require expensive simulations to propagate uncertainty in the estimates when they are used in any subsequent analyses (e.g. to infer primordial non-Gaussianity from delensed CMB data). Additionally, non-Gaussianity of the lensing potential has been shown to contain significant amounts of information on e.g. neutrino mass and dark energy (Namikawa 2016). Reconstructing the lensing convergence bispectrum is therefore of considerable interest. However, determining non-Gaussian statistics of the lensing potential using conventional estimators can be challenging due to having to model complicated noise biases (Kalaja et al. 2023), although this may be somewhat alleviated with alternative estimators (Namikawa, Hanson & Takahashi 2013b).

More recently, the problem of CMB delensing and lensing reconstruction has received a fully Bayesian treatment by Millea, Anderes, and Wandelt (Millea et al. 2019, 2020), which we will refer to as MAW [see also Anderes, Wandelt & Lavaux (2015) for earlier work using only temperature modes]. Their approach to this high-dimensional statistical problem relies on having access to the gradient of the CMB lensing posterior, allowing for the use of a Hamiltonian Monte Carlo (HMC) within Gibbs sampler that simultaneously samples cosmological parameters (the amplitude $A_\phi$ of the lensing potential power spectrum and the tensor-to-scalar ratio $r$), the delensed CMB data, and the lensing potential. The ability to sample from the posterior enables the accurate determination of uncertainties when working with the delensed data or the inferred lensing potential. This powerful method has since been used to improve the analysis of data from SPT-3G (Millea et al. 2021).

One of the crucial assumptions for MAW being able to write down and evaluate the posterior is that the lensing potential is Gaussian. However, in reality, this assumption does not hold since the large-scale structure of the universe is non-Gaussian. More accurate reconstruction of the lensing potential will thus require going beyond this Gaussian prior assumption, but this is not straightforward with existing methods. At the same time, the HMC method typically requires long sampling chains to ensure convergence of the lensing maps, which scales inefficiently with the size of the data. Although this issue can be partially ameliorated with suitable approximations (Millea 2022; Millea & Seljak 2022), it is worth investigating alternative simulation-based inference solutions (Alsing et al. 2019; Cranmer, Brehmer & Louppe 2020).

Here, we explore the use of probabilistic machine learning models in reconstructing the lensing potential, alleviating some of the limitations of the HMC approach. Specifically, we consider the use of score-based generative models (SGMs; Song et al. 2020), a variant of what is commonly referred to as diffusion models. Using techniques inspired by the stochastic nature of diffusion, these models can draw uncorrelated samples from complicated high-dimensional probability distributions, without requiring an explicit expression for this distribution, making them particularly suitable for problems involving images (Ho, Jain & Abbeel 2020). Simply put, such models learn to map samples from a simple distribution (i.e. a Gaussian) into samples of a more complicated distribution, iteratively turning pure noise into meaningful data. Diffusion models have been successfully applied to various problems in the context

of astrophysics and cosmology (e.g. Adam et al. 2022; Karchev et al. 2022; Heurtel-Depeiges et al. 2023; Legin et al. 2023; Remy et al. 2023; Rouhiainen et al. 2023; Schanz, List & Hahn 2023; Ono et al. 2024). Here, we will demonstrate the use of such models in reconstructing the CMB lensing convergence map given noisy CMB observations, enabling sampling of the Bayesian lensing posterior. While MAW solves the problem with explicitly specified likelihood and priors, our method implements an implicit inference approach where likelihood and prior are represented in terms of pairs of lensing and data maps. In our case, these are generated through simulations, making our work an example of simulation-based inference. Because of this, we are able to reconstruct non-Gaussian lensing potentials with accurate non-Gaussian statistics for the first time, requiring only a change in training data. Our model, trained on non-Gaussian lensing maps from N-body simulations, generates samples with the correct non-Gaussian statistics, as demonstrated by the one-point probability distribution function (PDF) and bispectrum. On signal-dominated scales, the sampled maps show excellent agreement with the true non-Gaussian lensing potential, while on noise-dominated scales, they revert to the learned non-Gaussian prior. Notably, models trained only on Gaussian lensing maps fail to capture the full non-Gaussian structure. Furthermore, our method allows for fast and uncorrelated sampling of the lensing potential posterior, circumventing the convergence issues and inefficiencies faced by traditional Markov chain Monte Carlo methods. These advancements open up new possibilities for efficient delensing and precise constraints on large-scale structure from upcoming CMB surveys.

The paper is organized as follows. In Section 2, we go over the necessary basics of CMB lensing. In Section 3, we introduce the Bayesian approach to CMB delensing by MAW. In Section 4, we introduce some basic aspects of score-based generative diffusion models. To validate the use of these models for Bayesian lensing reconstruction tasks, we perform a detailed comparison against the method and code of MAW in Section 5. Finally, we demonstrate that our model can be used for the accurate reconstruction of non-Gaussian lensing potentials in Section 6.

## 2 CMB LENSING BASICS

The intervening matter distribution gravitationally bends the paths of CMB photons that travel through it, such that we observe deflected photons. This induces a remapping of the true CMB temperature ($T$) and polarization (Stokes $Q$, $U$) modes to the observed, lensed quantities ($\tilde{T}$, $\tilde{Q}$, $\tilde{U}$) (for simplicity neglecting a small phase correction; Challinor & Chon 2002):

$$ \begin{aligned} \tilde{T}(\hat{n}) &= T(\hat{n} + \alpha(\hat{n})), \\ (\tilde{Q} \pm i\tilde{U})(\hat{n}) &= (\tilde{Q} \pm iU)(\hat{n} + \alpha(\hat{n})), \end{aligned} \qquad (1) $$

where $\hat{n}$ is the unit vector that denotes the position in the sky and $\alpha$ is the deflection angle that encodes the remapping. This deflection angle is given by the gradient of the lensing potential $\phi$:

$$ \alpha = \nabla\phi, \qquad (2) $$

which itself is a two-dimensional projection of the three-dimensional gravitational potential $\psi$:

$$ \phi(\hat{n}) = -2 \int d\chi \frac{\chi_{\text{CMB}} - \chi}{\chi_{\text{CMB}}\chi} \psi(\hat{n}\chi, \chi), \qquad (3) $$

where $\chi$ denotes the comoving distance and $\chi_{\text{CMB}}$ is the comoving distance to the last-scattering surface, where the CMB was emitted (at redshift $z \approx 1100$). In the flat-sky approximation, these fields
can also be written in Fourier space, parametrized by wavevector $\ell$. Another useful quantity is the CMB lensing convergence $\kappa$:

$$ \kappa(\hat{n}) = -\frac{1}{2} \nabla^2 \phi(\hat{n}), \qquad (4) $$

which is most easily obtained in Fourier space as

$$ \kappa_{\ell} = \frac{|\ell|^2}{2} \phi_{\ell}. \qquad (5) $$

Using a Taylor expansion of equation (1), one can determine the lowest order effect of lensing on the CMB fluctuations. This can then be used to construct a quadratic estimator for the CMB lensing convergence (Zaldarriaga & Seljak 1999; Hu & Okamoto 2002; Okamoto & Hu 2003). Additionally, iterating this estimator yields the MAP estimate of the lensing convergence (Hirata & Seljak 2003a, b; Carron & Lewis 2017).

## 3 BAYESIAN DELENSING

In Millea et al. (2019, 2020), MAW present a Bayesian approach to lensing reconstruction and delensing of the CMB fluctuations. The approach takes as a starting point the posterior probability distribution of the true CMB $f$ (i.e. $T, Q, U$), the lensing potential $\phi$, and cosmological parameters $\theta$ (the amplitude of the lensing power spectrum $A_{\phi}$ and the tensor-to-scalar ratio $r$), given observed data $d$ (i.e. $\tilde{T}, \tilde{Q}, \tilde{U}$): $P(f, \phi, \theta|d)$. This data constitutes lensed, noised, and masked CMB data:

$$ d = A \mathcal{L}(\phi) f + n \qquad (6) $$

where $\mathcal{L}(\phi)$ denotes the lensing operation, $n$ is the noise, and $A$ is a linear transformation that includes instrumental effects such as the beam and sky mask. Using Bayes' theorem, this posterior can be expressed in terms of the likelihood of the data, and a prior (following the concise notation of MAW):

$$ P(f, \phi, \theta|d) \propto P(d|f, \phi, \theta)P(f, \phi, \theta). \qquad (7) $$

Under the assumption that the noise is a Gaussian random field with covariance $C_n$, the likelihood is proportional to

$$ P(d|f, \phi, \theta) \propto \exp \left( - \frac{(d - A \mathcal{L}(\phi) f)^2}{2C_n} \right), \qquad (8) $$

and the priors on the $f$ and $\phi$ fields are taken to be Gaussian:

$$ P(f, \phi, \theta) \propto \exp \left( -\frac{f^2}{2C_f} \right) \exp \left( -\frac{\phi^2}{2C_\phi} \right) P(\theta). \qquad (9) $$

Since we are after the pixels of the fields $f$ and $\phi$, this constitutes a high-dimensional posterior distribution that is intractable with conventional Metropolis-Hastings Monte Carlo methods. MAW overcome this challenge by numerically implementing the posterior in a fully differentiable fashion (see CMBLENSING.JL<fnref n="1" />), providing access to the derivatives of the posterior with respect to all parameters (i.e. field pixels and $\theta$). This enables both MAP estimation through gradient-based optimization (Millea et al. 2019) and HMC sampling (Millea et al. 2020) of the posterior. Crucially, even if one is only after the lensing potential, they find that it is easier to sample the full posterior $P(f, \phi, \theta|d)$, rather than the marginalized posterior $P(\phi, \theta|d)$.

Although the Gaussian prior on the lensing potential does not imply that the posterior samples are necessarily Gaussian, it is not

the maximally informative prior for data that have been lensed with a non-Gaussian potential, and we expect samples in this case to have inaccurate non-Gaussian statistics. Including such a non-Gaussian prior in the approach of MAW would require forward modelling and evaluation of the posterior, which is expensive and hard to implement. Additionally, HMC methods can be inefficient, due to subsequent samples being correlated. In practice, the sample chains of the large-scale modes of the lensing potential have long autocorrelation lengths, resulting in a low yield of independent samples of these scales, especially when a mask is included (Carron & Lewis 2017; Millea et al. 2020).

In the rest of this work, we will demonstrate the use of generative machine learning models in reconstructing the CMB lensing convergence to overcome both of these limitations of the HMC approach, albeit at the expense of losing provable long-run convergence to the true posterior.

## 4 SCORE-BASED GENERATIVE MODELS

For our probabilistic machine learning approach, we adopt SGMs (Song et al. 2020). These models can produce samples of a learned data distribution $p(x)$ through a reverse-diffusion process. Starting from data $x$, we incrementally perturb the data using Gaussian noise with an increasing variance schedule $\sigma^2(t)$ over many time-steps $t \in [0, 1]$, until it is pure noise. This diffusion process is described by a stochastic differential equation (SDE; Song et al. 2020):

$$ d\mathbf{x} = f(\mathbf{x}, t)dt + g(t)d\mathbf{w}, \qquad (10) $$

where $f$ is called the drift coefficient, $g$ is the diffusion coefficient, and $d\mathbf{w}$ denotes a Wiener (noising) process. Remarkably, the diffusion process has a reverse SDE that turns pure Gaussian noise into a sample of $p(x)$ by reversing the noising process at each time-step; the denoising process (Anderson 1982):

$$ d\mathbf{x} = \left[ f(\mathbf{x}, t) - g(t)^2 \nabla_x \log p_t(\mathbf{x}) \right] dt + g(t)d\mathbf{w}. \qquad (11) $$

This requires access to the score $\nabla_x \log p_t(\mathbf{x})$, where $p_t(\mathbf{x})$ is the probability distribution of $\mathbf{x}$ at time $t$. Furthermore, if we are after the conditional distribution $p(\mathbf{x}|\mathbf{y})$ (i.e. the posterior of parameters $\mathbf{x}$ given data $\mathbf{y}$), we instead need the conditional score $\nabla_x \log p_t(\mathbf{x}|\mathbf{y})$. In our case, $\mathbf{x}$ represents the lensing convergence map, while $\mathbf{y}$ represents observed lensed CMB data (e.g. $Q$ and $U$ maps). We use a neural network, more specifically a U-Net, to approximate this conditional score by training it to predict random noise added to an image $\mathbf{x}$ at a random time-step, given conditioning data $\mathbf{y}(\mathbf{x})$ (Batzolis et al. 2021).

Typically in diffusion models, the noising process $d\mathbf{w}$ uses white noise

$$ z_w \sim \mathcal{N}(0, I), \qquad (12) $$

but we find that both training and sampling are faster and more accurate when using reddened Gaussian noise with a power spectrum similar to the target $\kappa$ maps:

$$ z_{\kappa}(z_w) = \Sigma_{\kappa}^{1/2} z_w, \qquad (13) $$

where $\Sigma_{\kappa}$ is the theory covariance matrix of the lensing convergence maps, and the multiplication is performed in Fourier space, where this covariance matrix is diagonal (i.e. the power spectrum $C_{\ell}^{\kappa\kappa}$). This type of noise allows the model to affect all scales of the data at every time-step, instead of different scales at different time-steps, resulting in higher quality samples with fewer diffusion steps. Furthermore, we employ the Variance Exploding variant of the algorithm (Song

<footnote n="1">¹https://github.com/marius311/CMBLensing.jl</footnote>
<table><caption>**Table 1.** Configuration used in our experiments. The set-up follows that of the 2PARAM set-up of Millea et al. (2020).</caption><thead><tr><th>Parameter</th><th>Configuration</th></tr></thead><tbody><tr><td>CMB data</td><td>Q, U</td></tr><tr><td>Map size (pixels)</td><td>256 × 256</td></tr><tr><td>Pixel width</td><td>2 arcmin</td></tr><tr><td>Total area</td><td>73 deg<sup>2</sup></td></tr><tr><td>White noise level in P</td><td>1 μK arcmin</td></tr><tr><td>(ℓ<sub>knee</sub>, α<sub>knee</sub>)</td><td>(100, 3)</td></tr><tr><td>Beam FWHM</td><td>2 arcmin</td></tr><tr><td>Fourier masking</td><td>2 &lt; ℓ &lt; 5000</td></tr><tr><td>Pixel masking</td><td>0.4° border + 0.6° apod</td></tr></tbody></table>

et al. 2020):

$$f(\mathbf{x}, t) = 0, \quad g(t) = \sqrt{\frac{d\sigma^2(t)}{dt}}, \quad \sigma(t) = \sigma_{\min} \left( \frac{\sigma_{\max}}{\sigma_{\min}} \right)^t, \quad (14)$$

where $\sigma_{\min}$ and $\sigma_{\max}$ set the lowest and highest noise scales of the noising process. Given this set-up, we use the noise schedule $\sigma_{\min} = 0.01$ and $\sigma_{\max} = 100$. This provides sufficient noise to properly swamp the image, as well as sufficient low-noise steps to fine-tune the sample. Finally, to not underemphasize small scales in the model's training objective, the score-matching loss to be optimized by the network is still defined at the level of the white noise:

$$L = \mathbb{E}_{t, x_0, z_w} \|z_w - s_{\theta}(x_0 + \sigma(t)z_{\kappa}(z_w), t)\|^2, \quad (15)$$

where $s_{\theta}(x_t, t)$ is the U-Net's noise prediction at time $t$, given the noisy image $x_t$.

After training, we sample from the learned conditional distribution by discretizing the reverse SDE with the Euler-Maruyama method:

$$x_{t-\Delta t} = x_t - g(t)^2 (\Sigma_{\kappa}^{1/2} s_{\theta}(x_t, t)) \Delta t + g(t) \sqrt{\Delta t} z_{\kappa}, \quad (16)$$

and applying this equation iteratively starting from pure noise at time $t=1$. We use 1000 time-steps to draw samples, but good results can be achieved with fewer steps.

More details on the implementation of the model and U-Net can be found in Appendix A.

## 5 GAUSSIAN LENSING RECONSTRUCTION

To validate our method, we first perform a benchmark comparison against the HMC algorithm, using CMBLENSING.JL. As noted, this algorithm can be used to simultaneously sample the true CMB, the lensing potential $\phi$, and parameters $A_{\phi}$ (the amplitude of the lensing potential power spectrum) and $r$ (the tensor-to-scalar ratio). In this work, we will only focus on reconstructing the lensing potential, thus implicitly marginalizing over the true CMB. Ideally, we would marginalize over both the amplitude and tensor-to-scalar ratio. However, we found that the score-based generative algorithm, as well as variants thereof (e.g. Denoising Diffusion Probabilistic Models), struggle to yield samples with the correct amplitude if we vary the lensing amplitude. We comment more on this in Appendix B. For now, we limit our set-up to fixed amplitude $A_{\phi} = 1$, but we do marginalize over $r$.

We mimic the 2PARAM survey configurations in MAW's work and use CMBLENSING.JL to generate training data, consisting of 32768 pairs of Gaussian $\kappa$ and lensed $(Q, U)$ maps, including masks and noise with specifications given in Table 1 and random $r \in [10^{-6}, 10^{-1}]$ (quadratically spaced). The $(Q, U, \kappa)$ maps are each normalized by the mean and standard deviation over the entire data set. Furthermore, we apply a cut-off $\ell < 3000$ to the $\kappa$ maps

<figure><img src="image_2.png" /><figcaption>**Figure 1.** Validation data used in the CMBLENSING.JL benchmark following the configuration presented in Table 1. Top: The lensed $(Q, U)$ observation. Bottom left: The true lensing convergence map $\kappa$ that lensed the $(Q, U)$ observation. Bottom right: The MAP estimate of the lensing convergence given the $(Q, U)$ observation, using the optimization algorithm of CMBLENSING.JL.</figcaption></figure>

during training, so the network only learns to reconstruct modes below this cut-off (the $(Q, U)$ maps still have $\ell < 5000$). In our set-up, modes beyond this cut-off are severely noise-dominated in the reconstruction samples and therefore not informative. Finally, since these maps have been generated with periodic boundary conditions, we apply periodic padding in our U-Net.

We validate our generative model by comparing 2048 samples of our model and a CMBLENSING.JL chain of 17000 samples (we drop the first 1000 samples and keep every fifth sample, leaving us with 3200 samples from the chain), for a validation case with parameters $(A_{\phi} = 1, r = 0.04)$, shown in Fig. 1. In Fig. 2, we show the mean and standard deviation over posterior samples drawn using the two methods, while in Fig. 3 we show the mean power spectrum of the samples. Additionally, we show the mean power spectrum of the reconstruction residual of samples $(\kappa_{\text{samples}} - \kappa_{\text{true}})$ in dashed lines, which quantifies the signal-to-noise ratio of the reconstruction on different scales. We see that for this configuration the reconstruction is signal-dominated up to $\ell \approx 800$. As a consequence, on small, noise-dominated scales the power spectrum of samples becomes increasingly prior driven, and the mean power spectrum of samples reduces to the theory spectrum, as observed in the bottom panel. We can further investigate the quality of our samples by comparing the correlation coefficient between samples and the true lensing convergence map as a function of scale:

$$C_{\ell} = \frac{C_{\ell}^{\kappa \text{samples}^{\kappa \text{true}}}}{\sqrt{C_{\ell}^{\kappa \text{true} \kappa \text{true}} C_{\ell}^{\kappa \text{samples}^{\kappa \text{samples}}}}}, \quad (17)$$

shown in Fig. 4. The various metrics demonstrate good agreement between the CMBLENSING.JL chain samples and SGM samples.

We conclude that our SGM has learned an accurate approximation of the lensing posterior, and can therefore be used to draw samples of the reconstructed lensing convergence. Our SGM model has the additional advantage that samples are uncorrelated, resulting in a more effective sampling of large angular scales compared to the
<figure><img src="image_3.png" /><figcaption>**Figure 2.** Results for the validation sample of the CMBLENSING.JL benchmark. Top: Mean over 2048 posterior samples from the SGM (left) and the CMBLENSING.JL chain (right). Bottom: Standard deviation of posterior samples from 1024 diffusion samples (left) and the CMBLENSING.JL chain (right).</figcaption></figure>

<figure><img src="image_4.png" /><figcaption>**Figure 3.** Top: Mean power spectrum of SGM samples, CMBLENSING.JL samples, and the true validation $\kappa$ map. The reconstruction noise curves give the mean power spectrum of reconstruction residuals ($\kappa_{sample} - \kappa_{true}$). The shaded regions and error bars demark the $1\sigma$ spread of the power spectra. Bottom: Mean power spectrum of samples divided by the theory power spectrum (at $A_{\phi} = 1$).</figcaption></figure>

CMBLENSING.JL chains, that struggle to converge here for masked data. In fact, for the CMBLENSING.JL chain presented here, the largest scale has not yet converged, explaining the large discrepancy of the largest scale bandpower with our SGM result, which more closely follows the truth.

<figure><img src="image_5.png" /><figcaption>**Figure 4.** Cross-correlations (equation 17) of validation samples with the true lensing convergence map. Solid lines denote cross-correlation of the mean of posterior samples from the SGM and CMBLENSING.JL chain. Dashed lines denote the mean of individual cross-correlations of samples.</figcaption></figure>

# 6 NON-GAUSSIAN LENSING RECONSTRUCTION

Contrary to the MAW approach based on an explicitly specified posterior, our approach is based on a posterior that is implicitly specified through the training data. Our generative model therefore allows for a straightforward extension to non-Gaussian lensing potentials; we simply need to train our model on data consisting of non-Gaussian lensing maps and lensed CMB data. To demonstrate this, we train it using a set of simulated non-Gaussian lensing potentials provided by Takahashi et al. (2017). This data set consists of 108 high-resolution full-sky lensing convergence maps generated by ray-tracing dark matter N-body simulations.

We use 100 of their full-sky $\kappa$ maps to lens 100 independent full-sky ($Q, U$) maps using the lensing algorithm implemented in the LENSPYX code (Reinecke, Belkner & Carron 2023). We generate 256 flat-sky observations (and the corresponding patches of the lensing convergence) per full-sky map, with the same properties as before (i.e. Table 1). These 25 600 flat-sky patches serve as training data for our model, while 1 of the remaining 8 full-sky maps is used to generate independent validation data. We keep cosmological parameters fixed ($A_{\phi} = 1, r = 0$) for all data for simplicity, but in principle, $r$ can be varied in the training data as in the previous section.

Once trained, we draw 4096 samples from our model for a validation case, the results of which are shown in Fig. 5. Both the samples and the mean over posterior samples show a clear correlation with the true non-Gaussian $\kappa$ map. The power spectrum and cross-correlation coefficient for the non-Gaussian validation case are shown in Figs 6 and 7.

Different from the Gaussian case of the previous section, the non-Gaussian lensing convergence maps contain information beyond the power spectrum. Such non-Gaussian statistics of the lensing field have been studied previously in Liu et al. (2016), in particular the one-point PDF and peak counts. Additionally, we can study higher order correlation functions, such as the bispectrum (Namikawa 2016).

In Fig. 8, we show the one-point PDF of validation $\kappa$ samples, and that of both Gaussian and non-Gaussian $\kappa$ maps, clearly showing that our model produces one-point statistics that are consistent with
<figure><img src="image_6.png" /><figcaption>**Figure 5.** Results for the validation case in the *N*-body non-Gaussian set-up. Top left: Mean over 4096 posterior samples. Top right: The true lensing convergence. Bottom: Two independent samples of the convergence map.</figcaption></figure>

<figure><img src="image_7.png" /><figcaption>**Figure 6.** Top: Mean power spectrum of SGM samples for the *N*-body validation case presented in Fig. 5. The reconstruction noise curve gives the mean power spectrum of reconstruction residuals $\kappa_{sample} - \kappa_{true}$. Bottom: Mean power spectrum of samples divided by the theory power spectrum (at $A_{\phi} = 1$).</figcaption></figure>

the mean over many non-Gaussian maps, but slightly shifted towards the statistics of the true $\kappa$ map.

Next, we investigate the bispectrum of the samples. We compute the bispectrum using 17 bins with a width of four times the largest mode, $\ell_{min} \approx 43$, of the flat-sky patches, resulting in 597 triangle configurations up to $\ell_{max} \approx 2911$. The left column of Fig. 9 shows the mean bispectrum of samples and many *N*-body non-Gaussian maps. We plot the triangle configurations logarithmically

<figure><img src="image_8.png" /><figcaption>**Figure 7.** Cross-correlations of validation samples with the true lensing convergence map for the non-Gaussian case. Solid lines denote the cross-correlation of the mean of posterior samples from the SGM. Dashed lines denote the mean of individual cross-correlations of samples.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>**Figure 8.** One-point PDF of $\kappa$, using 15 linear bins between $\kappa \in (-0.05, 0.05)$. We show the PDF of a set of the *N*-body non-Gaussian $\kappa$ maps, the PDF of a set of Gaussian $\kappa$ maps generated with the same power spectrum as the non-Gaussian maps, and the PDF of a set of samples using our generative model. The error bars for the latter denote the standard deviation per bin. Dots denote the one-point PDF of the true validation $\kappa$ map.</figcaption></figure>

to emphasize triangles on large scales, where the reconstruction is signal-dominated (as in Fig. 6). On large scales (low triangle index), the mean bispectrum of samples (blue) follows the bispectrum of the true lensing convergence map (grey) more closely than it follows the theory bispectrum of non-Gaussian maps (green). On smaller scales, where the reconstruction is noise-dominated, the bispectrum of samples more closely follows the theory bispectrum (i.e. falls back to the learned non-Gaussian prior). These results show that our model has indeed learned a non-Gaussian posterior and prior and generates accurate samples of the non-Gaussian lensing convergence map.

To demonstrate the error induced by drawing samples from the less informed posterior (with Gaussian lensing prior), we show the bispectrum for samples of the validation case drawn from a model that was trained using only Gaussian lensing maps (with the same power spectrum as the non-Gaussian *N*-body lensing maps). In the right column of Fig. 9, we show the mean bispectrum of
<figure><img src="image_10.png" /><figcaption>**Figure 9.** Bispectrum results for an SGM trained on non-Gaussian (left) or Gaussian (right) lensing data. Top: Bispectrum of the true validation lensing convergence $\kappa$ map, mean bispectrum of many different non-Gaussian $\kappa$ maps, and SGM samples of the true $\kappa$ map (blue). Bottom: Mean bispectrum of SGM samples of many different $\kappa$ maps. The mean bispectrum of many different maps is an approximation to the theory bispectrum.</figcaption></figure>

4096 validation samples from this Gaussian model. We expect such a model to be able to capture non-Gaussian information only on highly signal-dominated regions where the posterior is not dominated by the Gaussian prior (that is implicitly learned from the Gaussian training data). Indeed, the top panel shows that on the largest, signal-dominated scales, the model produces a bispectrum that follows the true validation map and is similar to that of the SGM trained on non-Gaussian maps. However, on smaller, noise-dominated scales the bispectrum vanishes, as the model falls back on the internally learned Gaussian prior, unlike our non-Gaussian model. Furthermore, in the bottom panel, we see that the mean bispectrum over 4096 samples of different lensing realizations does not closely follow the theory bispectrum.

## 7 CONCLUSION AND OUTLOOK

In this work, we demonstrated the use of SGMs in CMB lensing reconstruction tasks. We have shown that such models can effectively learn the conditional posterior distribution of the lensing convergence map given observed CMB data, by training on corresponding pairs of these. The trained model can be used to rapidly draw lensing posterior samples. For the case of Gaussian lensing convergence maps, we have validated our methodology by comparing the results of our model against the established HMC methods of Millea, Anderes, and Wandelt (Millea et al. 2021).

Since our model learns the posterior distribution entirely from simulated data, and the prior and likelihood are modelled through the processes encoded in the simulations, it can be used in situations where the analytical form of the posterior is not known, such as in the case of non-Gaussian lensing convergence maps. Using realistic non-

Gaussian lensing maps obtained by ray-tracing N-body simulations, we demonstrate that our model can indeed be trained to yield samples with the correct non-Gaussian statistics (i.e. one-point probability distribution and bispectrum). As we have shown, the resulting samples of our model can be used to estimate the CMB lensing bispectrum, which is a difficult task using traditional estimators such as the quadratic estimator, because of large noise biases (Kalaja et al. 2023), and the estimated lensing bispectrum can be used together with the lensing power spectrum to tighten cosmological constraints (Namikawa 2016).

Another important advantage of our model as compared to the HMC method is speed and efficiency. Our model, with neural network configuration as described in Appendix A and using 1000 denoising steps, generates a single sample in 12 s, which can be vectorized up to 128 simultaneous samples in about 225 s (~1.75 s per sample) on an A100 GPU, and can be easily parallelized over multiple GPUs. Significant speed improvements can be expected from optimizing neural network complexity and noise schedule (i.e. number of denoising steps), which is beyond the scope of this paper. Notably, the samples of our model are completely uncorrelated. Meanwhile, drawing a single sample with CMBLENSING.JL (in the configuration used in this paper, varying only *r*) takes ~1.6 s on an A100 GPU. However, HMC samples are correlated. Taking the autocorrelation length of the bandpower on the least efficiently sampled scale of the chain (which is found to be about 290) shows that the effective independent sample size of the chain is more than two orders of magnitude smaller. It should be noted that improving the HMC analysis with suitable approximations used in real lensing analyses is likely to again yield an order of magnitude improvement in speed for the configuration considered in this work (Millea 2022; Millea &
Seljak 2022). Finally, for larger maps, the HMC analysis requires more samples to converge, thus scaling inefficiently with problem size. For our model, besides the obvious increase in training and sampling time due to the larger data size, there is no such chain convergence issue, thus promising more efficient scaling to larger problems. We conclude that generative models allow for a more efficient approach to Bayesian CMB lensing reconstruction, which translates to faster convergence of analysis pipelines.

The methods described in this paper, as well as those in Millea et al. (2019, 2020, 2021), are so far limited to only using a small and flat patch of the sky at a time. Although adaptation of our methods to curved-sky geometry using spherical convolutional neural networks (e.g. Defferrard et al. 2020; Esteves, Slotine & Makadia 2023) is in principle straightforward, such an analysis is likely to be computationally prohibitive at present. Nevertheless, we can expect ongoing advancements in machine learning and computational technology to make such analyses more feasible. Additionally, it is worth noting that validation of conventional methods, such as the quadratic or iterative estimator, also requires many evaluations on simulated data to determine the uncertainty. In our methods, as well as the HMC method, this is inherently included in the analysis.

Looking ahead, ideally, the model would include the amplitude of the CMB lensing power spectrum as well. As mentioned earlier, we have been unsuccessful at this so far. Nevertheless, we are optimistic that with an appropriate modification of the model or algorithm, this can be overcome. We share some of our thoughts on this issue in Appendix B. Alternatively, sampling the cosmological parameters can also be left as separate Gibbs-like steps using another simulation-based inference pipeline conditioned on the data and our model's lensing samples. We leave an investigation into this for future work.

Though in this work we have focused on the non-Gaussianity of the lensing potential, our model's flexibility straightforwardly allows for other extensions that would similarly complicate the posterior evaluation, such as the inclusion of inhomogeneous and anisotropic noise in the data (Hanson, Rocha & Górski 2009). Once again, this would require only a modification of the simulated data that the model is trained on. Along these lines, obvious directions for future work are to assess biases due to foregrounds (Engelen et al. 2012; Osborne, Hanson & Doré 2014; MacCrann et al. 2024) and correlated point-source masks (Lembo et al. 2022) that complicate the reconstruction of the lensing potential and reconstruction of lensing curl modes (Cooray, Kamionkowski & Caldwell 2005; Namikawa, Yamauchi & Taruya 2012, 2013a; Pratten & Lewis 2016; Robertson & Lewis 2023; Belkner et al. 2024).

The results of this work pave the way for more efficient, accurate, and realistic lensing reconstruction. The resulting samples of our model can be used in subsequent analyses of the CMB or as an unbiased tracer of the total matter density distribution while being able to appropriately propagate errors due to uncertainty in the lensing reconstruction.

## ACKNOWLEDGEMENTS

The authors would like to thank Julien Carron, Shirley Ho, François Lanusse, Ronan Legin, Florian List, Daan Meerburg, Marius Millea, and Sebastian Wagner-Carena for useful discussions. We furthermore thank Marius Millea for providing support with CMBLENSING.JL. TF thanks the Flatiron Institute's Center for Computational Astrophysics, where part of this research was conducted, for its hospitality. We thank the Center for Information Technology of the University of Groningen for providing access to the Hábrók high-performance computing cluster. TF was supported by the Fundamentals of the

Universe research programme at the University of Groningen. BDW acknowledges support from the DIM ORIGINES project 'INFINITY NEXT - Calcul Hybride et IA en Astrophysique'. The Flatiron Institute is supported by the Simons Foundation.

## DATA AVAILABILITY

Our model and the code used to generate the results of this paper are publicly available at https://github.com/tsfloss/CMBLensingDiffusion/tree/paper.

## REFERENCES

* Abazajian K. N. et al., 2015, Astropart. Phys., 63, 66

* Abazajian K. N. et al., 2016, preprint (arXiv:1610.02743)

* Achúcarro A. et al., 2022, preprint (arXiv:2203.08128)

* Adam A., Coogan A., Malkin N., Legin R., Perreault-Levasseur L., Hezaveh Y., Bengio Y., 2022, in 36th Conference on Neural Information Processing Systems: Machine Learning and the Physical Sciences Workshop. p. E1, preprint (arXiv:2211.03812)

* Ade P. A. R. et al., 2014a, Phys. Rev. Lett., 113, 021301

* Ade P. A. R. et al., 2014b, A&A, 571, A17

* Ade P. A. R. et al., 2016, A&A, 594, A15

* Ade P. et al., 2019, J. Cosmol. Astropart. Phys., 2019, 056

* Alsing J., Charnock T., Feeney S., Wandelt B., 2019, MNRAS, 488, 4440

* Anderes E., Wandelt B., Lavaux G., 2015, ApJ, 808, 152

* Anderson B. D., 1982, Stochastic Processes and Their Applications, Vol. 12. p. 313

* Ange J., Meyers J., 2023, J. Cosmol. Astropart. Phys., 2023, 045

* Babich D., Zaldarriaga M., 2004, Phys. Rev. D, 70, 083005

* Batzolis G., Stanczuk J., Schönlieb C.-B., Etmann C., 2021, preprint (arXiv:2111.13606)

* Belkner S., Carron J., Legrand L., Umiltà C., Pryke C., Bischoff C., 2024, ApJ, 964, 148

* Benson B. A. et al., 2014, in Holland W. S., Zmuidzinas J., eds, Proc. SPIE Conf. Ser. Vol. 9153, Millimeter, Submillimeter, and Far-Infrared Detectors and Instrumentation for Astronomy VII. SPIE, Bellingham, p. 91531P

* Bradbury J. et al., 2018, JAX: composable transformations of Python + NumPy programs, http://github.com/google/jax

* Caldeira J. a., Wu W. L. K., Nord B., Avestruz C., Trivedi S., Story K. T., 2019, Astron. Comput., 28, 100307

* Carron J., Lewis A., 2017, Phys. Rev. D, 96, 063510

* Challinor A., Chon G., 2002, Phys. Rev. D, 66, 127301

* Cole S., Efstathiou G., 1989, MNRAS, 239, 195

* Cooray A., Kamionkowski M., Caldwell R. R., 2005, Phys. Rev. D, 71, 123527

* Coulton W. R., Meerburg P. D., Baker D. G., Hotinli S., Duivenvoorden A. J., van Engelen A., 2020, Phys. Rev. D, 101, 123504

* Coulton W., Miranthis A., Challinor A., 2023, MNRAS, 523, 825

* Cranmer K., Brehmer J., Louppe G., 2020, Proc. Natl. Acad. Sci. USA, 117, 30055

* Das S. et al., 2011, Phys. Rev. Lett., 107, 021301

* Defferrard M., Milani M., Gusset F., Perraudin N., 2020, preprint (arXiv:2012.15000)

* Esteves C., Slotine J.-J., Makadia A., 2023, preprint (arXiv:2306.05420)

* Green D., Meyers J., van Engelen A., 2017, J. Cosmol. Astropart. Phys., 2017, 005

* Guzman E., Meyers J., 2021, Phys. Rev. D, 104, 043529

* Hadzhiyska B., Sherwin B. D., Madhavacheril M., Ferraro S., 2019, Phys. Rev. D, 100, 023547

* Hanson D., Rocha G., Górski K., 2009, MNRAS, 400, 2169

* Heurtel-Depeiges D., Burkhart B., Ohana R., Régaldo-Saint Blancard B., 2023, preprint (arXiv:2310.16285)

* Heurtel-Depeiges D., Margossian C. C., Ohana R., Régaldo-Saint Blancard B., 2024, preprint (arXiv:2402.19455)
Hirata C. M., Seljak U., 2003a, Phys. Rev. D, 67, 043001

Hirata C. M., Seljak U., 2003b, Phys. Rev. D, 68, 083002

Ho J., Jain A., Abbeel P., 2020, Denoising Diffusion Probabilistic Models

Hotinli S. C., Meyers J., Trendafilova C., Green D., van Engelen A., 2022, J. Cosmol. Astropart. Phys., 2022, 020

Hu W., Okamoto T., 2002, ApJ, 574, 566

Kalaja A., Orlando G., Bowkis A., Challinor A., Meerburg P. D., Namikawa T., 2023, J. Cosmol. Astropart. Phys., 2023, 041

Karchev K., Anau Montel N., Coogan A., Weniger C., 2022, in 36th Conference on Neural Information Processing Systems: Workshop on Machine Learning and the Physical Sciences

Legin R., Ho M., Lemos P., Perreault-Levasseur L., Ho S., Hezaveh Y., Wandelt B., 2023, MNRAS, 527, L173

Lembo M., Fabbian G., Carron J., Lewis A., 2022, Phys. Rev. D, 106, 023525

Lewis A., Challinor A., 2006, Phys. Rep., 429, 1

Li P., Ilayda Onur I., Dodelson S., Chaudhari S., 2022, preprint (arXiv:2205.07368)

Linder E. V., 1990, MNRAS, 243, 353

Liu J., Hill J. C., Sherwin B. D., Petri A., Böhm V., Haiman Z., 2016, Phys. Rev. D, 94, 103501

MacCrann N. et al., 2024, ApJ, 966, 138

Metcalf R. B., Silk J., 1997, ApJ, 489, 1

Millea M., 2022, preprint (arXiv:2209.10512)

Millea M., Seljak U., 2022, Phys. Rev. D, 105, 103531

Millea M., Anderes E., Wandelt B. D., 2019, Phys. Rev. D, 100, 023509

Millea M., Anderes E., Wandelt B. D., 2020, Phys. Rev. D, 102, 123542

Millea M. et al., 2021, ApJ, 922, 259

Namikawa T., 2016, Phys. Rev. D, 93, 121301

Namikawa T., Yamauchi D., Taruya A., 2012, J. Cosmol. Astropart. Phys., 2012, 007

Namikawa T., Yamauchi D., Taruya A., 2013a, Phys. Rev. D, 88, 083525

Namikawa T., Hanson D., Takahashi R., 2013b, MNRAS, 431, 609

Okamoto T., Hu W., 2003, Phys. Rev. D, 67, 083002

Ono V., Park C. F., Mudur N., Ni Y., Cuesta-Lazaro C., Villaescusa-Navarro F., 2024, preprint (arXiv:2403.10648)

Osborne S. J., Hanson D., Doré O., 2014, J. Cosmol. Astropart. Phys., 2014, 024

Prabhu K. et al., 2024, preprint (arXiv:2403.17925)

Pratten G., Lewis A., 2016, J. Cosmol. Astropart. Phys., 2016, 047

Reinecke M., Belkner S., Carron J., 2023, A&A, 678, A165

Remy B., Lanusse F., Jeffrey N., Liu J., Starck J.-L., Osato K., Schrabback T., 2023, A&A, 672, A51

Robertson M., Lewis A., 2023, J. Cosmol. Astropart. Phys., 2023, 048

Rouhiainen A., Gira M., Münchmeyer M., Lee K., Shiu G., 2023, preprint (arXiv:2311.05217)

Schanz A., List F., Hahn O., 2023, preprint (arXiv:2310.06929)

Schmittfull M., Seljak U., 2018, Phys. Rev. D, 97, 123540

Seljak U., 1996, ApJ, 463, 1

Seljak U., 2009, Phys. Rev. Lett., 102, 021302

Seljak U., Hirata C. M., 2004, Phys. Rev. D, 69, 043005

Song Y., Sohl-Dickstein J., Kingma D. P., Kumar A., Ermon S., Poole B., 2020, preprint (arXiv:2011.13456)

Story K. T. et al., 2015, ApJ, 810, 50

Takahashi R., Hamana T., Shirasaki M., Namikawa T., Nishimichi T., Osato K., Shiroyama K., 2017, ApJ, 850, 24

Trendafilova C., Hotinli S. C., Meyers J., 2023, preprint (arXiv:2312.02954)

van Engelen A. et al., 2012, ApJ, 756, 142

Yan Y.-P., Wang G.-J., Li S.-Y., Yan Y.-J., Xia J.-Q., 2023, ApJ, 952, 15

Zaldarriaga M., Seljak U., 1999, Phys. Rev. D, 59, 123507

## APPENDIX A: NEURAL NETWORK SPECIFICS

The U-Net used in our SGM employs double convolution blocks, each of which consists of a 3 × 3 convolution with 64 channels, a sinusoidal time embedding, group normalization, and SiLU activation, followed (preceded) by a Fourier pooling (upsampling) layer in the encoding (decoding) process. The Fourier pooling layer

performs the downsampling of the image in the Fourier domain, conserving the data's frequency content more accurately than typical pooling or stride convolutions. Similarly, the Fourier upsampling (or interpolation) layer upsamples the data in the Fourier domain, by filling up the modes of a higher resolution image with the modes of the lower resolution image, again exactly preserving the frequency content. We furthermore employ residual connections (rescaled with $\sqrt{2}$) and skip connections. At the network bottleneck, the data has a size of $8^2$ pixels. The network input is a channel-wise concatenation of the noise'd image $x_i$ and conditioning data $y$ (Batzolis et al. 2021). The model is trained using the AdamW optimizer with gradient norms clipped to 1, and a batch size of 128. We first train with a learning rate of 0.0002, followed by a learning rate of $10^{-5}$ to ensure convergence. Additionally, we use exponential moving average (EMA) weights, with a decay rate of 0.9999, to further stabilize the model during inference. Each model is trained for a total of around 200 000 steps, taking around 16 h on two A100 GPUs. Our model is implemented using JAX (Bradbury et al. 2018).

## APPENDIX B: VARYING THE LENSING AMPLITUDE

In Section 5, we only varied the tensor-to-scalar ratio $r$. Ideally, we would like to be able to marginalize over the amplitude $A_\phi$ of the lensing power spectrum too. However, we found that the diffusion algorithm as is, is unable to yield samples with the correct amplitude on small scales in this case, even though the phases of the samples are accurate [as quantified by the cross-correlation coefficient in equation (17), which does not contain the amplitude]. We find that distinct models converge to different results. An example of this is visualized in Fig. B1, where we show results of the mean power spectrum of samples for three different realizations and parameters, from two models (solid and dotted) trained on data that simultaneously vary $A_\phi$ and $r$, compared to those of CMBLENSING.JL chains. These two models have also been trained for a significantly longer time than

<figure><img src="image_11.png" /><figcaption>**Figure B1.** Mean power spectrum, normalized by the theory power spectrum at $A_\phi = 1$, of samples from different models (solid and dotted) and CMBLENSING.JL chains for three different realizations and parameters.</figcaption></figure>
the models presented in the main text, to ensure convergence of their
EMA weights.

Although we are uncertain as to what exactly causes this issue, we
think that it could be related to the nature of the diffusion sampling
algorithm. The algorithm assumes that the signal-to-noise ratio of
the noised image at different time-steps follows the noise schedule
σ(t) in equation (14). However, when varying the amplitude of the
underlying data, at a fixed time this ratio varies significantly among
the data, especially at lower noise levels. Meanwhile, during the
sampling process, at every denoising step, the network makes a
guess for the amplitude that will not be perfect. Thus, at the next

time-step, there is a mismatch between the actual signal-to-noise
ratio and that expected by the model. The accumulation of many of
these mismatches could the wrong amplitude, especially on noise-
dominated scales (where the value of the amplitude is most relevant),
even though the phase structure is accurately resolved. From this
perspective, a blind denoising approach, such as that presented in
Heurtel-Depeiges et al. (2024), might prove useful.

This paper has been typeset from a TeX/LATeX file prepared by the author.