[논문 리뷰] DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization

Haowei Zhu, Dehua Tang외 13명|arXiv (Cornell University)|2024. 10. 22.

Optical Systems and Laser Technology인용 수 4

한 줄 요약

DiP-GO는 확산 모델을 위한 재학습 불필요한 새로운 절단 방법을 제안하며, 모델 절단을 미분 가능 SubNet 탐색 문제로 공식화한다. 특징 유사도를 바탕으로 한 보조 연결을 갖춘 SuperNet을 도입하고, 최적화된 손실 함수를 가진 플러그인 절단기와 함께, 몇 번의 그래디언트 최적화를 통해 최적의 SubNet을 식별한다. 이로 인해 Stable Diffusion 1.5에서 정확도 손실 없이 4.4배의 추론 속도 향상을 달성한다.

ABSTRACT

Diffusion models have achieved remarkable progress in the field of image generation due to their outstanding capabilities. However, these models require substantial computing resources because of the multi-step denoising process during inference. While traditional pruning methods have been employed to optimize these models, the retraining process necessitates large-scale training datasets and extensive computational costs to maintain generalization ability, making it neither convenient nor efficient. Recent studies attempt to utilize the similarity of features across adjacent denoising stages to reduce computational costs through simple and static strategies. However, these strategies cannot fully harness the potential of the similar feature patterns across adjacent timesteps. In this work, we propose a novel pruning method that derives an efficient diffusion model via a more intelligent and differentiable pruner. At the core of our approach is casting the model pruning process into a SubNet search process. Specifically, we first introduce a SuperNet based on standard diffusion via adding some backup connections built upon the similar features. We then construct a plugin pruner network and design optimization losses to identify redundant computation. Finally, our method can identify an optimal SubNet through few-step gradient optimization and a simple post-processing procedure. We conduct extensive experiments on various diffusion models including Stable Diffusion series and DiTs. Our DiP-GO approach achieves 4.4 x speedup for SD-1.5 without any loss of accuracy, significantly outperforming the previous state-of-the-art methods.

연구 동기 및 목표

• 확산 모델 추론의 높은 계산 비용, 특히 다중 단계 노이즈 제거로 인한 비용을 해결하기 위해.
• 시간대별로 변화하는 특징 유사성을 활용하지 못하는 정적이고 비적응형 절단 전략의 한계를 극복하기 위해.
• 미리 학습된 확산 모델을 재학습하지 않고도 효율적이고 고정밀도의 모델 절단을 가능하게 하기 위해.
• SuperNet에서 최적의 하위망원(하위망)을 식별하는 미분 가능이고 그래디언트 최적화된 절단 프레임워크를 개발하기 위해.
• 다양한 확산 아키텍처에서 높은 속도 향상과 함께 이미지 생성 품질을 유지하는 것을 목표로 하기 위해.

제안 방법

• 연속적인 노이즈 제거 단계 간의 특징 유사도를 기반으로 보조 연결을 추가하여 표준 확산 모델을 확장한 SuperNet을 구축한다.
• 최적의 SubNet이 SuperNet으로부터 유도되는 SubNet 탐색 문제로 절단 과정을 공식화한다.
• 작업에 특화된 최적화 손실 함수를 사용하여 불필요한 계산 경로를 식별하는 플러그인 절단기 네트워크를 설계한다.
• 완전한 재학습과 비교해 계산 비용을 줄이기 위해 몇 번의 그래디언트 최적화를 적용하여 최적의 SubNet을 효율적으로 탐색한다.
• 사용자 지정한 절단 비율과 구조적 제약 조건을 충족시키기 위해 경량의 후처리 단계를 통합한다.
• 확장된 노이즈 제거 단계를 거치는 백워드 전파 중 메모리 문제를 완화하기 위해 그래디언트 체크포인팅을 활용한다.

Figure 1: Overview of the SuperNet and SubNet. Standard diffusion models execute the full inference path step by step. In our framework, we propose a SuperNet based on the original flow and integrate backup connections to facilitate block removal. This allows the partial inference SubNet to efficien

실험 결과

연구 질문

• RQ1 재학습을 피하기 위해 확산 모델 절단을 미분 가능한 SubNet 탐색 문제로 공식화할 수 있는가?
• RQ2 연속적인 노이즈 제거 단계 간의 특징 유사성을 효과적으로 활용하여 불필요한 계산을 줄일 수 있는가?
• RQ3 몇 번의 그래디언트 최적화 전략이 기존 절단 방법보다 높은 성능의 SubNet을 더 효율적으로 식별할 수 있는가?
• RQ4 확산 모델에서 생성 품질을 손상시키지 않고 얼마나 빠른 추론 속도 향상을 달성할 수 있는가?
• RQ5 제안된 방법이 Stable Diffusion와 DiT와 같은 다양한 확산 아키텍처로 일반화 가능한가?

주요 결과

• DiP-GO는 재학습이 필요 없고 정확도 손실 없이 Stable Diffusion 1.5에서 추론 시간을 4.4배 빠르게 한다.
• 시간대별로 변화하는 특징 패턴을 효과적으로 활용함으로써 이전 최고 수준의 절단 기법을 초월한다.
• Stable Diffusion 1.5, 2.1, XL 및 DiT를 포함한 여러 확산 모델에 대해 일반화 가능하며 광범위한 적용 가능성을 입증한다.
• 최적화된 손실 함수를 가진 플러그인 절단기는 불필요한 구성 요소를 성공적으로 식별하여 효율적인 SubNet 탐색을 가능하게 한다.
• 몇 번의 그래디언트 최적화를 통해 고성능 SubNet으로의 신속한 수렴이 가능하여 탐색 비용을 감소시킨다.
• 후처리 단계는 최종 SubNet이 절단 비율 제약 조건을 준수하면서도 모델의 능력을 유지하도록 보장한다.

Figure 2: Overview of our diffusion pruner. a) DiP-GO employs a pruner network to learn the importance scores of blocks in the diffusion sampling process. It takes $N\times T$ queries as input and passes them through stacked self-attention (SA) and fully connected (FC) layers to capture the structur