Tìm hiểu về Flow Matching - Giải thuật Shortcut models

- 13 mins
Tìm hiểu về Flow Matching - Giải thuật Shortcut models

Bài viết này dành cho nhóm bạn đọc: người đã quen với deep learning, Gen AI khi nói về diffusion và lý thuyết học máy thống kê cơ bản.

TL;DR

  • Flow Matching viết bài toán sinh dữ liệu thành phép regression MSE, không cần simulation. Đây là backbone mặc định của gen AI — ảnh, video, giọng nói, robotics — nhờ bốn tầng tối ưu, không chỉ vì “nhanh hơn”.
  • Giới hạn cứng: ngay cả ở optimum, sinh một bước vẫn thất bại với phân phối multimodal.
  • Shortcut Models1 vượt giới hạn bằng cách điều kiện hoá mạng theo cả $t$ và step size $d$, huấn luyện end-to-end với self-consistency bootstrap. Không teacher, không schedule.
  • Paper thực nghiệm thuyết phục nhưng không có phân tích lý thuyết — chính tác giả gọi lời giải thích là “a blind hypothesis”.
  • Phân tích của bài viết chỉ ra: self-consistency là contraction mapping khi $Ld < 1$ — điều kiện paper chỉ áp đặt gián tiếp qua weight decay, và chưa ai đo $L$.

1. Flow Matching: viết lại bài toán sinh dữ liệu thành một phép regression

Continuous Normalizing Flows — mô hình hoá quá trình sinh như vector field $v_\theta(x,t)$ rồi tích phân ODE — không scale được vì phải simulate ODE ngay trong lúc train. Diffusion giải được bài toán scale nhưng đổi lại bộ máy cồng kềnh: SDE thuận/nghịch, noise schedule thiết kế, score ước lượng.

Diffusion Process. Diffusion Process.

Flow Matching2 (song song với Rectified Flow3 và Stochastic Interpolants4) đặt lại câu hỏi: nếu biết trước vector field mục tiêu, tại sao phải simulate?

Lấy $x_0 \sim \mathcal{N}(0,I)$ là noise, $x_1 \sim \mathcal{D}$ là data, định nghĩa:

\[x_t = (1-t)\,x_0 + t\,x_1, \qquad v_t = x_1 - x_0\]

Với mỗi cặp $(x_0, x_1)$, velocity biết trước và là hằng số. Chỉ cần regression:

\[\mathcal{L}^{\text{FM}}(\theta) = \mathbb{E}_{x_0, x_1, t}\Big[\big\|\bar{v}_\theta(x_t, t) - (x_1 - x_0)\big\|^2\Big]\]

Điểm tinh tế: marginal vector field không tính được dạng đóng, nhưng gradient của conditional loss bằng gradient của marginal loss. Regression lên các target dễ tính cho ta nghiệm của bài toán khó. Toàn bộ huấn luyện quy về simulation-free MSE regression.

Giá trị đóng góp của Flow Matching

  1. Objective. Least squares, target dạng đóng, không thành phần ngẫu nhiên. Gradient variance thấp, ít hyperparameter, ổn định ở quy mô lớn.
  2. Transport cost. Linear interpolant tạo quỹ đạo ODE thẳng nhất — kết nối với Optimal Transport. Quỹ đạo càng thẳng, discretization error của Euler càng nhỏ với cùng số bước.
  3. Inference. ODE deterministic, không phải SDE. Dùng solver bậc cao, ánh xạ invertible, cùng seed cho cùng kết quả.
  4. Kinh tế hệ thống. Diffusion cần 50–1000 bước; FM hạ xuống 20–30; shortcut/consistency đẩy xuống 1–4. Mỗi bước là một forward pass của transformer hàng tỉ tham số.

Ứng dụng chính

  • Ảnh/video — ràng buộc chi phí. Stable Diffusion 35 và FLUX dùng rectified flow; WAN 2.x, Movie Gen dùng flow matching.
  • Giọng nói — ràng buộc latency. Voicebox, Matcha-TTS, F5-TTS đưa flow/rectified flow vào TTS few-step.
  • Robotics — ràng buộc thời gian thực. π₀6, SmolVLA sinh action sequence bằng flow matching; control loop 30–50Hz.

