Flow Matching을 위한 보상 역전파(Reward Backpropagation)의 설계 공간 탐색

직접적인 보상 역전파 (direct reward backpropagation)를 통해 텍text-to-image flow matching 모델을 인간의 선호도에 맞추는 방식은 샘플 효율적이지만, 두 가지 잘 알려진 병리적 현상 (pathologies)으로 인해 어려움을 겪습니다. 즉, 현대적인 모델 규모에서는 전체 샘플링 궤적 (sampling trajectory)에 걸쳐 활성화 값 (activations)을 저장할 수 없으며, 단계별로 연결된 자코비안 곱 (chained Jacobian products)이 초기 인덱스로 역전파될 때 보상 그래디언트 (reward gradient)를 팽창시킨다는 점입니다. LeapAlign과 같은 커넥터 기반 (Connector-based) 방법들은 전체 역방향 궤적 (backward trajectory)을 짧은 고정 경로 (pinned path)로 대체함으로써 이러한 문제를 해결하며, 샘플링과 최적화 사이의 유용한 디커플링 (decoupling)을 강조합니다. 그러나 결과적인 그래디언트의 품질은 이 짧은 경로가 전체 롤아웃 (full rollout)을 얼마나 정확하게 근사하는지, 특히 긴 간격에 대해 어떻게 근사하는지에 달려 있습니다. 우리는 역방향 궤적 자체를 설계 객체로 취급하는 통합 대리 궤적 (surrogate-trajectory) 프레임워크인 FlowBP를 제안합니다. FlowBP는 샘플링을 위해 그래디언트가 없는 캐시된 롤아웃 (no-gradient cached rollout)을 유지한 다음, 캐시된 속도와 선택적으로 재전향 (re-forwarded)된 속도로부터 가벼운 역방향 대리 모델 (backward surrogate)을 구축합니다. 이러한 관점은 보상 모델 입력 (reward-model input), 활성 집합 (active set), 적분 가중치 (integration weights), 그리고 브리지 결합 (bridge coupling)이라는 네 가지 선택지를 분리하며, 기존의 직접 그래디언트 (direct-gradient) 방법들을 특정 설정으로서 복원합니다. 이 프레임워크 내에서 우리는 세 가지 변형을 구현합니다: FlowBP-Sparse는 희소 오일러 재구성 (sparse Euler reconstruction)을 사용하고, FlowBP-Bridge는 제어된 브리지 결합 (controlled bridge coupling)을 추가하며, FlowBP-Lagrange는 leap quadrature의 차수를 높입니다. 세 가지 방식 모두 활성 집합 (active-set) 크기에 의해 메모리를 제한하며, 그래디언트 체이닝 (gradient chaining)을 최대 하나의 자코비안 인자 (Jacobian factor)로 제한합니다. 선호도, 품질 및 구성적 지표 (compositional metrics)에 대해 SD3.5-M, FLUX.1-dev, FLUX.2-Klein-base를 대상으로 테스트한 결과, 세 가지 변형 모두 대부분의 지표에서 직접 그래디언트 베이스라인 (direct-gradient baselines)보다 향상된 성능을 보였습니다.