기술 요약: 반지도 학습 하의 확산 모델에 대한 다목적 학습

1. 문제 제기

확산 모델은 이질적인 프롬프트 도메인 간의 텍스트-이미지 합성이나 여러 환경에 걸친 로봇 제어와 같은 다양한 작업을 위한 조건부 생성기로 점점 더 널리 배포되고 있습니다. 이러한 시나리오들은 본질적으로 단일 모델이 여러 목표 분포에 걸쳐 파레토 최적 성능을 달성해야 하는 다목적 학습 (Multi-Objective Learning, MOL) 문제를 구성합니다.

이 설정에서 핵심적인 과제는 "범용 (generalist)" 모델을 훈련시키는 통계적 비용입니다. 여러 목표를 처리하기 위해 범용 모델 클래스 ($\mathcal{F}$) 는 일반적으로 개별 작업에 필요한 "전문가 (specialist)" 모델 클래스 ($\mathcal{H}_k$) 보다 훨씬 큰 용량을 필요로 합니다. 완전 지도 학습 설정에서 범용 모델의 훈련에 필요한 표본 복잡도는 종종 prohibitively(부적절하게) 높은 자체 복잡도 ($C_\mathcal{F}$) 에 비례하여 증가합니다.

또한, 실제 배포 환경은 종종 비대칭 데이터 체제에 직면합니다:

• 쌍을 이룬 데이터 $(x, y)$ (예: 큐레이션된 텍스트 - 이미지 쌍 또는 전문가 로봇 시연) 는 희소하며 획득 비용이 높습니다.
• 비쌍을 이룬 조건 데이터 $y$ (예: 원시 텍스트 프롬프트 또는 상태 관찰) 는 풍부합니다.

본 논문은 다음과 같은 질문에 답합니다: 약한 지도 (레이블 없는 조건) 데이터를 활용하여 표본 효율적인 방식으로 파레토 최적 확산 모델을 학습할 수 있는가?

2. 방법론

저자들은 높은 용량의 범용 모델 필요성과 쌍을 이룬 데이터의 희소성 사이의 요구를 조화시키는 2 단계 반지도 학습 프레임워크를 제안합니다.

1 단계: 전문가 훈련

$K$개의 목표 (목표 분포 $P_k$) 각각에 대해, 가용한 제한된 쌍을 이룬 데이터 $\{(x_i^k, y_i^k)\}_{i=1}^n$만을 사용하여 경량의 전문가 스코어 함수 $h_k \in \mathcal{H}_k$가 훈련됩니다.

• 전문가 클래스 $\mathcal{H}_k$는 범용 클래스보다 복잡도가 현저히 낮다고 ($C_{\mathcal{H}_k} \ll C_\mathcal{F}$) 가정됩니다.
• 훈련은 쌍을 이룬 데이터에서 노이즈 제거 스코어 매칭 목적 함수를 최소화합니다.

2 단계: 의사 샘플링을 통한 범용 모델 훈련

모든 목표에 걸쳐 파레토 최적 집합을 근사하도록 대규모 범용 모델 $f \in \mathcal{F}$가 훈련됩니다.

• 의사 샘플 생성: 훈련된 전문가들 $\{h_k\}$은 의사 쌍을 이룬 데이터를 생성하는 데 사용됩니다. 풍부한 레이블 없는 조건 $y_i^k \sim P_Y^k$ 각각에 대해, 전문가가 의사 타겟 $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$을 생성합니다.
• 목적: 범용 모델은 원래 레이블이 있는 데이터와 생성된 의사 쌍을 이룬 데이터의 조합으로 훈련됩니다. 훈련 목적 함수는 $K$개 목표에 걸친 스코어 매칭 손실의 스칼라화 $S$를 최소화합니다:
  $$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$$
  여기서 $\tilde{L}_k$는 전문가의 편향을 보정하기 위한 보정 항을 포함하여 의사 샘플 위에서 계산된 경험적 손실입니다.

순차적 의사 결정으로의 확장

이 프레임워크는 모방 학습의 확산 정책으로 확장됩니다. 여기서 "조건"은 상태 $s$이고 "타겟"은 행동 $a$입니다. 저자들은 온-정책 롤아웃에 내재된 분포 이동을 해결합니다: 전문가 상태 분포를 직접 샘플링하기 어렵기 때문에, 2 단계는 환경 역학 하에서 학습된 전문가 정책에 의해 생성된 궤적을 사용하여 범용 정책을 위한 의사 데이터를 수집합니다.

3. 주요 이론적 기여

