Stochastic Super-resolution of Cosmological Simulations with Denoising Diffusion Models

TL;DR

This work introduces denoising diffusion models as a powerful generative model for super-resolving cosmic large-scale structure predictions (as a first proof-of-concept in two dimensions) and demonstrates that this model not only produces convincing super-resolution images and power spectra consistent at the percent level, but is also able to reproduce the diversity of small-scale features consistent with a given low-resolution simulation.

Abstract

In recent years, deep learning models have been successfully employed for augmenting low-resolution cosmological simulations with small-scale information, a task known as "super-resolution". So far, these cosmological super-resolution models have relied on generative adversarial networks (GANs), which can achieve highly realistic results, but suffer from various shortcomings (e.g. low sample diversity). We introduce denoising diffusion models as a powerful generative model for super-resolving cosmic large-scale structure predictions (as a first proof-of-concept in two dimensions). To obtain accurate results down to small scales, we develop a new "filter-boosted" training approach that redistributes the importance of different scales in the pixel-wise training objective. We demonstrate that our model not only produces convincing super-resolution images and power spectra consistent at the percent level, but is also able to reproduce the diversity of small-scale features consistent with a given low-resolution simulation. This enables uncertainty quantification for the generated small-scale features, which is critical for the usefulness of such super-resolution models as a viable surrogate model for cosmic structure formation.

Figures (9)

• Figure 1: Comparison between the original displacement field $\boldsymbol{{\mathit \Psi}} (\textit{top})$ and filtered displacement field $f(\boldsymbol{{\mathit \Psi}})$ (bottom) using $\gamma = 1$. Directly evident are much more pronounced small-scale features in the filtered displacement field. The diffusion model is therefore forced to gain a good understanding of these small-scale features in order to accurately estimate the noise added to the filtered fields in the forward process.
• Figure 2: Visual comparison between a randomly selected LR simulation, its HR counterpart, and a SR sample generated by our model (trained on filtered data using $\gamma = 1$). The top row represents the $x$-displacement $\bm{{\mathit \Psi}}_x$, and the bottom row the density field. The rightmost column shows the residuals between SR and HR (that indicate the difference between two possible realisations of small scales given the LR constraint). The SR density field is computed from the output of our model by using cloud-in-cell interpolation. In order to increase the contrast of the density fields, we plot $\log(1 + \delta + \epsilon)$ with $\epsilon = 0.2$. The LR density is computed from the Fourier-upscaled LR displacement field (matching the resolution of the HR simulation).
• Figure 3: Zoom into the white box of the density fields presented in Figure \ref{['im:comp']}. The top row represents the density field from the LR simulation (upsampled with Fourier interpolation to match the resolution of the HR simulation), the middle row from the HR simulation and bottom row from the SR equivalent. The inset showcases that the SR authentically reproduces the high-density nodes and even discreteness artifacts of the individual $N$-body particles.
• Figure 4: Top panel: Dimensionless power spectrum $\Delta^2$ of the SR based on one LR/HR pair for different filters. We tested applying no filter ($\gamma = 0$), a generalized Laplacian ($\gamma = 1$), and Laplacian operator ($\gamma = 2$). The original LR and HR power spectra are represented by the gray dashed and black dashed lines, respectively. The vertical lines show the Nyquist wave numbers. The shaded areas represent the 1$\sigma$ error bars generated from 100 SR and HR simulations based on one single LR simulation. Bottom panel: Ratio between all respective power spectra in comparison to the HR simulation (= unity).
• Figure 5: RGB composite image (left column) generated from 3 different SR simulations (second column) that belong to the same LR image. Each SR simulation is represented by one color channel (red, green, and blue). We upsampled the image via Fourier interpolation to a resolution of 1,024$\times$1,024 in order to suppress particle noise and highlight the structural differences between the realizations. When plotted side by side, identifying the differences between the SR images is challenging; however, the composite image highlights subtle differences along the filamentary structures and at various high density nodes.