GUD: Generation with Unified Diffusion

We thank the referee for the helpful comments. First we would like to apologize for the late response. Upon the suggestion of the reviewer, we were trying out larger datasets, and trying to get as far as possible with the new numerical results before we respond. In what follows we will address the questions.

Limited Experimental Scope: The experiments are primarily conducted on CIFAR-10, a relatively small dataset. Datasets with larger images would better validate the findings and demonstrate the framework’s effectiveness in diverse settings.

First, we will update the results for CIFAR-10 with significantly improved performance. We would also like to point to the fact that our manuscript contains reports on a second dataset, namely the PCAM dataset. Following your suggestion, we are currently training our setup for the 256 sized CelebA-HQ dataset using the FFT basis and scanning over varying levels of softness. We hope to wrap up the new numerical experiments before the deadline for submitting an updated version. We believe the new experiment will further demonstrate the framework’s applicability in diverse settings.

Insufficient Explanation of Unified Diffusion and Autoregressive Generation: The explanation on how the framework unifies standard diffusion and autoregressive generation lacks clarity. Providing a specific example of the component-wise noise schedule would enhance understanding and illustrate this unification more concretely.

We thank the reviewer for the suggestion. Following your suggestion, we will edit our discussion in section 5.2 of the column-wise schedule as a specific example of the component-wise noise schedule, which illustrates the framework and in particular the meaning of the autoregressive limit. In this example, the softer the schedule, the more columns are noised together. In the opposite extreme, each column would be noised sequentially and independently, amounting to an autoregressive model in which a column is generated depending on the previous ones and in isolation. In our framework, analogous autoregressive generation can be carried out in any components (columns, Fourier modes, PCA components; any linear subspace).

I will first read the other reviewers' comments and authors' feedbacks to re-evaluate the score as I have not received any direct feedback to me.

We thank the reviewer for the suggestion to investigate larger sized image datasets.

The experiments are primarily conducted on CIFAR-10, a relatively small dataset. Datasets with larger images would better validate the findings and demonstrate the framework’s effectiveness in diverse settings.

As announced in a previous comment, we have explored the Celeb-A dataset and have trained a score network conditioned on $\gamma$ for a family of schedules. In the revised manuscript, we include in appendix E (highlighted in blue) training results and examples of generated samples. We primarily hope this demonstrates the applicability of our GUD framework to larger datasets, and in particular the feasibility to train $\gamma$-conditional networks. Due to time constraints, we have not finished tuning training and schedule hyperparameters, as well as architectural modifications, to allow for a rigorous numerical study of sample quality.

We thank the reviewer again for the helpful comments and suggestions, and for recognizing the contributions. We have uploaded a new version of the manuscript, in which we hope to have addressed the points you raised. Below are the main points.

• Following your suggestion, we added a sentence (in blue) in the first paragraph of section 5.2, and enhanced Figure 7 and the text explanation around it. We hope the improved explanation of this sequential example makes it clearer how the framework unifies standard diffusion and autoregressive generation.

Please let us know if you have further questions or would like to see further improvements.

We thank the reviewer for their comments and questions. First, we agree that the paper will benefit from a more extensive discussion on related work. We will expand this part of the paper in the revised version, and in particular include a discussion on [2] and combine it with the existing analyses on [1].

Regarding the question (and comments in weakness 1) the reviewer raised on quantifying the impact of different choices, we would first like to point out that our experiments show that the performance depends significantly on our schedule parameters and the optimum does not appear to coincide with standard diffusion. In particular, we refer to Figure 3 (which will be updated with strengthened results in the revised version) and Figure 5 of the manuscript. In Figure 3, the difference between i) the standard isotropic Gaussian (unwhitened) prior and the variance-matching Gaussian (whitened) prior and ii) different component-wise schedules with different levels of "autoregressiveness" is displayed quantitatively. Similarly in Figure 5, where this level as well as the ordering of components is varied.

Next, we would like to make one additional remark. The design space of GUD, particularly the component-wise schedule, is vast. Each chosen component-wise schedule might generally benefit from a different tuning of training hyperparameters, which complicates comparisons among them. We believe, in practice, our interpretable framework makes it possible for users to choose a suitable basis, schedule, and prior for their specific application, and subsequently tune the model with those choices and obtain a better-performing diffusion model in this way, as is indicated to be possible by the experiments mentioned above.

We thank the reviewer again for the helpful comments and suggestions. We have uploaded a new version of the manuscript, in which we hope to have addressed the points you raised. Below are the main points.

• We have expanded section 2 on related works, and in particular include a discussion on [2] and combined it with the existing analyses on [1].
• We have updated Figure 3, reflecting the results of new runs, which shows that variance-matching Gaussian prior and a more autoregressive noise schedule appears to perform better than standard diffusion as far as FID is concerned, at least given our noising-state-conditioned network architecture and choices of hyperparameters.

Please let us know if you have further questions or would like to see further improvements.

We thank the reviewer for the thoughtful questions and helpful suggestions.

