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/