Diffusion Transformer Captures Spatial-Temporal Dependencies: A Theory for Gaussian Process Data

We thank the reviewer for the valuable comments and suggestions. We would appreciate it if you could champion our paper during the discussions!

Although the overall writing is relatively clear, there’re some suggestions: 1.a) The object of this paper is sequential data, which has the axis of time (real world time, h). while in the diffusion process, there’s another axis of diffusion step (t), and authors also name it as ‘time’. To avoid confusion, authors should use some other descriptions to make it clearer. 1.b) The symbols should be consistent. E.g., the definitions of vt, D and v0t the use of subscripts of data x are confusing. 1.c) The font size of Figures could be larger, e.g. Fig.2, otherwise the equations and words are difficult to recognize under normal physical paper size.

We thank the reviewer for the advice on the meaning of “time” and “diffusion timesteps”, notation consistency and font size. We’ve changed all the “time” related to diffusion process to “diffusion timesteps” (always denoted as $t$, $t_0$ or $T$), and we refer the all the “time”s that appear in “time index” or “time embedding” to Gaussian Process (always denoted as numbers from $1$ to $N$ or $h_1,h_2,\dots,h_N$). We have also adjusted the font size of the figures for a better readability.

Although the effectiveness of the theory has been proved on synthesized Gaussian process data and semi-synthetic video data. More real data with complex contents should also be tested on.

We agree that experiments on real data offer stronger evidence of the effectiveness of our theory, the scale of real-data experiments like video generation is too large and beyond the scope of a theory-based work. The success of DiT in sequential data (like video data) has been empirically verified by a large number of works, such as SORA, openSora, kling AI, e.t.c. It's worth mentioning that we've conducted new experiments comparing DiTs with Unets and examining the training dynamics of DiTs, which is shown in the general response and Appendix A for the reviewer's interest.

The theorems proposed in the paper are based on the authors’ assumptions, then: 1.a) Does it mean that the theory is only valid under these assumptions, or does the model only work on the data satisfying these assumptions? b) How to ensure that the assumptions can represent the case in real sequential data?

We believe our theoretical framework is applicable to a broad class of data distribution. Gaussian Process is one of the most common ways of modeling sequential data. Please see our general response for a detailed discussion on the assumptions.

Our theory can be applied well beyond our assumptions and simplification. Indeed, our results suggest that as long as the score function can be represented by an efficient algorithm and transformers can implement the algorithm, then it is expected that DiT will capture the sequential distribution well. We believe this analytical framework can be applied to nonstationary Gaussian process and broad sequential data distributions. However, rigorous theoretical analysis requires further investigation.

How do authors derive these assumptions? E.g., in Assumption 1, why the definition of the covariance function has that formula, are there any theoretical basis for it?

The covariance function is derived from the definition of the stationary Gaussian process, which has been widely assumed in statistics literature when studying the inference and prediction of sequential data . The structure of the covariance matrix is required for our approximation theory (Theorem 1), but not the core intuition of our proof, i.e, algorithm unrolling. Please also see our general response for a detailed discussion about the assumptions of the Gaussian process.

We thank the reviewer for the valuable comments and suggestions.

While the results convincingly show that transformers can learn spatial-temporal dependencies, the experiments only show examples of properly learned tasks. Given that theoretical boundaries were derived, I would have loved more experiments investigating these limits. For example something like Figure 1 but with a degradation for higher length. The same goes for Section 6.2. What I am missing is a clear connection between model complexity and performance. Lack of comparison to other methods: I think it would be helpful to see how typical diffusion models, including Unet-based architectures, perform on the task of learning Gaussian processes. Clear results could help in shifting the design towards transformer-based diffusion architectures.

We thank the reviewer for the advice on experiments. We conducted new experiments with different sequence length (N) and model size (the number of transformer blocks in DiT) and also compare DiTs with UNets. Please see our general response and Appendix A in our updated manuscript for details.

Generalization to real-world problems: There are no experiments showing how well the insights generated by the experiments translate to real-world problems. You present the ball experiment but even there you have to consider how well your experiments align with the assumption of using Gaussian process data (cmp. l. 504-506). Real-world problems will often not follow Gaussian process dynamics Regarding the practical implication of our theory, please also see our general response. It would be nice if you could state some take-home messages for the reading. Ideally some general information such as design choices for the model architecture.

