Flow matching for generative modeling

ICLR2023 Notable top 25% (10/8/8/5), Stable diffusion 3

Motivation

(Left) the data distribution (checkerboard) to the noise distribution (standard gaussian) (Right) noise to data, with diffusion and flow-matching models. NFE stands for number of function estimation. Top row is the typical diffusion model, and middle and bottom row are the flow-matching models.

어떤 data-distribution 에서 simple-distribution (e.g. standard gaussian) 으로 변화하는 path (e.g. forward-diffusion process) 를 좀 더 잘 정의해서, 그것의 inverse (image generation via the diffusion model) 또한 더 잘 되도록 하고싶다 

• Simple diffusion process (adding a “simple“ gaussian noise) leads to rather confined space of sampling probability paths. (라고 표현하고 있는데, 여기서 설계한 flow도 마찬가지 아닌가? 하는 생각은 든다...)
• Meanwhile… the "Normalizing flow" is capable of modeling arbitrary probability path. 
• What if the given neural network can regress such “flow“ rather than simple “noise” (Diffusion) ??
• However, such “flow“ is somehow intractable → How can we define the tractable flow?

• 위 문제들을 해결하면, 결론적으로 diffusion 보다 더 좋은 sampling path를 얻는다. 샘플링 과정에서 inference 횟수가 더 적어지고, 학습도 더 빨라진다.

---

Flow-matching in a nutshell: flow-matching vs diffusion

노이즈를 더해가는 방식 비교: (Top) Diffusion forward process, (Middle) Diffusion Flow-matching, (Bottom) Optimal transport flow-matching. Diffusion의 경우에는 매 step 마다 더해지는 노이즈가 랜덤하여 지지직 하는 형태로 노이즈가 더해진다. 반면에, flow-matching에서는 image에서 noise로의 변환이 매끄럽게 이어진다. 어떤 Noise로 이동할지 정해놓고, raw-image에서 해당 noise로 이동하는것이다.

 

• Diffusion에서는, forward-diffusion에서 timestep 하나하나마다 stochastic한 어떤 값이 더해진다. 즉, 매 timestep마다 randomness가 있다 (매 step마다 더해지는 noise가 다르다). 반면에, Flow-based 모델에서는 개념적으로는 매 timestep마다 deterministic하게 이동한다. 즉, 어떤 데이터에 대해 단 하나의 end-point (probably sampled from the standard gaussian) 가 있다 (매 step마다 더해지는 noise가 같다).
  • 따라서, Diffusion (SDE modeling) 과 달리, Flow-matching (ODE modeling; no stochastic term) 하에서 이미지를 생성하고 싶다고 할 때, 처음 pure-gaussian noise $x_0 \sim \mathcal{N}(0, \mathbf{I}^d)$ 를 결정하고나면, 해당 $x_0$ 로 부터 이미지로 가는 path는 결정적이다.
• 위와 같은 flow-matching의 deterministic한 매커니즘을 기술하는 formulation으로 normalizing flow (비슷한 것) 를 사용한다. Normalizing flow는, 확률분포를 계속 “더 단순한 분포” 로 변경시켜나가는 좀 더 일반화된 과정이다 (forward-diffusion과 유사). 이후 그것을 거꾸로 변환시켜 나간다 (backward-diffusion과 유사).
  • 개념적으로 diffusion 대신 flow를 쓰겠다는것이지… Ideal하게 이것을 직접적으로 구현할 수는 없다. “확률분포의 변환” 을 무엇으로 정의할 것이며, 이것의 “역변환” 은 어떻게 구하는가?… 따라서 이것을 tractable 하게 구할수 있는 셋팅을 논문에서 고안한다.

---

Background

flow-matching의 formulation에 사용되는 normalizing flow와, 확률변수의 변환에 대해서 간단히 설명한다.

 

(Normalizing) flow

• 위의 확률분포의 변환 ($f_i$) 을 어떤 분포 $p_i(z_i)$에 대해 지속적으로 해주는 어떤 모델을 설계한다.
  • 여기서, $f$를 디자인할때, $f$는 역함수가 존재하는 어떤 연산이어야 한다 (flow model의 정의상 “거꾸로 되돌릴 수 있어야“ 하기 때문)
