arXiv:2608.01298cs.CVcs.AI2026-08

提出UDT架构,让扩散Transformer兼具U-Net的编码解码优势和数据自适应的高效下采样。

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

论文配图:UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction
图 1 · 摘自论文原文
  • 用数据自适应的令牌合并实现下采样与上采样,保持令牌维度不变
  • 在256x256 ImageNet上40轮内达到比SiT快40倍的收敛速度(FID 7.9)
  • 结合正则化后生成质量达FID 1.35,适合高效率图像生成任务

扩散Transformer(DiTs)因其可扩展性和多模态适应性成为生成建模的核心架构。但其各层逐步学习导致后期关注细节重建,使表征质量下降且编码-解码行为失衡。现有方法如表示对齐(REPA)通过训练正则化增强早期表征,而基于U-Net结构的DiT则引入显式多尺度编码-解码机制,但依赖可学习的空间下采样算子,不适用于变压器架构,带来效率低下和与交叉注意力等组件的兼容性问题。本文提出UDT:一种结合DiT表征能力与U-Net编码-解码优势的扩散变压器,通过数据自适应的令牌合并实现下采样与上采样,同时保持令牌维度。基线模型优于现有U-Net DiT,在所有模型规模下性能接近REPA。进一步结合架构优化与REPA,UDT在256x256 ImageNet上以XL模型规模仅用40轮即达成与SiT 7.9 FID(无分类引导)相当的结果(快40倍),并使用CFG生成时取得FID 1.38(320轮,SD-VAE)与1.35(500轮,VA-VAE),为扩散变压器提供新的强大骨干。

原文摘要 · Abstract (English)

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.

扩散模型图像生成TransformerU-Net

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。