기술 요약: 롱테일 확산 모델을 위한 대조적 조건부-무조건부 정렬 (CCUA)

문제 정의

클래스 조건부 이미지 합성(class-conditional image나 synthesis)을 위한 학습 데이터는 흔히 "헤드(head)" 클래스에 비해 "테일(tail)" 클래스의 이미지가 매우 적은 롱테일 분포를 나타냅니다. 이러한 불균형은 확산 모델이 테일 클래스의 이미지를 생성할 때 모드 붕괴(mode collapse)와 다양성 감소를 겪게 만듭니다. 표준 디노이징 확산 확률 모델(DDPM) 및 스코어 기반 확산 모델(SBDM)을 포함한 기존 방법들은 헤드 클래스의 품질을 저해하지 않으면서, 과소 대표된 클래스들에 대해 고충실도(high-fidelity) 및 고다양성 이미지를 생성하는 데 어려움을 겪습니다. 대조 학습(contrastive learning)이 롱테일 인식 및 GAN 기반 생성에서 효과적임이 입증되었음에도 불구하고, 이를 확산 모델의 디노이징 과정 특성에 맞춰 적용하여 클래스 불균형 생성을 해결하기 위해서는 새로운 적응 방식이 필요합니다.

방법론: 대조적 조건부-무조건부 정렬 (CCUA)

저자들은 헤드 클래스의 품질을 유지하면서 테일 클래스 이미지의 다양성과 충실도를 개선하기 위해 설계된 두 가지 시너지 효과가 있는 손실 함수 프레임워크인 CCUA를 제안합니다. 이 방법은 U-Net 기반 아키텍처와 Diffusion Transformer(DiT/SiT) 모두와 호환됩니다.

1. 무조건부 생성을 위한 비지도 대조 손실 (UCL)

모드 붕괴를 해결하기 위해, 저자들은 무조건부 생성(unconditional generation) 잠재 변수에 적용되는 비지도 대조 손실을 도입합니다.

• 메커니즘: 노이즈 추정 네트워크를 잠재 인코더 $e(\cdot)$와 디코더 $d(\cdot)$로 분리합니다. 인코더는 노이즈가 섞인 이미지 $x_t$를 저차원 잠재 변수 $h$로 매핑합니다.
• 손실 공식: 이 손실은 음의 샘플(negative samples)만을 사용하는 InfoNCE 스타일의 공식을 사용합니다. 미니 배치 내에서 각 이미지는 앵커(anchor) 역할을 하며, 배치의 다른 모든 이미지는 음의 샘플 역할을 합니다. 데이터 증강을 통한 양의 쌍(positive pairs)은 생성되지 않으며, 앵커 자체가 자신의 양의 벡터로 취급됩니다.
• 목적: 이는 잠재 공간 내에 "반발력(repulsive force)"을 생성하여, 클래스와 관계없이 합성된 이미지 간의 거리를 최대화합니다. 이는 모델이 무조건부 잠재 변수를 다양화하도록 강제하여, 특정 상수 벡터로 붕도되거나 제한된 학습 샘플에 과적합되는 것을 방지합니다.
• 이론적 근거: 대조 손실을 조건부 잠재 변수에 직접 적용하면, 모델이 클래스 간 다양성(클래스 분리)은 극대화하면서 클래스 내부의 다양성은 무시하는 사소한 해(trivial solution)에 빠질 수 있습니다. 이 손실을 무조건부 잠재 변수에 적용함으로써, 모델은 모든 샘플을 다양화하도록 강제되며, 이는 두 번째 손실을 통해 조건부 생성으로 전이(distill)됩니다.

2. 조건부-무조건부 정렬을 위한 정렬 손실 (AL)

두 번째 구성 요소는 다양화된 무조건부 생성과 조건부 생성을 정렬합니다.

• 메커니즘: 이 손실은 조건부 분포 $p_\theta(x_{t-1}|x_t, c)$와 무조건부 분포 $p_\theta(x_{t-1}|x_t)$ 사이의 KL 발산을 최소화합니다. 실제 구현에서는 추정된 노이즈 $\epsilon_\theta(x_t, t, c)$와 $\epsilon_\theta(x_t, t, \emptyset)$ 사이의 MSE 손실으로 단순화됩니다.
• 시간 가중치(Time-Weighting): 결정적으로, 이 정렬은 큰 타임스텝(역과정의 초기 단계)에서 더 무겁게 가중치가 부여됩니다.
• 동기: 저자들은 초기 디노이징 단계에서 서로 다른 클래스의 이미지(및 조건부와 무조건부 생성 간)의 저주파 성분이 매우 유사하다는 점에 주목했습니다. 이러한 초기 단계를 정렬함으로써, 모델이 헤드 클래스(풍부한 데이터)로부터 테일 클래스(희소한 데이터)로의 **지식 공유(knowledge sharing)**를 촉진할 수 있습니다. 이를 통해 테일 클래스는 생성 초기 단계에서 헤드 클래스가 학습한 풍부한 다양성의 혜택을 받을 수 있습니다.

