BGRem: A background noise remover for astronomical images based on a diffusion model

Technical Summary of "BGRem: A background noise remover for astronomical images based on a diffusion model"

Problem Statement
Astronomical imaging faces a fundamental challenge: maximizing signal capture while minimizing noise interference to achieve a high signal-to-noise ratio (SNR). While enhancing SNR directly on detectors is technically and financially prohibitive, post-processing techniques are essential for catalog-building tasks. Traditional methods, such as bilateral filtering, wavelets, and non-local means, are often used, but recent machine learning (ML) approaches have shown promise. However, a significant challenge remains in generalizing ML models trained on simulated data to real-world observations, which contain complex, uncontrolled noise. Furthermore, existing tools like SExtractor rely on background subtraction methods that may not fully optimize the detection of faint sources in noisy environments.

Methodology
The authors introduce BGRem, a machine learning tool designed to remove background noise from astronomical images using a supervised diffusion model framework.

• Architecture: The backbone of BGRem is an Attention U-Net (Oktay et al. 2018), which combines the U-Net structure with attention mechanisms to improve prediction accuracy with minimal computational cost.
• Training Framework: Unlike standard Denoising Diffusion Probabilistic Models (DDPM) which iteratively add Gaussian noise to destroy an image, BGRem employs a supervised additive corruption scheme.
  • The model is trained on pairs of "clean" source-only images ($y$) and background noise ($x$).
  • During training, a noisy image $x_t$ is constructed by adding a scaled portion of the background noise to the clean image: $x_t = y + \sigma_t \cdot x$.
  • The noise scale $\sigma_t$ is governed by a cosine-based schedule, interpolating between high signal and high noise regimes.
  • The network is trained to predict the noise component ($\hat{x}$) using a Mean Absolute Error (MAE) loss function, rather than the Mean Squared Error (MSE) typically used in DDPM.
• Data Preparation:
  • Optical Data: Training data consists of simulated full-field images from the MeerLICHT telescope, generated using GalSim. These include spatially varying sky backgrounds, detector effects (read noise, dark current), and Poisson noise. Images are cropped into $256 \times 256$ sub-images.
  • Gamma-Ray Data: To test cross-wavelength applicability, the model was retrained on simulated Fermi-LAT $\gamma$-ray sky maps (1–2 GeV range). These simulations include Active Galactic Nuclei (AGNs), Pulsars (PSRs), and the Interstellar Emission (IEM) background.
• Pre- and Post-Processing: Input images undergo median subtraction to center noise and normalization. During inference, the model iteratively removes noise. Post-processing involves setting negative pixel values to zero and denormalization.

Key Contributions

1. Novel Diffusion Application: BGRem represents a novel application of diffusion models for background removal in astronomy, utilizing a supervised additive noise approach rather than the standard destructive forward process.
2. Cross-Wavelength Adaptability: The study demonstrates the model's ability to be retrained for distinct noise characteristics, successfully applying the framework to both optical (MeerLICHT) and $\gamma$-ray (Fermi-LAT) data.
3. Integration with Existing Pipelines: The authors position BGRem not as a standalone photometry tool, but as a pre-processing step to enhance traditional source detection pipelines (specifically SExtractor).

Results

• Optical Performance (MeerLICHT): When used as a pre-processing step for SExtractor, BGRem increased the number of true positive source detections by approximately 7% compared to using SExtractor alone. The model effectively removed background noise while preserving point sources.
• Generalizability (Zero-Shot): BGRem was tested on optical images from the Canada-France–Hawaii Telescope (CFHT) and the DESI Legacy Survey (DECam) without retraining (zero-shot inference). In both cases, it improved source detection efficiency over SExtractor's native background removal, although it struggled to preserve low-surface-brightness (LSB) extended structures (e.g., M31), interpreting them as noise.
• Gamma-Ray Performance: On simulated Fermi-LAT data, BGRem successfully removed the IEM background, allowing SExtractor to identify more than double the number of correct sources compared to the standard method.
• Limitations Identified:
  • Photometric Bias: Setting negative pixels to zero introduces a bias, making BGRem unsuitable for standalone high-precision photometry. The authors propose using BGRem outputs solely to generate detection masks, which are then applied to original data for flux integration.
  • Extended Sources: The model, trained primarily on point sources, tends to suppress extended, diffuse structures.
  • Poisson Noise in $\gamma$-rays: While effective against IEM, the current model fails to properly handle observational Poisson noise in $\gamma$-ray data, leading to under-prediction of counts in bright patches and over-prediction in faint patches.

Significance and Claims
The paper claims that BGRem offers a robust, general-purpose framework for background removal that can enhance the accuracy of source detection in traditional pixel-based methods. By successfully applying a diffusion-based approach to both optical and $\gamma$-ray domains, the authors highlight its potential as a versatile pre-processing step for multi-wavelength astronomical surveys. The work serves as a proof-of-concept for integrating diffusion models into the ASID (Astronomical Source Identification and Detection) pipeline, aiming to complement rather than replace established tools like SExtractor. The authors maintain a modest tone, acknowledging that while BGRem improves detection efficiency, it currently requires further optimization (e.g., handling Poisson noise and preserving extended morphology) before being suitable for all scientific applications, particularly those requiring surface photometry.