Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

기술 요약: 개념 발견을 통한 확산 모델의 테스트 시간 구성적 일반화

문제 정의

구성적 일반화 (熟悉한 부분을 새로운 구성으로 재조합하는 능력) 는 인간 지능의 특징이지만, 딥러닝 모델에게는 여전히 중요한 과제입니다. 신경망은 종종 훈련 분포 내에서 잘 보간하지만, 익숙한 원시 요소 (primitives) 가 낯선 조합으로 나타날 때는 어려움을 겪습니다. 확산 모델에서 구성적 생성에 대한 기존 접근법은 일반적으로 새로운 구성의 생성을 안내하기 위해 사전 정의된 개념 라이브러리, 조건부 신호, 또는 외부 레이블에 의존합니다. 이는 개념 발견과 구성적 일반화 사이의 간극을 만듭니다: 기존 방법들은 관련 원시 요소가 이미 알려져 있고 사용 가능하다고 가정하여, 고정된 어휘나 외부 감독 없이 테스트 시간에 질의에 특화된 개념을 처음부터 발견해야 하는 시나리오를 다루지 못합니다.

방법론

본 논문은 사전 훈련된 확산 모델을 테스트 시간에 개념을 발견하고 구성하기 위한 시간 인덱싱 밀도 모드의 암시적 계층 구조로 재사용하는 프레임워크를 제안합니다. 이 방법은 세 가지 주요 단계로 작동합니다:

1. 모드 상승을 통한 개념 발견:
   확산 모델을 역 샘플링에만 사용하는 대신, 저자들은 학습된 시간 인덱스 스코어 함수 $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ 를 활용하여 노이즈가 있는 주변 분포 $p_t(x_t)$ 의 국소 밀도 모드를 찾습니다. 단일 분포 외 (OOD) 질의 $x_q$ 에서 시작하여, 이 방법은 여러 노이즈 시간 단계 $t$ 에서 무조건부 스코어에 대한 경사 상승을 수행하여 국소 모드를 복원합니다.

   • 이론적 기반: 트위디 (Tweedie) 공식을 활용하여, 저자들은 이러한 모드가 노이즈가 제거된 "개념 원형 (concept prototypes)"에 해당하며, 노이즈 수준에 따라 다양한 추상화 수준에서 이산적인 중심점 역할을 함을 증명합니다.
   • 표현: 발견된 각 모드는 $x_0$ 공간으로 매핑되어 국소 가우시안 전문가 $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$ 로 근사됩니다. 여기서 평균은 모드에서 유도되고 공분산은 허친슨 (Hutchinson) 프로브를 통해 추정됩니다.

2. 전문가 곱 (Product-of-Experts, PoE) 구성:
   발견된 가우시안 전문가들은 Product-of-Experts (PoE) 공식을 사용하여 단일 분석적 교사 모델로 결합됩니다.

   • 선택: 탐욕적 서브모듈러 최적화 알고리즘이 질의와 가장 관련성이 높은 원형의 부분 집합을 선택하여, 차원별 로그 가능도 커버리지 목적 함수를 최대화합니다.
   • 구성: 선택된 전문가들은 가중치 PoE 분포 $q_T$ 로 결합됩니다. 전문가들이 가우시안이므로, 결과적인 구성도 폐쇄형 스코어 함수를 가진 가우시안이 됩니다. 이를 통해 새로운 구성에 대한 훈련된 조건부 토큰 없이도 분석적 노이즈 예측 $\epsilon_{q_T}(x_t, t)$ 를 유도할 수 있습니다.

3. 구성적 샘플링 및 증류:

   • 분석적 샘플링: PoE 교사의 분석적 스코어는 분류기 없는 안내 (CFG) 프레임워크에 통합됩니다. 모델은 발견된 구성을 향한 샘플링을 안내하기 위해 PoE 예측과 무조건부 예측을 대조합니다.
   • LoRA 증류: 매니폴드 내 생성을 개선하고 재사용 가능한 조건부를 만들기 위해, 이 방법은 선택적으로 PoE 교사를 기본 모델로 증류합니다. 이는 PoE 교사에 의해 생성된 샘플 풀에서 새로운 클래스 임베딩과 저랭크 적응 (LoRA) 어댑터를 훈련하는 과정을 포함합니다. 이 단계는 발견된 구성을 모델의 스코어 매니폴드로 다시 흡수합니다.

