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

• URL: http://arxiv.org/abs/2608.01298v1
• Date: Sun, 02 Aug 2026 15:07:21 GMT
• Title: UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction
• Abstract summary: Diffusion Transformers (DiTs) have emerged as a core architecture in generative modeling.<n>We propose UDT, a U-Net diffusion transformer that combines the representation power of DiTs with the encoding-decoding benefits of U-Nets.
• Score: 4.511561231517167
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• 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.

Related papers

• MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model [6.401952350014671]
  Diffusion Transformers (DiTs) have become widely used in diffusion and flow-matching models.<n>DiTs process the same number of patchified tokens in every block, leading to relatively heavy computation during training process.<n>We introduce a multi-patch transformer design in which early blocks operate on larger patches to capture coarse global context, while later blocks use smaller patches to refine local details.
  arXiv  Detail & Related papers  (2026-03-27T12:30:10Z)
• Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders [74.72147962028265]
  Representation Autoencoders (RAEs) have shown distinct advantages in diffusion modeling on ImageNet.<n>We investigate whether this framework can scale to large-scale, freeform text-to-image (T2I) generation.
  arXiv  Detail & Related papers  (2026-01-22T18:58:16Z)
• Diffusion Transformers with Representation Autoencoders [35.43400861279246]
  Latent generative modeling, where a pretrained autoencoder maps pixels into a latent space for the diffusion process, has become the standard strategy for Diffusion Transformers (DiT)<n>Most DiTs continue to rely on the original VAE encoder, which introduces several limitations.<n>In this work, we explore replacing the VAE with pretrained representation encoders paired with trained decoders, forming what we term Representation Autoencoders (RAE)
  arXiv  Detail & Related papers  (2025-10-13T17:51:39Z)
• DDT: Decoupled Diffusion Transformer [51.84206763079382]
  Diffusion transformers encode noisy inputs to extract semantic component and decode higher frequency with identical modules.<n>textbfcolorddtDecoupled textbfcolorddtTransformer(textbfcolorddtDDT)<n>textbfcolorddtTransformer(textbfcolorddtDDT)<n>textbfcolorddtTransformer(textbfcolorddtDDT)
  arXiv  Detail & Related papers  (2025-04-08T07:17:45Z)
• Dynamic Diffusion Transformer [67.13876021157887]
  Diffusion Transformer (DiT) has demonstrated superior performance but suffers from substantial computational costs. We propose Dynamic Diffusion Transformer (DyDiT), an architecture that dynamically adjusts its computation along both timestep and spatial dimensions during generation. With 3% additional fine-tuning, our method reduces the FLOPs of DiT-XL by 51%, accelerates generation by 1.73, and achieves a competitive FID score of 2.07 on ImageNet.
  arXiv  Detail & Related papers  (2024-10-04T14:14:28Z)
• U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers [28.936553798624136]
  We propose a series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the extraordinary performance of U-DiT models. The proposed U-DiT could outperform DiT-XL/2 with only 1/6 of its cost computation.
  arXiv  Detail & Related papers  (2024-05-04T18:27:29Z)
• TCCT-Net: Two-Stream Network Architecture for Fast and Efficient Engagement Estimation via Behavioral Feature Signals [58.865901821451295]
  We present a novel two-stream feature fusion "Tensor-Convolution and Convolution-Transformer Network" (TCCT-Net) architecture. To better learn the meaningful patterns in the temporal-spatial domain, we design a "CT" stream that integrates a hybrid convolutional-transformer. In parallel, to efficiently extract rich patterns from the temporal-frequency domain, we introduce a "TC" stream that uses Continuous Wavelet Transform (CWT) to represent information in a 2D tensor form.
  arXiv  Detail & Related papers  (2024-04-15T06:01:48Z)
• DiffiT: Diffusion Vision Transformers for Image Generation [88.08529836125399]
  Vision Transformer (ViT) has demonstrated strong modeling capabilities and scalability, especially for recognition tasks. We study the effectiveness of ViTs in diffusion-based generative learning and propose a new model denoted as Diffusion Vision Transformers (DiffiT) DiffiT is surprisingly effective in generating high-fidelity images with significantly better parameter efficiency.
  arXiv  Detail & Related papers  (2023-12-04T18:57:01Z)
• AICT: An Adaptive Image Compression Transformer [18.05997169440533]
  We propose a more straightforward yet effective Tranformer-based channel-wise auto-regressive prior model, resulting in an absolute image compression transformer (ICT) The proposed ICT can capture both global and local contexts from the latent representations. We leverage a learnable scaling module with a sandwich ConvNeXt-based pre/post-processor to accurately extract more compact latent representation.
  arXiv  Detail & Related papers  (2023-07-12T11:32:02Z)
• Rich CNN-Transformer Feature Aggregation Networks for Super-Resolution [50.10987776141901]
  Recent vision transformers along with self-attention have achieved promising results on various computer vision tasks. We introduce an effective hybrid architecture for super-resolution (SR) tasks, which leverages local features from CNNs and long-range dependencies captured by transformers. Our proposed method achieves state-of-the-art SR results on numerous benchmark datasets.
  arXiv  Detail & Related papers  (2022-03-15T06:52:25Z)
• Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers [149.78470371525754]
  We treat semantic segmentation as a sequence-to-sequence prediction task. Specifically, we deploy a pure transformer to encode an image as a sequence of patches. With the global context modeled in every layer of the transformer, this encoder can be combined with a simple decoder to provide a powerful segmentation model, termed SEgmentation TRansformer (SETR) SETR achieves new state of the art on ADE20K (50.28% mIoU), Pascal Context (55.83% mIoU) and competitive results on Cityscapes.
  arXiv  Detail & Related papers  (2020-12-31T18:55:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.