• $p_k(z_k)$ 에 대한 negative-log-likelihood를 최소화한다 (다른 많은 생성 모델의 loss function 처럼…)

 

확률변수에 대한 변환과 change of variables

앞서 언급했듯, Flow-matching은 diffusion과 달리 확률변수를 지속적으로, 결정적 (deterministic) 으로 변환시켜 나가는 과정으로 기술된다. 여기서는 그러한 변환이 어떤식으로 정의되는지를 다룬다. 이 부분을 생략하더라도 flow-matching을 이해하는데에 큰 문제는 없다.

 

Univariate Random Variable (R.V) 에 대한 확률변수의 변환: $X$ to $Y$ 로의 변환을 한다고 생각해보자. 여기서 변환은 어떤 함수 $g$ 에 의해 일어난다: $Y=g(X)$. 또한, 함수 $g$가 invertable하고 bijective하다면, 해당 변환의 역변환도 정의할 수 있다.

$$f_X(x) = f_Y(y) |\frac{dx}{dy}| = f_X(g^{-1}(Y)) |\frac{dx}{dy}|$$

Multi-variate RV 에 대한 변환은 위 $dx/dy$가 Determinant-of-Jacobian 으로 변한 형태이다.

여담: 티스토리의 latex support는 정말 끔찍하구나??

(Q) Determinant of Jacobian의미 (Informal): 확률분포를 X에서 Y로 변경했으면, “확률분포의 정의 = 적분 값이 1” 을 만족하도록 서로 스케일을 맞추어줘라 → 해당 스케일을 맞추는것이 위의 determinant이다. ($X=2Y$ 같은 변환이 있을때, $dx/dy=1/2$ 를 곱해주지 않으면 Y를 적분해버리면 1/2 가 되어버릴것이다. 이는 확률은 적분하면 1이 되어야 한다는 공리에 어긋난다.)

이러한 확률변수의 변환을 여러번 하여, 단순한 확률분포를 우리가 원하는 복잡한 분포로 바꾸거나, 거꾸로 복잡한 분포를 단순하게 만들거나 할 수 있다. 실제로, 확률변수의 변환을 직접적으로 활용하는 또다른 생성 모델인 normalizing flow 에서는, 확률변수에 대한 변환을 여러개를 구성하고, 그것들의 역변환을 구성한다.

---

Flow matching

Something more about Flow: constructing the transformation with the vector field

Background에서 설명했던 normalizing flow에서, discrete한 형태로 표현되었던 확률분포의 변환들은 본 논문에서는 continuous하게 표현된다 (매 timestep $t \in [0-1]$ 마다 어떤 변환이 있다). 또한, 본 논문에서는 이러한 확률분포의 변환을 vector-field 라는 개념으로 만들어낸다. 우선, 위의 normalizing flow에서의 확률변수의 변환 자체를 k개의 고정된 어떤 함수가 아니고 0-1 사이의 연속적인 시간 t 에서 일어난다고 생각해보자. 이 때,

• probability density path, $p_t$ : t 시간에서의 X의 PDF.
• flow, $\phi_t$ : t 시간에서의 “확률변수 X의 변환”
• time-dependent vector field, $v_t$ : t 시간에서의 “flow“ 를 변화시키는 값 → flow의 시간에 대한 미분 (즉 vector field를 알면, 시간에 따라 확률분포를 어떻게 바꾸어 나갈건지 알 수 있고, 따라서 확률분포도 알 수 있다.)

우리의 probability distribution $p_0$ 와 확률변수 X가 있다고 생각해보자 (보통은 sampling을 할 수 있는 어떤 단순한 분포 일 것이다 e.g. standard normal). 그리고 해당 분포를 구성하는 "particle" (논문의 표현을 빌리자면) 이 있다고 가정해보자. 이 가상의 data-point, particle 들이 무수히 모혀서 $p_0$ 를 구성한다고 생각해보는것. 그럼 이 particle들은 0 -> 1 로의 timepoint마다 vector field $v_t$ 에 의해서 아래 그림처럼 이동한다.

