LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers

Introduction

Diffusion models (Ho, Jain, and Abbeel 2020; Rombach et al. 2022a; Song et al. 2020; Song and Ermon 2019; Dhariwal and Nichol 2021; Zhan et al. 2024b) have become dominant in image generation research, attributable to their remarkable performance. U-Net (Ronneberger, Fischer, and Brox 2015) is a widely used backbone in diffusion models, while transformers (Vaswani et al. 2017) are increasingly proving to be a strong alternative. Compared to U-Net, transformer-based diffusion models have demonstrated superior performance in high-fidelity image generation (Peebles and Xie 2023; Bao et al. 2023), and their efficacy extends to video generation as well (Lu et al. 2023; Chen et al. 2023; Lab and etc. 2024; Zheng et al. 2024). This highlights the versatility and potential of transformers in advancing generative tasks across different media. Despite the notable scalability advantages of transformers, diffusion transformers face major efficiency challenges. The high deployment costs and the slow inference speeds create the significant barriers to their practical applications (Zhan et al. 2024a, c, 2021; Wu et al. 2022; Li et al. 2022; Yang et al. 2023a), which motivates us to explore their acceleration methods.

The increased sampling cost in diffusion models stems from two main components: the numerous timesteps required and the computational expense associated with each inference step. To improve sampling efficiency, existing methods generally fall into two categories: reducing the total number of sampling steps (Song, Meng, and Ermon 2020; Liu et al. 2022; Bao et al. 2022; Zhan et al. 2024b) or lowering the computational cost per step (Yang et al. 2023b; He et al. 2023). Several works (Yin et al. 2024; Luo et al. 2024; Salimans and Ho 2022) employ distillation techniques to reduce the number of sampling steps. Conversely, works (Li et al. 2023c; Kim et al. 2023; Fang, Ma, and Wang 2023; Li et al. 2023b) utilize compression techniques to streamline diffusion models. Recently, some studies have introduced caching mechanisms into the denoising process (Ma, Fang, and Wang 2024; Wimbauer et al. 2023) to accelerate the sampling. However, previous compression approaches of diffusion models have primarily focused on optimizing U-Net, leaving transformer-based models largely unexplored.

Leveraging characteristic of uniquely structured, prior compression works (Zhang et al. 2024; Raposo et al. 2024; Fan, Grave, and Joulin 2019; Kong et al. 2022, 2023; Zhang et al. 2022; Li et al. 2023d; Zhao et al. 2024; Shen et al. 2024d, a, c, b, 2023b) have concentrated on techniques such as layer pruning and width pruning. However, we observe that removing certain layers results in a significant performance drop. This indicates the redundancy in diffusion transformers primarily occurs between sampling steps rather than the model architecture. This finding forms basis for exploring methods to reduce frequency of layer usage, aiming to decrease computational costs and accelerate the diffusion.

In this paper, we propose LazyDiT, a cache-based approach designed to dynamically reduce computational costs and accelerate the diffusion process. We begin by analyzing the output similarity between the current and previous steps, identifying that the lower bound of this similarity is notably high during the diffusion process. Then, we delve deeper into the similarity using a Taylor expansion around the current input, revealing that the similarity can be linearly approximated. Building on the theoretical analysis, we implement a lazy learning framework by introducing linear layers before each Multi-Head Self-Attention (MHSA) and pointwise feedforward (Feedforward) module. These added layers are trained with the proposed lazy loss to learn whether the subsequent module can be lazily bypassed by leveraging the previous step’s cache. Compared to the DDIM sampler, extensive experiments demonstrate that our method achieves superior performance with similar computational costs. As shown in Figure 1, by lazily skipping 50% of the computations, our method achieves nearly the same performance as the original diffusion process. We also profile the latency of the diffusion process on mobile devices to offer a detailed comparison with the DDIM sampler. Our results show our superior image generation quality than DDIM with similar latency. Our main contributions are summarized as follows,

We explore the redundancy in diffusion process by evaluating the similarity between module outputs at consecutive steps, finding that the lower bound of the similarity is notably high.

We establish that the lazy skip strategy can be effectively learned through a linear layer based on the Taylor expansion of similarity.

