通过动态图像块采样与块跳过实现扩散 Transformer 的显存高效微调
Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping
面向 DiT 个性化微调,按扩散时间步动态调整图像块尺寸并跳过非关键 Transformer 块、预计算残差特征,可直接迁移到显存受限下的 T2I 定制与少样本适配。
原文依据 integrating timestep-aware dynamic patch sampling and block skipping by precomputing residual features
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论扩散 Transformer(DiT)显著提升了文本到图像(T2I)生成质量,使高质量个性化内容创作成为可能。然而,微调这些模型需要大量计算复杂度与显存,限制了其在资源受限条件下的实际部署。为应对这些挑战,我们提出了一种显存高效的微调框架 DiT-BlockSkip,通过预计算残差特征将时间步感知的动态图像块采样与块跳过相结合。我们的动态图像块采样策略根据扩散时间步调整图像块尺寸,再将裁剪出的图像块缩放至固定的较低分辨率。该方法降低了前向与反向传播的显存占用,同时使模型能够在较高时间步捕获全局结构、在较低时间步捕获细粒度细节。块跳过机制选择性地微调关键的 Transformer 块,并为被跳过的块预计算残差特征,从而显著降低训练显存。为识别对个性化至关重要的块,我们引入了基于交叉注意力掩码的块选择策略。评估表明,我们的方法在定性与定量上均取得有竞争力的个性化性能,同时大幅降低显存占用,朝着大规模扩散 Transformer 在设备端(如智能手机、物联网设备)可行部署迈进一步。
英文原文摘要 1406 字符
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.