Gradient-Free Noise Optimization for Reward Alignment in Generative Models

기술 요약: ZeNO (Zeroth-order Noise Optimization)

문제 제기
생성 모델 (예: 확산 모델 및 흐름 모델) 을 위한 기존 보상 정렬 방법들은 일반적으로 다단계 확률적 궤적을 최적화에 의존합니다. 이러한 설정에서는 효과적이지만, 동일한 궤적 구조를 자연스럽게 노출하지 않는 결정론적 또는 단일 단계 생성기 (예: 일관성 모델, 증류된 확산 모델) 로는 그러한 접근법을 확장하기 어렵습니다. 자연스러운 대안은 입력 노이즈를 직접 최적화하는 것입니다. 그러나 표준 경사 기반 노이즈 최적화는 생성기와 보상 함수 모두를 통한 역전파를 필요로 합니다. 이는 미분 가능한 파이프라인에만 적용 가능성을 제한하며, 경사가 의미 있는 의미론적 개선보다는 가짜 국소 민감도를 반영할 때 신뢰성이 떨어질 수 있습니다. furthermore, 많은 실용적인 보상 함수 (예: 단백질 설계 가능성, 구성적 텍스트 - 이미지 정렬) 는 미분 불가능하거나 블랙박스이므로, 경사 기반 방법을 실행 불가능하게 만듭니다.

방법론: ZeNO
저자들은 ZeNO(Zeroth-order Noise Optimization) 를 제안합니다. 이는 경로 적분 제어 (PIC) 문제로 노이즈 최적화를 공식화하는 경사 프리 프레임워크입니다. 핵심 아이디어는 입력 노이즈의 진화를 제어 가능한 확률 과정으로 간주하고, 오직 영차수 (zeroth-order) 보상 평가만을 사용하여 이를 최적화하는 것입니다.

1. 경로 적분 제어 공식화:
   이 방법은 노이즈 최적화를 제어된 확률 미분 방정식 (SDE) 으로 모델링합니다. 목표는 제어 노력을 정규화하면서 노이즈 궤적을 높은 보상을 주는 말단 상태가 도출되는 영역으로 유도하는 것입니다. 최적 제어 $u^*(z, t)$는 경로 적분 형태로 유도됩니다:
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   결정적으로, 이 추정기는 말단 비용 평가 (보상) 에만 의존하며 생성기 $G_\theta$나 보상 함수 $r$을 통한 미분을 필요로 하지 않습니다.

2. 참고 동역학 (Ornstein–Uhlenbeck):
   업데이트된 노이즈가 사전 훈련된 생성기와 호환되도록 보장하기 위해, ZeNO 는 참조 동역학으로 Ornstein–Uhlenbeck (OU) 과정을 채택합니다. 분산이 폭발하는 SDE 와 달리, OU 과정은 대부분의 사전 훈련된 생성기의 사전 분포와 일치하는 가우시안 정상 분포 $N(0, I)$을 유지합니다. 이 선택은 노이즈 업데이트를 보상 기울어진 분포 $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$를 목표로 하는 랑주뱅 동역학 (Langevin dynamics) 과 암묵적으로 연결합니다.

3. 선형화된 추정기 및 분산 감소:
   표준 경로 적분 추정기는 지수 가중치로 인해 높은 분산을 겪습니다. 저자들은 평균 보상 $\bar{r}$을 기준선으로 빼고 작은 보상 분산 가정 하에 지수 항을 선형화함으로써 선형화된 근사 (명제 1) 를 도입합니다:
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   이 추정기는 자연 진화 전략 (Natural Evolution Strategies) 에서 사용되는 스코어 함수 추정기와 일치하며, 실험적으로 지수 변형보다 우수한 성능을 보입니다.

4. 단일 단계 후퇴 시야 (One-Step Receding-Horizon):
   장기 시야 롤아웃의 계산 비용을 피하기 위해, ZeNO 는 단일 단계 후퇴 시야 근사 ($H=1$) 를 사용합니다. 각 반복에서 $N$개의 교란된 노이즈 후보가 참조 동역학에서 샘플링되고, 보상 모델에 의해 평가되며, 선형화된 제어 추정을 사용하여 현재 노이즈 상태를 업데이트합니다. 이는 생성기를 수정하지 않고도 파티클 수 ($N$) 또는 반복 횟수 ($M$) 를 증가시킴으로써 추론 시간 확장 (inference-time scaling) 을 가능하게 합니다.

