ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

기술 요약: 확산 샘플링을 위한 ART

문제 제기
점수 기반 확산 모델은 학습된 역시간 동역학을 적분하여 샘플을 생성합니다. 이 과정에서 중요한 수치적 선택은 시간 범위를 시간 단계 그리드로 이산화하는 것입니다. 표준 접근법은 일반적으로 균일 그리드나 수동으로 설계된 스케줄 (예: EDM, DPM-Solver) 을 사용합니다. 저자들은 고정된 스케줄이 함수 평가 횟수 (NFE) 가 고정된 경우 종종 최적이지 않다고 주장합니다. 균일 그리드는 역동역학이 궤적을 따라 크게 변할지라도 계산 노력을 균등하게 배분하는 반면, 수동으로 설계된 그리드는 특정 데이터 분포나 점수 함수의 복잡성에 적응하지 못할 수 있습니다. 목표는 샘플링 궤적을 따라 계산 노력을 체계적으로 재분배하여 전체 오일러 이산화 오차를 최소화하는 것입니다.

방법론
본 논문은 최적 시간 단계 스케줄링을 위한 제어 이론적 프레임워크인 **적응형 재매개변수화 시간 (Adaptive Reparameterized Time, ART)**을 소개합니다.

1. 제어 공식화 (ART):

   • ART 는 그리드 점을 직접 선택하는 대신 확산 과정의 "클록 속도"를 제어 변수로 취급합니다.
   • 연속적인 매핑 $\psi(t)$를 통해 원래 확산 시간 $\tau$에 매핑된 재매개변수화된 시간 변수 $t$를 도입합니다.
   • 제어 $\theta(t) = \dot{\psi}(t)$는 확산 시간의 순간 변화율을 나타냅니다.
   • 목적은 총 시간 예산이 보존된다는 제약 조건 ($\int_0^T \theta(t) dt = T$) 하에 전체 국소 오일러 이산화 오차를 최소화하도록 $\theta(t)$를 선택하는 것입니다.
   • 국소 오차 대리 변수는 $|Q(x, \psi)| \theta(t)^2$에 비례하는 것으로 유도되며, 여기서 $Q$는 점수 함수와 그 도함수에 의존합니다.

2. 강화 학습 솔루션 (ART-RL):

   • 결정론적 ART 문제는 고차원 상태 공간과 해밀턴 - 야코비 - 벨만 (HJB) 방정식에 대한 폐쇄형 해의 부재로 인해 직접 해결하기 어렵습니다.
   • 저자들은 연속 시간 강화 학습 (CTRL) 접근법인 ART-RL을 제안합니다.
   • 제어 $\theta$를 가우시안 정책 $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$에서 샘플링하는 무작위화된 보조 문제를 도입합니다.
   • 이 공식화는 Jia 와 Zhou [14] 의 프레임워크를 특히 적응시킨 액터 - 크리틱 알고리즘을 사용하여 정책 평균 $\mu$와 가치 함수를 확률적 근사를 통해 학습할 수 있게 합니다.

3. 이론적 연결고리:

   • 핵심 이론적 기여는 결정론적 ART 최적해와 ART-RL 보조 문제의 최적 가우시안 정책 사이의 양방향 연결고리를 확립하는 것입니다.
   • 충분성: 결정론적 최적 제어 $\mu^*$는 평균이 정확히 $\mu^*$인 최적 가우시안 정책으로 승격됩니다.
   • 필요성: 보조 문제에 대한 임의의 최적 가우시안 정책은 결정론적 ART 제어 $\mu^*$를 회복하는 평균을 가져야 합니다.
   • 이 결과는 액터를 가우시안 평균으로 매개변수화하고 훈련 후 결정론적으로 배포하는 것을 정당화하며, 학습된 스케줄이 원래 제어 문제에 대한 진정한 최적해임을 엄격하게 보장합니다.

4. 증류:

   • 오프라인 훈련 후, 확률적 정책은 결정론적이고 시간만의 그리드로 증류됩니다. 이는 추론 중 단계별 액터 네트워크 평가나 분산 계산의 필요성을 제거하여 표준 스케줄러에 비해 추가적인 계산 비용이 전혀 없음을 보장합니다.