We propose a lazy learning framework to optimize the diffusion process in transformer-based models by lazily bypassing computations using the previous step’s cache.

Experiments show that the proposed method achieves better performance than DDIM sampler. We further implement our method on mobile devices, showing that our method is a promising solution for real-time generation.

Related Work

Recent works such as GenVit (Yang et al. 2022), U-Vit (Bao et al. 2023), DiT (Peebles and Xie 2023), LlamaGen (Sun et al. 2024), and MAR (Li et al. 2024a) have incorporated transformers (Vaswani et al. 2017) into diffusion models, offering a different approach compared to the traditional U-Net architecture. GenViT incorporates the ViT (Dosovitskiy et al. 2021; Li et al. 2024e, d) architecture into DDPM, while U-ViT further enhances this approach by introducing long skip connections between shallow and deep layers. DiT demonstrates the scalability of diffusion transformers, and its architecture has been further utilized for text-to-video generation tasks, as explored in works (OpenAI 2024). LlamaGen introduces autoregressive models to image generation, verifying the effectiveness of the ’next-token prediction’ in this domain. Thus, it is crucial to explore efficient designs for those large models to accelerate the diffusion process.

Acceleration for Diffusion Models.

High-quality image generation with diffusion models necessitates multiple sampling steps, leading to increased latency (Gong et al. 2024; Shen et al. 2023a). To enhance efficiency, DDIM (Song, Meng, and Ermon 2020) extends original DDPM to non-Markovian cases when DPM-Solver (Lu et al. 2022) advances the approximation of diffusion ODE solutions. Regarding the works that require fine-tuning, such as (Lin, Wang, and Yang 2024; Yin et al. 2024), they employ distillation techniques to effectively reduce the number of sampling steps. Additionally, reducing the computational workload for each diffusion step is a widely adopted, strategy to enhance the efficiency of the diffusion process. Various approaches have been explored, such as works (Fang, Ma, and Wang 2023; Castells et al. 2024; Wang et al. 2024a; Zhang et al. 2024) that adopt weight pruning techniques, works (He et al. 2023; Li et al. 2023b) that employ quantization techniques, and even works (Kim et al. 2023; Li et al. 2023c) that redesign the architecture of diffusion models.

Methodology

Diffusion Formulation.

Diffusion models (Ho, Jain, and Abbeel 2020; Song et al. 2020) operate by transforming a sample $x$ from its initial state within a real data distribution $p_{\beta}(x)$ into a noisier version through diffusion steps. For a diffusion model $\epsilon_{\theta}(\cdot)$ with parameters $\theta$, the training objective (Sohl-Dickstein et al. 2015) can be expressed as follows,

where $t$ denotes the timestep; $\epsilon$ denotes the ground-truth noise; $z_{t}=\alpha_{t}\cdot x+\sigma_{t}\cdot\epsilon$ denotes the noisy data; $\alpha_{t}$ and $\sigma_{t}$ are the strengths of signal and noise.

For comparison purposes, this paper adopts Denoising Diffusion Implicit Models (DDIM) (Song, Meng, and Ermon 2020) as sampler. The iterative denoising process from timestep $t$ to the previous timestep $t^{\prime}$ is described as follows,

where $z_{t^{\prime}}$ is iteratively fed to $\epsilon_{\theta}(\cdot)$ until $t^{\prime}$ becomes .

Latent Diffusion Models.

The Latent Diffusion Model (LDM) (Rombach et al. 2022b) decreases computational demands and the number of steps with the latent space, which is obtained by encoding with a pre-trained variational autoencoder (VAE) (Sohl-Dickstein et al. 2015). Besides, the classifier-free guidance (CFG) (Ho and Salimans 2022) is adopted to improve quality as follows,

where $\epsilon_{\theta}(t,z_{t},c_{\phi})$ denotes unconditional prediction with null text; $w$ denotes guidance scale which is used as control of conditional information and $w\geq 1$.

2 Similarity Establishment

Meanwhile, we define broadcasted matrices to represent the scaling and shifting factors, ensuring the alignment with the implementation of diffusion transformers as follows,

