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

Authors: Hengyu Fu, Zehao Dou, Jiawei Guo, Mengdi Wang, Minshuo Chen
Primary area: generative models
Keywords: Diffusion model, transformers, sequential data, spatial-temporal dependency, sample complexity

Abstract:
Diffusion Transformer, the backbone of Sora for video generation, successfully scales the capacity of diffusion models, pioneering new avenues for high-fidelity sequential data generation. Unlike static data such as images, sequential data consists of consecutive data frames indexed by time, exhibiting rich spatial and temporal dependencies. These dependencies represent the underlying dynamic model and are critical to validate the generated data. In this paper, we make the first theoretical step towards bridging diffusion transformers for capturing spatial-temporal dependencies. Specifically, we establish score approximation and distribution estimation guarantees of diffusion transformers for learning Gaussian process data with covariance functions of various decay patterns. We highlight how the spatial-temporal dependencies are captured and affect learning efficiency. Our study proposes a novel transformer approximation theory, where the transformer acts to unroll an algorithm. We support our theoretical results by numerical experiments, providing strong evidence that spatial-temporal dependencies are captured within attention layers, aligning with our approximation theory.

Decision: accept (poster)
Reviewer scores: 5, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1729466859862

thread_type:
official_review

<thread_object>
creator: reviewer_fvay

<content>
summary:
In this paper, authors explore the reasons for Diffusion Transformer (DiT)’s capability of capturing spatial-temporal dependencies within the sequential data. They theoretically establish score approximation and distribution estimation guarantee of DiT for learning Gaussian process data. Specifically, they replace the score function in diffusion models with gradient descent process defined by truncated covariance, which are further implemented with transformer models. Authors conduct some experiments to prove the effectiveness of their theory under their assumptions.

strengths:
1. This paper studies the principle behind the Diffusion Transformer and provides a theoretical explanation of its capability of modeling spatial-temporal correlations within data, which is the very first trial in related fields. 
2. Authors propose a novel score function approximation theory as well as the sample complexity bound. All the lemmas and theorems are well defined with sufficient mathematical proofs. 
3. The correctness of the proposed theory has been evaluated on both synthesized and real data through experiments.

weaknesses:
1. 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 $v_t$, $\mathcal{D}$ and $v_0^t$, the use of subscripts of data $\mathbf{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.   
2. 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.

questions:
1. 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? 
1.b) How to ensure that the assumptions can represent the case in real sequential data? 
1.c) 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730287878446

thread_type:
official_review

<thread_object>
creator: reviewer_zott

<content>
summary:
This paper provides a theoretical study on the limits of the learning capabilities of transformer-based diffusion models on learning spatial-temporal data. The authors achieve this by reformulating the task of learning the score function of Gaussian process data as a gradient descent algorithm. This algorithm can then be unrolled using a transformer archtitecture where the attention layers can be studied to evaluate how spatial-temporal dependencies are learned. 
The results are thoroughly motivated and dervied and validated using a few numerical examples. 
 
I believe that the theoretical results could have an impact in terms of understanding the learning of temporal dependencies but currently the practical results are not extensive and the take away messages are not clear enough

strengths:
- studying gaussian processes to asses spatio-temporal dependencies is inherently reasonable and could be used as an important theoretical framework to study neural architectures.

- The method mathematically is well grounded and derived

- High-quality figures visualizing key components of the contribution

weaknesses:
- 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.

- 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

- 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.

questions:
- 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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730756177994

thread_type:
official_review

<thread_object>
creator: reviewer_fqsv

<content>
summary:
This paper scrutinises diffusion transformers under a simplified setting where the input time series is assumed to follow a Gaussian process. The paper carefully derives the analytical forms of the dynamics and provides score approximation as well as data-distribution estimation guarantees in the simplified setting of interest. The authors cast this paper as a first step towards building more sophisticated theories.

strengths:
I really like the framing in this paper. While sounding like toy settings, GPs can be quite general and could be a good starting point of analysis. I especially appreciate the analytical forms, for example in the score approximator, and how the covariance function relates to positional embeddings. The paper is also quite well written. I have gone through parts of the proofs and they made sense - although I have not exhaustively verified every step.

weaknesses:
- 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.
- Ln. 397: What's a "properly transformer" architecture? 
- Could the authors give more elaborations on what motivates the choices of $\epsilon$ and $T$ in Thm. 2? 
- What about empirically testing the learning on GPLVMs? Wouldn't this be a natural fit? 
- 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?

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730829159217

thread_type:
official_review

<thread_object>
creator: reviewer_xnmz