( Top ) $P_t$, ( Middle ) $v_t$, ( Bottom ) $\phi_t(x)$ : transformed datapoints and their initial positions. ( NOTE ) $p_t$ as just a density with all possible data point $x$, after transformation via $phi_t$ and $v_t$.

달리 말해서, time t 에서의 확률변수 $X$의 변환 $\phi_t$는 vector field $v_t$를 꾸준히 더해감으로 (적분해감으로) 정의된다. 논문에서의 formulation은 아래처럼 정의된다. 

$$ \frac{d}{dt} \phi_t (x) = v_t (\phi_t(x))$$

$$ \phi_0 (x) = x $$

그럼, 우리가 원하던 변환된 확률변수에 대한 분포 $p_t$ 는 아래와 같이 정의된다

$$ p_t = [\phi_t]_* p_0 $$

where $[]_*$ is the push-forward operator, defined as

$$ [\phi_t]_* p_0 (x) = p_0 (\phi^{-1} (x)) \text{det} [\frac{\partial \phi_t^{-1}}{\partial x}(x)] $$

아까 봤던 normalizing-flow의 likelihood 랑 유사하다. 다만, flow-matching이라는 이름과 diffusion의 alternative라는 점에서 추측할 수 있듯이, 이러한 likelihood가 우리의 training-objective가 되지는 않는다. 실제로는 vector-field $v_t$ 를 매 $t$에 대해서 맞추는것이 training objective가 될 것이다.

---

Now… lets match such flow

이제 대강 flow가 뭔지, 그리고 거기에 대응하는 vector field가 뭔지를 알겠다. 헌데 우리는 실제로 접근할 수 있는 “분포”는 t=1 에서의 “진짜 이미지 데이터” 와, t=0에서의 아주 단순한 특정 확률분포 (e.g. standard normal distribution) 뿐이다. 중간에 변환 (flow) 이 어떻게 생겼는지 모르고, 당연히 그에 상응하는 vector field도 뭔지 모른다.

 

우선 이것들을 무시하고, 개념적으로 시간 $t$ 에서, 특정 데이타 $x$에 대한 “vector field“ $u_t(x)$ 를 안다고 가정했을 때, 해당 $u_t(x)$에 대한 regression task (diffusion 에서의 noise 맞추기 처럼…) 를 정의해보자. 이것이 flow-matching model 에서의 (ideal) training objective가 될 것이다.

 

마치 diffusion의 objective 와 유사하고, neural-network output이 현재 "particle" $x$ 에 대한 true vector field $u_t(x)$ 를 모사해나가는 objective이다.

 

다만, 아까 언급했듯 우리는 true vector field $u_t$ 에 대해서 알 수 없다… 그러면 어떻게 해야할까?… 논문에서는 conditional probability path 라는 것을 정의해서 무언가 tractable한 $u_t$를 정의하고자 한다.

Approximate $u_t$ with per-sample: conditional probability paths

이제 우리의 목적은 ideal하고 intractable 했던 flow-matching loss를 사용할 수 없기 때문에, 적절한 정의를 통해 tractable 한 flow를 찾아내는 것이다. 저자는 이에 대한 해답으로 conditional probability path라는 것을 제안한다.

 

Conditional probability path, $p_t(x|x_k)$ : 특정 데이터 샘플 $x_k$ 에 대해서, 아래와 같은 성질을 만족한다

1. $p_0(x|x_k) = p(x)$ : 시간 0에서는, 모든 데이터들의 확률 분포와 똑같다.

2. $p_1(x|x_k)$ to be a distribution concentrated around $x = x_k$, $e.g., p_1(x|x_k) = N(x_k, I)$ 즉, 시간 1에서는 “에지간하면 $x_k$, 스스로가 나온다”.

그리고 $p_t(x)$ 를 위와 같은 conditional probability path 들의 mixture라고 생각하자. 아래의 그림 (from "Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport", TMLR 2024) 처럼, timestep에 따라 밑에서 위로 분포가 변해갈것이다.

이렇게 설계하면, flow-matching loss를 실제로 구현상에서 활용할 수 있으며, 이에 더해 아래의 두가지 훌륭한 이론적 특징을 보장한다:

 