Then, we deliver the demonstration showing that there exist $y_{t},y_{t-1}$ such that, after scaling and shifting, the distance between inputs $X_{l,t-1}^{\Phi},X_{l,t}^{\Phi}$, defined in the Left Hand Side of the following Eq. (1), can be constrained within a small bound. Given $a_{t}$ and $b_{t}$ are both linear transformation of $y_{t}$, the problem reduces to demonstrating the existence of vectors $a_{t}$, $b_{t}$, $a_{t-1}$, and $b_{t-1}$ that satisfy following conditions,

where $\eta\in(0,0.1)$. And Eq. (1) is equivalent as follows,

where $A:=A_{t-1},B:=-A_{t},C:=B_{t-1}-B_{t}$.

We identify that there exists $a$, $b$ and $c$ such that Eq. (2) holds, and the detailed demonstration and explanation are included in Lemma 12 at Appendix C.2. Subsequently, we generate the following theorem,

There exist time-variant and condition-variant scalings and shiftings such that the distance between two inputs at consecutive steps for MHSA or Feedforward is bounded.

Similarity Lower Bound.

To leverage the cache from the previous step, we begin by investigating the cache mechanism in transformer-based diffusion models. This analysis focuses on the similarity between the current output and the preceding one, providing insights into the efficacy of reusing cached information. One typical transformer block consists of two primary modules: MHSA module and Feedforward module. Both modules are computationally expensive, making them significant contributors to the overall processing cost. Thus, we aim to examine the output similarities between the current and previous steps for both modules. By identifying cases of high similarity, we can skip redundant computations, thereby reducing the overall computational cost. In practice, we employ the cosine similarity $f(\cdot,\cdot)$ for the computation of similarity as follows,

where $C$ is the Lipschitz constant related to the module.

Subsequently, with Theorem 1, we integrate Eq. (1) and derive the bound of the similarity as follows,

for $\alpha:=O(C^{2}\eta^{2})$ and $\eta$ is sufficiently small in practice.

Thus, we deliver Theorem 2 as below, which asserts that the lower bound of the output similarity between the two consecutive sampling steps is high.

The lower bound of the similarity $f(Y^{\Phi}_{l,t-1},Y^{\Phi}_{l,t})$ between the outputs at timestep $t-1$ and timestep $t$ is high.

Linear Layer Approximation.

The similarity can be approximated using the inputs from either the current step $Z_{l,t}^{\Phi}$ or previous one $Z_{l,t-1}^{\Phi}$, due to its mathematical symmetry according to Eq. (3). We then apply the Taylor expansion around $Z_{l,t}^{\Phi}$ as follows,

where the detailed proof is included in Appendix C.5 Eq.(9).

Then, we generate the Theorem 3 as follows,

The similarity function $f(\cdot,\cdot)$ can be approximated by a linear layer with respect to the current input, i.e. $f(Y^{\Phi}_{l,t-1},Y^{\Phi}_{l,t})=\langle W^{\Phi}_{l},Z_{l,t}^{\Phi}\rangle$ where $W^{\Phi}_{l}$ is the weight of a linear layer for MHSA or Feedforward in the $l$-th layer of diffusion model.

3 Lazy Learning

As illustrated in Figure 2, we incorporate lazy learning linear layers before each MHSA module and Feedforward module to learn the similarity. The MHSA module or Feedforward module is bypassed and replaced with the cached output from the previous step if the learned similarity is below 0.5. The input scale, input shift, output scale, and residual connections remain unchanged from the normal computation. The training details and the calculation of lazy ratio are outlined in the following paragraphs.

We then define the forward pass of the MHSA module or Feedforward module at $l$-th layer and $t$-th step with the input $X_{l,t}^{\Phi}$ during the training progress as follows,

Backward Loss.

Alongside the diffusion loss for a given timestep $t$ during training, we introduce a lazy loss to encourage the model to be more lazy—relying more on cached computations rather than diligently executing the MHSA modules or Feedforward modules, as follows,

Accelerate Sampling.

After finishing the lazy learning with a few steps, we then accelerate the sampling during the diffusion process as follows,

Experimental Results

