针对DiT模型的深度修剪方法TinyFusion：通过端到端学习去除冗余层，以减少模型的参数量和提高推理效率

新加坡国立大学的研究人员推出一个针对DiT模型的深度修剪方法TinyFusion，旨在通过端到端学习去除冗余层，以减少模型的参数量和提高推理效率。DiT架构在图像生成领域展现出了卓越的能力，但通常伴随着大量的参数和较高的推理开销，限制了其在实际应用中的部署。TinyFusion通过学习可微的层掩码采样技术，以及一个共同优化的参数来模拟未来的微调过程，从而识别出在修剪后仍具有高恢复力的模型，确保模型在微调后能保持竞争力能。

• GitHub：https://github.com/VainF/TinyFusion

例如，在一个图像生成任务中，我们可能需要生成一张“阳光下的宁静山湖”的图片。传统的扩散变换器虽然能够生成高质量的图像，但可能需要大量的计算资源。TinyFusion通过深度修剪，可以创建一个更浅层的模型，这个模型在微调后能够以更少的计算资源生成同样高质量的图像。

主要功能和特点

1. 深度修剪：通过去除模型中的冗余层来减少参数量和计算成本。
2. 端到端学习：通过可微的采样技术和联合优化的参数，使修剪过程可学习，直接优化修剪模型的恢复力。
3. 高恢复力：修剪后的模型在微调后能快速恢复到接近原始模型的性能。
4. 跨架构通用性：TinyFusion适用于多种架构，如DiTs、MARs和SiTs。

工作原理

TinyFusion的工作原理包括以下几个关键步骤：

1. 可微采样：将修剪过程建模为层掩码的可微采样过程，使得修剪决策可以通过梯度下降进行优化。
2. 联合优化：在修剪的同时，通过一个共同优化的参数来模拟未来的微调过程，以预测修剪后模型的恢复力。
3. 恢复力建模：通过优化后微调性能的模型，而不是仅仅最小化修剪后的损失，从而更准确地评估修剪的效果。
4. 知识蒸馏：使用知识蒸馏技术进一步提升修剪后模型的性能，通过教师-学生模型的结构对齐来传递知识。

具体应用场景

1. 图像生成：在条件图像生成任务中，使用TinyFusion修剪后的模型可以生成与文本描述相匹配的图像。
2. 移动设备上的图像处理：由于TinyFusion可以显著减少模型大小，因此适用于在计算资源受限的移动设备上进行图像处理和生成。
3. 实时内容创作：在需要快速生成图像的应用中，如游戏或虚拟现实，TinyFusion可以提供实时的图像生成能力。
4. 数据集增强：在机器学习训练中，可以使用TinyFusion生成的图像来扩充数据集，提高模型的泛化能力。