关于 KL 散度和变分推断的 ELBO

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 性质

非负性：KL 散度总是非负的，$\text{KL}(P \| Q) \geq 0$。
不对称性：KL 散度不是对称的，即 $\text{KL}(P \| Q) \neq \text{KL}(Q \| P)$。
零点：当 $P$ 和 $Q$ 完全相同时，$\text{KL}(P \| Q) = 0$。
不满足三角不等式：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/