Lời mở đầu

Gần đây, chúng ta đã chứng kiến sự phát triển chóng mặt của các mô hình sinh ảnh từ văn bản. Các mô hình như Mid-journey, Dall-e 3, Imagen, Stable Diffusion, StyleGAN-T,… có thể tạo ra những hình ảnh chân thực và sống động đến mức điên rồ.

Dù vậy, các mô hình đều có lượng thời gian suy diễn (hay sinh ảnh) rất lớn do cần lặp đi lặp lại quá trình khử nhiễu như ở các mô hình Diffusion. Các nghiên cứu gần đây ứng dụng quá trình chắt lọc tri thức (knowledge distillation) đã có thể giảm số bước xuống rất nhỏ (còn khoảng < 20 bước) để cho ra những bức ảnh có chất lượng tốt với mô hình Stable Diffusion tuy nhiên chưa từng có phương pháp được kiểm chứng đầy đủ nào có thể giảm số bước trong quá trình infer xuống chỉ còn “1 bước duy nhất”.

Ảnh 1

Ảnh 1

Hãy thử tưởng tượng từ một bức ảnh nhiễu, chúng ta có thể khử hoàn toàn trong một bước để sinh ra một bức ảnh có chất lượng tốt, điều đó có thể tiết kiệm cho chúng ta gấp nhiều lần thời gian khi muốn sinh ảnh từ văn bản. Nghe thật tuyệt đúng không nào?

Tuy nhiên khi ta cố giảm số step xuống một bước bằng quá trình chắt lọc tri thức, mô hình student học được cho ra kết quả vô cùng tệ. Điều này có thể lí giải bởi mô hình student chỉ có khả năng bắt chước mô hình SD teacher mà không có nhiều thông tin về mối liên hệ “có đôi có cặp” giữa ảnh nhiễu và ảnh gốc, cũng như quỹ đạo phức tạp mà ảnh nhiễu được vận chuyển (Bản chất Stable Diffusion có thể được mô hình hoá bởi phương trình vi phân, việc giải ptvp từ thời điểm T (ảnh nhiễu) đến 0 (ảnh gốc) tương ứng với quá trình sinh ảnh, sự di chuyển của điểm dữ liệu theo ptvp là quỹ đạo đang nói tới). Như ví dụ bên tay trái của ảnh phía dưới đây, quá trình distillation khiến cho kết quả vừa lẫn lộn, mờ và khó nhận ra được context của ảnh.

Ảnh 2

Ảnh 2

InstaFlow là một mô hình được phát triển từ Stable Diffusion có khả năng giải quyết thách thức đó. Bằng cách “ghép cặp” đúng và hiệu quả các điểm dữ liệu thuộc 2 phân phối khác nhau trước khi áp dụng quá trình distillation, InstaFlow đã có thể “một bước lên tiên” biến từ một ảnh nhiễu thành ảnh có chất lượng tốt, rõ ràng về nội dung. Vậy mô hình này đã sử dụng phép thuật gì hay ho ở đây? Hãy cùng mình đi sâu hơn về nó nhé!

Rectified flow và Reflow

Rectified flow

Ảnh 3

Ảnh 3

Mô hình luồng chỉnh lưu (Rectified Flow) là một chuỗi các biến đổi ánh xạ một điểm từ không gian này sang không gian khác hay từ một phân phối này sang phân phối khác (ví dụ như $\pi_0$ đến phân phối $\pi_1$).

<aside> ❓ Cho trước các quan sát của 2 phân phối $\pi_0$ và $\pi_1$, tìm ánh xạ $T\colon\mathbb{R}^d\to\mathbb{R}^d$, thoả mãn $Z_1 := T(Z_0)\sim \pi_1$ khi $Z_0 \sim \pi_0$

</aside>

Ví dụ trong quá trình sinh ảnh, $\pi_0$ thường là những phân phối đơn giản như phân phối gaussian và $\pi_1$ là phân phối của dữ liệu ảnh. Biết rằng $X_0 \sim \pi_0, X_1 \sim \pi_1$, dựa trên 2 quan sát này, Rectified Flows cố gắng học cách biến đổi dữ liệu $Z_0$ được lấy mẫu từ $\pi_0$ thành $Z_1$ thuộc phân phối $\pi_1$ tương ứng bằng phương trình vi phân thường (ODE- Ordinary Differential Equations) có dạng

$$ \begin{align} \frac{dZ_t}{dt} = v(Z_t, t) \end{align} $$

Phương trình này có dạng thật quen đúng không nào, nếu ta xem $Z$ là quãng đường thì $v(Z_t,t)$ là vận tốc tức thời tại thời điểm t. Lúc này bài toán chúng ta đang giải là tìm ra con đường dẫn từ một địa điểm $Z_0$ đến $Z_1$ bằng cách đi từng bước nhỏ theo hướng của $v(Z_t,t)$. Hay ví dụ khác thể hiện rõ sự khác biệt về tính chất hay phân phối của các điểm dữ liệu: Tưởng tượng chúng ta đang muốn biến đổi một hình vuông 🟦 thành một hình tròn 🔵, ta có thể áp dụng một chuỗi các biến đổi nhỏ hơn ví dụ như từng bước gọt đi các góc nhọn của đa giác rồi cứ lặp lại như vậy cho đến khi hình thu được gần như tròn hoàn toàn.

Ảnh 4

Ảnh 4

Oke, có thể bây giờ bạn đã hiểu rõ hơn về ODE, và bạn sẽ nhận ra đây vẫn là một bài toán khó nhằn! Một trong những cách để tìm ra $v$ chính là đi cực tiểu hoá khoảng cách giữa $\rho_1^v$ - phân phối của $Z_1$ khi đi theo luồng $v$ với phân phối $\pi_1$, ví dụ như phân kì KL (KL Divergence) $D_{KL}(\rho_1^v;~ \pi_1)$, tuy nhiên giống như có vô số con đường có thể đi giữa 2 điểm, cũng sẽ có vô số phương trình vi phân nên việc tính toán $\rho_1^v$ là hết sức tốn kém.

Ảnh 5

Ảnh 5