AnyFlow: On-Policy Flow Map Distillation을 이용한 Any-Step 비디오 확산 모델

일련의 비디오 확산 모델에서 흔히 쓰이는 consistency distillation은 테스트 시점에 샘플링 스텝이 늘어나면 성능이 저하되는 경향이 있다. AnyFlow는 flow map 기반으로 임의의 시간 쌍 간 전이를 학습해 어떤 스텝 수에서도 안정적으로 샘플링 가능하도록 한다. 두 단계 학습(Forward flow map training + On-policy flow map distillation)으로 PF-ODE 경로를 보존하면서도 discretization error와 exposure bias를 감소시킨다. 1.3B~14B 파라미터 규모의 bidirectional/causal 아키텍처에서 기존 CM 기반 방법과 비슷하거나 더 나은 few-step 성능을 보이고, 샘플링 스텝이 늘어날수록 성능이 계속 향상된다.

단계적 흐름(flow) 이론에 기반한 비디오 확산 모델의 일반화 방향을 제시한다. 기존의 endpoint-centric consistency 모델은 z_t에서 z_0로의 단일 엔드포인트 매핑을 학습하기 때문에 다중 샘플링에서 경로가 PF-ODE 경로에서 벗어나기 쉽다. AnyFlow는 z_t에서 z_r로의 임의 시간 쌍 간 전이(f_theta(z_t, t, r) ≈ z_r)를 학습하는 flow map 모델로 일반화하고, 구성 성질을 이용해 경로를 여러 단계로 분해하는 backward simulation을 도입한다. 이 방식은 크게 네 가지로 구분된다. 1) Forward flow map training으로 초기화 학습, 2) Interpolated timestep conditioning 및 Time sampler 도입으로 포스팅-트레이닝의 안정성 확보, 3) Guidance-fused training으로 CFG를 학습에 통합하되 테스트 시에는 CFG를 제거 가능, 4) On-Policy flow map distillation의 flow map backward simulation으로 rollout 드ift를 보정하고, 다양한 step budget에 대해 함께 학습하되 연산 비용은 감소시키는 shortcut 전이를 활용한다. 이 결과, few-step에서의 품질을 유지/향상시키고, 더 많은 NFEs에서의 성능 증가를 달성한다.

1. Flow Map Training: Forward flow map 학습은 MeanFlow objective를 사용해 f_theta(z_t, t, r) ≈ z_r를 학습한다. 입력으로 z_t, t, r를 주고, z_T에서 z_t로의 초기 샘플링 후 속도 벡터 v(z_t, t)와 du_theta/dt를 이용해 utgt를 계산한다. 손실은 L(theta) = E[ || u_theta(z_t, r, t) − sg(u_tgt) ||^2_2 ]로 정의되며, sg는 stop-gradient다. 역微分 항의 수식 도출은 d u(z_t, r, t)/dt를 해석적으로 다루기 어려워, 근사치를 사용한다(Δt 차분 등). 타임스텝 conditioning은 emb(t)와 emb′(r)의 선형 조합으로 처리하며, g를 0.25로 고정해 boundary case t=r일 때 pretrained embedding과의 연속성을 보장한다. Time sampler로 t와 r을 Uniform에서 max/min으로 재정렬하고, w(t)로 노이즈 수준에 따라 손실 가중치를 조정한다. Guidance-fused training은 CFG를 u에 직접 적용해 테스트 시 CFG 비용을 제거할 수 있도록 한다. Adaptive loss reweighting은 t=r 경계 샘플의 손실을 기준으로 t≠r 샘플의 손실 규모를 안정화한다. 2) On-Policy Flow Map Distillation: self-rollout 상태를 z_T → z_t → z_r → z_0의 흐름으로 역전 시퀀스화하고, Flow Map Backward Simulation으로 rollout을 분해해 중간 단계에서의 그래디언트를 흐름 체인 전체로 역전시키는 KL 기반 DMD를 적용한다. 흐름 체인의 구성은 T, t, r, N으로 제어되며, f_theta(z_T, T, t) → z_t → z_r → z_0의 경로를 만들고 z_0에서의 KL_gradients를 역전 전달한다. 이로써 test-time discretization error 및 노출 편향을 낮추는 한편, 다양한 inference budgets에 대해 동일한 학습 파이프라인으로 조정이 가능하다.