We validate the effectiveness of our method on both the DiT (Peebles and Xie 2023) and LargeDiT (Zhang et al. 2023) model families. Specifically, our experiments utilize the officially provided models including DiT-XL/2 (256$\times$256), DiT-XL/2 (512$\times$512), Large-DiT-3B (256$\times$256), and Large-DiT-7B (256$\times$256).

Lazy Learning.

We freeze the original model weights and introduce linear layers as lazy learning layers before each MHSA and Feedforward module at every diffusion step. For various sampling steps, these added layers are trained on the ImageNet dataset with 500 steps, with a learning rate of 1e-4 and using the AdamW optimizer. Following the training pipeline in DiT, we randomly drop some labels, assign a null token for classifier-free guidance, and set a global batch size of 256. The training is conducted on 8$\times$NVIDIA A100 GPUs within 10 minutes.

Penalty Regulation.

Evaluation.

To evaluate the effectiveness of our method, we primarily compare our method to the DDIM (Song, Meng, and Ermon 2020), varying the sampling steps from 10 to 50. Visualization results are generated with DiT-XL/2 model in 256$\times$256 and 512$\times$512 resolutions. For quantitative analysis, the 50,000 images are generated per trial with classifier-free guidance in our experiments. We adopt the Fréchet inception distance (FID) (Heusel et al. 2017), Inception Score (IS) (Salimans et al. 2016), sFID (Nash et al. 2021), and Precision/Recall (Kynkäänniemi et al. 2019) as the evaluation metrics. The computation cost as TMACs is calculated with the work (Zhu 2022).

Testing Bed.

We implement our acceleration framework on mobile devices, specifically, we use OpenCL for mobile GPU backend. LazyDiT is built upon our existing DNN execution framework that supports extensive operator fusion for various DNN structures. We also integrated other general DNN inference optimization methods similar to those in (Chen et al. 2018; Abadi et al. 2016), including memory layout and computation graph. Results are obtained using a smartphone with a Qualcomm Snapdragon 8 Gen 3, featuring a Qualcomm Kryo octa-core CPU, a Qualcomm Adreno GPU, and 16 GB of unified memory. Each result take 50 runs, with average results reported as variance is negligible.

2 Results on ImageNet

We present the results generated with DiT officially released models compared to DDIM in Table 1. Full results with more model sizes and lazy ratios are included in Table 5 of Appendix A.1. Due to the addition of lazy learning layers, the computational cost of our method is slightly higher than that of DDIM. Our experiments demonstrate that our method can perform better than the DDIM on DiT models with 256$\times$256 and 512$\times$512 resolutions. Particularly, for sampling steps fewer than 10, our method demonstrates a clear advantage over DDIM at both resolutions, highlighting the promise of our approach. For larger models with 3B and 7B parameters, we present the results in Table 2. Compared to the DiT-XL/2 model with 676M parameters, Large-DiT models with a few billion parameters exhibit more redundancy during the diffusion process. Full results for Large-DiT models are included in Table 4 at Appendix A.1. Experiments demonstrate that at 50% lazy ratio, our method significantly outperforms the approach of directly reducing sampling steps with DDIM. We further visualize the images generation results in Figure 1. We also compare with other cache-based method Learn2Cache (Ma et al. 2024) which adopts input independent cache strategy and requires full training on ImageNet, the results are in Table 7 at Appendix A.4. For each sampling step, Learn2Cache only has one cache strategy, whereas our method outperforms it with less training cost, demonstrating both the effectiveness and the flexibility of our method.

3 Generation on Mobile

We present the latency profiling results on mobile devices in Table 3. Our method achieves better performance with less computation cost compared to DDIM. Additionally, when computational costs and latency are similar, our method significantly outperforms DDIM in terms of performance. Notably, with 10 sampling steps and a 30% lazy ratio, our method produces significantly higher image quality in 256$\times$256 resolution than the DDIM sampler. Besides, we visualize the images generated on mobile in Figure 3. The images in the last row, generated with our method, exhibit higher quality compared to the second row, which are generated without the laziness technique under similar latency. Therefore, our method is especially beneficial for deploying diffusion transformers on mobile devices, offering a promising solution for real-time generation on edge platforms in the future. Meanwhile, the latency results tested on GPUs are included in Table 6 at Appendix A.2. Our method delivers much better performance with faster latency on GPUs, especially when the number of sampling steps are fewer than 10. Moreover, with almost the same latency, our method performs much better than DDIM.

