Diffusion Models for Polarimetric Reconstruction of Circumstellar Environments in Correlated Speckle Noise

Abstract

High-contrast polarimetric imaging of circumstellar disks is severely limited by stellar leakage and speckle noise. Building upon the RHAPSODIE framework for polarimetric inverse problems, diffusion-based methods have shown promising results by replacing classical regularization with learned priors. However, RHAPSODIE's white noise assumption fails to capture the spatial correlation structure of atmospheric and instrumental speckles. We extend this framework to handle speckle fluctuations, which are modeled as a Gaussian field with a stationary covariance matrix with a Gaussian spectral density. To efficiently apply the dense precision matrix that includes correlated speckle and heteroskedastic photon noise, we develop a dedicated preconditioner for conjugate gradient iterations. The diffusion prior is trained exclusively on disk-only data. While limitations remain at very low SNRs, our method significantly outperforms Tikhonov regularization and competes favorably with a Plug-and-Play baseline in realistic regimes, successfully recovering morphological features masked by correlated noise.