(1) 이와 같이 정의할 경우에, vector field로 부터 probability distribution 을 얻어내는 특성을 여전히 만족한다.

정확히는, all possible data point at the timepoint 1 (the last timepoint), $x_1$, 에 대해 marginalize한 vector field가 우리가 알고싶은 true-vector field (not, conditional, vector vield) 와 같아진다.

 

 

(2) Conditional flow 를 맞추는 loss function을 디자인 할 경우, 원래 맞추려고 했지만 intractable 하여서 맞출 수 없었던 “true flow“ 에 대한 loss랑 동일한 gradient 를 가진다 (at least for the expectation)

 

 

이제 true-flow에 대한 적절한 근사치 conditional flow 가 위 두가지 훌륭한 속성을 만족하는 것을 알았다. 이러한 conditional flow를 실제로 구체화 해보자.

---

Constructing conditional probability paths

이전까지 Intractable 했던 flow를, conditional flow라는 것을 정의하여 실제 flow를 모사했다. 그럼 이제conditional flow의 정의 아래에서, probability path $p$를 실제로 정의하여 conditional path (우리의 “label” 이 될 대상) 을 구체화시켜보자.

(1) Probability path p 정의

언제나 그렇듯, 가장 단순하고 tractable 한 것은 gaussian distribution이다. 

$$p_t (x|x_1) = \mathcal{N}(x| \mu_t(x_1), \sigma_t(x_1)^2I)$$

여기서, $\mu_t$ 와 $\sigma_t$ 를 아래와 같이 제한해보자

• $\mu_0(x) = 0,\qquad  \mu_1(x) = x$
• $\sigma_0(x) = 1, \qquad \sigma_1(x) = \text{TINY_NUMBER (sufficiently small)}$

(2) Probability flow $\phi$ 정의

$$ \phi_t (x) = \sigma_t(x_1)x + \mu_t(x_1)$$

x는 어떤 Normal-distribution을 따르는 RV 이다. 즉, normal distribution의 affine-transformation이 우리가 정의한 probability flow이다.

(3) $p$와 $\phi$를 통해, 우리의 label, conditional vector field 구하기

위 처럼 flow를 normal distribution의 affine transformation이라고 생각하면, vector field는 아래 theorem 3 처럼 계산이된다.

 

이제 아래처럼, 위에 정의된 vector field를 “맞추는” Loss를 줄이면 된다.

막간, 그리고 Recap: Diffusion과의 차이점

앞서 말했던것 처럼, Diffusion의 경우는 timestep마다 어디로 움직일지 non-deterministic하다 (SDE-modeling). 반면에, 여기서 정의한 flow-matching은 한번 particle ($t_0$ 에서의 sample) 이 결정되고 나면, 어디로 움직여야할지 deterministic하다. (사실 구현적으론 아주 큰 차이가 있지는 않다.)

---

Example of conditional vector fields: Gaussian conditional vector field

앞서서 flow-matching을 정의하기 위해서는 세가지가 정의되어야 함을 이야기했다: probability path, probability flow 그리고 vector field (to generate the probability flow). 여기서부터는 위 세가지를 하나씩 정의해 나가면서 실제 구현상에서 쓰이는 flow-matching의 형태를 잡아본다. 실제 구현상의 코드를 먼저 보고싶다면, 본문 맨 아래의 Implementation을 확인하면 된다.

Preliminaries: three components we have defined

저자는 우선적으로 probability path를 다음과 같이 어떤 time-dependent mean-function, $\mu_t$ 과 std-function, $\sigma_t$ 를 변수로 갖는 gaussian으로 정의한다. Diffusion 에서와 유사하게, input $x_1$과 무관하게 처음에는 standard gaussian이 된다: $\mu_0(x_1) = 0$, $\sigma_0 (x_1)=1$. 좀 헷갈리는점은 Diffusion에서는 $t=1$ 일때가 pure-noise였던것 같은데, 반대로 써져있는 느낌이다 (기억이 맞다면...)

 

Probability path

$$p_t(x|x_1) = \mathcal{N} (x|\mu_t(x_1) , \sigma_t(x_1)^2 I)$$

 

