Paper: On gauge freedom, conservativity and intrinsic dimensionality estimation in diffusion models (ICLR-2024)

Reviewer comment (Theory):
The main finding — that gauge invariance as in (11) is necessary and sufficient for the time marginals of the learned process to be identical to those of the true process — is not formally stated as a mathematical result, and the paper would benefit from a clear and precise mathematical statement of this main finding.

Author response:
Indeed, we agree that a formal statement would be appropriate. Please refer to our official comment: https://openreview.net/forum?id=92KV9xAMhF&noteId=0bECf5o20B.

Reviewer comment (Theory):
Theorem 1 on the intrinsic dimension estimation problem states that conservativity implies their intrinsic dimension estimation procedure is consistent (under their assumptions), but their gauge invariance finding (equation (11)) states that the law of the reverse time process is unaffected under this gauge invariance, so it would be more natural to expect that the intrinsic dimension estimation procedure is consistent under gauge invariance, making conservativity unnecessary.

Author response:
We agree that the relationship between the empirics and the gauge freedom requires some further thought, especially if one solely understands them separately as a way to estimate the intrinsic dimensionality (which is a difficult task on its own) and as a principle to understand under which conditions a diffusion model is exact, respectively. In principle, a remainder term fulfilling the gauge freedom conditions must also lead to the exact intrinsic dimensionality - we agree with the reviewer on that.
However, from a practical point of view, fulfilling the gauge freedom conditions is challenging, and the nature of the diffusion model $s_{\theta}$ will affect the trajectories approaching the data manifold. This is exactly what we observed in the experiments: even though a conservative $s_{\theta}$  does not match $s$ exactly, the trajectories approaching the data manifold are sufficiently well-behaved to allow the estimation of the intrinsic dimensionality. This observation and the new corollary (a) (which states that the gauge freedom term $r$ is $0$ if and only if the vector field is conservative) suggest enforcing conservativity. This establishes a new link between theorem 1 and the implications from section 3.

Reviewer comment (Evaluation):
In the discussion about Figure 3, the authors claim that their empirics demonstrate that conservativity leads to accuracy of the intrinsic dimension estimation procedure, while models that aren't constrained to be conservative have inaccurate intrinsic dimension estimation; this finding seems odd given the main observation of gauge invariance, which would suggest enforcing conservativity is not important.

Author response:
We agree that the relationship between the empirics and the gauge freedom requires some further thought, especially if one solely understands them separately as a way to estimate the intrinsic dimensionality (which is a difficult task on its own) and as a principle to understand under which conditions a diffusion model is exact, respectively. In principle, a remainder term fulfilling the gauge freedom conditions must also lead to the exact intrinsic dimensionality - we agree with the reviewer on that.
However, from a practical point of view, fulfilling the gauge freedom conditions is challenging, and the nature of the diffusion model $s_{\theta}$ will affect the trajectories approaching the data manifold. This is exactly what we observed in the experiments: even though a conservative $s_{\theta}$  does not match $s$ exactly, the trajectories approaching the data manifold are sufficiently well-behaved to allow the estimation of the intrinsic dimensionality. This observation and the new corollary (a) (which states that the gauge freedom term $r$ is $0$ if and only if the vector field is conservative) suggest enforcing conservativity. This establishes a new link between theorem 1 and the implications from section 3.

Reviewer comment (Evaluation):
While there is no formal contradiction between these findings (since models not enforced to be conservative may not satisfy gauge invariance), the result appears to somewhat contradict the essence of the paper.

Author response:
We agree that the relationship between the empirics and the gauge freedom requires some further thought, especially if one solely understands them separately as a way to estimate the intrinsic dimensionality (which is a difficult task on its own) and as a principle to understand under which conditions a diffusion model is exact, respectively. In principle, a remainder term fulfilling the gauge freedom conditions must also lead to the exact intrinsic dimensionality - we agree with the reviewer on that.
However, from a practical point of view, fulfilling the gauge freedom conditions is challenging, and the nature of the diffusion model $s_{\theta}$ will affect the trajectories approaching the data manifold. This is exactly what we observed in the experiments: even though a conservative $s_{\theta}$  does not match $s$ exactly, the trajectories approaching the data manifold are sufficiently well-behaved to allow the estimation of the intrinsic dimensionality. This observation and the new corollary (a) (which states that the gauge freedom term $r$ is $0$ if and only if the vector field is conservative) suggest enforcing conservativity. This establishes a new link between theorem 1 and the implications from section 3.

Reviewer comment (Theory):
For Figure 3, it is unclear why the true singular values must be 1, as they should depend on the local geometric properties of the manifold, the true data density, and the map.

Author response:
The reviewer correctly states that the true singular values are not generally 1 but depend on the local geometric properties, the true data density, and the map. However, in this case, the manifold is a 2-D plain, the true data density a Gaussian, and the map a simple padding. Therefore, the singular values are indeed 1.

Reviewer comment (Presentation):
For Figure 3, the colors are difficult to make out; fewer lines should be used (e.g., removing 2 of the 3 non-material singular values), and more visual differentiation (e.g., different line types) should be applied.

