Suy diễn Biến phân (Phần 1)

Cập nhật (13/03/2024) Thay thế họ phân phối biến phân sử dụng trong ví dụ từ phân phối lũy thừa sang phân phối chuẩn.

Bối cảnh

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 (exponential distribution) với hằng số \(\lambda > 0\). Tuy nhiên, thực tế 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 bằ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 \(x > 0\), ước lượng tuổi thọ thực tế của nó.

Mô hình

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*} Z &\sim \text{Exp}(\lambda) \\ X | Z=z &\sim \mathcal{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}(x, z) = p_{X|z}(x|z) p_{Z}(z) = \mathcal{N}(x; z, \sigma^2)\text{Exp}(z; \lambda)\]

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) \end{align*}\]

Việc ước lượng tuổi thọ thực tế có thể xem như một dạng của bài toán suy diễn Bayes (Bayesian inference). Ở đây, chúng ta muốn tính phân phối hậu nghiệm (posterior distribution) của tuổi thọ thực tế dựa trên (các) quan sát về tuổi thọ thực nghiệm. Tức là, chúng ta muốn tính

\[p_{Z|x}(z) = \dfrac{p_{X, Z}(x, z)}{p_X(x)} = \dfrac{p_{X, Z}(x, z)}{\int p_{X, Z}(x, z') dz'}\]

trong đó ta đã dùng công thức Bayes và xác suất toàn phần.

Trong bài viết này, chúng ta sẽ thảo luận về một số cách để tính phân phối hậu nghiệm này.

Đặt vấn đề

Trong nhiều trường hợp (đặc biệt là trong ứng dụng thực tế), phân phối hậu nghiệm không thể tính toán được một cách trực tiếp dễ dàng, trừ khi chúng ta có một công thức dạng đóng (closed-form) cho nó.

Hãy thử ngay trong tình huống này:

\[\begin{align*} p_{Z|x}(z) & \propto p_{X, Z}(x, z) \\ & = \mathcal{N}(x; z, \sigma^2)\text{Exp}(z; \lambda) \\ & \propto \exp \left[ -\dfrac{1}{2\sigma^2} (x - z)^2 \right] \exp(-\lambda z) \mathbb{1}_{[0, \infty)}(z) \\ & = \exp \left[ - \dfrac{1}{2\sigma^2} \left( x^2 - 2(x - \sigma^2 \lambda)z + z^2 \right) \right] \mathbb{1}_{[0, \infty)}(z) \\ & \propto \exp \left[ - \dfrac{1}{2\sigma^2} \left( (x - \sigma^2 \lambda)^2 - 2(x - \sigma^2 \lambda)z + z^2 \right) \right] \mathbb{1}_{[0, \infty)}(z) \\ & \propto \exp \left[ - \dfrac{1}{2\sigma^2} \left( x - \sigma^2 \lambda - z \right)^2 \right] \mathbb{1}_{[0, \infty)}(z) \end{align*}\]

Nói cách khác, phân phối hậu nghiệm trong trường hợp này là một dạng phân phối chuẩn cụt (truncated normal distribution), vốn không dễ dàng tính toán.

Bên lề: Xích Markov Monte-Carlo

Một trong những cách tiếp cận phổ biến nhất để giải quyết vấn đề này là sử dụng xích Markov Monte-Carlo (Monte-Carlo Markov Chain). Ý tưởng chính của MCMC là tạo ra một xích Markov sao cho phân phối ổn định (stationary distribution) của xích này là phân phối hậu nghiệm cần tìm. Một trong những thuật toán MCMC phổ biến nhất là thuật toán Metropolis-Hastings.

Một điểm cần lưu ý là MCMC đòi hỏi nhiều thời gian và tài nguyên tính toán: một xích Markov phải được tạo ra và chạy đủ lâu để đảm bảo rằng phân phối ổn định của nó là phân phối hậu nghiệm cần tìm. Tuy nhiên, khi đã tiệm cận phân phối ổn định thì MCMC lại cho ra những kết quả với độ chính xác cao. Điều này có nghĩa là thuật toán này không phù hợp cho những tình huống mà đòi hỏi tính toán nhanh và có thể chấp nhận được sai số.

Trong khuôn khổ của bài viết này, chúng ta sẽ không đi sâu vào MCMC mà thay vào đó sẽ xem xét một phương pháp khác nhanh hơn, tuy không chính xác bằng. Phương pháp này phù hợp để đưa ra một xấp xỉ ban đầu để có hình dung sơ khởi về phân phối hậu nghiệm, để tử từ cải thiện bằng các phương pháp khác có độ chính xác cao hơn nếu cần.

Suy diễn biến phân

Suy diễn biến phân (variational inference) là một phương pháp để ước lượng phân phối hậu nghiệm bằng cách tìm một phân phối xấp xỉ phân phối hậu nghiệm. Cụ thể hơn, ý tưởng chính của phương pháp này là giải bài toán tối ưu biến phân (variational optimization): tìm phân phối \(q_{Z}\) thuộc một họ phân phối \(\mathcal{Q}\) sao cho gần nhất với phân phối hậu nghiệm \(p_{Z\\|x}\) theo một tiêu chí nào đó.

Tiêu chí Kullback-Leibler

Một trong những tiêu chí phổ biến nhất là phân kỳ Kullback-Leibler (KL divergence):

\[\text{KL}(q_{Z} \| p_{Z\\|x}) = \int q_{Z}(z) \log \dfrac{q_{Z}(z)}{p_{Z\\|x}(z)} dz\]

Nói cách khác, bài toán suy diễn Bayes trở thành bài toán tối ưu sau:

\[q_{Z}^* = \arg\min_{q_{Z} \in \mathcal{Q}} \text{KL}(q_{Z} \| p_{Z\\|x})\]

Tuy nhiên, nếu để ý thì để tính được phân kỳ KL, ta cần biết phân phối hậu nghiệm \(p_{Z\\|x}\), quay lại chính vấn đề mà chúng ta đang cố gắng giải quyết. Như vậy, tại sao mọi người lại quan tâm đến phân kỳ KL? Thực tế, chúng ta có thể biến đổi phân kỳ KL thành một dạng có thể tính toán được:

\[\begin{align*} \text{KL}(q_{Z} || p_{Z\\|x}) &= \int q_{Z}(z) \log \dfrac{q_{Z}(z)}{p_{Z\\|x}(z)} dz \\ &= \int q_{Z}(z) \log \dfrac{q_{Z}(z) p_X(x) }{p_{X, Z}(x, z)} dz \\ &= \int q_{Z}(z) \log \dfrac{q_{Z}(z)}{p_{X, Z}(x, z)} dz + \log p_X(x) \\ &=: - \text{ELBO}(q_{Z}, p_{X, Z}) + \log p_X(x) \end{align*}\]

trong đó, ta đã định nghĩa hàm ELBO (Evidence Lower BOund) như sau:

\[\begin{align*} \text{ELBO}(q_{Z}, p_{X, Z}) &= \int q_{Z}(z) \log \dfrac{p_{X, Z}(x, z)}{q_{Z}(z)} dz \\ &= \mathbb{E}_{q_{Z}}\left[ \log p_{X, Z}(x, Z) - \log q_{Z}(Z) \right] \end{align*}\]

Do \(\log p_X(x)\) không phụ thuộc vào \(q_{Z}\), nên việc tìm \(q_{Z}\) để cực tiểu hóa phân kỳ KL cũng tương đương với việc cực đại hóa hàm ELBO trên. Trong nhiều trường hợp, tuy vẫn là một tích phân, nhưng tích phân này mang ý nghĩa của một kỳ vọng, nên có thể tính toán được một cách dễ dàng hơn.

Họ phân phối

Để tìm một phân phối trong tất cả các phân phối có thể thì sẽ là một khó khăn khác. Do đó, ta muốn giới hạn họ phân phối \(\mathcal{Q}\) sao cho việc tìm tìm kiếm trở nên dễ dàng hơn. Một cách giới hạn phổ biến là cân nhắc các họ phân phối được tham số hóa (parametric distribution). Như vậy phân phối biến phân \(q_{Z}\) được giả định là thuộc một họ phân phối được tham số hóa bởi một số tham số \(\phi\), ký hiệu \(\mathcal{Q}_{\phi}\).

Từ đây, có thể viết phân phối biến phân dưới dạng \(q_{Z}(z; \phi)\). Bài toán tối ưu trở thành

\[\phi^* = \arg\max_{\phi} \text{ELBO}(\phi; q_{Z}, p_{X, Z})\]

Giải quyết bài toán

Như vậy, một cách tiếp cận bài toán đầu đề là chọn một họ phân phối được tham số hóa và giải quyết bài toán tối ưu. Trong trường hợp này, việc phân phối hậu nghiệm là một phân phối chuẩn cụt gợi ý ta chọn họ phân phối biến phân là phân phối chuẩn với tham số \(\mu\) và \(s^2 \in (0, \infty)\).

\[q_{Z}(z; \mu, s^2) = \mathcal{N}(z; \mu, s^2) = \dfrac{1}{\sqrt{2\pi s^2}} \exp\left[ -\dfrac{1}{2s^2} (z - \mu)^2 \right]\]

Việc còn lại là tìm các tham số biến phân tối ưu sao cho hàm ELBO đạt cực đại. Với giả định này, hàm ELBO trở thành

\[\begin{align*} &\text{ELBO}(\mu, s^2; q_{Z}, p_{X, Z}) \\ &= \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)}\left[ \log p_{X, Z}(x, Z) - \log q_{Z}(Z; \mu, s^2) \right] \\ &= \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)}\left[ \log \mathcal{N}(x; Z, \sigma^2) + \log \text{Exp}(Z; \lambda) - \log \mathcal{N}(Z; \mu, s^2) \right] \\ &\stackrel{c}{=} \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)}\left[ - \dfrac{1}{2\sigma^2}(x - Z)^2 - \lambda Z + \dfrac{1}{2} \log s^2 + \dfrac{1}{2s^2} (Z - \mu)^2 \right] \\ &\stackrel{c}{=} \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)}\left[ -\dfrac{1}{2\sigma^2} Z^2 + \dfrac{x}{\sigma^2} Z - \lambda Z + \dfrac{1}{2} \log s^2 + \dfrac{1}{2s^2} Z^2 - \dfrac{1}{s^2} \mu Z + \dfrac{1}{2s^2} \mu^2 \right] \end{align*}\]

