Ở Phần 1, chúng ta đã xem xét một ví dụ đơn giản của suy diễn biến phân. Trong bài viết ngày hôm nay, chúng ta sẽ xem xét một ví dụ phức tạp hơn, đòi hỏi các giả định mạnh hơn để giải quyết.
Một bóng đèn có tuổi thọ thực tế \(Z\) là một biến ngẫu nhiên tuân theo phân phối lũy thừa với hằng số \(\Lambda\). Khác với phần trước, ở đây \(\Lambda\) là một biến ngẫu nhiên tuân theo phân phối Gamma với hằng số \(\alpha > 0\) và \(\beta > 0\). Khi sử dụng, tuổi thọ thực nghiệm \(X\) của bóng đèn có tồn tại một nhiễu ngẩu nhiên với phân phối chuẩn (Gaussian) có kỳ vọng \(0\) và phương sai \(\sigma^2 \in (0, \infty)\). Giả sử ta có một bóng đèn với tuổi thọ thực nghiệm là \(x > 0\), ước lượng tuổi thọ thực tế của nó.
Mô hình sinh của tuổi thọ thực nghiệm là một mô hình phân tầng (hierarchical model) được mô tả như sau
\[\begin{align*} \Lambda &\sim \text{Gamma}(\alpha, \beta) \\ Z|\Lambda = \lambda &\sim \text{Exp}(\lambda) \\ X|Z = z &\sim \text{N}(z, \sigma^2) \end{align*}\]Dựa trên mô hình này, phân phối kết hợp (joint distribution) của tuổi thọ thực nghiệm và tuổi thọ thực tế là
\[p_{X, Z, \Lambda}(x, z, \lambda) = p_{X|z}(x|z) p_{Z|\lambda}(z) p_{\Lambda}(\lambda) = \mathcal{N}(x; z, \sigma^2)\text{Exp}(z; \lambda) \text{Gamma}(\lambda; \alpha, \beta)\]trong đó
\[\begin{align*} \mathcal{N}(x; z, \sigma^2) &= \dfrac{1}{\sqrt{2\pi\sigma^2}} \exp\left[-\dfrac{(x - z)^2}{2\sigma^2}\right] \\ \text{Exp}(z; \lambda) &= \lambda \exp(-\lambda z) \mathbb{1}_{[0, \infty)}(z) \\ \text{Gamma}(\lambda; \alpha, \beta) &= \dfrac{\beta^\alpha}{\Gamma(\alpha)} \lambda^{\alpha - 1} \exp(-\beta\lambda) \mathbb{1}_{(0, \infty)}(\lambda) \end{align*}\]Tương tự như ví dụ trước, ta mong muốn tính được phân phối hậu nghiệm kết hợp (joint posterior distribution) của cả tham số tỉ lệ \(\Lambda\) lẫn tuổi thọ thực tế \(Z\) dựa trên dữ liệu quan sát được:
\[p_{Z, \Lambda|x}(z, \lambda) = \dfrac{p_{X, Z, \Lambda}(x, z, \lambda)}{p_X(x)}\]Việc phải ước lượng cả hai tham số \(\Lambda\) và \(Z\) cùng lúc khiến cho bài toán trở nên khó khăn hơn. Không những thế, hai biến ngẫu nhiên này còn không độc lập! Điều này đồng nghĩa với việc ta không thể phân rã phân phối hậu nghiệm kết hợp thành các thành phần không phụ thuộc vào nhau, và việc tối ưu cũng không còn dễ dàng nữa.
Suy diễn biến phân với trường trung bình (mean-field variational inference) là một trong những cách tiếp cận phổ biến trong các trường hợp như vậy.
Ý tưởng chính của phương pháp này là giả định họ phân phối biến phân (variational distribution family) là họ trường trung bình (mean-field family). Đây là họ các phân bố phân rã được (factorizable). Nói cách khác, các biến ngẫu nhiên được giả định là độc lập với nhau, mỗi biến tuân theo một phân phối phần tử (factor) riêng. Cụ thể, các phân phối trong thuộc họ trường trung bình trên không gian vector ngẫu nhiên \(m\)-chiều \(\vec{\Theta} = (\Theta_1, \dots, \Theta_m)\) có dạng
\[q_{\vec{\Theta}}(\vec{\theta}) = \prod_{i=1}^m q_{\Theta_i}(\theta_i)\]Hãy nhớ lại về định nghĩa ELBO:
\[\text{ELBO}(q_{\vec{\Theta}}, p_{X, \vec{\Theta}}) = \mathbb{E}_{q_{\vec{\Theta}}}\left[ \log p_{X, \vec{\Theta}}(x, \vec{\Theta}) - \log q_{\vec{\Theta}}(\vec{\Theta}) \right]\]Trong đó, phân phối kết hợp có thể được phân tách như sau
\[p_{X, \vec{\Theta}}(x, \vec{\Theta}) = p_{X}(x) p_{\vec{\Theta}|x}(\vec{\Theta}) % = p_{X}(x) \prod_{j = 1}^{m} p_{\Theta_j|x, \Theta_{1:j-1}}(\Theta_j)\]và do đó
\[\begin{align*} \mathbb{E}_{q_{\vec{\Theta}}}\left[ \log p_{X, \vec{\Theta}}(x, \vec{\Theta}) \right] & = \log p_{X}(x) + \mathbb{E}_{q_{\vec{\Theta}}} \left[ \log p_{\vec{\Theta}|x}(\vec{\Theta}) \right] \end{align*}\]Với việc sử dụng họ phân phối trường trung bình,
\[\begin{align*} \mathbb{E}_{q_{\vec{\Theta}}}\left[ \log q_{\vec{\Theta}}(\vec{\Theta}) \right] & = \int q_{\vec{\Theta}}(\vec{\theta}) \log q_{\vec{\Theta}}(\vec{\theta}) d\vec{\theta} \\ & = \int \left( \prod_{j=1}^m q_{\Theta_j}(\theta_j) \right) \left( \sum_{j = 1}^{m} \log q_{\Theta_j}(\theta_j) \right) d\vec{\theta} \\ & = \sum_{j = 1}^{m} \int \left( \prod_{k = 1}^{m} q_{\Theta_k}(\theta_k) \right) \log q_{\Theta_j}(\theta_j) d\vec{\theta} \\ & = \sum_{j = 1}^{m} \left( \int \prod_{k \neq j} q_{\Theta_k}(\theta_k) d\vec{\theta}_{-j} \right) \left( \int q_{\Theta_j}(\theta_j) \log q_{\Theta_j}(\theta_j) d\theta_j \right) \\ & = \sum_{j = 1}^{m} \mathbb{E}_{q_{\Theta_j}}\left[ \log q_{\Theta_j}(\Theta_j) \right] \end{align*}\]Kết hợp 2 điều trên, ELBO trở thành
\[\begin{align*} \text{ELBO}(q_{\vec{\Theta}}, p_{X, \vec{\Theta}}) = \log p_{X}(x) + \mathbb{E}_{q_{\vec{\Theta}}} \left[ \log p_{\vec{\Theta}|x}(\vec{\Theta}) \right] - \sum_{j=1}^{m} \mathbb{E}_{q_{\Theta_j}} \left[ \log q_{\Theta_j}(\Theta_j) \right] \end{align*}\]Một điểm đáng lưu ý ở đây là với việc sử dụng họ phân phối trường trung bình, tại một tọa độ \(k \in [m]\) bất kì,
\[\begin{align*} \mathbb{E}_{q_{\vec{\Theta}}} \left[ \log p_{\vec{\Theta}|x}(\vec{\Theta}) \right] &= \mathbb{E}_{q_{\Theta_k}} \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-j}}(\Theta_k) + \log p_{\vec{\Theta}_{-k}|x}(\vec{\Theta}_{-k}) \right] \\ & = \mathbb{E}_{q_{\Theta_k}} \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\Theta_k) \right] + \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\vec{\Theta}_{-k}|x}(\vec{\Theta}_{-k}) \right] \end{align*}\]Như vậy, tại một tọa độ bất kì, ELBO có thể được tách được thành
\[\begin{align*} \text{ELBO}(q_{\Theta}, p_{X, \vec{\Theta}}) & = \log p_X(x) \\ & + \mathbb{E}_{q_{\Theta_k}} \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\Theta_k) \right] + \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\vec{\Theta}_{-k}|x}(\vec{\Theta}_{-k}) \right] \\ & - \mathbb{E}_{q_{\Theta_k}} \left[ \log q_{\Theta_k}(\Theta_k) \right] - \sum_{j \neq k} \mathbb{E}_{q_{\Theta_j}} \left[ \log q_{\Theta_j}(\Theta_j) \right] \end{align*}\]Việc có thể viết được ELBO dưới dạng trên cho thấy rằng việc tìm \(q_{\vec{\Theta}}\) tối ưu ELBO có thể được thực hiện thông qua việc tối ưu tuần tự từng thành phần khi giữ cố định các thành phần còn lại. Phương pháp này được gọi là phương pháp gia tăng từng tọa độ (coordinate ascent).
Xét tọa độ \(k \in [m]\). Khi giữ cố định \(q_{\Theta_{-k}}\), ta có thể tìm \(q_{\Theta_k}\) để tối ưu ELBO bằng cách xét đạo hàm bằng không.
ELBO khi chỉ xét các thành phần liên quan đến \(q_{\Theta_k}\) là
\[\text{ELBO}(q_{\Theta_k}; q_{\vec{\Theta}}, p_{X, \vec{\Theta}}) \stackrel{c}{=} \mathbb{E}_{q_{\Theta_k}} \left\{ \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\Theta_k) \right] - \log q_{\Theta_k}(\Theta_k) \right\}\]Đạo hàm theo \(q_{\Theta_k}(\theta_k)\) là
\[\dfrac{\partial}{\partial q_{\Theta_k}(\theta_k)} \text{ELBO}(q_{\Theta_k}; q_{\vec{\Theta}}, p_{X, \vec{\Theta}}) = \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\theta_k) \right] - \log q_{\Theta_k}(\theta_k) - 1\]Đặt đạo hàm bằng không, ta có
\[\log q^\star_{\Theta_k}(\theta_k) \stackrel{c}{=} \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\theta_k) \right]\]hay
\[q^\star_{\Theta_k}(\theta_k) \propto \exp\left\{ \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{\Theta_k|x, \vec{\Theta}_{-k}}(\theta_k) \right] \right\}\]\(p_{\Theta_k\\|x, \vec{\Theta}_{-k}}\) được gọi là phân phối điều kiện hoàn toàn (complete conditional) của biến \(\Theta_k\). Trong một số trường hợp, ta có thể giả định thêm về nó để thuận tiện cho việc tính toán. Tuy nhiên, việc này nằm quá giới hạn của bài viết này và sẽ được thảo luận ở phần sau (nếu có).
Ở đây, ta tiếp cận nó theo một hướng khác, sử dụng việc
\[p_{\Theta_k|x, \vec{\Theta}_{-k}}(\theta_k) = \dfrac{ p_{X, \vec{\Theta}}(x, \theta_k, \vec{\Theta}_{-k}) }{ p_{X, \vec{\Theta}_{-k}}(x, \vec{\Theta}_{-k}) }\]với \(p_{X, \vec{\Theta}_{-k}}(x, \vec{\Theta}_{-k})\) không phụ thuộc vào \(\theta_k\), ta có
\[q^\star_{\Theta_k}(\theta_k) \propto \exp\left\{ \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{X, \vec{\Theta}}(x, \theta_k, \vec{\Theta}_{-k}) \right] \right\}\]Trong một số trường hợp, ta có thể tính được kỳ vọng này một cách dễ dàng.
Tóm lại, thuật toán gia tăng từng tọa độ để tối ưu ELBO khi sử dụng họ phân phối trường trung bình, gọi tên là CAVI (Coordinate Ascent Variational Inference), có thể được mô tả như sau:
Lặp lại cho đến khi hội tụ:
Với mỗi \(k = 1, \dots, m\), cập nhật
\[q_{\Theta_k}(\theta_k) = \exp\left\{ \mathbb{E}_{q_{\vec{\Theta}_{-k}}} \left[ \log p_{X, \vec{\Theta}}(x, \theta_k, \vec{\Theta}_{-k}) \right] \right\}\]Chuẩn hóa \(q_{\Theta}\)
Trong trường hợp này, ta giả định họ phân phối trường trung bình cho cả hai biến ngẫu nhiên \(\Lambda\) và \(Z\):
\[q_{Z, \Lambda}(z, \lambda) = q_{Z}(z) q_{\Lambda}(\lambda)\]Một điều đặc biệt là ta không hề giả định gì về từng thành phần \(q_{Z}\) và \(q_{\Lambda}\) cả, khác với phần trước khi ta giả định \(q_{Z}\) là một phân phối chuẩn.
Theo như thuật toán trên, các cập nhật sẽ là
\[\begin{align*} q_{Z}(z) &\propto \exp\left\{ \mathbb{E}_{q_{\Lambda}} \left[ \log p_{X, Z, \Lambda}(x, z, \Lambda) \right] \right\} \\ q_{\Lambda}(\lambda) &\propto \exp\left\{ \mathbb{E}_{q_{Z}} \left[ \log p_{X, Z, \Lambda}(x, Z, \lambda) \right] \right\} \end{align*}\]Phân phối kết hợp Để tính được bước cập nhật cho từng tọa độ, ta cần biến đổi phân phối kết hợp \(p_{X, Z, \Lambda}\).
\[\begin{align*} p_{X, Z, \Lambda}(x, z, \lambda) &= p_{X|Z, \Lambda}(x|z, \lambda) p_{Z|\Lambda}(z|\lambda) p_{\Lambda}(\lambda) \\ &= \mathcal{N}(x; z, \sigma^2) \text{Exp}(z; \lambda) \text{Gamma}(\lambda; \alpha, \beta) \\ &\propto \exp\left[ -\dfrac{1}{2\sigma^2} (x - z)^2 \right] \lambda \exp(- \lambda z) \lambda^{\alpha - 1} \exp(- \beta\lambda) \end{align*}\]Từ đó, ta có thể tính được
\[\begin{align*} \log p_{X, Z, \Lambda}(x, z, \lambda) &\stackrel{c}{=} -\dfrac{1}{2\sigma^2} (x - z)^2 - \lambda z + \alpha \log \lambda - \beta \lambda \end{align*}\]Cập nhật của \(q_{Z}\) Lược bỏ đi các thành phần không liên quan đến \(z\), ta có
\[\begin{align*} \mathbb{E}_{q_{\Lambda}} \left[ \log p_{X, Z, \Lambda}(x, z, \Lambda) \right] &\stackrel{c}{=} -\dfrac{1}{2\sigma^2} (x - z)^2 - \mathbb{E}_{q_{\Lambda}}[\Lambda] z \\ &\stackrel{c}{=} -\dfrac{1}{2\sigma^2} z^2 + \left( \dfrac{x}{\sigma^2} - \mathbb{E}_{q_{\Lambda}}[\Lambda] \right) z \\ & = -\dfrac{1}{2\sigma^2} \left[ z^2 - 2(x - \sigma^2 \mathbb{E}_{q_{\Lambda}}[\Lambda])z \right] \\ &\stackrel{c}{=} -\dfrac{1}{2\sigma^2} \left[ z - (x - \sigma^2 \mathbb{E}_{q_{\Lambda}}[\Lambda]) \right]^2 \end{align*}\]Như vậy, cập nhật cho \(q_{Z}\) sẽ là
\[\begin{align*} q_{Z}(z) &\propto \exp\left\{ -\dfrac{1}{2\sigma^2} \left[ z - (x - \sigma^2 \mathbb{E}_{q_{\Lambda}}[\Lambda]) \right]^2 \right\} \end{align*}\]Nếu để ý, ta sẽ nhận ra rằng
\[q_{Z}(z) = \mathcal{N}(z; x - \sigma^2 \mathbb{E}_{q_{\Lambda}}[\Lambda], \sigma^2)\]Ở đây, \(\mathbb{E}_{q_{\Lambda}}[\Lambda]\) đang được bỏ ngõ, do ta chưa có đủ thông tin về \(q_{\Lambda}\). Ta sẽ tính nó ngay sau đây.
Cập nhật của \(q_{\Lambda}\) Lược bỏ đi các thành phần không liên quan đến \(\lambda\), ta có
\[\begin{align*} \mathbb{E}_{q_{Z}} \left[ \log p_{X, Z, \Lambda}(x, Z, \lambda) \right] &\stackrel{c}{=} \alpha \log \lambda - (\beta + \mathbb{E}_{q_{Z}}[Z]) \lambda \end{align*}\]Như vậy, cập nhật cho \(q_{\Lambda}\) sẽ là
\[\begin{align*} q_{\Lambda}(\lambda) &\propto \exp\left\{ \alpha \log \lambda - (\beta + \mathbb{E}_{q_{Z}}[Z]) \lambda \right\} = \lambda^{\alpha} \exp(- (\beta + \mathbb{E}_{q_{Z}}[Z]) \lambda) \end{align*}\]Nếu để ý, ta sẽ lại nhận ra rằng
\[q_{\Lambda}(\lambda) = \text{Gamma}(\lambda; \alpha + 1, \beta + \mathbb{E}_{q_{Z}}[Z])\]Lại một lần nữa, \(\mathbb{E}_{q_{Z}}[Z]\) đang được bỏ ngỏ. Tuy nhiên, dựa vào kết quả ở trên, ta biết \(q_{Z}\) là một phân phối chuẩn, và do đó \(\mathbb{E}_{q_{Z}}[Z]\) có thể được tính toán một cách dễ dàng. Song song đó, ta cũng có thể tính \(\mathbb{E}_{q_{\Lambda}}[\Lambda]\) một cách dễ dàng.
Tóm lại, thuật toán để giải quyết bài toán đầu đề có thể được mô tả như sau:
Với mỗi \(t = 0, 1, \dots\):
Cập nhật
\[\mu_{Z}^{(t + 1)} = x - \sigma^2 \dfrac{\alpha_{\Lambda}^{(t)}}{\beta_{\Lambda}^{(t)}}\]và theo đó
\[q_{Z}^{(t + 1)}(z) = \mathcal{N} \left( z; \mu_{Z}^{(t + 1)}, \sigma^2 \right)\]Cập nhật
\[\begin{align*} \alpha_{\Lambda}^{(t + 1)} &= \alpha_{\Lambda}^{(t)} + 1 \\ \beta_{\Lambda}^{(t + 1)} &= \beta_{\Lambda}^{(t)} + \mu_{Z}^{(t + 1)} \end{align*}\]và theo đó
\[q_{\Lambda}^{(t + 1)}(\lambda) = \text{Gamma}\left(\lambda; \alpha_{\Lambda}^{(t + 1)}, \beta_{\Lambda}^{(t + 1)}\right)\]Nếu quan tâm, mọi người có thể thử tại đây để so sánh giữa phân phối hậu nghiệm chính xác và phân phối biến phân tại các bước của thuật toán.
Hình trên vẽ mặt cắt \(h = 0.005\) của các phân phối. Trong đó, theo thứ tự màu đen, đỏ, xanh lam, xanh lục, cam, tím lần lượt là phân phối hậu nghiệm chính xác và các phân phối biến phân theo CAVI tại các bước từ 1 đến 5.
Có thể thấy giả định trường trung bình đã giới hạn khả năng xấp xỉ của phân phối biến phân. Bên cạnh đó, do ta đang tối ưu phân kỳ KL \(\text{KL}(q_{Z, \Lambda} \| p_{Z, \Lambda\\|x})\), phân phối biến phân không cố gắng che phủ (như trong trường hợp tối ưu phân kỳ KL hướng ngược lại) mà thay vào đó muốn nằm gọn trong phân phối hậu nghiệm chính xác.
Trong bài viết này, chúng ta đã thảo luận về một ví dụ phức tạp hơn của suy diễn biến phân, và cách tiếp cận nó thông qua phương pháp suy diễn biến phân với trường trung bình. Bài viết này có thể sẽ được cập nhật để thêm phần cài đặt và mô phỏng kết quả.
Trong các phần tiếp theo (nếu có), chúng ta sẽ xem xét một số ví dụ khác, đôi chút phức tạp hơn, nhưng vẫn có thể giải quyết được bằng phương pháp này. Hy vọng được gặp lại!