[논문 리뷰] Successfully Applying Lottery Ticket Hypothesis to Diffusion Model

Chao Jiang, Bo Hui외 2명|arXiv (Cornell University)|2023. 10. 28.

Advanced Neuroimaging Techniques and Applications인용 수 4

한 줄 요약

이 논문은 확산 모델에 대해 라affle 티켓 가설(Lottery Ticket Hypothesis, LTH)을 최초로 적용하여, 최대 99%의 희박성(sparsity)을 가진 부분망(subnetworks)이 전체 모델과 동일한 성능을 달성할 수 있음을 입증한다. 모듈 유사도(CKA를 통해 측정)를 기반으로 층별로 다른 희박성을 도입함으로써, 성능 손실 없이 더 희박하고 효율적인 워닝 티켓을 식별하였으며, 이로 인해 FLOPs와 메모리 사용량이 감소하였다.

ABSTRACT

Despite the success of diffusion models, the training and inference of diffusion models are notoriously expensive due to the long chain of the reverse process. In parallel, the Lottery Ticket Hypothesis (LTH) claims that there exists winning tickets (i.e., aproperly pruned sub-network together with original weight initialization) that can achieve performance competitive to the original dense neural network when trained in isolation. In this work, we for the first time apply LTH to diffusion models. We empirically find subnetworks at sparsity 90%-99% without compromising performance for denoising diffusion probabilistic models on benchmarks (CIFAR-10, CIFAR-100, MNIST). Moreover, existing LTH works identify the subnetworks with a unified sparsity along different layers. We observe that the similarity between two winning tickets of a model varies from block to block. Specifically, the upstream layers from two winning tickets for a model tend to be more similar than the downstream layers. Therefore, we propose to find the winning ticket with varying sparsity along different layers in the model. Experimental results demonstrate that our method can find sparser sub-models that require less memory for storage and reduce the necessary number of FLOPs. Codes are available at https://github.com/osier0524/Lottery-Ticket-to-DDPM.

연구 동기 및 목표

• 라플렛 티켓 가설이 확산 모델, 특히 노이즈 제거 확산 확률 모델(Denoising Diffusion Probabilistic Models, DDPM)에 적용 가능한지 조사한다.
• 구조적 프루닝을 통해 확산 모델의 높은 계산 및 메모리 비용을 줄인다.
• 다양한 층에서 프루닝 비율을 다르게 적용할 경우, 균일한 희박성보다 더 희박하고 효율적인 워닝 티켓을 도출할 수 있는지 탐색한다.
• 원래 가중치 초기화를 사용한 크기 기반 프루닝을 통해 식별된 부분망(subnetworks)이 전체 모델 성능을 따라할 수 있음을 경험적으로 검증한다.

제안 방법

• CIFAR-10, CIFAR-100, MNIST 데이터셋에서 U-Net 기반 DDPM에 대해 크기 기반 반복 프루닝을 적용한다.
• 일회성 프루닝 전략을 사용하며, 반복적 재학습과 가중치 재초기화를 통해 워닝 티켓을 식별한다.
• 여러 워닝 티켓 반복 과정에서 중심 커널 정렬(Centered Kernel Alignment, CKA)을 통해 모듈 유사도를 분석함으로써 층별로 다른 희박성을 도입한다.
• 관측된 유사도 추세에 따라 상류 층에서는 낮은 프루닝 비율(높은 희박성)과 상대적으로 높은 프루닝 비율을 하류 층에 적용하도록 설정한다.
• 표준 DDPM 하이퍼파라미터(학습률, 에포크 수, 시간 단계 수)를 사용하며, 8개의 A100 GPU에서 학습한다.
• FID 및 IS 지표를 사용해 성능을 검증하고, 프루닝된 부분망을 전체 모델과 비교한다.

실험 결과

연구 질문

• RQ1 라플렛 티켓 가설은 확산 모델, 특히 DDPM에 성공적으로 적용될 수 있는가? 이로써 희박하고 높은 성능을 내는 부분망을 식별할 수 있는가?
• RQ2 LTH 기반 프루닝에서 모든 층에 대해 균일한 희박성을 적용할 경우, 확산 모델에서는 최적 또는 비최적의 워닝 티켓을 도출하는가?
• RQ3 확산 모델의 다양한 층에서 워닝 티켓 간의 유사도는 어떻게 변화하는가?
• RQ4 층별로 다른 희박성 전략은 성능 저하 없이 더 희박한 워닝 티켓을 개선할 수 있는가?
• RQ5 프루닝은 프루닝된 확산 모델의 FLOPs와 메모리 사용량에 어떤 영향을 미치는가?

주요 결과

• CIFAR-10, CIFAR-100, MNIST에서 DDPM에 대해 최대 99%의 희박성을 가진 워닝 티켓이 성능 저하 없이 성공적으로 식별되었다.
• 중심 커널 정렬(CKA)을 통해 측정한 결과, 상류 층에서 워닝 티켓 간의 유사도가 하류 층보다 항상 높았다.
• 상류 층에서는 낮은 프루닝 비율을 적용함으로써 층별로 다른 희박성 전략을 도입한 결과, 균일한 희박성 대비 더 희박한 워닝 티켓을 도출하였으며 성능은 유지되었다.
• 제안된 방법은 최대 90%의 FLOPs 감소와 메모리 사용량 감소를 달성하여 더 효율적인 학습 및 추론을 가능하게 하였다.
• 결과적으로, 모듈 유사도를 기반으로 한 구조적이고 비균일한 프루닝 전략이 확산 모델에서 더 효율적인 부분망을 도출할 수 있음을 확인하였다.