Author response:
We agree with the unfortunate choice of colors. We apologize for this inconvenience and will accommodate it in the camera-ready version.

Reviewer comment (Theory):
For Figure 3, as it is constructed (and given the difficulty of making out details), it appears there are more than 2 singular values (green, black, yellow) converging to 1, which seems to contradict Theorem 1 and the discussion in the first paragraph after "Intrinsic dimensionality estimation using diffusion models" where it is stated "2 of them converge to 1, whereas the remaining 3 diverge."

Author response:
Indeed, this misconception is due to our unfortunate choice of colors and line types. Only two converge to 1. We apologize for this inconvenience and will update the figure following the reviewer's suggestions.

Reviewer comment (Theory):
Do the authors have an intuitive interpretation of the conditions in 12?

Author response:
Indeed, it is an interesting observation and provides intuition for the old equations (12a) and (12b).

Reviewer comment (Theory):
How can one practically realise the constraints of equation 12? Imposing the orthogonality constraint in combination with the divergence free part seems as difficult as just estimating the exact score.

Author response:
Because of the very insightful comments of reviewer h3ie, we will drop the individual conditions (see old equations (12a) and (12b)) while only focusing on the gauge freedom condition (see eq. (10)). Indeed, the orthogonality perspective provides a very intuitive interpretation of the gauge freedom condition.

Reviewer comment (Theory):
If one takes no drift (i.e. the ode does not move the distribution), then it becomes clear that the constraints force the dynamics to be such that 1) particles are constrained move along the same level-set of the distribution (i.e. through time a particle remains at \(p_t(x)=c\) for all \(t\)), the orthogonal constraint, as long as 2) the dynamics do not alter the density of those level sets, the divergence free constraint. Is this perspective interesting when combined with a non-zero drift?

Author response:
Indeed, it is an interesting observation and provides intuition for the old equations (12a) and (12b).

Reviewer comment (Theory):
I do not understand the point of the first part of section 4. I understand the argument that given a path \(x_t\), if one adds a divergence free term to the score preserved density estimation. However, why would one assume they have a path? Without the exact score, or one satisfying the additional orthogonality constraint, how would one be able to generate a correct path from the model?

Author response:
Indeed the first part of Section 4 assumes a correct path is given, which cannot be tested without knowing the exact score. This observation should be understood as a motivation for remark 1.2 where we explain the difference between a gauge freedom for the instantaneous change of variable formula and a gauge freedom for the Fokker-Planck equation. We think it is not apparent why only the latter yields correct sampling and density estimation, and therefore, we think the reader will benefit from this observation.

Reviewer comment (Experiments):
Section 4.1, it would be highly instructive of the authors to plot examples of the trajectories with and without the additional \(r\) term I think. From my own plots, one ends up with a series of trajectories that at each instance are "spinning" around the mean, following the shape of the covariance ellipses. One can multiply the \(r\) term by a constant to speed up said spin. This related to my previous point also. Even if one initialises the ODE with the reference measure, the particles can still move as they can follow the covariance ellipses at constant speed on each ellipse.

Author response:
This is a great idea and would provide additional intuition. We thank the reviewer for this suggestion.

Reviewer comment (Experiments):
The experiments of section 5 do not seem to be related to theorem 1. Theorem 1 says that you can correctly estimate the data dimensionality if the extra \(r\) field satisfies the gauge freedom constraints. However the experiments only use a fully unconstrained score, and a score constrained to be conservative. There is not experiment showing that one can estimate the dimension correctly if the vector field is non-conservative, but does satisfy the gauge constraints.

Author response:
We don't agree with the reviewer on that point. Theorem 1 does not say that the data dimensionality can be estimated if the extra field satisfies the gauge freedom constraints. Indeed, the gauge freedom is simply an assumption. We show in Theorem 1 that if in addition, the score is conservative, then the intrinsic dimensionality can be estimated. Thus, conservativity is a sufficient condition (given all the other assumptions from the Theorem). Therefore, it makes sense to compare a conservative vs. a non-conservative vector field using Theorem 1 to estimate the intrinsic dimensionality.

Reviewer comment (Theory):
I did not find equations (12a) and (12b) to be helpful, but rather confusing. Rather than the Helmholtz decomposition, the right decomposition is that any vector field $s \in L^2(p_t)$ can be decomposed uniquely as $s = \nabla \phi + r$ where $\phi$ is a scalar field and $r$ satisfies the gauge freedom equation (10) (this decomposition then depends on $p_t$). Furthermore, this decomposition is orthogonal, in the sense that $\mathbb E[||s(x_t)||^2] = \mathbb E[||\nabla \phi(x_t)||^2 + ||r(x_t)||^2]$. This is readily checked with the integration by parts identity $\mathbb E[\nabla \phi(x_t) \cdot r(x_t)] = -\mathbb E[\phi(x_t) \left( \nabla \cdot r(x_t) + \nabla \log p_t(x_t) \cdot r(x_t) \right)]$ for any vector field $r$ and scalar field $\phi$, which shows that conservative fields $\nabla\phi$ are orthogonal in $L^2(p_t)$ to vector fields $r$ which satisfy (10). More abstractly, the adjoint of the gradient operator $\nabla$ is the "gauge freedom" operator $-(\nabla + \nabla \log p_t) \cdot$, showing that the image of the former is the orthogonal complement of the kernel of the latter.