표본 복잡도 상한

본 논문은 제안된 반지도 접근법에 대한 비점근적 일반화 상한을 수립합니다.

• 주요 결과 (정리 3.2): 목표 오차 $\epsilon$을 달성하기 위해 필요한 쌍을 이룬 (레이블 있는) 샘플의 수는 전문가 클래스의 복잡도 ($C_{\mathcal{H}_k}$) 에만 의존하며, 범용 클래스 ($C_\mathcal{F}$) 에는 의존하지 않습니다.
  • 필요한 레이블 있는 샘플: $\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$.
  • 필요한 레이블 없는 조건: $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$.
• 비교: 반면, 범용 모델의 일반 완전 지도 학습은 $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$개의 레이블 있는 샘플을 필요로 합니다. 이는 $C_\mathcal{F} \gg C_\mathcal{H}$일 때 표본 효율성의 상당한 개선을 보여줍니다.
• 선형 스칼라화 (정리 3.5): 선형 (가중 합) 스칼라화와 볼록 함수 클래스의 경우, 저자들은 더 날카로운 속도를 유도하여 근사 오차에 대한 상한을 추가로 정제합니다.

분포 추정 및 정책 보장

• 분포 추정: 본 논문은 스코어 매칭 오차와 분포 추정 (총 변이 거리) 간의 간극을 연결합니다. 학습된 범용 모델이 목표 조건부 분포를 근사하는 데 파레토 최적 성능을 달성함을 증명합니다.
• 확산 정책 (정리 3.6): 저자들은 반지도 설정 하의 모방 학습에서 확산 정책의 비최적성 간격에 대한 최초의 이론적 보장을 제공합니다. 그들은 이 프레임워크가 온-정책 롤아웃으로 인한 분포 이동을 성공적으로 처리하면서도 전문가에서 범용으로 이어지는 파이프라인의 표본 효율성 이점을 유지함을 보여줍니다.

4. 실험 결과

저자들은 두 가지 다른 도메인에서 이론을 검증합니다:

1. 로봇 조작 (StackCube 및 PickCube):

   • 설정: 장면, 조명, 카메라 포즈와 같은 다양한 도메인 무작위화 하에서 확산 정책을 훈련합니다.
   • 결과: 반지도 MOL 방법은 특히 분포 외 (Out-of-Distribution, OOD) 설정에서 레이블만 있는 기준선보다 훨씬 뛰어난 성능을 보입니다 (예: StackCube Level 2 에서 성공률 +16%, Level 3 에서 +14%). 이는 의사 롤아웃이 분포 이동을 처리하는 데 중요한 지도를 제공함을 시사합니다.

2. 이미지 복원 (CelebA-HQ 인페인팅):

   • 설정: 세 가지 다른 마스크 유형 (교차 선, 얼굴 영역, 넓은 스트로크) 으로 얼굴을 인페인팅합니다.
   • 결과: 반지도 접근법은 모든 마스크 유형에 대해 레이블만 있는 기준선 대비 PSNR, SSIM, LPIPS 를 개선하여 생성적 이미지 작업에서 방법의 효과성을 입증했습니다.

5. 중요성 및 주장

본 논문은 반지도 체제 하의 확산 모델에 대한 다목적 학습을 위한 원칙적인 통계적 프레임워크를 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 이론과 실무의 연결: "전문가에서 범용으로" 이어지는 증류 파이프라인을 공식화하여, 레이블 데이터가 희소할 때 직접적인 다중 작업 학습보다 왜 이 접근법이 더 표본 효율적인지에 대한 이론적 근거를 제공합니다.
• 표본 효율성: 복잡한 범용 모델을 학습하는 통계적 비용이 자체 용량에서 분리될 수 있음을 증명하며, 대신 제한된 데이터로 훈련된 경량 전문가들의 복잡도에 의존함을 보여줍니다.
• 분포 이동 처리: 이러한 보장을 순차적 의사 결정 (모방 학습) 으로 확장하여 로봇공학의 중요한 실용적 과제를 해결합니다.

저자들은 겸손한 입장을 유지하며, 그들의 보장이 실현 가능성 가정 (진정한 스코어 함수가 전문가 클래스 내에 존재한다는 가정) 에 의존하며 현재 테스트된 특정 도메인 무작위화 수준을 넘어선 분포 외 (OOD) 일반화에는 아직 적용되지 않는다고 인정합니다. 그들은 이 작업을 생성 모델에 대한 반지도 MOL 의 통계적 속성을 이해하기 위한 기초적인 단계로 봅니다.