Our take-home message is that diffusion transformers capture spatial-temporal correlation in sequential data through its unique attention mechanism. For the guidance on model architecture, as shown in our theory, when learning sequential data with length $N$ and a decaying temporal dependency, it is suggested to equip the diffusion transformers with $O(\log N)$ transformer blocks for the best performance as predicted by our theory and experiments. Moreover, according to the roles of each components in transformers, if we want to improve the temporal correlations among tokens, we may try making more emphasis on the attention block (such as raising the latent dimension of query and key weight matrices), when we want to improve the single-frame quality, we may try emphasizing the value weight matrices and the residual blocks, or improving the latent encoder outside the DiT architecture. Also, our truncation technique (described in Figure 2) shows that we may impose a window-limited attention with window length depending on the temporal correlation to reduce inter-timestep computation without losing too many temporal dependencies.

The learned temporal dependencies presented in Figure 5 look very similar to what you get when you just plot the attention between two uncorrelated matrices with positional embedding (l. 300). What is the difference between these maps and the ones you present, e.g., in Fig. 5

We remark that the decay of the temporal dependencies can be different from that of two uncorrelated matrices because our theory is applicable to general $\ell$ and $\nu$. In the training process, we actually observe that the temporal dependencies of generated samples are gradually adapting to the temporal dependencies of the truth data. Please see Figure 9 in Appendix A for details.

Dear Reviewer zotT,

Thank you for your valuable feedback on our paper. We believe our revisions have fully addressed your concerns and significantly improved our work.

If you have any additional questions, we're available for discussion.

Your insights have been instrumental in enhancing our paper. Thank you again for your thorough review.

Best regards,

The Authors

We thank the reviewer for the valuable comments and suggestions. We would appreciate it if you could champion our paper during the discussions!

Can we have a section in the appendix where all the assumptions and simplifications are listed? This would really ease the conclusions that the reader can draw from this paper. There are some duplicate looking sections in the appendix, for example relating to the size of the transformer blocks. Can't we gather those in a big table instead of repeating them after each relevant section?

Thanks for the suggestion. We have added a short paragraph in Appendix C to list assumptions and simplifications we have made. We have also presented the network sizes using tables at corresponding places. The overall network size is listed in Table 3, which combines the results from Tables 1 and 2.

Ln. 397: What's a "properly transformer" architecture?

By proper transformer architecture, we refer to the constructed transformer $\mathcal{T}(D, L, M, B, R_t)$ in Theorem 1. With this choice, we establish sample complexity in Theorem 2. Note that, the configuration parameters in $\mathcal{T}$ are chosen dependent on the sample size, ensuring that the transformer is expressive enough to accurately represent the score function and succinct meanwhile for efficient statistical learning.

Could the authors give more elaborations on what motivates the choices of ϵ and T in Thm. 2?

We choose $\epsilon$ depending on the sample size due to bias-variance trade-off when learning the score function. We note from the first inequality at the beginning of the Proof of Proposition 2. that $\epsilon$ is the bias (transformer approximation error) and all the remaining terms are statistical variance when learning the score function using $n$ i.i.d. samples. Given $\epsilon = \epsilon_\delta$, we observe that smaller $\epsilon$ leads to large variance and large $\epsilon$ leads to small variance but large bias. Therefore, we optimally choose $\epsilon$ so that the bias and variance have the same order of dependence on sample size $n$. The choice of $T$ is relevant to the forward mixing error in inequality $(17)$. As we terminate the forward process at a finite time $T$, the terminal distribution is not exactly standard Gaussian. However, as long as $T$ is sufficiently large, the terminal distribution is close to Gaussian and we show that the convergence speed is exponential. Therefore, it suffices to choose $T = \log n$ to match the $1/\sqrt{n}$ statistical error.

What about empirically testing the learning on GPLVMs? Wouldn't this be a natural fit?

From our understanding (if correct), GPLVM is a multi-output regression model which serves as a dimension reduction tool. However, GPLVM assumes the data are sampled from one Gaussian process, while we are considering learning and generating the whole Gaussian process. Could you provide more information on how to empirically test the learning on GPLVMs if possible? Thanks!

Thank you for supporting our work and explaining the GPLVM!

We think GPLVM could be a viable tool for modeling the Gaussian process data in our work, given its flexibility to learn low-dimensional representation of high-dimensional sequences. We note that GPLVM would work differently compared to diffusion models, where the latter capture the sequential data distribution by estimating the score function. Given stacked sequential data $X \in \mathbb{R}^{d \times N}$, GPLVM aims to find a mapping $f$ that can transform a low-dimensional latent Gaussian process data $Z \in \mathbb{R}^{q \times N}$ to match the distribution of $X$, where $q \ll d$. The training of GPLVM is often likelihood based. Our understanding is that GPLVM works similarly to variational autoencoder, with the ability to handle sequential data.

It is a very interesting direction to consider the observed sequential data being generated by a GPLVM. In this case, the score function may be more complicated and involve the differentiation of the mapping $f$. We believe that the temporal and spatial correlation should still be learned by the attention mechanism in a DiT as long as $f$ assumes certain smoothness regularity. A rigorous analysis requires future investigation though.

