Paper: $\infty$-Diff: Infinite Resolution Diffusion with Subsampled Mollified States (ICLR-2024)

Reviewer comment (Novelty):
A major weakness of this paper is how the authors frame their contributions within the existing literature. For example, in the abstract the authors write "Unlike prior infinite-dimensional models, which use point-wise functions requiring latent compression, our method employs non-local integral operators to map between Hilbert spaces, allowing spatial context aggregation." Similarly, in Section 1, the authors write "We introduce a new Gaussian diffusion model defined in an infinite-dimensional state space that allows infinite resolution data to be generated (see Fig. 2)". However, numerous recent works have posed infinite-dimensional diffusion models which precisely use these integral operators and a similar theoretical framework [1, 2, 3, 4, 5, 6, 7].

Author response:
> **Numerous recent works have posed infinite-dimensional diffusion models… [The authors] do not appropriately frame their contributions**

The ICLR guidelines state that these works are considered contemporaneous, and we are not required to compare our paper to them. Specifically, it states that if a paper has not been published at a peer-reviewed venue by May 28th then it is considered contemporaneous. The papers we discuss as concurrent work either exist only on arXiv, or were published after this date (e.g. at NeurIPS 2023). Additionally, our paper was published on arXiv during the months the reviewer states as the “earliest of these works”, in fact, some of these works cite our paper as concurrent work. As such, we believe it is fair to describe our paper as a concurrent work.

Reviewer comment (Novelty):
The authors are clearly aware of this prior work (see Section 6) but do not appropriately frame their contributions. Many of these prior works develop the theory of infinite-dimensional models which is very closely related to the proposed theory, but the authors of this submission do not state this. The authors claim that these works are concurrent, but the earliest of these works appeared in December 2022 [4], February 2023 [1], and March 2023 [2, 3], more than 6 months prior to the ICLR submission deadline. In addition, the authors are missing a citation to Kerrigan et al. [4] who previously developed a theory for discrete-time diffusion models which is closely related to the theory the authors propose in Section 3, as well as references to several other infinite-dimensional diffusion models in continuous time [5, 6, 7].

Author response:
> **“The authors are missing a citation to Kerrigan et al. [4]” and “The contributions regarding the theory are over-stated”**

Thanks for pointing out this missing citation; we were unaware of it and did not intend to overstate our contributions. While our approach was developed entirely independently, with our paper published on arXiv early this year before the conference proceedings of this work, we believe it is important to properly credit prior work. As such, in the introduction, we mention that concurrent works also developed diffusion models in infinite-dimensions; at the start of Section 3 we state that these concurrent works exist, and that they delve deeper into the theory than our paper; and finally, we have expanded the discussion on the concurrent works in Section 6. Throughout the paper we have reduced the strength of our contribution claims and/or also mentioned concurrent work (including those pointed out by the reviewer), and in parts where details are limited compared to these concurrent works (such as on the mentioned Radon-Nikodym theorem and regularity assumptions) we have added references to appropriate papers and stated their contributions.

Nonetheless, in terms of the practical aspect, we believe our approach offers substantial contributions extending these works since it is a major challenge to efficiently scale such models to enable practical modelling of complex data. Kerrigan et al. evaluate their approach only on small 2D curve datasets, while the other concurrent works evaluate on simple datasets such as Gaussian mixtures or MNIST (with FID scores substantially worse than baseline finite-dimensional approaches). Our work instead focuses on scaling to substantially more complex, high resolution datasets, developing a practical architecture, subsampling coordinates to achieve speedup/memory gains, and achieving state-of-the-art results for infinite-dimensional models by a wide margin, which presents a very large challenge as highlighted by the other reviewers. As such, with emphasis on such different aspects, we believe our work and each of these concurrent works are very complementary.

Reviewer comment (Theory):
The theory in Section 3 is imprecise to the degree of incorrectness. The authors write "The Radon-Nikodym theorem tells us the density for a measure $\nu$ absolutely continuous with respect to a base measure $\mu$". This is not what the Radon-Nikodym theorem says -- the Radon-Nikodym theorem states the *existence* of a density given the absolute continuity. To actually compute this density, you require stronger results, such as the Cameron-Martin and Feldman-Hajek theorems. See [1, 4] for a discussion of these theorems in the context of diffusion models.

Author response:
> **“The authors are missing a citation to Kerrigan et al. [4]” and “The contributions regarding the theory are over-stated”**

