A Data-Driven Prism: Multi-View Source Separation with Diffusion Model Priors

We would like to thank the reviewer for their detailed comments and feedback. Below we address the concerns and questions brought up within the review:

Initialization

As the reviewer points out, the results of EM can depend heavily on the quality of the initialization. EM reaches a local but not necessarily global optimum. To mitigate this issue, we follow [1] and start by running an EM procedure that optimizes the parameters of a simple Gaussian diffusion model (mean and covariance). These approximate samples are then used to train our initial full diffusion model, which we then plug into our EM loop. We agree that exploring the importance of this initialization procedure helps contextualize our results, so we ran an ablation study on the Grassy MNIST experiment where we test random initialization compared to initialization using a Gaussian diffusion model optimized over 4, 16, and 64 (the default) EM laps. Below we show the results of our ablation study:

We find that the number of Gaussian EM laps used to generate the initial samples (and therefore train the initial diffusion model) has minimal impact on performance. Only starting from a randomly initialized model considerably reduces performance. While this result may seem surprising, we believe it can be explained by our ability to directly sample from the posterior given our current diffusion model prior (Eq.17). Because the likelihood is often constraining, this enables us to return high-quality posterior samples even when the prior specified by our diffusion model is not an optimal fit to the empirical distribution. As evidence for this point, we note the large gap in performance between posterior and prior samples on the Grassy MNIST experiment. This means that even if the initial prior fit is poor, the resulting posterior samples will likely not be random noise. We included this ablation study in the appendix of our work and pointed to it in Sections 5.2 and 6. We also reworked the introduction of our method in Section 4 and the pseudocode to highlight the use of the Gaussian denoiser initialization.

EM Laps

We agree that a full exploration of the performance of our model as a function of EM laps is important, especially given the higher computational cost of our model compared to the baseline methods. Beyond the performance on the one-dimensional manifold problem shown in Figures 1 and 2, we also added a comparison of our model’s performance on the Grassy MNIST experiment as a function of EM laps to the appendix. We summarize the results below:

We find that the model achieves good performance after as few as 8 EM iterations and that longer training leads to a slight degradation in performance. In the table comparing performance to baselines we report the results for the model trained for the default number (64) of EM laps.

For completeness, we also added a Figure to the appendix showing how the visual performance of the galaxy experiment varies as a function of EM Laps. However, given the restrictions of the rebuttal, we cannot reproduce that here.

Dataset Size

The reviewer brings up a valid concern about the scale of the dataset required to successfully fit a diffusion model to a distribution. Some previously explored solutions to reduce the dataset size requirement include data augmentation [2], patch-wise diffusion training [3], latent-space diffusion models [4,5], and transfer learning [6]. It would be possible to incorporate these methods during our maximization step.

To better quantify our dependence on the dataset size, we ran a set of ablation studies on our performance on the full-resolution Grassy MNIST experiment. We train our method using 1/4th, 1/16th, and 1/64th of the original grass and grass+MNIST datasets. The 1/16th and 1/64th runs are given 1/4th as many EM Laps as the original training setup to account for the smaller dataset, but all other hyperparameters are unchanged:

There is a clear degradation in performance as the grass and MNIST datasets are reduced in size. However, even with 1/16th of the original dataset (2048 grass images, 512 MNIST digits) our method still outperforms the baselines run on the full dataset. We added the dataset scaling study to the appendix and referenced it in the Discussion and Limitations section.

In terms of whether these dataset sizes are realistic for real world applications, we think the galaxy experiment serves as a relevant example. Using just the HST COSMOS field, we have millions of random cutouts and ~80k galaxy image examples. If we incorporate ground-based telescopes or future space-based telescopes, we can access orders-of-magnitude more random and galaxy datapoints [7].

Additional notes

• To ensure fair comparison, we ran a hyperparameter sweep for the baseline methods. Similarly, we made small architectural changes to improve the diffusion models used in our method and reran all experiments. Therefore, some of the metrics quoted in the tables above changed from our original submission.

[1] Rozet, François, et al. "Learning diffusion priors from observations by expectation maximization" (2024)

[2] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models" (2022)

[3] Wang, Zhendong, et al. "Patch diffusion: Faster and more data-efficient training of diffusion models" (2023)

