분포 기반 보상으로 시각 생성 모델 최적화

샘플 단위 보상으로 강화학습을 적용하면 보상 해킹과 모드 붕괴가 발생하여 생성물의 다양성과 전체적 분포 적합도가 악화될 수 있다. 본 논문은 분포 단위의 보상을 효율적으로 추정하고 이를 RL 신호로 활용함으로써 시각 생성 모델의 지각 품질과 분포 커버리지를 동시에 개선했다. 결과적으로 사후 병합과 결합한 접근은 SDE-기반 훈련과 ODE-기반 추론 간 불일치를 완화하는 실제적 대안을 제시한다.

기존 강화학습 기반의 생성 모델 미세조정은 보상 함수를 샘플 단위로 평가하는 관행에 의존해왔다. 이러한 샘플별 보상은 각 샘플이 독립적으로 높은 보상을 얻도록 수렴시키는 경향이 있어 보상 해킹을 유발하고 결과적으로 생성 분포의 다양성 및 모드 커버리지를 손상시켰다. FID나 MMD 같은 분포 기반 지표는 집합 수준의 분포 적합도를 반영하므로 이러한 문제를 근본적으로 완화할 잠재력을 지닌다.

분포 기반 보상을 직접 RL 신호로 사용하려면 전체 분포에 대한 통계적 측정을 반복해야 하므로 계산 비용이 매우 크다. 이 논문은 이 문제를 해결하기 위해 레퍼런스 집합을 만들고 훈련 중 소규모 서브셋만 새 샘플로 교체한 뒤 교체된 집합의 FID를 계산해 해당 서브셋에 보상을 부여하는 방법을 제시한다. 이렇게 하면 전체 집합에 대한 근사적 분포 정보가 각 서브셋에 전이되고, 샘플별 보상보다 더 풍부하고 안정된 신호를 얻을 수 있다.

또한 SDE 기반의 확률적 롤아웃에서 얻은 이득이 ODE 기반의 결정적 추론으로 잘 이전되지 않는 훈련-추론 불일치 문제가 확인되었다. 이를 해결하기 위해 체크포인트 풀을 병합하는 계수를 RL로 탐색하는 사후 병합(post-hoc model merging) 접근을 도입했다. 이 방식은 RL 최적화가 직접적으로 ODE 기반 샘플러 아래에서 작동하도록 하여 SDE-ODE 불일치로 인한 성능 저하를 회피한다.

전체 접근은 두 축으로 구성된다. 첫째는 subset-replace 전략을 통해 분포 기반 보상을 효율적으로 계산하는 방법이며, 둘째는 사후 모델 병합 계수를 RL로 최적화하는 방법이다. subset-replace는 먼저 초기 pretrained 모델로 클래스 균형을 맞춘 레퍼런스 세트 G(크기 N)를 생성하고, 각 라운드에서 이 세트의 작은 서브셋 g(n개)를 동일한 클래스의 새 샘플 g'로 교체한 뒤 교체된 집합의 FID를 계산한다. 보상은 교체된 집합의 FID에 음수를 취한 값으로 정의되고, 이 값들은 배치 수준에서 정규화되어 advantage로 변환된 후 정책 그래디언트 업데이트에 사용된다.

정규화와 샘플 선택이 중요한 역할을 한다. 논문에서는 batch-level advantage normalization을 채택하고, 전역 랭킹에 따른 상위 25%의 샘플만 파라미터 업데이트에 사용하여 최적화를 안정화했다. 정책 최적화 목표는 KL 규제 항을 포함한 정책 그래디언트 목표로 표기되며, 중요도 비례 클리핑(clip)과 KL 벌점을 함께 사용한다. 레퍼런스 세트는 최신 모델로 주기적으로(예: 10 스텝마다) 갱신하여 참조 분포와 모델 분포 간 괴리를 완화한다.

사후 모델 병합에서는 훈련 궤적에서 선택한 N_c개의 체크포인트 {M_i}에 대해 가중치 벡터 w를 적용해 M_merge=Σ w_i M_i를 구성한다. 작은 MLP(EMANet)가 입력 임베딩 z로부터 각 계수의 평균과 표준편차를 출력하고, 이로부터 가우시안 샘플링으로 w를 생성한다. 각 w에 대해 병합 모델을 생성하고 subset-replace를 통해 보상 컬렉션을 얻은 뒤 배치 수준 이점 정규화를 통해 EMANet 파라미터를 업데이트한다. 이 과정에서는 ODE 샘플러를 사용해 효율적이고 결정적인 생성 절차를 유지한다.

