11Mô hình khuếch tán#

Mô hình khuếch tán phá dần dữ liệu bằng nhiễu cho tới khi chỉ còn nhiễu, rồi học cách đảo ngược từng bước. Mỗi bước đảo ngược là một bài khử nhiễu nhỏ, và như ta sẽ thấy, huấn luyện cả mô hình quy về một bài hồi quy thông thường: không có trò chơi giữa hai mạng như GAN, và mất mát giảm nghĩa là mô hình tốt lên.

11.1Sinh mẫu bằng nhiều bước khử nhiễu#

Sinh một ảnh từ nhiễu trong một bước là bài toán rất khó. Nhưng khử một chút nhiễu khỏi một ảnh đã hơi nhiễu là bài toán khử nhiễu quen thuộc, tức một bài hồi quy. Mô hình khuếch tán (Sohl-Dickstein và cộng sự, 2015; Ho, Jain và Abbeel, 2020) ghép nhiều bước khử nhiễu nhỏ, thường 1 000 bước, thành một quá trình sinh. Cách tư duy này dùng được ở nhiều nơi: khi một phép biến đổi quá khó để học trực tiếp, tìm cách viết nó thành một chuỗi nhiều phép biến đổi dễ.

11.2Quá trình thuận và dạng đóng#

Quá trình thuận thêm nhiễu Gauss từng chút một:

$$q(x_t \mid x_{t-1}) = \mathcal{N}\big(\sqrt{1 - \beta_t}\,x_{t-1},\; \beta_t I\big),$$

với $\beta_t$ nhỏ, tăng dần theo một lịch định trước. Hệ số $\sqrt{1 - \beta_t}$ thu nhỏ tín hiệu một chút ở mỗi bước, vừa đủ để nếu $x_{t-1}$ có phương sai 1 thì $x_t$ cũng có phương sai 1.

Mệnh đề 11.1 (Dạng đóng của quá trình thuận)

Đặt $\alpha_t = 1 - \beta_t$ và $\bar\alpha_t = \prod_{s=1}^{t}\alpha_s$. Khi đó

$$q(x_t \mid x_0) = \mathcal{N}\big(\sqrt{\bar\alpha_t}\,x_0,\; (1 - \bar\alpha_t)\,I\big),$$

tức $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1 - \bar\alpha_t}\,\varepsilon$ với $\varepsilon \sim \mathcal{N}(0, I)$.

Chứng minh

Quy nạp theo $t$. Với $t = 1$ mệnh đề là định nghĩa. Giả sử $x_{t-1} = \sqrt{\bar\alpha_{t-1}}\,x_0 + \sqrt{1 - \bar\alpha_{t-1}}\,\varepsilon'$. Khi đó $x_t = \sqrt{\alpha_t}\,x_{t-1} + \sqrt{\beta_t}\,\varepsilon'' = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{\alpha_t(1 - \bar\alpha_{t-1})}\,\varepsilon' + \sqrt{1 - \alpha_t}\,\varepsilon''$. Tổng của hai biến Gauss độc lập có trung bình 0 là biến Gauss có phương sai bằng tổng các phương sai: $\alpha_t - \bar\alpha_t + 1 - \alpha_t = 1 - \bar\alpha_t$.

Dạng đóng là điều làm cho huấn luyện khả thi: muốn có $x_t$ với $t$ bất kỳ, lấy trực tiếp từ $x_0$ trong một bước, không phải mô phỏng $t$ bước trung gian.

Dạng đóng kiểm chứng được bằng mô phỏng. Thí nghiệm dùng lịch tuyến tính $\beta$ từ $10^{-4}$ tới $0{,}02$ qua 1 000 bước, giá trị ban đầu $x_0 = 2$, mô phỏng 200 000 quỹ đạo đi từng bước rồi so trung bình và phương sai với dạng đóng. Chỉ số bước $t$ đánh từ 0 như trong mã, tức $t = 0$ là sau bước nhiễu đầu tiên.

