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

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

왜 중요한가

핵심 기여

StitchVM의 핵심 아이디어

Pretrained pixel-space reward 모델의 tail과 frozen diffusion backbone의 head를 stitching layer로 연결해, 노이즈 latent에서 직접 보상-가치 추정을 가능하게 한다. 이는 적은 파인튜닝으로 두 영역의 지식을 결합하는 모델-스티칭의 원리 위에 세워졌다.

인터페이스 선택의 효율적 수학적 해결

diffusion head의 층(i)와 reward tail의 층(j) 사이의 매핑을 선형 변환 W로 근사하고, 많은 후보 쌍에 대해 closed-form 최소화(W★i,j)로 적합한 접합 지점을 찾는다. 이 탐색은 빠르고, 스티칭 후 Stage-2 재학습으로 잔차를 보정한다.

저비용의 스티칭-학습

Stage 2에서 sψ(스티칭 layer)와 r≥jφ를 소량의 unlabeled 이미지(z0)로 미세조정하며, Lvalue = E[z0,t,ε]{V(i★,j★)ω(zt) − φ(z0)}^2 와 같은 비지도/오프-폴리시 손실로 가치 함수를 추정한다.

추론 시간 정렬의 대폭적 효율성 향상

DPS의 gradient를 패치-단위로 직접 계산하는 대신 StitchVM의 노이즈 latent 공간에서의 gradient를 활용해, denoiser/VAE 전 과정을 역전하지 않고도 guided sampling이 가능하다. 이로써 메모리 사용량이 대폭 감소하고 샘플링 속도가 최대 3.2× 빨라진다.

훈련 시간 정렬의 가속 및 품질 유지

AlignProp/DRaFT, DiffusionNFT 등에서 중간 노이즈 단계에서의 StitchVM 값을 supervision으로 사용해 전체 denoising 경로를 종료시키고, GPU-시간을 22–26% 감소시키거나 55% 이상 감소를 달성하며, DRaFT의 생성 품질도 향상될 수 있다.

교차 백본 일반화

작은 백본에서 학습된 StitchVM이 더 큰 생성기에도 거의 손실 없이 가이드를 제공하며, 𝑀-스케일링에서의 비용-효율 증가를 가능하게 한다.

핵심 아이디어 이해하기

현 기법은 프리트레인된 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 방법의 적용 가능성이 제시되나 추가 연구가 필요하다.

실무 활용

StitchVM은 기존의 pixel-space reward 모델을 노이즈 latent에서도 직접 사용할 수 있게 함으로써, 샘플링과 학습시 보상 기반 정렬의 효율을 크게 향상한다.

inference-time alignment에서 DPS/FK의 처리 속도와 메모리 효율 개선
training-time alignment에서 DiffusionNFT/DRaFT의 학습 시간을 단축하고 품질 유지
다양한 diffusion backbones에 대한 reward 모델의 일반화 및 재사용
교차 백본에서의 빠른 가이드링으로 대형 생성기 시스템의 비용 절감

코드 공개 여부: 비공개

키워드

diffusion modelsreward modelsnoisy latentsmodel stitchingStitchVMCLIPAesthetic score

용어 해설

스티칭 레이어(Stitching Layer): — 두 개의 사전 학습된 모델 간 표현 공간의 차이를 연결하기 위해 diffusion 백본의 중간 층 출력과 reward 모델의 출력 공간을 매핑하는 학습 가능한 모듈.
diffusion 백본(Diffusion Backbone): — 노이즈가 포함된 latent 공간을 처리하도록 설계된 diffusion 모델의 앞부분(헤드)을 의미하며, 보상 모델의 출력 공간과의 연결 지점을 제공한다.
픽셀-스페이스 보상 모델(Pixel-Space Reward Model): — 깨끗한 이미지에 대해 학습된 보상 모델로, StitchVM은 이를 노이즈 Latent 공간에 전달 가능하도록 이전-지식으로 활용한다.
노이즈가 있는 latent(Noisy Latents): — 확산 과정 중 중간 시점의 Latent 공간으로, 보상 신호가 정의된 청정 샘플(x0) 대신 이 공간에서 가치 평가를 수행한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환