Quan trọng nhất: Flow Matching thống nhất: Diffusion, rectified flow, stochastic interpolants đều là trường hợp riêng của cùng một khung.

2. Giới hạn cứng: vì sao một bước vẫn thất bại

Trong huấn luyện, các cặp $(x_0, x_1)$ được ghép ngẫu nhiên — tại điểm $x_t$, nhiều $x_1$ khác nhau đều hợp lệ. Nghiệm tối ưu của least squares là conditional expectation $\mathbb{E}[v_t \mid x_t]$, không phải velocity thật của một quỹ đạo. Hệ quả: tại $t=0$, mạng nhận pure noise và trả về vector trỏ vào trung bình của dataset.

Lưu ý: đây không phải vấn đề capacity. Kể cả ở optimum, one-step generation vẫn thất bại với mọi phân phối multimodal. Lỗi nằm ở objective, không nằm ở mô hình.

Learn a step-size conditioned model that is grounded to the data at a small step-size, and trained via self-bootstrapping at larger step-sizes. Hình 2 — Learn a step-size conditioned model that is grounded to the data at a small step-size, and trained via self-bootstrapping at larger step-sizes..

3. Ba nhánh giải pháp

  Cần teacher? Số lần train Bootstrap Cần schedule? Giữ many-step?
Distillation (Progressive Distillation7, Reflow) Có 2+ — Không Không
Consistency Training8 Không 1 $T$ Có Hạn chế
Shortcut Models Không 1 $\log_2 T$ Không Có

Shortcut Models khác biệt ở hai điểm: (1) lấy $x’_{t+d}$ từ chính ODE đã học, tránh bias tích luỹ của consistency models; (2) chỉ cần $\log_2 T$ bootstrap thay vì $T$. Chi phí: một bước huấn luyện shortcut tốn thêm 16% so với diffusion thường.

4. Cơ chế: điều kiện hoá theo $d$ và self-consistency

Gọi $\mathcal{D}_M = {1/M, \dots, 1}$ là tập step size rời rạc. Shortcut model học:

