提出TreeQ框架，实现扩散Transformer的4比特高效量化部署。
TreeQ: Pushing the Quantization Boundary of Diffusion Transformer via Tree-Structured Mixed-Precision Search
- 采用树形搜索机制，快速定位最优混合精度配置。
- 首次实现扩散Transformer在4比特下的近无损性能。
- 适合需要低资源部署图像生成模型的研究与开发者。
扩散Transformer（DiT）作为图像生成的高效骨干网络，具有优异的可扩展性与性能，但其实际部署受制于高计算与内存需求。混合精度量化（MPQ）在提升U-Net模型压缩效果方面表现卓越，已实现亚4比特量化，但对DiT架构的应用仍不充分。本文提出TreeQ框架，针对该问题：首先，设计树状结构搜索（TSS），利用DiT的线性特性以O(n)时间遍历解空间，并通过对比剪枝提高目标准确性；其次，提出环境噪声引导（ENG），仅用一个超参数统一后训练量化（PTQ）与量化感知训练（QAT）配置；最后，引入通用雄鹰分支（GMB），通过结构化稀疏设计缓解超低比特下的信息瓶颈。大量实验表明，TreeQ在DiT-XL/2上实现了W3A3和W4A4 PTQ/PEFT设置下的最先进性能，首次达成扩散模型4比特下的近无损量化效果。代码与模型将公开于https://github.com/racoonykc/TreeQ。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have emerged as a highly scalable and effective backbone for image generation, outperforming U-Net architectures in both scalability and performance. However, their real-world deployment remains challenging due to high computational and memory demands. Mixed-Precision Quantization (MPQ), designed to push the limits of quantization, has demonstrated remarkable success in advancing U-Net quantization to sub-4bit settings while significantly reducing computational and memory overhead. Nevertheless, its application to DiT architectures remains limited and underexplored. In this work, we propose TreeQ, a unified framework addressing key challenges in DiT quantization. First, to tackle inefficient search and proxy misalignment, we introduce Tree Structured Search (TSS). This DiT-specific approach leverages the architecture's linear properties to traverse the solution space in O(n) time while improving objective accuracy through comparison-based pruning. Second, to unify optimization objectives, we propose Environmental Noise Guidance (ENG), which aligns Post-Training Quantization (PTQ) and Quantization-Aware Training (QAT) configurations using a single hyperparameter. Third, to mitigate information bottlenecks in ultra-low-bit regimes, we design the General Monarch Branch (GMB). This structured sparse branch prevents irreversible information loss, enabling finer detail generation. Through extensive experiments, our TreeQ framework demonstrates state-of-the-art performance on DiT-XL/2 under W3A3 and W4A4 PTQ/PEFT settings. Notably, our work is the first to achieve near-lossless 4-bit PTQ performance on DiT models. The code and models will be available at https://github.com/racoonykc/TreeQ
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。