기술 요약: 확산 샘플링을 위한 ART (Adaptive Reparameterized Time)

문제 정의

확산 모델은 학습된 역방향 시간 역학(reverse-time dynamics)을 통해 노이즈를 샘플로 변환함으로써 데이터를 생성한다. 스코어(또는 디노이징) 모델의 학습은 잘 정립되어 있으나, 샘플링 단계는 유한한 타임스텝 그리드 상에서 역방향 과정을 수치적으로 이산화하는 것에 의존한다. 현재의 표준 관행은 균등 그리드(uniform grids) 또는 수작업으로 설계된 스케줄(예: EDM, DPM-Solver)을 활용한다. 이러한 접근 방식은 대개 임의적이고 고정된 처방이며, 특정 궤적의 수치적 특성에 적응하지 못한다. 결과적으로, 이들은 계산 자원을 최적으로 할당하지 못하여, 단계가 너무 성기면 상당한 이산화 오차(discretization error)가 발생하거나, 역학이 매끄러운 영역에서 단계가 너무 조밀하면 불필요한 함수 호출(function evaluations)이 낭비되는 문제를 초래할 수 있다. 핵심 문제는 고정된 계산 예산(함수 호출 횟수) 하에서 샘플 품질을 최적화하기 위한 체계적이고 원칙적인 타임스텝 할당(timestep allocation) 방법을 공식화하는 것이다.

방법론

1. 적응형 재매개변수화 시간 (Adaptive Reparameterized Time, ART)

저자들은 타임스텝 선택을 제어 문제로 취급하는 연속 시간 최적 제어 프레임워크인 ART를 제안한다.

• 시간 재매개변수화: 물리적 시간 $\tau \in [0, T]$ 상에서 역방향 과정을 진화시키는 대신, ART는 보조적인 "샘플링 시계(sampling clock)" $t \in [0, T]$를 도입한다. 물리적 시간은 이 시계의 함수인 $\tau = \psi(t)$이다.
• 제어 변수: 샘플링 시계에 대한 물리적 시간의 진행 속도 $\theta(t) = \dot{\psi}(t)$를 제어 변수로 취급한다.
• 목표: **리딩 오더 오일 이산화 오차(leading-order Euler discretization error)**로부터 유도된 목적 함수를 최대화하는 것이다. 오일 스킴의 국소 오차는 $|Q(x, \psi)| \theta^2$에 비례하며, 여기서 $Q$는 확률 흐름 벡터 필드의 국소적 강성(stiffness)을 포착한다. 제어 문제는 총 시간 예산 제약 조건 $\int_0^T \theta(t) dt = T$를 준수하면서 누적 오차를 최소화하는 것을 목표로 한다.
• 공식화: 이는 상태 변수에 샘플 궤적과 누적 물리 시간이 포함되며, 제어 변수 $\theta$가 물리적 시간을 통해 얼마나 빠르게 또는 느리게 이동할지를 결정하는 결정론적 최적 제어 문제로 귀결된다.

2. ART-RL: 무작위 제어 공식화

해밀턴-자코비-벨만(HJB) 방정식을 통해 결정론적 ART 문제를 직접 해결하는 것은 상태 공간의 고차원성(예: 이미지 차원) 때문에 실행 불가능하다. 이를 극복하기 위해 저자들은 보조적인 무작위 제어 문제인 ART-RL을 도입한다.

• 가우시안 정책 (Gaussian Policies): 결정론적 제어 $\theta$는 확률적 정책 $\pi(\lambda)$, 구체적으로 가우시안 분포 $N(\mu, \lambda/|Q|)$로 대체된다. 분산은 국소 강성 $|Q|$에 반비례하며, 이는 강성이 높은 영역에서 무작위성을 줄여준다.
• 이론적 등가성: 논문은 ART-RL 공식화의 최적 가우시안 정책의 평균이 원래의 결정론적 ART 문제의 최적 제어 해와 동일함을 증명한다. 따라서 무작위 문제를 해결함으로써 결정론적 문제의 해를 회복할 수 있다.
• 액터-크리틱 학습 (Actor–Critic Learning): 최근의 연속 시간 강화 학습(CTRL) 발전을 활용하여, 저자들은 구현 가능한 액터-크리틱 업데이트 규칙을 도출한다.
  • 크리틱 (Critic): 마팅게일 기반 모멘트 조건을 사용하여 가치 함수(정책 평가)를 추정한다.
  • 액터 (Actor): 크리틱의 피드백을 바탕으로 가치를 개선하기 위해 정책의 평균 $\mu$를 업데이트한다(정책 개선).
  • 라그랑주 승수 (Lagrange Multiplier): 총 시간 예산 제약을 강제하기 위해 승수 $\gamma$를 업데이트한다.

3. 구현 및 증류 (Distillation)

학습된 정책은 시간, 상태, 누적 시간의 피드백 함수이다. 그러나 경험적 결과에 따르면 실제로는 최적 제어 $\theta$가 주로 시간에 의존한다. 저자들은 학습된 신경망 정책을 고정된 결정론적 시간 전용 스케줄(제어의 경험적 평균)로 교체하는 증류(distillation) 단계를 제안한다. 이를 통해 학습된 스케줄을 단순한 사전 계산된 그리드로 배포할 수 있으며, 수작업으로 설계된 스케줄과 비교했을 때 추론 시 추가적인 오버헤드가 발생하지 않는다.