4 Ablation Study

We perform ablation studies on the laziness of MHSA and Feedforward modules separately by regulating the corresponding penalty ratios to determine the maximum applicable laziness for each, thereby exploring the redundancy within both components. We present the results generated with DDIM 20 steps on DiT-XL/2 (256$\times$256) in the upper figure in Figure 5. The analysis indicates that the maximum applicable lazy ratio is 30% for MHSA and 20% for Feedforward modules. The identification reveals that applying laziness individually to either MHSA or Feedforward network is not the most effective lazy strategy, which motivates us to apply the laziness to both modules simultaneously in our experiments.

Lazy Strategy.

To optimize laziness in both MHSA and Feedforward modules for optimal performance, we fix the laziness in one module and regulate the penalty ratio of the other, varying the lazy ratio from 0% to 40%. Specifically, we separately fix 30% lazy ratio to MHSA or 20% lazy ratio to Feedforward modules, and analyzed the model performance by regulating the lazy ratio of another module with DDIM 20 steps on DiT-XL/2 (256$\times$256). The results, as presented in the lower figure of Figure 5, reveal that the model achieves optimal performance when the same lazy ratio is applied to both MHSA and Feedforward. Thus, we adopt the same penalty ratio for both modules in our experiments to achieve the best performance.

Layer-wise Laziness.

To investigate the layer-wise importance during the diffusion process, we examined the laziness of each layer over 20 sampling steps with DiT-XL/2 model in 256$\times$256 resolution with 8 images. The results, visualized in Figure 4, illustrate the layer-wise lazy ratio distribution and highlight key patterns in layer importance. The analysis reveals that, for MHSA, the latter layers are more critical, whereas for Feedforward layers, the initial layers hold greater importance. This is evidenced by the decreasing lazy ratio in MHSA and the increasing lazy ratio in MLP as going deeper. Moreover, all layers contribute to the process, as there is no such layer that has a 100% lazy ratio, meaning no layer is completely bypassed. Therefore, strategies such as removing layers or optimizing model structure are not applicable for transformer-based diffusion models.

Conclusion and Limitation

In this work, we introduce the LazyDiT framework, designed to accelerate the diffusion process in transformer-based diffusion models. Specifically, we first highlight redundancy in the diffusion process by showing that the lower bound of similarity between consecutive steps is notably high. Then, we design the lazy learning framework, which incorporates a lazy skip strategy inspired by the Taylor expansion of similarity. Experimental results validate the effectiveness of our method, indicating that it performs better than the DDIM sampler. We further implement our method on mobile devices, achieving better generation performance than DDIM sampler with lower latency. For the limitation, the paper also has some shortcomings, such as the additional computation overhead brought by the lazy learning layers.

References

Appendix A More Results

We present the full results with Large-DiT models in Table 5 and DiT models including DiT-XL and DiT-L in Table 4. A diverse range of models with varying sampling steps has confirmed the effectiveness of our method. Besides, as the number of sampling steps decreases, our method proves to be even more effective compared to the standard 50-step sampling with the DDIM sampler. As there is no officially released DiT-L model, thus we pretrain it for one million steps. The results with DiT-L model here to show the generalization of our method to smaller models.

A.2 Latency on GPU

We present the latency profiling results with 8 images per batch in classifier-free guidance on a single A5000 GPU in Table 6. Our method achieves better performance with faster latency on GPUs. Meanwhile, for similar latency cost, our method outperform DDIM on the image generation quality.

A.3 Compare to Other Cache-Base Method

We further compare our method to the Learn2Cache method (Ma et al. 2024), we reproduce the training and evaluation with the given hyperparameter recipes. Notably, Learn2Cache requires a full epoch of training on ImageNet, which is significantly more intensive compared to our method, which only requires 500 training steps. Meanwhile, Learn2Cache does not support bigger laziness, offering only a single caching strategy for each specified number of sampling steps. The results are shown in Table 7. Our results achieve better performance than Learn2Cache method in most cases. Especially, our method performs much better with less computation cost when the resolution becomes 512$\times$512.

