RewardFlow: 보상 최적화를 통한 이미지 생성 및 편집 프레임워크

기존의 확산 모델 기반 이미지 편집은 원본의 정체성을 잃거나 특정 영역만 정밀하게 수정하는 데 한계가 있었다. RewardFlow는 별도의 학습이나 복잡한 역전 과정 없이 추론 단계에서 여러 보상 함수를 결합해 사용자의 의도를 정확히 반영하면서도 이미지의 구조를 완벽하게 보존한다.

기존 확산 모델의 이미지 편집은 텍스트 프롬프트와 이미지 사이의 전역적인 정렬에만 의존하여 특정 부분만 고치려 해도 전체적인 구도가 바뀌는 'Semantic Leakage' 문제가 빈번했다. RewardFlow는 이를 해결하기 위해 이미지 생성 과정을 '여러 목표를 동시에 만족시켜야 하는 최적화 문제'로 재정의한다.

핵심은 Langevin Dynamics를 활용해 잠재 공간의 궤적을 사용자가 원하는 보상(Reward)이 높은 방향으로 미세하게 조정하는 것이다. 이때 Gradient Descent 개념을 사용하여 각 보상 함수가 생성 중인 이미지에 대해 주는 점수를 높이는 방향으로 잠재 변수를 갱신한다. 예를 들어 '빨간 사과'라는 보상이 있다면, 현재 이미지가 사과 모양을 갖추고 빨간색을 띨수록 더 강한 가이드를 제공하게 된다.

특히 이 과정에서 LLM이 프롬프트를 분석해 '무엇을 추가하고 무엇을 유지할지'를 미리 판단하고, 샘플링 단계마다 각 보상의 중요도를 바꾼다. 초기에는 전체적인 구도를 잡는 보상에 집중하고, 후기에는 세부 묘사를 다듬는 보상에 집중하는 방식이다. 결과적으로 원본 이미지의 틀은 유지하면서도 사용자가 지시한 부분만 정밀하게 수정된 이미지를 얻을 수 있다.

RewardFlow는 사전 학습된 Flow-matching 모델의 추론 궤적을 멀티 보상 신호로 가이드한다. 전체 프로세스는 프롬프트 분석, 보상 계산, 적응적 가중치 조절, 그리고 Langevin 업데이트의 단계로 구성된다.

먼저 GPT-5와 같은 LLM을 사용하여 입력 프롬프트 p에서 독립적인 의미 단위인 Semantic Primitives(SP)를 추출한다. 각 SP에 대해 Global Alignment(SigLIP), Perceptual Alignment, Region-level Grounding(RegionCLIP), Object Consistency(SAM2) 보상을 계산한다. 추가적으로 전체 프롬프트에 대해 Human Preference(HPSv2)와 VQA 보상을 적용하여 세밀한 논리적 일치성을 확보한다.

Langevin 업데이트 식은 z(k+1) = z(k) + ηk(fk + gRtot,k + gKL,k) + ξk 형태로 정의된다. 여기서 fk는 모델의 기본 드리프트이며, gRtot,k는 여러 보상의 기울기를 통합한 값이다. ηk는 적응적 스텝 사이즈로, 현재 총 보상 값 Rtot에 로지스틱 함수를 적용하여 계산한다. [현재 보상 값 입력 → 로지스틱 함수 연산 → 스텝 사이즈 출력] 순으로 계산되며, 목표에 가까워질수록(보상이 높을수록) 스텝 사이즈를 줄여 정밀한 최적화를 수행한다.

정체성 보존을 위해 도입된 KL Tether 항 gKL,k는 현재 예측된 클린 잠재 변수와 원본 잠재 변수 z0 사이의 차이를 계산한다. [현재 예측값과 원본값 입력 → 차이 벡터 계산 및 자코비안 매핑 → 기울기 출력] 과정을 거쳐, 최적화가 원본의 구조에서 너무 멀어지지 않도록 잡아주는 닻 역할을 수행한다.

PIE-BENCH 이미지 편집 벤치마크에서 Flux 백본을 사용했을 때 기존 SOTA 모델 대비 Distance 지표를 7.3% 개선(7.78 vs 8.39)했으며, PSNR과 SSIM에서도 각각 5.3%, 2.6% 향상된 수치를 기록했다. 이는 배경과 정체성을 더 잘 보존하면서도 선명한 편집이 가능함을 의미한다. 편집 정확도 측면에서도 Whole accuracy 29.44%, Edited accuracy 26.62%를 달성하여 모든 Flux 기반 편집 모델을 능가했다.

T2I-COMPBENCH를 이용한 조합적 이미지 생성 실험에서도 PixArt-α, Flux, Qwen Image 등 모든 백본 모델에서 성능 향상을 보였다. 특히 Qwen Image 모델에 적용했을 때 비공간적(Non-spatial) 및 복잡한 조합(Complex) 카테고리에서 약 12.8%의 성능 향상을 기록하며, 여러 객체와 속성이 얽힌 복잡한 프롬프트에 대한 대응 능력을 입증했다.

Ablation Study 결과, VQA 보상이 가장 강력한 세밀한 감독 신호를 제공하며, KL Tether가 구조적 왜곡을 방지하는 데 결정적인 역할을 함이 확인됐다. 또한 동적 가중치 조절을 제거했을 때 텍스트 일관성이 저하되는 현상이 관찰되어, 제안된 적응적 정책의 유효성이 증명됐다.

RewardFlow의 핵심 아키텍처는 Langevin SDE의 이산화(Discretization)를 통해 유도된다. 이는 프롬프트에 의해 편향된 밀도(Prompt-tilted density) ρt(z|p, x)를 타겟으로 하는 샘플링 과정으로 해석될 수 있다. 수학적으로는 로그 확률 밀도의 기울기를 모델의 스코어 함수와 보상 함수의 기울기 합으로 분해하여 업데이트에 반영한다.

특히 VQA 보상 Rvqa는 Qwen-2.5-VL과 같은 시각 언어 모델을 활용한다. 프롬프트에서 생성된 질문-답변 쌍에 대해 모델이 출력하는 토큰 로그 확률(Logits)을 기반으로 Cross-entropy 손실을 계산하고, 이를 미분하여 이미지 생성 가이드로 사용한다. 이는 기존의 CLIP 기반 보상이 포착하지 못하는 세밀한 속성(예: 객체의 개수, 특정 부위의 상태)을 제어하는 데 탁월한 성능을 보인다.

구현 측면에서는 PyTorch의 Automatic Mixed Precision(AMP)을 사용하여 메모리 사용량과 지연 시간을 최적화했다. 1024x1024 고해상도 이미지에 대해 NVIDIA A100 GPU 환경에서 효율적으로 작동하며, 35~50 스텝의 적은 샘플링 횟수로도 고품질 결과를 생성할 수 있도록 설계되었다.

RewardFlow의 성능은 내부적으로 사용되는 보상 모델(VQA, SAM2 등)의 성능에 의존한다. 예를 들어, VQA 모델이 아주 작은 객체의 개수를 정확히 세지 못할 경우 해당 보상 신호가 부정확해져 편집 실패로 이어질 수 있다. 또한 여러 보상 모델을 동시에 실행해야 하므로 순수 확산 모델 추론보다는 계산 비용이 높다.