9Bộ tự mã hoá biến phân#

Bộ tự mã hoá nén dữ liệu tốt nhưng không sinh được mẫu mới. Bộ tự mã hoá biến phân (variational autoencoder, VAE; Kingma và Welling, 2014) sửa điểm này bằng cách ép không gian ẩn có một phân phối biết trước. Ràng buộc đó xuất hiện trong hàm mục tiêu dưới dạng một số hạng KL, và hệ số $\beta$ đứng trước số hạng này làm nhiều việc hơn cách nó thường được mô tả: như thí nghiệm ở Mục 9.3 cho thấy, nó quyết định mô hình dùng bao nhiêu chiều ẩn.

9.1Bộ tự mã hoá và hạn chế khi sinh mẫu#

Bộ tự mã hoá (autoencoder) gồm bộ mã hoá nén $x$ thành một vector ẩn $z$ ngắn hơn, và bộ giải mã dựng lại $\hat x$ từ $z$, huấn luyện để $\hat x$ gần $x$:

$$x \;\xrightarrow{\text{mã hoá}}\; z \;\xrightarrow{\text{giải mã}}\; \hat x, \qquad \min\; \|x - \hat x\|^2.$$

Với bộ mã hoá và giải mã tuyến tính, nghiệm tối ưu chính là PCA (Chương 14 của Nền tảng). Để sinh mẫu mới, cần lấy một $z$ ngẫu nhiên rồi giải mã. Nhưng lấy từ phân phối nào? Bộ tự mã hoá không ràng buộc hình dạng của tập các $z$ mà nó dùng: chúng có thể tụ thành những cụm rời rạc hoặc nằm trên một đường cong mảnh. Một điểm chọn ngẫu nhiên trong không gian ẩn nhiều khả năng rơi vào vùng bộ giải mã chưa từng thấy, và cho ra kết quả vô nghĩa. Vấn đề không phải nén kém mà là không gian ẩn không có cấu trúc để lấy mẫu.

9.2VAE và chặn dưới ELBO#

VAE xem $z$ là biến ẩn ngẫu nhiên với tiên nghiệm cố định $p(z) = \mathcal{N}(0, I)$ và bộ giải mã $p_\theta(x \mid z)$. Hợp lý của dữ liệu $p_\theta(x) = \int p_\theta(x \mid z)\,p(z)\,dz$ không tính được, nên VAE dùng thêm bộ mã hoá $q_\phi(z \mid x) = \mathcal{N}\big(\mu_\phi(x), \operatorname{diag}(\sigma^2_\phi(x))\big)$ để xấp xỉ hậu nghiệm $p_\theta(z \mid x)$.

Mệnh đề 9.1 (Chặn dưới của bằng chứng)

Với mọi phân phối $q(z \mid x)$,

$$\begin{aligned} \log p_\theta(x) &= \underbrace{\mathbb{E}_{q(z \mid x)}\big[\log p_\theta(x \mid z)\big] - \mathrm{KL}\big(q(z \mid x)\,\|\,p(z)\big)}_{\text{ELBO}} \\ &\quad + \mathrm{KL}\big(q(z \mid x)\,\|\,p_\theta(z \mid x)\big), \end{aligned}$$

và vì KL không âm, $\log p_\theta(x) \ge \text{ELBO}$.

Chứng minh

Theo định lý Bayes, $\log p_\theta(x) = \log p_\theta(x \mid z) + \log p(z) - \log p_\theta(z \mid x)$ với mọi $z$. Cộng và trừ $\log q(z \mid x)$ rồi lấy kỳ vọng theo $q(z \mid x)$ (vế trái không phụ thuộc $z$):

$$\begin{aligned} \log p_\theta(x) &= \mathbb{E}_q[\log p_\theta(x \mid z)] - \mathbb{E}_q\Big[\log\frac{q(z \mid x)}{p(z)}\Big] \\ &\quad + \mathbb{E}_q\Big[\log\frac{q(z \mid x)}{p_\theta(z \mid x)}\Big], \end{aligned}$$

trong đó hai kỳ vọng cuối chính là hai phân kỳ KL.

VAE tối đa ELBO theo cả $\theta$ và $\phi$. Thêm hệ số $\beta$ cho số hạng KL (Higgins và cộng sự, 2017) được hàm mục tiêu

