Accelerating Redshift-Conditioned Galaxy Image Synthesis with One-step Generative Modeling

Technical Summary: Accelerating Redshift-Conditioned Galaxy Image Synthesis with One-step Generative Modeling

Problem Statement
Understanding galaxy morphology evolution across cosmic time requires models capable of generating realistic galaxy populations conditioned on redshift. While large photometric surveys (e.g., SDSS, DES) provide millions of images, obtaining redshift-matched machine-learning datasets is challenging due to the difficulty of cross-matching imaging and spectroscopic catalogs. Existing generative models, such as Diffusion Probabilistic Models (DDPMs) and Flow Matching (FM), have shown promise in synthesizing astrophysical images but are intrinsically multi-step. This results in high computational costs and latency, which limits their utility for large-scale cosmological surveys, data augmentation, and iterative scientific inference where models may need to be evaluated millions of times. The core problem addressed is the trade-off between generation fidelity and computational efficiency in redshift-conditioned galaxy image synthesis.

Methodology
The authors investigate efficient generative modeling using two primary approaches: accelerated multi-step diffusion samplers and a one-step pixel-MeanFlow (p-MF) model.

1. Dataset and Setup: The study utilizes the GalaxiesML-64 dataset, consisting of galaxy images and continuous redshift conditions. The goal is to learn a conditional generator $p(x|c)$ where $x$ is the image and $c$ is the redshift.
2. Baselines (Multi-step): The authors compare against a standard DDPM baseline (following Lizarraga et al.) sampled using four different solvers:
   • Standard DDPM (1000 steps).
   • DDIM (60 steps, deterministic).
   • DPM++2M (30 steps, second-order).
   • DEIS-AB2 (30 steps, second-order).
3. Proposed Method (One-step): The authors adapt the pixel-MeanFlow (p-MF) model to a continuously conditioned setting.
   • Architecture: A latent-free U-Net (attention-augmented NCSN style) implemented in JAX/Flax.
   • Training Objective: The model learns the average velocity over a time interval $[r, t]$ rather than instantaneous noise prediction. It incorporates classifier-free guidance during training to improve conditional generation. The training target is derived from the MeanFlow identity, optimizing a velocity regression loss.
   • Inference: Generation requires a single Neural Function Evaluation (NFE), calculating the image directly from noise and the condition: $\hat{x} = \epsilon - v_\theta(\epsilon, 0, 1, w, c)$.
4. Evaluation Metrics: Performance is assessed using morphology-based statistics: ellipticity, semi-major axis, Sérsic index, and isophotal area. Metrics include:
   • Unconditional Jensen–Shannon Divergence (JSD): Measures distributional agreement between real and generated samples.
   • Conditional Bin Center Differences (BCD): Measures the agreement of morphology trends across redshift bins.

Key Contributions

1. Adaptation of p-MF: The first application of one-step diffusion models (specifically pixel-MeanFlow) to astrophysics, adapted for continuous redshift conditioning and enhanced with classifier-free guidance.
2. Efficiency-Accuracy Trade-off Analysis: A systematic evaluation comparing standard DDPM, accelerated second-order diffusion solvers (DEIS, DPM++), and one-step p-MF on astrophysical image generation.
3. Ablation Studies: Investigation of design choices (Huber loss vs. L2 loss, redshift label smoothing) specific to the p-MF framework, revealing that strategies optimal for DDPMs do not always transfer directly to one-step models.
4. Comprehensive Review: A concise synthesis of the theoretical connections between score-based diffusion, Flow Matching, and one-step generative models to inspire future astrophysical applications.

Results

• Fidelity vs. Cost: A clear trade-off exists. The standard DDPM (1000 steps) achieves the highest distributional fidelity (lowest JSD and BCD) but incurs the highest computational cost (33.5s inference, 7020 GFLOPs).
• Accelerated Solvers: Second-order samplers (DEIS-AB2 and DPM++2M) using 30 steps significantly outperform the first-order DDIM (60 steps) in the low-NFE regime, offering a strong balance of speed and accuracy.
• One-Step Performance: The p-MF model (25M parameters, 1 NFE) achieves competitive performance on several statistics (e.g., isophotal area and ellipticity) compared to accelerated solvers, with inference times comparable to 30-step solvers (~1.5s) but at a fraction of the GFLOPs (154 GFLOPs).
• Limitations of One-Step: While p-MF captures coarse morphology well, it struggles with fine-grained structural statistics, particularly the Sérsic index. Qualitative analysis indicates issues such as oversmoothing, PSF-like blurring, and loss of compact central cores, suggesting mild mode collapse in high-detail regimes.
• Ablation Findings: For the 25M p-MF model, label smoothing on the redshift condition improves conditional alignment (BCD) but slightly degrades unconditional realism (JSD). The L2 loss provided the most balanced performance for conditional generation compared to the Huber loss.

Significance and Claims
The paper claims that one-step generative models can recover key galaxy morphology statistics at orders-of-magnitude lower computational cost compared to standard DDPM sampling. While the one-step p-MF does not yet match the fine-grained structural fidelity of 1000-step DDPM, its ability to generate statistically consistent galaxy populations in a single step opens a path toward efficient conditional simulators. This advancement is significant for enabling large-scale deployment in cosmological surveys, facilitating data augmentation, and supporting simulation-based scientific inference where rapid, repeated generation is required. The authors position this work as a meaningful step toward practical, high-throughput generative modeling in astrophysics, though they acknowledge that further improvements in model scale or training objectives are needed to fully capture fine-grained structural details.