주요 기여

1. 발견과 구성의 연결: 이 연구는 사전 훈련된 확산 모델을 밀도 모드의 암시적 계층 구조로 취급함으로써 개념 발견과 구성적 일반화를 연결합니다. 이는 사전 정의된 원시 요소 풀 없이 테스트 시간에 이산적인 개념 원형을 복원할 수 있게 합니다.
2. 테스트 시간 구성 메커니즘: 저자들은 서브모듈러 가능도 목적 함수를 사용하여 관련 모드를 탐욕적으로 선택하고, 폐쇄형 스코어를 가진 PoE 교사 모델로 결합하는 방법을 도입하여, 특정 구성에 대한 기본 모델을 재훈련하지 않고도 직접 샘플링을 가능하게 합니다.
3. 실증적 검증: 이 프레임워크는 ColorMNIST 와 CelebA 두 가지 벤치마크에서 평가되었으며, 발견된 원형과 그들의 PoE 구성이 질의만 샘플링하거나 훈련된 가장 가까운 클래스 검색보다 미시적 구성에 대한 더 풍부한 정보를 인코딩함을 입증했습니다.

결과

ColorMNIST (숫자 정체성, 색상, 배경) 와 CelebA (얼굴 속성) 에 대한 실험은 제안된 방법이 여러 베이스라인보다 우수함을 보여줍니다:

• 베이스라인: 이 방법은 "질의만 (Query-only)"(고정 분산으로 질의 주변 샘플링), "Top-1/Top-3"(가장 가까운 훈련된 클래스 검색), 그리고 표준 CFG 와 비교됩니다.
• 성능: 분석적 PoE 샘플러와 LoRA 증류 변형 모두 베이스라인에 비해 우수한 FID, CLIP 유사도, 그리고 정밀도/재현율 지표를 달성합니다.
  • ColorMNIST 에서 LoRA 변형은 충실도 (Faithfulness) 기준 F1 점수 87.6% 와 일반화 (Generalization) 기준 76.3% 를 달성하여 Top-1 베이스라인 (~6.5% F1) 을 크게 앞섰습니다.
  • CelebA 에서 PoE 방법은 최고의 일반화 FID(140.3) 와 CLIP 점수를 달성했으며, LoRA 는 커버리지와 충실도의 균형을 추가로 개선하여 최고의 일반화 재현율과 F1 을 달성했습니다.
• 정성적 분석: 시각적 결과는 베이스라인이 종종 속성의 특정 조합을 유지하지 못함 (예: 숫자 정체성 변경 또는 다른 속성 조합으로 이동) 에 비해, 제안된 방법이 OOD 구성을 성공적으로 재구성하고 동일한 클래스의 다른 인스턴스로 일반화함을 나타냅니다.

중요성과 주장

본 논문은 확산 모델의 시간 인덱스 스코어 기하학이 사전 정의된 개념 라이브러리 없이 테스트 시간 구성적 생성을 지원하는 재사용 가능한 밀도 모드 개념을 포함한다고 주장합니다.

• 이론적 통찰: 결과는 확산 모델이 암묵적으로 개념의 계층 구조를 인코딩하며, 노이즈가 있는 주변 분포의 모드가 다양한 추상화 수준에서 원형으로 작용함을 시사합니다.
• 실용적 함의: 이 방법은 단일 OOD 질의에서 새로운 구성을 생성할 수 있게 하여, 외부 감독이나 사전 지정된 어휘에 의존하지 않는 구성적 일반화를 위한 경로를 제공합니다.
• 한계: 저자들은 프로토콜이 원시 요인이 알려져 있다고 가정하며 (표준 구성 설정), 이 방법이 사전 훈련된 백본의 귀납적 편향에 의존한다고 겸손하게 지적합니다. 알려진 원시 요소 간 보간은 가능하지만, 훈련 지원 범위 (예: 어두운 배경만 훈련된 모델에서 흰색 배경) 를 넘어선 진정한 새로운 원시 요소를 창출할 수는 없습니다.

이 연구는 확산 모델을 개념 발견과 구성을 위해 재사용할 수 있음을 이론적 및 실증적 증거로 제시하며, 언어 모델링과 로봇 제어 등 더 넓은 응용 분야로 확장될 가능성을 보여줍니다.