기술 요약: 확산 모델의 테스트 타임 스케일링을 위한 Lookahead Sample Reward Guidance (LiDAR)

문제 정의

확산 모델(Diffusion models)은 이미지, 비디오, 언어 분야에서 놀라운 생성 성능을 달 achievement 했으나, 생성된 샘플이 인간의 의도와 완전히 일치하지 못하는 경우가 빈번합니다. 인간의 선호도를 포착하는 보상 함수(Reward functions)가 존재함에도 불구하고, 이를 샘플링 과정에 통합하는 것은 상당한 과제를 안겨줍니다. 테스트 타임 스케일링(Test-time scaling)의 핵심 목표는 보상 기울어진 분포(Reward-tilted distribution) $p_\theta^r(x_0|c) \propto p_\theta(x_0|c) \exp(\lambda \cdot r(x_0, c))$로부터 샘플링하는 것입니다.

이를 달성하기 위해서는 중간 입자(Intermediate particles) $x_t$에 대한 기대 미래 보상(Expected Future Reward, EFR), $r_t^\lambda(x_t, c)$를 계산해야 하며, 이는 역방향 확산 과정(Reverse diffusion process)을 가이드하는 데 사용됩니다. 기존 방식들은 다음과 같은 결정적인 한계에 직면해 있습니다:

1. 역방향 롤아웃(Backward Rollout): EFR을 추정하기 위해 $x_t$에서 $x_0$까지 직접 샘플링하는 것은 매 타임스텝마다 반복적인 디노이징(Denoising)이 필요하므로 계산 비용이 매우 높습니다.
2. Tweedie 기반 근사(Tweedie-based Approximations): Gradient Guidance나 순차적 몬테카를로(Sequential Monte Carlo, SMC)와 같은 방법들은 EFR을 추정하기 위해 1차 테일러 근사(Twedie's formula)에 의존합니다. 이러한 방식은 특히 보상 신호의 강도($\lambda$)가 높을 때 편향(Bias)을 유발하며, 신경망을 통한 역전파(Backpropagation)가 필요하거나(Gradient Guidance), 고차원 공간에서 샘플 붕괴(Sample collapse) 현상을 겪기도 합니다(SMC).

방법론: LiDAR 샘플링

저자들은 중간 입자 $x_t$와 EFR 사이의 신경망 의존성을 분리하여, 폐쇄형(Closed-form) 및 미분 불필요(Derivative-free) 가이딩을 가능하게 하는 LiDAR(Lookahead Sample Reward Guidance) 프레임워크를 제안합니다.

1. 기대 미래 보상의 재정의

본 논문은 EFR의 새로운 정식화(Theorem 3.1)를 도출합니다. 이는 현재의 입자 $x_t$를 신경망을 통해 조건화하는 대신, 사전 학습된 확산 모델의 **주변 샘플(Marginal samples)**과 순방향 섭동 커널(Forward perturbation kernel)만을 사용하여 보상 기댓값을 표현합니다:
$$r_t^\lambda(x_t, c) = \log \mathbb{E}_{p_\theta(x_0|c)} \left[ \frac{p(x_t|x_0)}{\mathbb{E}_{p_\theta(x_0|c)}[p(x_t|x_0)]} \exp(\lambda \cdot r(x_0, c)) \right]$$
이 "순방향 롤아웃(Forward rollout)" 접근 방식은 미리 생성된 주변 샘플 집합 $\{\hat{x}_0^i\}$를 사용하여 어떤 $x_t$에 대해서도 EFR을 계산할 수 있게 하며, 가이던스 단계에서 스코어 네트워크 $s_\theta$나 보상 모델을 통한 역전파를 제거합니다.

2. 효율적인 Lookahead 샘플링

모든 프롬프트에 대해 주변 샘플을 생성하는 비용을 완화하기 위해, 저자들은 효율적인 룩어헤드 샘플링(Lookahead sampling) 전략을 도입합니다. 느린 전체 확산 과정을 사용하는 대신, 몇 단계의 ODE 솔버(예: DPM-Solver)나 증류된 모델(Distilled model)과 같은 빠른 생성기 $q(x_0|c)$를 사용하여 $n$개의 룩어헤드 샘플 $\{\hat{x}_0^i\}_{i=1}^n$을 생성합니다. 이 샘플들은 보상 값 $r(\hat{x}_0^i, c)$가 함께 주석 처리됩니다.

3. 폐쇄형 미분 불필요 가이던스

유한한 룩어헤드 샘플 집합을 사용하여, 저자들은 경험적 룩어헤드 보상의 그래디언트에 대한 폐쇄형 표현(Theorem 3.3)을 도출합니다. 가이던스 항은 현재의 입자 $x_t$를 높은 보상을 가진 룩어헤드 샘플 쪽으로 유도하고, 낮은 보상을 가진 샘플으로부터는 밀어냅니다:
$$\nabla_{x_t} \hat{r}_t^\lambda(x_t, c) = \sum_{i=1}^n (w_i^r - w_i) \frac{\hat{x}_0^i}{\sigma_t^2}$$
여기서 $w_i^r$와 $w_i$는 보상 값과 $x_t$와 룩어헤드 샘플 $\hat{x}_0^i$ 사이의 거리에 기반한 소프트맥스 가중치입니다. 이 공식은 신경망 미분을 요구하지 않습니다.

4. 이론적 스케일링 법칙

본 논문은 이 방법의 스케일링 법칙을 확립합니다:

• 룩어헤드 정확도($\delta$): 타겟 분포와 LiDAR로 생성된 분포 사이의 총 변동 거리(Total variation distance)는 룩어헤드 솔버의 이산화 단계 수인 $\delta$에 따라 $O(1/\sqrt{\delta})$로 감소합니다.
• 샘플 크기($n$): 경험적 보상 추정치의 분산은 $O(1/\sqrt{n})$으로 감소하며, 이는 룩어헤드 샘플의 수를 늘릴수록 근사 정확도가 향려됨을 의미합니다.

주요 결과

SD v1.5, SDXL, FLUX (Flow matching), UDLM (Discrete diffusion)에 대해 ImageReward, CLIP, HPS, GenEval 지표를 사용하여 실험을 진행했습니다.

• 성능 대비 효율성: SDXL에서 LiDAR는 최신 그래디언트 가이던스 방법(DATE)과 대등한 GenEval 성능을 달성하면서도, 추론 시간에서 9.5배의 속도 향상을 보여주었습니다.
• 메모리 효율성: 역전파로 인해 약 3배 더 많은 메모리를 요구하는 그래디언트 가이던스 방법과 달리, LiDAR의 메모리 사용량은 바닐라 샘플링과 거의 동일합니다. 이는 타겟 샘플링 단계(Algorithm 2)에서 그래디언트 계산이 필요하지 않기 때문입니다.
• 보상 해킹(Reward Hacking)에 대한 강건성: 그래디언트 가이던스 방법들이 ImageReward를 최적화할 때 CLIP이나 HPS 같은 지표에서 성능이 저하되는 것과 달리, LiDAR는 모든 지표를 동시에 개선합니다. 이는 룩어헤드 샘플 내에서 ImageReward와 다른 지표들 간의 상관관계가 샘플러를 균형 잡힌 방향으로 자연스럽게 유도하기 때문입니다.
• 직교적 개선(Orthogonal Improvements): LiDAR는 DPO(Direct Preference Optimization)로 이미 미세 조정된 모델에도 적용되어 추가적인 성능 향상을 제공하며, 이는 학습 기반 정렬(Alignment)과의 직교성을 입증합니다.
• 확장성: 이 방법은 룩어헤드 샘플 수($n$)와 솔버 단계($\delta$)에 따라 효과적으로 확장되며, SMC 및 Best-of-N과 같은 다른 기술들과 결합될 수 있습니다.

의의 및 주장

본 논문은 LiDAR가 기존의 테스트 타임 스케일링 방법들이 가진 효율성 및 보상 해킹 문제를 체계적으로 해결한다고 주장합니다. EFR 계산을 순방향 롤아웃과 주변 샘플에 의존하도록 재정의함으로써, 저자들은 계산적으로 효율적이고 이론적 근거가 확실한 **미분 불필요 가이던스(Derivative-free guidance)**를 가능하게 했습니다.

저자들은 LiDAR가 기존 베이스라인들에 비해 우수한 효율성-성능 트레이드오프를 달성한다는 점을 강조합니다. 구체적으로, 신경망 역전파의 계산 오버헤드와 SMC의 샘플 붕괴 문제를 피하면서, 테일러 근사에 기반한 방법들보다 더 정확한 EFR 추정을 제공합니다. 이 연구는 보상 신호를 중간 입자의 신경망 의존성으로부터 분리하는 것이 확장 가능한 인간 정렬 생성 모델링을 향한 유효한 경로임을 시사합니다.