주요 기여

• 경사 프리 정렬: ZeNO 는 생성기나 보상 모델을 통한 역전파 없이 결정론적 및 단일 단계 생성기에 대한 보상 정렬을 가능하게 합니다.
• 블랙박스 호환성: 이 방법은 미분 불가능하고 블랙박스인 보상 함수에 적용 가능하며, 기존 경사 기반 노이즈 최적화 방법 (예: ORIGEN, ReNO) 이 실패하는 설정에서 작동합니다.
• 이론적 연결: 이 프레임워크는 영차수 노이즈 최적화와 보상 기울어진 분포를 목표로 하는 랑주뱅 샘플링 사이의 이론적 연결을 수립하여 업데이트 규칙에 대한 원칙적인 정당성을 제공합니다.
• 추론 시간 확장: ZeNO 는 계산 예산 ($N$ 및 $M$을 통해) 과 향상된 보상 성능 사이의 트레이드오프를 가능하게 하는 유연한 메커니즘을 도입합니다.

실험 결과
저자들은 다양한 생성기 (SDXL-Turbo, DMD2, LCM, DFGAN) 와 보상 함수 (미적분 점수, PickScore, GenEval, 단백질 설계 가능성) 에 대해 ZeNO 를 평가했습니다.

• 이미지 생성: ZeNO 는 테스트된 모든 생성기와 보상 유형에 걸쳐 보상 점수를 일관되게 향상시킵니다. 경사 기반 베이스라인 (ORIGEN, $\Psi$-Sampler) 과 Best-of-N 샘플링과 비교하여 경쟁력 있거나 우수한 성능을 달성합니다. 특히 LCM 생성기에서 경사 기반 방법은 통제되지 않은 베이스라인 이하로 성능이 저하되었으나, ZeNO 는 개선을 유지하여 국소 경사 민감성에 대한 영차수 방법의 견고성을 강조합니다.
• 미분 불가능한 보상: 보상이 미분 불가능한 GenEval 작업 (객체 카운팅 및 위치 지정) 에서, ZeNO 는 카운팅 작업에서 Best-of-N 을 능가하여 동일한 계산 예산으로부터 반복적 노이즈 최적화가 독립적 샘플링보다 더 많은 가치를 추출함을 입증했습니다.
• 단백질 구조 생성: 단일 단계 리만 평균 흐름 (RMF) 단백질 백본 생성기에 적용된 ZeNO 는 설계 가능성 (자기 일관성 RMSD 로 측정) 을 크게 향상시켰습니다. 이는 단일 흐름 평가만으로 5 단계 RMF 베이스라인과 비교 가능한 설계 가능성 수준을 달성하여, 노이즈 최적화를 통해 감소된 샘플링 단계를 보상할 수 있음을 보여줍니다.
• 다양성: Vendi Score 를 사용한 분석은 ZeNO 가 샘플 다양성을 유지하면서 보상을 향상시키고, 경사 기반 접근법에서 자주 관찰되는 단일 모드로의 붕괴를 방지함을 나타냅니다.

의의 및 주장
이 논문은 ZeNO 가 효율적이고 결정론적이며 단일 단계인 생성 모델의 성장하는 클래스를 보상 함수와 정렬하기 위한 실용적이고 이론적으로 근거 있는 해결책을 제공한다고 주장합니다. 그 주요 의의는 블랙박스 보상 모델이나 복잡한 생물학적 설계 작업과 같이 경사를 사용할 수 없거나 신뢰할 수 없는 설정에서 작동할 수 있는 능력에 있습니다. 노이즈 최적화를 경로 적분 제어 문제로 공식화함으로써, ZeNO 는 재학습이나 기본 생성 모델 수정 없이 추론 시간 정렬을 위한 확장 가능하고 생성기 무관한 프레임워크를 제공합니다. 저자들은 한계를 지적합니다: 이 방법은 업데이트 단계당 여러 번의 보상 평가를 필요로 하므로, 보상 모델 자체가 비싸다면 계산적으로 요구 사항이 높을 수 있습니다.