Capturing Conditional Dependence via Auto-regressive Diffusion Models

TL;DR

The paper tackles the limitation of vanilla diffusion models in learning conditional dependencies by introducing autoregressive (AR) diffusion across data patches. It provides the first theoretical results showing KL convergence guarantees for conditional distributions under AR diffusion, and demonstrates that AR diffusion incurs only a modest increase in inference cost while better capturing conditional structure when such dependencies exist. Empirically, AR diffusion outperforms vanilla diffusion on tasks with clear conditional dependencies and underperforms when dependencies are absent or disrupted, with extensive synthetic and real-data experiments validating the theory. The work thus offers a principled framework for incorporating conditional dependence into diffusion models and clarifies when AR diffusion provides practical gains.

Abstract

Diffusion models have demonstrated appealing performance in both image and video generation. However, many works discover that they struggle to capture important, high-level relationships that are present in the real world. For example, they fail to learn physical laws from data, and even fail to understand that the objects in the world exist in a stable fashion. This is due to the fact that important conditional dependence structures are not adequately captured in the vanilla diffusion models. In this work, we initiate an in-depth study on strengthening the diffusion model to capture the conditional dependence structures in the data. In particular, we examine the efficacy of the auto-regressive (AR) diffusion models for such purpose and develop the first theoretical results on the sampling error of AR diffusion models under (possibly) the mildest data assumption. Our theoretical findings indicate that, compared with typical diffusion models, the AR variant produces samples with a reduced gap in approximating the data conditional distribution. On the other hand, the overall inference time of the AR-diffusion models is only moderately larger than that for the vanilla diffusion models, making them still practical for large scale applications. We also provide empirical results showing that when there is clear conditional dependence structure in the data, the AR diffusion models captures such structure, whereas vanilla DDPM fails to do so. On the other hand, when there is no obvious conditional dependence across patches of the data, AR diffusion does not outperform DDPM.

Figures (9)

• Figure 1: Visualization of Task 1 and Task 2. In Task 1, a patch size of 16 ensures that correlated features (sun and shadow) are located in different patches. In contrast, Task 2 uses a patch size of 8, which disrupts the dependency between the square's lengths by segmenting them into different patches.
• Figure 2: Comparison of AR Diffusion and DDPM Performance on Task 1.\ref{['fig:task1_tr']} demonstrates the validity of the evaluation method using the training data as a baseline. \ref{['fig:task1_ar']} and \ref{['fig:task1_ddpm']} illustrate the performance of AR Diffusion and DDPM during the inference phase, showing that AR Diffusion better captures inter-feature dependencies with a higher $R^2$. \ref{['fig:task1_loss']} presents the difference in training loss between DDPM and AR Diffusion, denoted as $\delta$. For most training steps, AR Diffusion's training loss is lower than that of DDPM, with $\delta > 0$.
• Figure 3: Comparison of AR Diffusion and DDPM Performance on Task 2.\ref{['fig:task2_tr']} demonstrates the validity of the evaluation method using the training data as a baseline. \ref{['fig:task2_ar']} and \ref{['fig:task2_ddpm']} illustrate the performance of AR Diffusion and DDPM during the inference phase, showing that AR Diffusion captures inter-feature dependencies less effectively, with a lower $R^2$. \ref{['fig:task2_loss']} presents the difference in training loss between DDPM and AR Diffusion, denoted as $\delta$. For most training steps, AR Diffusion's training loss is higher than that of DDPM, with $\delta \leq 0$. .
• Figure 4: The training loss variance with a sliding window of $500$. We observed that at the end of training, AR Diffusion and DDPM exhibit training stability with loss variation less than $1e-5$.
• Figure 5: Comparison of AR Diffusion with U-Net and DDPM with U-Net Performance on Sun-Shadow Setting.\ref{['fig:task1_ar_unet_app']} and \ref{['fig:task1_ddpm_app']} illustrate the performance of AR Diffusion with U-Net and DDPM with U-Net during the inference phase, showing that AR Diffusion better captures inter-feature dependencies with a higher $R^2$. \ref{['fig:task1_loss_unet']} presents the difference in training loss between DDPM and AR Diffusion, denoted as $\delta$. For most training steps, AR Diffusion's training loss is lower than that of DDPM, with $\delta > 0$.

Theorems & Definitions (45)

• Lemma 3.1
• Remark 1
• Lemma 3.2
• Lemma 3.3
• Remark 2
• Lemma 4.1
• Lemma 4.2
• Theorem 4.3
• Remark 3
• Remark 4