3. 시너지

프레임워크는 표준 DDPM 손실과 제안된 CCUA 손실을 결합합니다:
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
시너지는 UCL이 무조건부 잠재 공간을 다양화하고, AL이 조건부 잠재 공간을 이 다양화된 무조건부 잠재 공간에 정렬한다는 점에 있습니다. 이러한 조건부 잠재 변수의 간접적인 다양화는 조건부 잠재 변수에 직접 대조 손실을 적용하는 것보다 경험적으로 더 효과적임이 입증되었습니다.

주요 기여

1. CCUA 프레임워크: 무조건부 잠재 변수에 대한 비지도 대조 학습과 조건부-무조건부 정렬 손실을 결합한 클래스 불균형 확산 모델을 위한 새로운 방법론입니다.
2. 정렬 전략: 저주파 성분의 유사성을 활용하고 헤드 클래스에서 테일 클래스로의 지식 전달을 가능하게 하기 위해, 특히 큰 타임스텝(초기 디노이징 단계)에서 조건부 및 무조건부 생성을 정렬할 것을 제안합니다.
3. 비지도 대조 정규화: 대조 쌍에 클래스 레이블을 요구하지 않고 무조건부 잠재 변수에 음의 샘플만을 사용하여 클래스 내 및 클래스 간 다양성을 최대화하는 비지도 대조 손실을 사용합니다.
4. 아키텍처 불가지론(Architecture Agnosticism): 이 방법은 U-Net 기반 아키텍처와 Diffusion Transformer(SiT) 모두에서 효과적임을 입증했습니다.

실험 결과

본 방법은 ImageNet-LT (256×256), TinyImageNet-LT (64×64), Places-LT (64×64), 그리고 CIFAR-LT 데이터셋에서 평가되었습니다.

• 정량적 성능:
  • ImageNet-LT에서 CCUA는 SiT, CBDM, REPA, Dispersive Loss를 포함한 베이스라인들을 크게 앞질렀습니다.
  • 900k 학습 단계에서 CCUA는 FID 15.1(SiT의 19.9 대비) 및 FID$_{tail}$ 22.5(SiT의 33.9 대비)를 달성하였으며, 이는 테일 클래스 충실도에서 약 11.4 포인트의 개선을 의미합니다.
  • 또한 이 방법은 전반적인 Recall(다양성) 및 Precision(충실도) 지표를 개선했습니다.
  • TinyImageNet-LT에서 CCUA는 최고의 FID(12.0)와 FID$_{tail}$(20.8)를 달성하여 DDPM, CBDM, OCLT를 능가했습니다.
• 정성적 결과: 시각화 결과, CCUA는 표준 SiT가 종종 반복적이거나 저품질의 이미지를 생성하는 것과 비교하여, 테일 클래스(예: 'red wine', 'espresso', 'window shade')에 대해 더 높은 다양성과 충실도를 가진 이미지를 생성함을 보여줍니다.
• 절제 연구(Ablation Studies):
  • 두 손실의 결합이 필수적입니다. 하나만 사용할 경우 최적의 결과를 얻을 수 없습니다.
  • 정렬 손실을 큰 타임스텝에 적용하는 것(선형 가중치)이 상수 또는 이차 가중치를 사용하는 것보다 우수합니다.
  • 트랜스포머의 N/4-th 블록에서 추출한 잠재 표현을 사용하는 것이 전체 FID와 공간적 FID/Recall 사이의 가장 좋은 절충안을 제공합니다.

의의 및 주장

본 논문은 저자들이 롱테일 GAN의 원리(특히 낮은 해상도 또는 초기 단계에서의 무조건부 학습)를 확산 패러다임에 적응시켜 롱테일 생성 문제를 효과적으로 해결한다고 주장합니다. 다양화 작업(무조건부 대조 학습이 담당)과 조건화 작업(정렬이 담당)을 분리함으로써, 직접적인 조건부 대조 정규화의 함정을 피합니다.

저자들은 자사의 접근 방식이 다음과 같음을 강조합니다:

• 헤드 클래스의 품질을 저해하지 않으면서 테일 클래스 생성을 개선합니다.
• 기존 DDPM 및 SBDM 파이프라인에 구현하기 쉽습니다.
• 극단적인 클래스 불균형 상황(예: 불균형 계수가 0.001인 ImageNet-LT)에서 Dispersive Loss 및 DiffROP과 같은 동시대 연구들보다 뛰어난 성능을 보입니다.
• 역과정의 초기 단계에서 조건부 및 무조건부 생성을 정렬하는 것이 롱테일 환경에서 지식 공유에 왜 유익한지에 대한 이론적, 경험적 근거를 제공합니다.