주요 기여

1. 제어 이론적 프레임워크: 본 논문은 타임스텝 할당을 휴리스틱한 그리드 설계를 넘어 연속 시간 최적 제어 문제(ART)로 공식화한 최초의 연구이다.
2. 이론적 가교 (ART-RL): 결정론적 제어 문제와 무작위 강화 학습 공식화 사이의 엄격한 연결 고리를 구축한다. 최적의 무작위 정책의 평균이 원래의 결정론적 문제를 해결함을 증명하여 액터-크리틱 알고리즘의 사용을 가능하게 한다.
3. 알고리즘 개발: 타임스텝 재매개변수화를 위한 구체적인 정책 평가 및 개선 특성을 도출하여, 최적 스케줄을 학습하기 위한 실용적이고 데이터 기반의 액터-크리틱 알고리즘(Algorithm 1)을 개발하였다.
4. 일반화 및 전이 가능성: 학습된 스케줄이 특정 데이터셋(CIFAR-10)에서 학습되었더라도, 재학습 없이 다른 데이터셋(AFHQv2, FFHQ, ImageNet), 다른 해상도(64x64에서 512x512), 다른 파이프라인(EDM에서 EDM2), 다른 표현 공간(픽셀에서 잠재 공간)으로 전이될 수 있음을 입증하였다.

실험 결과

저자들은 스코어 모델, 솔버, 파이프라인은 고정하고 타임스텝 그리드만 변경하여 다양한 설정에서 ART-RL을 평가하였다:

• 저차원 분석 사례: 알려진 스코어 함수를 가진 1D 확산 모델을 사용하여, ART-RL은 Wasserstein-2 오차 측면에서 Uniform, DPM, EDM 스케줄보다 일관되게 우수한 성능을 보였다. 이는 이미지용으로 설계된 수작업 스케줄이 다른 도메인에서는 실패할 수 있음을 보여준다.
• CIFAR-10 (EDM 파이프라인):
  • Heun 샘플링: ART-RL은 모든 예산 범위에서 가장 우수한 Fréchet Inception Distance(FID)를 달성하였으며, 가장 강력한 구성(NFE=35)에서도 SOTA인 EDM 스케줄(1.82 vs 1.85)을 능가하였다.
  • Euler 샘플링: 1차 솔버를 사용하는 경우에도 유사한 개선이 관찰되어, 그리드를 학습하는 것의 이점을 확인하였다.
• MNIST (소규모 모델): 사전 학습되지 않은 경량 스코어 모델과 RK4 솔버를 사용하더라도 ART-RL은 베이스라인을 크게 앞질렀으며, 이는 모델 용량에 대한 방법론의 견고함을 보여준다.
• 전이 학습 (Transfer Learning):
  • 단계 수 간의 전이: $K=18$에서 학습된 스케줄은 로그-선형 보간/외삽을 통해 다른 단계 수로 효과적으로 전이된다.
  • 데이터셋 간의 전이: CIFAR-10에서 학습된 스케줄은 재학습 없이 AFHQv2, FFHQ, ImageNet-64의 FID를 개선한다.
  • 파이프라인 및 해상도 간의 전이: 학습된 스케줄은 ImageNet-512(잠재 공간, 고해상도)의 EDM2 파이프라인으로 성공적으로 전이되어, EDM 및 DPM 베이스라인보다 우수한 성능을 보였다.

의의 및 주장

본 논문은 ART-RL이 기존의 임의적이고 고정된 휴리스틱 그리드에 대한 원칙적이고 데이터 기반의 대안이라고 주장한다.

• 재사용성: 주요한 실무적 의의는 재사용 가능한 타임스텝 스케줄을 학습할 수 있다는 점이다. 한 번 학습 및 증류된 스케줄(예: CIFAR-10)은 오프라인 학습 비용을 분담하면서, 재학습 없이 다양한 다운스트림 작업(다른 데이터셋, 해상도, 파이프라인)에 적용될 수 있다.
• 드롭인 교체 (Drop-in Replacement): 이 방법은 스코어 모델 구조, 솔버 구현, 또는 학습 프로토콜의 변경을 요구하지 않는다. 순수하게 타임스텝 그리드를 대체하는 방식으로 작동하여, 기존 확산 파이프라인에 모듈형 개선을 제공한다.
• 이론적 근거: 기존의 휴리스틱 접근 방식과 달리, ART는 최적 제어 이론과 연속 시간 강화 학습에 근거하고 있어 확산 샘플링의 이산화 오차를 다루는 체계적인 방법을 제시한다.

저자들은 현재의 분석이 확률 흐름 ODE 샘플링에 국한되어 있으며, 시간 전용 스케줄로의 증류가 상태 의존성을 제거한다는 점을 언급하며, 더 풍부한 상태 조건부 스케줄이 더 큰 이득을 줄 수 있는지에 대한 과제를 남겨두었다. 그러나 현재의 결과는 체계적인 타임스텝 설계를 위한 강력한 토대를 마련하였다.