[Github]
[arXiv](2023/11/16 version v3)
Abstract
Stable Diffusion에서 residual block과 attention block 제거, 증류를 통해 모델 크기를 줄인 BK-SDM 제안
Compact U-Net architecture
Pruning 후 generation score의 차이를 측정하여 민감도 분석. 높을수록 제거가 가능함을 의미한다.
Fewer blocks in the down and up stages
Down stages
첫 번째 RA pair가 변경된 공간 정보를 처리하기 때문에 더 중요하다. 민감도 분석과도 일치.
따라서 두 번째 RA pair를 제거.
Up stages
Down stage의 두 번째 RA pair와 잔차 연결된 RA pair 제거.
Removal of the entire mid-stage
민감도 분석에 따라 기여도가 낮은 Mid stage를 제거. 피라미터를 11% 줄일 수 있다.
Further removal of the innermost stages
민감도 분석에 따라 가장 작은 해상도의 Up/Down block 제거.
Distillation-based retraining
학생 모델 S가 ground truth를 예측하고
교사 모델의 출력을 모방하고
각 해상도의 마지막 단계에서 feature map이 일치하는지 확인한다.
Total Loss:
Experiments
