让扩散模型按需分配算力,动态调整生成质量与速度
One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers
- 用可变长度潜变量序列解耦图像大小与计算量
- 推理时灵活调整潜变量数量,实现算力与质量平衡
- 无需修改主干模型,通用性强适合各类扩散架构
扩散变换器(DiT)虽生成质量高,但计算量固定于图像分辨率,难以实现合理的延迟-质量权衡,且对输入空间标记均匀分配计算资源,浪费在不重要区域。本文提出弹性潜变量接口变换器(ELIT),一种即插即用、兼容DiT的机制,可解耦输入图像尺寸与计算需求。该方法引入一个可学习的变长潜变量序列,标准变换器块在其上操作。轻量级读写交叉注意力层在空间标记与潜变量间传递信息,并优先处理重要区域。通过随机丢弃尾部潜变量进行训练,ELIT学会生成重要性排序的表示:早期潜变量捕捉全局结构,后期包含细节修复信息。推理时可动态调整潜变量数量以匹配算力约束。ELIT设计极简,仅增加两个交叉注意力层,保持修正流目标与DiT堆栈不变。在多个数据集和架构(DiT、U-ViT、HDiT、MM-DiT)上均取得稳定提升。在ImageNet-1K 512px上,FID与FDD得分平均分别提升35.3%和39.6%。
原文摘要 · Abstract (English)
Diffusion transformers (DiTs) achieve high generative quality but lock FLOPs to image resolution, limiting principled latency-quality trade-offs, and allocate computation uniformly across input spatial tokens, wasting resource allocation to unimportant regions. We introduce Elastic Latent Interface Transformer (ELIT), a drop-in, DiT-compatible mechanism that decouples input image size from compute. Our approach inserts a latent interface, a learnable variable-length token sequence on which standard transformer blocks can operate. Lightweight Read and Write cross-attention layers move information between spatial tokens and latents and prioritize important input regions. By training with random dropping of tail latents, ELIT learns to produce importance-ordered representations with earlier latents capturing global structure while later ones contain information to refine details. At inference, the number of latents can be dynamically adjusted to match compute constraints. ELIT is deliberately minimal, adding two cross-attention layers while leaving the rectified flow objective and the DiT stack unchanged. Across datasets and architectures (DiT, U-ViT, HDiT, MM-DiT), ELIT delivers consistent gains. On ImageNet-1K 512px, ELIT delivers an average gain of $35.3\%$ and $39.6\%$ in FID and FDD scores. Project page: https://snap-research.github.io/elit/
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。