위 같이 정의된 probability path를 생성할수 있는 probability flow는 무수히 많다. 저자는 여기서 가장 단순한 형태의 flow를 생각하기로 한다. (무언가 복잡한 형태더라도, 대부분이 쓸모없는 연산 으로 이루어졌을거라 저자는 언급한다. 예를들어, 최종 PDF가 rotation-invariant하여 $x$에 대한 rotation을 하더라도 $p_t(f(x)|x_1)$--ignore the determinant term--이 같다고 하면, 해당 연산을 포함하는 flow 는 불필요한 회전변환을 포함하고 있는것이다.)

 

Probability flow

$$\psi_t(x) = \sigma_t(x_1)x + \mu_t (x_1)$$

이제 vector field $u_t$는 flow의 $t$에 대한 미분으로 아래처럼 얻을 수 있다.

 

Vector field

$$\frac{d}{dt} \psi_t(x) = u_t(\psi_t(x)|x_1)$$

위에서 정의한 flow를 통해서 $u_t$를 계산하면 아래 Theorem 3과 같은 결과가 나온다

 

Note that $f'$ denotes the derivative w.r.t. time $f'=\tfrac{d}{dt}f$.

Gaussian conditional flow matching loss

이렇게 구해진 것들을 원래의 conditional flow-matching objective $L_{CFM}(\theta) = \mathbb{E}_{t, q(x_1), p_t(x|x_1)} \| v_t(x) - u_t(x|x_1)\|^2$ 에 넣으면 아래와 같아진다

 

Special instances of gaussian conditional probability paths

위 formulation에서 한가지 정의되지 않은것은 mean and std function $\mu_t(x_1)$, $\sigma_t(x_1)$ 이다. 이는 어떤 differentiable한 function 으로도 정의할 수 있다. 논문에서는 두가지 예시를 제안하는데, Diffusion의 probability path를 그대로 가져온 버전과 Optimal transport probability path이다. 

Diffusion conditional VF

논문에서는 Variance Exploding (VE) case와 Variance Preserving (VP) case를 모두 다루지만, 여기서는 VE case만 다루도록 하겠다 (VP case가 좀 더 보이기에 복잡하게 전개된다). VE는 아래처럼, mean and std function을 $\mu_t(x_1)=x_1$, $\sigma_t(x_1) = \sigma_{1-t}$ 로 둔다 (where $\sigma_t$ is an increasing function and $\sigma_0 = 0$)

 

Probability path, from the diffusion (Variance exploding diffusion)

$$p_t(x) =\mathcal{N}(x|x_1, \sigma^2_{1-t}I)$$

 

이후 위에서 언급된 Theorem 3에 mean and std function을 대입하면 아래와 같은 결과를 얻는다.

 

Probability flow (mu and sigma) and vector field.