A.4 Ablation for Skipping MHSA or Feedforward

We present the ablation study for skipping either MHSA or Feedforward layers individually, utilizing the sparse learning weights trained with the combined skipping of both components at same sparsity. The results are shown in Figure 6.

A.5 More Visualization

More visualization is provided in Figure 7 in 512$\times$512 resolution compared to DDIM sampler with 50% lazy ratio.

Appendix B Theoretical Preliminary

In this section, we provide the preliminary of our theoretical results.

B.2 Facts

In this section, we provide several facts we use.

We have the following properties of trace:

where the first step follows from Fact 6, the second step follows from the basic algebra, the third and the last step follow from Fact 6. ∎

B.3 Definitions

In this section, we provide key definitions we use for the proofs. We define the attention module and feedforward module as follows.

where (1) $V:=XW_{V}$, (2) $W:=W_{Q}W_{K}^{\top}$, (3) $A:=\exp(XWX^{\top})$, and (4) $D:=\operatorname{diag}(A\cdot{\bf 1}_{n})$.

Appendix C Theoretical Results

For simplicity of proofs, we assume the batch size $B=1$. This will only affect the results by a constant factor of $B$. In this section, we provide the theoretical part of our paper.

The concept of lazy update has emerged as a powerful strategy in addressing computational challenges. Lazy updating defers certain computations or updates until they are absolutely necessary, which can significantly reduce computational overhead in large-scale problems. This approach has applications to several fundamental theoretical frameworks such as Linear Programming (LP) and Semi-Definite Programming (SDP) (Anstreicher 2000; d’Aspremont et al. 2006; Amini and Wainwright 2009; Diakonikolas et al. 2019; Dong, Hopkins, and Li 2019; Cohen, Lee, and Song 2021; Jambulapati, Li, and Tian 2020; Jiang et al. 2020a; Gu and Song 2022; Song, Ye, and Zhang 2023). The application of lazy update strategies extends beyond LP and SDP to other important areas in computer science and machine learning. Dynamic maintenance of data structures and algorithms is one such area where lazy updates have proven valuable (Cohen, Lee, and Song 2021; Lee, Song, and Zhang 2019; Brand 2020; Jiang et al. 2020b; Brand et al. 2020; Jiang et al. 2020a, c; Song and Yu 2021; Dong, Lee, and Ye 2023; van den Brand 2020; Huang et al. 2021; Gu and Song 2022). This approach allows for efficient updates to solutions as input data changes, without the need for complete recomputation. In the realm of machine learning, Empirical Risk Minimization (ERM) is a fundamental principle that has benefited from lazy update techniques. ERM has been extensively studied in various contexts (Nesterov 1983; Vapnik 1991; Polyak and Juditsky 1992; Bartlett, Bousquet, and Mendelson 2005; Bottou and Bousquet 2007; Nemirovski et al. 2009; Moulines and Bach 2011; Feldman et al. 2012; Nesterov 2013; Johnson and Zhang 2013; Vapnik 2013; Shalev-Shwartz and Zhang 2013; Défossez and Bach 2014; Frostig et al. 2015; Zhang and Xiao 2017; Jin et al. 2018; Lee, Song, and Zhang 2019; Lianke et al. 2023; Bian, Song, and Yin 2023), with recent work focusing on improving its efficiency through lazy updates. Support Vector Machines (SVMs), a popular machine learning model, have also seen advancements through the application of lazy update strategies. Recent research has explored ways to optimize SVM training and inference using these techniques (Chang and Lin 2001; Joachims 2006; Tarzanagh et al. 2023; Brand, Song, and Zhou 2023; Li, Song, and Zhou 2023; Gao, Mahadevan, and Song 2023; Gu, Song, and Zhang 2023; Gao et al. 2023a), leading to improved performance in both time and space complexity.

Lazy Machine Learning.