<content>
summary:
This paper explores transformers as denoisers within diffusion models for generating spatiotemporal data, motivated by recent advancements in video foundation models, such as SORAH.
The approach begins by assuming a Gaussian process as the underlying data distribution and simplifies it to a scale-varying Gaussian process where the covariance changes only in scale over time. In this framework, the score function is related to the inverse of a modified covariance matrix. Since directly calculating this inverse is computationally prohibitive for high-dimensional latent vectors (dimension = d * N), the paper introduces a convex quadratic relaxation solvable via gradient descent. The authors then design a transformer architecture inspired by unrolled gradient descent steps. They also derive sample complexity bounds for the transformer’s ability to learn the covariance based on the smoothness and decay properties of the kernel matrix.
The paper includes basic experiments with synthetic Gaussian process data and a simple ball-motion scenario to validate the proposed method.

strengths:
- The paper addresses a timely and relevant problem, aiming to understand why spatiotemporal transformers effectively learn distributions in spatiotemporal data.

- The writing is clear, making the methodology and findings accessible.

weaknesses:
- 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 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.

-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.

questions:
- 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?

- 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.

Sahiner, A., Ergen, T., Ozturkler, B., Pauly, J., Mardani, M., & Pilanci, M. (2022, June). Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In the International Conference on Machine Learning (pp. 19050-19088). PMLR.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732265999691

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We would like to thank the reviewers for their valuable comments and suggestions. Here, we elaborate on the contributions of our theory and demonstrate how transformers efficiently learn spatial-temporal dependencies  through further experiments.

**1.The generality of our theory (not limited in our assumption)**

We agree with Reviewers XNMz and  zotT that directly assuming real-world data to follow a Gaussian process is a strong assumption. However, as discussed in our paper, our theory also applies to **latent diffusion models** (e.g., [1], [2]), where an autoencoding model maps real-world data to a low-dimensional latent space, making the latent distribution more suitable for score-based generative models. With this autoencoder, the distribution of the low-dimensional representation could be simpler, so the Gaussian process assumption is milder in this case. Additionally, some variants of VAEs perform inference on a latent Gaussian process model, assuming the prior distribution of the latent representation follows a Gaussian process for convenience in sequential data modeling (such as [3] and [4]). These works align well with our assumptions.

Moreover, our theoretical insights and analytical tools developed in the paper can be **extended to broader applications**. Firstly, we theoretically elucidate **the role of each transformer block** (see the discussion following Theorem 1 and Figure 3) in learning the score function of sequential data. These observations are strongly corroborated by the experiments in the paper (Section 5). Secondly, our **algorithm unrolling theory** of diffusion transformers applies as long as there is an efficient algorithm for computing the score function and each iteration of the algorithm can be implemented by transformer blocks. Therefore, our analytical tools can be generalized to time series, kinetic dynamics, etc.

**2.Further experiments on DiTs and comparison with UNets**

Please see **Appendix A** for additional numerical results. We summarize our major discoveries.

a) **Comparison with UNet-based diffusion**

We implement a UNet-based diffusion model with a similar size as the DiT to learn Gaussian Process data. After training, we independently generate samples to test the performance. We show that the spatial-temporal correlation generated by the DiT is close to the ground truth. While the UNet-based diffusion model struggles in capturing the spatial dependencies and exhibit clear inaccuracy for learning temporal dependencies.

b) **Sequence length and transformer size v.s. performance**

We examine the influence of the sequence length $N$ and the number of transformer blocks $L$ on the performance of DiT for learning Gaussian process data. We demonstrate that the estimation error increases mildly as the sequence length increases, which aligns with our $\sqrt{N}$-dependence in Theorem 2. We also show that the performance of DiT improves as the number of transformer blocks increases, yet at a marginally diminishing speed if $L$ is sufficiently large. This supports our construction of a transformer with $O(\log N)$ blocks.




[1] Blattmann et, al. “Stable video diffusion: Scaling latent video diffusion models to large datasets”

[2] Wang, et,al. Lavie: “High-quality video generation with cascaded latent diffusion models”

[3] Casale, et al. “Gaussian Process Prior Variational Autoencoders”, NeurIPS 2018.

[4] Fortuin, et al. “GP-VAE: Deep Probabilistic Time Series Imputation”, AISTATS 2020.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper develops a theoretical framework to explain how diffusion transformers (DiTs) capture spatial-temporal dependencies in sequential data, using Gaussian processes as a simplified model. It introduces a novel score function approximation and establishes sample complexity bounds, supported by mathematical proofs and experiments on synthetic Gaussian process data and semi-synthetic tasks. Strengths include a strong theoretical contribution to understanding transformers, well-founded mathematical derivations, and high-quality visualizations. Weaknesses include the reliance on Gaussian process assumptions, limited practical experiments on real-world data, and insufficient comparisons to alternative architectures like UNets. Missing elements include more diverse datasets and guidelines for architectural design beyond theoretical insights. The paper’s significant theoretical contributions and clear presentation justify acceptance, with suggestions to expand empirical validation and real-world applicability in future work.