Discrete Spatial Diffusion: Intensity-Preserving Diffusion Modeling

We thank the reviewer for their careful assessment of our results.

4.1 [Theoretical results] We agree that the core derivation of DSD builds on the framework of Santos et al. (2023). That said, we submit that the paper includes theoretical contributions, albeit specific to DSD. These include the derivation of first-reaction rates (Section 3.3 and Appendix A), as well as the Green’s function (propagator) for the DSD process under periodic boundary conditions (Appendix C). These mathematical results are critical for enabling efficient and exact computation within the DSD framework.

4.2 [FID comparisons] We agree that this helps comparability. We have found the following preliminary FIDs:

We have also performed additional quantitative metrics for our models, as well as comparisons to existing techniques, please see points 1.2 and 2.2.

4.3 [Lengthy paragraphs] We appreciate the feedback. If accepted we will clearly expand the explanation of how DSD links to Fokker–Planck and heat equations using some of the additional space, and we will reorganize lengthy paragraphs for better clarity. We are also willing to include other detailed derivations and/or explanations that the reviewer requests in the appendices.

4.4 [Fixing Fig. 1(c)] We are unsure about the specific concern the reviewer is raising regarding Fig. 1c, and would appreciate clarification to better address this comment.

4.5 [Differences in generations of Fig. 4] Thank you for this helpful question. In response, we have added a new visualization that shows the distribution of the training data for the rock microstructure dataset. Specifically, we plot the sampling distribution of $\mu$, along with $\mu \pm 2\sigma$ and $\mu \pm 3\sigma$ to better illustrate the spread of the generated samples. This updated figure more highlights our model's ability to generate samples across the full range of the dataset’s variability, maintaining consistent structure even at the distributional tails. The updated figure can be found here.

We also contrast our approach with recent work, led by Professor Martin Blunt. Their model lacks explicit control over porosity, and so it produces out-of-distribution generated samples, with longer tails in porosity distribution (see their Figure 3E). Our DSD offers a clear advantage when precise field-measurement tolerances of multiple significant digits have to be met. We will revise the manuscript to briefly include this very recent work and clarify how our model addresses these limitations.

4.6 [Comparison with mass conditioning approaches] In order to respond to this point, we have evaluated a baseline discrete diffusion model conditioned on the target total intensity (or mass). The model performs reasonably well near the mean of the dataset, where the target intensity falls clearly within the training distribution. However, it tends to reach this point by producing small negative values in the final image—despite being trained on images with pixel values in the [0-1] range. When we threshold these negative values to zero and discretize the image to the [0, 255] range, the intensity error increases. This degradation becomes especially pronounced in the tails of the intensity distribution. In these regions, the Gaussian diffusion model often fails to honor the conditioning and produces non-compliant samples. In contrast, our method (DSD) guarantees intensity compliance across the full range, including the tails of the distribution. This difference is demonstrated in this Figure. Please also see response 1.4 where we compared to another approach.

4.7 [Summary] We have addressed the reviewer’s clarification to the best of our ability, added further quantitative comparisons, and better explanation of the significance for Fig. 4. We believe that addressing these comments substantially improves the manuscript.

We thank the reviewer for their constructive comments; we feel that the reviewer has understood the purpose of the work and offers concrete avenues for improvement, all of which we are able to address in our revision.

3.1 [computational complexity] First, let us remark on the computational requirements, per iteration, involved in DSD. Noising the images involves sampling using a random walk propagator that is precomputed–this is cheap enough to be amortized by dataloader workers on CPU, so training DSD is not any more expensive than training NCSN++. At test time with our current non-parallel code, about half the time is spent moving the particles around and half the time at network evaluation. The overall time for training is similar to other diffusion models. Sampling time depends on the noise schedule and CFL tolerance parameter. Good samples can presently be accomplished in <4 minutes per celebA sample per GPU without any special software engineering and A5000 hardware (this includes batching of images). In terms of cost breakdown, on celebA (the most challenging for DSD in this work), roughly equal amount of time is required for moving particles as performing NCSN++ evaluation; the extra computational cost for DSD is approximately a factor of 2, and in our inference code this (CPU-based) cost is presently hidden by running multiple inferences in parallel processes. We will revise the manuscript to describe this information.

3.2 [high resolution] Regarding scaling, our method scales effectively to larger resolutions, and as a demonstration, we have trained it on 1000×1000 Leopard sandstone images: realizations are shown in the attached repository as well as included in the revision. Leopard sandstones feature larger critical percolation diameters and more dispersed clay regions. Our model captures these morphological nuances, demonstrating both fine-scale and macro-scale structural details.