주요 기여

• 공식화: 적응형 시간 재매개변수화를 통한 확산 시간 단계 스케줄링을 위한 최초의 제어 이론적 프레임워크인 ART 를 제안합니다.
• 이론: 결정론적 최적 제어 문제와 가우시안 정책을 가진 연속 시간 액터 - 크리틱 학습 절차 사이의 양방향 동등성을 확립합니다. 이는 휴리스틱 RL 적용을 결정론적 최적해를 찾는 원칙적인 경로로 변환합니다.
• 알고리즘: 데이터 주도적으로 적응형 시간 스케줄을 학습하는 액터 - 크리틱 알고리즘인 ART-RL 을 개발합니다.
• 일반화: 학습된 스케줄이 재훈련 없이 다른 데이터셋과 시간 단계 수 간에 증류되어 이전될 수 있음을 입증합니다.

실험 결과
저자들은 공식 EDM 파이프라인 내에서 CIFAR-10, AFHQv2, FFHQ, ImageNet 에 대해 ART-RL 을 평가합니다.

• CIFAR-10: ART-RL 은 광범위한 NFE 예산 범위에서 균일, DPM-Solver log-SNR, 표준 EDM 스케줄보다 일관되게 우수한 성능을 보입니다.
  • 낮은 예산 (예: NFE=3) 에서 ART-RL 은 기준선들이 여전히 노이즈 근처에 있을 때 인식 가능한 객체를 생성합니다.
  • NFE=35 에서 ART-RL 은 FID 1.82 를 달성하여 EDM(1.85) 과 DPM-logSNR(1.89) 보다 개선된 성능을 보입니다.
  • 여러 시드 (seeds) 를 통해 통계적 유의성이 확인되었으며, ART-RL 의 최악의 실행이 EDM 의 최상위 실행을 능가했습니다.
• 일반화:
  • 보간/외삽: CIFAR-10 에서 $K=18$단계로 학습된 스케줄은 다른 단계 수 ($K \in \{4, 9, 12, 15, 20\}$) 로 보간되거나 외삽될 때 성능 우위를 유지합니다.
  • 교차 데이터셋 이전: CIFAR-10 에서 증류된 스케줄은 재훈련 없이 AFHQv2, FFHQ, ImageNet 에 적용됩니다. 이는 특히 낮고 중간 NFE 에서 EDM 과 DPM-logSNR 보다 FID 를 일관되게 개선합니다.
• 1 차원 분석: 분석적 점수 함수를 가진 설정에서 ART-RL 은 수동으로 설계된 스케줄에 비해 워슈틴 거리 ($W_2$) 를 현저히 감소시켜, 개선이 점수 추정 오차가 아닌 더 나은 시간 할당에서 비롯됨을 확인합니다.

의의 및 주장
본 논문은 적응형 시간 재매개변수화를 통해 확산 시간 단계 스케줄링을 제어 문제로 공식화한 최초의 연구이며, 결정론적 최적해와 액터 - 크리틱 학습 절차 사이의 엄격한 양방향 연결고리를 제공한 최초의 연구라고 주장합니다.

저자들은 그들의 방법이 데이터 주도적이고 **상각 (amortized)**적이라고 강조합니다:

• 데이터 주도적: 스케줄은 수동으로 설계된 것이 아니라 목적 함수와 데이터 동역학에서 학습됩니다.
• 상각된 이전: 오프라인 훈련 비용 (CIFAR-10 의 경우 Colab T4 에서 약 1~2 시간) 은 일회성 비용입니다. 결과적으로 생성된 결정론적 그리드는 추가 추론 비용 없이 다른 데이터셋과 단계 수 간에 재사용 가능하며, 균일이나 EDM 과 같은 표준 스케줄러의 실행 시간과 일치합니다.

이 연구는 적응형 시간 재매개변수화를 생성형 확산 모델에서 체계적인 스케줄 설계를 위한 신흥하지만 유망한 방향으로 위치시키며, 향후 연구를 위한 제어 이론적 출발점을 제공합니다.