주요 벤치마크 및 ablation 결과를 수록한다. 1) Table 2의 Bidirectional/CAusal 비디오 확산 모델에서 Flow Map Backward Simulation (Ours) 조합이 32 NFEs 조건에서 가장 높은 종합 성능(Overall)을 기록한다. 예를 들어 Bidirectional에서 32 NFEs일 때 Flow Map Backward Simulation의 Overall은 83.96으로 나타났고, Causal에서도 동일 조건에서 83.96으로 보고된다. 2) AnyFlow-FAR-Wan2.1-14B는 4 NFEs에서 84.05, 32 NFEs에서 84.41의 VBench 점수를 달성해 Krea-Realtime-14B(83.25, 4 NFEs)보다 우수하다. 3) Text-to-Video(VBench)에서 AnyFlow-FAR-Wan2.1-14B는 4 NFEs에서 87.87의 I2V 점수로 Wan2.1-14B의 87.71 대비 소폭 상회한다. 이미지-투-비디오(VBench-I2V)에서도 AnyFlow-FAR-Wan2.1-14B가 4 NFEs에서 87.87로 최상위에 근접한다. 4) Time sampler ablation에서 Beta(2,1.5) 가중치와 Interpolated Time Embedding이 가장 우수하며, Uniform 가중치는 32 NFEs 시나리오에서 가장 열등하다. 16 NFEs에서 Flow Map Backward Simulation은 Consistency Backward Simulation에 비해 학습 비용을 43.4% 감소시키고 Bidirectional/CAusal에서 각각 47.0% 감소시키는 등 효율성 측면에서도 개선된다.

1. 전체 아키텍처: 사전 학습된 비디오 확산 백본 위에 flow map 모델 f_theta(z_t, t, r)을 두 시간대 플로우 맵으로 학습시키고, Stage 2에서 온-정책 distillation으로 역전 드리븐 KL 손실을 통해 rollout 드리프트를 보정한다. 2) Flow Map Training의 수학적 기반: PF-ODE d z_t/dt = v(z_t, t)에서 z_r = Φ_{r←t}(z_t)이고, f_theta(z_t, t, r) ≈ z_r로 근사한다. MeanFlow 목표는 L(theta) = E[ || u_theta(z_t, r, t) − sg(u_tgt) ||^2_2 ]를 최소화하는 것으로, u_tgt = v(z_t, t) − (t − r) du_theta(z_t, r, t)/dt 로 정의된다. 3) 역전 파이프라인 구성: Flow Map Backward Simulation은 f_theta(z_T, T, t) → z_t, f_theta(z_t, t, r) → z_r, f_theta(z_r, r, 0) → z_0를 통해 z_0에서 KL gradient를 계산하는 역전 경로를 구성한다. 4) 학습 기법의 구체: Interpolated timestep conditioning(emb(t) + (1 − g)·emb′(r), g=0.25), Time sampler(Uniform에서 재정렬), Guidance-fused Training(CFG을 prediction에 적용), Adaptive loss reweighting(경계 샘플과 비경계 샘플의 손실 규모를 조정)으로 안정성을 확보한다. 5) On-Policy Distillation의 구현 방식: 샘플 스텝 s를 Uniform로 선택하고, t를 Uniform(1..s), r = t − T/s로 설정한 후 z_T에서 z_0까지의 rollout 상태를 생성하고 s_real, s_fake에 대한 DMD 경로로 KL 그래디언트를 계산한다.

외부 데이터셋에 의존한 flow map 학습의 한계가 존재하며, 학습 분포와 base 모델의 분포 차이로 경미한 분포 변화가 생길 수 있다. 이 경우 더 매끄러운 텍스처 등으로 이어질 수 있으며, 이러한 문제는 pretraining 데이터와 동일한 데이터로 AnyFlow를 적용하거나 도메인 맞춤형 파인튜닝으로 보완 가능하다는 점이 논의된다.