SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer

TL;DR

This work novelly unleashes the self-supervised discrimination knowledge to boost DiT training by encoding discriminative pairs with student and teacher DiT encoders, designed to encourage the inter-image alignment in the selfsupervised embedding space.

Abstract

Diffusion Transformer (DiT) has emerged as the new trend of generative diffusion models on image generation. In view of extremely slow convergence in typical DiT, recent breakthroughs have been driven by mask strategy that significantly improves the training efficiency of DiT with additional intra-image contextual learning. Despite this progress, mask strategy still suffers from two inherent limitations: (a) training-inference discrepancy and (b) fuzzy relations between mask reconstruction & generative diffusion process, resulting in sub-optimal training of DiT. In this work, we address these limitations by novelly unleashing the self-supervised discrimination knowledge to boost DiT training. Technically, we frame our DiT in a teacher-student manner. The teacher-student discriminative pairs are built on the diffusion noises along the same Probability Flow Ordinary Differential Equation (PF-ODE). Instead of applying mask reconstruction loss over both DiT encoder and decoder, we decouple DiT encoder and decoder to separately tackle discriminative and generative objectives. In particular, by encoding discriminative pairs with student and teacher DiT encoders, a new discriminative loss is designed to encourage the inter-image alignment in the self-supervised embedding space. After that, student samples are fed into student DiT decoder to perform the typical generative diffusion task. Extensive experiments are conducted on ImageNet dataset, and our method achieves a competitive balance between training cost and generative capacity.

Figures (7)

• Figure 1: Conceptual comparison between (a) our SD-DiT and (b) MaskDiT. MaskDiT integrates generative diffusion process with mask reconstruction auxiliary task, and the whole DiT encoder plus decoder are jointly optimized for the two tasks. In contrast, our SD-DiT frames mask modeling on the basis of discrimination knowledge distilling in a self-supervised manner, pursuing the inter-image alignment in the joint embedding space of teacher and student encoder. DiT encoder and decoder are decoupled to separately tackle discriminative and generative diffusion objectives.
• Figure 2: The overview of our SD-DiT. During training, the student view is diffused with regular noise as in EDM formulation, while the teacher view is derived from fixed minimum noise of the consistency function that is close to real data distribution. SD-DiT feeds the discriminative pair into teacher and student DiT encoders to perform self-supervised discriminative process within the joint embedding space. Meanwhile, only the student DiT encoder and DiT decoder undertake the generative diffusion process. At inference, all patches are fed into student branch for sampling.
• Figure 3: Training speed (training steps per second) vs. generative performance (FID-50K score) for our SD-DiT, MDT, DiT, and MaskDiT on 8 $\times$ A100 GPUs. We also label each run with the number of input patches.
• Figure 4: Comparison of convergence speed with SOTA DiT-based approaches in DiT-XL backbone (batch size: 256). The results of DiT and MaskDiT are directly cited from MaskDiT zheng2023fast. Our SD-DiT-XL/2 consistently outperforms DiT-XL/2 and MaskDiT-XL/2 across training steps, leading to better training convergence.
• Figure 5: FID vs. mask ratio on SD-DiT-S/2 with 400k steps.