메인 벤치마크에서 SiT-XL 모델에 대해 subset-replace 기반 RL을 적용하면 FID-50K가 초기 8.30에서 RS 적용 시 6.98로, RL 적용 시 5.77로 개선되었다. FD DINOv2 지표도 230.39에서 RL 후 164.88로 낮아져 다양한 특징 표현 공간에서도 분포 적합도가 향상되었다. EDM2 계열에서는 사후 모델 병합에 RL을 적용해 EDM2-XS의 FID를 3.74에서 3.52로, EDM2-S를 2.57에서 2.52로 추가 개선했다.

광범위한 ablation에서 레퍼런스 세트 크기는 성능과 안정성 사이의 트레이드오프를 보였으며, 실험에서는 5,000 샘플 구성이 대표성과 안정성의 균형을 가장 잘 달성했다. 교체 서브셋 크기는 50이 계산 비용과 신호 분산 측면에서 최적으로 나타났고, 배치 수준 정규화가 그룹 수준 정규화보다 빠른 수렴을 보였다. 또한 롤아웃 샘플 선택은 전역 상위 25%를 사용하는 방식이 가장 우수했고, RS로 먼저 파인튜닝한 뒤 RL을 추가하는 방식은 과적합으로 이어져 순수 RL 단독 적용이 더 나은 결과를 보였다.

SDE 기반 훈련과 ODE 기반 평가 간의 불일치 실험에서 SDE로 훈련된 모델은 SDE 평가에서는 지속적 개선을 보였지만 동일한 설정의 ODE 평가에서는 개선이 미진했다. 이 문제를 회피하기 위해 제안한 사후 병합 최적화는 ODE 기반 생성 절차에서 직접 계수를 탐색하므로 훈련-추론 불일치 완화에 기여했다.

전체 아키텍처는 정책 네트워크 형태의 diffusion policy와 EMANet라 명명한 3-layer MLP로 구성된다. diffusion policy는 denoising 과정을 MDP로 다루며 각 시간 단계의 행동은 다음 단계의 복원 샘플을 생성하는 확률 분포로 모델링된다. EMANet는 입력 임베딩 z를 받아 각 병합 계수의 평균과 표준편차를 출력하고 가우시안 샘플링으로 계수를 생성한다.

subset-replace의 수학적 기반은 레퍼런스 세트 G에서 일부 서브셋 g를 g'로 교체한 집합의 FID를 계산하고 보상 R(g')=-FID[(G\g)∪g', Ḡ]로 정의하는 것이다. 이 보상은 배치 수준에서 평균과 표준편차로 정규화되어 advantage ĤA_i=(R_i-mean)/std로 변환되고, 각 denoising trajectory의 확률비 r_t^{i,j}(θ)=p_θ(x_{t-1}|x_t,c)/p_{θ_old}(...)와 중요도 클리핑을 포함한 정책 그래디언트 항으로 결합된다. 추가로 KL(π_θ||π_ref) 항을 β로 가중해 안정성을 확보한다.

구현 세부에서는 denoising reduction을 사용해 훈련 시 50 스텝, 평가 시 250 스텝 샘플링을 적용했고 레퍼런스 세트는 5,000 샘플로 설정해 10 스텝마다 최신 모델로 갱신했다. 훈련은 16개의 NVIDIA Hopper GPU에서 병렬로 수행되었고 최고 성능 실험은 약 20시간 소요되었다. 배치 수준 정규화와 전역 상위 25% 샘플 필터링은 안정적 수렴을 위해 핵심 하이퍼파라미터로 채택되었다.

레퍼런스 세트 크기 선택은 비단순한 trade-off를 보이며 논문에서는 5,000 샘플 구성이 대표성과 안정성의 균형을 제공한다고 보고되었지만 일부 중간 크기(예: 7,500)에서는 성능 저하와 불안정성이 관찰되었다. subset-replace는 전체 FID 계산보다 비용을 절감하지만 여전히 주기적인 레퍼런스 세트 갱신과 여러 샘플 대체로 인한 계산 오버헤드가 존재한다. RS 후에 추가로 RL을 적용하면 과적합이 발생할 수 있어 RS→RL 병행 전략은 항상 유리하지 않다는 제한이 논문에 명시되어 있다.