Skip to content
Preprint

UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction

Aug 2026 · 0 citations · 62 references
Computer Science

TL;DR

UDT is proposed, 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.

Abstract

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.

View source

Similar papers

Preprint Aug 2026

Semantic Steering for Controllable Generation: Tuning-Free Concept Erasure in Multimodal Diffusion Transformers

This work proposes to erase concepts by directly manipulating the model's internal representations by operating exclusively on the sparse text-branch tokens and leveraging the straight sampling trajectory of rectified flow, achieving effective concept erasure with negligible overhead and without any training.

Qiao Li, Xiaomeng Fu, Yuanshu Zhao et al. · 0 citations
Preprint Aug 2026

DiverseDiT++: Quantifying, Analyzing, and Promoting Representation Diversity in Diffusion Transformers

Recent advances in Diffusion Transformers (DiTs) have enabled remarkable progress in visual synthesis, benefiting from their superior scalability. To facilitate DiTs'capability of capturing meaningful internal representations, recent works such as REPA incorporate external pretrained encoders for representation alignment. However, the underlying mechanisms governing representation learning within DiTs remain poorly understood in the community. To this end, this paper first presents a systematic analysis of the representation dynamics of DiTs via quantifying the diversity of block-wise representations. Specifically, we introduce a novel metric, termed the Weighted Diversity Score (WDS), to measure the representational discrepancies across different blocks. Through extensive investigations on the evolution and influence of internal representations under various settings, we reveal that representation diversity across blocks is a critical factor for effective representation learning in DiTs. More importantly, WDS exhibits a strong correlation with synthesis quality across diverse settings, model scales, and training stages (Pearson's $r=-0.869$ with $\log(\text{FID})$), suggesting its potential as an indicator to reflect model performance and a principled guide for model optimization. Based on this key finding, we propose DiverseDiT++, a novel framework that explicitly promotes diverse representation learning. Concretely, our method incorporates long residual connections to diversify input representations across blocks and a representation diversity loss to encourage blocks to learn distinct features. Extensive experiments on ImageNet $256\times256$ and $512\times512$ demonstrate that our DiverseDiT++ yields consistent performance gains and convergence acceleration when applied to different backbones with various sizes,...

Binglei Li, Mengping Yang, Zhiyu Tan et al. · 0 citations
Review Jul 2026

Pixel-Space Diffusion Transformers

Latent diffusion models (LDMs) enable efficient high-resolution image synthesis by denoising in a VAE-compressed latent space. However, fixed visual tokenizers can discard fine textures and structural details, while separate representation and diffusion training creates a mismatch between reconstruction and generation objectives. These limitations have renewed interest in pixel-space diffusion, which models raw pixels directly, removes the VAE bottleneck, and supports end-to-end optimization. This formulation better matches the demands of high-fidelity generation but introduces challenges in high-dimensional modeling, including noise scheduling, loss weighting, token efficiency, and scalable architecture design. Pixel-space modeling also offers a promising basis for unified multimodal systems: raw pixels, text, and task conditions can be represented in a shared token space and jointly processed by a single Transformer, narrowing the gap between visual understanding and generation. This paper reviews Pixel-Space Diffusion Transformers (pDiTs) from the perspectives of model architecture, continuous generative mechanisms, and unified multimodal modeling. We summarize representative methods, identify key technical challenges, and discuss future directions toward high-fidelity, end-to-end vision foundation models that integrate generation and understanding.

Renye Yan, Jikang Cheng, You Wu et al. · 3 citations
Open access Aug 2026

GAESA-iFormer: a graph attention-enhanced scale-aware inverted transformer for performance characterization of aircraft engine compressors

Accurate determination of aircraft engine compressor characteristics is vital for stability and performance optimization. However, experimental data acquisition is often hindered by high costs, resulting in a strong reliance on limited test data. To address this challenge, this paper proposes the Graph Attention Enhanced Scale-Aware Inverted Transformer (GAESA-iFormer), a novel surrogate modeling framework integrating a Graph Attention Network (GAT) encoder, a Multi-scale Feature Enhancement (MSFE) module, and a cross-attention decoder based on the inverted Transformer architecture. The GAT encoder explicitly models relational dependencies among neighboring operating points along each speed line, capturing local continuity while reducing computational complexity. The MSFE module employs parallel convolutional kernels of sizes 1, 3, and 5 to extract multi-granularity features, effectively broadening the receptive field. Furthermore, the cross-attention decoder enables explicit knowledge transfer across different speed lines, leveraging structural similarities between well-sampled and sparsely sampled speeds—a critical advantage when data is scarce. Comprehensive evaluations demonstrate that GAESA-iFormer achieves optimal performance with a feature dimension of 64 and four encoder layers. When trained on transformed secondary data, the proposed model significantly outperforms state-of-the-art baselines, reducing the RMSE by 27.03% and MAE by 31.78%. The results indicate that the proposed GAESA-iFormer model achieves improved prediction accuracy on the current compressor dataset under limited experimental data.

Xiang Gao, Zhanxue Wang, Zhenyu Tao · 0 citations