Category-based Galaxy Image Generation via Diffusion Models

1. Problem Statement

Galaxy image generation is crucial for astronomical research, particularly for data augmentation in classification tasks, testing classification algorithms, and simulating observational data for future surveys (e.g., Euclid, CSST). However, existing methods face significant limitations:

• Traditional Methods: Semi-analytical models and hydrodynamic simulations rely heavily on explicit physical assumptions and parameter tuning, making them computationally expensive and inflexible.
• Generative Models (VAEs & GANs):
  • VAEs: Often produce blurred images lacking fine structural details.
  • GANs: Suffer from training instability, mode collapse, and a tendency to bias outputs toward the training set, failing to synthesize rare features absent from the data.
• Diffusion Models: While promising for high-quality generation, existing applications in astronomy (e.g., Astrddpm) are often unconditional, computationally costly, and struggle to generate specific galaxy morphologies without training separate models for each class. Furthermore, they often lack mechanisms to explicitly integrate physical constraints or category information effectively into the network architecture.

The core challenge addressed is developing a data-driven, category-conditioned diffusion model that generates high-quality, physically consistent galaxy images for specific morphological classes while maintaining computational efficiency and avoiding the need for separate models per class.

2. Methodology

The authors propose GalCatDiff, a novel framework leveraging Denoising Diffusion Probabilistic Models (DDPM) with specific architectural enhancements tailored for astronomy.

A. Model Backbone: Enhanced U-Net

The framework utilizes a modified U-Net architecture as the backbone for the reverse diffusion process. Key enhancements include:

• Wider Architecture: Incorporates multiple Astro-Residual Attention Blocks (Astro-RAB) in each down-sampling and up-sampling stage.
• Enhanced Skip Connections: Introduces additional skip connections from every convolutional or Astro-RAB layer in the down-sampling path to the corresponding up-sampling layer to facilitate multi-scale feature fusion and reduce detail loss.

B. Novel Component: Astro-RAB (Astro-Residual Attention Block)

This is the core innovation of the paper. The Astro-RAB is designed to dynamically balance local feature extraction with global consistency.

• Structure: It combines a standard Convolutional Block (for local spatial features) with an Attention Fusion Unit (using window attention for global dependencies).
• Dynamic Weighting: The block learns optimal weights ($w_{conv}$ and $w_{attn}$) during training to combine the outputs of the convolution and attention mechanisms.
• Embedding Integration:
  • Category Embedding: Discrete galaxy class indices are mapped to continuous vectors via an MLP and injected into the block.
  • Time Embedding: Sinusoidal time-step embeddings are also integrated.
• Function: This design allows the model to capture global structural consistency (via attention) while preserving fine-grained local details (via convolution), ensuring the generated galaxies adhere to physical properties.

C. Training Strategy

• Dataset: Utilized the Galaxy6 DECaLS dataset (a refined subset of Galaxy10 DECaLS), containing 14,497 images across 6 morphological classes (e.g., Round Smooth, Barred Spiral, Edge-on with/without Bulge).
• Preprocessing: Images were downsampled to $64 \times 64$ pixels.
• Regularization: Category embeddings were randomly dropped with a 10% probability to encourage the model to learn latent data features and improve diversity.
• Loss: Standard DDPM noise prediction loss ($L_\theta$).

3. Key Contributions

1. First Category-Conditioned Diffusion Framework in Astronomy: GalCatDiff is the first framework to leverage both image features and astrophysical category information within the network design of diffusion models for galaxy generation, eliminating the need for separate models per class.
2. Astro-RAB Module: Introduction of a novel block that dynamically fuses attention and convolution mechanisms. This ensures global consistency (morphology) and local fidelity (texture/structure), addressing the specific physical requirements of galaxy images.
3. Enhanced U-Net Architecture: A wider U-Net with dense skip connections that significantly improves the reconstruction of high-resolution details compared to standard diffusion backbones.
4. Physical Consistency Verification: The study goes beyond visual quality, rigorously evaluating the physical properties (color, size, flux) of generated images against real data.

4. Results

The model was evaluated using the Galaxy6 DECaLS test set (2,902 images) against a baseline Astrddpm (Smith et al. 2022) and an Astrddpm with Class Embedding.

A. Image Quality (FID Score)

• Performance: GalCatDiff achieved a Total FID of 5.354, significantly outperforming the baseline Astrddpm (30.493) and the class-embedded baseline (11.735).
• Improvement: This represents a ~54% improvement in overall FID over the class-embedded baseline.
• Ablation: Removing Astro-RAB increased the FID to 5.420, suggesting that while the U-Net architecture is the primary driver of quality, Astro-RAB provides marginal but consistent improvements.

B. Physical Consistency (Wasserstein-1 Distance)

The authors defined metrics for Color-WD (color distribution) and HLR-WD (half-light radius distribution), combined into a Weighted Synthetic Galaxy Distance (WSGD).

• WSGD: GalCatDiff achieved a WSGD of 0.063, a 38.8% improvement over the class-embedded baseline (0.103) and a 74% improvement over the original Astrddpm (0.243).
• Stability: The inclusion of Astro-RAB significantly reduced the standard deviation of the metrics (e.g., WSGD std dev reduced by 45.45%), indicating more robust and stable generation.
• Distribution Matching: Histograms showed that GalCatDiff closely matched the peak positions and tail distributions of real data for color indices (g-r, r-z) and half-light radii, whereas baselines showed significant shifts or underestimations.

C. Ablation Studies on Astro-RAB

• Weight Dynamics: Analysis of learned weights revealed a global-to-local integration pattern. In early layers (down-sampling/bottleneck start), attention weights ($w_{attn}$) dominate to capture global structure. In later layers (up-sampling/bottleneck end), convolution weights ($w_{conv}$) dominate to refine local details.
• Criticality of Bottleneck: Ablating the late bottleneck (layers 15–16) caused the most severe performance degradation (WSGD increased by 187%), confirming that the transition from global modeling to local detail refinement in the bottleneck is crucial.

5. Significance and Future Outlook

• Data Augmentation: GalCatDiff serves as a powerful data augmentor, capable of generating realistic, class-specific galaxy images. This is vital for training deep learning models on imbalanced datasets (e.g., rare galaxy types) and for downstream tasks like classification, detection, and redshift prediction.
• Physical Reliability: Unlike previous generative models that often produce "hallucinated" features, GalCatDiff ensures that generated images adhere to physical constraints (size, color, morphology), making them suitable for scientific simulation.
• Scalability: By using category embeddings, the framework avoids the computational cost of training separate models for each of the hundreds of potential galaxy morphologies.
• Future Potential: The authors suggest that integrating dynamical equations into this framework could further enhance the generation of complex, rare morphologies (e.g., merging or disturbed galaxies), supporting the data needs of upcoming large-scale surveys like Euclid and CSST.

In conclusion, GalCatDiff represents a significant step forward in astronomical generative modeling, successfully bridging the gap between high-fidelity image synthesis and physical consistency through a novel, category-aware diffusion architecture.