3.3 [scalability] The reviewer’s point that DSD may be difficult to scale to images which are both high resolution and high bit depth (high-bit-depth) images. As mentioned (3.1), sampling noised images is not expensive–and so this is a concern only for inference throughput. We do address 24-bit CelebA images and separately 1000x1000 rock images (3.2). We feel that DSD is applicable to a variety of datasets, and future innovations may make it possible to scale the approach further. It may only be a question of implementation that is solvable with a GPU kernel for particle motion sampling.

3.4 [quantifying performance] Regarding quantifying performance, we have primarily focused on domain-specific measures (see also 2.2) because these binary microstructure images differ from typical computer vision benchmarks. Nonetheless the reviewer’s point about comparability is well-taken, and so we now include FID metrics for MNIST and CelebA. It is important to take these in context, as although the performance of these models with respect to FID is not state-of-the-art, our intent is not to replace the state of the art, but to address the niche where conservation of intensity/particles is necessary. To that end, we also offer to demonstrate the ability for a conditional standard gaussian diffusion model to produce given overall intensities on MNIST (1.4). The conditioning works well within the typical range of intensities for the dataset, but falters very significantly as one approaches the tails of intensity. Even still, the conditioned model sometimes produces negative intensities in order to meet its objective for total intensity; this failure mode is not possible with DSD. All of this will also be described in our revision.

3.5 [standard benchmarks] To further address the reviewer’s comment that they would like to see more standard benchmarks: We would like to argue that the non-standard tasks we introduce (intensity constrained generation and intensity constrained inpainting) provide value to the field as well; while one goal for ML research is to advance the state of the art for identified important tasks, another important goal is to advance the scope of (well-motivated) tasks for which ML techniques are available.

3.6 [summary] Through these revisions and clarifications, we are confident that our revised version will address the reviewer’s concerns about quantitative comparability, computational costs, and scalability. While our models do not advance state-of-the art FID scores, DSD is reasonably affordable, solves a well-motivated novel scientific task (intensity constrained generative modeling) far better than simple extensions of existing work (input-based conditioned models, and c.f. 1.4 regarding another approach), and produces very high quality domain metrics (2.2, 4.2, 4.6)

We thank the reviewer for providing a clear review that shows understanding of the work, constructive criticism for improvement, and good questions for clarification.

2.1 [incorporate quantitative metrics such as FID] We emphasize that DSD primary advantage is for scientific applications, whereas the FID metric implicitly focuses on human-centric datasets (the Frechet distance is computed as the latent embedding of the Inception V3 model, which was pre-trained on the ImageNet dataset). That said, we have performed the computation of FID as requested, achieving results in rebuttal section 4.2. We do not have any rational or theoretical foundation that the FID metric would be meaningful when applied to scientific datasets, however, for completeness sake we will now provide FID to the scientific images as rendered in RGB in the paper. In addition to FID, please see also 1.4 regarding tests comparing Gaussian models to ours with regard to total intensity generation targets.

2.2 [scientific applications… supported by qualitative results] The reviewer contends that our quantitative results are limited. In addition to the FID metrics included in revision 2.1, we do already provide several quantitative measurements highly relevant for the scientific data. For one, porosity (phase volume fraction) is a key variable because permeability of porous media typically scales by the inverse cube of the porosity, and so quantifying our ability to match the total intensity is a relevant metric for scientific applications. Additionally, we present several quantitative scientific evaluations in appendices (due to space constraints). These come from earth science, materials science, and energy storage literature, and include two-point correlation functions and pore size distributions (Appendix F.1), as well as interface length, triple-phase boundary, and relative diffusivity (Appendix G.2). For context on the relevance of these metrics, please see the paper Gayon-Lombardo et al 2020. Thus we argue that the quantitative evaluation of our models has not been quite limited as the reviewer has claimed. Perhaps this was not emphasized enough in the main text, and we will revise it to emphasize the domain science metrics.

2.3 [motivation for loss function] Your question about the loss functions is a good one. We propose to add the following explanation in revision: We remark that Eq. (5) is a heuristic approach to match the predicted rate and the ground-truth one, which is a popular approach for building loss functions in machine learning (similar approaches include flow-matching and score-matching); the motivation for using this loss is simplicity and analogy to existing work. Eq. (7) is a more principled statistical approach, derived in Santos 2023 by designing a maximum likelihood loss using the analytical reverse transition rates. While less intuitive, it is easy to verify by taking ordinary derivatives that the loss is an absolute minimum when the predicted and true reverse rates are equal, the process of which also reveals that this Eq. (7) is essentially the integral of the mean absolute percentage error. We have tested both approaches, and did not observe a significant difference, which demonstrates the robustness of the DSD framework. We hope this clarifies the nature of the logarithmic loss function.

