Skip-DiT: Stabilizing and Accelerating Diffusion Transformers with Long-Skip-Connections and Spectral Constraints

Introduction

Diffusion Transformers (DiT) have become a go-to architecture for high-quality image and video generation, but they often suffer from dynamic feature instability during cached inference. This instability can lead to error amplification and degraded output quality. A new approach, Skip-DiT, addresses these issues by introducing Long-Skip-Connections (LSCs) and spectral constraints. In this article, we'll dive into the technical details of Skip-DiT, its architecture, and how it improves both training and inference efficiency.

What Changed: Introducing Skip-DiT

Skip-DiT is a variant of DiT that incorporates LSCs to stabilize feature dynamics during generation tasks. The key changes include:

• Long-Skip-Connections (LSCs): These connections link shallow and deep layers, allowing the model to preserve long-range features and reduce perturbation sensitivity.
• Spectral Constraints: These constraints help control the spectral norm of the model's layers, further stabilizing feature propagation.

Why It Matters

For practitioners, Skip-DiT offers several significant benefits:

• Stability in Caching: LSCs ensure that cached features remain stable across different timesteps, reducing error amplification.
• Training Acceleration: The architecture leads to 4.4x faster training and quicker convergence.
• Inference Efficiency: Skip-DiT achieves a 1.5-2x speedup in inference with minimal quality loss, making it suitable for real-time applications.

Technical Details

Architecture Overview

Skip-DiT modifies the vanilla DiT architecture by adding long-skip-connections between shallow and deep layers:

• Long-Skip-Connections (LSCs): These connections allow information to flow directly from early layers to later layers, bypassing intermediate computations. This helps maintain feature stability and reduces the risk of error propagation.
• Spectral Constraints: By controlling the spectral norm of the model's layers, these constraints ensure that the model remains stable and predictable.

Implementation Details

• Feature Similarity: Skip-DiT demonstrates higher feature similarity between standard and cache-accelerated outputs compared to other methods like FORA and Faster-Diff. This indicates better stability after caching.
• Caching Mechanism: The efficient statical caching mechanism reuses deep features across timesteps while updating shallow components, reducing computational overhead.

Benchmarks

• Training Efficiency:

  • Acceleration: Skip-DiT achieves a 4.4x speedup in training compared to vanilla DiT.
  • Convergence: Faster convergence rates lead to shorter training times and more efficient model tuning.

• Inference Performance:

  • Speedup: Inference is accelerated by 1.5-2x, making it suitable for real-time applications.
  • Quality: Despite the speedup, Skip-DiT maintains high output quality with minimal loss in fidelity compared to vanilla DiT and other caching methods.

Experimental Results

Extensive experiments on image and video generation tasks show that Skip-DiT outperforms existing DiT caching methods across various quantitative metrics:

• Image Generation:

  • Higher feature similarity between cached and standard outputs.
  • Faster convergence and training times.

• Video Generation:

  • Maintains higher PSNR (Peak Signal-to-Noise Ratio) even at greater speedup factors.
  • Consistent quality across different timesteps.

Conclusion

Skip-DiT represents a significant step forward in stabilizing and accelerating diffusion transformers. By incorporating long-skip-connections and spectral constraints, it addresses the inherent instability of DiT models during cached inference, leading to faster training, efficient inference, and high-quality outputs. For developers working on image and video generation tasks, Skip-DiT is a promising approach that combines stability and efficiency.