확산 모델의 유동 맵을 이용한 테스트 시간 스케일링

초록

사용자가 지정한 보상에 대해 샘플이 높은 점수를 얻도록 확산 모델을 테스트 타임에 개선하는 일반적인 방법은 보도의 기울기를 확산 과정 자체의 역학에 도입하는 것입니다. 이 절차는 사용자가 지정한 보상이 일반적으로 생성 말단의 데이터 분포에서만 명확하게 정의되기 때문에 종종 잘못 설정됩니다. 이 문제에 대한 일반적인 해결책은 디노이저를 사용하여 생성 말단에서 샘플이 어떻게 되었을지 추정하는 것이지만, 우리는 플로우 맵을 직접 활용하는 간단한 해법을 제안합니다. 플로우 맵과 순간 이동을 제어하는 속도장 사이의 관계를 활용하여, 우리는 보도 기울기를 포함하는 표준 테스트 타임 방법보다 이론적으로 보상 상승을 더 효과적으로 수행하는 알고리즘인 Flow Map Trajectory Tilting(FMTT)을 구성합니다. 이 접근법은 중요도 가중치를 통한 정확한 샘플링을 수행하거나 보도 기울어진 분포의 지역 최대점을 식별하는 원칙적 탐색에 사용될 수 있습니다. 우리는 다른 선행 탐색 기법들과 비교하여 우리 접근법의 효용성을 입증하고, 플로우 맵이 어떻게 비전 언어 모델과의 인터페이싱을 통해 새로운 형태의 이미지 편집을 가능하게 하는 복잡한 보상 함수와의 연동을 가능하게 하는지 보여줍니다.

English

A common recipe to improve diffusion models at test-time so that samples score highly against a user-specified reward is to introduce the gradient of the reward into the dynamics of the diffusion itself. This procedure is often ill posed, as user-specified rewards are usually only well defined on the data distribution at the end of generation. While common workarounds to this problem are to use a denoiser to estimate what a sample would have been at the end of generation, we propose a simple solution to this problem by working directly with a flow map. By exploiting a relationship between the flow map and velocity field governing the instantaneous transport, we construct an algorithm, Flow Map Trajectory Tilting (FMTT), which provably performs better ascent on the reward than standard test-time methods involving the gradient of the reward. The approach can be used to either perform exact sampling via importance weighting or principled search that identifies local maximizers of the reward-tilted distribution. We demonstrate the efficacy of our approach against other look-ahead techniques, and show how the flow map enables engagement with complicated reward functions that make possible new forms of image editing, e.g. by interfacing with vision language models.

확산 모델의 유동 맵을 이용한 테스트 시간 스케일링

Test-time scaling of diffusions with flow maps

초록

Support