In the quadratic formulation in Equation (3), the score function estimation remains non-separable over timesteps (i.e., across x1,t, ..., xN,t). Would further decomposition allow for a distributed solution with minimal inter-timestep communication?

Although the score function estimation remains non-separable, we remark that our truncation technique (described in Figure 2) actually serves as a tool to reduce inter-timestep communication, which is similar to adding a sliding window in the attention block that deactivate the interaction between tokens that have small correlations.

Prior work, such as [Sahiner et al., 2022], has already shown that attention layers capture correlations and reflect low-rank priors, so it’s unsurprising that attention proves effective for correlated and stationary data.

We thank the reviewer for mentioning this work. Although both [Sahiner et al., 2022] and our work aims to interpret attention mechanism, the setting and viewpoint are different. We make emphasis on learning sequential data with transformer-based diffusion models, and leverage an algorithm-unrolling strategy for approximation and generalization theory, while [Sahiner et al., 2022] makes emphasis on vision transformers and provides analysis in a more optimization-based viewpoint. We will add a discussion of [Sahiner et al., 2022] in the revised version of our paper.

We appreciate the suggestions on the presentation of our work. We will emphasize the generality of our work and revise the confusing sentences for better presentation. We would be grateful if the reviewer could reconsider the evaluation based on our contributions. If you have any further questions, please do not hesitate to comment.

We thank the reviewer for the valuable comments and suggestions.

The goal of the paper lacks clarity. If the aim is to demonstrate the expressiveness of transformers for correlated or spatiotemporal data, this doesn’t necessarily relate to diffusion models and could be examined independently in a more simplified setting.

The goal of our paper is to theoretically investigate how diffusion transformers capture spatial-temporal dependencies for efficiently learning sequential data distributions. This is the core question raised in the introduction section on Page 2. To answer the question, we establish a sample complexity bound in Theorem 2, which relies on our developed transformer approximation theory for the score function in Gaussian process data. In other words, our efficient learning guarantees are a consequence of the modeling power of both diffusion processes and transformers. We acknowledge that transformers are known to be universal approximators in the existing literature, however, our study presents a very different perspective. In diffusion models, learning data distributions is conducted by estimating the score function, which is distinct from using a transformer directly for sequential data modeling in an autoregressive manner. To be more specific, for Gaussian process data, we closely examine the structures in the score function, and show that transformers can unroll a gradient descent algorithm for approximating the score function. This result is highly relevant in the context of diffusion models, yet goes beyond the conventional study of expressiveness of transformers.

The assumption of a Gaussian process with scale-wise stationarity is a strong simplification, limiting the applicability to realistic distributions. Consequently, the results offer limited insights into transformer design for broader, more complex distributions.

Please see our general response for a discussion about the assumption of Gaussian process and the practical implications of our theoretical results. Our paper is theory driven and derives the first sample complexity bound for diffusion transformers in learning Gaussian process data. The data assumption aims to provide clear theoretical insights while maintaining technical cleanness. In particular, under the assumption, we precisely characterize the spatial and temporal dependencies and show how these dependencies are learned in a transformer by performing algorithmic unrolling. Such theoretical insights generalize to broader and more complex distributions such as time series and kinetic dynamics, as long as the two essential ingredients are established: 1) the score function of the distribution is representable by an efficient algorithm, and 2) a transformer can approximate each iteration of the algorithm. Yet, rigorous theoretical analysis requires further investigation beyond the current paper.

For sequential data modeling, Gaussian processes are used frequently [3, 4]. Meanwhile, we have discussed the relevance of our results and assumptions to real applications in the last paragraph on Page 3. Our study aligns with the commonly used latent diffusion models, where raw data will be transformed into a low-dimensional embedding, presumably close to Gaussian [1, 2].

We recognize the importance of transformer architecture design for better sequential data modeling. We believe that understanding how transformers represent spatial-temporal dependencies is the first step towards a principled answer to architectural innovation.

The experiments are weak and unconvincing. To substantiate the theory, more realistic data should be used, and comparisons with alternative architectures, such as UNet (3D or 2+1D), known for capturing spatiotemporal correlations, would strengthen the argument.

We conducted new experiments on Unets, where we find that DiT outperforms Unets in capturing spatial-temporal correlations especially for long-horizon dependencies. Please see our general response and Appendix A in our updated manuscript for the details of our experimental results.

Dear Reviewer XNMz,

Thank you for your valuable feedback on our paper. We believe our revisions have fully addressed your concerns and significantly improved our work.

If you have any additional questions, we're available for discussion. Otherwise, we would greatly appreciate it if you could reconsider your score based on our revisions.

Your insights have been instrumental in enhancing our paper. Thank you again for your thorough review.

Best regards,

The Authors