\[s_\theta : \mathcal{X} \times [0,1] \times \mathcal{D}_M \to \mathcal{X}, \qquad x'_{t+d} = x_t + s_\theta(x_t, t, d)\cdot d\]

Ý tưởng: nếu mạng biết trước nó nhảy xa bao nhiêu, nó có thể tính đến độ cong phía trước và nhảy tới đúng điểm.

Target cho $d$ lớn lấy từ self-consistency: một bước nhảy $2d$ bằng hai bước nhảy $d$ liên tiếp. Ghép với flow matching làm base case $d=0$:

\[\mathcal{L}^{\text{SC}}(\theta) = \mathbb{E}\Big[\underbrace{\|s_\theta(x_t, t, 0) - (x_1 - x_0)\|^2}_{\text{Flow-Matching}} + \underbrace{\|s_\theta(x_t, t, 2d) - s_{\text{target}}\|^2}_{\text{Self-Consistency}}\Big]\] \[s_{\text{target}} = \texttt{stopgrad}\Big(\tfrac{1}{2} s_\theta(x_t, t, d) + \tfrac{1}{2} s_\theta(x'_{t+d}, t+d, d)\Big)\]

Giải thuật huấn luyện được tóm tắt như sau:

\[\boxed{ \begin{array}{l} \text{Algorithm 1 Shortcut Model Training} \\ \text{while not converged do} \\ \quad x_0 \sim \mathcal{N}(0, I),\; x_1 \sim D,\; (d,t) \sim p(d,t) \\ \quad x_t \leftarrow (1-t)x_0 + t x_1 \quad \triangleright \text{Noise data point} \\ \quad \text{for first } k \text{ batch elements do} \\ \quad\quad s_{\text{target}} \leftarrow x_1 - x_0 \quad \triangleright \text{Flow-matching target} \\ \quad\quad d \leftarrow 0 \\ \quad \text{end for} \\ \quad \text{for other batch elements do} \\ \quad\quad s_t \leftarrow s_\theta(x_t, t, d) \quad \triangleright \text{First small step} \\ \quad\quad x_{t+d} \leftarrow x_t + s_t d \quad \triangleright \text{Follow ODE} \\ \quad\quad s_{t+d} \leftarrow s_\theta(x_{t+d}, t+d, d) \quad \triangleright \text{Second small step} \\ \quad\quad s_{\text{target}} \leftarrow \text{stopgrad}(s_t + s_{t+d})/2 \quad \triangleright \text{Self-consistency target} \\ \quad \text{end for} \\ \quad \theta \leftarrow \nabla_\theta \| s_\theta(x_t, t, 2d) - s_{\text{target}} \|^2 \\ \text{end while} \end{array} }\] \[\boxed{ \begin{array}{l} \text{Algorithm 2 Sampling}. \\ x \sim \mathcal{N}(0, I) \\ d \leftarrow 1/M \\ t \leftarrow 0 \\ \text{for } n \in [0, \ldots, M-1] \text{ do} \\ \quad x \leftarrow x + s_\theta(x, t, d) d \\ \quad t \leftarrow t + d \\ \text{end for} \\ \text{return } x \end{array} }\]

Samples from Celeba library Samples from Celeba library.

Ba chi tiết quyết định: stopgrad trên cả hai lần đánh giá; EMA weights khi sinh target; weight decay $\lambda = 0.1$ — tác giả gọi là crucial. Toàn bộ phân tích lý thuyết dưới đây xoay quanh câu hỏi: ba mẹo này đang thay thế cho điều kiện toán học nào?

5. Tác giả đã bỏ ngỏ những gì

(a) Không có phân tích lý thuyết. Không hypothesis space, không error bound, không điều kiện hội tụ.

(b) Kết quả trái trực giác bị bỏ lửng. Trên CelebAHQ-256, Shortcut Models đạt 128-step FID 6.9, tốt hơn Flow Matching 7.3 — dù mang thêm bootstrap bias. Tác giả gọi đây là “a blind hypothesis” về implicit regularization.

(c) Ổn định huấn luyện bằng mẹo, không bằng điều kiện. stopgrad + EMA + weight decay giữ training không phân kỳ, nhưng tại sao thì không ai chứng minh.

(d) Khoảng cách many-step ↔ one-step vẫn còn (6.9 → 20.5), và ánh xạ noise→data bị quyết định bởi kỳ vọng trên dataset — khác GAN hay VAE.

(e) $p(d)$ — nhân đôi nhị phân, cách đều trên thang log — không được biện luận.

Ba câu hỏi lý thuyết trung tâm:

  1. Hypothesis space $\mathcal{H}_{\text{SC}}$ có chứa shortcut lý tưởng $s^*$ không?
  2. Toán tử huấn luyện có hội tụ về $s^*$ không, và dưới điều kiện nào?
  3. Mở rộng $\mathcal{H}_{\text{FM}} \to \mathcal{H}_{\text{SC}}$ làm thay đổi generalization error ra sao?

6. Điều tra A — phân tích bằng statistical learning theory

6.1 Hypothesis space

Với $\mathcal{E}(h)$ là expected loss, $\mathcal{E}^* = \inf_h \mathcal{E}(h)$:

\[\underbrace{\mathcal{E}(h_D)}_{\text{total}} = \underbrace{\mathcal{E}^*}_{\text{optimal}} + \underbrace{\big[\mathcal{E}(h^*_{\mathcal{H}}) - \mathcal{E}^*\big]}_{\text{approximation}} + \underbrace{\big[\mathcal{E}(h_D) - \mathcal{E}(h^*_{\mathcal{H}})\big]}_{\text{estimation}}\]

Mở rộng $\mathcal{H}$: approximation giảm, estimation tăng. Điều thú vị: Shortcut Models làm cả ba hạng tử đều dịch chuyển.

Khi $d \to 0$, $s_\theta(x_t,t,0) = \bar{v}_\theta(x_t,t)$, nên $\mathcal{H}_{\text{FM}} \subset \mathcal{H}_{\text{SC}}$. Nhưng mở rộng này không tự do. Nếu mỗi $d$ có bộ tham số riêng $\theta_d$, ta có $\mathcal{H}_{\text{free}} = \mathcal{H}_{\text{FM}}^{|\mathcal{D}_M|+1}$. Ràng buộc parameter sharing $\theta_d \equiv \theta$ áp đặt inductive bias. Ta được $\mathcal{H}_{\text{FM}} \subset \mathcal{H}_{\text{SC}} \subset \mathcal{H}_{\text{free}}$.

Sơ đồ Venn ba lớp lồng nhau: H_free (ngoài), H_SC (giữa), H_FM (trong). Chấm đỏ s* nằm trong H_free nhưng ngoài H_SC. Cấu trúc hypothesis space.

6.2 Bayes floor bị nâng lên

Flow Matching có hai thành phần trong $\mathcal{E}^*$: multimodal stochasticity và discretization error $O(1/M)$. Shortcut Models thêm thành phần thứ ba:

Bootstrap bias. Tại vòng lặp $k$, target $s_{\text{target}} = f(\theta^{(k)})$ là tự sinh. Dù với dữ liệu vô hạn, nếu mô hình sai số $\varepsilon_k$, target mang sai số $O(\varepsilon_k)$. Do đó $\mathcal{E}^*_{\text{SC}} = \mathcal{E}^*_{\text{FM}} + \varepsilon_{\text{bootstrap}}$, $\varepsilon_{\text{bootstrap}} > 0$.

Bias - Variance analysis Bias - Variance analysis.

Kết quả $6.9 < 7.3$ không mâu thuẫn với bất đẳng thức trên: FID đo total error, không đo Bayes floor. $\mathcal{E}^*_{\text{SC}} > \mathcal{E}^*_{\text{FM}}$ nhưng approximation + estimation của shortcut nhỏ hơn đủ để bù. Đó là hình dạng của một regularizer: nâng sàn, hạ trần.

6.3 Approximation và estimation

Approximation — được: velocity model thuần tuý không thể biểu diễn bước nhảy lớn đúng. $\mathcal{H}_{\text{SC}}$ thêm expressiveness này.

Approximation — mất: một $\theta$ duy nhất phải đúng cho mọi $d$, tạo khoảng cách $\delta_{\text{shared}} \ge 0$ so với $\mathcal{H}_{\text{free}}$.

Estimation — moving target: $s_{\text{target}} = f(\theta^{(k-1)})$ phụ thuộc tham số hiện tại — đồng dạng với TD-learning trong RL, vốn bất ổn. Chứng minh hội tụ chuẩn của SGD không áp dụng.

Estimation — error propagation: $M=128$, $\log_2 128 = 7$ bậc lũy thừa 2. Sai số $\varepsilon_k$ ở tầng $k$ thừa kế vào tầng $k+1$.

6.4 Kết quả chính: contraction condition $Ld < 1$

Toán tử self-consistency. $\mathcal{T} : \mathcal{H}_{\text{SC}} \to \mathcal{H}_{\text{SC}}$: \((\mathcal{T}s)(x_t, t, 2d) = \tfrac{1}{2} s(x_t, t, d) + \tfrac{1}{2} s\big(x_t + s(x_t,t,d)\cdot d,\; t+d,\; d\big)\). Huấn luyện là tìm fixed point $\mathcal{T}s^* = s^*$.

Định lý. Giả sử $s^*_\theta$ thoả $|s^*_\theta(x,t,d) - s^*_\theta(y,t,d)| \le L|x-y|$. Khi đó $\mathcal{T}$ là contraction mapping với hệ số $\kappa = \frac{1+Ld}{2}$. Theo Banach fixed-point theorem9, $\mathcal{T}$ có fixed point duy nhất và hội tụ khi $\kappa < 1$, tức $\boxed{Ld < 1}$.

Chứng minh. Đặt $x’_{t+d} = x_t + s(x_t,t,d)d$ và tương tự cho $\tilde{s}$:

\[\|x'_{t+d} - \tilde{x}'_{t+d}\| = d\|s - \tilde{s}\|_\infty\]

Áp dụng triangle inequality cho hai nửa của $\mathcal{T}$:

\[\|(\mathcal{T}s) - (\mathcal{T}\tilde{s})\| \le \tfrac{1}{2}\|s - \tilde{s}\|_\infty + \tfrac{1}{2}\big(\|s - \tilde{s}\|_\infty + L d \|s - \tilde{s}\|_\infty\big) = \tfrac{1+Ld}{2}\|s-\tilde{s}\|_\infty\]

Lấy supremum. $\square$

Ba hệ quả:

  • Điều kiện khắt khe nhất ở one-step ($d=1$): $L < 1$ — mạng phải là contraction theo input.
  • Weight decay là proxy, không phải chứng minh. Paper không đo $L$; dùng weight decay và EMA gián tiếp chặn $L$.
  • Norm mismatch: chứng minh dùng $\ell_\infty$, huấn luyện dùng $\ell_2$.

Sơ đồ chuỗi bốn hộp: d=0 (ε₀) → d=1/64 (κ·ε₀) → d=1/32 (κ²·ε₀) → ... → d=1 (κ⁷·ε₀). Điều kiện contraction khi Ld < 1. Error propagation qua 7 tầng.

6.5 Generalization bound

Hướng tiếp cận Đánh giá
Complexity-based (VC dim, Rademacher) ❌ Với DiT-B hàng trăm triệu tham số, bound > 1
Algorithmic stability10 ✅ Khả thi nhất
PAC-Bayes ⚠️ Vướng phụ thuộc vòng
Sharpness-based ⚠️ Tính được sau train

Với stability: $\beta = \frac{\eta\,\delta}{1 - \kappa} = \frac{\eta\,\delta}{(1 - Ld)/2}$.

7. Kết

Flow Matching kéo NFE từ hàng trăm về hàng chục, bằng cách đơn giản hoá. Shortcut Models kéo từ hàng chục về một — vẫn một mạng, một lần huấn luyện.

Nhưng đây là phương pháp bootstrap, và nó không sụp đổ mà không ai biết chắc tại sao. Phân tích ở trên đưa ra một câu trả lời: nếu $Ld < 1$, self-consistency là contraction mapping, fixed point tồn tại và duy nhất, sai số giảm geometric qua 7 tầng. Nếu không: mọi thứ đứng vững nhờ weight decay và may mắn.

Ba mẹo engineering — stopgrad, EMA, weight decay — hoá ra là ba cách gián tiếp giữ $L$ đủ nhỏ. Thứ tác giả để ngỏ không phải phương pháp, mà là lý do phương pháp hoạt động. Khoảng cách giữa hai thứ đó rộng đúng bằng một phép đo spectral norm mà chưa ai thực hiện.

Tham khảo

  1. K. Frans, D. Hafner, S. Levine, P. Abbeel. One Step Diffusion via Shortcut Models. ICLR, 2025. arXiv:2410.12557 ↩

  2. Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, M. Le. Flow Matching for Generative Modeling. ICLR, 2023. arXiv:2210.02747 ↩

  3. X. Liu, C. Gong, Q. Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. ICLR, 2023. arXiv:2209.03003 ↩

  4. M. S. Albergo, E. Vanden-Eijnden. Building Normalizing Flows with Stochastic Interpolants. ICLR, 2023. arXiv:2209.15571 ↩

  5. P. Esser và cộng sự. Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. ICML, 2024. arXiv:2403.03206 ↩

  6. K. Black và cộng sự. π₀: A Vision-Language-Action Flow Model for General Robot Control. 2024. arXiv:2410.24164 ↩

  7. T. Salimans, J. Ho. Progressive Distillation for Fast Sampling of Diffusion Models. ICLR, 2022. arXiv:2202.00512 ↩

  8. Y. Song, P. Dhariwal, M. Chen, I. Sutskever. Consistency Models. ICML, 2023. arXiv:2303.01469 ↩

  9. S. Banach. Sur les opérations dans les ensembles abstraits et leur application aux équations intégrales. Fundamenta Mathematicae, 3(1):133–181, 1922. ↩

  10. O. Bousquet, A. Elisseeff. Stability and Generalization. JMLR, 2:499–526, 2002. ↩

Machine Learning AI