Flow Maps: Khám phá "Tích phân" của Mô hình Diffusion để Tăng tốc Tạo mẫu AI

Flow Maps: Khám phá "Tích phân" của Mô hình Diffusion để Tăng tốc Tạo mẫu AI

Việc tạo mẫu (sampling) từ một mô hình Diffusion hiện nay về cơ bản là một quá trình tính toán tích phân lặp đi lặp lại. Tại mỗi bước, mô hình khử nhiễu (denoiser) ước lượng hướng tiếp tuyến của một đường dẫn trong không gian đầu vào, và chúng ta di chuyển theo hướng này bằng các bước nhỏ. Quá trình này dần dần chuyển đổi các mẫu từ một phân phối nhiễu đơn giản thành phân phối dữ liệu mục tiêu.

Tuy nhiên, câu hỏi đặt ra là: Liệu chúng ta có thể huấn luyện các mạng nơ-ron để dự đoán trực tiếp tích phân này — tức là dự đoán ngay lập tức điểm đến trên đường dẫn — nhằm tăng tốc độ tạo mẫu hay không? Câu trả lời là có, và chúng ta hãy cùng bước vào thế giới của Flow Maps.

Sơ đồ minh họa quá trình chuyển đổi từ nhiễu sang dữ liệu

Từ Diffusion đến Flow Maps

Kể từ khi mô hình Diffusion bùng nổ, các nhà nghiên cứu luôn tìm cách làm cho chúng nhanh hơn và rẻ hơn. Một trong những công cụ chính là "chưng cất" (distillation), giúp giảm số bước lặp cần thiết. Flow Maps là một bước tiến xa hơn trong hướng tiếp cận này.

Trong khi mô hình Diffusion truyền thống mô tả các đường dẫn giữa nhiễu và dữ liệu bằng cách dự đoán hướng tiếp tuyến tại mỗi điểm (một cái nhìn cục bộ), Flow Maps lại có khả năng dự đoán bất kỳ điểm nào trên một đường dẫn từ bất kỳ điểm nào khác trên cùng đường dẫn đó (một cái nhìn toàn cục).

Hai phương pháp tạo mẫu: Ngẫu nhiên và Xác định

Để hiểu Flow Maps, ta cần nhìn lại các thuật toán tạo mẫu của Diffusion. Chúng chia làm hai loại: ngẫu nhiên (stochastic) và xác định (deterministic).

Phương pháp ngẫu nhiên, như DDPM gốc, tạo ra các đường dẫn gồ ghề và khó lường. Ngược lại, phương pháp xác định (như DDIM hoặc dựa trên ODE) tạo ra các đường dẫn mượt mà. Điều quan trọng là cả hai phương pháp đều bảo toàn các phân phối biên (marginal distributions), nhưng phương pháp xác định thiết lập một ánh xạ một-một (bijection) duy nhất giữa một mẫu nhiễu và một mẫu dữ liệu.

Minh họa sự khác biệt giữa đường dẫn ngẫu nhiên và xác định

Điều này có nghĩa là với một mẫu nhiễu cụ thể, luôn có một đường dẫn duy nhất dẫn đến một mẫu dữ liệu cụ thể. Flow Maps tận dụng tính chất này để "nhảy" dọc theo đường dẫn đó thay vì đi từng bước nhỏ.

Cơ chế hoạt động của Flow Maps

Hãy ký hiệu $f(\mathbf{x}_t, t)$ là một mô hình khử nhiễu truyền thống. Flow Maps, được ký hiệu là $F(\mathbf{x}_s, s, t)$, nhận vào một điểm tại thời gian $s$ và dự đoán vị trí của nó tại thời gian $t$ trên cùng một đường dẫn.

Về mặt toán học, Flow Maps có thể được xây dựng bằng cách tích hợp vận tốc (tangent direction) trên một khoảng thời gian:

$$F(\mathbf{x}_s, s, t) = \mathbf{x}s + \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau$$

Trong đó $v(\mathbf{x}_t, t)$ là vận tốc tại điểm đó. Nếu ta đặt $t=0$ (mức dữ liệu sạch), Flow Maps có thể đưa ta từ bất kỳ điểm nào trên đường dẫn trực tiếp đến đích, cho phép tạo mẫu chỉ trong một bước duy nhất. Đây chính là ý tưởng cốt lõi của các Consistency Models (Mô hình nhất quán), và Flow Maps là sự khái quát hóa của ý tưởng này.

Ba quy tắc nhất quán để huấn luyện

Việc huấn luyện một Flow Maps phức tạp hơn nhiều so với mô hình Diffusion thông thường vì nó đòi hỏi sự hiểu biết toàn cục về đường dẫn. Các nhà nghiên cứu đã đề xuất ba quy tắc "nhất quán" (consistency rules) để giải quyết vấn đề này:

1. Tính kết hợp (Compositionality): Đi từ điểm $s$ đến $u$ phải cho kết quả giống như đi từ $s$ đến $t$ rồi từ $t$ đến $u$.
2. Nhất quán Lagrangian: Tập trung vào việc thay đổi đích ($t$). Đầu ra của Flow Maps phải thay đổi theo vận tốc khi ta thay đổi thời gian đích.
3. Nhất quán Eulerian: Tập trung vào việc thay đổi điểm xuất phát ($s$). Nếu ta thay đổi điểm bắt đầu nhưng giữ đích cố định, kết quả dự đoán của Flow Maps không được thay đổi.

Từ các quy tắc này, các hàm mất mát (loss functions) được xây dựng để buộc mạng nơ-ron học được các ánh xạ toàn cục này.

Minh họa tính chất kết hợp trong Flow Maps

Thách thức và Cơ hội

Mặc dù Flow Maps mang lại lợi ích to lớn về tốc độ suy luận (inference) — có thể giảm từ hàng chục bước xuống chỉ còn một bước — nhưng chúng không phải là "bữa trưa miễn phí".

Việc huấn luyện Flow Maps thường phức tạp hơn nhiều và đòi hỏi nhiều tài nguyên tính toán hơn. Hầu hết các phương pháp hiện nay đều "bootstrap" từ một mô hình Diffusion đã được huấn luyện trước (dùng mô hình Diffusion làm thầy để dạy Flow Maps). Ngoài ra, để ổn định quá trình huấn luyện, các kỹ thuật như stop-gradient (ngắt gradient) hoặc sử dụng trung bình động mũ (EMA) thường được áp dụng.

Kết luận

Flow Maps đại diện cho một sự chuyển dịch tư duy từ việc dự đoán hướng đi cục bộ sang lập bản đồ toàn cục cho quá trình tạo mẫu AI. Chúng mở ra cánh cửa cho việc tạo hình ảnh và dữ liệu thời gian thực với chi phí tính toán thấp hơn nhiều, đồng thời cung cấp các công cụ mạnh mẽ mới cho việc học dựa trên phần thưởng (reward-based learning) và khả năng điều khiển mẫu (steerability).

Trong tương lai, chúng ta có thể mong đợi thấy越来越多的 các mô hình áp dụng Flow Maps hoặc các biến thể của nó, đưa sức mạnh của AI tạo sinh đến gần hơn với các ứng dụng thực tế đòi hỏi tốc độ cao.

---

Tham khảo thêm: Bài viết gốc của Sander Dieleman trên sander.ai