Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping

Each language version is independently generated for its own context, not a direct translation.

这篇论文介绍了一种名为 DiT-BlockSkip 的新方法，旨在解决一个非常棘手的问题：如何在手机或普通电脑上，用很少的内存，把强大的 AI 绘画模型（Diffusion Transformers）“教会”画特定的东西（比如画你自家的猫，或者你设计的独特玩偶）。

为了让你更容易理解，我们可以把训练 AI 模型想象成**“教一个超级大厨做一道新菜”**。

1. 背景：为什么现在很难“教”AI？

现状：现在的 AI 绘画模型（如 FLUX）非常强大，像是一个拥有整个图书馆食谱的超级大厨。如果你想让它学会画“你家的猫”，通常的做法是给它看几张家猫的照片，让它重新学习。
问题：这个“重新学习”的过程（微调）非常烧钱、烧电、烧内存。
- 这就好比，为了教大厨做一道新菜，你不得不把整个图书馆的食谱全部搬进厨房，还要把每一页都翻一遍、记一遍。
- 结果就是：只有拥有昂贵显卡（大厨房）的实验室才能做这件事，普通人的手机（小厨房）根本放不下，一运行就“内存爆炸”。

2. 核心方案：DiT-BlockSkip（两步走策略）

为了解决这个问题，作者提出了两个聪明的“偷懒”技巧，既省内存，又不让大厨变笨。

技巧一：动态切片采样 (Dynamic Patch Sampling)

——“看山看水，远近有别”

传统做法：无论教什么，都一直盯着高清大图看。这就像教大厨做菜时，无论看哪一步，都要用 4K 显微镜看食材纹理，非常费眼（费内存）。
新方法：根据**“噪音程度”（可以理解为学习的阶段）来调整看图的“焦距”**。
- 刚开始学（高噪音/大步骤）：这时候画面很模糊，看不清细节。这时候我们把视野拉远，只看大块的色块和整体轮廓（比如先看清这是只猫，而不是先看清猫毛）。这就好比用广角镜头看全景，不需要太高的分辨率，省内存。
- 快学完时（低噪音/小步骤）：这时候画面清晰了，需要画细节。这时候我们把视野拉近，只看局部特写（比如猫耳朵的绒毛）。这就好比用微距镜头看细节。
比喻：就像你教人画画，先让他看整体构图（不用太清晰），最后再让他画眼睛和胡须（需要清晰）。这样既不用一直盯着高清大图，又能学会所有东西。

技巧二：块跳过与残差预计算 (Block Skipping with Residual Feature)

——“只记重点，跳过废话”

传统做法：AI 模型由很多层“大脑皮层”（Transformer 块）组成。传统微调时，每一层都要重新计算、重新记忆，非常累。
新方法：
1. 找重点：作者发现，并不是每一层大脑都对“画猫”这件事重要。有些层只负责画背景，有些层负责画猫。通过一种“注意力测试”，他们发现中间层对画猫最关键，而头尾的层相对不那么重要。
2. 跳过废话：在训练时，直接跳过那些不重要的层（头尾），只训练中间的关键层。
3. 预存“答案”：既然跳过了，那被跳过的层原本该产生的变化怎么办？作者想了一个绝招：提前算好。
  - 在正式训练前，先算出被跳过的那些层“输入”和“输出”的差值（就像提前算好这道菜里盐放了多少，不用每次都重新称）。
  - 正式训练时，直接把算好的“差值”加上去，就不用让那些层再工作一次了。
比喻：就像你背课文，有些段落（头尾）是通用的，不用每次都背；只有中间那段（关键情节）需要重点背。而且，你提前把通用段落的“总结”写在便签上，复习时直接看便签，不用把整本书再读一遍。

3. 效果如何？

