Contrastive regularization improves diffusion transformer training with new anchoring method
Unlocking the Potential of Contrastive Regularization for Diffusion Transformer Training
Computer Vision and Pattern RecognitionMachine Learning
Summary
Training diffusion transformer models to generate images can be improved by a technique called contrastive regularization, but usual methods don't help much because of noisy data tricks. The authors found that standard contrastive learning lets the model cheat by focusing on noise rather than true signal. To fix this, they introduced a new method called CRAN that uses clean, noise-free examples to guide the model’s learning. Their experiments show that CRAN helps the model learn better representations and improves performance on a big image dataset.
What this means in practice
- •For machine learning engineers: Improve training of diffusion transformers for image generation by integrating CRAN contrastive regularization to get better model performance.
- •For computer vision developers: Enhance diffusion-based image generation tools by reducing noise shortcuts with CRAN, leading to cleaner and more accurate outputs.
Authors
Yi Tan, Beier Zhu, Sirui Chen, Yunzhe Lin, Yuan Zhou, Cheng Chen, Xingyu Zhu, Tianrun Chen, Yanbin Hao, Hanwang Zhang
Abstract
Representation regularization has emerged as an effective way to improve diffusion transformer training. Yet one of the most canonical regularizers--InfoNCE-style contrastive learning--remains largely underexplored in this setting. We naturally ask a question: can standard contrastive learning serve as an effective regularizer for diffusion transformers? Surprisingly, we find that naive InfoNCE regularization brings only marginal gains. Our analysis reveals that the objective is easily dominated by a shortcut: near-orthogonal noise components in diffusion inputs provide a bypass for enlarging negative-pair distances while weakening positive-pair alignment, especially at high noise levels. To unlock the benefit of contrastive regularization, we propose Clean Representation-ANchored Contrastive (CRAN) regularization. Specifically, CRAN anchors contrastive dispersion with clean-sample representation distances, transferring clean relational semantic structure to noisy diffusion representations to suppress the noise shortcut. Extensive experiments on ImageNet show that CRAN consistently improves diffusion transformer training across model scales and training budgets, outperforming internal feature-alignment and dispersive regularization baselines while matching external DINO-guided REPA. These results re-establish InfoNCE-based contrastive learning as an effective regularizer for diffusion training, once its noise-driven over-dispersion is properly constrained.