通过动态采样与跳过模块,实现扩散Transformer的低内存微调。
Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping
- 根据扩散步数动态调整图像块大小,降低显存占用。
- 跳过非关键模块并预计算残差特征,显存减少超60%。
- 适合资源受限设备上的个性化图像生成任务。
扩散Transformer(DiTs)显著提升了文本到图像生成质量,支持高质量个性化内容创作。然而,微调这些模型需要巨大的计算复杂度和内存,限制了在资源受限场景下的实际部署。为此,我们提出名为DiT-BlockSkip的低内存微调框架,结合时间步感知的动态图像块采样和块跳过机制,并通过预计算残差特征实现优化。动态块采样根据扩散时间步调整块大小,再将裁剪后的块缩放到固定低分辨率,从而在高时间步保留全局结构,在低时间步捕捉细粒度细节,同时降低前向与反向传播的内存开销。块跳过机制选择性地微调关键Transformer块,并为跳过的块预计算残差特征,大幅减少训练内存。为识别对个性化至关重要的块,我们引入基于交叉注意力掩码的块选择策略。实验表明,该方法在定性和定量评估中均达到良好个性化性能,同时显著降低内存使用,推动大规模扩散模型在手机、物联网等设备上的本地化可行性。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have significantly enhanced text-to-image (T2I) generation quality, enabling high-quality personalized content creation. However, fine-tuning these models requires substantial computational complexity and memory, limiting practical deployment under resource constraints. To tackle these challenges, we propose a memory-efficient fine-tuning framework called DiT-BlockSkip, integrating timestep-aware dynamic patch sampling and block skipping by precomputing residual features. Our dynamic patch sampling strategy adjusts patch sizes based on the diffusion timestep, then resizes the cropped patches to a fixed lower resolution. This approach reduces forward & backward memory usage while allowing the model to capture global structures at higher timesteps and fine-grained details at lower timesteps. The block skipping mechanism selectively fine-tunes essential transformer blocks and precomputes residual features for the skipped blocks, significantly reducing training memory. To identify vital blocks for personalization, we introduce a block selection strategy based on cross-attention masking. Evaluations demonstrate that our approach achieves competitive personalization performance qualitatively and quantitatively, while reducing memory usage substantially, moving toward on-device feasibility (e.g., smartphones, IoT devices) for large-scale diffusion transformers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。