기술적 요약: 샘플 기반 변분 추론을 통한 소수 단계 생성 모델의 정렬

문제 정의

보상 함수나 인간 선호도와 같은 특정 목적과 소수 단계 생성 모델 (예: GAN, 일관성 모델, 흐름 지도, 드리프트 모델) 을 정렬하는 것은 상당한 도전 과제를 제시합니다. 기존 정렬 프레임워크는 종종 소수 단계 생성기가 위반하는 제한적인 구조적 가정에 의존합니다:

1. 처리 가능한 가능도 (Tractable Likelihoods): 많은 방법들은 GAN 과 같은 암시적 모델이나 1 단계 매핑에서는 사용할 수 없는 탈노이즈 궤적에 대한 명시적 가능도 평가를 요구합니다.
2. 특정 솔버: 확률적 최적 제어 (SOC) 또는 강화 학습 (RL) 기반 접근법은 종종 특정 상미분 방정식 (ODE) 또는 확률 미분 방정식 (SDE) 솔버와 순차적 탈노이즈 과정을 가정합니다.
3. 모델 특이성: 공식들은 종종 확산 모델과 같은 특정 모델 계열에 맞춰져 있어 VAE 나 흐름 지도와 같은 다양한 소수 단계 아키텍처로 일반화되지 못합니다.

결과적으로, 핵심 이점인 빠르고 단일 단계 추론을 유지하면서 다양한 소수 단계 생성기를 미세 조정할 수 있는 일반적인 정렬 프레임워크가 부재합니다.

방법론: FAV

저자들은 생성 모델 계열 및 샘플링 역학으로부터 정렬을 분리하는 일반적인 프레임워크인 **FAV (Sample-based Variational Inference 를 통한 소수 단계 생성 모델 정렬)**를 제안합니다. FAV 는 정렬을 기준 분포에 고정된 **보상 기울어진 분포 (reward-tilted distribution)**로부터의 샘플링으로 공식화합니다:
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
여기서 $p_{\text{ref}}$는 기준 분포 (예: 사전 훈련된 생성기 또는 경험적 데이터) 이고 $r(x)$는 목적 함수입니다.

이 방법은 세 가지 주요 단계로 진행됩니다:

1. SVGD 를 통한 샘플 기반 근사

FAV 는 **Stein 변분 경사 하강 (SVGD)**을 사용하여 목표 분포 $q^*$를 근사합니다. 밀도 평가를 요구하는 대신, SVGD 는 파티클로 분포를 표현하고 이를 목표 방향으로 반복적으로 수송합니다. 최적의 Stein 속도장 $\phi^*$는 KL 발산의 감소를 최대화하도록 유도됩니다:
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
이 장치는 세 가지 구성 요소를 포함합니다:

• 사전 정렬: 샘플을 기준 분포 $p_{\text{ref}}$ 방향으로 유도합니다.
• 보상 안내: 샘플을 높은 보상 영역으로 밀어냅니다.
• 다양성 강제: 커널 기울기 항을 통해 모드 붕괴를 방지합니다.

2. 비모수 점수 추정 (KDE)

주요 과제는 경험적 데이터나 암시적 생성기의 경우 점수 $\nabla \log p_{\text{ref}}$가 종종 처리 불가능하다는 점입니다. FAV 는 **커널 밀도 추정 (KDE)**을 사용하여 점수를 비모수적으로 추정함으로써 이를 해결합니다. 기준 분포를 가우시안 커널과 컨볼루션함으로써, 점수는 기준 샘플과 현재 파티클 간의 차이의 가중 평균으로 근사되어 명시적 가능도에 대한 필요성을 효과적으로 우회합니다.

3. 고정점 회귀를 통한 상각 (Amortization)

소수 단계 모델의 빠른 추론 속도를 유지하기 위해, FAV 는 추론 시 반복적인 SVGD 업데이트를 수행하지 않습니다. 대신, 고정점 회귀를 통해 생성기 매개변수 $\theta$로 수송을 **상각 (amortizes)**합니다. 생성기는 "수송된" 목표와 출력 간의 거리를 최소화하도록 훈련됩니다:
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
여기서 수송된 샘플은 안정적인 회귀 목표 역할을 합니다. 이를 통해 생성기는 단일 순전파 (forward pass) 에서 정렬된 샘플을 생성하도록 학습할 수 있습니다.

