【AI论文】Skrr: 跳过并重用文本编码器层以实现内存高效的文本到图像生成

摘要：在文本到图像（T2I）扩散模型中，大规模的文本编码器在根据文本提示生成高质量图像方面表现出了卓越的性能。与依赖多个迭代步骤的去噪模块不同，文本编码器只需一次前向传递即可生成文本嵌入。然而，尽管文本编码器对总推理时间和浮点运算数（FLOPs）的贡献很小，但其内存使用量却显著更高，是去噪模块的八倍之多。为了解决这一低效问题，我们提出了跳过并重用层（Skrr）策略，这是一种专为T2I扩散模型中的文本编码器设计的简单而有效的剪枝策略。Skrr通过针对T2I任务定制的方式，选择性地跳过或重用Transformer块中的某些层，从而利用Transformer块中固有的冗余来减少内存消耗，同时不牺牲性能。大量实验表明，即使在高度稀疏的情况下，Skrr也能保持与原始模型相当的图像质量，优于现有的块级剪枝方法。此外，Skrr在多个评价指标（包括FID、CLIP、DreamSim和GenEval分数）上均保持了性能，同时实现了最优的内存效率。Huggingface链接：Paper page，论文链接：2502.08690

1. 引言

• 文本到图像生成模型的发展：
  • 文本到图像（T2I）扩散模型在生成高质量图像方面取得了显著进展，能够紧密匹配文本提示。这些模型，如Stable Diffusion、PixArt-Σ等，利用现代架构和大规模文本编码器，生成了高质量的图像。
  • 然而，这些模型需要大量的计算资源，特别是内存资源，这限制了它们的部署和可扩展性。
• 现有优化方法：
  • 为了提高T2I扩散模型的效率，研究人员提出了多种策略，包括知识蒸馏（KD）、剪枝（Pruning）和量化（Quantization）。
  • 这些方法主要关注于去噪模块的优化，而文本编码器作为模型的重要组成部分，其优化却相对被忽视。
• 文本编码器的重要性与挑战：
  • 文本编码器在T2I扩散模型中占据大量参数（超过70%），但对FLOPs的贡献却很小（小于0.5%）。
  • 文本编码器的高内存使用量是部署大规模T2I模型的主要瓶颈之一。

2. 方法：Skrr策略

• Skrr策略概述：
  • Skrr（Skip and Re-use layers）是一种专为T2I扩散模型中的文本编码器设计的简单而有效的剪枝策略。
  • 它通过选择性地跳过或重用Transformer块中的某些层，以减少内存消耗，同时保持性能。
• Skip阶段：
  • 在Skip阶段，使用T2I扩散定制的差异度量来评估每个Transformer子块（包括多头注意力MHA和前馈网络FFN）的重要性。
  • 采用基于beam search的算法来探索多个剪枝路径，以优化剪枝过程，同时考虑块之间的相互作用。
• Re-use阶段：
  • 在Re-use阶段，评估通过重用相邻未跳过块来恢复被跳过块的可能性，以最小化性能损失。
  • 通过理论分析证明，Re-use可以在某些条件下提供比仅Skip更紧的错误界。

3. 实验与结果

• 实验设置：
  • 在多个T2I扩散模型（如PixArt-Σ、Stable Diffusion 3、FLUX.1-dev）上进行了广泛的实验。
  • 使用CC12M数据集的子集作为校准数据集，以识别不重要的块进行剪枝。
  • 评估指标包括FID、CLIP分数、DreamSim分数和GenEval分数。
• 定量结果：
  • 在不同稀疏度下，Skrr在保持图像质量和文本对齐方面优于现有的块级剪枝方法（如ShortGPT、LaCo、FinerCut）。
  • 在高稀疏度（>40%）下，Skrr仍能生成与原始模型相当质量的图像，同时在多个评价指标上保持领先。
• 定性结果：
  • 通过视觉比较，Skrr生成的图像在高度稀疏的情况下仍能保持与原始模型相似的图像质量和文本对齐。
  • Skrr在压缩多个文本编码器（如Stable Diffusion 3中的CLIP-L、CLIP-G和T5-XXL）时也表现出色，生成的图像质量几乎无损。

4. 讨论

• Skrr策略的优势：
  • Skrr通过利用Transformer块中的固有冗余，实现了高效的内存使用，同时保持了模型的性能。
  • 与现有的块级剪枝方法相比，Skrr在多个评价指标上均表现出色，特别是在高稀疏度下。
