FlowBender: 자기 보정형 조건부 Flow 모델

조건부 샘플링에서 조건 신호(y)와 데이터 분포 간의 불일치를 줄이는 것이 어렵다. FlowBender는 alignment error를 입력으로 받아 샘플링 경로를 학습하는 폐쇄 루프 방식으로, 기존의 Open-loop 학습과 추론-时의 가이드 방식이 갖는 한계를 극복한다. gradient-based 및 non-differentiable 환경에서 모두 작동하는 보정 정책을 학습하고, prior-step shortcut으로 추론 비용을 거의 증가시키지 않는다. image-to-image translation, restoration, 3D mesh texturing에서 Fidelity와 Plausibility를 동시 향상시키는 것이 가능하다.

출발점은 Flow Matching 모델이 conditioning signal에 의해 샘플링 경로를 좁히는 과정에서 정합성의 손실이 생긴다는 점이다. FlowBender는 alignment error를 네트워크의 입력으로 제공하고, look-ahead를 통해 초깃값을 추정한 뒤 이 오차를 바탕으로 두 번째 패스에서 보정 벡터를 생성한다. 이 보정은 gradient-based 피드백이나 zero-order 피드백 중 하나를 선택적으로 활용하며, 기존의 가이드 방식이 가진 단순 가중합의 한계를 넘는 비선형 정책으로 샘플의 경로를 데이터 분포에 더 잘 맞추게 한다. 또한 prior-step shortcut으로 추론 비용을 크게 낮춘다.

FlowBender는 two-pass 학습: PASS 1에서 s_t를 0으로 두고 look-ahead velocity v_LA를 예측한다. x̂1은 x_t와 v_LA를 이용해 추정된 clean signal이다. 이후 s_t를 계산하고, Pass 2에서 v_ref를 계산한다. s_t의 정의는: First-Order Feedback의 경우 ∇{x_t}L(H(x̂1), y) 혹은 ∇{x̂1}L(H(x̂1), y)를 사용한다. Zero-Order Feedback의 경우 s_t^{err} = R(H(x̂1), y)로 정의되고, 두 피드백을 합친 Hybrid Feedback도 가능하다. Inference 시에는 two-pass를 유지하되 t_thresh에 따라 Look-ahead를 건너뛰고 이전 step의 x̂1_prev를 재활용하는 Prior-Step Shortcut를 적용할 수 있다. 학습 손실은 L_FA = || v_ref − u_t ||^2로 정의되며, sg[·]로 교차-그라운딩 없이 두 패스를 학습한다. CFG를 추가 가능하며, CFG를 통해 추가적인 샘플링 제어를 할 수 있다.

Image-to-Image Translation에서 four tasks(SR, Depth→RGB, Edge→RGB, JPEG restoration)에서 FlowBender 변형이 baselines를 능가했다. 예를 들어 JPEG restoration에서 Zero-order가 Standard FT 대비 PSNR 2.57dB 향상(26.29→28.86), SSIM은 크게 증가하고 LPIPS는 감소했다. Depth→RGB, Edge→RGB에서 FlowBender 변형은 FID가 더 낮고 Fidelity 지표도 상향되었으며, SR task에서도 PSNR/SSIM/LPIPS가 개선되었다. 3D Mesh Texturing에서 Objaverse/Toys4K 데이터셋에 대해 Our(w.r.t. x^1) 방식이 MV-MPSNR, MV-SSIM 및 FID에서 우수했고, FID도 낮아 다-view 일관성이 개선되었다. Ablation으로 p_un=0.1이 최적값으로 확인되었고, Prior-Step Shortcut의 t_thresh 조절에 따라 추론 속도와 성능의 트ade-off가 나타났다.

Two-pass closed-loop 아키텍처; PASS 1에서 look-ahead를 통한 x̂1 추정, PASS 2에서 s_t를 이용한 보정 벡터 v_ref 산출. First-Order Feedback(gradient-based)와 Zero-Order Feedback(피드백 잔여(residual) 기반) 변형, Hybrid variant 가능. Training은 sg[·]로 stop-gradient 처리된 s_t를 이용하여 L_FA를 최소화하며, p_un를 통해 conditioning dropout을 적용. Inference는 t_thresh에 따라 Previous step의 x̂1_prev를 재사용하는 Prior-Step Shortcut 적용 가능. Baselines: Standard FT, FT+L_align, IT Guidance 및 FlowBender 변형들을 비교. 실험은 이미지-투-이미지 변환과 3D 텍스처링에서 수행되었고, 각 Task마다 PSNR/SSIM/LPIPS/FID 등 다수의 평가 지표를 사용한다.

학습 단계에서 추가적인 피드백 신호를 도출하기 위해 한 단계의 연산이 더 필요하므로 학습 비용이 증가한다. CFG를 추가하면 성능이 더 향상되기도 하지만 더 많은 하이퍼파라미터 의존성이 생길 수 있다. t_thresh가 0으로 설정될 경우 2-pass 구조의 이점을 일부 잃을 수 있다.