블랙박스 보상으로의 확장: 미분 불가능한 보상의 경우, FAV 는 $\nabla r(x)$를 근사하기 위해 0 차 순 미분 추정 (가우시안 평활화) 을 사용하여 보상 함수가 블랙박스일 때에도 정렬이 가능하도록 합니다.

주요 기여

1. 일반 정렬 프레임워크: FAV 는 처리 가능한 가능도, 특정 ODE/SDE 솔버, 또는 다단계 궤적이 필요 없이 VAE, GAN, 일관성 모델, 흐름 지도를 포함한 다양한 소수 단계 생성기를 정렬하는 최초의 프레임워크입니다.
2. 샘플 기반 VI 공식화: SVGD 와 KDE 를 사용하여 보상 기울어진 분포로부터의 샘플링으로 정렬을 제시함으로써, FAV 는 정렬 절차를 생성기의 내부 역학으로부터 분리합니다.
3. 효율성 유지: 고정점 회귀를 통한 상각을 통해 FAV 는 기본 생성기의 단일 단계 추론 효율성을 유지하며, 반복적 샘플링 또는 증류 과정의 계산 오버헤드를 피합니다.
4. 블랙박스 호환성: 0 차 순 미분 추정 통합을 통해 FAV 는 임의의 블랙박스 보상을 최적화할 수 있습니다.

실험 결과

저자들은 FAV 를 로봇 조작 및 이미지 생성기 정렬이라는 두 가지 주요 도메인에서 평가했습니다.

1. 로봇 조작 (오프라인 및 오프라인 - 온라인 RL)

• 벤치마크: 56 개 오프라인 RL 작업 (OGBench 및 D4RL) 과 30 개 오프라인 - 온라인 RL 작업.
• 성능: FAV 는 재파라미터화 정책 경사, Adjoint Matching, 증류된 생성 정책 (예: FQL, IFQL) 을 포함한 10 개 베이스라인을 능가하며 모든 작업에서 최첨단 성능을 달성했습니다.
• 주요 발견: FAV 는 보조 생성 모델이나 다단계 탈노이즈가 없더라도 우수한 정책 추출 능력과 오프라인에서 온라인 학습으로의 전환 성능을 보여주었습니다.

2. 이미지 생성기 정렬

• 규모: 실험은 ImageNet-256(256x256) 에서 고해상도 텍스트 - 이미지 합성 (1024x1024) 까지 다양했습니다.
• 모델: StyleGAN-XL, 드리프트 모델, iMeanFlow, SANA-Sprint 에서 테스트되었습니다.
• 지표: 미적 점수, 인간 선호도 점수 (HPS), 안전 감지기 (NSFW) 에 대해 평가되었습니다.
• 성능:
  • FAV 는 DRaFT 및 Adjoint Matching 과 같은 베이스라인보다 보상 과최적화를 완화하면서 보상 최적화와 샘플 다양성/품질 간의 가장 좋은 파레토 프론티어를 일관되게 달성했습니다.
  • 추론 효율성: FAV 는 경쟁력 있는 이미지 품질을 유지하면서 Best-of-N, ReNO 와 같은 추론 시 정렬 방법보다 약 180 배에서 280 배 더 빠르게 정렬된 샘플을 생성했습니다.
  • 확장성: 이 방법은 256x256 에서 1024x1024 해상도로 성공적으로 확장되었으며 다양한 모델 아키텍처에서 일반화되었습니다.

중요성 및 주장

이 논문은 FAV 가 소수 단계 생성 모델을 정렬하기 위한 통합적이고 모델에 구애받지 않는 솔루션을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 구조적 장벽 제거: 처리 가능한 가능도나 특정 샘플링 궤적에 대한 필요성을 제거하여 더 넓은 범위의 암시적 생성 모델에 정렬을 접근 가능하게 만듭니다.
• 효율성과 품질의 균형: 고보상 생성과 소수 단계 추론의 계산 효율성 사이의 트레이드오프를 성공적으로 조정하며, 이전 방법들은 종종 성능 저하나 복잡한 증류를 요구했던 과제를 해결합니다.
• 견고성: 이 프레임워크는 저차원 제어부터 고해상도 이미지 합성까지 다양한 작업에서 견고성을 보여주며, 미분 가능 및 블랙박스 보상 함수 모두를 처리합니다.

저자들은 FAV 가 RL 의 정책 추출과 이미지 생성의 보상 정렬을 위한 강력하고 범용적인 접근법을 제공하며, 더 복잡하고 가정이 많은 방법들에 대한 실용적인 대안이라고 결론지었습니다.