기술 요약: 확산 도메인 확장 (DDE)

문제 정의

사전 훈련된 확산 모델은 이미지, 오디오 및 기타 데이터 모달리티 생성에서 최첨단 결과를 달성했습니다. 그러나 이러한 모델을 처음부터 다시 훈련하여 더 큰 출력 도메인 (예: 파노라마 이미지, 긴 오디오 트랙) 이나 더 복잡한 조건 입력 (예: 여러 객체 위치) 을 처리하도록 하는 것은 계산 비용이 많이 듭니다. 생성 도메인을 확장하기 위한 기존 접근 방식은 종종 추가 훈련이 필요하지 않지만 특정 작업으로 제한되고 유연성이 부족하며 고정된 알고리즘적 조정 메커니즘 (예: MultiDiffusion) 에 의존합니다. 동적 디퓨저 모듈을 훈련하는 것과 같은 다른 방법들은 훈련 시 고정된 모달리티를 요구하거나 훈련 중 관찰된 도메인보다 큰 도메인으로 일반화하지 못하는 경우가 많습니다.

이 논문이 다루는 핵심 과제는 기본 모델을 재훈련하지 않으면서도 사전 훈련된 확산 모델의 생성 도메인을 더 큰 객체 생성 및 더 복잡한 조건 처리로 효율적으로 확장하는 동시에, 조정자의 훈련 중 관찰된 인스턴스보다 큰 인스턴스에 대한 솔루션이 일반화되도록 보장하는 방법입니다.

방법론: 확산 도메인 확장 (DDE)

저자들은 사전 훈련된 확산 모델의 출력을 조율하기 위해 컴팩트하고 훈련 가능한 "조정자 (coordinator)" 네트워크를 사용하는 확산 도메인 확장 (DDE) 프레임워크를 제안합니다.

일반적 공식화

소거 네트워크 $D(x, y, t)$를 가진 사전 훈련된 조건부 확산 모델 $p(x|y)$가 주어졌을 때, 목표는 확장된 입력 $Y^{[L]}$에 조건부로 확장된 객체 $X^{[L]}$을 생성하는 것이며, 여기서 $L$은 스케일 (예: 패치 수 또는 조건 레이블 수) 을 나타냅니다.

1. 분해: 확장된 쌍 $(X^{[L]}, Y^{[L]})$은 매핑 함수 $F$를 사용하여 $L$개의 더 작은 부분 $([x_1, \dots, x_L], [y_1, \dots, y_L])$으로 분해됩니다. 이러한 부분들은 기본 모델의 원래 도메인 내에 속합니다.
2. 기본 생성: 각 부분 $x_i$는 사전 훈련된 소거기 $D(x_i, y_i, t)$에 의해 처리됩니다.
3. 조정: 훈련 가능한 조정자 네트워크 $C[\cdot]$는 소거된 출력 시퀀스 $[D(x_i, y_i, t)]_{i=1}^L$와 조건 정보 $[y_i]_{i=1}^L$를 입력으로 받습니다. 이는 이러한 출력들을 조율하여 전체 확장된 객체 $X^{[L]}(0)$에 대한 일관된 추정을 생성하도록 학습됩니다.
4. 훈련 목적: 조정자는 크기 $L_{train}$의 데이터셋에서 복합 소거기 $D^{[L]}$의 소거 오차를 최소화하도록 훈련됩니다. 중요한 점은 이 방법이 훈련 테스트 크기 $L_{test} \ge L_{train}$에 대해 일반화되도록 설계되었다는 것입니다.

조정자 아키텍처

저자들은 Peebles & Xie (2023) 에서 적응된 비전 트랜스포머 (ViT) 아키텍처를 조정자에 활용합니다.

• 입력 처리: 사전 훈련된 모델에서 나온 소거된 패치들은 "패치화 (patchified)"되고 더 큰 객체에 대한 글로벌 위치 인코딩이 추가됩니다. 저자들은 더 큰 공간 차원을 효과적으로 처리하기 위해 회전 위치 임베딩 (RoPE) 을 사용합니다.
• 화합: 순방향 과정에서 트랜스포머는 모든 패치를 처리합니다. 연속성을 보장하기 위해 사용되는 겹치는 패치를 처리하기 위해, 최종 출력은 MultiDiffusion 샘플링 방법에서 영감을 받아 겹치는 영역의 값을 평균화하여 구성됩니다.
• 효율성: 조정자는 기본 모델보다 훨씬 작습니다 (예: 음악 생성에서 66M 대 405M 파라미터) 그리고 더 적은 훈련 에포크를 요구합니다.

