학습 경로와 플로우를 함께 훈련해 이미지 생성 품질을 높인다

무엇인가

flow matching은 소스 점과 데이터 점을 보간하는 경로를 미리 고정해 두고, 그 경로의 접선을 맞추도록 속도장만 학습한다. 대규모 이미지 모델에서는 보통 선형 경로 γ_lin(t) = (1-t)x_0 + t x_1이 쓰인다. 그런데 이 경로는 수동적인 구현 세부가 아니라, 모델이 학습하는 상태와 예측해야 할 벡터 타깃을 동시에 결정한다. 이 논문은 "좋은 경로란 무엇이고, 그것을 플로우와 함께 학습할 수 있는가"라는 질문을 던지고, 그 답으로 path-flow alignment를 제안한다.

어떻게 동작하나

방법의 골격은 잔차 경로 네트워크다. 경로를 γ_ψ(t; x_0, x_1, y) = γ_lin(t; x_0, x_1) + t(1-t) R_ψ(t; x_0, x_1, y)로 매개화하는데, t(1-t) 봉투가 양 끝점에서 0이 되므로 경로 내부만 휘어지고 시작점과 끝점은 정확히 보존된다. 학습은 교대 최적화로 진행된다. 플로우 단계에서는 경로를 고정하고 속도장 v_φ를 경로 접선에 회귀시키고, 경로 단계에서는 v_φ를 고정한 채 같은 정합 손실로 γ_ψ를 갱신한다. 경로 단계 손실에는 경로의 운동 에너지 항이 계수 β로 더해져 경로를 직선 쪽으로 잡아당긴다. 논문의 Proposition 1은 임의의 endpoint-preserving 경로에 대해 이 회귀 문제의 population 최소화자가 경로 주변분포의 정확한 주변 속도가 되고, 그 속도장이 만드는 ODE가 p_0를 p_1로 운반한다는 것을 보인다. 즉 경로를 갱신할 때마다 다음 플로우 단계는 여전히 유효한 flow matching 문제다.

무엇과 다른가

문제는 여기서 발생한다. 저자들은 정렬 손실이 계속 줄어드는데도 FID가 처음에는 좋아지다가 이후 급격히 나빠지는 현상을 관찰하고 path overfitting이라고 부른다. 진단 도구는 학습된 속도장의 평균 divergence다. Proposition 2에 따르면 어떤 확률 흐름에서든 속도장의 기대 divergence는 미분 엔트로피의 시간 변화율과 정확히 같다. 따라서 t=1 근처에서 큰 양의 divergence 스파이크가 나타난다는 것은 샘플러가 끝점 부근에서 엔트로피를 급격히 늘린다는 뜻이고, 그 전에 과도하게 수축한 구간이 있었다는 뜻이다. 저자들은 이를 저엔트로피 병목이라고 부른다. 시간 구간을 10개로 나눠 Hutchinson trace 추정으로 누적 divergence를 정규화한 통계 G_φ(t)를 그리면, 정규화 없는 ρ_t=1 실행은 t≈0.9 부근에서 0 아래로 내려가 끝점 엔트로피보다 낮은 상태를 통과했음을 보여준다.

어떻게 쓰나

처방은 소스 정보 은닉이다. 경로 네트워크가 보는 x_0 입력의 일부를 확률적으로 가리는 스케줄 ρ_t를 도입하고, σ_t = (1-t)√(1-ρ_t²) 크기의 가우시안 섭동을 경로에 더한다. Theorem 1은 독립 엔드포인트 커플링 아래에서 이 확률적 경로의 중간 주변분포가 어떤 분포 q와 가우시안의 합성곱이 되며, 그 결과 엔트로피가 d/2·log(2πe σ_t²) 이상으로 하한이 잡히고 Fisher information이 d/σ_t² 이하로 상한이 잡힌다는 것을 보인다. Proposition 3은 여기에 경로의 운동 에너지 L_euc를 결합해 순간 엔트로피 감소율의 하한까지 제시한다. 스무딩이 클수록 score가 작아져 급격한 엔트로피 붕괴가 막히고, 운동 에너지 항이 그 스무딩된 분포 위에서 속도 크기를 제한한다는 구도다. 다만 이 하한은 t→1에서 무의미해지는데, 끝점 보존 때문에 스무딩이 끝점에서 사라져야 하므로 불가피하다고 저자들은 밝힌다. 경로 학습을 끄면 R_ψ≡0이 되어 ρ_t와 무관하게 선형 경로로 돌아가므로, 이 정규화는 학습된 잔차 경로가 있을 때만 효과가 있다.