[4] Zhang, Zhaoyu, et al. “Training Diffusion-Based Generative Models with Limited Data” (2025)

[5] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models" (2022)

[6] Hur, Jiwan, et al. "Expanding expressiveness of diffusion models with limited data via self-distillation based fine-tuning" (2024)

[7] Mellier, Y., et al. "Euclid. I. Overview of the Euclid mission" (2024)

We would like to thank the reviewer for their detailed comments and feedback. Below we address the concerns and questions brought up within the review:

Additional Results

Running on a fourth dataset was not feasible for the rebuttal, but we ran a wider set of baselines and evaluation metrics to better compare the models performance:

PSNR is the peak signal-to-noise ratio, PQM is the pqmass metric introduced in [4], FID is the FID score [5] using our MNIST classifier, and SD is the sinkhorn divergence. Note that we’ve replaced the CVAE [6] baseline with CLVM-VAE [2] since the latter is a superset of the former and achieves better performance. Across the experiments, we find that our method substantially outperforms the baselines.

Runtime and Computational Cost

In addition, we provide a detailed comparison of the computational time for each method in the table below:

All timing was done on NVIDIA A100 GPUs. Our method is more computationally expensive than the baselines, almost entirely due to the cost of sampling from the diffusion model. However, this computational cost comes with significant improvements in performance and sample quality. In our ablation studies below we also compare the performance as a function of EM laps, which can be used to control the tradeoff between training time and sample quality for our method.

Effectiveness on Linear Source Separation

The effectiveness of our method is driven by the ability to directly sample from the posterior given our current diffusion model prior (Eq.17). Because the likelihood is often constraining, this enables us to return high-quality posterior samples even when the prior specified by our diffusion model is not an optimal fit to the empirical distribution. As evidence for this point, we note the large gap in performance between posterior and prior samples on the Grassy MNIST experiments.

Ablation Studies

Per reviewer’s suggestion, in order to clarify the importance of individual components of our proposed method, we conducted an ablation study. In particular, we explored the effect of the following components on the performance of our method: (1) diffusion model architecture; (2) training length (in terms of the number of EM iterations), (3) initialization strategy, and (4) dataset size:

We find that the model achieves good performance after as few as 8 EM iterations and that longer training leads to a slight degradation in performance. In the table comparing performance to baselines we report the results for the model trained for the default number of EM laps. For dataset size, we train our method using 1/4th, 1/16th, and 1/64th of the original grass and grass+MNIST datasets. The 1/16th and 1/64th runs are given 1/4th as many EM Laps as the original training to account for the smaller dataset, but all other hyperparameters are unchanged. There is clear degradation in performance as the grass and MNIST datasets are reduced in size. However, even with 1/16th of the original dataset (2048 grass images, 512 MNIST digits) our method still outperforms the baselines run on the full dataset. Surprisingly, the number of Gaussian EM laps used to generate the initial samples (and therefore train the initial diffusion model) has minimal impact on performance. Only starting from a randomly initialized model considerably reduces performance.

We also find that model architecture has significant impact on performance: using an MLP-based architecture (5 hidden layers with 2048 hidden features per layer) gives poor results across all metrics. However, scaling down the UNet model (channels per level: (32, 64, 128) -> (16, 32), residual blocks per level: (2, 2, 2) -> (2, 2), embedding features: 64 ->16, attention head moved up one level) does not degrade performance appreciably.

However, the inductive bias of the UNet architecture alone is unlikely to be the cause of our model’s strong performance. In our extended baseline comparison above we contrast our performance with a CLVM implementation that uses neural networks as encoders and decoders. After running an extensive hyperparameter sweep, we found that using MLPs over UNet-style convolutional encoders / decoders led to a substantial improvement in performance. If the UNet inductive bias was essential, it would be surprising to find the CLVM did not benefit from it. We added a summary of this discussion in Section 6.

Linear Mixing and Known Mixing

We agree with the reviewer that the principle limitation of our method is that it requires linear mixing. As it stands, our framework cannot easily be extended to the non-linear case without substituting our joint sampling for Gibbs sampling. It is likely possible to generalize our method further under a reasonable set of assumptions, but we cannot explore those extensions within the timeframe of the rebuttal. However, we think the linear assumption is a reasonable constraint for the current work given that our method is already more flexible than previous baseline methods. For example, we make no source assumptions, do not require contrastive observations, and can (theoretically) separate an arbitrary number of sources within a single observation.

