Long skip connections improve features in diffusion transformer models
Why DiT Models Underperform as Representation Learners without Long Skip Connections
Computer Vision and Pattern Recognition
Summary
Diffusion Transformer (DiT) models can create impressive images but often produce weaker features needed for recognizing and understanding images. This paper shows that adding long skip connections (LSCs) helps the model keep the important parts of the image separate from noise, improving feature quality. The authors explain that noise in these models contains useful timing information for generation, but mixing it too much with image data reduces clarity. Using LSCs, the noise can be passed directly where needed, making the main processing clearer and the features stronger.
What this means in practice
- •For computer vision engineers: Improve image feature extraction in diffusion transformer models by incorporating long skip connections to better separate noise and image signals.
- •For machine learning platform developers: Enhance unified generation and vision understanding workflows by adopting skip connection designs informed by diffusion noise behavior.
Authors
Benyuan Meng, Yiliang Zhang, Jin-Wen Wu, Qianqian Xu, Qingming Huang, Longtao Huang
Abstract
Diffusion feature, i.e., intermediate activations extracted from diffusion model backbones, has emerged as a promising approach for unifying generation and vision understanding (discrimination). Despite the advancement of diffusion backbones into Diffusion Transformers (DiT), it has been found in previous literature that DiT models tend to yield weaker features despite their stronger generation capability. This poses a threat to good generation-understanding unification. In this work, we show via controlled from-scratch training that long skip connections (LSCs) affect feature quality: a DiT with LSCs can produce better features. Furthermore, through detailed diagnosis, we reveal the mechanism causing this effect. Specifically, the input to a diffusion backbone contains both the latent image and the diffusion noise, where the diffusion noise, though appearing in the form of noises, actually reflects the information of the diffusion trajectory. To make the final prediction of a noisy variable ($v$ or $ε$), the diffusion backbone needs both signals. However, carrying both signals through the main backbone effectively "pollutes" the image signal with noises, leading to poor feature quality. In comparison, LSCs can shortcut diffusion noises directly to where they are needed, leaving much weaker noise signals in the main backbone, thus improving feature quality. To name the broader impact, this reveals that LSCs are not a mere trick but closely related to the diffusion objective. The code is at https://github.com/Darkbblue/dit-lsc.