Auto-Encoding Variational Bayes

Hướng dẫn đọc

Vấn đề và động lực

Các mô hình xác suất có hướng với biến ẩn liên tục $\mathbf{z}$ rất hữu ích cho mô hình sinh (generative modeling), học biểu diễn không giám sát và suy luận biên. Tuy nhiên, việc học chúng thông qua phương pháp hợp lý cực đại hoặc các phương pháp Bayes gặp phải những trở ngại tính toán nghiêm trọng khi:

1. Hậu nghiệm bất khả quy (Posterior intractability): Đối với các mối quan hệ phi tuyến (ví dụ: mạng nơ-ron sâu), tích phân hợp lý biên $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ và phân phối hậu nghiệm $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ không có nghiệm giải tích dạng đóng, khiến thuật toán Expectation-Maximization (EM) chuẩn và mean-field variational Bayes không thể áp dụng.
2. Tập dữ liệu lớn: Các phương pháp lấy mẫu truyền thống như Monte Carlo EM đòi hỏi quá trình lấy mẫu MCMC lặp chậm chạp cho từng điểm dữ liệu trên toàn bộ tập dữ liệu, khiến việc tối ưu hóa mini-batch quy mô lớn trở nên quá tốn kém.

Ước lượng gradient hàm điểm (score-function gradient estimator) ngây thơ chuẩn ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ có phương sai cực kỳ cao, cản trở việc tối ưu hóa ổn định bằng phương pháp leo đồi gradient ngẫu nhiên (stochastic gradient ascent).

Ý tưởng cốt lõi

• Reparameterization Trick: Thay vì lấy mẫu $\mathbf{z}$ trực tiếp từ phân phối tham số hóa bất khả quy $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$, biến này được viết lại dưới dạng một phép biến đổi tất định của một biến nhiễu độc lập $\mathbf{ϵ}\sim p(\mathbf{ϵ})$: $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$. Điều này cho phép lấy đạo hàm theo $\mathbf{\varphi }$ thông qua lan truyền ngược chuẩn với phương sai giảm đi đáng kể.
• Amortized Inference (Recognition Model): Thay vì tối ưu hóa các tham số biến phân riêng biệt cho từng điểm dữ liệu, một mạng nơ-ron encoder $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ sẽ dự đoán các tham số biến phân (chẳng hạn như kỳ vọng và phương sai của một phân phối Gaussian ẩn) có điều kiện theo $\mathbf{x}$.
• Auto-Encoding Variational Bayes (AEVB): Tối ưu hóa đồng thời các tham số sinh $\mathbf{\theta }$ (probabilistic decoder) và các tham số biến phân $\mathbf{\varphi }$ (probabilistic encoder) bằng cách sử dụng stochastic gradient ascent trên chặn dưới biến phân đã được tái tham số hóa.

Phương pháp

Mục tiêu biến phân tối đa hóa Evidence Lower Bound (ELBO) $\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$: $$\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)})\ge \mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=-D_{KL}(q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})\parallel p_{\mathbf{\theta }}(\mathbf{z}))+{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})}\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|\mathbf{z})\right]$$

Áp dụng reparameterization trick ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ với ${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$, ước lượng SGVB tổng quát (Estimator A) có dạng: $${\tilde{\mathcal{L}}}^A(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=\frac{1}{L}\sum _{l=1}^L\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i,l)})-\log q_{\mathbf{\varphi }}({\mathbf{z}}^{(i,l)}|{\mathbf{x}}^{(i)})\right]$$

Khi tiên nghiệm $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ và hậu nghiệm xấp xỉ $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ đều là Gaussian, phân kỳ KL có thể tính toán giải tích, mang lại Estimator B: $$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})\simeq \frac{1}{2}\sum _{j=1}^J\left(1+\log (({\sigma }_j^{(i)}{)}^2)-({\mu }_j^{(i)}{)}^2-({\sigma }_j^{(i)}{)}^2\right)+\frac{1}{L}\sum _{l=1}^L\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|{\mathbf{z}}^{(i,l)})$$ trong đó ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ và ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$.

Đối với các mini-batch có kích thước $M$, ước lượng chặn dưới đầy đủ là: $${\tilde{\mathcal{L}}}^M(\mathbf{\theta },\mathbf{\varphi };{\mathbf{X}}^M)=\frac{N}{M}\sum _{i=1}^M\tilde{\mathcal{L}}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$$ Trong thực tế, việc đặt $L=1$ mẫu cho mỗi điểm dữ liệu là đủ khi kết hợp với kích thước mini-batch như $M=100$.

Kết quả

• Cài đặt mô hình: Các mô hình sinh được đánh giá trên MNIST (decoder Bernoulli MLP, 500 đơn vị ẩn) và Frey Face (decoder Gaussian MLP, 200 đơn vị ẩn), được tối ưu hóa bằng Adagrad với tốc độ học trong khoảng $\{0.01,0.02,0.1\}$ cùng weight decay.
• Tối ưu hóa chặn dưới: Trên nhiều số chiều biến ẩn khác nhau ($N_{\mathbf{z}}\in \{2,5,10,20,400\}$ trên Frey Face và MNIST), AEVB hội tụ nhanh hơn đáng kể và đạt chặn dưới cao hơn so với thuật toán wake-sleep.
• Khả năng chống overfitting: Tăng số chiều của $\mathbf{z}$ không gây ra overfitting nhờ tác dụng điều chuẩn của số hạng phân kỳ KL dạng giải tích hướng về phân phối tiên nghiệm.
• Ước lượng hợp lý biên: Trong các thử nghiệm không gian ẩn 3 chiều trên MNIST sử dụng bộ ước lượng hợp lý biên dựa trên HMC, AEVB hội tụ nhanh hơn wake-sleep và vượt xa tính khả thi của Monte Carlo EM (thuật toán quá chậm để chạy trên toàn bộ tập dữ liệu MNIST).
• Thời gian chạy: Quá trình huấn luyện mất khoảng 20–40 phút cho mỗi triệu mẫu trên một CPU Intel Xeon (~40 GFLOPS).

Hạn chế và các câu hỏi mở

• Biến ẩn liên tục: Reparameterization trick phụ thuộc trực tiếp vào tính khả vi đối với các tham số, khiến việc áp dụng trực tiếp cho các phân phối biến ẩn rời rạc gặp nhiều khó khăn.
• Khoảng cách xấp xỉ (Approximation Gap): Việc giả định phân phối Gaussian đường chéo cho $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ giới hạn năng lực biểu diễn khi hậu nghiệm thực tế là đa mốt hoặc có tương quan chặt chẽ.
• Mở rộng sang Full Variational Bayes: Đánh giá thực nghiệm chủ yếu tập trung vào ước lượng tham số MAP/ML cho $\mathbf{\theta }$; suy luận biến phân trên các tham số toàn cục $\mathbf{\theta }$ đã được xây dựng công thức trong phần phụ lục nhưng để dành cho việc kiểm chứng thực nghiệm trong tương lai.