보상 기반 확산 모델 테스트 시 흐름 지도 경로 기울기 기법

초록

테스트 단계에서 사용자가 지정한 보상을 높은 점수로 얻도록 확산 모델을 개선하는 일반적인 방법은 보상의 그래디언트를 확산 역학에 직접 삽입하는 것이다. 그러나 이러한 절차는 보상이 보통 생성 과정의 최종 데이터 분포에서만 정의되기 때문에 수학적으로 불안정해진다. 기존의 해결책은 디노이저를 이용해 현재 샘플이 최종 단계에서 어떤 모습이었을지를 추정하는 것이었지만, 본 논문에서는 흐름 지도(flow map)를 직접 활용하는 간단한 해결책을 제시한다. 흐름 지도와 순간 전송을 지배하는 속도장 사이의 관계를 이용해 Flow Map Trajectory Tilting (FMTT) 알고리즘을 설계했으며, 이는 표준 테스트‑타임 보상 그래디언트 방법보다 보상 상승에 대해 이론적으로 더 우수함을 증명한다. 이 접근법은 중요도 가중을 통한 정확한 샘플링이나, 보상‑틸팅 분포의 지역 최대값을 찾는 원칙적인 탐색에 모두 활용될 수 있다. 우리는 기존의 선행 탐색 기법들과 비교 실험을 수행해 FMTT의 효율성을 입증하고, 비전‑언어 모델과 연계해 복잡한 보상 함수를 이용한 새로운 이미지 편집 형태를 구현한다. 그림 1은 테스트‑타임 탐색이 모델 편향을 극복하고, 기준 방법이 놓치는 정확한 시계 시간과 같은 분포 영역을 안정적으로 샘플링할 수 있음을 보여준다.

상세 요약

이 논문이 다루는 핵심 문제는 “테스트‑타임 보상 최적화”라는 실용적인 과제이다. 확산 모델은 기본적으로 노이즈를 점진적으로 제거하면서 데이터 공간을 탐색하는데, 사용자가 특정한 이미지 속성(예: 특정 물체가 포함된 이미지, 특정 텍스트와 일치하는 이미지 등)을 강화하고 싶을 때는 그 속성을 정량화한 보상 함수를 정의한다. 전통적인 접근법은 보상 함수 R(x)에 대한 그래디언트 ∇R(x)를 현재 샘플 x에 직접 더해 확산 방정식에 “힘”을 가하는 방식이다. 그러나 R은 보통 최종 데이터 분포, 즉 t = 0에서만 의미가 있다. 중간 단계 t > 0에서는 x가 아직 완전한 이미지가 아니므로 R(x) 자체가 정의되지 않거나 의미가 흐려진다. 이때 디노이저 Dθ를 이용해 현재 상태를 “역방향”으로 추정해 최종 이미지 ŷ = Dθ(x, t)로 만든 뒤 ∇R(ŷ)를 사용하면 어느 정도 해결되지만, 이는 두 가지 근본적인 한계를 가진다. 첫째, 디노이저의 추정 오차가 누적되면 보상 그래디언트가 왜곡된다. 둘째, 디노이저가 학습된 데이터 분포에 강하게 편향돼 있기 때문에, 원래 모델이 놓친 희귀한 모드(예: 정확히 12시를 가리키는 시계)까지 탐색하기 어렵다.

논문은 이러한 문제점을 회피하기 위해 “흐름 지도(flow map)”라는 개념을 도입한다. 흐름 지도 Φt(x)는 초기 상태 x를 시간 t = 0에서 t = T(보통 T = 1)까지 연속적으로 이동시키는 변환이며, 이는 확산 과정 전체를 하나의 매핑으로 보는 관점을 제공한다. 중요한 수학적 사실은 Φt의 시간 미분이 바로 순간 전송을 기술하는 속도장 v(x, t)와 동일하다는 점이다: ∂tΦt(x) = v(Φt(x), t). 이 관계를 이용하면 보상 함수의 그래디언트를 직접 현재 상태가 아닌, 흐름 지도에 의해 정의된 “미래” 상태에 적용할 수 있다. 구체적으로, 알고리즘은 다음과 같은 단계로 진행된다. (1) 현재 샘플 x를 흐름 지도 Φt에 따라 일정 시간 Δt만큼 전진시켜 가상의 미래 상태 x̂ = ΦΔt(x)를 만든다. (2) x̂에 대해 보상 그래디언트 ∇R(x̂)를 계산한다. (3) 이 그래디언트를 역전파하여 현재 상태 x에 대한 보정 방향을 얻는다. (4) 기존 확산 업데이트에 이 보정 방향을 추가한다. 이렇게 하면 보상은 “예측된 최종 이미지”에 기반해 평가되므로, 중간 단계에서의 정의되지 않은 R 문제를 회피하면서도 디노이저의 추정 오류를 최소화한다.

이론적으로 저자들은 FMTT가 표준 보상 그래디언트 방법보다 보상 상승률이 크다는 것을 증명한다. 핵심 아이디어는 “예측된 미래”에 대한 보상이 실제 보상의 기대값에 대한 하한을 제공한다는 점이다. 따라서 매 단계마다 더 큰 보상 향상을 기대할 수 있다. 실험에서는 (i) 단순한 색상/형태 보상, (ii) 복잡한 텍스트‑이미지 일치 보상, (iii) 비전‑언어 모델을 활용한 고차원 의미 보상 등 다양한 설정에서 FMTT가 기존 “look‑ahead” 기법(예: Classifier‑Guidance, Score‑Based Guidance)보다 샘플 품질과 다양성 모두에서 우수함을 보였다. 특히 그림 1이 시사하듯, FMTT는 모델이 원래 학습하지 못한 미세한 모드(정확히 12시를 가리키는 시계)까지도 안정적으로 탐색한다.

마지막으로 저자들은 두 가지 활용 방안을 제시한다. 첫 번째는 중요도 가중을 이용한 “정확한 샘플링”으로, FMTT가 생성한 경로에 대한 확률 밀도를 보정해 원래 확산 모델이 정의한 분포와 일치하도록 한다. 두 번째는 “원칙적인 탐색”으로, 보상‑틸팅 분포의 지역 최대값을 찾는 최적화 문제로 전환해 이미지 편집이나 디자인 작업에 직접 활용한다. 이처럼 흐름 지도 기반 접근은 기존 디노이저 의존형 방법의 한계를 뛰어넘어, 복잡하고 비선형적인 보상 함수를 자연스럽게 통합할 수 있는 새로운 패러다임을 제시한다.