전제와 한계

실험은 ImageNet-256×256 클래스 조건부 생성, SiT 백본(SiT-S/B/L/XL)에서 수행된다. 모든 실행은 동일한 사전학습 SiT 체크포인트에서 출발하며, 베이스라인은 표준 flow matching을 12k 스텝 더 학습시키고 제안 방법은 3번의 교대 사이클(경로 3k 업데이트 + 플로우 4k 업데이트)을 돈다. 12k 플로우 업데이트 예산 안에서 달성한 최고 50k-FID를 비교하면 모든 스케일에서 개선되며, 가장 큰 이득은 SiT-L의 19.1%, SiT-XL의 12.7% 상대 개선이다. 특히 SiT-XL에서는 표준 파인튜닝이 시작 체크포인트를 개선하지 못하는 반면 이 방법은 개선한다. 하이퍼파라미터는 β=0.05, ρ_t = c(1-t)이며 SiT-S/B/L은 c=1, SiT-XL은 c=0.7을 쓴다. Section 5의 모델 가이던스 확장도 같은 교대 갱신 틀에 유효 속도 ṽ_φ^ω를 끼워 넣는 방식으로 적용되어, SiT-B-MG에서 최고 FID가 7.940에서 5.585로, SiT-XL-MG에서 1.494에서 1.322로 좋아진다. FlowDCN-B에서도 이득이 보고되며, 컴퓨트를 맞춘 연속 파인튜닝 대비, 그리고 서로 다른 세 개의 생성 시드 집합에서도 개선이 유지된다고 한다. recall과 Inception Score도 함께 오른다.

개발자 입장에서 이 방법의 실용적 매력은 추론 시점을 전혀 건드리지 않는다는 점이다. 경로 네트워크는 학습 중 감독 신호를 만드는 역할만 하고, 샘플링에는 기존 속도장 v_φ와 기존 샘플러가 그대로 쓰인다. 따라서 이미 돌리고 있는 SiT 계열 체크포인트에 post-training 단계로 얹을 수 있고, 샘플링 비용도 변하지 않는다. 도입 전에 확인할 것은 세 가지다. 첫째, 가우시안 스무딩과 엔트로피 보장은 독립 엔드포인트 커플링을 전제로 하며 optimal-transport 계열 커플링에는 그대로 적용되지 않는다. 둘째, ρ_t 스케줄과 β, 교대 사이클 수가 새로 튜닝해야 할 하이퍼파라미터다. 셋째, 학습 중 평균 divergence 곡선을 로깅해 두면 경로 과적합이 시작되는 시점을 조기에 감지하는 실무 지표로 쓸 수 있다.

저자들이 명시한 한계도 분명하다. 이론은 엔트로피 병목과 생성 품질 사이의 인과관계를 확립하지 않으며, 경로 수준에서 얻은 엔트로피 제어가 불완전하게 학습된 플로우로 전이된다는 보장도 없다. 논문은 경로 주변분포에 대한 보장과 실제 학습된 샘플러의 divergence 측정을 구분해서 다루고, 후자는 이론적 결론이 아니라 경험적 증거로 해석해야 한다고 못 박는다. 또한 엔트로피 하한은 t→1에서 약해지고, 가우시안 스무딩 결과는 독립 커플링 조건에 의존한다. 사전학습된 이미지 생성기 너머로 공동 경로-플로우 학습을 확장하는 일은 남은 과제로 남겨 둔다.