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


01 KL 散度

Kullback-Leibler (KL) 散度,是一种描述 一个概率分布 P 相对于另一个概率分布 Q 的非对称性差异的概念。

KL 散度是非负的;当且仅当两个分布相同时,它为零。

1.1 定义

对于离散概率分布,PQ 的 KL 散度定义为:

KL(PQ)=xP(x)logP(x)Q(x)

对于连续概率分布,定义为:

KL(PQ)=p(x)logp(x)q(x)dx

其中,p(x)P 的概率密度函数,q(x)Q 的概率密度函数。

1.2 性质

  1. 非负性:KL 散度总是非负的,KL(PQ)0
  2. 不对称性:KL 散度不是对称的,即 KL(PQ)KL(QP)
  3. 零点:当 PQ 完全相同时,KL(PQ)=0
  4. 不满足三角不等式:KL 散度不满足传统意义上的三角不等式。

1.3 变分推断中的 KL 散度

在变分推断中,KL 散度用于衡量一个变分分布 q(z|x) 与真实后验分布 p(z|x) 之间的差异,即:

KL(q(z|x)  p(z|x))

通过最小化这个差异,我们可以得到一个对后验分布 p(z|x) 的良好近似。

  • 先验:没有任何信息,先猜一波 latent 分布, p(z)
  • 后验:给定结果,猜猜我是基于什么 latent 做的, p(z|x)

然而,直接最小化 KL 散度可能很困难,因为它涉及到对真实后验分布 p(z|x) 的直接计算。变分下界(如 ELBO)提供了一种通过下界来间接最小化 KL 散度的方法,使得优化过程更加可行。

02 变分下界(证据下界 Evidence Lower Bound, ELBO)

变分下界(Variational Lower Bound)是变分推断中的一个概念。在复杂概率模型中,ELBO 用于近似难以直接计算的量,如互信息或其他后验分布。

2.1 变分下界的含义

在变分推断中,我们通常有一个复杂的概率模型,它包含观测数据 x 和一些隐变量 z。我们希望找到隐变量的后验分布 p(z|x),比如给定轨迹 x 后,该轨迹对应的 task z 的分布。

由于计算复杂性,这个分布往往难以直接计算。变分下界提供了一种近似后验分布的方法,通过优化一个简化的变分分布 q(z)

变分下界基于 Kullback-Leibler (KL) 散度的概念,KL 散度衡量了两个概率分布之间的差异。

在变分推断中,我们希望找到 q(z|x),使得它与真实后验分布 p(z|x) 尽可能接近:最小化它们之间的 KL 散度:

KL(q(z|x)  p(z|x))=zq(z|x)logq(z|x)p(z|x)dz

然而,直接最小化 KL 散度可能很困难,因为它涉及到对 p(z|x) 的直接计算。变分下界提供了间接最小化 KL 散度的方法,通过最大化 KL 散度的下界。

我们考察两个后验概率分布的 KL 散度,得到:

KL(q(z|x)  p(z|x))=logp(x)+KL(q(z|x) || p(z))Ezq(z|x)[logp(x|z)]

  • 该式的证明:按定义写一遍,然后只对概率分布 p 用贝叶斯公式变换一下, p(x,z)=p(z)p(x|z)=p(x)p(z|x) ,即可发现该式正确)

  • 贴一个证明:

  • DKL(qϕ(z|x)pθ(z|x))=qϕ(z|x)logqϕ(z|x)pθ(z|x)dz=qϕ(z|x)logqϕ(z|x)pθ(x)pθ(z,x)dz;Because p(z|x)=p(z,x)/p(x)=qϕ(z|x)(logpθ(x)+logqϕ(z|x)pθ(z,x))dz=logpθ(x)+qϕ(z|x)logqϕ(z|x)pθ(z,x)dz;Because q(z|x)dz=1=logpθ(x)+qϕ(z|x)logqϕ(z|x)pθ(x|z)pθ(z)dz;Because p(z,x)=p(x|z)p(z)=logpθ(x)+Ezqϕ(z|x)[logqϕ(z|x)pθ(z)logpθ(x|z)]=logpθ(x)+DKL(qϕ(z|x)pθ(z))Ezqϕ(z|x)logpθ(x|z)

现在,重新排列等式的左右两侧,得到

logp(x)KL(q(z|x)  p(z|x))=Ezq(z|x)[logp(x|z)]KL(q(z|x)  p(z))

