今天看论文，里面用了TVAE（Transformer VAE）来把大量的运动状态序列压缩成一个连续、结构化的 16 维 latent action space，并据此为后续的力矩解码和强化学习提供可复用的动作表示。我突然发现自己对于VAE的理解一直有些偏差，这里简单梳理一下。不涉及复杂的KL散度和ELBO的推理
要点
我们用最简单的MNIST手写数字作为例子。VAE的Encoder根据输入图像x，输出latent variable z的概率分布参数。标准 VAE 中，Encoder为latent space的每个维度输出一个均值和方差（实际实现中通常输出），从而定义近似后验分布： 。
随后，我们从这个概率分布中采样一个具体的latent vector z。实际训练时通常采用重参数化技巧：得到具体的z后，再将其送入Decoder，重建输入图像x。
训练VAE时，Loss主要由两部分组成：。其中，Reconstruction Loss要求重建结果尽可能接近原始输入x，从而保证latent variable z中保留了足够的信息用于重建。KL Loss则衡量Encoder得到的近似后验分布：。与我们预先规定的先验分布：(0,I)之间的差异，即：。KL Loss会约束不要偏离先验分布太远，使latent space保持相对规整并且可以进行采样。
因此，VAE 的训练过程可以简单理解为：。训练完成之后，我们不再必须经过 Encoder，而可以直接从先验分布中采样：。然后把采样得到的z送入 Decoder： 。从而生成新的数据。
注意 1： Encoder 输出的并不是一个确定的latent vector，而是用于描述latent分布的一组参数和。随后才从这个分布中采样得到具体的z。
注意 2： KL Loss比较的是近似后验分布与先验分布之间的差异，而不是简单比较两个输入样本之间的差异。
注意 3： 更严谨地说，标准VAE并不是强制整个latent space精确服从，而是通过KL Loss鼓励每个不要距离先验太远，从而使 latent representation 与预先规定的先验分布更加匹配，使得从先验分布中直接采样并生成数据成为可能。
代码的核心流程
前面已经知道，VAE 的核心流程是：
对应到 PyTorch，核心其实只有下面几行：
其中： 而重参数化技巧：
使得随机采样过程仍然可以参与梯度反向传播。
VAE的Loss由Reconstruction Loss和KL Loss两部分组成：
其中 KL Loss 对应：.
尤其需要注意KL Loss的代码，它其实只是 两个高斯分布之间 KL 散度公式的直接实现。这里简单展开一下
第一次看到这段代码可能会觉得有些莫名其妙，但它其实只是 两个高斯分布之间 KL 散度公式的直接实现。
前面我们已经知道，对于一张输入图像x，Encoder 会输出latent variable的均值和方差 ，从而定义近似后验分布：
而 VAE 预先规定的先验分布通常是：
也就是说，我们希望 Encoder 得到的 不要偏离标准正态分布太远。
对于这两个高斯分布，KL 散度存在解析解：
其中i表示latent space的第i个维度。
例如，如果latent dimension为4，那么实际上就是分别计算四个维度的KL，再将它们相加：
上面的 KL 公式重新排列一下，可以写成：
注意：在实际实现中，Encoder 通常不会直接输出 ，而是输出：
训练完成后，如果只需要生成新数据，则不再需要 Encoder，可以直接从先验分布中采样：
即：
因此，整个 VAE 最核心的代码逻辑可以总结为：
这几行代码基本已经覆盖了 VAE 最重要的三个概念：概率编码、重参数化采样和 KL 正则化。
义父，请我喝杯蜜雪冰城吧。
- 作者:LIU Xiao
- 链接:http://liuxiao916.com/article/3ba58a7f-6b9f-8064-aa9b-fa11029c9c91
- 声明:本文采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。