If we place no constraints on the mixing matrix, the problem becomes fully degenerate. Blind source separation methods must make strong assumptions about the sources in order to break this degeneracy [7]. However, it could be possible to incorporate a known prior distribution on the mixing matrix. One could then approximate each evaluation of Eq. 17 with a single Monte Carlo sample, and rely on the large dataset size and multiple training epochs to expose the model to the full mixing matrix distribution. Unfortunately, we were unable to run an experiment along these lines in the time allotted to the rebuttal.

Additional notes

• To ensure fair comparison, we ran a hyperparameter sweep for the baseline methods. Similarly, we made small architectural changes to improve the diffusion models used in our method and reran all experiments. Therefore, some of the metrics quoted in the table changed from our original submission.

[1] Li, Didong, et al. "Probabilistic contrastive principal component analysis” (2020)

[2] Severson, Kristen A., et al. “Unsupervised learning with contrastive latent variable models” (2019)

[3] Heurtel-Depeiges, David, et al. "Listening to the noise: Blind denoising with Gibbs diffusion" (2024)

[4] Lemos, Pablo, et al. "Pqmass: Probabilistic assessment of the quality of generative models using probability mass estimation" (2024)

[5] Heusel, Martin, et al. "GANs trained by a two time-scale update rule converge to a local Nash equilibrium" (2017)

[6] Abid, Abubakar, et al. “Contrastive Variational Autoencoder Enhances Salient Features” (2019)

[7] Hyvärinen, Aapo, et al. "Independent component analysis: algorithms and applications" (2000)

We would like to thank the reviewer for their detailed feedback. Below we address the concerns and questions brought up within the review:

Additional Comparisons

We agree that the work would benefit from more baselines, metrics, and a fuller exploration of the baselines. To address this we ran three baselines, PCPCA, CLVM-Linear, and CLVM-VAE, on all the contrastive experiments. To our knowledge, these are the best baselines available. We have run the Gibbs sampling baseline where computationally feasible. Where applicable, we track PSNR, the peak signal-to-noise ratio, PQM, the pqmass metric introduced in [1], FID, the FID score [2] using our MNIST classifier, and SD, the Sinkhorn divergence.

To ensure fair comparison, we ran an extensive hyperparameter sweep for the baseline methods. Similarly, we made small architectural changes to improve the diffusion models used in our method and reran all experiments. We did not run hyperparameter sweeps on our method due to time constraints for the rebuttal, but will do so for the camera-ready paper. Note that some of the metrics quoted in the original submission have changed, including improvements in the PCPCA performance that concerned the reviewer. However, there remains a substantial gap between the performance of the baselines and our model.

Runtime and Computational Cost

We provide a detailed comparison of the computational cost for each method in the table below:

All timing was done on NVIDIA A100 GPUs. Our method is more computationally expensive than the baselines, almost entirely due to the cost of sampling from the diffusion model. However, this computational cost comes with significant improvements in performance and sample quality. In our ablation studies below we also compare the performance as a function of EM laps which can be used to control the tradeoff between training time and sample quality for our method.

We added a section to the appendix that includes the full timing comparison above and point to that section in our discussion and limitations section and in the results for each experiment.

Ablation Studies

Per the reviewer’s suggestion, in order to clarify the importance of individual components of our method, we conducted an ablation study. We explored the effect of the following components on the performance of the method: (1) diffusion model architecture; (2) training length (in terms of the number of EM iterations), (3) initialization strategy, and (4) dataset size:

We find that model architecture is important for performance: using an MLP-based architecture (5 hidden layers with 2048 hidden features per layer) gives poor results across all metrics. However, scaling down the UNet model (channels per level: (32,64,128)->(16,32), residual blocks per level: (2,2,2)->(2,2), embedding features: 64->16, attention head moved up one level) does not degrade performance appreciably.

We also find that the model achieves good performance after as few as 8 EM iterations and that longer training leads to a slight degradation in performance. In the table comparing our performance to the baselines we report the results for the model trained for the default number of EM laps.