省内存：这种方法能把训练时的内存占用减少 46% 到 65%。
- 比喻：原本需要 22GB 内存（像一个大仓库）才能跑的训练，现在 11GB（像个小衣柜）就能搞定。这意味着普通的高端显卡甚至未来的手机都能跑起来。
质量没降：虽然“偷懒”了，但画出来的猫依然很像，文字描述（比如“在森林里”）也能准确执行。
速度更快：因为计算量少了，训练时间也缩短了。

4. 总结：这对我们意味着什么？

这篇论文就像给 AI 绘画领域装上了一个**“节能模式”**。

以前，你想定制一个专属的 AI 画师，得去租用昂贵的云端服务器。现在，有了 DiT-BlockSkip，未来你完全有可能在自己的手机上，花几分钟时间，上传几张自家宠物的照片，就训练出一个能画出你家宠物各种姿势的专属 AI，而且手机不会发烫、不会卡顿。

一句话总结：通过**“看远看近灵活切换”和“只背重点、提前存答案”**，让 AI 模型在极小的内存下也能学会画任何你想画的东西，让个性化 AI 从“实验室”走向“你的口袋”。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与问题 (Problem)

随着基于扩散 Transformer (DiT, Diffusion Transformer) 的文本到图像 (T2I) 模型（如 FLUX、SANA）的出现，图像生成质量显著提升。然而，在这些大规模模型上进行个性化微调 (Personalization) 面临着巨大的挑战：

资源消耗巨大：微调 DiT 模型需要极高的计算复杂度和显存（VRAM），导致难以在资源受限的设备（如智能手机、IoT 设备）上部署。
现有方法的局限性：
- 参数高效微调 (PEFT)（如 LoRA）：虽然冻结了基座模型权重，但仍需对整个网络进行反向传播，导致激活值和梯度的显存占用依然很高。
- 无梯度/零阶优化方法：虽然减少了显存，但往往存在优化不稳定、收敛慢（需要数万次迭代）的问题。
- 基于 U-Net 的现有优化：现有的高效微调技术（如 HollowedNet）主要针对 U-Net 架构，直接应用于 DiT 架构时效果不佳，因为 DiT 的块功能分布与 U-Net 不同。

核心目标：在保持个性化生成质量的同时，显著降低 DiT 模型的微调显存占用，使其具备在边缘设备部署的可行性。

2. 方法论 (Methodology)

作者提出了名为 DiT-BlockSkip 的框架，包含两个核心组件：动态 Patch 采样 (Dynamic Patch Sampling) 和 带残差特征预计算的块跳过 (Block Skipping with Residual Feature Precomputation)。

2.1 动态 Patch 采样 (Dynamic Patch Sampling)

原理：利用扩散过程中的时间步 (timestep) 特性。
- 高时间步 (High Timesteps)：图像噪声大，主要学习全局结构。此时采样较大的 Patch（覆盖更广区域）。
- 低时间步 (Low Timesteps)：图像噪声小，主要学习细粒度细节。此时采样较小的 Patch（聚焦局部）。
实现：
1. 根据当前时间步 $t$ 动态调整裁剪的 Patch 大小。
2. 将不同大小的 Patch 统一 Resize 到一个固定的低分辨率（如 $256 \times 256$ ）输入模型。
优势：在不改变模型架构和去噪目标的前提下，显著减少了前向和反向传播的显存占用，同时让模型能从低分辨率输入中同时学习全局结构和局部细节。

2.2 带残差特征预计算的块跳过 (Block Skipping with Residual Feature Precomputation)

核心思想：在微调过程中，只更新对个性化至关重要的 Transformer 块，跳过其余块的计算，但需解决跳过导致的特征不匹配问题。
步骤：
1. 关键块选择策略 (Block Selection Strategy)：
  - 通过交叉注意力掩码 (Cross-Attention Masking) 实验发现，DiT 模型中的中间层块对主体信息的编码最为关键，而首尾层块影响较小。
  - 基于此，提出算法自动选择跳过首部的 $n$ 个块和尾部的 $m$ 个块，保留中间的块进行微调。
