기술 요약: 확산 정렬을 위한 스티치드 가치 모델 (StitchVM)

1. 문제 제기

확산 (Diffusion) 및 흐름 기반 (flow-based) 생성 모델은 데이터 분포를 포착하는 데 뛰어나지만, 실제로 유용하기 위해서는 작업별 보상 (예: 프롬프트 충실도, 미적 선호도, 인간 선호도) 과의 정렬이 종종 필요합니다. 이러한 정렬의 핵심 과제는 보상이 일반적으로 깨끗한 출력 이미지 ($x_0$) 에 대해 정의되는 반면, 정렬 절차 (추론 및 훈련 시간 모두) 는 노이즈가 있는 중간 잠재 변수 ($z_t$) 에 의해 유도된 깨끗한 샘플의 기대 보상을 추정해야 한다는 점입니다. 이러한 추정은 가치 함수 (value function) 로 알려져 있습니다.

이 가치 함수를 근사화하는 기존 방법들은 상당한 트레이드오프를 겪습니다:

• 트위디 (Tweedie) 근사: 깨끗한 샘플의 사후 평균을 추정하고 그 대리 모델에서 보상을 계산합니다. 효율적이지만 고노이즈 영역에서 상당한 편향을 도입하며, 디노이저와 VAE 디코더의 추가적인 고비용 평가를 필요로 합니다.
• 몬테카를로 (MC) 근사: $z_t$ 에서 여러 디노이징 궤적을 전개하고 보상을 평균냅니다. 더 정확하지만 여러 전체 궤적 전개가 필요하여 계산 비용이 prohibitive(부담스러울 정도로) 합니다.
• 학습된 가치 모델: 노이즈가 있는 잠재 변수에 대한 가치 함수를 직접 훈련하는 것은 해결되지 않은 문제입니다. 기존 시도들은 종종 소규모 훈련이나 좁은 선호도 데이터셋에 의존하여, 기초 규모 픽셀 공간 보상 모델에 비해 열악한 정확도와 일반화를 보입니다. 또한, 새로운 확산 백본이나 보상 모델마다 새로운 가치 모델을 훈련하는 것은 계산적으로 지속 가능하지 않습니다.

2. 방법론: StitchVM

저자들은 StitchVM (Stitched Value Model) 을 제안합니다. 이는 사전 훈련된 고품질 픽셀 공간 보상 모델의 능력을 최소한의 계산 비용으로 노이즈가 있는 잠재 영역으로 이전하는 프레임워크입니다. 이 접근법은 서로 다른 신경망의 레이어를 재결합하는 기술인 모델 스티칭 (model stitching) 을 활용합니다.

핵심 아키텍처

StitchVM 은 경량 스티칭 레이어를 통해 잘려진 프리즈된 확산 백본 ("헤드") 을 슬라이스된 사전 훈련된 보상 모델 ("테일") 에 연결하여 하이브리드 모델을 구성합니다.

• 헤드: 레이어 $i$ 에서 잘려진 확산 백본 (예: SD 3.5, FLUX) 입니다. 이 구성 요소는 본질적으로 노이즈가 있는 잠재 변수를 처리하고 노이즈 인지 특징을 추출합니다.
• 테일: 레이어 $j$ 부터 시작하는 사전 훈련된 보상 모델 (예: CLIP, HPSv2, Aesthetic Predictor) 입니다. 이 구성 요소는 대규모 사전 훈련 동안 학습된 견고한 작업별 보상 능력을 유지합니다.
• 스티칭 레이어 ($s_\psi$): 확산 헤드의 특징 공간과 보상 테일을 정렬하는 경량 모듈입니다.

훈련 절차

StitchVM 의 훈련은 두 단계로 이루어집니다:

1. 인터페이스 선택 (폐형식): 저자들은 확산 특징 (노이즈가 있는 잠재 변수에서) 과 보상 특징 (깨끗한 이미지에서) 이 가장 호환되는 최적의 스티칭 인덱스 $(i, j)$ 를 탐색합니다. 이는 깨끗한 이미지의 작은 프로브 세트와 해당 노이즈가 있는 잠재 변수에서 추출된 짝지어진 특징들 사이의 선형 매핑 $W$ 를 피팅함으로써 달성됩니다. 특징 매칭 손실을 최소화하는 쌍이 선택됩니다.
2. 경량 파인튜닝: 인터페이스가 고정되면, 스티칭 레이어와 잘려진 보상 테일은 레이블이 없는 깨끗한 이미지를 사용하여 파인튜닝됩니다. 목표는 깨끗한 이미지의 원래 보상 점수 $r(z_0)$ 로 가치 함수 $V(z_t)$ 를 회귀시키는 것입니다. 이 오프-폴리시 목적은 전체 궤적 전개 없이 노이즈가 있는 잠재 변수에 조건부인 깨끗한 샘플의 기대 보상을 학습할 수 있게 합니다.

정렬에의 적용

StitchVM 은 추론 시간 및 훈련 시간 정렬에서의 근사화를 대체하는 직접적인 가치 함수 추정기로 작용합니다:

• 추론 시간:
  • DPS (Diffusion Posterior Sampling): 디노이저와 VAE 를 통한 역전파의 필요성을 제거하고, Tweedie 기반 기울기를 StitchVM 의 직접 기울기 ($\nabla_{z_t} V(z_t)$) 로 대체합니다.
  • FK Steering: 각 입자가 여러 지역 제안 (local proposals) 을 생성하는 "제안 스케일링"을 가능하게 합니다. 이러한 제안들은 Tweedie 근사로는 너무 비용이 많이 드는 과정을 거쳐 StitchVM 에 의해 저렴하게 점수 매겨져 최상의 후보가 선택됩니다.
• 훈련 시간:
  • 직접 보상 파인튜닝 (예: AlignProp, DRaFT): 훈련을 중간 노이즈 잠재 변수 $z_\tau$ 에서 중단하고 StitchVM 예측을 감독 목표로 사용하여 전체 궤적 전개와 긴 역전파 체인을 피할 수 있게 합니다.
  • DiffusionNFT: 최종 깨끗한 보상을 중간 가치 추정치로 대체하여 forward 과정에서의 효율적인 온라인 RL 을 가능하게 합니다.

3. 주요 기여

• 가치 모델링을 위한 새로운 패러다임: StitchVM 은 강력한 픽셀 공간 보상 모델이 대규모 재훈련 없이 단순히 확산 백본에 스티칭됨으로써 효과적으로 노이즈가 있는 잠재 영역으로 이전될 수 있음을 보여줍니다.
• 효율성과 정확성: 이 방법은 가치 추정에서 최첨단 성능을 달성하면서도 매우 경량입니다. 예를 들어, CLIP ViT-L 을 SD 3.5 Medium 가치 모델에 스티칭하는 데는 약 10 GPU 시간이 소요됩니다.
• 정렬에서의 성능 향상:
  • 추론: DPS 는 3.2$\times$ 빨라지고 피크 GPU 메모리는 50% 감소합니다. StitchVM 을 사용한 FK Steering 은 특히 효율적인 지역 제안 선택을 가능하게 함으로써 표준 FK Steering 및 Best-of-$N$ 샘플링을 크게 능가합니다.
  • 훈련: DiffusionNFT 는 2.3$\times$ 빨라지고, 직접 보상 파인튜닝 (DRaFT) 은 1.3$\times$ 빨라지며 고노이즈 단계에서의 효과적인 감독으로 인해 더 높은 생성 품질 (예: GenEval +30%) 을 달성합니다.
• 일반화: 이 접근법은 다양한 확산 백본 (SD 3.5 Medium/Large, FLUX) 과 보상 모델 (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor) 에서 작동합니다.

4. 실험 결과

이 논문은 여러 벤치마크에 걸쳐 StitchVM 을 평가합니다:

• 가치 함수 품질: 다양한 노이즈 수준 ($\sigma \in [0.1, 0.9]$) 의 노이즈가 있는 잠재 변수에서 StitchVM 은 제로샷 검색 (MSCOCO, Flickr30K), 선호도 예측 (HPDv2, ImageReward), 미적 점수 매기기 (AVA) 에서 원래 깨끗한 보상 모델의 성능과 밀접하게 일치합니다. 고노이즈에서 붕괴되는 VAE-스티칭이나 노이즈가 있는 데이터에서 처음부터 훈련된 모델 (예: NoisyCLIP, DiNa-LRM) 과 같은 베이스라인을 크게 능가합니다.
• 추론 정렬: DrawBench 에 대한 DPS 및 FK Steering 실험에서 StitchVM 강화 방법은 ImageReward, HPSv2, GenEval 과 같은 지표를 지속적으로 개선하면서 계산 비용을 극적으로 줄입니다. 주목할 점은 StitchVM 을 사용한 FK Steering 에서 지역 제안 수 ($M$) 를 늘리는 것이 입자 수 ($N$) 를 늘리는 것보다 더 효율적이라는 것입니다.
• 훈련 정렬: 훈련 곡선은 StitchVM 변형들이 표준 DRaFT 및 DiffusionNFT 에 비해 적은 GPU 시간으로 더 높은 최종 점수에 도달함을 보여줍니다. 정성적 결과는 StitchVM 강화 모델이 더 빠르게 수렴하고 훈련 초기에 더 선명하고 고품질의 이미지를 생성함을 나타냅니다.

5. 중요성과 주장

이 논문은 StitchVM 이 어느 쪽의 광범위한 사전 훈련도 희생하지 않고 잠재 확산 백본과 픽셀 공간 피드포워드 모델을 결합하기 위한 범용 템플릿을 제공한다고 주장합니다. 강력한 보상 모델을 잠재 공간으로 끌어올림으로써, 노이즈가 있는 잠재 변수에 대한 올바른 가치 함수를 한 번 구성하고 많은 샘플과 반복에 걸쳐 분산 (amortized) 하는 새로운 유형의 확산 정렬을 가능하게 합니다.

저자들은 StitchVM 을 기초 모델 사전 훈련의 대체물이 아니라, 이러한 모델을 특정 작업에 적응시키기 위한 실용적이고 저비용의 메커니즘으로 위치시킵니다. 그들은 이 접근법이 가치 모델 기반 정렬을 접근 가능하고 효율적으로 만들어 안전하고 신뢰할 수 있는 생성 모델링 연구에 가속을 줄 수 있음을 강조합니다. 이 연구는 사전 훈련된 구성 요소의 "스티칭"이 가치 함수를 처음부터 훈련하거나 편향된 근사화에 의존하는 prohibitive 비용에 대한 실행 가능한 대안임을 시사합니다.