• 对模型指导的见解：
  • 实验结果表明，剪枝或合并层不仅可以减少模型复杂性，还可以在某些情况下提高某些性能指标（如FID分数）。
  • 这可能与模型指导方法（如分类器自由指导CFG）有关，这些方法通过修改去噪网络来近似无条件分数。
• 未来工作：
  • 进一步提高Skrr在极端稀疏度下的性能，特别是在保持文本对齐方面。
  • 探索将Skrr与其他优化策略（如量化）结合使用，以实现更高的内存和计算效率。

5. 结论

• Skrr策略的有效性：
  • Skrr是一种专为T2I扩散模型中的文本编码器设计的有效剪枝策略。
  • 它通过选择性地跳过或重用Transformer块中的某些层，显著减少了内存消耗，同时保持了模型的性能。
• 广泛的适用性：
  • Skrr在多个T2I扩散模型和评价指标上均表现出色，证明了其广泛的适用性。
  • 未来，Skrr有望成为部署大规模T2I模型的重要工具，降低计算资源要求，促进其在更多场景中的应用。

6. 具体实施细节

• Skip算法：
  • 算法1详细描述了Skip阶段的实现过程。
  • 使用MSE（均方误差）作为差异度量来评估剪枝后的模型与原始模型之间的输出差异。
  • 通过beam search算法来探索多个剪枝路径，以优化剪枝过程。
• Re-use算法：
  • 算法2详细描述了Re-use阶段的实现过程。
  • 对于每个被跳过的块，评估重用相邻未跳过块的可能性，并选择差异最小的块进行重用。
  • 通过迭代更新重用字典来优化重用过程。
• 实验配置：
  • 在PixArt-Σ模型上进行了详细的实验配置和参数设置。
  • 使用了不同的稀疏度水平（如24.3%、32.4%、41.9%）来评估Skrr的性能。
  • 通过比较Skrr与基线方法（如ShortGPT、LaCo、FinerCut）在多个评价指标上的表现来验证Skrr的有效性。

7. 局限性

• 极端稀疏度下的性能下降：
  • 在极端稀疏度下（>50%），Skrr的性能可能会显著下降，这可能是由于文本编码器的表示能力受到严重限制。
  • 未来工作可以探索将Skrr与其他优化策略结合使用，以在极端稀疏度下保持模型性能。
• 性能提升的局限性：
  • 尽管Skrr在多个评价指标上表现出色，但它并未实现超越原始密集模型的性能提升。
  • 未来工作可以进一步探索如何通过剪枝策略来提高T2I扩散模型的性能。

8. 实际应用前景

• 降低部署成本：
  • Skrr通过减少文本编码器的内存使用量，显著降低了T2I扩散模型的部署成本。
  • 这使得大规模T2I模型能够在更多资源受限的场景中部署和应用。
• 促进模型创新：
  • Skrr为T2I扩散模型的研究人员提供了一种有效的工具来探索模型的优化和剪枝策略。
  • 这将促进T2I扩散模型的持续创新和发展，推动其在更多领域的应用。

9. 对相关领域的贡献

• 对T2I扩散模型的贡献：
  • Skrr为T2I扩散模型提供了一种新的优化策略，显著提高了模型的内存效率。
  • 这将促进T2I扩散模型在更多实际应用中的部署和应用。
• 对剪枝策略的贡献：
  • Skrr提出了一种新的块级剪枝策略，通过选择性地跳过或重用Transformer块中的某些层来减少内存消耗。
  • 这为其他类型的深度学习模型提供了一种新的剪枝思路和方法。

10. 总结与展望

• 总结：
  • Skrr是一种专为T2I扩散模型中的文本编码器设计的简单而有效的剪枝策略。
  • 它通过利用Transformer块中的固有冗余来减少内存消耗，同时保持了模型的性能。
  • 大量实验表明，Skrr在多个评价指标上均表现出色，特别是在高稀疏度下。
• 展望：
  • 未来工作可以进一步探索Skrr在极端稀疏度下的性能提升策略。
  • 同时，也可以将Skrr与其他优化策略结合使用，以实现更高的内存和计算效率。
  • 此外，Skrr的思想和方法也可以应用于其他类型的深度学习模型中，以推动深度学习领域的持续创新和发展。