trong đó \(\stackrel{c}{=}\) là chênh lệch hằng số cộng vào (additive constant). Ta có thể làm thế do khi giải bài toán tối ưu, ta chỉ quan tâm đến các thành phần có liên quan đến biến tối ưu.

Vì \(q_{Z}\) là phân phối chuẩn với tham số \(\mu\) và \(s^2\), nên

\[\begin{align*} \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)} & Z = \mu \\ \mathbb{E}_{q_{Z}(\cdot; \mu, s^2)} & Z^2 = s^2 + \mu^2 \end{align*}\]

Do đó, hàm ELBO trở thành

\[\begin{align*} \text{ELBO}(\mu, s^2; q_{Z}, p_{X, Z}) \stackrel{c}{=} -\dfrac{1}{2\sigma^2} (s^2 + \mu^2) + \left( \dfrac{x}{\sigma^2} - \lambda \right) \mu + \dfrac{1}{2} \log s^2 \end{align*}\]

Đạo hàm của hàm ELBO theo \(\mu\) và \(\sigma^2\) là

\[\begin{align*} \dfrac{\partial}{\partial \mu} \text{ELBO}(\mu, s^2; q_{Z}, p_{X, Z}) &= -\dfrac{1}{\sigma^2} \mu + \left( \dfrac{x}{\sigma^2} - \lambda \right) \\ \dfrac{\partial}{\partial s^2} \text{ELBO}(\mu, s^2; q_{Z}, p_{X, Z}) &= -\dfrac{1}{2\sigma^2} + \dfrac{1}{2} \dfrac{1}{s^2} \end{align*}\]

