확산 정합용 스티치드 가치 모델(StitchVM)

diffusion/flow 기반 생성모델은 프롬프트 적합성이나 미적 선호도 같은 보상에 맞춰 조정되기 어렵다. 보상은 일반적으로 깨끗한 이미지에 대해 정의되고, 노이즈가 있는 latent에서의 가치 평가가 필요하다. Tweedie 또는 Monte Carlo 접근은 편향 혹은 비용 이슈를 동반한다. StitchVM은 pixel-space reward models를 노이즈 latent으로 이식해, 고정된 head와 작은 finetuning으로 보상 모델의 강점을 그대로 유지하면서 latent 공간에서의 정렬을 가능하게 한다.

현 기법은 프리트레인된 pixel-space reward 모델의 표현 공간과 diffusion backbone의 노이즈-처리 표현 공간 사이의 차이를 스티칭 레이어로 매핑한다. 1) StitchVM은 후보 인터페이스(i, j)를 선형 매핑으로 평가하고, 가장 낮은 매핑 오차를 주는 인터페이스를 선택한다. 2) Stage 2에서 스티칭 레이어와 reward suffix를 소량의 unlabeled 이미지로 미세조정해 잔차를 보정한다. 3) Inference-time에서 StitchVM은 노이즈 latent에서 직접 가치 함수를 평가하므로 Tweedie/MC의 편향이나 다중 denoiser/decoder 경로를 피하고, PKG에서의 제안들로 더 빠르고 안정적인 샘플링을 제공한다. 4) Training-time에서도 중간 노이즈 단계에서 가치 함수를 supervises로 이용해 전체 역과정을 중단시키고, DiffusionNFT/DRaFT의 학습 효율을 높인다. 5) 이 접근은 pixel-space reward의 강력한 일반화 능력을 latent-space로 확장하되, 재학습 비용은 최소화한다.

주요 결과는 다음과 같다. Inference-time에서 DPS의 gradient를 Tweedie 대신 StitchVM으로 대체해, SD 3.5 Medium/Large 및 FLUX 백본에서 대부분의 보상-지표에서 성능을 유지하거나 개선하고, 메모리 사용량은 약 50% 감소하며 샘플링 시간은 최대 3.2× 단축된다. 예를 들어 HPSv2 보상 하에서 StitchVM은 ImageReward를 0.93(기준)에서 1.10으로 향상시키고 GenEval은 0.62에서 0.69로 올린다. FK Steering에서도 StitchVM은 GenEval과 ImageReward에서 상향 효과를 보이며, 특히 SD3.5 Medium에서 큰 개선을 보인다. 훈련 시 정렬에서도 StitchVM 도입으로 DRaFT와 DiffusionNFT의 GPU-시간이 각각 22–26% 및 55% 이상 감소하며, DRaFT의 생성 품질은 향상된다. Cross-backbone 실험(Table 6)에서 SD3.5 Medium으로 학습한 StitchVM이 SD3.5 Large를 가이드하되 성능 저하가 미미하며, backbones 간의 transfer 비용을 크게 줄일 수 있음을 보인다.

StitchVM은 (i, j) 인터페이스를 찾기 위해, forward 경로에서 paired features u≤iθ(zt) 와 r≤j−1φ(z0)의 매칭을 수행하고, 선형 매개변수 W를 최소 제곱으로 구한다. 이 값이 최소가 되는 (i*, j*)를 선택한다. Stage 2에서 stitching layer sψ(h) = UpFψ(h) + Gψ(UpFψ(h))를 사용하여 diffusionhead의 특징과 reward model의 출력 간의 비선형 차이를 보정한다. Fψ은 1×1 컨볼루션으로 채널 차이를 맞추고, Up은 해상도 보정, Gψ은 2-layer MLP로 비선형 보정을 수행한다. 보상 모델 suffix r≥jφ는 고정된 pretrained 모델을 유지하며, Lvalue = E[z0,t,ε]{V(i*, j*)ω(zt) − φ(z0)}^2를 최소화하도록 finetune한다. Inference-time에서 gradient guidance는 ∇ztVt(z t) 대신 StitchVM의 직접 계산된 gradient를 사용하며, 이를 통해 denoiser/decoder를 통과하지 않고도 샘플링을 수행한다. Training-time에서는 LV-DRFT(θ)와 LV-NFT(θ)로 두 가지(StitchVM 버전)를 도입하여 짧은 중간 노이즈 지점에서 감독 신호를 제공하고, 전체 denoising 경로의 롤아웃 필요성을 제거한다.

StitchVM은 feedforward reward 모델로 구현된 보상에 한정되며, 비피드포워드 보상에는 직접 적용하기 어렵다. 학습/추론 시 timesteps-aware 방법의 적용 가능성이 제시되나 추가 연구가 필요하다.