그룹 디퓨전: 교차 샘플 협업을 통한 이미지 생성 향상

초록

본 논문은 기존 확산 모델이 샘플을 독립적으로 생성하는 한계를 넘어, 동일 조건 하에 여러 이미지를 동시에 denoise하도록 설계된 “그룹 디퓨전(Group Diffusion)”을 제안한다. 이미지 패치를 배치 차원으로 재배열해 크로스‑샘플 어텐션을 적용함으로써 intra‑와 inter‑image 관계를 동시에 학습한다. 실험 결과, 그룹 크기가 커질수록 어텐션 강도가 증가하고 FID가 크게 개선되며, ImageNet‑256에서 최대 32.2%의 FID 향상을 달성한다.

상세 분석

그룹 디퓨전은 확산 트랜스포머(DiT) 구조에 최소한의 변형만을 가해 크로스‑샘플 어텐션을 구현한다. 구체적으로, 배치 안의 N개의 이미지 패치를 (N·L)×C 형태로 reshape한 뒤 기존의 Multi‑Head Self‑Attention(MHSA) 연산에 투입한다. 어텐션 연산 후 다시 (N, L, C) 형태로 복원함으로써 각 패치가 동일 배치 내 다른 이미지의 패치와도 상호작용하도록 만든다. 이때 이미지 구분을 위해 각 이미지에 고유한 learnable sample embedding을 더해, 모델이 “어느 패치가 어느 이미지에 속하는가”를 인식하게 한다.

학습 단계에서는 이미지 유사도(주로 CLIP/DINO 임베딩 기반 코사인 유사도)를 활용해 같은 그룹에 포함될 샘플을 쿼리한다. τ_img=0.7 정도의 임계값을 두어 의미·시각적으로 연관된 이미지들을 모은 뒤, 각 이미지마다 독립적인 타임스텝을 샘플링하되 그룹 내 타임스텝 변동 σ_tv 이하로 제한한다. 손실은 기존 확산 손실과 동일하게 각 이미지별 노이즈 예측 오차를 평균한 L_group을 사용한다.

추론 단계에서는 두 가지 변형을 제시한다. GroupDiff‑f는 조건부와 무조건부 두 스코어 모두를 그룹 어텐션으로 계산하고, Classifier‑Free Guidance(CFG)를 기존 방식대로 적용한다. 반면 GroupDiff‑l은 무조건부 스코어만 그룹 어텐션으로 처리하고, 조건부 스코어는 기존 단일 이미지 방식으로 유지한다. 후자는 학습 비용이 크게 증가하지 않으면서도 대부분의 성능 향상을 제공한다는 점에서 실용적이다.

실험에서는 ImageNet‑256을 대상으로 DiT‑XL/2와 SiT을 베이스라인으로 사용했으며, 배치 크기 256을 유지하면서 그룹 크기(N)를 1,2,4,8 등으로 변동시켰다. FID, IS, Precision‑Recall 등 다양한 지표에서 그룹 크기가 커질수록 성능이 개선되는 스케일링 효과가 명확히 관찰되었다. 특히, 그룹 어텐션을 무조건부 모델에만 적용한 GroupDiff‑l이 가장 효율적인 트레이드오프를 보였으며, 사전 학습된 체크포인트에서 재학습할 경우 32.2%의 FID 감소를 기록했다.

또한, 저자들은 “Cross‑Sample Attention Strength”(CSA)라는 정량적 지표를 정의해, 어텐션 매트릭스에서 그룹 내 이미지 간 평균 어텐션 가중치를 측정하고 이를 FID와 상관관계 분석하였다. CSA가 높을수록 FID가 낮아지는 강한 음의 상관관계(r≈‑0.87)를 보였으며, 이는 모델이 얼마나 효과적으로 이미지 간 정보를 교환하는지를 직접적으로 설명한다.

한계점으로는 그룹 내 이미지가 충분히 유사하지 않을 경우 어텐션 노이즈가 증가할 가능성, 그리고 메모리 사용량이 배치 크기와 그룹 크기에 비례해 증가한다는 점을 들 수 있다. 하지만 샘플링 단계에서 배치 재구성을 통해 메모리 오버헤드를 최소화할 수 있다는 제안이 포함되어 있다. 전반적으로 그룹 디퓨전은 기존 확산 모델의 인퍼런스 파이프라인에 거의 비용 없이 크로스‑샘플 협업 메커니즘을 도입함으로써, 대규모 이미지 생성 품질을 크게 끌어올릴 수 있음을 입증한다.