Thanks for pointing out this missing citation; we were unaware of it and did not intend to overstate our contributions. While our approach was developed entirely independently, with our paper published on arXiv early this year before the conference proceedings of this work, we believe it is important to properly credit prior work. As such, in the introduction, we mention that concurrent works also developed diffusion models in infinite-dimensions; at the start of Section 3 we state that these concurrent works exist, and that they delve deeper into the theory than our paper; and finally, we have expanded the discussion on the concurrent works in Section 6. Throughout the paper we have reduced the strength of our contribution claims and/or also mentioned concurrent work (including those pointed out by the reviewer), and in parts where details are limited compared to these concurrent works (such as on the mentioned Radon-Nikodym theorem and regularity assumptions) we have added references to appropriate papers and stated their contributions.

Nonetheless, in terms of the practical aspect, we believe our approach offers substantial contributions extending these works since it is a major challenge to efficiently scale such models to enable practical modelling of complex data. Kerrigan et al. evaluate their approach only on small 2D curve datasets, while the other concurrent works evaluate on simple datasets such as Gaussian mixtures or MNIST (with FID scores substantially worse than baseline finite-dimensional approaches). Our work instead focuses on scaling to substantially more complex, high resolution datasets, developing a practical architecture, subsampling coordinates to achieve speedup/memory gains, and achieving state-of-the-art results for infinite-dimensional models by a wide margin, which presents a very large challenge as highlighted by the other reviewers. As such, with emphasis on such different aspects, we believe our work and each of these concurrent works are very complementary.

Reviewer comment (Theory):
The authors miss key regularity assumptions necessary on the Gaussian noise in order to obtain a well-posed model. See again [1, 4] for a discussion of these requirements.

Author response:
> **“The authors are missing a citation to Kerrigan et al. [4]” and “The contributions regarding the theory are over-stated”**

Thanks for pointing out this missing citation; we were unaware of it and did not intend to overstate our contributions. While our approach was developed entirely independently, with our paper published on arXiv early this year before the conference proceedings of this work, we believe it is important to properly credit prior work. As such, in the introduction, we mention that concurrent works also developed diffusion models in infinite-dimensions; at the start of Section 3 we state that these concurrent works exist, and that they delve deeper into the theory than our paper; and finally, we have expanded the discussion on the concurrent works in Section 6. Throughout the paper we have reduced the strength of our contribution claims and/or also mentioned concurrent work (including those pointed out by the reviewer), and in parts where details are limited compared to these concurrent works (such as on the mentioned Radon-Nikodym theorem and regularity assumptions) we have added references to appropriate papers and stated their contributions.

Nonetheless, in terms of the practical aspect, we believe our approach offers substantial contributions extending these works since it is a major challenge to efficiently scale such models to enable practical modelling of complex data. Kerrigan et al. evaluate their approach only on small 2D curve datasets, while the other concurrent works evaluate on simple datasets such as Gaussian mixtures or MNIST (with FID scores substantially worse than baseline finite-dimensional approaches). Our work instead focuses on scaling to substantially more complex, high resolution datasets, developing a practical architecture, subsampling coordinates to achieve speedup/memory gains, and achieving state-of-the-art results for infinite-dimensional models by a wide margin, which presents a very large challenge as highlighted by the other reviewers. As such, with emphasis on such different aspects, we believe our work and each of these concurrent works are very complementary.

Reviewer comment (Writing):
The discussion of the "diffusion autoencoder framework" in Section 5 could use a short description for those unfamiliar with the work, as it was not entirely clear to me how the model is trained based on the description provided alone.

Author response:
> **Discussion on diffusion autoencoder framework**

We have added more details on the diffusion autoencoder framework to this section to explain how they can be used to reduce stochasticity and allow training with smaller batch sizes. Specifically, that an encoding network is introduced (which in our case operates on sparsely sampled coordinates) to map data points to small latent vectors; our mollified diffusion process is then conditioned on these latent vectors meaning that at the high time steps when the input to the denoising network is mostly noise, the latent provides additional information that allows better estimation of the denoised data, reducing gradient stochasticity. Subsequently, a small diffusion model with MLP backbone can be quickly trained to model these latents to allow unconditional sampling.

Reviewer comment (Experiments):
The works based on similar ideas are mentioned only in the Discussion section, but the proposed model is not compare against them in an experimental way.

Author response:
Our approach is designed to address the problem of efficient training on high resolution datasets, which we achieve through the combination of our mollified diffusion framework and efficient sparse architecture design. In contrast, the architectures alone used in these concurrent works do not scale well, making training on these impractical (in terms of memory and compute), where these works only consider very simple toy datasets such as Gaussian mixtures or MNIST with FID scores substantially worse than baseline finite-dimensional approaches. These works are considered contemporaneous, which we are not required to compare against.

Reviewer comment (Evaluation):
The lack of information regarding the computational costs of both training and inference for the proposed model.

