Flow Matching

Flow matching là phương pháp huấn luyện model tạo sinh biến nhiễu ngẫu nhiên thành dữ liệu bằng cách đi theo trường vận tốc đã học. Một mạng nơ-ron học hướng và tốc độ di chuyển của một điểm tại mọi vị trí và thời điểm từ 0 đến 1; tạo mẫu nghĩa là bắt đầu từ nhiễu Gaussian rồi tích phân vận tốc đó bằng bộ giải phương trình vi phân thông thường (ODE) cho đến thời điểm 1. Yaron Lipman và các cộng sự tại Meta AI cùng Viện Khoa học Weizmann giới thiệu phương pháp này trong Flow Matching cho mô hình tạo sinh tại ICLR 2023; hai bài báo đồng thời khác cũng đi đến gần như cùng mục tiêu dưới tên gọi rectified flow và stochastic interpolants. Flow matching huấn luyện flow chuẩn hóa liên tục mà không mô phỏng ODE trong quá trình huấn luyện, bao hàm model diffusion như một trường hợp đặc biệt, và là mục tiêu huấn luyện đằng sau Stable Diffusion 3, FLUX.1, Movie Gen của Meta cùng nhiều model giọng nói và robot gần đây.

Cách flow matching hoạt động#

Quá trình huấn luyện không bao giờ tạo ra dữ liệu. Mỗi bước lấy một mẫu dữ liệu x1, một mẫu nhiễu x0 từ phân phối Gaussian chuẩn và một thời điểm t từ 0 đến 1, rồi đặt một điểm trên đường thẳng nối chúng: xt = (1 - t) * x0 + t * x1. Điểm trên đường thẳng đó di chuyển với vận tốc không đổi x1 - x0, vì vậy mạng v(xt, t) được huấn luyện bằng sai số bình phương trung bình thông thường để dự đoán x1 - x0. Mỗi cặp nhiễu-dữ liệu chỉ xác định một đường đi, nhưng nhiều đường đi có thể qua cùng một điểm và vận tốc của toàn bộ flow từ nhiễu đến dữ liệu tại điểm đó là giá trị trung bình của các đường đi này. Bài báo chứng minh rằng hồi quy trên các mục tiêu đơn giản theo từng cặp cho cùng gradient như hồi quy trên giá trị trung bình khó tính toán đó; vì vậy, huấn luyện chỉ là hàm mất mát hồi quy thông thường, không cần mô phỏng ODE.

Quá trình tạo tích phân dx/dt = v(x, t) từ t = 0 đến t = 1, thường bằng các bước của bộ giải Euler hoặc bộ giải bậc cao hơn. Đoạn mã PyTorch này huấn luyện một mạng vận tốc nhỏ trên hai cụm điểm 2D rồi lấy mẫu với 1, 2 và 8 bước Euler:

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")

Kết quả in ra 1-step sampling: 0% of samples land near a cluster, sau đó là 47% với 2 bước và 94% với 8 bước. Tại t = 0, mọi điểm nhiễu đều có cùng đích đến kỳ vọng là giá trị trung bình của dữ liệu, nên chỉ một bước duy nhất đưa mọi mẫu đến khoảng trống giữa hai cụm. Mỗi đường huấn luyện đều thẳng, nhưng các đường bắt đầu gần nhau lại kết thúc ở những cụm khác nhau, và flow học được lấy trung bình các đường đó vẫn có dạng cong. Vì vậy, bộ tạo chỉ cần vài bước bổ sung các kỹ thuật như reflow, tức huấn luyện lại trên các cặp nhiễu-mẫu do chính model tạo ra để làm thẳng flow, hoặc chưng cất thành model nhất quán.

Flow matching và các phương pháp liên quan#

• Model diffusion: Huấn luyện diffusion làm nhiễu dữ liệu theo lịch nhiễu và dạy mạng dự đoán nhiễu được thêm vào, sau đó lấy mẫu bằng cách đảo ngược quy trình này. Bài báo về flow matching cho thấy mục tiêu của phương pháp này chấp nhận mọi đường đi xác suất Gaussian, bao gồm cả những đường đi được model diffusion sử dụng, và nhận thấy đây là cách huấn luyện ổn định hơn. Đường đi thẳng là một lựa chọn khác, đơn giản hơn. Cả hai họ phương pháp đều dùng chung các kiến trúc như transformer diffusion (DiT), hướng dẫn không cần classifier và bộ lấy mẫu ODE.
• Rectified flow: Sử dụng cùng phép nội suy đường thẳng và loss vận tốc, vì vậy hai tên gọi phần lớn đồng nghĩa khi dùng đường đi tuyến tính. Bài báo về rectified flow bổ sung quy trình reflow để làm thẳng quỹ đạo hơn nữa.
• Flow chuẩn hóa: Flow chuẩn hóa rời rạc xếp chồng các lớp khả nghịch và huấn luyện bằng likelihood chính xác. Flow chuẩn hóa liên tục, một dạng ODE nơ-ron, xác định phép biến đổi bằng ODE nhưng được huấn luyện bằng hợp lý cực đại, phương pháp giải ODE ở mọi bước huấn luyện. Flow matching huấn luyện cùng loại model ODE bằng hồi quy, không cần giải ODE.
• Optical flow: Không liên quan dù tên gọi giống nhau. Optical flow là chuyển động biểu kiến của pixel giữa các frame video, còn flow matching di chuyển các mẫu trong không gian dữ liệu khả dĩ.

Ứng dụng của flow matching#

• Tạo ảnh: Stability AI huấn luyện Stable Diffusion 3 bằng transformer rectified flow và lấy mẫu timestep huấn luyện từ phân phối logit-normal ưu tiên các mức nhiễu trung gian; FLUX.1 của Black Forest Labs là transformer rectified flow có 12 tỷ tham số. Hình ảnh từ các model chuyển văn bản thành hình ảnh này được dùng làm dữ liệu tổng hợp, bổ sung vật thể hiếm, thời tiết và góc nhìn vào tập dữ liệu huấn luyện model phát hiện YOLO26.
• Tạo video: Movie Gen của Meta là transformer có 30 tỷ tham số, được huấn luyện bằng flow matching, có thể tạo video 1080p dài tối đa 16 giây với tốc độ 16 frame mỗi giây.
• Giọng nói: Voicebox của Meta huấn luyện model flow matching phi tự hồi quy trên hơn 50.000 giờ giọng nói; trong tác vụ chuyển văn bản thành giọng nói zero-shot, model đạt tỷ lệ lỗi từ 1.9% so với 5.9% của VALL-E, đồng thời chạy nhanh hơn tới 20 lần. F5-TTS sau đó kết hợp flow matching với transformer diffusion.
• Điều khiển robot: π0 của Physical Intelligence bổ sung một “chuyên gia hành động” flow matching 300 triệu tham số vào model thị giác-ngôn ngữ 3 tỷ tham số, tạo các nhóm gồm 50 hành động tương lai với 10 bước tích phân và điều khiển robot ở tần số lên đến 50 Hz.