Giải đạo hàm bằng 0, ta có

\[\begin{align*} \mu^\star &= x - \sigma^2 \lambda \\ (s^\star)^2 &= \sigma^2 \end{align*}\]

Như vậy, ta đã tìm được tham số biến phân tối ưu và từ đó, phân phối biến phân tối ưu và có thể sử dụng nó thay thế cho phân phối hậu nghiệm khi cần suy luận gì với dữ liệu quan sát.

\[p_{Z\\|x}(z) \approx q_{Z}(z; \mu^\star, (s^\star)^2) = \mathcal{N}(z; \mu^\star, (s^\star)^2)\]

Một điểm đáng lưu ý là tham số biến phân tối ưu này phụ thuộc vào dữ liệu quan sát \(x\), vốn dễ hiểu vì chính bản thân phân phối hậu nghiệm cũng phụ thuộc vào dữ liệu này. Ngoài ra, mức độ phù hợp của xấp xỉ phụ thuộc lớn vào họ phân phối được chọn.

Minh họa

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 ưu.

Trong hình trên, \(\lambda = 0.2, \sigma^2 = 0.16, x = 0.5\), đường xanh lá tương ứng với phân phối hậu nghiệm chính xác, đường đỏ tương ứng với phân phối biến phân tối ưu.

Hứa hẹn

