보상 모델링을 통한 이미지 생성의 공간적 이해 능력 향상

기존 이미지 생성 AI는 사물 간의 왼쪽, 오른쪽, 뒤쪽과 같은 복잡한 공간 관계를 정확히 표현하지 못하는 고질적인 한계가 있었다. 이 논문은 전용 데이터셋과 보상 모델을 통해 AI가 사물의 배치를 논리적으로 이해하고 생성하도록 개선하여, 정교한 이미지 제어가 필요한 디자인 및 시뮬레이션 분야의 실무적 난제를 해결한다.

기존 Diffusion Model은 Attention Mechanism을 통해 텍스트와 이미지의 연관성을 학습하지만, 주로 색상이나 질감 같은 시각적 특징에 치중하여 사물 간의 상대적 위치와 같은 기하학적 논리를 놓치는 경우가 많다. 이는 학습 데이터의 캡션이 공간적 제약 조건을 충분히 반영하지 못하기 때문에 발생하는 문제이다. 이 논문은 이를 해결하기 위해 '공간적 정답'을 판별하는 전용 보상 모델인 SpatialScore를 도입한다. 마치 바둑 AI가 승률을 예측하듯, 생성된 이미지가 프롬프트의 위치 지시를 얼마나 잘 따랐는지 점수를 매긴다. 특히 GPT-5를 활용해 의도적으로 위치를 틀린 '오답 이미지'를 생성하고 이를 '정답 이미지'와 비교 학습시켜 미세한 위치 차이를 구분해낸다. 학습된 보상 모델은 강화학습의 피드백 엔진이 되어, 모델이 이미지를 생성할 때마다 더 정확한 배치를 구현하는 방향으로 가중치를 업데이트하게 유도한다. 결과적으로 AI는 단순한 픽셀 조합을 넘어 사물 간의 물리적 거리를 이해하는 능력을 갖추게 된다.

SpatialReward-Dataset 구축 및 보상 모델 학습 단계에서는 GPT-5를 사용하여 복잡한 공간 관계 프롬프트를 생성한 뒤, 관계를 하나씩 비튼 변형 프롬프트를 생성한다. Qwen-Image 등 최신 모델로 이미지를 생성하여 정답/오답 쌍을 구성하고, Qwen2.5-VL-7B를 Backbone으로 Bradley-Terry Model 기반의 Loss를 사용하여 SpatialScore를 학습시킨다. 온라인 강화학습 파이프라인에서는 FLUX.1-dev 모델을 대상으로 GRPO 알고리즘을 적용한다. 한 프롬프트당 여러 장의 이미지를 생성하고, SpatialScore가 매긴 점수를 그룹 내에서 정규화하여 Advantage를 계산한다. [생성 이미지 점수 → 그룹 평균 차감 및 표준편차 정규화 → Advantage 산출] 과정을 통해 모델이 상대적으로 더 나은 배치를 학습하도록 유도한다. 또한 Top-k Filtering 전략을 통해 그룹 내 상위 k개와 하위 k개 샘플만 선택하여 학습에 사용함으로써 프롬프트 난이도에 따른 Advantage Bias를 해결하고 학습 효율을 높인다.

SpatialScore 모델은 자체 벤치마크에서 95.8%의 정확도를 기록하며 GPT-5(89.0%)와 Gemini 2.5 Pro(95.1%)를 능가하는 성능을 보였다. 기존의 HPSv3(65.2%)나 PickScore(50.9%) 같은 일반 보상 모델들이 공간 관계 평가에서 낮은 성능을 보인 것과 대조적이다. FLUX.1-dev 모델에 적용한 결과, DPG-Bench의 공간 관계 점수가 2.18에서 7.81로 대폭 상승했다. 또한 TIIF-Bench와 UniGenBench++ 등 외부 벤치마크에서도 기본 모델 및 기존 강화학습 방식 대비 일관된 성능 향상을 입증했다. Ablation Study를 통해 Backbone 모델 크기가 커질수록 성능이 향상됨을 확인했으며, Top-k Filtering을 통해 학습 속도를 가속화하면서도 최종 성능 저하 없이 효율적인 최적화가 가능함을 확인했다.

아키텍처는 Qwen2.5-VL-7B를 기반으로 하며, 마지막 레이어에 특수 토큰을 삽입하여 텍스트와 이미지 임베딩을 동시에 참조하도록 설계했다. 이 토큰의 임베딩을 MLP 헤드에 통과시켜 가우시안 분포의 평균(μ)과 표준편차(σ)를 출력하며, 이를 통해 보상 점수의 불확실성을 모델링한다. 데이터셋 구축 시 Adversarial Perturbation 기법을 사용하여 GPT-5가 원본 프롬프트에서 최소한의 수정을 가해 오답 쌍을 생성함으로써, 모델이 시각적 품질이 아닌 순수하게 공간적 논리에만 집중하여 학습하도록 유도했다. RL 학습 시 Flow Matching 기반 모델인 FLUX를 위해 결정론적 ODE 샘플링을 확률적 SDE 샘플링으로 변환하여 정책 탐색을 가능하게 했다. Euler-Maruyama 스킴을 사용하여 SDE를 이산화하고, GRPO의 그룹 내 상대적 보상 체계를 통해 KL Divergence 페널티와 함께 안정적인 업데이트를 수행한다.

현재 연구는 정적인 이미지의 공간 관계에 집중하고 있어, 비디오 생성과 같이 시간에 따른 동적인 공간 변화를 다루는 데에는 한계가 있다. 또한 로봇 시뮬레이션과 같은 실세계 물리 법칙과의 완벽한 정렬은 향후 과제로 남아 있다.