Flow Matching

플로 매칭은 학습된 속도장을 따라 무작위 노이즈를 데이터로 변환하는 생성 모델 학습 방법입니다. 뉴럴 네트워크는 시간 0과 1 사이의 임의의 시점과 임의의 지점에서 해당 지점이 이동해야 할 방향과 속도를 학습합니다. 샘플을 생성하려면 가우시안 노이즈에서 시작해 일반 미분 방정식(ODE) 솔버로 시간 1까지 속도를 적분합니다. Meta AI와 Weizmann Institute of Science의 Yaron Lipman과 동료들은 ICLR 2023에서 Flow Matching for Generative Modeling을 발표하며 이 방법을 소개했습니다. 같은 시기에 발표된 두 논문은 정류 플로와 확률적 보간이라는 이름으로 거의 동일한 목적 함수에 도달했습니다. 플로 매칭은 학습 중 ODE를 시뮬레이션하지 않고 연속 정규화 플로를 학습하며, 디퓨전 모델을 특수한 경우로 포함합니다. 또한 Stable Diffusion 3, FLUX.1, Meta의 Movie Gen, 최근의 여러 음성 및 로보틱스 모델을 학습하는 목적 함수입니다.

플로 매칭 작동 방식#

학습 과정에서는 어떤 것도 생성하지 않습니다. 각 단계에서 데이터 샘플 x1, 표준 가우시안 분포에서 뽑은 노이즈 샘플 x0, 0과 1 사이의 시점 t를 추출한 뒤 둘 사이의 직선 위에 점을 배치합니다: xt = (1 - t) * x0 + t * x1. 직선 위의 점은 일정한 속도 x1 - x0로 이동하므로, 네트워크 v(xt, t)는 단순 평균 제곱 오차를 사용해 x1 - x0을 예측하도록 학습됩니다. 각 노이즈-데이터 쌍은 하나의 경로만 정의하지만, 임의의 한 점을 통과하는 경로는 여러 개이며 해당 지점에서 전체 노이즈-데이터 플로의 속도는 이 경로들의 평균입니다. 논문은 간단한 쌍별 목표를 회귀하는 것이 계산하기 어려운 평균값을 회귀하는 것과 동일한 그래디언트를 제공함을 증명합니다. 따라서 학습은 ODE 시뮬레이션 없이 단순 회귀 손실 함수를 사용합니다.

생성 시에는 t = 0에서 t = 1까지 dx/dt = v(x, t)을 적분하며, 일반적으로 Euler 또는 고차 솔버 단계를 사용합니다. 이 PyTorch 코드는 2D 점 군집 두 개로 소규모 속도 네트워크를 학습한 다음 Euler 단계를 1, 2, 8회 적용해 샘플링합니다:

import torch
from torch import nn

torch.manual_seed(0)
# Velocity network v(x, t) for 2D points; the data are two clusters at (2, 2) and (-2, -2)
net = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
centers = torch.tensor([[2.0, 2.0], [-2.0, -2.0]])

