Sparse-to-Sparse Training of Diffusion Models

기술 요약: 확산 모델(Diffusion Models)의 Sparse-to-Sparse 학습

문제 정의
확산 모델(Diffusion Models, DMs)은 이미지 합성, 자연어 처리, 시계열 데이터 모델링을 포함한 생성 작업에서 최첨단 성능을 달anim 성취했습니다. 그러나 이들의 광범 fact한 채택은 학습과 추론 과정 모두에서 발생하는 상당한 계산 비용에 의해 저해되고 있습니다. 기존 연구는 주로 (증류 또는 빠른 샘플링 등을 통한) 추론 단계의 가속화에 집중해 왔으나, 학습 과정은 여전히 방대하고 자원 집약적입니다. 따라서 초기 학습 단계부터부터 계산 비용과 메모리 점유율을 줄일 수 있는 방법이 필요합니다.

방법론
본 논문은 사전 학습된 조밀한(dense) 네트워크를 압축하는 대신, 희소(sparse) 네트워크를 처음부터 학습시키는 것을 목표로 하는 sparse-to-sparse 학습 패러다임을 확산 모델에 도입합니다. 본 연구는 비구조적 희소성(unstructured sparsity)에 초점을 맞추며, 두 가지 최첨단 DM 아키텍처인 Latent Diffusion(연속적인 이미지 데이터용)과 ChiroDiff(이산적인 시공간 스케치 데이터용)에 적용되는 세 가지 특정 학습 전략을 평가합니다.

제안된 세 가지 방법은 다음과 같습니다:

1. Static-DM (정적 희소 학습): 네트워크 연결성이 희소도 비율에 따라 초기화 시점에 고정됩니다. 저자들은 비제로(non-zero) 가중치를 할당하기 위해 Erdős–Rényi (ER) 및 ERK 전략을 사용하며, 이를 통해 더 큰 레이어가 더 높은 희소성을 갖도록 보장합니다. 연결 패턴은 학습 내내 변경되지 않습니다.
2. RigL-DM (동적 희소 학습): 이 방법은 동적인 가지치기(prune) 및 성장(grow) 체계를 활용합니다. 가중치는 크기에 따라 가지치기되며, 새로운 연결은 그래디언트(gradient)의 크기에 기반하여 재성장(gradient growth)됩니다.
3. MagRan-DM (동적 희소 학습): 가지치기 방식(크기 기반)은 RigL-DM과 유사하지만, 재성장은 그래디언트가 아닌 무작위(randomly)로 결정됩니다.

실험은 다양한 희소도 수준($S \in \{0.1, 0.25, 0.5, 0.75, 0.9\}$)에 걸쳐 6개의 데이터셋(LSUN-Bedrooms, CelebA-HQ, Imagenette, QuickDraw, KanjiVG, VMNIST)을 다룹니다. 저자들은 또한 가지치기 및 재성장률($p$)과 탐색 빈도($\Delta T_e$)가 미치는 영향을 조사합니다.

주요 기여

• 확산 모델에 대한 Sparse-to-Sparse 학습의 첫 적용: 본 논문은 정적 및 동적 전략을 모두 포함하여, 희소한 연결성을 사용하여 무조건적(unconditional) 확산 모델을 처음부터 학습시키는 것의 타당성을 입증합니다.
• 포괄적인 경험적 평가: 본 연구는 두 가지 서로 다른 모델 아키텍처와 6개의 데이터셋에 대해 세 가지 희소 학습 방법을 조밀한(dense) 베이스라인과 벤치마킹하여, 희소성이 DM 성능에 미치는 영향에 대한 폭넓은 관점을 제공합니다.
• 최적의 하이퍼파라미터 식별: 연구 결과, 동적 희소 학습(DST)이 일반적으로 잘 작동하지만, 가지치기 및 재성장률의 선택이 매우 중요하다는 것을 확인했습니다. 다른 맥락에서 자주 사용되는 공격적인 비율($p=0.5$)보다 보수적인 비율($p=0.05$)이 높은 희소도 영역에서 더 나은 결과를 낳습니다.
• 성능 대 효율성 트레이드오프: 본 연구는 희소 DM이 훈련 가능한 파라미터 수와 부동 소수점 연산량(FLOPs)을 크게 줄이면서도 조밀한 모델의 성능과 대등하거나 이를 능가할 수 있음을 보여줍니다.

실험 결과

• 성능: 대부분의 경우, 적어도 하나의 희소 구성이 조밀한 대응 모델의 성능을 능가했습니다. 예를 들어, CelebA-HQ 데이터셋에서 25% 희소도를 가진 RigL-DM은 파라미터를 25% 줄이고 FLOPs를 10% 줄이면서도 조밀한 모델(FID 32.74)보다 낮은 FID 점수(32.12)를 달성했습니다. QuickDraw의 경우, 90% 희소도를 가진 MagRan-DM은 FLOPs를 88% 줄이면서도 조밀한 모델과 유사한 성능을 유지했습니다.
• 희소도 수준: 중간 정도의 희소도 수준(25%~50%)이 일반적으로 성능과 효율성 사이의 가장 좋은 균형을 보였습니다. 극도로 높은 희소도(예: 90%)는 종종 성능 저하를 초래했으나, 이는 데이터셋과 모델 아키텍처에 따라 달랐습니다.
• 동적 vs 정적: 동적 희소 학습(DST) 방식(RigL-DM, MagRan-DM)은 중간 희소도 영역에서 일반적으로 Static-DM보다 우수한 성능을 보였으나, 공격적인 $p=0.5$ 비율을 사용할 때 Static-DM은 매우 높은 희소도 설정($S > 0.5$)에서 경쟁력 있거나 우월한 결과를 보였습니다. 그러나 보수적인 $p=0.05$ 비율을 사용할 때는 DST 방식이 높은 희소도 영역에서도 우위를 탈환했습니다.
• 학습 역학: 희소 모델은 조밀한 모델과 유사하게 안정적인 학습 역학을 보였으며, 이는 sparse-to-sparse 패러다임이 확산 학습 과정을 불안정하게 만들지 않음을 시사합니다.
• 조건부 생성: 클래스 조건부 모델(ImageNet 및 QuickDraw)에 대한 예비 실험은 sparse-to-sparse 학습이 조건부 설정에서도 효과적임을 나타냈으나, ImageNet에서는 조밀한 모델이 약간의 우위를 점했습니다.

의의 및 주장
본 논문은 sparse-to-sparse 학습이 생성 품질을 희생하지 않으면서 확산 모델의 효율성을 개선할 수 있는 유망한 경로임을 주장합니다. 희소 네트워크를 처음부터 학습함으로써, 이 접근 방식은 생성 비용과 추론 비용을 모두 해결하며, 잠재적으로 DM과 관련된 탄소 발자국과 메모리 요구 사항을 줄일 수 있습니다. 저자들은 적절한 방법(Static vs. Dynamic), 희소도 수준, 그리고 하이퍼파라미터(특히 가지치기/재성장률)를 선택하는 것이 성공의 핵심이라고 강조합니다. 본 연구는 적절한 구성을 갖춘다면 희소 DM이 조밀한 모델의 실행 가능한 대안이 되어, 생성형 AI 연구에서 더 넓은 구현과 빠른 반복을 가능하게 할 수 있음을 시사합니다. 저자들은 현재의 하드웨어가 비구조적 희소성의 실제적인 속도 향상을 제한하는 경우가 많지만, 파라미터 수와 FLOPs의 감소는 특히 희소 연산에 대한 하드웨어 지원이 발전함에 따라 중요한 이론적 및 실질적 이점이 된다고 언급합니다.