Hai cách tính khớp nhau trong phạm vi sai số lấy mẫu của 200 000 quỹ đạo.

11.3Lịch nhiễu và tỉ số tín hiệu trên nhiễu#

$\bar\alpha_t$ cho biết phần tín hiệu còn lại sau $t$ bước. Đại lượng thường dùng hơn là tỉ số tín hiệu trên nhiễu

$$\mathrm{SNR}(t) = \frac{\bar\alpha_t}{1 - \bar\alpha_t}.$$

Ở bước cuối chỉ còn 0,64% biên độ tín hiệu, tức $x_T$ gần như là nhiễu Gauss thuần tuý. Điều này cần thiết vì khi sinh mẫu, quá trình bắt đầu từ $\mathcal{N}(0, I)$, và điểm xuất phát đó phải khớp với điểm kết thúc của quá trình thuận.

SNR cũng cho biết mô hình học gì ở mỗi bước. Ở $t$ nhỏ, SNR rất cao, ảnh gần như nguyên vẹn, và việc khử nhiễu liên quan tới các chi tiết nhỏ. Ở $t$ lớn, SNR rất thấp, gần như chỉ còn nhiễu, và việc khử nhiễu liên quan tới bố cục và màu sắc tổng thể. Khi sinh mẫu đi ngược từ $t = T$ về $t = 0$, mô hình vì vậy quyết định bố cục trước rồi mới thêm chi tiết. Nichol và Dhariwal (2021) chỉ ra lịch tuyến tính làm SNR giảm quá nhanh ở cuối, khiến nhiều bước cuối gần như lãng phí, và đề xuất lịch cosine.

11.4Huấn luyện bằng hồi quy#

Từ dạng đóng, $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1 - \bar\alpha_t}\,\varepsilon$. Mô hình khuếch tán huấn luyện một mạng $\varepsilon_\theta(x_t, t)$ để dự đoán lại nhiễu $\varepsilon$ từ $x_t$ và $t$:

$$\mathcal{L} = \mathbb{E}_{x_0,\,\varepsilon,\,t}\Big[\big\|\varepsilon - \varepsilon_\theta(x_t, t)\big\|^2\Big].$$

Đây là bình phương sai số, một bài hồi quy thông thường như ở Chương 4 của Nền tảng. Không có trò chơi giữa hai mạng, không có điểm yên ngựa; mất mát giảm nghĩa là mô hình tốt lên. Thủ tục huấn luyện gồm bốn bước:

1. Lấy một mẫu thật $x_0$.
2. Chọn ngẫu nhiên một bước $t \in \{1, \dots, T\}$.
3. Lấy $\varepsilon \sim \mathcal{N}(0, I)$ và tính $x_t$ bằng dạng đóng.
4. Tính bình phương sai số giữa $\varepsilon$ và $\varepsilon_\theta(x_t, t)$, cập nhật tham số.

Hàm mất mát này là dạng đơn giản hoá của một chặn dưới biến phân cho log hợp lý, tương tự ELBO của VAE; Ho, Jain và Abbeel (2020) thấy bỏ các trọng số theo $t$ của chặn dưới cho chất lượng mẫu tốt hơn. Vì biết $\varepsilon$ thì suy ra được $x_0$ từ dạng đóng, dự đoán $\varepsilon$ và dự đoán $x_0$ tương đương về mặt toán học; dự đoán $\varepsilon$ là một lựa chọn thực nghiệm. Chẳng hạn với $x_0 = 2$ ở bước 200 của lịch tuyến tính, $\sqrt{\bar\alpha_t} = 0{,}810$ và $\sqrt{1 - \bar\alpha_t} = 0{,}586$. Nếu nhiễu rút ra là $\varepsilon = 0{,}5$ thì $x_t = 0{,}810 \cdot 2 + 0{,}586 \cdot 0{,}5 = 1{,}913$, và đó là đầu vào của mạng, còn đáp án mạng phải đoán là 0,5. Biết đáp án đó thì $x_0 = (1{,}913 - 0{,}586 \cdot 0{,}5)/0{,}810 = 2$. Khi sinh mẫu, mỗi bước dùng $\varepsilon_\theta$ để ước lượng giá trị trung bình của $x_{t-1}$, rồi cộng thêm một lượng nhiễu nhỏ.

