Flow Matching, Conditional Flow Matching và Code

Hãy tưởng tượng bạn muốn dạy máy tính vẽ tranh giống tranh thật (hoặc tạo ảnh, tạo giọng nói, tạo video, tạo chuyển động robot…).

Có mấy cách phổ biến mọi người hay dùng:

• GAN → hai mạng đánh nhau
• Diffusion → từ nhiễu trắng → từ từ “làm sạch” thành ảnh đẹp
• Flow → học cách “kéo dãn, uốn éo” không gian để biến đám nhiễu thành dữ liệu thật

Flow Matching là cách học Flow mới, hiện đang cực kỳ hot vì:

• Train nhanh hơn diffusion rất nhiều
• Không cần simulation dài dòng
• Tạo mẫu deterministic (chạy cùng seed cho cùng kết quả)
• Dễ scale lên ảnh lớn, video, 3D, âm thanh, protein…

Ý tưởng cốt lõi

Bạn có hai đám điểm:

• Đám A (t=0): đám mây tròn Gaussian nhiễu trắng (dễ sample)
• Đám B (t=1): đám điểm dữ liệu thật (ảnh mèo, giọng nói, pose robot…)

Ta muốn tạo một “dòng chảy” liên tục từ A → B trong khoảng thời gian t từ 0 → 1.

Mỗi điểm trong không gian sẽ di chuyển theo một vector vận tốc (velocity) → giống như bạn vẽ đường đi cho từng hạt bụi bay từ đám nhiễu thành hình con mèo.

Mạng neural chỉ việc học:
Tại vị trí x, thời điểm t → nên di chuyển theo hướng nào, nhanh bao nhiêu?
→ uθ(x, t) = vector vận tốc dự đoán

Vấn đề: nếu học thẳng vector cho toàn bộ đám đông → rất khó vì phải biết mật độ xác suất ở mọi t.

Vận tốc (velocity field) chính là “linh hồn” của Flow Matching

Hãy nghĩ Flow Matching như một bộ phim hoạt hình mà ta dạy máy tính cách vẽ từng khung hình.

• Ở t = 0: màn hình toàn nhiễu trắng (Gaussian)
• Ở t = 1: màn hình là ảnh thật (mèo, người, xe…)

Nhưng thay vì vẽ từng khung một (như diffusion), ta chỉ cần dạy hướng di chuyển + tốc độ cho từng điểm ảnh ở mọi thời điểm t.

→ Vận tốc chính là vector chỉ hướng + độ lớn mà mỗi điểm phải di chuyển tại thời điểm t.

Mạng neural học một hàm:

vθ(x, t) hoặc uθ(x, t) → vector (thường là 512, 1024, 4096 chiều tùy model)

Ý nghĩa vật lý rất trực quan:

• vθ(x, t) giống như vận tốc gió tại vị trí x vào lúc t
• Nếu bạn thả một hạt bụi vào vị trí x lúc t, hạt đó sẽ bay theo hướng và tốc độ mà vθ(x, t) chỉ định
• Sau thời gian Δt nhỏ, vị trí mới ≈ x + vθ(x, t) ⋅ Δt

Khi ta tích phân (integrate) vận tốc này từ t=0 đến t=1 → ta được đường đi của hạt bụi → từ nhiễu → thành dữ liệu thật.

Tại sao phải học vận tốc thay vì học xác suất trực tiếp?

Học trực tiếp mật độ xác suất p(x) rất khó vì:

• Phải chuẩn hóa (integral = 1) → tính toán nặng
• Không biết cách sample dễ dàng

Còn học vận tốc thì:

• Chỉ là regression đơn giản (dự đoán vector)
• Khi đã có vận tốc → tự động sinh được phân phối bằng cách giải ODE (chạy bộ tích phân)
• Có thể sample nhanh và xác định (deterministic) nếu dùng solver tốt

Conditional Flow Matching – cách học vận tốc “thông minh” nhất

Vấn đề: nếu ta cố học vận tốc chung cho cả đám (marginal velocity field u(x,t)) thì rất khó, vì phải biết mật độ ở mọi t.

Giải pháp CFM: chỉ dạy từng hạt một, theo kiểu “giả sử”:

1. Lấy một mẫu thật x₁ (ảnh mèo thật)
2. Lấy một nhiễu ngẫu nhiên x₀ ~ N(0,I)
3. Quyết định rằng hạt này sẽ đi thẳng từ x₀ → x₁ theo đường tuyến tính (đường phổ biến nhất):
   x(t) = (1−t) ⋅ x₀ + t ⋅ x₁
   → vận tốc thật của hạt này là hằng số: v thật(t) = dx(t)/dt = x₁ − x₀ (không phụ thuộc t!)
4. Bắt mạng neural đoán gần bằng vận tốc thật này:
   Loss = || vθ( x(t), t ) − (x₁ − x₀) ||²
   → Đây là conditional velocity v(x,t | x₁) mà ta đang giả định

