Score-matching neural networks for improved multi-band source separation

We present the implementation of a score-matching neural network that represents a data-driven prior for non-parametric galaxy morphologies. The gradients of this prior can be incorporated in the optimization of galaxy models to aid with tasks like deconvolution, inpainting or source separation. We demonstrate this approach with modification of the multi-band modeling framework scarlet that is currently employed as deblending method in the pipelines of the HyperSuprimeCam survey and the Rubin Observatory. The addition of the prior avoids the requirement of non-differentiable constraints, which can lead to convergence failures we discovered in scarlet. We present the architecture and training details of our score-matching neural network and show with simulated Rubin-like observations that using a data-driven prior outperforms the baseline scarlet method in accuracy of total flux and morphology estimates, while maintaining excellent performance for colors. We also demonstrate significant improvements in the robustness to inaccurate initializations. The trained score models used for this analysis are publicly available at https://github.com/SampsonML/galaxygrad.