Lazy update strategies have been extensively explored in the context of machine learning to enhance algorithmic efficiency by reducing iterative loops. For instance, (Sprechmann et al. 2018) proposed a memory-based parameter adaptation (MbPA) method, leveraging lazy updates to optimize model parameters. In the domain of reinforcement learning, (Jacq et al. 2022) introduced lazy-MDPs, incorporating lazy update techniques to streamline the learning process. Furthermore, (Jagielski et al. 2023) investigated the interplay between forgetting and privacy protection during model training, highlighting the potential of lazy updates in balancing computational efficiency and data privacy. (Chen et al. 2021) proposed a meta-learning approach via in-context tuning, demonstrating the effectiveness of lazy update principles in adapting models to new tasks. Recently, (Du et al. 2023) provided a comprehensive overview of the shortcut learning problem in large language models for natural language understanding tasks, emphasizing the role of lazy update strategies in mitigating this challenge.

Attention Theory.

Attention mechanisms have been extensively studied and developed over the years, becoming a fundamental component in various neural network architectures. (Deng, Li, and Song 2023; Li et al. 2023a; Gu et al. 2024a) investigate the softmax regression problem, while (Song, Ye, and Zhang 2023) propose and analyze the attention kernel regression problem. The rescaled hyperbolic functions regression is examined by (Gao, Song, and Yin 2023). Following this, (Song, Wang, and Yin 2023; Li et al. 2023e) delve into two-layer attention regression problems. (Gu et al. 2024b) demonstrate that attention layers in transformers learn two-dimensional cosine functions. Additionally, (Deng et al. 2023) explore data recovery using attention weights, and (Kacham, Mirrokni, and Zhong 2023; Song, Xu, and Yin 2023) investigate the replacement of the softmax unit with a polynomial unit. Moreover, some works theoretically explore variations or combinations of the attention mechanism with other techniques, such as quantum attention (Gao et al. 2023b, 2024), tensor attention (Alman and Song 2024; Liang et al. 2024e), and differentially private attention (Liang et al. 2024d; Gao et al. 2023c; Gu et al. 2024c) and other applications such as (Clark et al. 2019; Tenney, Das, and Pavlick 2019; Hewitt and Liang 2019; Vig and Belinkov 2019; Belinkov 2022; Wang et al. 2023b, a, 2024b; Brand, Song, and Zhou 2023; Liang et al. 2024f; Chen et al. 2024c, b, a; Shrivastava, Song, and Xu 2023; Qin, Song, and Sun 2023; Deng, Mahadevan, and Song 2023; Li et al. 2024c; Liang et al. 2024b; Li et al. 2024b; Liang et al. 2024c, e, a).

C.2 Effect of Scaling Factor

In this section, we examine the scaling factor used in DiT architecture. We first prove the vector case.

Then, we can show that there exist $a,b,c$ such that

Let $a=\frac{0.5\eta}{\|x_{1}\|_{2}}\cdot{\bf 1}_{D}$, $b=\frac{0.5\eta}{\|x_{2}\|_{2}}\cdot{\bf 1}_{D}$, and $c=0$.

where the first step follows from the way we pick $a,b,c$, the second step follows from the triangle inequality, the last step follows from basic algebra. ∎

It can be proven that for any matrix, there exist scaling and shifting factors that can adjust the input such that the distance between them is bounded.

We define $M_{1}:=\max_{i\in[N]}\|X_{1,i}\|_{2}$ for $X_{1}$.

We define $M_{2}:=\max_{i\in[N]}\|X_{2,i}\|_{2}$ for $X_{2}$.

Then, we can show that there exist $a,b,c$ such that

Let $a={\bf 1}_{D}\cdot\frac{0.5\eta}{NM_{1}}$, $b={\bf 1}_{D}\cdot\frac{0.5\eta}{NM_{2}}$, and $c=0$. Let $\widetilde{A}=A\circ X_{1}$ and $\widetilde{B}=B\circ X_{2}$ and $\widetilde{C}=C$.

where the first step follows from the definition of $\widetilde{A}$, the second step follows from the definition of $a$, the third step follows from the definition of $M_{1}$ and $X_{1,i}$, the last step follows from the definition of $\eta_{0}$.