$$u_t(x|x_1) = -\frac{\sigma'_{1-t}}{\sigma_{1-t}} {x - x_1}$$

Optimal Transport conditional VF

논문에서는 좀 더 단순한 형태의 mean and std function을 제안하는데, time에 따라 linear 하게 바뀌는 형태이다.

$$\mu_t(x)=tx_1, \text{ and } \sigma_t(x) = 1 - (1 - \sigma_{min})t$$

Theorem 3를 통해 vector field를 계산하면 아래와 같다

$$ u_t(x|x_1) = \frac{x_1 - (1-\sigma_{min})x}{1-(1-\sigma_{min})t} $$

그리고 위 vector field는 $t \in [0, 1]$ 에 대해 실제로 계산할 수 있는데 (실제로 t에 대해 해당 구간에 대해 적분), 이는 다음과 같다:

$$\psi_t(x) = (1-(1-\sigma_{min}) t)x + tx_1$$

이에 따라 위에서 구했던 conditional flow matching loss에 대입하면 아래와 같은 결과를 얻는다

Why “Optimal transport?“

• 위 conditional flow는 두 gaussian distributions 사이의 optimal-transport mapping 이다.
• 직관적으로, 해당 mapping은 “particle“ 들을 직선상의 trajectory에서 등속으로 옮기는 맵핑이다.
• (주의) 아까 정의했던 “Conditional PDF“ 상에서 optimal transport solution 이라는것이지, original PDF (marginalized conditional PDF) 에 대해 optimal transport 라는 것은 아니다
•  

---

Implementation example

실제 구현이 될 때는, noisy-input $\psi(x_1)$ 를 만들어내는 코드와 conditional flow-matching loss만 구현하면 된다. 여기서는 Optimal-transport case를 생각한다. 

torch로 구현된 특정 neural network 모델이 있다고 하자. 해당 모델의 output을 argument로 받는 flow-matching loss 를 계산하는 Mixin class 코드는 아래처럼 구현할 수 있다. 사실 구현상으로는 Diffusion의 noise matching task와 이 flow matching task가 거의 유사하다.

"""
Dongpin Oh, dhehdqls@gmail.com
"""

from typing import Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from diffusers.utils.torch_utils import randn_tensor


class FlowMixin:
    """Mixin class for flow-based models."""

    MIN_STD: float = 1e-4  # minimum size of std for the flow matching
    CLAMP_CONTINUOUS_TIME: float = 1e-7

    def get_noisy_input(self,
                        input: torch.Tensor,
                        device: str = 'cuda') -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Generate noisy input, based on the optimal-transport flow.

        Args:
            input (torch.Tensor): Input tensor.
            device (str, optional): Device to use for computation. Defaults to 'cuda'.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Tuple containing noise, noisy input, and timestep.
                - note that, the timestep t is continuous timestep ranged [0, 1] which cannot directly be used
                 for the timestep embedding (need to be discritized by self.discritize_timestep() )
        """
        b = input.shape[0]
        noise = randn_tensor(input.shape, device=device)
        t = torch.rand((b, ), device=device)
        t = t.clamp(0 + self.CLAMP_CONTINUOUS_TIME, 1 - self.CLAMP_CONTINUOUS_TIME)  # for numerical stability
        t = t.unsqueeze(1).unsqueeze(1).unsqueeze(1)
        noisy_input = (1 - t) * input + (self.MIN_STD + (1 - self.MIN_STD) * t) * noise
        return noise, noisy_input, t.squeeze()

    def flow_loss(self, input: torch.Tensor, noise: torch.Tensor, t: torch.Tensor,
                            preds: torch.Tensor) -> torch.Tensor:
		"""
        Args:
            input (torch.Tensor):
            noise (torch.Tensor): 
            t (torch.Tensor): Timestep tensor , [0, 1].
            preds (torch.Tensor): Model output

        Returns:
            torch.Tensor: Rectified flow loss.
        """

        # Matching dimension for broadcasting
        t = t.reshape(t.shape[0], *[1 for _ in range(len(input.shape) - len(t.shape))])

        target_flow = input - (1 - self.MIN_STD) * noise
        loss = F.mse_loss(preds.float(), target_flow.float(), reduction='none')
        loss = loss.mean()
        return loss

    def discritize_timestep(self, t: Union[torch.Tensor, float], n_timesteps: int = 1000) -> torch.Tensor:
        """
        Discretize the continuous timestep.

        Args:
            t (Union[torch.Tensor, float]): Continuous timestep.
            n_timesteps (int, optional): Number of discrete timesteps. Defaults to 1000.

        Returns:
            torch.Tensor: Discretized timestep tensor.
        """
        return (t * n_timesteps).round().long()

 

flow-matching으로 학습시킬 모델은 다음과 같이 구현할 수 있다.

class Model(torch.nn.Module, FlowMixin):
	...
    def forward(self, images: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the  model.

        Args:
            images (torch.Tensor): Input images tensor, [0-1] ranged.

        Returns:
            torch.Tensor: flow matching loss.
        """
        images = self.normalize_img(images)

        # 1. Get noisy input via the rectified flow
        noise, noisy_image, t = self.get_noisy_input(latents)
        timesteps = self.discritize_timestep(t, NUM_TIMESTEPS)

        # 2. Forward pass through the dit
        preds = self.model(noisy_image, timesteps)

        # 3. Rectified flow matching loss
        loss = self.flow_loss(noisy_image, noise, t, preds)

        return loss