2.4 [SSIM-based scheduler] Indeed, our SSIM-based approach is a heuristic way to construct the time scheduler, which we have mentioned in the manuscript (lines 216-217). We share the sentiment that it would be nice to have a theoretically sound way to construct the scheduler for DSD, but we are not aware of an existing approach which fulfills this criteria, and it is not an innovation that we are able to offer at this time. In revision we will remark on the pursuit of a deeper mathematical support for the time sampling schedule as a good candidate for future work. Ultimately, the design of schedules (noise schedule, time schedule, learning rate schedule) is a grand challenge for which there is no existing framework with a solid and complete mathematical justification (as far as we know).

2.5 [summary] We appreciate the reviewers comments, and will be able to address them in revision; we will improve the quantitative comparison using FID score, we will clarify the loss function, and remark further on the heuristic nature of our scheduling and prospects for future work. We also gently argue that our present results are not only qualitative, as we have applied several significant domain-based quantitative metrics. Although we concede that the CelebA faces fall short of the state of the art (and now quantifyably so), we believe that by addressing the rest of the reviewer’s concerns we have substantially improved the manuscript.

We thank the reviewer for their constructive feedback and questions.

1.1 [Difficult to understand] We have tried our best to make it accessible and attempted to follow the conventions of the theoretical paper. We will be happy to provide clarifications. It would be helpful if the reviewer can be more specific regarding confusing terms, missing definitions or explanations.

1.2 [No comparison] Our main point of the paper is to impose the hard constraint, strictly constraining the particle numbers. As such, the most important objective is to generate configurations which are exactly conditioned on the total intensity. In this sense, the rest of the models are failing the objective. The commonly adopted metrics, such as the FID score, are difficult to perform for exact conditioning for the approximate methods, as there is only a small set of images in the test set that has exactly a given total intensity, where FID between two given sets requires a large sample number.

1.3 [Markov process] We used a consistent notation as defined in the main text (line 156 of the text): “We treat a digital image with discretized intensities $I_{x,y,c}$ at pixel $(x,y)$ in color channel $c$”. $p(I_t \vert I_{t-1})$ are conditional distributions and a common notation used in the diffusion literature (e.g. Fig. 2, equation (1) in https://rb.gy/3v28z7). We have also provided a clear description of the continuous-state Gaussian generative diffusion models (line 41-48, left column), as well as state-of-the-art discrete-state generative models based on discrete-state Markov chains (line 72-98, left column), highlighting the difference to the Gaussian generative diffusion models in line 83-88, left column. We specified what the discreteness is referring to in line 147, right column: “...Discrete Spatial Diffusion (DSD), noting the 'discreteness' refers to both the discretized intensity units and the discreteness of the spatial lattice... where the particles are allowed to reside."

1.4 [Alternative approach] Thank you. The guidance function in the provided paper makes sense for reinforcement-learning tasks - which the paper is about - where a cost function $J$ is naturally defined. For our generative tasks, it is not well-defined---in fact, it has to be learned either additionally as stated in our cited reference about a posteriori sampling, or it has to be built into the diffusion model and learned during training. This is because the "guidance function" is exact only when it is $\partial_{I_t} \log p(C\vert I_t)$ where $C$ is the condition. Here, $\log p(C\vert I_t)$ is not known. While one can argue that we could try to impose heuristic guidance functions, our additional numerical experiments - which imposes a MSE between the generated and target total intensity - suggest that this is not a valid approach. We provide the Jupyter Notebooks (/codes/Guided*.ipynb in https://rb.gy/x1fjvy) generating the statistics of the total intensity of the generated samples with various guided functions (under /figs/). It is difficult to explain the experimental details with the 5K-character limit of the rebuttal, but we hope the notebooks are clear and we look forward to future discussion with the reviewer. In addition, we have implemented a conditional diffusion model taking the target total intensity as an additional input of the NN during training (/codes/mass*.ipynb in https://rb.gy/x1fjvy). The results (see also 4.6) suggest that (1) the total intensity can be statistically conditioned within the intensities in the data distributions, but not exactly, (2) after rounding and clipping the samples generated by Gaussian diffusion to uint8, a systematic additional bias is introduced, indicating that the conditioner is introducing physically impossible negative intensities to enforce the constraints, and (3) outside the data distribution, the conditioner completely fails. In comparison, DSD always has the required total intensity.

1.5 [$r$] The transition rate is a standard term referring to the transition probability per time. Cf. the cited textbooks, Van Kampen and Gardiner, or the published discrete-state generative model such as https://rb.gy/d7vqid and https://rb.gy/z7vfjv).

1.6 [Timestep] The diffusion model generating the schematic diagram has 1K steps and the snapshots were taken every 200 steps.

1.7 [Mass] We apologize. The mass stands for total particle numbers. We will revise.

1.8 [Quant. performance] Please refer to 1.2, 1.4, and 4.2.

1.9 [Additional exp.] Please refer to 1.4.

1.10 [Summary] We hope to have adequately addressed the points of confusion in the manuscript. We have also performed additional experiments following the suggestion to improve the work. These experiments help to show the benefits of our approach, and we thank the reviewer for the suggestions.