For dataset size, we train our method using 1/4th, 1/16th, and 1/64th of the original grass and grass+MNIST datasets. The 1/16th and 1/64th runs are given 1/4th as many EM laps as the original training to account for the smaller dataset, but all other hyperparameters are unchanged. There is clear degradation in performance as the grass and MNIST datasets are reduced in size. However, even with 1/16th of the original dataset (2048 grass images, 512 MNIST digits) our method still outperforms the baselines run on the full dataset.

Surprisingly, the number of Gaussian EM laps used to generate the initial samples (and train the initial diffusion model) has minimal impact on performance. Only starting from a randomly initialized model considerably reduces performance.

We added these ablation studies to the appendix, referred to them in Section 5.2, and discussed them in Section 6. If the reviewer would like additional ablations, we can incorporate them in the camera-ready version.

Linear Mixing and Gaussian Noise Assumptions

As it stands, our framework cannot easily be extended to the non-linear case without substituting our joint sampling for Gibbs sampling. It is likely possible to generalize our method further under a reasonable set of assumptions about the form of the mapping (i.e. a small perturbative expansion around linear), but we leave this for future work. We think the linear assumption is a reasonable constraint for the current paper given that our method is already more flexible than previous baseline methods.

The Gaussian noise assumption, however, can be relaxed within our framework. It would require access to “noise-only” views or the ability to sample from a known noise distribution. Essentially, one can consider the noise as an additional source and pre-train its corresponding diffusion model with the noise samples. From there, the rest of our machinery operates identically to the experiments presented in the paper, with the exception that the weights of the noise diffusion model are never updated. The only limitation is that the linear mixing assumption must be preserved. That is, noise that depends on the signal (for example, Poisson noise mentioned by the reviewer) cannot be modeled in this way.

Theoretical Bounds

To our knowledge, there are three relevant convergence guarantees for this method. First, in the limit of perfect optimization, infinite training examples, and a sufficiently expressive architecture, the score-based diffusion model is guaranteed to converge to the correct score [6,7]. In the limit of infinite sampling steps, our sampling method is guaranteed to sample from the learned prior [8,9]. Finally, substituting the prior score for the posterior score in Equation 6 is guaranteed to return posterior samples. However, we do not meet the requirements for any of these guarantees due to finite training dataset, finite sampling steps, and limitations of the MMPS approximation. We are unfortunately not aware of any relevant convergence bounds.

[1] Lemos, Pablo, et al. (2024) [2] Heusel, Martin, et al. (2017) [3] Li, Didong, et al. (2020) [4] Severson, Kristen A., et al. (2019) [5] Heurtel-Depeiges, David, et al. (2024) [6] Hyvärinen, A. (2005) [7] Song, Yang et al. (2019) [8] Sohl-Dickstein, Jascha, et al. (2015) [9] Song, Yang, et al. (2020)

We would like to thank the reviewer for their thoughtful feedback. Below we address the concerns and questions brought up within the review.

Runtime and Computational Cost

We provide a detailed comparison of the computational time for each method in the table below:

CLVM-Linear denotes a contrastive latent-variable model with a linear transformation from observation to latents, and CLVM-VAE replaces the linear model with a neural network encoder / decoder. All timing was done on NVIDIA A100 GPUs. As the reviewer notes, our method is more computationally expensive than the baselines, almost entirely due to the cost of sampling from the diffusion model (since sampling must be done both during the EM loop and at inference time). However, this computational cost comes with significant improvements in performance. Below is the comparison of our model to the baselines across several experiments and metrics:

PSNR is the peak signal-to-noise ratio, PQM is the pqmass metric introduced in [4], FID is the FID score [5] using our MNIST classifier, and SD is the sinkhorn divergence. Note that we’ve replaced the CVAE [6] baseline with CLVM-VAE [2] since the latter is a superset of the former and achieves better performance. The substantive gap between our method and the baselines justifies the additional computational resources required to train and sample the diffusion model.

In terms of the feasibility for large datasets, the galaxy images experiment serves as a good baseline. The HST images we use are much higher resolution that ground-based observations [7,8] and equivalent in resolution to other current space missions [9]. Training the random and galaxy images took 48 GPU hours, and sampling from the posterior of an individual observation takes 1.5 seconds. For the ~80k galaxy images in the catalog, this amounts to under 34 GPU hours to draw one posterior sample for each galaxy. This compute budget is realistic for an academic lab.

