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.
Diffusion Transformers (DiTs) have emerged as a core architecture in generative modeling due to their scalability and adaptability to multimodal tasks. DiTs comprise isotropic transformer blocks, and learn representations progressively across depth, where the denoising objective drives later layers to focus on fine-detail reconstruction. This results in degraded representation quality and an imbalanced encoder-decoder behavior. Prior approaches such as representation alignment (REPA) mitigate this by encouraging stronger early representations via training regularization. Alternatively, U-Net-style DiT architectures introduce explicit multi-scale encoder-decoder structures for improved convergence. But they build on standard U-Net wisdom via learnable operators for spatial downsampling, which are not well-suited to transformer architectures, introducing inefficiencies and compatibility issues with components such as cross-attention and representation regularization. In this work, we propose UDT, a U-Net diffusion transformer that combines the representation power of DiTs with the encoding-decoding benefits of U-Nets, through data-adaptive token merging for downsampling and upsampling, while preserving the DiT token dimension. Our baseline UDT architecture outperforms existing U-Net DiTs and achieves performance comparable to REPA across all model sizes. Furthermore, using architectural optimization and REPA, UDT outperforms SiT's 7.9 FID at 1400 epochs (w/o CFG) within 40 epochs (~ 40x faster convergence) for XL model size on 256x256 ImageNet. Finally, it achieves strong image generation performance with CFG, reaching FID of 1.38 (320 epochs) with SD-VAE and 1.35 (500 epochs) with VA-VAE, providing a new backbone for DiTs with strong empirical benefits.
Diffusion Transformers (DiTs) are a dominant backbone for high-fidelity text-to-image generation due to strong scalability and alignment at high resolutions. However, quadratic self-attention over dense spatial tokens leads to high inference latency and limits deployment. We observe that denoising is spatially non-uniform with respect to aesthetic descriptors in the prompt. Regions associated with aesthetic tokens receive concentrated cross-attention and show larger temporal variation, while low-affinity regions evolve smoothly with redundant computation. Based on this insight, we propose AccelAes, a training-free framework that accelerates DiTs through aesthetics-aware spatio-temporal reduction while improving perceptual aesthetics. AccelAes builds AesMask, a one-shot aesthetic focus mask derived from prompt semantics and cross-attention signals. When localized computation is feasible, SkipSparse reallocates computation and guidance to masked regions. We further reduce temporal redundancy using a lightweight step-level prediction cache that periodically replaces full Transformer evaluations. Experiments on representative DiT families show consistent acceleration and improved aesthetics-oriented quality. On Lumina-Next, AccelAes achieves a 2.11× speedup and improves ImageReward by +11.9% over the dense baseline. Code is available at https://github.com/xuanhuayin/AccelAes.
Xuanhua Yin, Chuanzhi Xu, Haoxian Zhou +2
School of Computer Science, The University of Sydney, NSW 2006, Australia
Diffusion Transformers (DiTs) have achieved state-of-the-art performance in image and video generation, but their success comes at the cost of heavy computation. This inefficiency is largely due to the fixed tokenization process, which uses constant-sized patches throughout the entire denoising phase, regardless of the content's complexity. We propose dynamic tokenization, an efficient test-time strategy that varies patch sizes based on content complexity and the denoising timestep. Our key insight is that early timesteps only require coarser patches to model global structure, while later iterations demand finer (smaller-sized) patches to refine local details. During inference, our method dynamically reallocates patch sizes across denoising steps for image and video generation and substantially reduces cost while preserving perceptual generation quality. Extensive experiments demonstrate the effectiveness of our approach: it achieves up to 3.52× and 3.2× speedup on FLUX-1.Dev and Wan 2.1, respectively, without compromising the generation quality and prompt adherence.