OBS-Diff: 원샷에서 디퓨전 모델을 위한 정확한 가지치기
OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot
October 8, 2025
저자: Junhan Zhu, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang
cs.AI
초록
대규모 텍스트-이미지 확산 모델은 강력한 성능을 자랑하지만, 과도한 계산 비용이 발생한다는 문제가 있습니다. 기존의 원샷 네트워크 프루닝 방법들은 확산 모델의 반복적인 노이즈 제거 특성으로 인해 이를 직접 적용하기 어렵습니다. 이러한 격차를 해소하기 위해, 본 논문은 대규모 텍스트-이미지 확산 모델의 정확하고 학습 없이도 압축을 가능하게 하는 새로운 원샷 프루닝 프레임워크인 OBS-Diff를 제안합니다. 구체적으로, (i) OBS-Diff는 고전적인 Optimal Brain Surgeon(OBS)을 재활용하여 현대 확산 모델의 복잡한 아키텍처에 적합하도록 조정하고, 비정형, N:M 준정형, 그리고 구조화된(MHA 헤드 및 FFN 뉴런) 희소성을 포함한 다양한 프루닝 세분화를 지원합니다; (ii) 프루닝 기준을 확산 과정의 반복적 역학과 일치시키기 위해, 오류 누적 관점에서 문제를 검토하여 로그 감소 가중치 방식을 통합한 새로운 시간 단계 인식 헤시안 구성을 제안합니다. 이는 초기 시간 단계에 더 큰 중요성을 부여하여 잠재적인 오류 누적을 완화합니다; (iii) 또한, 계산적으로 효율적인 그룹별 순차적 프루닝 전략을 제안하여 비용이 많이 드는 보정 과정을 분산시킵니다. 광범위한 실험을 통해 OBS-Diff가 확산 모델에 대한 최첨단 원샷 프루닝을 달성하고, 시각적 품질의 최소한의 저하로 추론 가속화를 제공함을 보여줍니다.
English
Large-scale text-to-image diffusion models, while powerful, suffer from
prohibitive computational cost. Existing one-shot network pruning methods can
hardly be directly applied to them due to the iterative denoising nature of
diffusion models. To bridge the gap, this paper presents OBS-Diff, a novel
one-shot pruning framework that enables accurate and training-free compression
of large-scale text-to-image diffusion models. Specifically, (i) OBS-Diff
revitalizes the classic Optimal Brain Surgeon (OBS), adapting it to the complex
architectures of modern diffusion models and supporting diverse pruning
granularity, including unstructured, N:M semi-structured, and structured (MHA
heads and FFN neurons) sparsity; (ii) To align the pruning criteria with the
iterative dynamics of the diffusion process, by examining the problem from an
error-accumulation perspective, we propose a novel timestep-aware Hessian
construction that incorporates a logarithmic-decrease weighting scheme,
assigning greater importance to earlier timesteps to mitigate potential error
accumulation; (iii) Furthermore, a computationally efficient group-wise
sequential pruning strategy is proposed to amortize the expensive calibration
process. Extensive experiments show that OBS-Diff achieves state-of-the-art
one-shot pruning for diffusion models, delivering inference acceleration with
minimal degradation in visual quality.