for step in range(2000):
    x1 = centers[torch.randint(0, 2, (256,))] + 0.2 * torch.randn(256, 2)  # data samples
    x0 = torch.randn(256, 2)  # noise samples
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # point on the straight path from noise to data
    loss = ((net(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the path's velocity
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate: start from noise and follow the learned velocity field with Euler steps
for steps in (1, 2, 8):
    x = torch.randn(1000, 2)
    with torch.no_grad():
        for i in range(steps):
            x = x + net(torch.cat([x, torch.full((1000, 1), i / steps)], dim=1)) / steps
    near = (torch.cdist(x, centers).min(dim=1).values < 0.6).float().mean()
    print(f"{steps}-step sampling: {near:.0%} of samples land near a cluster")

출력 결과는 1-step sampling: 0% of samples land near a cluster이며, 2단계에서는 47%, 8단계에서는 94%입니다. t = 0에서는 모든 노이즈 점의 기대 목적지가 데이터의 평균값으로 동일하므로, 한 단계만 사용하면 모든 샘플이 두 군집 사이의 빈 공간으로 이동합니다. 각 학습 경로는 직선이지만 서로 가까이에서 시작하는 경로가 서로 다른 군집에서 끝날 수 있으며, 그 평균을 내는 학습된 플로는 여전히 휘어집니다. 따라서 단계를 적게 사용하는 생성기는 모델 자체의 노이즈-샘플 쌍으로 재학습해 플로를 곧게 만드는 리플로 또는 증류를 통해 일관성 모델로 변환하는 등의 기법을 추가합니다.

플로 매칭과 관련 방법 비교#

• 디퓨전 모델: 디퓨전 학습은 노이즈 스케줄에 따라 데이터를 오염시키고 네트워크가 추가된 노이즈를 예측하도록 학습한 다음, 이 과정을 역으로 수행해 샘플을 생성합니다. 플로 매칭 논문은 목적 함수가 디퓨전 모델에서 사용하는 경로를 포함해 모든 가우시안 확률 경로에 적용되며, 디퓨전 모델을 더 안정적으로 학습하는 방법이라는 점을 보였습니다. 직선 경로는 별도로 선택할 수 있는 더 단순한 방법입니다. 두 방법 모두 디퓨전 트랜스포머(DiT)와 같은 아키텍처, 분류기 없는 가이던스, ODE 샘플러를 공유합니다.
• 정류 플로: 동일한 직선 보간과 속도 손실을 사용하므로 선형 경로에서는 두 용어가 대체로 동의어입니다. 정류 플로 논문은 궤적을 더 곧게 만들기 위해 리플로 절차를 추가합니다.
• 정규화 플로: 이산 정규화 플로는 가역 계층을 쌓고 정확한 우도로 학습합니다. 뉴럴 ODE의 한 형태인 연속 정규화 플로는 ODE로 변환을 정의하지만 최대 우도로 학습하므로 각 학습 단계에서 ODE를 풀어야 합니다. 플로 매칭은 같은 종류의 ODE 모델을 회귀로 학습하며 ODE를 풀 필요가 없습니다.
• 옵티컬 플로: 이름만 같을 뿐 관련이 없습니다. 옵티컬 플로는 비디오 프레임 사이에서 픽셀이 겉으로 움직이는 현상이며, 플로 매칭은 샘플을 가능한 데이터 공간을 따라 이동시킵니다.

플로 매칭의 활용 분야#

• 이미지 생성: Stability AI는 정류 플로 트랜스포머로 Stable Diffusion 3를 학습하고 중간 노이즈 수준에 가중치를 두는 로짓 정규 분포에서 학습 시점을 샘플링했습니다. Black Forest Labs의 FLUX.1은 매개변수가 120억 개인 정류 플로 트랜스포머입니다. 이러한 텍스트-이미지 모델로 만든 이미지는 합성 데이터로 활용되어 YOLO26 탐지 모델 학습용 데이터셋에 희귀 객체, 날씨, 시점을 추가합니다.
• 비디오 생성: Meta의 Movie Gen은 플로 매칭으로 학습한 매개변수 300억 개 규모의 트랜스포머로, 초당 16프레임의 1080p 비디오를 최대 16초까지 생성합니다.
• 음성: Meta의 Voicebox는 50,000시간이 넘는 음성 데이터로 비자기회귀 플로 매칭 모델을 학습했습니다. 제로샷 텍스트 음성 변환에서 단어 오류율은 1.9%로, VALL-E의 5.9%보다 낮았으며 실행 속도는 최대 20배 빨랐습니다. 이후 F5-TTS는 플로 매칭과 디퓨전 트랜스포머를 결합했습니다.
• 로봇 제어: Physical Intelligence의 π0는 매개변수 30억 개 규모의 비전-언어 모델에 매개변수 3억 개 규모의 플로 매칭 "액션 전문가"를 추가해, 미래 동작 50개의 청크를 10회 적분 단계로 생성하고 최대 50Hz로 로봇을 제어합니다.