關於 KL 散度和變分推斷的 ELBO

MoonOut發表於2024-06-23

01 KL 散度

Kullback-Leibler (KL) 散度，是一種描述 一個機率分佈 \(P\) 相對於另一個機率分佈 \(Q\) 的非對稱性差異的概念。

KL 散度是非負的；當且僅當兩個分佈相同時，它為零。

1.1 定義

對於離散機率分佈，\(P\) 和 \(Q\) 的 KL 散度定義為：

\[\text{KL}(P \| Q) = \sum_{\mathbf{x}} P(\mathbf{x}) \log \frac{P(\mathbf{x})}{Q(\mathbf{x})} \]

對於連續機率分佈，定義為：

\[\text{KL}(P \| Q) = \int p(\mathbf{x}) \log \frac{p(\mathbf{x})}{q(\mathbf{x})} d\mathbf{x} \]

其中，\(p(\mathbf{x})\) 是 \(P\) 的機率密度函式，\(q(\mathbf{x})\) 是 \(Q\) 的機率密度函式。

1.2 性質

  1. 非負性：KL 散度總是非負的，\(\text{KL}(P \| Q) \geq 0\)。
  2. 不對稱性：KL 散度不是對稱的，即 \(\text{KL}(P \| Q) \neq \text{KL}(Q \| P)\)。
  3. 零點：當 \(P\) 和 \(Q\) 完全相同時，\(\text{KL}(P \| Q) = 0\)。
  4. 不滿足三角不等式：KL 散度不滿足傳統意義上的三角不等式。

1.3 變分推斷中的 KL 散度

在變分推斷中，KL 散度用於衡量一個變分分佈 \(q(\mathbf{z}|\mathbf{x})\) 與真實後驗分佈 \(p(\mathbf{z}|\mathbf{x})\) 之間的差異，即：

\[\text{KL}\big(q(\mathbf{z}|\mathbf{x}) ~\|~ p(\mathbf{z}|\mathbf{x})\big) \]

透過最小化這個差異，我們可以得到一個對後驗分佈 \(p(\mathbf{z}|\mathbf{x})\) 的良好近似。

  • 先驗：沒有任何資訊，先猜一波 latent 分佈， \(p(\mathbf{z})\) ；
  • 後驗：給定結果，猜猜我是基於什麼 latent 做的， \(p(\mathbf{z}|\mathbf{x})\) 。

然而，直接最小化 KL 散度可能很困難，因為它涉及到對真實後驗分佈 \(p(\mathbf{z}|\mathbf{x})\) 的直接計算。變分下界（如 ELBO）提供了一種透過下界來間接最小化 KL 散度的方法，使得最佳化過程更加可行。

02 變分下界（證據下界 Evidence Lower Bound, ELBO）

變分下界（Variational Lower Bound）是變分推斷中的一個概念。在複雜機率模型中，ELBO 用於近似難以直接計算的量，如互資訊或其他後驗分佈。

2.1 變分下界的含義

在變分推斷中，我們通常有一個複雜的機率模型，它包含觀測資料 \(\mathbf{x}\) 和一些隱變數 \(\mathbf{z}\)。我們希望找到隱變數的後驗分佈 \(p(\mathbf{z}|\mathbf{x})\)，比如給定軌跡 \(\mathbf{x}\) 後，該軌跡對應的 task \(\mathbf{z}\) 的分佈。

由於計算複雜性，這個分佈往往難以直接計算。變分下界提供了一種近似後驗分佈的方法，透過最佳化一個簡化的變分分佈 \(q(\mathbf{z})\)。

變分下界基於 Kullback-Leibler (KL) 散度的概念，KL 散度衡量了兩個機率分佈之間的差異。

在變分推斷中，我們希望找到 \(q(\mathbf{z}|\mathbf{x})\)，使得它與真實後驗分佈 \(p(\mathbf{z}|\mathbf{x})\) 儘可能接近：最小化它們之間的 KL 散度：

\[\text{KL}\big(q(\mathbf{z}|\mathbf{x}) ~\|~ p(\mathbf{z}|\mathbf{x})\big) = \int_\mathbf{z} q(\mathbf{z}|\mathbf{x}) \log \frac{q(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})} d\mathbf{z} \]

然而，直接最小化 KL 散度可能很困難，因為它涉及到對 \(p(\mathbf{z}|\mathbf{x})\) 的直接計算。變分下界提供了間接最小化 KL 散度的方法，透過最大化 KL 散度的下界。

我們考察兩個後驗機率分佈的 KL 散度，得到：

