핵심 개념
본 논문에서는 단일 교사 모델에서 여러 학생 모델로 지식을 증류하는, Diffusion Model을 가속화하기 위한 새로운 지식 증류 프레임워크인 O2MKD(One-to-Many Knowledge Distillation)를 제안합니다. O2MKD는 Diffusion Model이 서로 다른 타임스텝에서 다른 동작을 보인다는 관찰을 활용하여, 모든 타임스텝에 걸쳐 교사의 지식을 각각 연속적이고 겹치지 않는 타임스텝 전용으로 설계된 여러 학생에게 증류합니다.
초록
개요
본 연구는 단일 교사 모델에서 여러 학생 모델 그룹으로 지식을 전이하는 O2MKD(One-to-Many Knowledge Distillation)라는 새로운 지식 증류 프레임워크를 제안하여 Diffusion Model의 가속화를 목표로 합니다. 연구진은 Diffusion Model이 각기 다른 타임스텝에서 다른 양상을 보인다는 점에 착안하여, 전체 타임스텝에 걸친 교사 모델의 지식을 겹치지 않는 연속적인 타임스텝에 할당된 여러 학생 모델에게 전달하는 방식을 택했습니다.
배경 및 문제 제기
Diffusion Model은 고품질 이미지 생성 능력에도 불구하고, 기존 생성 모델에 비해 상당한 계산 오버헤드로 인해 실시간 생성에 어려움을 겪고 있습니다. 이러한 문제를 해결하기 위해 최근 연구에서는 향상된 샘플링 기술이나 단계 증류를 통해 샘플링 단계 수를 줄여 Diffusion Model을 가속화하려는 시도가 있었습니다. 그러나 각 타임스텝의 계산 비용을 줄이는 방법은 상대적으로 연구가 미흡한 실정입니다.
O2MKD의 핵심 아이디어
본 연구는 Diffusion Model이 서로 다른 타임스텝에서 다른 입력 분포, 특징 분포, 생성 정보 유형을 보인다는 점에 주목합니다. 이러한 특성은 충분한 학습 용량을 갖춘 교사 모델에게는 큰 문제가 되지 않지만, 매개변수가 제한적인 학생 모델에게는 어려움을 야기합니다. O2MKD는 이러한 문제를 해결하기 위해 여러 학생 모델을 활용하여 각 학생 모델이 인접한 타임스텝의 하위 집합에 대한 교사 지식을 학습하도록 합니다.
O2MKD의 장점
- 다른 가속화 기술과의 호환성: O2MKD는 샘플링 단계 수를 줄이는 대신 Diffusion Model 내 UNet의 계산과 관련된 오버헤드를 줄이도록 설계되어 DDIM과 같은 다른 가속화 기술과 호환됩니다.
- 기존 KD와의 호환성: O2MKD는 여러 학생과 단일 교사를 포함하는 새로운 프레임워크를 도입하여 대부분의 기존 지식 증류 방법에 직접 적용할 수 있습니다.
- 모델 압축 그 이상의 활용: O2MKD는 일반적인 Diffusion Model 학습에도 활용되어 일반화 성능을 향상시킬 수 있으므로 모델 압축 이외의 설정에도 활용될 수 있습니다.
실험 결과
CIFAR10, LSUN Church, CelebA-HQ, COCO30K 데이터셋을 사용한 실험 결과, O2MKD는 다양한 지식 증류 방법과 빠른 샘플링 기술과 함께 사용될 때 FID 점수를 크게 향상시키는 것으로 나타났습니다. 또한, O2MKD는 여러 학생을 사용하여 발생하는 추가 메모리 오버헤드가 미미하며, 학습 후 여러 학생을 하나로 병합하는 모델 병합 기술을 사용하여 추가로 줄일 수 있음을 보여주었습니다.
결론
본 연구는 Diffusion Model을 가속화하기 위한 새로운 지식 증류 프레임워크인 O2MKD를 제안했습니다. O2MKD는 여러 학생 모델을 활용하여 각 학생 모델이 특정 타임스텝 범위에 대한 교사 지식을 전문적으로 학습하도록 하여 성능을 향상시킵니다. 또한, 지식 증류가 Diffusion Model 학습에 안정적인 지도를 제공한다는 사실을 발견했으며, 이는 향후 Diffusion Model 학습 방법론 개발에 새로운 관점을 제시할 수 있습니다.
통계
CIFAR10 데이터셋에서 O2MKD를 사용하여 1.8배의 가속화와 0.18 FID 개선을 달성했습니다.
O2MKD를 사용하여 학습된 학생 모델은 지식 증류 없이 학습된 학생 모델에 비해 평균 2.57 낮은 FID를 보였습니다.
LSUN Church 및 CelebA-HQ에서 O2MKD는 지식 증류 없이 학습된 학생 모델에 비해 각각 평균 24.12 및 14.82 FID 감소를 보였습니다.
COCO30K에서 O2MKD는 기존 KD로 학습된 학생 모델에 비해 FID가 2.91 감소했으며 IS 및 CLIP 점수는 각각 4.23 및 0.0218 향상되었습니다.
O2MKD는 DDIM과 같은 빠른 샘플링 방법과 함께 사용할 때 다양한 샘플링 단계 수에서 FID를 지속적으로 감소시켰습니다.
인용구
"In brief summary, the above observations indicate that an ideal diffusion model should be capable of handling different input distributions, feature distributions, and the generation of different kinds of information."
"Fortunately, we also derive the following observation from Figure 1(b): The transition in distributions within diffusion models at different timesteps occurs gradually rather than abruptly."
"This paper brings two insights as the takeaway for readers: We find that the effectiveness of O2MKD can be explained from the perspective of distillation from a general model (i.e the teacher for all the timesteps) into multiple domain experts (i.e. multiple students for different timestep ranges), where a hyper-parameter should be utilized to balance the domain knowledge and the general knowledge, as discussed in Figure 7."