11.5Tăng tốc lấy mẫu#

Nhược điểm chính của mô hình khuếch tán là lấy mẫu chậm: sinh một mẫu cần chạy mạng $T$ lần, và với $T = 1\,000$ là chậm hơn GAN khoảng một nghìn lần. Bảng dưới tóm tắt các cách khắc phục chính.

Dòng cuối là cơ sở của nhiều mô hình sinh ảnh chạy được trên máy cá nhân, và là một ví dụ kết hợp hai họ mô hình: bộ tự mã hoá lo việc nén ảnh xuống không gian ẩn nhỏ, mô hình khuếch tán lo việc sinh trong không gian đó.

11.6Sinh có điều kiện#

Các ứng dụng như sinh ảnh từ mô tả văn bản cần mô hình $p(x \mid c)$ với $c$ là điều kiện. Mạng $\varepsilon_\theta(x_t, t, c)$ nhận thêm $c$, thường là embedding của câu mô tả do một bộ mã hoá văn bản tạo ra, đưa vào qua cross-attention (Mục 9.1 của Học sâu).

Kỹ thuật classifier-free guidance (Ho và Salimans, 2022) làm cho mẫu bám điều kiện chặt hơn: huấn luyện cùng một mạng với cả $c$ và với điều kiện rỗng (bỏ $c$ ngẫu nhiên khoảng 10% số lần), rồi khi sinh dùng

$$\tilde\varepsilon = \varepsilon_\theta(x_t, t, \varnothing) + w\,\big(\varepsilon_\theta(x_t, t, c) - \varepsilon_\theta(x_t, t, \varnothing)\big),$$

với $w > 1$ đẩy mẫu theo hướng mà điều kiện $c$ tạo ra. Chẳng hạn nếu ở một toạ độ, dự đoán không điều kiện là 0,2 và dự đoán có điều kiện là 0,5, thì với $w = 3$ ta dùng $0{,}2 + 3 \cdot 0{,}3 = 1{,}1$: không dừng ở dự đoán có điều kiện mà đi xa hơn theo cùng hướng. $w$ lớn cho mẫu khớp mô tả hơn nhưng kém đa dạng hơn. Chương 14 của Ứng dụng LLM bàn về việc dùng các mô hình sinh ảnh trong ứng dụng.

11.7Tóm tắt#

Quá trình thuận thêm nhiễu Gauss từng chút một, và nhờ dạng đóng $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1 - \bar\alpha_t}\,\varepsilon$, ta lấy được $x_t$ ở bước bất kỳ trong một phép tính; mô phỏng 200 000 quỹ đạo khớp với dạng đóng trong phạm vi sai số lấy mẫu. Với lịch tuyến tính, ở bước cuối chỉ còn 0,64% biên độ tín hiệu, nên quá trình sinh xuất phát được từ nhiễu Gauss thuần tuý. SNR giảm dần theo bước, nên khi sinh, mô hình quyết định bố cục trước rồi mới thêm chi tiết. Huấn luyện là hồi quy dự đoán nhiễu, ổn định như mọi bài hồi quy. Lấy mẫu chậm được khắc phục bằng DDIM, chưng cất và khuếch tán trong không gian ẩn, còn classifier-free guidance làm mẫu bám điều kiện chặt hơn.

Với ba họ mô hình sinh, ta đã biết cách tạo ra nội dung mới. Nhưng một mô hình ngôn ngữ sinh văn bản trôi chảy chưa chắc làm đúng điều người dùng muốn, như ví dụ câu hỏi về thủ đô của Pháp ở Chương 1. Phần cuối giáo trình bàn về việc điều chỉnh hành vi đó, và công cụ đầu tiên cần có là học tăng cường, cách học từ phần thưởng thay vì từ đáp án đúng (Chương 12).