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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Zhu, Rui, Pan, Yingwei, Li, Yehao, Yao, Ting, Sun, Zhenglong, Mei, Tao, Chen, Chang Wen
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911812426924032
author Zhu, Rui
Pan, Yingwei
Li, Yehao
Yao, Ting
Sun, Zhenglong
Mei, Tao
Chen, Chang Wen
author_facet Zhu, Rui
Pan, Yingwei
Li, Yehao
Yao, Ting
Sun, Zhenglong
Mei, Tao
Chen, Chang Wen
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2403_17004
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer
Zhu, Rui
Pan, Yingwei
Li, Yehao
Yao, Ting
Sun, Zhenglong
Mei, Tao
Chen, Chang Wen
Computer Vision and Pattern Recognition
Multimedia
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.
title SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer
topic Computer Vision and Pattern Recognition
Multimedia
url https://arxiv.org/abs/2403.17004