Could the authors explain how to select in the large design space the various parameters/hyperparameters?

We agree that this is an important question and one we are currently investigating further. As discussed in this paper and confirmed in the numerical experiments, we believe that the ordering of noising of the different components, and how sharply (autoregressively) we noise these has a significant impact. In this regard, it is important to stress that the optimal choice will depend on the dataset. To demonstrate the impact, in this work we experimented with choices informed by well-understood properties of natural images such as their power spectrum and multi-scale nature, including the PCA and Fourier basis and sequentially noising the respective components from small variance to large variance.

Can the authors briefly position wrt the works like the ones cited in the weaknesses section?

We will discuss these two previous papers in the updated version. In short, the noising processes in these works can be regarded as special cases of the unifying framework we propose.

Minor: Figure 7 is qualitatively difficult to interpret from someone not specialized in the field. I suggest the authors to either add some extra comments or produce a similar image for a dataset which is more understandable for a generic reader.

We will add an additional figure and more detailed explanations in the updated version to better explain the procedure.

Architecturally, GUD’s design adds complexity by requiring modifications like cross-attention mechanisms for conditioning on component-wise noise states. These additions complicate the implementation and increase the risk of instability during training, as standard architectures like U-Nets are not inherently optimized for GUD’s intricate conditioning needs. This limitation might however be not so crucial.

We thank the reviewer for the thoughtful comment on the architecture complexity. In what follows we clarify the following three points.

• noising state conditioned network As GUD allows for choosing the SDE forward process that connects the data distribution (approximately) to the Gaussian prior by choosing a path for all components of the noising state vector $\gamma$ (as illustrated in Figure 1), it is indeed natural for us to consider a $\gamma$-conditioned network. In particular, conditioning on gamma instead of on time (for some given schedule) allows for one score network to (approximately) capture a range of different noising schedules, as we explain in the second paragraph of section 4.5. In the work we explored how an expressive $\gamma$-conditioned network may look like, in particular using cross attention. We have not encountered severe limitations in our experiments. In particular, we note that for the experiments in our paper it would in fact suffice to condition the network on just two or three scalar variables as opposed to the whole $\gamma$ vector, as in our experiments we limit ourselves to one or two dimensional families of schedules.

• for a given component-wise schedule We would like to emphasize that, once a (component-wise) schedule is chosen, any standard score network conditioned on time can be used and $\gamma$-conditioning will not be necessary to run a GUD model with a fixed schedule.

• supporting numerical experiments In part to demonstrate the above points, we have re-done the experiment of Figure 3 using a minimally modified version of Song et al's [2011.13456] NSCN++ architecture without cross attention. Having trained for more training steps than previously, we now obtain significantly better FID scores.

We thank the reviewer again for the helpful comments and suggestions. We have uploaded a new version of the manuscript, in which we hope to have addressed the points you raised. Below are the main points.

• We have expanded section 2 on related works and in particular commented on the work by Voleti et al and Das et al.
• We have added a second panel (the right panel) to Figure 7 and expanded and edited the text around it to explain the procedure better.
• In view of the comments on the architecture, we added a final paragraph to the Introduction and a final paragraph to section 4.5 (both in blue text), to clarify the points we wrote in the previous response.

Please let us know if you have further questions or would like to see further improvements.

We thank the referee for the thoughtful comments. In what follows we will address the questions.

On "Weakness 1":

It is not clear how scalable learning the covariance matrix for a full high dimensional data distribution would be or if it would even be beneficial.

We would like to clarify that our work does not rest on learning the covariance matrix. In many cases this can be estimated directly from the dataset, and where this is not feasible one could use a power-law approximation to the power spectrum. More generally, our model does not rely on the noise scale exactly matching the data distribution, and such a choice may indeed not be the optimal one.

The time-dependent diffusion scale term has not been investigated in much detail as far as I am aware and I believe this should be the main focus of the paper or at least more attention.

We agree that the component-wise schedule is the most important element of the GUD framework, and it has indeed been by far the main focus on our research effort. We will highlight this more clearly in the paper. The reason why we included the prior and the basis in the discussion of the design choices of the GUD framework is that they are inseparable and the schedule cannot be described without specifying the other two: the diagonal basis dictates what the "components" are in the component-wise schedule, and the prior dictates the endpoint of the distribution path that the component-wise schedule leads to. We do not see the inclusion of the other two choices as weakness. We also thank the referee to remind us of the two papers which used a "non-standard" basis or prior. We will include a discussion on them in the manuscript.

On "Weakness 2": We agree with the referee that the FID reported in Figure 3 is high. As mentioned in the footnote, we abruptly ran out of compute available to us. We have in the meantime resumed training with a somewhat modified network and achieved a much lower FID as a result. This part of the manuscript will be updated accordingly.

On "Weakness 3":

It is not clear to me the theoretical soundness of using the autoregressive approach for extending existing images i.e. changing dimension from previously trained model. It seems the generative process is no longer related to the time reversal of an SDE given the dimension changes. Can this be formalised?