where the first step follows from the definition of $\widetilde{B}$, the second step follows from the definition of $b$, the third step follows from the definition of $M_{2}$ and $X_{2,i}$, the last step follows from the definition of $\eta_{0}$.

where the first step follows from definition of Frobenius norm, the second step follows from Cauchy–Schwarz inequality, the third step follows from Eq.(C.2) and Eq.(C.2), and the last step follows from definition of $\eta_{0}$.

Taking the square root of the both side of the above equation, this can be further simplified as:

Using the above idea, we can prove Theorem 13.

Then, we can show there exist $A_{t-1},A_{t},B_{t-1},B_{t}$

where the first step follows from basic algebra, the second step follows from Fact 5, the third step follows from we define $A:=A_{t-1},B:=-A_{t},C:=B_{t-1}-B_{t}$, the last step follows from Lemma 12. ∎

C.3 Lipschitz of Model Outputs

We first prove the Lipschitz for attention layer.

Assume $\|W\|\leq R,\|W_{V}\|\leq R,\|X\|\leq R$.

Using the previous results from (Deng et al. 2023), we can prove the Lipschitz for attention layer.

Assume $\|W\|\leq R,\|W_{V}\|\leq R,\|X\|\leq R$.

Then, we have the attention module (Definition 9) satisfying:

where the first step follows from Fact 4, the second step follows the definition of $\|\cdot\|_{\infty}$, and the second step follows from Lemma 14. ∎

We then prove the Lipschitz for feedforward layer.

Then, we have the feedforward module (Definition 10) satisfying:

Combining the Lipschitz of two modules, we have the following lemma.

Combining Lemma 15 and 16, we finish the proof. ∎

C.4 Similarity Lower Bound

In this section, we prove that the similarity of outputs between two diffusion timestep can be bounded.

We define $Y_{l,t}^{\Phi}:=\mathcal{F}_{l}^{\Phi}(X_{l,t}^{\Phi})$ and assume $\|Y_{l,t}^{\Phi}\|_{F}=1$ for any $t$.

Then, the similarity between $Y_{l,t-1}^{\Phi}$ and $Y_{l,t}^{\Phi}$ can be bounded by:

Suppose we use cosine similarity. We derive a new form of $f$ first.

where the first step follows from the definition of cosine similarity, the second step follows from Fact 7, the third step follows from $\|Y_{l,t}^{\Phi}\|_{F}=1$ for any $t$.

where the first step follows from Eq. (C.4), the second step follows from Fact 5, the third step follows from Lemma 17, the fourth step follows from $\|X_{l,t-1}^{\Phi}-X_{l,t}^{\Phi}\|\leq R_{2}$, and the last step follows from we define $\alpha:=0.5C^{2}R_{2}^{2}\min\{N,D\}$.

C.5 Approximating Similarity Function

In this section, we prove that the similarity function can be approximated by a linear linear with certain error.

The similarity function $f$ is assumed to be sufficiently smooth such that its derivatives up to at least the second order are well-defined and continuous over the range of interest.

We define $Y_{l,t}^{\Phi}:=\mathcal{F}_{l}^{\Phi}(X_{l,t}^{\Phi})$ and assume $\|Y_{l,t}^{\Phi}\|_{F}=1$ for any $t$.

There exist weights $W_{l}$ of a linear layer in the $l$-th layer of diffusion model such that $f(Y_{l,t-1}^{\Phi},Y_{l,t}^{\Phi})=\langle W^{\Phi}_{l},X_{l,t}^{\Phi}\rangle$.

Suppose we use cosine similarity. We have

where the first step follows from the definition of cosine similarity, the second step follows from $\|Y_{l,t}^{\Phi}\|_{F}=1$ for any $t$, the third steps from Fact 6, the fourth step follows from Fact 8 and we expand $Y_{l,t}^{\Phi}$ around , the fifth step follows from the simple algebra, the sixth follows from the definition of inner product, the last step follows from we define $W^{\Phi}_{l}:=(Y_{l,t-1}^{\Phi})^{\top}\cdot J$. ∎