$$\mathcal{L} = \underbrace{\mathbb{E}_{q_\phi(z \mid x)}\big[\log p_\theta(x \mid z)\big]}_{\text{tái dựng}} - \beta\,\underbrace{\mathrm{KL}\big(q_\phi(z \mid x)\,\|\,p(z)\big)}_{\text{kéo về tiên nghiệm}},$$

với $\beta = 1$ là VAE gốc.

Hai số hạng có ý nghĩa rõ ràng. Số hạng tái dựng yêu cầu giải mã từ $z$ phải dựng lại được $x$; với bộ giải mã Gauss, đó là bình phương sai số tái dựng với dấu trừ, cộng hằng số. Số hạng KL yêu cầu phân phối của $z$ ứng với mỗi $x$ gần $\mathcal{N}(0, I)$. Nếu mọi $q(z \mid x)$ đều gần tiên nghiệm, lấy mẫu $z \sim \mathcal{N}(0, I)$ sẽ rơi vào vùng bộ giải mã đã quen, và việc sinh mẫu trở nên hợp lệ. Với hai phân phối Gauss, số hạng KL có dạng đóng: mỗi chiều đóng góp $\tfrac12\big(\mu^2 + \sigma^2 - 1 - \ln\sigma^2\big)$, bằng 0 khi và chỉ khi $\mu = 0$, $\sigma = 1$. Chẳng hạn một chiều có $\mu = 1$ và $\sigma = 0{,}5$ đóng góp $\tfrac12(1 + 0{,}25 - 1 + 1{,}386) \approx 0{,}82$ nat: hậu nghiệm của nó lệch khỏi gốc một đơn vị và hẹp hơn tiên nghiệm, và đó là chi phí phải trả để chiều này mang thông tin về $x$. Hai số hạng kéo theo hai hướng ngược nhau, và $\beta$ điều chỉnh sự cân bằng giữa chúng.

Lưu ý (Tái tham số hoá)

Lấy mẫu $z \sim q_\phi(z \mid x)$ là phép ngẫu nhiên, và gradient không truyền qua phép lấy mẫu được. Cách giải quyết là viết $z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon$ với $\varepsilon \sim \mathcal{N}(0, I)$: phần ngẫu nhiên nằm ở $\varepsilon$, không phụ thuộc tham số, nên gradient truyền qua $\mu$ và $\sigma$ như bình thường. Đây là kỹ thuật cài đặt để huấn luyện VAE bằng gradient descent, không phải ý tưởng chính của VAE.

9.3Hệ số β và số chiều ẩn còn mang thông tin#

Lời giải thích thường gặp là "$\beta$ cân bằng giữa tái dựng và regularization". Thí nghiệm sau cho một mô tả cụ thể hơn. Dữ liệu gồm 600 điểm trong $\mathbb{R}^4$, sinh từ đúng 2 yếu tố ẩn Gauss qua một phép biến đổi tuyến tính, cộng nhiễu nhỏ. VAE được cho 6 chiều ẩn, nhiều hơn cần thiết, để xem nó làm gì với các chiều thừa. Bộ mã hoá và giải mã đều tuyến tính, nên VAE ở đây có liên hệ trực tiếp với PCA xác suất (Tipping và Bishop, 1999) và mọi hiện tượng đều giải thích được bằng công thức.

Một chiều ẩn được xem là còn mang thông tin nếu KL của riêng chiều đó, lấy trung bình trên dữ liệu, lớn hơn 0,01 nat: hậu nghiệm $q(z_i \mid x)$ còn khác tiên nghiệm $\mathcal{N}(0, 1)$. Chiều có KL bằng 0 có hậu nghiệm trùng tiên nghiệm với mọi $x$, tức không mang thông tin gì về $x$.

Ba hàng $\beta = 0$, $\beta = 1$ và $\beta = 16$ cho thấy rõ cơ chế. Ở $\beta = 0$, cả 6 chiều đều mang thông tin. Không có ràng buộc nào kéo hậu nghiệm về tiên nghiệm, nên mô hình dùng mọi chiều được cho. Đây là bộ tự mã hoá thường ở Mục 9.1: tái dựng tốt nhất (0,1219) nhưng không gian ẩn không có cấu trúc để lấy mẫu.