We have added a section to the appendix that includes the full timing comparison above and point to that section in our discussion and limitations section and in the results for each experiment.

Performance as a Function of Dataset Size

We agree with the reviewer that, by using diffusion models to fit our prior distributions, our method will be sensitive to dataset size and likely perform poorly in the low data regime. Some work has explored solutions to reduce dataset size requirements, including data augmentation [11], patch-wise diffusion training [12], latent-space diffusion models [10, 13], and transfer learning [14]. It would be possible to incorporate these methods during our maximization step.

To better quantify our dependence on the dataset size, we ran a set of ablation studies on our performance on the full-resolution Grassy MNIST experiment. We train our method using 1/4th, 1/16th, and 1/64th of the original grass and grass+MNIST datasets. The 1/16th and 1/64th runs are given 1/4th as many EM laps as the original training to account for the smaller dataset, but all other hyperparameters are unchanged:

As the reviewer suggested, there is clear degradation in performance as the grass and MNIST datasets are reduced in size. However, even with 1/16th of the original dataset (2048 grass images, 512 MNIST digits) our method still outperforms the baselines run on the full dataset. We have added the dataset scaling study to the appendix and referenced it in the Discussion and Limitations section. We have also added a brief discussion of strategies that can be used to improve the performance of diffusion models in the low-data regime.

Linear Mixing and Gaussian Noise Assumptions

The reviewer is correct that our current method assumes both a linear observation model and Gaussian noise. As it stands, our framework cannot easily be extended to the non-linear case without substituting our joint sampling for Gibbs sampling. It is likely possible to generalize our method further under a reasonable set of assumptions, but we reserve that for future work.

The Gaussian noise assumption can be relaxed within our framework. It would require access to “noise-only” views or the ability to sample from a known noise distribution. One can consider the noise as an additional source and pre-train its corresponding diffusion model with the noise samples. From there, the rest of our machinery operates identically to the experiments presented in the paper, with the weights of the noise diffusion model frozen. However, the linear mixing assumption must be preserved. Noise that depends on the signal (for example Poisson noise from photon counts) cannot be modeled in this way. We leave the extension to non-linear noise models for future work.

Additional notes

• For the posterior metrics, we are comparing posterior samples to the original signal, meaning that they are not independent draws. Therefore, it is possible to get large positive pqmass values and large negative sinkhorn divergence values. The posterior metrics are meant to compare the reconstruction produced by different methods, and should not be interpreted in absolute terms. We have added a sentence to the main body of the paper to clarify this.
• To ensure fair comparison, we ran a hyperparameter sweep for the baseline methods. Similarly, we made small architectural changes to improve the diffusion models used in our method and reran all experiments. Therefore, some of the metrics quoted in the table changed from our original submission.

[1] Li, Didong, et al. "Probabilistic contrastive principal component analysis” (2020) [2] Severson, Kristen A., et al. “Unsupervised learning with contrastive latent variable models” (2019) [3] Heurtel-Depeiges, David, et al. "Listening to the noise: Blind denoising with Gibbs diffusion" (2024) [4] Lemos, Pablo, et al. "Pqmass: Probabilistic assessment of the quality of generative models using probability mass estimation" (2024) [5] Heusel, Martin, et al. "GANs trained by a two time-scale update rule converge to a local Nash equilibrium" (2017) [6] Abid, Abubakar, et al. “Contrastive Variational Autoencoder Enhances Salient Features” (2019) [7] Koekemoer, Anton M., et al. "The COSMOS survey: Hubble space telescope advanced camera for surveys observations and data processing" (2007) [8] Ivezić, Željko, et al. "LSST: from science drivers to reference design and anticipated data products" (2019) [9] Gardner, Jonathan P., et al. “The James Webb Space Telescope Mission” (2023). [10] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models" (2022) [11] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models" (2022) [12] Wang, Zhendong, et al. "Patch diffusion: Faster and more data-efficient training of diffusion models" (2023) [13] Zhang, Zhaoyu, et al. “Training Diffusion-Based Generative Models with Limited Data” (2025) [14] Hur, Jiwan, et al. "Expanding expressiveness of diffusion models with limited data via self-distillation based fine-tuning" (2024)