Paper: Nearly $d$-Linear Convergence Bounds for Diffusion Models via Stochastic Localization (ICLR-2024)

Reviewer comment (Writing):
Previous to stating the main theorem, the main object of interest \( p_{t_N} \) has never been defined.

Author response:
We thank the reviewer for their positive and detailed comments on our work. Regarding the reviewer's remarks on the writing and notation, we acknowledge the points of confusion and apologize for the lack of clarity in the original version. We will take on board the suggestions regarding the definitions of $p_{t_N}$, $\mathbf{x}$, $\mathbf{x}_\ast$, and $\mathbf{a} (U_s)$ and update the notation in the camera-ready version. We will provide an updated manuiscript reflecting the changes below by tomorrow.

Reviewer comment (Writing):
The use of bold \( \mathbf{x} \) for realization when defining conditional objects is a little confusing when at the same time \( \mathbf{x}_* \) is a random variable. This is also in conflict with the use in equation 3 where \( \mathbf{x}_t \) is a random variable drawn from \( q_t \). It would be helpful if there was a consistent distinction between realization and random object throughout and if this were defined early in the manuscript.

Author response:
We thank the reviewer for their positive and detailed comments on our work. Regarding the reviewer's remarks on the writing and notation, we acknowledge the points of confusion and apologize for the lack of clarity in the original version. We will take on board the suggestions regarding the definitions of $p_{t_N}$, $\mathbf{x}$, $\mathbf{x}_\ast$, and $\mathbf{a} (U_s)$ and update the notation in the camera-ready version. We will provide an updated manuiscript reflecting the changes below by tomorrow.

Reviewer comment (Writing):
Following Lemma 1, it would not hurt to expound a little on the significance of the covariance matrix and why it will be useful in the sequel. It would be nice if some version of what is stated "inside" the proof of theorem 1, second paragraph, appears near this lemma.

Author response:
The covariance matrix $\mathbf{\Sigma_t}$ is significant since it can be related to the Jacobian of the score function (see Lemma 4). Thus, controlling $\mathbf{\Sigma_t}$ as in Lemma 1 allows us to control the Jacobian of the score function. In turn, this allows us to control time-discretization terms of the form $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ via Ito calculus (see Lemmas 2 and 3). We will include a discussion of this intuition for Lemma 1, coming from the proof of Theorem 1, after the statement of Lemma 1 in the camera-ready version.

Reviewer comment (Theory):
With regard to the second part of assumption 2—you mention that your analysis is not dependent on "it"—I have a clarifying question: would the result incur extra "condition numbers" or does it still work even if this covariance matrix has 0 or near 0 eigenvalues?

Author response:
No, we do not incur extra condition numbers, and the result still holds if the covariance matrix has zero or near zero eigenvalues. For an arbitrary covariance matrix, our bounds may occur a linear dependence on $M_2 = \mathbb{E}[||x_0||^2]$ instead of $d$, but are essentially very similar in spirit. For a more detailed analysis of what exactly must be changed in our proofs and the form of the resulting bound, see our response to Reviewer 73Lh's point number 3.

Reviewer comment (Theory):
Could you expand on the steps taken in equation (9)? I do not immediately see how this follows from the standard Ito Isometry relating squares of stochastic integrals and integrals of their squares.

Author response:
Of course; our additional exposition was probably too brief. First, writing $\textup{d} (e^t \nabla \log q_{T-t}(Y_t)) = \sqrt{2} e^t \nabla^2 \log q_{T-t}(Y_t) \cdot \textup{d} B_t'$ in integral form we get
\begin{equation*}
    e^t \nabla \log q_{T-t}(Y_t) - e^s \nabla \log q_{T-s}(Y_s) = \int_s^t  \sqrt{2} e^r \nabla^2 \log q_{T-r}(Y_r) \cdot \textup{d} B_r'