Ở $\beta = 1$, đúng 2 chiều mang thông tin, bằng số yếu tố thật sinh ra dữ liệu, còn bốn chiều thừa có KL bằng 0. Số hạng KL đã hoạt động như một cơ chế tự chọn số chiều: mỗi chiều mang thông tin phải trả một "chi phí" KL, và chỉ những chiều giảm sai số tái dựng nhiều hơn chi phí đó mới được giữ lại.

Ở $\beta = 16$, không chiều nào mang thông tin. Đây là hiện tượng sụp hậu nghiệm (posterior collapse): $z$ không còn liên quan tới $x$, và bộ giải mã chỉ dự đoán được giá trị trung bình của dữ liệu. Sai số tái dựng tăng lên 15,28, gấp khoảng 125 lần so với $\beta = 0$. Sụp hậu nghiệm cũng gặp khi bộ giải mã quá mạnh, ví dụ một mô hình tự hồi quy, vì khi đó bộ giải mã tự mô hình hoá được dữ liệu mà không cần $z$ (Bowman và cộng sự, 2016).

9.4Nguyên nhân mẫu của VAE bị mờ#

Với bộ giải mã Gauss, số hạng tái dựng là bình phương sai số (Mục 10.2 của Nền tảng), và hàm cực tiểu hoá kỳ vọng bình phương sai số là kỳ vọng có điều kiện:

$$\hat x(z) = \mathbb{E}[x \mid z].$$

Với ảnh, nếu cùng một $z$ ứng với nhiều ảnh thật khả dĩ, ví dụ cùng một khuôn mặt nhưng đường viền tóc lệch nhau vài điểm ảnh, thì trung bình của chúng là một ảnh nhoè ở các đường viền. Trung bình của nhiều ảnh sắc nét lệch nhau không phải một ảnh sắc nét. Lấy hai "ảnh" một chiều bốn điểm ảnh, $(0, 0, 1, 1)$ và $(0, 1, 1, 1)$, cùng ứng với một $z$ với xác suất như nhau: cùng một cạnh, lệch nhau một điểm ảnh. Dự đoán tối ưu theo bình phương sai số là trung bình của chúng, $(0;\ 0{,}5;\ 1;\ 1)$, có một điểm ảnh xám không xuất hiện trong ảnh thật nào. Sai số kỳ vọng của nó là $0{,}5^2 = 0{,}25$, trong khi dự đoán bằng một trong hai ảnh sắc nét có sai số kỳ vọng $\tfrac12 \cdot 0 + \tfrac12 \cdot 1 = 0{,}5$. Hàm mất mát bình phương vì vậy thưởng cho ảnh mờ. Độ mờ vì vậy không chủ yếu do mô hình thiếu dung lượng hay huấn luyện chưa đủ; nó đến từ việc chọn hàm mất mát bình phương, tức giả định nhiễu Gauss độc lập trên từng điểm ảnh. Đổi hàm mất mát thì đổi hiện tượng, và đó là điều GAN làm: thay hàm mất mát viết sẵn bằng một hàm mất mát học được.

9.5Tóm tắt#

Bộ tự mã hoá không sinh được mẫu vì không gian ẩn của nó không có cấu trúc để lấy mẫu. VAE xem $z$ là biến ẩn có tiên nghiệm $\mathcal{N}(0, I)$ và tối đa ELBO, một chặn dưới của log hợp lý gồm số hạng tái dựng và số hạng KL kéo hậu nghiệm về tiên nghiệm; mẹo tái tham số hoá cho phép huấn luyện bằng gradient descent. Thí nghiệm với dữ liệu sinh từ 2 yếu tố cho thấy $\beta$ quyết định số chiều ẩn còn mang thông tin: 6 chiều ở $\beta = 0$, đúng 2 chiều ở $\beta = 1$, và không chiều nào ở $\beta = 16$, khi hậu nghiệm sụp. Mẫu của VAE mờ vì hàm mất mát bình phương có nghiệm tối ưu là trung bình có điều kiện của nhiều ảnh sắc nét.

Nếu độ mờ đến từ một hàm mất mát viết sẵn, hướng sửa tự nhiên là để mô hình tự học hàm mất mát. Chương 10 trình bày GAN, mô hình làm đúng điều đó bằng một mạng thứ hai.