기술 요약: 개선된 비혼합 확산 (Improved Immiscible Diffusion)

1. 문제 정의

확산 모델(Diffusion Model)의 학습은 모델 크기와 데이터셋 복잡도에 따라 크게 증가하는 상당한 계산 비용과 시간 요구 사항으로 인해 어려움을 겪고 있습니다. 본 연구에서 식별된 근본적인 병목 현상은 **궤적 혼합성(trajectory miscibility)**입니다. 바닐라 확산 학습에서는 이미지에 무작위 가우시안 노이즈가 추가되는데, 이로 인해 서로 다른 이미지의 확산 경로(flows)가 노이즈 공간 내의 동일한 지점에서 서로 섞이게 되며, 특히 노이즈 수준이 높아질수록($t \to T$) 이러한 현상이 심화됩니다.

여러 개의 서로 다른 이미지 궤적이 동일한 노이즈 지점 $x_t$로 수렴할 때, 디노이징 네트워크 $\epsilon_\theta(x_t, t)$는 모호한 학습 목표에 직면하게 됩니다. 모델은 고유한 이미지를 복구하기 위해 특정 노이즈 벡터 하나를 예측하는 대신, 상충하는 여러 노이즈 벡터들의 평균을 예측하도록 강요받습니다. 노이즈 수준이 최대치에 도달함에 따라 디노이징 타겟은 상수(데이터셋의 평균)로 수렴하며, 이는 특정 이미지를 복구하기 위한 학습 과정을 무력화합니다. 이러한 "섞임(mixing)" 현상은 되돌릴 수 없고 혼란스러운 궤적을 생성하여 학습 효율을 저하시킵니다.

배치 단위의 최적 운송(Optimal Transport, OT)을 적용하거나 선형 할당(linear assignment)을 통해 이미지와 더 가까운 노이즈 지점을 쌍으로 묶는 등의 기존 시도들은 제한적인 성공만을 보여주었습니다. 이러한 방법들은 이미지-노이즈 거리를 약 2% 감소시키고 디노이징 함수의 표준 편차를 약 4% 감소시키는 데 그쳤습니다. 또한, 선형 할당 방식은 $O(n^3)$의 계산 복잡도를 가지므로 대규모 배치 크기에는 부적합하며, 이미지-노이즈 쌍을 제한하는 것이 생성되는 이미지의 다양성을 해칠 수 있다는 우려를 낳았습니다.

2. 방법론

저자들은 **개선된 비혼합 확산(Improved Immiscible Diffusion)**을 제안합니다. 이는 '비혼합성(immiscibility)'의 개념을 단순히 특정 이미지-노이즈 쌍을 맞추는 전략이 아니라, 서로 다른 이미지로부터 오는 확산 궤적의 혼합을 전반적으로 줄이는 일반적인 개념으로 재정의한 정교한 프레임워크입니다.

2.1 이론적 토대: 안정적인 이미지-노이즈 상관관계

본 연구는 생성된 이미지가 그 기원이 되는 노이즈와 안정적으로 상관되어 있음을 먼저 입증합니다. 바닐라 DDIM에 대한 섭동 실험(perturbation experiments)을 통해, 저자들은 특정 노이즈 지점에 미세한 노이즈(최대 20% 섭동)를 추가하더라도 생성된 이미지가 세만틱하게 동일하게 유지됨을 보여주었습니다. 오직 상당한 수준의 섭동(30%)이 가해졌을 때만 객체 수준의 변화가 발생했습니다. 이 발견은 확산을 특정 노이즈 영역으로 제한하는 것(비혼합성)이 노이즈에서 이미지로의 매핑이 견고하기 때문에 생성 다양성을 본질적으로 파괴하지 않는다는 점을 검증합니다.

2.2 작동 메커니즘

핵심 메커니즘은 모델의 디노이징 목표를 명확히 하는 것입니다. 확산 궤적의 중첩을 줄임으로써 다음과 같은 효과를 얻습니다:

1. 노이즈 클러스터의 분리: 서로 다른 이미지에 할당된 노이즈 지점들은 노이즈 공간 내에서 겹치지 않는 뚜렷한 클러스터를 형성합니다.
2. 디노이징 목표의 명확화: 디노이징 함수는 더 이상 상충하는 타겟들을 평균 낼 필요가 없습니다. 높은 노이즈 수준($\tau = S$)에서도 모델은 특정 이미지에 추가된 특정 노이즈를 효과적으로 예측할 수 있습니다.
3. 특징 공간의 개선: t-SNE 분석 결과, 비혼합 확산 하에서는 디노이징 단계 전반에 걸친 예측된 이미지들이 뚜렷하고 교차하지 않는 경로를 형성하는 반면, 바닐라 확산 경로는 무질서하게 뒤엉킨 모습을 보였습니다.

2.3 구현 방식

본 논문은 계산 비용이 많이 드는 선형 할당을 넘어, 비혼합성을 달성하기 위한 일련의 구현 방식을 소개합니다:

• 배치 단위 선형 할당 (Baseline): $L_2$ 최소화를 통해 각 이미지를 가장 가까운 노이즈 지점에 할당합니다. 궤적을 분리하는 데는 효과적이지만 $O(n^3)$의 복잡도 문제를 가집니다.
• K-최근접 이웃(KNN) 노이즈 선택: 각 이미지에 대해 $k$개의 무작위 노이즈 지점을 샘플링하고, 그중 $L_2$ 거리가 최소인 것을 선택합니다. 이 방법은 계산 효율적($O(n)$)이며 대규모 배치 크기에도 잘 확장됩니다. 분석 결과, 결과적인 노이즈 분포는 거의 가우시안 형태를 유지합니다 (엄격한 가우시안과 비교했을 때 KL 발산이 무시할 수 있는 수준임).
• 이미지 스케일링 (Image Scaling): 이미지-노이즈 쌍을 직접 다루지 않는 새로운 접근법입니다. 확산 전 이미지 픽셀 값을 상수 계수 $>1$(예: 2 또는 4)로 곱함으로써, 서로 다른 이미지들의 확산 영역 중심 간의 $L_2$ 거리를 증가시킵니다. 이는 노이즈 공간을 변경하거나 명시적인 쌍을 맞출 필요 없이 자연스럽게 확산 궤적의 교차를 줄여줍니다.

3. 주요 기여

1. 개념적 정교화: 저자들은 비혼합 확산을 단순한 선형 할당 기법에서 확산 궤적의 혼합 감소로 정의되는 더 넓은 개념으로 확장했습니다. 이들은 이러한 감소가 이미지-노이즈 상관관계의 본질적인 안정성 덕분에 생성 다양성을 해치지 않음을 입증했습니다.
2. 효율적인 구현: 논문은 선형 할당의 효율적인 대안으로 KNN 노이즈 선택과 이미지 스케일링을 제안합니다. 이 방법들은 매우 낮은 계산 오버헤드로 비혼합성을 달성하여 대규모 학습에 사용할 수 있게 합니다.
3. 포괄적인 검증: 다양한 모델(Consistency Models, Flow Matching, DDIM, Stable Diffusion), 데이터셋(CIFAR-10, ImageNet, MSCOCO), 그리고 작업(무조건/조건부 생성, 미세 조정, 이미지 편집, 로보틱스 플래닝)에 걸쳐 광범위한 실증적 증거를 제공합니다.
4. 메커니즘 분석: 단계별 특징 분석을 통해, 비혼합성이 어떻게 학습을 가속화하는지 밝혀냈습니다. 즉, 바닐라 확산에서 가장 학습하기 어려운 부분인 가장 노이즈가 많은 레이어들의 디노이징 능력을 활성화함으로써 이를 달성합니다.

4. 실험 결과

실험은 일관되고 유의미한 학습 효율 및 최종 성능 향상을 보여줍니다:

• 학습 속도: 비혼합 확산 제품군은 바닐라 베이스라인과 동일한 성능 수준(FID)에 도달하기 위해 최대 4배 빠른 학습을 달성했습니다.
• 성능 향상:
  • 무조건부 생성: CIFAR-10에서 KNN 방식으로 구현된 비혼합 DDIM은 FID를 6.63(Vanilla)에서 4.45로 개선했습니다.
  • 조건부 생성: ImageNet 및 MSCOCO에 대한 Stable Diffusion 학습에서, 비혼합 확산은 처음부터 학습하는 경우와 미세 조정하는 경우 모두에서 바닐라 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • 이미지 편집: 인페인팅(In-painting) 및 아웃페인팅(Out-painting) 작업에서 더 낮은 FID 점수를 기록하여(예: In-painting FID가 18.35에서 17.32로 감소), 소스 정보 보존 능력이 향상되었음을 나타냈습니다.
  • 로보틱스 플래닝: PushT 작업에 적용했을 때, 비혼합 확산은 로봇의 평균 성공 영역 커버리지를 79.56%에서 82.83%로 개선했습니다.
• 효율성: KNN 방식은 노이즈 선택 실행 시간을 밀리초 단위(Linear Assignment)에서 서브 밀리초 단위(예: 배치 크기 256에서 0.2ms)로 단축하여 확장성을 확보했습니다.
• 다양성: CLIPScore 평가를 통해 비혼합 확산이 바닐라 모델과 대등한 수준의 프롬프트-이미지 대응 관계 및 생성 다양성을 유지함을 확인했습니다.

5. 의의 및 주장

본 논문은 **궤적 혼합성(trajectory miscibility)**이 플로우 매칭(flow matching)이나 리프티파이드 플로우(rectified flow)와 같은 다른 효율화 기술에 밀려 간과되어 온, 확산 학습의 근본적인 병목 현상이라고 주장합니다. 플로우 매칭이 궤적을 선형화할 수는 있지만, 궤적이 섞이는 것(혼합성)을 방지하지는 못합니다.

저자들은 디노이징의 모호함의 근본 원인인 궤적 혼합을 해결함으로써, 고효율 확산 학습을 위한 새로운 방향을 제시한다고 주장합니다. 이 연구는 "혼합성 문제"가 다양한 확산 아키텍처에 영향을 미치는 일반적인 이슈이며, 이를 줄이는 것이 더 빠른 수렴과 더 나은 성능을 위한 견고하고 모델 불가지론적인(model-agnostic) 경로임을 시사합니다. 결론적으로 본 논문은 이 접근 방식이 단순한 구조적 변화를 넘어 확산 과정의 역학을 근본적으로 바꾸는, 효율적인 확산 학습을 위한 새로운 연구 방향을 제공할 수 있음을 시사합니다.