为了最小化 KL 散度,我们希望最大化 上式的 RHS 】:

  • 第一项,最大化 Ezq(z|x)[logp(x|z)] ,相当于最大化 p(x|z) 的 log likelihood,希望学到变分分布 q(z|x) ,使得在 z 下生成的 x ,更符合我们观测到的 x 数据;
  • 第二项,最小化 KL(q(z|x)  p(z)) ,意味着我们希望变分分布 q 尽可能接近先验分布 p(z),从而确保 变分分布不会偏离 我们对隐藏变量的先验知识。

在变分贝叶斯方法中,这种最大化的形式称为 ELBO。ELBO 名字里的 “lower bound” 是因为,RHS 中的第二项 KL 散度始终是非负的,因此 RHS 是 Ezq(z|x)[logp(x|z)] 的下界。

2.2 省流

如果我们想最小化 KL 散度:

KL(q(z|x)  p(z|x))=zq(z|x)logq(z|x)p(z|x)dz

那么可以把优化目标写成,最大化:

J=Ezq(z|x)[logp(x|z)]KL(q(z|x)  p(z))

即,设计 [-上式] 为损失函数。

其中,第一项:最大化样本点 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 的形式: π(a|s,z) ,其中 z 代表一个 skill,策略基于这个 latent skill 来生成轨迹。

我们希望的策略,符合下面两个要求:

  • Predictable:各个 skill 下的 policy,不要都训成一样的;每个 skill 下的行为,可以被明显区分。
  • Diverse:所有 skill 下 policy 访问的状态,要尽可能覆盖整个状态空间。

为此,我们希望最大化 skill z 和 state s 的互信息 I(s;z)

I(s;z)=szp(s,z)logp(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)=xp(x)logp(x)dx

我们介绍一下互信息(Mutual Information,MI)。

  • 性质:
    • 对称性, I(s;z)=I(z;s)
    • 非负性, I(s;z)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ϕ(z|x)
  • (然后就使用 ELBO 嘛?DIAYN 好像原文不是这样写的,没细看,我也不太清楚了 😵💦

3.2 VAE 的 loss function

Autoencoder:核心思想是使用一个沙漏型网络,尽可能无损地 把大的数据(如图片)压缩到一个更小的 embedding 里,其损失函数是 MSE[原图片, 基于 embedding 复原的图片]。

VAE:是一种生成模型,它可以基于一些 latent 来生成数据,比如给一些自然语言的描述 来生成图片,或给一张图片 生成相似的图片。(diffusion 也是著名的生成模型)

VAE 跟 autoencoder 的思想不尽相同;对于一个输入图片 x,它不想把图片映射到一个固定的 embedding 向量 z,而是将其映射到一个分布 p(z|x) 中。

VAE 的组成部分:

  • 条件概率 pθ(x|z) 定义了一个生成模型,类似于 autoencoder 的解码器,即从 latent z 还原到原图片 x 的过程。
  • 近似函数 qϕ(z|x) 是概率编码器,输入是图片 x,输出是这张图片对应的 latent z 的分布。

VAE 的损失函数:

  • 对于编码器部分 qϕ(z|x) ,貌似采用了 ELBO 形式,即,最小化 KL 散度 KL(qϕ(z|x)  p(z|x)) → 最大化 Ezqϕ(z|x)[logp(x|z)]KL(qϕ(z|x)  p(z))
  • 对于解码器部分 pθ(x|z) ,貌似还是 autoencoder 的样本重构损失(?)具体技术细节我也不太清楚…

参考资料 / 博客:



本文作者:MoonOut

本文链接:https://www.cnblogs.com/moonout/p/18263752

版权声明:本作品采用知识共享署名-非商业性使用-禁止演绎 2.5 中国大陆许可协议进行许可。

posted @   MoonOut  阅读(865)  评论(0编辑  收藏  举报
点击右上角即可分享
微信分享提示
评论
收藏
关注
推荐
深色
回顶
收起
  1. 1 Sibelius: Violin Concerto in D Minor, Op. 47:III. Allegro, ma non tanto Jascha Heifetz / Chicago Symphony Orchestra
Sibelius: Violin Concerto in D Minor, Op. 47:III. Allegro, ma non tanto - Jascha Heifetz / Chicago Symphony Orchestra
00:00 / 00:00
An audio error has occurred.