Abstract
Diffusion models and flow-matching models have enabled generating diverse and realistic images by learning to transfer noise to data.
However, sampling from these models involves iterative denoising over many neural network passes, making generation slow and expensive.
Previous approaches for speeding up sampling require complex training regimes, such as multiple training phases, multiple networks, or fragile scheduling.
We introduce shortcut models, a family of generative models that use a single network and training phase to produce high-quality samples in a single or multiple sampling steps.
single network와 학습 방식을 사용
Shortcut models condition the network not only on the current noise level but also on the desired step size, allowing the model to skip ahead in the generation process. 현재의 noise x_t, timestep t 뿐만 아니라 원하는 step size d도 네트워크의 컨디션으로 사용한다.
Across a wide range of sampling step budgets, shortcut models consistently produce higher quality samples than previous approaches, such as consistency models and reflow
Compared to distillation, shortcut models reduce complexity to a single network and training phase and additionally allow varying step budgets at inference time.
Distillation과 다른 접근법을 쓰는 듯?
iterative denoising methods(diffusion, flow matching)는 여러 분야에서 좋았다. alpha fold 3도 사용했구나
단점은 iterative한 특성으로 인한 inference의 시간/비용
→ diffusion training의 장접을 가지면서 single step만으로 denoising을 끝내보겠다.
이미지로만 봤을 때는 one step에 저하가 약간 있지만, 기존의 flow matching 모델과 비교하면 차이가 크다.
CelebA-HQ and Imagenet-256
end-to-end 모델 구조, single model and single training run.
기존에 가장 비슷한건 원래 방식대로 학습한 좋은 diffusion model을 가지고 one-step 능력을 가진 모델을 distillation으로 학습하는 것.
Consistency model이 end-to-end에 좀 더 가깝지만, 미세한 임의의 튜닝이 필요하다.
Key insight : 학습 과정에서 noise level(아마 timestep t) 뿐만 아니라, desired step size도 condition으로 사용해서 학습한다 - training time에 self-distillation 된다고 볼 수도 있다. 그래서 학습을 한번만에 끝내버리는 것
→ 그럼 학습 자체가 더 어려워지진 않나? 기존보다 최대 16% 정도 학습이 더 필요했다고 한다.
CelebA-HQ와 Imagenet-256으로 테스트해봤고, single step이 아닌 경우의 퀄리티 저하는 딱히 없었다고 한다.
robotic control diffusion policy에도 테스트 해봤다고함. 여기에 쓰고 싶었나보다. Diffusion, flow matching 둘다 실험
기존의 flow matching으로 학습된 모델은 왜 few-step generation에서 실패하는가?
이미지에서는 학습 데이터가 2개라고 가정한다. 각 데이터로 학습할 때, target flow는 xt, t와 무관하게 항상 x1 - x0로 같다.
데이터가 하나라면 당연히 문제될게 없다. 하지만 수많은 데이터(여기선 2개) point에서 학습할 때의 inherent uncertainty로 인해(각기 다른 데이터를 학습하는데 xt, t가 같을 때도 있다. 하지만 출력값은 다르게 학습된다) 해당 학습으로 형성되는 ODE는 위 그림처럼
While flow-matching models learn a deterministic ODE, its paths are not straight and have to be followed closely.
학습은 deterministic하게 일어나지만 결과적으로 형성되는 path는 직선이 아니고 따라서 경로를 순차적으로 따라가면서 목적지에 도달해야한다.
그림에서 볼 수 있듯이 각 step에서의 predicted diration은 모든 데이터의 평균지점을 가리키는 경향이 생긴다.
특히 x0에서는 모든 데이터가 동일한 input으로 학습하게 되기 때문에, 1 step은 데이터셋의 mean을 가리킬 수 밖에 없다.
To sample from a flow model, a noise point x0 is first sampled from the normal distribution. This point is then iteratively updated from x0 to x1 following the the denoising ODE defined as following the learned flow model v¯θ(xt, t). In practice, this process is approximated using Euler sampling over small discrete time intervals.
다시 핵심 결론
Shortcut models condition the network not only on the current noise level but also on the desired step size, allowing the model to skip ahead in the generation process. 현재의 noise x_t, timestep t 뿐만 아니라 원하는 step size d도 네트워크의 컨디션으로 사용한다.
shortcut model 로 를 보내는 normalized direction을 표현한다.
따라서 flow-matching model은 오직 현재 시점에서의 velocity를 배우고, shortcut model은 larger jump를 위한 방향을 배운다.
모델 s가 d step만큼 가는 방향을 배운다고 해보자. 그럼 target flow는 아마 가 될텐데, 여기서 x_t+d를 알지 못하기 때문에 목적함수 계산이 어렵다. 그렇다고 기존 방법대로 sampling해서 x_t+d를 얻고 이걸 사용하려면, 매번 샘플링을 해야한다는 비용이 따라온다.
그래서 논문에서는 이렇게 계산함.
d=0일 때는 기존의 flow matching loss와 같은걸로 학습한다.
In principle, the trained model can handle any distribution of d ~ p(d)
but in practive, split the batch into a friction.
thus arrive at the combined shortcut model loss function
loss를 보면 학습을 d=0부터 하고 아래로 순차적으로 하는게 아니라 한번에 둘다 jointly 학습하는걸 볼 수 있다. s_target을 계산하는데 model을 사용하지만 gradient를 보내진 않는다.
그럼 논문 처음에서 말한 추가적인 16%의 연산은 각 loss 계산에서 추가되는 연산인가? 더 클 것 같긴한데. 그리고 학습 자체도 더 많아져야 할텐데
전체 스텝 수를 정해야하는데 128로 정했다.
근데 또 보니까 first k batch 학습이 따로 있다. 그래야 학습이 안정적일테니 이게 맞지만 그럼 크게 차이가 뭘까 기존의 distillation이랑..??
Exponential moving average weights
최근 diffusion model들에서 sample quality와 학습 안정성을 위해 exponential moving average weights를 쓴다.
EMA란?
수식
- 매 스텝 모델의 파라미터를 갱신
- 이후! EMA 파라미터를 업데이트
alpha는 0.999 ~ 0.9999로, 과거의 파라미터를 더 크게 반영하도록.
- 그리고 이 ema 파라미터는 sampling 과정에서 사용
여기서는 step size d의 target을 계산할 때도 사용
학습은 기존대로 파라미터를 업데이트 하고, ema는 샘플링 시에만 사용한다.
ema 가중치는 보다 천천히 변한다.
특히 diffusion에서는 학습시에 데이터 샘플/timestep에 따른 loss의 크기 차이가 크다(objective has inherent variance).
- 특히 diffusion loss는 variance(분산)이 커서 훈련 중 노이즈가 심합니다.
주로 타임스텝이 작은 경우에 그런데, 그럼 파라미터가 크게 변화하고 → 다른 타임스텝에서의 출력값이 크게 영향을 받는다. 이걸 완화하기 위해서 특히 diffusion에서 ema를 활용한다.
사실 얼마나 필수적인지는 잘 모르겠다. (출처 : https://arxiv.org/pdf/2101.08482) DDPM, DiT 등에서도 썼었구나 몰랐다.
본 논문에서 : Evaluating with EMA parameters provides a consistent improvement, but is not strictly neccessary.
여러 방식들과 비교
- reflow : 기존의 flow(nosie → data)가 아니라 반대방향(data → noise)을 학습시킨다.
- teacher 모델이 필요하다.
- noise → data는 trajectory가 다양하고 non-deterministic해서 학습이 어렵고 불안정하다. + High variance loss
- 그냥 128 step generation 성능도 일반적인 flow matching 방식으로 학습했을 때보다 더 좋게 나오는데, 이유가 뭘까?
아래에 설명이 나온다 : Interestingly, in our experiments shortcut models display a slightly better FID than a comparable flow-matching model.
A blind hypothesis is that the self-consistency loss acts as a form of implicit regularization of the model, although we leave this investigation to future work.
결론은 모르겠다인데, loss중 self-consistency loss가 regularization처럼 작동해서 그렇지 않을까라고 추측
one-step도 여전히 scaling이 적용된다.
one-step이기 때문에 initial noise에서의 interpolation이 뭔가 더 실제로 적용되는 느낌의 사진
computation
All experiments are run on TPUv3 nodes, and methods are implemented in JAX. While runtimes vary per method, each training run typically takes 1-2 days to complete.
Share article
Subscribe to our newsletter