주요 기여

1. 새로운 프레임워크: 작은 훈련 가능 조정자를 사용하여 사전 훈련된 확산 모델의 생성 능력을 더 큰 객체와 더 복잡한 조건으로 확장하는 DDE 방법의 도입.
2. 일반화 능력: 조정자가 훈련 데이터에 존재하는 것보다 더 큰 스케일의 객체 생성 및 조건 입력 처리로 일반화할 수 있음을 입증 ($L_{test} \ge L_{train}$).
3. 아키텍처 설계: 국소적 일관성을 겹침 평균화를 통해 강제하면서 장기 의존성을 학습하는 적합한 조정자로서 ViT 기반 아키텍처의 검증.
4. 효율성: 기본 모델을 재훈련하거나 더 큰 조정자 아키텍처를 사용하는 것에 비해 훨씬 적은 파라미터와 훈련 시간으로 고품질 생성을 달성.

실험 결과

저자들은 DDE 를 세 가지 도메인에서 평가했습니다.

1. 긴 오디오 트랙 생성 (Slakh2100)

• 작업: 길이 $l$ (12 초) 로 훈련된 기본 모델을 사용하여 길이 $10l$ (120 초) 의 음악 트랙 생성.
• 베이스라인: 연결 (Concatenation) 과 MultiDiffusion.
• 결과: DDE 는 4l 트랙에 대해 2.112의 프리체 오디오 거리 (FAD) 를 달성했고, 10l 트랙에 대해서는 2.142를 달성했습니다. 이는 10l 작업에서 MultiDiffusion (4.796) 과 연결 (4.596) 보다 훨씬 우수한 성능이었습니다. 조정자 (66M 파라미터) 는 기본 모델 (405M) 보다 훨씬 작았습니다.

2. 조건부 이미지 생성 (CLEVR)

• 작업: 기본 모델이 하나의 위치만 처리할 때, 여러 객체 위치 ($L$개의 레이블) 에 조건부로 이미지를 생성.
• 베이스라인: RRR (Reduce, Reuse, Recycle) 과 MultiDiffusion.
• 결과: DDE 는 여러 조건 제약 사항을 충족하는 데 있어 우수한 정확도를 보여주었습니다. 5 개 조건에 대해 DDE 는 (Heun 샘플러 기준) **44.5%**의 정확도를 달성했으며, 이는 RRR 의 33.6% 와 MultiDiffusion 의 36.3% 보다 높았습니다. 조정자는 $L_{train}=2$에서 훈련되어 $L_{test}=5$에서 테스트되는 것으로 일반화되었습니다.

3. 위성 이미지 생성

• 작업: 개략도 지도에 조건부로 큰 위성 이미지를 생성.
• 베이스라인: MultiDiffusion.
• 결과: DDE 는 분포 내 ($N=96$, FID 31.75 대 37.81) 와 분포 외 ($N=128$, FID 27.37 대 35.01) 크기 모두에서 MultiDiffusion 보다 낮은 FID 점수를 달성했습니다. 시각적 검토 결과, DDE 는 MultiDiffusion 샘플에서 관찰된 가시적인 이음새 없이 더 일관된 이미지를 생성했습니다.

중요성과 주장

이 논문은 DDE 가 재훈련의 금지적인 비용 없이 확산 모델을 확장하기 위한 실용적이고 효율적인 솔루션을 제공한다고 주장합니다. 사전 훈련된 프로세스를 조율하는 방법을 학습함으로써, 이 방법은 다음과 같은 것을 가능하게 합니다:

• 도메인 확장: 원래 모델의 능력을 넘어 더 큰 객체 (긴 오디오, 큰 이미지) 생성 및 더 복잡한 조건 (여러 레이블) 처리.
• 일반화: 조정자의 훈련 중 관찰된 것보다 더 큰 스케일로 외삽할 수 있는 능력.
• 품질과 일관성: ViT 기반 조정자가 장기 의존성을 효과적으로 학습하여 국소적 소거 단계 간의 충돌을 해결하고 전역적으로 일관된 출력을 생성.

저자들은 주요 한계는 조정자 모델을 훈련하기 위한 추가 데이터와 시간이 필요하다는 점이지만, 이 비용은 처음부터 전체 확산 모델을 훈련하는 것보다 훨씬 낮다고 지적합니다. 향후 작업으로는 단백질 서열 생성 및 짧은 클립에서 비디오 생성에 DDE 를 적용하는 것이 제안되었습니다.