A-SelecT: Automatic Timestep Selection for Diffusion Transformer Representation Learning
Authors: Changyu Liu, James Chenhao Liang, Wenhao Yang, Yiming Cui, Jinghao Yang, Tianyang Wang, Qifan Wang, Dongfang Liu, Cheng Han
Organizations: University of Missouri–Kansas City · U. S. Naval Research Laboratory · Lamar University · University of Florida · University of Texas Rio Grande Valley · University of Alabama at Birmingham · Meta AI · Rochester Institute of Technology
Diffusion models have significantly reshaped the field of generative artificial intelligence and are now increasingly explored for their capacity in discriminative representation learning. Diffusion Transformer (DiT) has recently gained attention as a promising alternative to conventional U-Net-based diffusion models, demonstrating a promising avenue for downstream discriminative tasks via generative pre-training. However, its current training efficiency and representational capacity remain largely constrained due to the inadequate timestep searching and insufficient exploitation of DiT-specific feature representations. In light of this view, we introduce Automatically Selected Timestep (A-SelecT) that dynamically pinpoints DiT's most information-rich timestep from the selected transformer feature in a single run, eliminating the need for both computationally intensive exhaustive timestep searching and suboptimal discriminative feature selection. Extensive experiments on classification and segmentation benchmarks demonstrate that DiT, empowered by A-SelecT, surpasses all prior diffusion-based attempts efficiently and effectively.
Figures & tables
Figure 1: A Preliminary Study on the impact of the High-Frequency Ratio (HFR) with classification performance on Oxford Flowers (a) and CUB (b). The green curve represents the HFR values, and the red curve are classification accuracies. We have two key observations: I. HFR values exhibit a positive correlation with classification accuracies. II. The highest classification accuracy is achieved when the HFR value reaches its maximum. More results in Appendix §S2
Figure 2: Overview of A utomatically Selec ted T imestep (A-SelecT). A-SelecT begins by simulating samplet at timestep t (see Eq. 1 ). This samplet is then processed through the diffusion backbone to extract the query feature Qt at each timestep t . Upon obtaining all Qt , their HFR are calculated. The timestep exhibiting the highest average HFR is subsequently selected for feature extraction. The query feature extracted at this optimal timestep t^ is then fed into targeted discriminative tasks.
Figure 3: Visualizations of High-frequency vs. Low-frequency Information. We present a decomposition of the original features extracted from SD 3.5 into components that exclusively contain high-frequency and low-frequency information. The second column is the original features extracted from the model. As seen, the high-frequency features are shown to contain more discriminative information ( i.e . , edge, texture, corner information from the black footed albatross can be clearly preserved) than their low-frequency counterparts. Inspired by this, we design HFR to assess the significance of high-frequency information (see § 3.3 ).
Figure 4: Comparison of HFR and Fisher Score across time- steps on Oxford Flowers (top) and CUB (bottom). They show strong alignment, indicating that HFR captures discriminative characteristics consistent with the Fisher Score and serves as a reliable label-free indicator of feature separability.
Figure 5: Impact of Feature and Block Selection. We present accuracy across the features Q , K , V , A , and O extracted from different transformer blocks on Oxford Flowers. Q and V achieve the highest accuracies (90.6% and 88.7%, respectively), while A , O , and K show comparatively lower performance. The middle transformer blocks yield the most discriminative representations, highlighting the importance of both feature and block selection for optimal performance. Additional experimental results on other datasets are provided in the Appendix §S7.
Figure 6: Feature Visualization from CUB and Oxford Flowers. Visualizations of four feature sets from the 9- th block of SD 3.5. Each column represents Q feature extracted at timestep t∈{1,50,100,150,200,250} . Notably, the differences among features from various timesteps are subtle and challenging to discern visually, showing that manually selecting the discriminative feature extraction is impractical. More examples in Appendix §S8.
Table 5: Discussions on A-SelecT w.r.t feature visualization and traversal search further confirm the significance of our approach, showing outstanding automatic discriminative feature selection and training efficiency.
Figure S1 : More Preliminary Results on the impact of the High-Frequency Ratio (HFR) with classification performance on Stanford Cars (a) and Aircraft (b).
Figure S2 : More Visualization of Feature Decomposition Examples from CUB, Oxford Flowers and Stanford Dogs . We present six sets of decomposition of the original features extracted from SD 3.5 into components that exclusively contain high-frequency and low-frequency information.
Figure S3 : More Feature Visualization Examples from CUB and Oxford Flowers. We present visualizations of three feature sets extracted from the 9- th block of SD 3.5. The columns represent features extracted at different timesteps.
Figure S4 : Impact of Feature and Block Selection. We show more block and feature performance on CUB, Aircraft, and Stanford Car. The figure shows the consistent result with Oxford Flower dataset.
Figure S5 : Feature Visualization Across Blocks on CUB and Oxford Flowers. The first column shows original input images and subsequent columns presents feature visualizations extracted from the 0- th to 23- th block of SD 3.5. The top three rows correspond to samples from CUB and the bottom three rows are from Oxford Flowers.
Figure S6 : More Comparison of HFR and Fisher Score across timesteps on Aircraft (top) and Stanford Cars (bottom). The results show that HFR and Fisher Score exhibit consistent trends.
Figure S7 : Comparison of HFR and Classification Accuracy across multiple DiT Models on Oxford Flowers. The alignment of peak accuracy with maximum HFR values demonstrates the consistent generalization of HFR.
Visual representation learning develops image and video features that capture useful structure for recognition, retrieval, spatial reasoning, and action.