Khi lặp hàng triệu lần với rất nhiều cặp (x₀, x₁) khác nhau → mạng học được vận tốc trung bình có điều kiện → khi marginal hóa (trung bình ngầm) lại, nó trở thành vận tốc đúng cho toàn bộ phân phối.

Kết quả:
Mặc dù ta chỉ dạy từng đường đi riêng lẻ → nhưng khi gom lại, nó tạo thành dòng chảy đúng cho toàn bộ phân phối dữ liệu.

Khi sample thì vận tốc được dùng như thế nào?

Sau khi train xong, ta không dùng vận tốc có điều kiện nữa, mà dùng vận tốc đã marginal hóa (tức là vθ(x,t) mà mạng học được).

Quy trình sample:

1. Sample x₀ ~ N(0,I)
2. Chạy ODE solver (Euler, Heun, Dopri5, …): x(t + Δt) ≈ x(t) + vθ(x(t), t) ⋅ Δt
3. Sau ~10–50 bước → x(1) chính là mẫu mới sinh

→ Toàn bộ quá trình chỉ là theo dõi vận tốc mà mạng dự đoán!

Tóm tắt một câu siêu ngắn về vận tốc

Flow Matching dạy máy tính dự đoán “gió thổi hướng nào, mạnh cỡ nào” tại mọi vị trí và mọi thời điểm → để đẩy đám hạt nhiễu bay thành hình ảnh/video/protein thật.

Conditional Flow Matching làm việc này dễ hơn bằng cách dạy từng hạt một đường đi riêng (thường là đường thẳng) → mạng tự học được “gió chung” hợp lý cho cả đám.

So sánh

Flow Matching = “Thay vì dạy cả đám đông cùng lúc, ta dạy từng người một đường đi thẳng từ nhiễu đến dữ liệu thật → mạng tự học được cách dẫn dắt cả đám một cách mượt mà.” Bạn nào đang làm diffusion mà thấy chậm quá thì thử chuyển sang Flow Matching xem, train nhanh gấp mấy lần đấy!

Bài thuyết trình

Đây là một video tutorial, hơi toán, nhưng chúng ta không cần hiểu hết phần toán để train vì đã có thư viện nên cái quan trọng nhất là hiểu được linh hồn của flow matching là vận tốc

Code của video trên

https://github.com/dome272/Flow-Matching

Để hiểu sâu hơn về nhiều kỹ thuật liên quan các bạn có thể xem bài thuyết trình sau tại NeurIPS

https://neurips.cc/virtual/2024/tutorial/99531

Code

Dưới đây là một số ví dụ triển khai Flow Matching / Conditional Flow Matching phổ biến và dễ tiếp cận trên GitHub. Mình chọn lọc những repo có code rõ ràng, có notebook minh họa, phù hợp cho người mới bắt đầu hoặc muốn hiểu sâu hơn về vận tốc và cách train.

Tutorial cơ bản nhất – Conditional Flow Matching trên MNIST

• Repo: https://github.com/Jackson-Kang/Pytorch-Conditional-Flow-Matching-Tutorial
• Nội dung chính:
• Notebook 01_conditional-flow-matching.ipynb
• Triển khai Conditional Flow Matching (CFM) cơ bản trên MNIST
• Dùng đường thẳng (straight path) → vận tốc hằng số (x₁ − x₀)
• Code ngắn gọn, có giải thích từng bước loss và sampling
• Có cả phần mở rộng MeanFlow
• Phù hợp nếu: Bạn muốn chạy thử ngay trên Colab, hiểu rõ conditional velocity field.

1. Mở Colab
2. Clone repo đầu tiên:

   !git clone https://github.com/Jackson-Kang/Pytorch-Conditional-Flow-Matching-Tutorial
   %cd Pytorch-Conditional-Flow-Matching-Tutorial

   !git clone https://github.com/Jackson-Kang/Pytorch-Conditional-Flow-Matching-Tutorial
   %cd Pytorch-Conditional-Flow-Matching-Tutorial

3. Chạy notebook 01_conditional-flow-matching.ipynb → thấy ngay model học vận tốc và generate số từ nhiễu.
4. Sau khi quen, chuyển sang TorchCFM để thử trên dữ liệu phức tạp hơn (image, conditional generation với class/label).

Thư viện TorchCFM

• Repo: https://github.com/atong01/conditional-flow-matching
• Nội dung chính:
• Thư viện TorchCFM chuyên về Conditional Flow Matching
• Hỗ trợ nhiều biến thể: CFM chuẩn, Rectified Flow, OT-Flow Matching, v.v.
• Có các notebook ví dụ:
  • examples/images/conditional_mnist.ipynb → MNIST conditional generation
  • Các ví dụ khác về image, toy data, optimal transport path
• Code sạch, có package torchcfm để install và tái sử dụng
• Phù hợp nếu: Bạn muốn code chất lượng cao, có thể mở rộng lên image lớn hơn hoặc kết hợp với U-Net.