Author response:
We are very thankful for the reviewer's insightful comments. In particular, we fully agree that the orthogonality perspective makes the message much more straightforward. We will, therefore, update the derivation of the gauge freedom conditions and present them as the condition required to obtain an orthogonal decomposition of the vector field; see our general comment: https://openreview.net/forum?id=92KV9xAMhF&noteId=0bECf5o20B. As pointed out by the reviewer, the consequence is that gauge freedom conditions and the conservativity condition imply that $r = 0$. We will update the text and Figure 3 accordingly.

Reviewer comment (Theory):
Because of the orthogonality property in the decomposition $s = \nabla \phi + r$, the score-matching loss can be decomposed $\mathbb E[|| \nabla \log p_t(x_t) - s_\theta(x_t) ||^2] = \mathbb E[|| \nabla \log p_t(x_t) - \nabla \phi(x_t) ||^2 + ||r(x_t)||^2]$. This shows that (denoising) score matching training naturally regularizes the $r$ part of the score to have a small norm (but it cannot be zero, see above). We also see that $s_\theta$ is conservative if and only if $r=0$.

Author response:
We are very thankful for the reviewer's insightful comments. In particular, we fully agree that the orthogonality perspective makes the message much more straightforward. We will, therefore, update the derivation of the gauge freedom conditions and present them as the condition required to obtain an orthogonal decomposition of the vector field; see our general comment: https://openreview.net/forum?id=92KV9xAMhF&noteId=0bECf5o20B. As pointed out by the reviewer, the consequence is that gauge freedom conditions and the conservativity condition imply that $r = 0$. We will update the text and Figure 3 accordingly.

Reviewer comment (Theory):
The decomposition $s = \nabla \phi + r$ simplifies the proof of Theorem 1 and removes an unnecessary assumption. Indeed, the assumption that $s_\theta$ is conservative is equivalent to $r = 0$ (at all times!). There is thus no need to assume that the divergence of $r$ is finite.

Author response:
We are very thankful for the reviewer's insightful comments. In particular, we fully agree that the orthogonality perspective makes the message much more straightforward. We will, therefore, update the derivation of the gauge freedom conditions and present them as the condition required to obtain an orthogonal decomposition of the vector field; see our general comment: https://openreview.net/forum?id=92KV9xAMhF&noteId=0bECf5o20B. As pointed out by the reviewer, the consequence is that gauge freedom conditions and the conservativity condition imply that $r = 0$. We will update the text and Figure 3 accordingly.

Reviewer comment (Experiments):
The main weakness of the paper is that it cannot yield real world practical results. More specifically, in order to test or enforce the gauge freedom condition (eq 11 or conditions 12a and 12b), one needs to have access to the true score, which is not generally possible. Diffusion models are mostly used for cases where the true score is not known, so it is not clear how knowing this condition would solve the question of conservativeness. The empirical results given in the paper are all toy examples for this very reason.

Author response:
We thank the reviewer for the constructive review. It seems that the reviewer's concerns are twofold. First, the practical implementations of the gauge freedom findings as it is necessary to have access to the true score. Second, the example in section 4.1 is trivial in 2D as the gauge freedom condition implies zero divergence in 2D.

Regarding the first concern, we refer to our general answer: https://openreview.net/forum?id=92KV9xAMhF&noteId=0JS67DC6Z8.

Reviewer comment (Theory):
In section 4.1, a two-dimensional example is discussed to show that conservativeness of score is neither necessary nor sufficient. The true score in this example has a divergence of zero, which makes it quite trivial. It is unclear if this is an inevitable consequence of conditions 12a and 12b in two dimensions. In other words, doesn't satisfaction of gauge freedom conditions imply that the true score must have zero divergence in 2D?

Author response:
We thank the reviewer for the constructive review. It seems that the reviewer's concerns are twofold. First, the practical implementations of the gauge freedom findings as it is necessary to have access to the true score. Second, the example in section 4.1 is trivial in 2D as the gauge freedom condition implies zero divergence in 2D.

Regarding the second concern. We disagree with the reviewer's statement, "The true score in this example has a divergence of zero, which makes it quite trivial." Actually, the divergence term of the true score $s(\mathbf{x},t)$ is equal to $-\sigma_{1}^{-2} - \sigma_{2}^{-2}$, which is not $0$ as both $\sigma_{1}$ and $\sigma_{2}$ are assumed to be positive. Therefore, we think this example serves as a non-trivial counter-example for conservativity not being necessary for exact sampling and density estimation. In case this does not answer your question, please don't hesitate to clarify.