Author response:
We agree that more information on computational costs, beyond the details in Section 5 and Table 3, would be useful. Particularly when applied to different resolutions, to assess how our approach scales. Sampling enough images for FID calculations is particularly slow at the highest resolutions, so will add more experiments into the final manuscript. Specifically, we will make a table with various combinations (e.g. data resolution, number of sampled coordinates, and inner grid resolution) and show run-time and memory usage.

Reviewer comment (Experiments):
I couldn’t find the information which is actual upscaling procedure limit, which will be significant from the practitioner point of view.

Author response:
Theoretically, the potential maximum resolution is infinite (technically bounded only by computation constraints and floating point precision), since the nature of the model is continuous. While at higher resolutions than the training data, you can’t expect results that contain substantially greater detail than the training data, we do find in Figure 7 sampling at higher resolutions does achieve good FID scores. However, as observed in Figure 2, we find that samples converge in terms of details and as observed by reviewer RCWm, can contain artefacts. Although this is, to the best of our knowledge, the first class of infinite-dimensional diffusion model that can capture non-blurry results higher than the training distribution. In the final manuscript we will add FID scores for even higher resolutions and observe how FID is affected.

Reviewer comment (Novelty):
The paper lists the extension of diffusion models to infinite-dimensional spaces as one of its contributions, but there are already many works on this topic [1-5] (of which only [2-4] are briefly mentioned in the paper). The authors should discuss the relationship of their framework to this related work in greater detail (and they should be mentioned before the discussion). Also, the mathematical treatment of the extension to infinite dimensions in the paper is lacking compared to [1-5] and thus cannot be considered a contribution.

Author response:
In the updated paper we have expanded the discussion on these works in Section 6, expanding on each paper, including a discussion on [1], and how our paper fits in. We have also added references throughout the paper to indicate additional mathematical treatment in these works.

Reviewer comment (Theory):
While the paper aims to work at any resolution, the use of "mollification" (I suggest to simply call this blurring which is a better-known term in the NeurIPS community) on the data distribution effectively limits the maximum resolution of the generated images. A close inspection of Figure 2 also reveals artifacts on the images generated at resolution larger than the training resolution. On a related note, the authors never discuss the mechanisms by which the score network should be expected to generate to higher resolutions while never having observed high-resolution images. Could the authors expand on this point?

Author response:
By mollifying, we in effect suppress high-frequency details beyond a specific scale yielding smoother functions with characteristics resembling those from a finite-dimensional subspace. However, it importantly facilitates the recovery of high-frequency details under certain conditions, an advantage over information loss induced by coarse discretisation—as showcased in our super-resolution results (Fig. 8). In practice, we only need to use a very small variance blur, which corresponds to very high resolutions, far exceeding the memory of the subsampled coordinates (up to 8x tested in Table 2), yet still lying in Hilbert space to allow for arbitrary resolution sampling (Figs. 2 and 7) and super-resolution (Fig. 8).

Reviewer comment (Reproducibility):
The central contribution of the paper seems to be its architecture. While its components are described rather abstractly in Section 4, I do not understand what is being implemented exactly. The discretization of the integrals is never discussed in detail, as well as the grid on which the various intermediate activations are represented. When using a regular grid, what is the difference between the proposed architecture and a regular CNN such as UNet? As another important point, how discretization invariance is obtained when using images with different resolutions is not discussed. Does the number of points used in the finite-sum approximation of the integrals changes?

Author response:
At the sparse level integrals are Monte-Carlo approximated: Eqn. 16 becomes $x(c) \approx \frac{|N(c)|}{|\mathbf{c}|}\sum_{y\in\mathbf{c}} \kappa(c-y)v(y)$. Sampled coordinates do vary across training points, and the same coordinates are used through the forward pass. We found global operators worse and be less scalable than local operators so we interpolate points to/from a regular grid where a model such as a U-NO (Rahman et al. 2022) can be used for global context.

When on a regular grid, there are some key differences to regular UNets. Primarily that each layer has a different formulation than typical CNN blocks, see Equation 14. The design of these layers is crucial for generalisation; for example, in Table 2 we show that if the operator kernel is non-linear (the overall operator is still non-linear), it is unable to generalise to different numbers of coordinates/resolutions.

Reviewer comment (Evaluation):
Super-resolution and inpainting are rather unfaithful to the original images: downsampling or cropping the generated image does not yield back the input image.

Author response:
Obtaining better faithfulness when inpainting with diffusion models is a popular research area. Diffusion only in masked areas leads to edge artefacts so approximation approaches such as DDIM inversion (Song et al. 2021) and guidance (Ho et al. 2022b; which we use) were proposed, avoiding artefacts at the expense of perfect reconstruction. More recent work substantially improving this [1] can be directly applied to our method.