\[\text{KL}(q(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z}|\mathbf{x})) = \log p(\mathbf{x}) + \text{KL}\big(q(\mathbf{z}|\mathbf{x})~||~p(\mathbf{z})\big) - \mathbb{E}_{\mathbf{z}\sim q(\mathbf{z}|\mathbf{x})} \big[\log p(\mathbf{x}|\mathbf{z})\big] \]

  • 該式的證明：按定義寫一遍，然後只對機率分佈 p 用貝葉斯公式變換一下， \(p(\mathbf{x},\mathbf{z})=p(\mathbf{z})p(\mathbf{x}|\mathbf{z})=p(\mathbf{x})p(\mathbf{z}|\mathbf{x})\) ，即可發現該式正確）

  • 貼一個證明：

  • \[\begin{aligned} &D_{\mathrm{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z}|\mathbf{x})) \\ &=\int q_\phi(\mathbf{z}|\mathbf{x})\log\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{z}|\mathbf{x})}d\mathbf{z} \\ &=\int q_\phi(\mathbf{z}|\mathbf{x})\log\frac{q_\phi(\mathbf{z}|\mathbf{x})p_\theta(\mathbf{x})}{p_\theta(\mathbf{z},\mathbf{x})}d\mathbf{z}& ;\mathrm{Because~}p(z|x)=p(z,x)/p(x) \\ &=\int q_\phi(\mathbf{z}|\mathbf{x})\big(\log p_\theta(\mathbf{x})+\log\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{z},\mathbf{x})}\big)d\mathbf{z} \\ &=\log p_\theta(\mathbf{x})+\int q_\phi(\mathbf{z}|\mathbf{x})\log\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{z},\mathbf{x})}d\mathbf{z}& ;\mathrm{Because~}\int q(z|x)dz=1 \\ &=\log p_\theta(\mathbf{x})+\int q_\phi(\mathbf{z}|\mathbf{x})\log\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{x}|\mathbf{z})p_\theta(\mathbf{z})}d\mathbf{z}& ;\mathrm{Because~}p(z,x)=p(x|z)p(z) \\ &=\log p_\theta(\mathbf{x})+\mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}|\mathbf{x})}[\log\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p_\theta(\mathbf{z})}-\log p_\theta(\mathbf{x}|\mathbf{z})] \\ &=\log p_\theta(\mathbf{x})+D_{\mathrm{KL}}(q_\phi(\mathbf{z}|\mathbf{x})\|p_\theta(\mathbf{z}))-\mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}|\mathbf{x})}\log p_\theta(\mathbf{x}|\mathbf{z}) \end{aligned} \]

現在，重新排列等式的左右兩側，得到

\[\log p(\mathbf{x}) - \text{KL}\big(q(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z}|\mathbf{x})\big) = \mathbb{E}_{\mathbf{z}\sim q(\mathbf{z}|\mathbf{x})} \big[\log p(\mathbf{x}|\mathbf{z})\big] - \text{KL}\big(q(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z})\big) \]

【 為了最小化 KL 散度，我們希望最大化 上式的 RHS 】：

  • 第一項，最大化 \(\mathbb{E}_{\mathbf{z}\sim q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x}|\mathbf{z})]\) ，相當於最大化 \(p(\mathbf{x}|\mathbf{z})\) 的 log likelihood，希望學到變分分佈 \(q(\mathbf{z}|\mathbf{x})\) ，使得在 \(\mathbf{z}\) 下生成的 \(\mathbf{x}\) ，更符合我們觀測到的 \(\mathbf{x}\) 資料；
  • 第二項，最小化 \(\text{KL}(q(\mathbf{z|x})~\|~p(\mathbf{z}))\) ，意味著我們希望變分分佈 \(q\) 儘可能接近先驗分佈 \(p(\mathbf{z})\)，從而確保 變分分佈不會偏離 我們對隱藏變數的先驗知識。

在變分貝葉斯方法中，這種最大化的形式稱為 ELBO。ELBO 名字裡的 “lower bound” 是因為，RHS 中的第二項 KL 散度始終是非負的，因此 RHS 是 \(\mathbb{E}_{\mathbf{z}\sim q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x}|\mathbf{z})]\) 的下界。

2.2 省流

如果我們想最小化 KL 散度：

\[\text{KL}\big(q(\mathbf{z}|\mathbf{x}) ~\|~ p(\mathbf{z}|\mathbf{x})\big) = \int_\mathbf{z} q(\mathbf{z}|\mathbf{x}) \log \frac{q(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})} d\mathbf{z} \]

那麼可以把最佳化目標寫成，最大化：

\[J = \mathbb{E}_{\mathbf{z}\sim q(\mathbf{z}|\mathbf{x})} \big[\log p(\mathbf{x}|\mathbf{z})\big] - \text{KL}\big(q(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z})\big) \]

即，設計 [-上式] 為損失函式。

其中，第一項：最大化樣本點 x 的 log likelihood，第二項：最小化 z 分佈與先驗 p(z) 的 KL 散度。

03 ELBO 的應用：skill discovery、VAE

3.1 skill discovery 的 loss function

Skill discovery 是一種無 reward function 的 online RL 任務，它透過無監督的方法，學習一組覆蓋狀態空間的、具有明顯差異的技能（skill）。

Policy 的形式： \(\pi(a|s,z)\) ，其中 z 代表一個 skill，策略基於這個 latent skill 來生成軌跡。

我們希望的策略，符合下面兩個要求：

  • Predictable：各個 skill 下的 policy，不要都訓成一樣的；每個 skill 下的行為，可以被明顯區分。
  • Diverse：所有 skill 下 policy 訪問的狀態，要儘可能覆蓋整個狀態空間。