\end{equation*}
By the standard form of Ito's isometry applied to this integral, we get (under suitable integrability conditions)
\begin{align*}
    \mathbb{E}\Big[\Big|\Big| \int_s^t  \sqrt{2} e^r \nabla^2 \log q_{T-r}(Y_r) \cdot \textup{d} B_r' \Big|\Big|^2 \Big] & = \mathbb{E}\Big[\int_s^t (\sqrt{2} e^r \nabla^2 \log q_{T-r}(Y_r) \cdot \textup{d} B_r')^T(\sqrt{2} e^r \nabla^2 \log q_{T-r}(Y_r) \cdot \textup{d} B_r') \Big] \\\\
    & = \mathbb{E}\Big[\int_s^t 2 e^{2r} || \nabla^2 \log q_{T-r} (Y_r)||^2_F \textup{d} r \Big]
\end{align*}
Combining both expressions and exchanging the order of itnegration on the RHS, we have
\begin{equation*}
    \mathbb{E}[||e^t \nabla \log q_{T-t}(Y_t) - e^s \nabla \log q_{T-s}(Y_s)||^2] = 2 e^{2r} \int_s^t \mathbb{E}[|| \nabla^2 \log q_{T-r} (Y_r)||^2_F] \textup{d} r.
\end{equation*}
Equation (9) then follows by differentiating with respect to $t$.

Reviewer comment (Theory):
Following the statement of proposition 4, could you be a little more precise/remind us about which stochastic transformation is used in the invocation of the data-processing inequality? It seems like it is here the early stopping idea is used but I find it hard to follow.

Author response:
We are applying the data processing inequality with the transformation $f : (\omega_t)  \mapsto \omega_{t_N}$ which takes a path $(\omega_t) \in C([0,t_N], \mathbb{R})$ on $[0,t_N]$ and returns its endpoint. Under this transformation, $Q$ and $P^{\pi_d}$ are mapped to $q_\delta$ and $p_{t_N}$ respectively. (This is a similar argument to that used in Chen et al. (2023d).)

Reviewer comment (Writing):
I encourage the authors to add "nearly linear" in the title. This is a standard terminology (at least in theoretical computer science) for hiding log factors.

Author response:
We agree this is a sensible suggestion, and will update the title to "Nearly $d$-Linear Convergence Bounds for Diffusion Models via Stochastic Localization".

Reviewer comment (Theory):
One dumb question: doesn't the second part of Assumption 2 imply its first part? Or am I mis-interpreting the meaning of second moment of a distribution on R^d?

Author response:
Yes, the first part of Assumption 2 is technically superfluous. We chose to state it in this way to separate the assumption that the data distribution has finite second moments (which is absolutely required by our methods) and the assumption that $p_{\textup{data}}$ is isotropic (which can be weakened, see next comment).

Reviewer comment (Theory):
It is claimed that the isotropic condition in Assumption 2 is not essential. However, this has never been formally justified, if I'm not missing it. Which nuisance of the proof needs to be changed if the covariance is general? In fact, I have some doubt on the claim that this assumption is immaterial. What if the covariance is not strictly positive definite? What if the covariance is **unknown**?

Author response:
The main changes required are in the proofs of Lemmas 5 and 6 and Proposition 4. Suppose that instead of assuming that $\textup{Cov} (p_{\textup{data}}) = I_d$, we only assume that we have finite second moments and let $M_2 = \mathbb{E}[||x_0||^2]$. Then, Lemmas 2, 3 and 4 hold unchanged. Also, the proof of (13) can be adapted, replacing each time we use $ \mathbb{E}[||x_0||^2] = d$ with $ \mathbb{E}[||x_0||^2] = M_2$; we find that all instances of $M_2$ cancel in the final result and so Lemma 5 also holds unchanged. The only change in the proof of Lemma 6 comes at the top of page 18, where we find instead that $\mathbb{E}[\textup{Tr}(\mathbf{\Sigma_t})] \leq M_2$ for $t \in [1,T]$. We deduce that Lemma 6 holds almost unchanged, except the RHS of (17) should contain an additional term of $\kappa M_2$. The final change that must be made is in Proposition 4, where the bound becomes $\lesssim (d + M_2) e^{-2T}$ (this can be easily seen from the proof). Putting this all together, we end up with the final result that $\textup{KL}(q_{\delta}||p_{t_N}) \lesssim \varepsilon^2_{\textup{score}} + \kappa^2 d N + \kappa d T + \kappa M_2 + (d + M_2) e^{-2T}.$ This bound holds even in the case where the covariance is not strictly positive definite, and in the case that it is unknown in advance. Since we expect $M_2$ to scale linearly in $d$, we consider this result to be in essentially the same spirit as Theorem 1, with only minor adaptations required.

Reviewer comment (Writing):
In the line following equation (9), there seems to be some typos in the equations for verifying the applicability of Fubini.

Author response:
Thank you for drawing our attention to this line. In fact, in the first version of the manuscript, we made a mistake when giving the relevant inequalities which permit the application of Fubini. As can be seen from writing out the application of Ito's isometry in full (see our response to Reviewer wcsw below), the correct justification should be $\int_s^t e^{2r} \mathbb{E}[||\nabla^2 \log q_{T-r}(Y_r)||._F^2] \textup{d} r < \infty$ (which follows from (13), as stated above equation (9)). We will correct this in the camera-ready version.

Reviewer comment (Writing):
Please introduce the notation $\lesssim$ somewhere if it's not been done yet. I assume that this means the LHS is upper bounded by the RHS multiplied by a numerical constant.

Author response:
We will include a definition in the camera-ready version.

Reviewer comment (Theory):
It's argued informally in section 4 (with some details in appendix G) that the linear dependence of the convergence rate on d is necessary. This seems like a nice observation. Is it possible to turn it into a formal proposition and specify how this lower bound depends on other parameters besides d? In fact, this observation is a little surprising to me since for product distribution, the problem is effectively 1-dimensional however denoising diffusion still incurs a dimension-dependence. Could the authors offer some more intuition?

Author response:
Intuitively, we see the dimension dependence as arising unavoidably from the use of KL divergence to measure the distance between the true and approximate data distribution. Even if the problem factorizes completely, so that it can essentially be viewed as a product of independent 1-dimensional approximation problems, the fact that the KL divergence tensorizes across dimensions means that the error will still scale linearly in the dimension. We considered trying to formalize a lower bound on the iteration complexity in terms of $d$ and the other problem parameters. However, we came to the conclusion that it would be challenging to supply a lower bound that had realistic dependence on both the data dimension and the step size, since the two appear to somewhat trade off against each other for reasons discussed in Section 4. Getting more precise about exactly how the dependence on the dimension and the step size trade off would likely require a much more thorough theoretical investigation.

Reviewer comment (Theory):
I wish that there was more discussion of the key lemma. Why is stochastic localization needed to prove it? How can we interpret it?

Author response:
> I wish that there was more discussion of the key lemma. Why is stochastic localization needed to prove it? How can we interpret it?

The covariance matrix $\mathbf{\Sigma_t}$ is important since it can be related to the Jacobian of the score function (see Lemma 4). Lemma 1 thus allows us to control the Jacobian of the score function, which in turn allows us to control time-discretization terms of the form $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ via Ito calculus (see Lemmas 2 and 3).

Lemma 1 could be proved without reference to stochastic localization. The proofs of Propositions 1 and 2 boil down to standard applications of Ito calculus, and these could be rephrased entirely in the language of diffusion models. However, we choose to present Lemma 1 in the context of stochastic localization in order to provide due credit, since this was the context in which we discovered these ideas, and because we feel that the relationship to stochastic localization provides additional intuition and motivation for our proofs. We will include some more discussion of Lemma 1 after the statement in the camera-ready version (see also our reply to Reviewer wcsw).

Reviewer comment (Presentation):
I have some issues with the presentation of the main theorem in the paper. This is important since about a third of the paper is dedicated to explaining the proof of Theorem 1. I found their explanation to be hard to follow. First, I would appreciate a roadmap explaining precisely what steps are new/novel and what steps are guided by past work. In particular, it seems that the work on the discretization error is new, so the authors are then able to apply the bound from stochastic localization. However, is the perspective taken where they view the single process $(Y_t)$ as solutions to (2) and (4) under different measures new? The biggest issue, though, in my opinion, is that much of the proof is left to the appendix. I understand this is necessary, but without more interpretation in the paper to help the reader understand what is happening, it is easy to get lost. Improving this explanation would greatly help the paper since it is primarily theoretical.

Author response:
> I have some issues with the presentation of the main theorem in the paper. This is important since about a third of the paper is dedicated to explaining the proof of Theorem 1. I found their explanation to be hard to follow. First, I would appreciate a roadmap explaining precisely what steps are new/novel and what steps are guided by past work. In particular, it seems that the work on the discretization error is new, so the authors are then able to apply the bound from stochastic localization.

We apologize for the lack of clarity surrounding this overarching structure in the original manuscript. We will provide an updated manuscript by tomorrow with a clearer roadmap at the start of Section 3 explaining which sections of our work are novel. We will also slightly modify the structure of Section 3 to make it easier to get an overview of the proof at a glance, including adding in lemmas summarising the key results of the sections on the bound on the discretization error and the application of Girsanov's theorem.

With the updated structure, our proof comes in three key steps. We view Step 1 as our main novel contribution, while Steps 2 and 3 closely follow Chen et al. (2023a;d).

**Step 1:** We bound the error from discretizing the reverse SDE. Previous works bound $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ using a Lipschitz assumption. We use a new Ito calculus argument to get a differential inequality for $E_{s,t}$ (Lemmas 2 and 3), relate the coefficients of this inequality to $\mathbf{m}_t$ and $\mathbf{\Sigma}_t$ using known results (Lemma 4), and bound the differential inequality coefficients using our key Lemma 1 (resulting in Lemmas 5 and 6).

**Step 2:** We bound the KL distance between the path measures of the true and approximate reverse process. We use an analogous Girsanov-based method to Chen et al. (2023a;d).

**Step 3:** We use the data processing inequality to bound $\textup{KL}(q_{\delta}||p_{t_N})$ in terms of the distance between reverse path measures and the distance between $q_T$ and $\pi_d$. We bound the former using Step 2 and the latter using the convergence of the OU process (Proposition 4), as in Chen et al. (2023a;d).

> However, is the perspective taken where they view the single process as solutions to (2) and (4) under different measures new?

No, this is also done in Chen et al. (2023a;d). We will make this clear in our updated introduction to Section 3.

> The biggest issue, though, in my opinion, is that much of the proof is left to the appendix. I understand this is necessary, but without more interpretation in the paper to help the reader understand what is happening, it is easy to get lost. Improving this explanation would greatly help the paper since it is primarily theoretical. Some suggestions are the following...

Hopefully, our improved introduction to Section 3 will help to mitigate any lack of clarity caused by the fact that some technical details must inevitably be left to the appendix. In particular, we have taken on board the reviewer's helpful suggestions to pull out the key results from the sections on the bound on the discretization error and the application of Girsanov's theorems, and will use these to add more context to the proofs.

Reviewer comment (Presentation):
In the bound on the discretization error section, include an initial lemma stating the final bound for this section so that the reader can put into context the computations being done. Then, put Lemmas 2–6 into context for how they help with the final bound as they are discussed.

Author response:
We apologize for the lack of clarity surrounding this overarching structure in the original manuscript. We will provide an updated manuscript by tomorrow with a clearer roadmap at the start of Section 3 explaining which sections of our work are novel. We will also slightly modify the structure of Section 3 to make it easier to get an overview of the proof at a glance, including adding in lemmas summarising the key results of the sections on the bound on the discretization error and the application of Girsanov's theorem.

With the updated structure, our proof comes in three key steps. We view Step 1 as our main novel contribution, while Steps 2 and 3 closely follow Chen et al. (2023a;d).

**Step 1:** We bound the error from discretizing the reverse SDE. Previous works bound $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ using a Lipschitz assumption. We use a new Ito calculus argument to get a differential inequality for $E_{s,t}$ (Lemmas 2 and 3), relate the coefficients of this inequality to $\mathbf{m}_t$ and $\mathbf{\Sigma}_t$ using known results (Lemma 4), and bound the differential inequality coefficients using our key Lemma 1 (resulting in Lemmas 5 and 6).

**Step 2:** We bound the KL distance between the path measures of the true and approximate reverse process. We use an analogous Girsanov-based method to Chen et al. (2023a;d).

**Step 3:** We use the data processing inequality to bound $\textup{KL}(q_{\delta}||p_{t_N})$ in terms of the distance between reverse path measures and the distance between $q_T$ and $\pi_d$. We bound the former using Step 2 and the latter using the convergence of the OU process (Proposition 4), as in Chen et al. (2023a;d).

Hopefully, our improved introduction to Section 3 will help to mitigate any lack of clarity caused by the fact that some technical details must inevitably be left to the appendix. In particular, we have taken on board the reviewer's helpful suggestions to pull out the key results from the sections on the bound on the discretization error and the application of Girsanov's theorems, and will use these to add more context to the proofs.

Reviewer comment (Presentation):
In the Girasanov section, state the final bound as a lemma so that the discussion can be put into context. This can then be pointed to throughout.

Author response:
We apologize for the lack of clarity surrounding this overarching structure in the original manuscript. We will provide an updated manuscript by tomorrow with a clearer roadmap at the start of Section 3 explaining which sections of our work are novel. We will also slightly modify the structure of Section 3 to make it easier to get an overview of the proof at a glance, including adding in lemmas summarising the key results of the sections on the bound on the discretization error and the application of Girsanov's theorem.

With the updated structure, our proof comes in three key steps. We view Step 1 as our main novel contribution, while Steps 2 and 3 closely follow Chen et al. (2023a;d).

**Step 1:** We bound the error from discretizing the reverse SDE. Previous works bound $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ using a Lipschitz assumption. We use a new Ito calculus argument to get a differential inequality for $E_{s,t}$ (Lemmas 2 and 3), relate the coefficients of this inequality to $\mathbf{m}_t$ and $\mathbf{\Sigma}_t$ using known results (Lemma 4), and bound the differential inequality coefficients using our key Lemma 1 (resulting in Lemmas 5 and 6).

**Step 2:** We bound the KL distance between the path measures of the true and approximate reverse process. We use an analogous Girsanov-based method to Chen et al. (2023a;d).

**Step 3:** We use the data processing inequality to bound $\textup{KL}(q_{\delta}||p_{t_N})$ in terms of the distance between reverse path measures and the distance between $q_T$ and $\pi_d$. We bound the former using Step 2 and the latter using the convergence of the OU process (Proposition 4), as in Chen et al. (2023a;d).

Hopefully, our improved introduction to Section 3 will help to mitigate any lack of clarity caused by the fact that some technical details must inevitably be left to the appendix. In particular, we have taken on board the reviewer's helpful suggestions to pull out the key results from the sections on the bound on the discretization error and the application of Girsanov's theorems, and will use these to add more context to the proofs.

Reviewer comment (Novelty):
Majority of the proof closely follows the analysis of Chen et al. (2023a;d).

Author response:
We thank the reviewer for their thoughtful remarks and suggestions.

> Majority of the proof closely follows the analysis of Chen et al. (2023a;d).

We agree that the structure of our proof closely resembles that of Chen et al. (2023a;d). We see our main contribution as being a more refined treatment of the time-discretization component of their bounds, and the relationship between this treatment and the stochastic localization literature. We are planning to update the introduction to Section 3 with a more thorough overview of which parts of our work are novel and which parts follow previous results; see our response to reviewer TDrD below for more details.

Reviewer comment (Theory):
The paper improves upon best existing bounds in the absence of uniform-Lipschitzness of the score function. However, more explanation is required to understand why Lipschitzness is a limiting assumption. For instance, while for discrete measures it is apparent that one must resort to early stopping on a corrupted version of the distribution, it requires further justification for measures absolutely continuous w.r.t the Lebesgue measure. Furthermore, the proofs in the algorithmic stochastic localization works such as El Alaoui et al. (2022), themselves rely on the Lipschitzness of the estimated score function.

Author response:
> However, more explanation is required to understand why Lipschitzness is a limiting assumption. For instance, while for discrete measures it is apparent that one must resort to early stopping on a corrupted version of the distribution, it requires further justification for measures absolutely continuous w.r.t the Lebesgue measure.

To clarify, our motivation for wishing to remove the Lipschitz assumption is two-fold. Firstly, as the reviewer comments, this assumption is inappropriate for discrete measures, but also for measures which are supported on a sub-manifold of $\mathbb{R}^d$. We expect that many data distributions of interest will be of this form.

Secondly, even if the score of the data distribution is Lipschitz, the Lipschitz constant will be very large when $t \rightarrow 0$ for many data distributions (e.g. those approximately supported on a sub-manifold). Crucially, in many cases we expect that the Lipschitz constant will scale with the data dimension (as is seen, for example, in Chen et al. (2023a), where they use early stopping to ensure that the scores of the marginal distributions are uniformly Lipschitz, but in doing so pick up an extra $d$ dependence in the Lipschitz constant). As such, previous assumptions of a uniform-in-time Lipschitz constant can hide the full dimension dependence. We therefore consider them misleading in many practical cases, even if the data distribution is smooth. As this intuition is a key part of the motivation for our work, we thank the reviewer for highlighting it and will clarify the presentation in the camera-ready version.

> For what classes of distributions is the uniform Lipshitzness assumption restrictive? What iteration complexity could one expect under smoothness assumptions non-uniform in time?

Hopefully, our response above clarifies that the uniform Lipschitzness assumption is restrictive for data supported on a sub-manifold, and also may hide sub-optimal dependence on the data dimension in many additional cases -- for example, any case where the Lipschitz constant implicitly scales with the data dimension, which may occur if the data distribution is approximately distributed on a sub-manifold. Using a Lipschitz constant which varies in time will not automatically avoid these problems; indeed, the work of Chen et al. (2023a) only assumes Lipschitz smoothness at $t=0$, but still obtains quadratic dependence in the data dimension. Fundamentally, it seems to us that any proof method which aims to control a quantity of the form $E_{s,t} := \mathbb{E}[||\nabla \log q_{T-t}(Y_t) - \nabla \log q_{T-s}(Y_s)||^2]$ using Lipschitzness of the scores - i.e. a supremum bound - rather than working in expectation will pick up an extra factor of $d$ and so be suboptimal compared to the results in our work.