one-step generative modeling을 목표로하는 효율적인 프레임워크를 제시한다.
average velocity to characterize flow fields. 기존 flow matching은 국소적인 범위에서의 veclocity 모델링이 목표였다.
requires no pre-training, distillation, or curriculum learning(distill 아니라고 해놓고 거의 원본 모델부터 학습하고 진행하는 형태가 있던데 그게 아니라는 뜻인가?)
Imagenet 256으로 학습, 1-step으로 FID 3.43
이것만 놓고보면 차이가 꽤 크다. 131M으로도
Consistency model은 네트워크가 같은 path에서 나온 input에 대해서 consistent한 output을 뱉도록 학습했다. 결과는 좋았지만, 항상 출력이 consistent로 고정되어있고(few-step이 잘 안됨), 이론적으로 GT를 정의하고 하는게 아니다보니 학습도 불안정했다.
평균 속도라는 개념이 뭘까? flow matching에서 training objective를 생각해보면 어쨌든 target velocity가 항상 같기 때문에.. 평균과 instanteneous가 같다고 생각되는??가?
well-defined, intrinsic relation between the average and instantaneous velocities
No extra consistency heuristic is needed.
GT target field가 존재한다는건 이론적으로 최적해로 가도록 네트워크를 학습시킬 수 있다는 것이기 때문에 네트워크 구조에 독립적이며, stable & robust 해진다.
Flow Map, Shortcut, Inductive Moment Matching
Background : Flow matching
조건부 velocity로 각 condition에 대해 학습하고, 그걸 marginalize하면(평균내면) 오른쪽 그림처럼 자연스러운 filed가 생긴다.
주어진 z와 v는 다양한 x와 쌍에서 나올 수 있기 때문에, Flow Matching은 본질적으로 모든 가능한 경우에 대한 기댓값을 모델링한다. = marginal velocity
즉, conditional flow는 직선으로 설계되었지만 이로 인해 실제로 만들어지는 marginal velocity field는 곡선이 될 수 있다.
이건 학습 특성으로 인한 당연한 성질이다. → 그럼 few-step으로 할수록 deiscretization 문제가 생긴다.
MeanFlow
meanflow의 핵심은 average velocity를 나타내는 새로운 필드를 모델링하는 것이다.
Average velocity = 두 timestep
t와 r 사이의 displacement(위치변화량)을 시간 간격으로 나눈 값. (원래 맞잖아 이거?)여기서 displacement는 순간 속도
v를 시간에 따라 적분해서 얻는다.따라서 식으로는 이렇게 된다.
주의할점!
u는 v의 함수로부터 계산된 결과이지, neural network가 직접 출력하는게 아니다. 즉! 뉴럴넷의 출력은 v인데 학습 objective는 이 v를 이용해서 u를 계산한다음 GT u와 비교해서 얻는다
u는 자연스럽게 특정 boundary condition을 만족한다.
- r → t 일 때 u → v 이다.
- r → t와 r → s → t가 같다.
즉, 위에서 average velocity를 정의한 방식이 올바른 수학적 구조를 갖고 있다는걸 의미한다. = 믿고 써도 된다.
하지만 위의 식으로 학습하는건 학습 과정에서 매번 model의 output을 적분해야하기 때문에 비효율적이다.
따라서 평균 속도 정의식을 적절히 변형해서 instantaneous velocity만 가지고 training objective를 만드는걸 목표로 한다.
근데 여기서 우항의 time derivative를 풀기 위해서 partial derivatives를 또 계산하는데 슬슬 헷갈린다. (t, r, z)가 전부 t와 관련이 있는 변수라고 가정하면 각 변수를 t로 미분한 값과 각 변수로 편미분한 값을 전부 더해서 계산(아래처럼)
그런데 여기서 dt/dt = 1, dr/dt = 0, dz/dt = v 이다. 따라서
결국 학습은 6번 식을 만족하도록 일어난다. 적분을 하지않고 target을 정의가능하다.
u의 derivative가 필요하긴함.
flow matching과 같은 방식으로 v는 단일 샘플 조건부로 계산한다. 즉, 기존 flow matching의 target velocity v_t = - x에다가 오른쪽 식만큼 보정되어있다고 볼 수 있다.
그러므로 t=r로 제한하면 사라지는걸 볼 수 있다.
target 계산에도 u(를 편미분한)가 들어가는데 이럼 back prop이 복잡해져서 gradient는 흘려보내지 않는다.
1. 🔧 목적은 uθ 학습이다
- 우리는 평균 속도 를 예측하는 네트워크 를 학습하고 싶음.
- 하지만 GT u 자체는 적분으로 정의되어 있어서 직접 학습 때 계산해서 쓰기 어려움.
- 그래서 적분 대신, 순간 속도 v만을 이용해 u의 대체 타겟 을 계산함. (모델의 출력은 여전히 average velocity)
여기서 복잡한건 우항의 인데, 이건 u의 jacobian matrix와 벡터의 dot product로 볼 수 있다. 그리고 파이토치와 jax에서 이걸
torch.func.jvp로 쉽게 지원한다! (왜지 많이 쓰는건가?)pseudo code
# x ~ training batch
# fn(z, r, t): function to predict u
t, r = sample_t_r()
e = randn_like(x)
z = (1 - t)*x + t*e
v = e - x
u, dudt = jvp(fn, (z, r, t), (v, 0, 1))
u_tgt = v - (t - r) * dudt
error = u - stopgrad(u_tgt)
loss = metric(error)inference
e = randn_like(x)
x = e - fn(e, r=0, t=1)With CFG
CFG를 적용가능한데, 기존과 같은 방식은 아니다. 기존 방식은 배치가 생겨서 사실 NFE 두번 하는 것과 같다(다르지 않나)
대신 학습 때 쓰는 Grount-truth field 자체에 cfg가 반영되도록 해서 학습하고, inference는 한번만에 끝낸다.
(그냥 2배하는 것도 뭐..)
Loss
단순 MSE(L2)보다 오차가 작을 때는 덜 민감하고, 오차가 클 때는 더 민감하게 가중치로 조절한다.
Conditioning on (r, t)
모델에 r을 주는게 아니라 t-r로 구간 길이를 주는게 더 직관적일 수 있다. 실제로 더 나아졌다. 아래에 비교했는데 작은 차이긴 함
디테일한 ablations on ImageNet 256
- t랑 r이 다른게 25%밖에 안된다..? 큰 차이는 없긴하다.
- 왜 비교한건지 모르겠음
- conditioning 방식 비교 차이 적음
- r과 t는 독립적으로 logit-norm에서 샘플링, 값은 lognorm(-0.4, 1.0)
- p=1.0
- CFG weight가 3.0일 때 제일 좋다. 4배가 차이난다.
누가 구현해둔 것도 있다.
MeanFlow
haidog-yaqub • Updated Jul 15, 2025
jacobian vector product가 한줄로 지원된다고는 하지만 내부 연산은 복잡한만큼 학습이 꽤 불안정한 것 같다(아마 특정 차원값이 너무 크거나 작아서?). speech로 해보려고 하니까 수렴이 안된다고 하는 git issue가 있음
그냥 MNIST로 완전 수렴 속도 보고싶다. 어짜피 ImageNet을 돌려도 수렴할 때까지 돌리는건 아니니까 뭔가 애매한 느낌?
개인 메모
Share article
Subscribe to our newsletter