為此，我們希望最大化 skill z 和 state s 的互資訊 \(I(s;z)\) ：

\[I(s;z)=\int_s\int_z p(s,z)\log\frac{p(s,z)}{p(s)p(z)} \\ =H(z)-H(z|s)=H(s)-H(s;z) \\ =H(s)+H(z)-H(s,z) \]

其中 H 是熵，定義為 \(H(x) = -\int_x p(x)\log p(x)dx\) 。

我們介紹一下互資訊（Mutual Information，MI）。

  • 性質：
    • 對稱性， \(I(s;z)=I(z;s)\) ；
    • 非負性， \(I(s;z)\ge 0\)，等於 0 當且僅當 s z 獨立。
  • 上面公式 10 的幾個等號，把熵的公式帶進去 就能得到。
  • 當兩個分佈完全相同 完全不獨立時，貌似 \(I(s;z)\) 取到最大值，最大值為 \(H(s)=H(z)\)。

怎麼最大化互資訊呢？

我們從最大化 \(I(s;z)=H(z)-H(z|s)\) 或 \(I(s;z)=H(s)-H(s;z)\) 的形式入手。具體的，

  • Reverse MI：
    • 最大化 \(I(s;z)=H(z)-H(z|s)\)，被稱為 Reverse MI（相關文章：Diversity is all you need）。
    • 其中，第一項最大化 \(H(z)\)，鼓勵學到多樣的 skill；
    • 第二項最小化 \(H(z | s)\)，希望看到 state 就推斷出 skill。
    • 多說一句，Diversity is all you need 的主要貢獻之一，貌似是這裡還會最大化 \(H[a|s,z]\) ，最大化給定 skill 後的策略的熵，旨在鼓勵 diversity。
  • Forward MI：
    • 最大化 \(I(s;z)=H(s)-H(s|z)\)，被稱為 Forward MI，一般用於 model-based RL（相關文章：Dynamics-Aware Unsupervised Discovery of Skills）。
    • 其中，第一項最大化 \(H(s)\)，鼓勵學到多樣的 state；
    • 第二項最小化 \(H(s | z)\)，鼓勵透過 state 和 z 推斷出 state'，這貌似是 model-based RL 學 env model 的一個魔改。

對於 reverse MI（Diversity is all you need），現在我們要最小化 \(H(z | s)\) 了。

  • 因此，對於後驗分佈 \(p(z|{x})\) ，我們需要搞一個引數化的近似分佈 \(q_\phi({z}|{x})\) 。
  • （然後就使用 ELBO 嘛？DIAYN 好像原文不是這樣寫的，沒細看，我也不太清楚了 😵💦

3.2 VAE 的 loss function

Autoencoder：核心思想是使用一個沙漏型網路，儘可能無損地 把大的資料（如圖片）壓縮到一個更小的 embedding 裡，其損失函式是 MSE[原圖片, 基於 embedding 復原的圖片]。

VAE：是一種生成模型，它可以基於一些 latent 來生成資料，比如給一些自然語言的描述 來生成圖片，或給一張圖片 生成相似的圖片。（diffusion 也是著名的生成模型）

VAE 跟 autoencoder 的思想不盡相同；對於一個輸入圖片 \(\mathbf{x}\)，它不想把圖片對映到一個固定的 embedding 向量 \(\mathbf{z}\)，而是將其對映到一個分佈 \(p(\mathbf{z}|\mathbf{x})\) 中。

VAE 的組成部分：

  • 條件機率 \(p_\theta(\mathbf{x}|\mathbf{z})\) 定義了一個生成模型，類似於 autoencoder 的解碼器，即從 latent \(\mathbf{z}\) 還原到原圖片 \(\mathbf{x}\) 的過程。
  • 近似函式 \(q_\phi(\mathbf{z}|\mathbf{x})\) 是機率編碼器，輸入是圖片 \(\mathbf{x}\)，輸出是這張圖片對應的 latent \(\mathbf{z}\) 的分佈。

VAE 的損失函式：

  • 對於編碼器部分 \(q_\phi(\mathbf{z}|\mathbf{x})\) ，貌似採用了 ELBO 形式，即，最小化 KL 散度 \(\text{KL}(q_\phi(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z}|\mathbf{x}))\) → 最大化 \(\mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}|\mathbf{x})} \big[\log p(\mathbf{x}|\mathbf{z})\big] - \text{KL}\big(q_\phi(\mathbf{z}|\mathbf{x})~\|~p(\mathbf{z})\big)\) 。
  • 對於解碼器部分 \(p_\theta(\mathbf{x}|\mathbf{z})\) ，貌似還是 autoencoder 的樣本重構損失（？）具體技術細節我也不太清楚…

參考資料 / 部落格：

  • Diversity is All You Need，https://arxiv.org/abs/1802.06070
  • lilian weng 的 VAE 部落格，https://lilianweng.github.io/posts/2018-08-12-vae/
  • lilian weng 的 diffusion 部落格，https://lilianweng.github.io/posts/2021-07-11-diffusion-models/


相關文章