2. 残差特征预计算 (Residual Feature Precomputation)：
  - 在微调前，预先计算被跳过块（从第 $i$ 块到第 $i+l$ 块）的输入 $f_i$ 和输出 $f_{i+l}$ 之间的残差 $\Delta = f_{i+l} - f_i$ 。
  - 在微调阶段，跳过这些块的计算，直接将预计算的残差 $\Delta$ 加到输入上，模拟被跳过块的输出。
3. 微调过程：仅更新保留的块中的 LoRA 参数，被跳过的块参数从 GPU 卸载，仅使用预计算的残差特征。

3. 主要贡献 (Key Contributions)

动态 Patch 采样策略：首次将时间步感知的动态采样引入 DiT 微调，通过调整 Patch 大小平衡全局与局部特征学习，大幅降低显存。
基于残差预计算的块跳过机制：提出了一种针对 DiT 架构的块跳过方法，通过预计算残差特征解决了跳过层导致的训练/推理不一致问题，并引入基于交叉注意力的块选择算法，精准定位关键层。
显著的性能提升：实验证明该方法在大幅降低显存（最高减少 65.8% 最大显存）和计算量 (TFLOPs) 的同时，保持了与全量 LoRA 微调相当的个性化生成质量（Subject Fidelity 和 Text Fidelity）。
边缘设备可行性：将 DiT 微调的显存需求降低到可接受范围（例如 FLUX 模型从 22.84 GiB 降至 10.42 GiB），为在消费级 GPU 甚至未来移动设备上的个性化生成铺平了道路。

4. 实验结果 (Results)

作者在 FLUX.1-dev 和 SANA 两个主流 DiT 模型上，使用 DreamBooth 和 CustomConcept101 数据集进行了评估。

显存与计算效率：
- FLUX 模型：在 50% 块跳过率下，训练显存从 LoRA 的 22.84 GiB 降至 10.42 GiB (减少约 54%)，最大显存峰值减少 65.8%。
- SANA 模型：显存从 8.35 GiB 降至 3.10 GiB。
- TFLOPs：相比基线方法显著降低，训练速度提升约 25%。
生成质量：
- 定量指标：在 DINO (主体相似度)、CLIP-I (图像相似度) 和 CLIP-T (文本对齐) 指标上，DiT-BlockSkip 的表现与标准 LoRA 微调非常接近，明显优于 HollowedNet、LISA 和 LoRA-FA 等基线。
- 定性分析：用户偏好测试显示，该方法在主体保持和文本遵循性上与 LoRA 相当，且优于其他高效微调方法。
消融实验：
- 证明了“中间层关键性”假设：跳过中间层会导致主体消失，而跳过首尾层影响较小。
- 证明了残差预计算的必要性：没有残差预计算，跳过块会导致特征漂移，生成质量严重下降。
- 证明了动态 Patch 采样的有效性：相比固定 Resize，动态采样能更好地保留纹理细节。

5. 意义与展望 (Significance)

技术突破：解决了 DiT 架构在资源受限场景下难以微调的瓶颈，填补了 U-Net 架构之外的高效微调技术空白。
应用价值：使得在单张消费级显卡（如 A100 甚至更小的卡）上微调最新的 FLUX 等模型成为可能，极大地降低了个性化内容创作的门槛。
未来方向：
- 虽然目前尚未在移动端实测，但显存的大幅降低为端侧 (On-Device) 个性化提供了算法基础。
- 该方法具有通用性，可与其他技术（如量化、DoRA 等）结合，进一步优化资源占用。
- 未来的工作可以探索如何优化预计算特征的存储策略，以平衡 ROM 和 VRAM 的占用，实现真正的端侧微调。

总结：DiT-BlockSkip 通过巧妙的“空间压缩”（动态 Patch）和“架构剪枝”（块跳过 + 残差补偿）策略，在几乎不牺牲生成质量的前提下，实现了 DiT 模型微调的内存效率革命，是迈向边缘 AI 生成的重要一步。