Reviewer comment (Writing):
What is the "pointwise" constraint that is mentioned in the abstract and conclusion? Perhaps related, could the authors elaborate what is meant by "coordinates are treated independently" in neural fields?

Author response:
Neural fields use an MLP to directly map coordinates to pixel values, e.g for $c \in [0,1]^2$, the corresponding pixel is $p=\text{MLP}(c)$ (Sec 2.2). Hence they are independent (pointwise) because the mapping function does not transform over multiple points unlike convolutions/transformers. We have clarified this in Section 2.2.

Reviewer comment (Writing):
In section 4.1, what does the notation $c \in {m \choose D}$ means?

Author response:
It is a subset of length $m$ from $D$, where $D=[0,1]^2$ is the set of all coordinates. In practice we uniformly sample m coordinates, but other methods exist (Sec 4.2). We have clarified this in the updated manuscript.

Reviewer comment (Writing):
In equation (13), I am confused by the use of both $x$ and $v_l$. I do not understand if the operator represents a function from $x$ to $s$ or $v_l$ to $v_{l+1}$.

Author response:
We use $\mathcal{F}$ to denote the entire neural operator, mapping from $\mathcal{X}$ to $\mathcal{S}$, and use $v$ to indicate the activations of the network. Therefore $x=v_0$ and $s=v_L$. We have clarified this in this section.

Reviewer comment (Writing):
I did not understand the second paragraph of section 5 on diffusion autoencoder framework. What do the authors mean by "using the first half of our architecture"?

Author response:
Following Preechakul et al., 2022, the encoder is the downsampling stage of our architecture with the upsampling and skip connections removed: looking at Figure 4, it is only the left half of the image. We have added more details on diffusion autoencoders to Section 5.

Reviewer comment (Reproducibility):
What do the authors mean by "each model being trained to optimize validation loss" in Appendix A? Is validation data used to train the networks on top of training data?

Author response:
We only train on the training data. We use a kept-out validation set to tune hyperparameters, in this case when to stop.

Reviewer comment (Writing):
The writing really needs to be polished. A lot of typos should be corrected, for example, Page 4 Section 3.2 Paragraph 1 Line 11 $\mu_\theta: \mathcal{H}, \mathbb{R} \rightarrow \mathcal{H}$ → $\mu_\theta: \mathcal{H} \times \mathbb{R} \rightarrow \mathcal{H}$; Page 5 Section 4 .1 Paragraph 2 Line 1-2 “respectfully” → “respectively”; etc.

Author response:
Thanks for spotting these typos, we have corrected them and revised similar cases in the updated paper.

Reviewer comment (Writing):
What does $x_t(x_0, \psi)$ mean in Equation (11)? What is the difference between $x_t(x_0, \psi)$ and $x_t$?

Author response:
$x_t(x_0, \xi)$ is a reparameterization of $x_t \sim q(x_t | x_0)$. That is, for a data point $x_0$ and Gaussian noise sample $\xi$, $x_t(x_0, \xi)=\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\xi$. We have added this definition to the updated paper.

Reviewer comment (Theory):
Could the authors provide any theoretical justification and guarantee on why mollification is necessary and suitable for diffusion models to analyze infinite resolution data?

Author response:
Gaussian white noise for finite-dimensional diffusion models is defined such that each dimension is independent. In infinite-dimensional Hilbert spaces, white noise defined similarly exhibits discontinuities that mean that it does not lie in this space. Specifically, this is because it does not satisfy the trace-class requirement, that is, that the covariance operator $C$ must satisfy $\int_{\mathcal{H}} \| x \|_{\mathcal{H}} d\mu(x) = \text{tr}(C) < \infty$ (Da Prato and Zabczyk, 2014). Mollification involves convolving white noise with a mollifier kernel  $k(s)>0$ that corresponds to a linear operator $T$, resulting in white noise that can be described by $\mathcal{N}(0, TT^*)$ (Higdon 2002). This alters the white noise to ensure that it lies within Hilbert space and satisfies the trace-class condition (Sec. 3.1). Without mollification, it is much more difficult to tune the noise hyperparameters, leading to the trained model being unable to generalise across different coordinate subsampling rates or different resolutions. We have clarified this in Section 3.2.

Reviewer comment (Writing):
Would the authors clarify what $v_0$ means above Equation (14)?

Author response:
$v_0$, which lies in Banach space, is the input to the deep neural operator-based network. Each subsequent $v_l$ are the activations of the network, which also lie in Banach space, with $v_l \mapsto v_{l+1}$ being a single neural operator layer. We have clarified this in the updated paper.