Đây chỉ là một ví dụ cực kỳ đơn giản (và cũng khá là “nhân tạo”) để mọi người có thể hình dung ý tưởng cũng như cách hoạt động của các thuật toán xấp xỉ dựa trên biến phân. Trong thực tế, việc tìm phân phối biến phân tối ưu có thể phức tạp hơn rất nhiều, và cần phải sử dụng các phương pháp tối ưu phức tạp hơn. Ở các phần tiếp theo (nếu có), mình sẽ thảo luận về các tình huống phức tạp hơn và các giải pháp trong những trường hợp đó (thường sẽ yêu cầu nhiều giả định hơn).

Đôi lời bên lề

Mình chợt nhận ra là cứ mỗi năm thì mình lại có (ít nhất) một bài viết. Không biết nếu giờ mà đọc lại thì có cảm nhận được sự thay đổi trong cách hành văn của mình không, vì mình có cảm giác cứ mỗi năm trôi qua thì khả năng viết của mình lại ngày càng đi xuống. Một phần vì dạo này mình dùng tiếng Anh cũng nhiều lên (nhưng không có nghĩa là khả năng viết tiếng Anh của mình cũng không thụt lùi). Phần còn lại vì mình cũng không viết nhiều như trước nữa (bây giờ khó mà viết được một bài cảm nhận phim 3 4 đoạn như xưa). Không những vậy, có những thứ mình từng hứa hẹn là sẽ viết mà cuối cùng lại không chấp bút. Lý do sâu xa có lẽ là vì mình không tìm được cảm hứng để tiếp tục, do mình có một quan điểm khá là cụ thể về bố cục và thành phần của một bài viết mà nếu không có được toàn bộ những thành phần đó thì cũng không đáng viết là bao. Lý do gần là mình lừi.

Bài viết này là một trong những thứ mình ấp ủ cũng đã lâu rồi. Do nó nằm trong không gian những thứ mà mình yêu thích nhất. Dạo này nhân tiện được làm một dự án liên quan nên cũng tìm hiểu đủ sâu và đủ nhiều để hội tụ đủ những yếu tố mà mình cần. Mình cũng không biết mình sẽ còn viết những kỳ sau không, hy vọng là có.

Tham khảo

  1. Video “Variational Inference: Simple Example (+ Python Demo)” của Machine Learning & Simulation ở Youtube