SD-DiT: Self-supervised Diffusion Transformer

• SD-DiT is a family of Diffusion Transformers that integrates self-supervised discrimination to decouple discriminative and generative objectives.
• It employs a teacher–student framework with EMA updates and an optimal mask ratio (~20%) to eliminate training–inference mismatch.
• The framework achieves state-of-the-art performance in large-scale image generation with fewer training steps and lower FID scores than prior methods.

SD-DiT refers to a family of Diffusion Transformer (DiT) frameworks that integrate self-supervised discrimination into the training of generative diffusion models. SD-DiT models are designed to enhance both convergence speed and generative quality beyond conventional mask-based approaches by decoupling the discriminative and generative objectives within a teacher-student architecture. The resulting frameworks address core inefficiencies of previous DiT models—especially training-inference mismatch and suboptimal objective coupling—and establish state-of-the-art results in terms of fidelity and efficiency for large-scale image generation tasks (Zhu et al., 2024).

1. Motivation and Key Innovations

Diffusion Transformers, which combine iterative denoising with Transformer architectures, exhibit slow convergence under standard training due to the lack of auxiliary signals beyond the conventional denoising loss. Mask-based DiT methods improve convergence by leveraging intra-image mask reconstruction losses. However, these approaches present two fundamental issues:

• Training-inference discrepancy: Mask tokens used during training are absent in inference, inducing a mismatch and lower sample quality.
• Interference of objectives: A single encoder–decoder DiT cannot properly align intra-image (mask reconstruction) and inter-image (denoising) learning objectives.

SD-DiT addresses these by introducing a self-supervised discrimination task that leverages a teacher–student framework. Here, the DiT encoder and decoder are decoupled:

• The encoder pair learns inter-image discrimination without influencing the generative process directly.
• The decoder optimizes the conventional generative diffusion objective. This architecture eliminates the training–inference mismatch and properly separates the representation and generative pathways (Zhu et al., 2024).

2. Model Architecture: Teacher–Student Discriminative Decoupling

SD-DiT models are architected as follows:

• Student Encoder $S_\theta$ receives partially noised and randomly masked images.
• Teacher Encoder $T_{\theta'}$ (an EMA copy of $S_\theta$) processes minimally noised, unmasked versions.
• Discriminative Head: Both encoder outputs are projected into a high-dimensional self-supervised embedding, with cross-entropy alignment performed between student and teacher embeddings.
• Student Decoder $G_\theta$ reconstructs the clean image from an assembly of student encoder embeddings (on visible tokens) and actual noised tokens for masked positions, following the standard diffusion reverse process.

Training workflow:

1. Generate two views from the same image along the Probability Flow ODE (PF-ODE): a highly noised, partially masked "student" input, and a minimally noised, unmasked "teacher" input.
2. Encode both; apply the discriminative cross-entropy on the projected embeddings.
3. The decoder receives a hybrid of visible-token embeddings and direct noised content, reconstructing the signal via the generative loss.
4. Only the student encoder and decoder are updated with gradients; the teacher is updated via EMA.

Inference:

At sampling time, the model reduces to a standard DiT with no mask tokens. Only the student encoder and decoder participate, and all generative pathways are identical to baseline sampling (Zhu et al., 2024).

3. Training Objectives and Formulations

The SD-DiT training objective is the sum of two losses, each acting on a different path:

1. Discriminative Loss ($\mathcal{L}_{\mathrm{disc}}$) For each visible token $i$,

$$\mathcal{L}_D(i) = -\sum_{k=1}^K P_{T_i}[k] \log P_{S_i}[k]$$

where $P_{S_i}$, $P_{T_i}$ are softmax-normalized temperature-scaled logits from the respective encoder heads.

The total discriminative loss sums over all visible tokens and the [CLS] token:

$$\mathcal{L}_{\mathrm{disc}} = \frac{1}{|\{\text{visible}\}|}\sum_{i\in\text{visible}} \mathcal{L}_D(i) + \mathcal{L}_D(\text{[CLS]})$$

2. Generative Loss ($\mathcal{L}_{\mathrm{gen}}$) Standard denoising (score matching), for student decoder:

$$\mathcal{L}_{\mathrm{gen}} = \mathbb{E}_{\mathbf{x}_0,\,\sigma_S,\,n_s}\|\epsilon - \epsilon_{\theta_d}(\mathbf{x}_0+n_s, \sigma_S)\|^2$$

where $\epsilon_{\theta_d}$ denotes the decoder’s prediction at noise level $\sigma_S$ (Zhu et al., 2024).

Total loss:

$$\mathcal{L} = \mathcal{L}_{\mathrm{disc}} + \mathcal{L}_{\mathrm{gen}}$$

Key practical notes:

• Teacher is updated via EMA (momentum 0.996–0.999); discriminative projection heads have fixed temperature schedule; no explicit mask tokens are used in the decoder.
• Mask ratio is optimized (best around 20%) via ablation.

4. Training Protocols and Hyperparameters

• Noise Scheduling:

Student noise $\sigma_S$ sampled log-normally over $[\sigma_{\min}, \sigma_{\max}]$, with mean $-1.2$ and std $1.2$; teacher noise fixed at $\sigma_{\min}=0.002$.

• Projection Head:

Dimensionality $K=8192$, temperatures $\tau_{\rm S}=0.1$, $\tau_{\rm T}$ warmed up from 0.09 to 0.099.

• Batch size: 256 on 8×A100 GPUs.
• Optimizer: Follows DiT baseline training (AdamW), equal weighting of loss terms.
• Inference: No mask tokens or teacher branch; standard DiT sampling loop.

5. Comparative Performance and Experimental Results

SD-DiT demonstrates significant improvements over baseline and mask-based DiT variants, both in efficiency and generation quality.

Model	Steps (k)	FID$_{50\text{K}}$ ↓	Speed (steps/s, S/2 backbone)
DiT-S/2	400	68.40	5.03
MDT-S/2	400	53.46	2.40
MaskDiT-S/2	600	50.30	9.47
SD-DiT-S/2	400	48.39	9.20

On DiT-XL/2 backbone:

• SD-DiT-XL/2 at 1300k steps achieves FID = 9.01, compared to DiT-XL/2 at 7000k steps (FID = 9.62) and MaskDiT-XL/2 at 2000k steps (FID = 12.15), indicating a 5× reduction in required update steps (Zhu et al., 2024).

Key ablations:

• Removing $\mathcal{L}_{\mathrm{disc}}$ increases FID from 53.72 (full) to 62.84.
• No masking (mask ratio = 0) increases FID to 58.92.
• Optimal mask ratio is ∼20% (FID = 48.4).
• Fixed teacher noise outperforms variable teacher noise.

In SOTA comparison (ImageNet-256):

• SD-DiT-XL/2 (FID = 7.21) outperforms DiT-XL/2 (FID = 9.62) and MaskDiT-XL/2 (FID = 5.69).

6. Practical Implications and Extensions

The SD-DiT framework:

• Eliminates the need for learnable mask tokens at inference, thus resolving training-inference mismatch.
• Enables better utilization of large batches and reduced training times without sacrificing generative performance.
• Facilitates the extension of the self-supervised discrimination paradigm to other diffusion backbones, including U-ViT.
• Is compatible with additional generative enhancements, such as classifier-free guidance (reported FID ≈ 3.2 with CFG).

Recommendations:

• Adopt mask ratio ≈ 20%; set teacher noise to the lowest possible value; update teacher parameters with slow EMA; decouple encoder/decoder to follow the described training objective (Zhu et al., 2024).

Potential extensions include longer schedules, larger models, improved projection head architectures, and applications to modalities beyond image generation.

• "SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer" (Zhu et al., 2024)