One does not need to change the dimension from previously trained model in the application of image extension discussed in the paper. In particular, the generative process still corresponds to the reverse time SDE, as we described in section C of the appendix. Instead, though the image gets larger in the process, the "active" dimension that participate in the generative process at any given diffusion time remains the same, while the already generated portion of the image remains frozen. In the updated manuscript, we will add additional figure and extra comments to explain this procedure better.

In summary, the soft autoregressive forward process only noises a subset of columns at any time. The score network thus only has to output the score for this relevant region, outside of which the SDE is either frozen or exactly generates Gaussian noise. Assuming the distribution only depends on a finite region (neighboring columns), we can repeat the generative process by extending the image with noise and sliding the input to the score network over the generated image. More formally we can think of the whole large image initialized with white noise, across we then slide the score network to compute the non-trivial part of the SDE.

Finally we thank the reviewer for pointing out reference [1]. We will include the citation in the manuscript in the part where relevant work is discussed.

In summary:

We believe that our main contributions are:

1. To explain and interpret the linear SDE ansatz in terms of the three degrees of freedom, unifying many specific choices discussed in previous works.
2. To give theoretical heuristic arguments for component-wise schedules in particular with regards to the discussion of soft autoregressiveness.
3. To give experimental evidence that such choices have a significant impact on performance.

We will expand on the discussions on related literature including the 3 papers the referee cited, and edit the manuscript to highlight the component-wise schedule as the referee suggested. Finally, we will provide an updated version of experimental results.

Please let us know if there is anything else we could further clarify or strengthen.

The time-dependent diffusion scale term has not been investigated in much detail as far as I am aware and I believe this should be the main focus of the paper or at least more attention. What are the benefits of this compared to cascading diffusion, can cascading be seen as a case of this?

We thank the reviewer for recognizing the novelty of this aspect of the our work and for giving us the opportunity to further explain the difference between our approach and cascading diffusion.

In cascaded diffusion models, images are upscaled by training conditional diffusion models that generate higher resolution images with a low resolution image as input. The complete high-resolution image is generated from a noise sample. In our framework, one can simply make the corresponding choice for the basis and component-wise schedule to achieve a similar effect, albeit in a fundamental different way as we will now explain. If the components are chosen to represent coarse to fine grained linear subspaces of the image (e.g. Fourier or wavelet components), then the image can be generated from low to high resolution with any desired (soft) autoregressive ordering in our model. The coarse degrees of freedom are frozen after they have been fully denoised by virtue of the schedule, and are generally not provided as a separate conditional input. As a result, the coarse degrees of freedom in the generated high resolution image exactly match the low resolution image genearated earlier in the generation process in our model. This is to be constrasted with the cascaded diffusion model, in which the coarse degrees of freedom in the high resolution image only implicitly match the low resolution image that one conditions, through the training target. Finally, we note that coarse to fine is only one class of generative processes captured by our framework, and the continuous choice of softness (i.e. to what degree generation of degrees of coarseness overlap) is manifest in our framework but not in cascaded diffusion models. This example highlight the flexible and unifying nature of our framework.

We thank the reviewer again for the helpful comments and suggestions. We have uploaded a new version of the manuscript, which we hope to have improved based on the points you raised. Below are the main points.

• On page 1, with the text in blue we highlight that the component-wise schedule is the most important element of the GUD framework.

• We have expanded the section 2 on related work. We have in particular expanded the discussions on blurring diffusion, and included references to [Generative Modelling with Inverse Heat Dissipation] and [Improved Techniques for Training Score-Based Generative Models].

• We have updated Figure 3 (page 8) to reflect the lower FID after further training. In relation to the comments on available codes, we also amended section 4.5 with text in blue to report on further experimentations in this direction. To clarify the nature and the main purpose of our numerical experiments, we added a small paragraph at the end of the introduction. In particular, we expect that given a fixed schedule one can use a simpler network and fine-tune it to improve the resulting performance. However, this would not serve the primary purpose of the experiments of this paper.

• Given the misunderstanding ("weakness 3") of the image extension experiments, we have further explained what we did with edited text in blue and by adding a new Figure (the right panel of Figure 7), on p.9-10.

Please let us know if you have further questions or would like to see further improvements.

We thank the reviewer for sharing the thoughts on this. Just for clarification 1. in figure 3, the most optimal point has FID 5.26. 2. We showed PCAM dataset for column-wise extension. 3. We added an appendix on celebA experiments, again focussing on the comparison between schedules and not on the absolute performance.

We would like to emphasize again that we consider our main contribution to be to offer a unifying framework in which many optimizations can take place that aree bound to lead to improvement, given a specific task. In this spirit, our experiments focus on the comparative study between different choices: to show that the choices make a difference in performance, and that the standard diffusion does not seem to be optimal (which is hardly surprising), and to corroborate the advantage of incorporating more autoregressiveness. Given choices, hyperparameters than need to be tuned individually to achieve better perforrmances.

We hope this clarifies why we think our paper makes valuable contribution to the study of diffusion generative models.