变分自编码器 (VAE)Variational Autoencoders

机器人经常需要生成东西，而不只是给输入打标签：想象下一帧相机画面会是什么样（用来在“脑子里”试动作），为一个没见过的物体提出几个可行的抓取，或者在“从左边绕”和“从右边绕”都对的时候，给出一个具体的动作而不是两者的平均。这些都要求一个能从中采样的概率模型 \(p(\mathbf x)\)，而且 \(\mathbf x\) 往往是高维的：一张图像、一段动作序列。

18.2 节的自编码器把图像压成一个低维编码 \(\mathbf z\) 再还原，学到的编码很有用，但它没有告诉我们编码服从什么分布：编码散落在隐空间里形状任意、中间有大片空洞，随便取一个 \(\mathbf z\) 解码，多半不像任何真实图像。VAE（Kingma 与 Welling，2014；Rezende 等，2014 同时提出）的做法是把整件事写成一个隐变量模型：先从一个简单的先验 \(p(\mathbf z)=\mathcal N(\mathbf 0,I)\) 采样，再由解码网络生成 \(\mathbf x\)。训练这个模型需要反推“哪个 \(\mathbf z\) 生成了这个 \(\mathbf x\)”（后验），这在高维里算不出来，于是再训练一个编码网络去近似它，并用一个可计算的下界 ELBO 代替算不出的似然。

这个框架在机器人学习里随处可见：ACT（22.2 节）用条件 VAE 生成动作块；World Models 用 VAE 把游戏画面压成隐向量再在隐空间里学控制；Dreamer（24.3 节）的世界模型是一个按时间展开的隐变量模型，训练目标就是一个 ELBO；扩散模型（19.6 节）也可以看成层数很多、编码器固定的 VAE。本节把这些背后的共同部件一个个拆开。

Robots often need to generate things, not just label their inputs: imagine what the next camera frame will look like (to try an action "in the head"), propose a few feasible grasps for an unseen object, or, when "go around on the left" and "go around on the right" are both right, output one concrete action rather than the average of the two. All of this needs a probabilistic model \(p(\mathbf x)\) you can sample from, and \(\mathbf x\) is usually high-dimensional: an image, a sequence of actions.

The autoencoder of Lesson 18.2 squeezes an image into a low-dimensional code \(\mathbf z\) and back. The codes are useful, but nothing says how they are distributed: they scatter over the latent space in an arbitrary shape with large holes in between, and decoding a \(\mathbf z\) picked at random usually gives something unlike any real image. The VAE (Kingma & Welling, 2014; proposed at the same time by Rezende et al., 2014) writes the whole thing as a latent-variable model: first sample from a simple prior \(p(\mathbf z)=\mathcal N(\mathbf 0,I)\), then let a decoder network generate \(\mathbf x\). Training this model requires inverting it, "which \(\mathbf z\) produced this \(\mathbf x\)?" (the posterior), which cannot be computed in high dimensions, so a second, encoder network is trained to approximate it, and a computable lower bound, the ELBO, stands in for the likelihood we cannot compute.

This framework is everywhere in robot learning: ACT (Lesson 22.2) generates action chunks with a conditional VAE; World Models compresses game frames into latent vectors with a VAE and learns control in that latent space; the world model of Dreamer (Lesson 24.3) is a latent-variable model unrolled in time whose training objective is an ELBO; diffusion models (Lesson 19.6) can be read as a VAE with many layers and a fixed encoder. This lesson takes their shared parts apart one at a time.

先看一个我们完全知道答案的生成模型。一台平面二连杆机械臂（连杆长 0.62 m 和 0.50 m，7.3 节的正运动学），它的关节角由一个二维隐变量决定：\(\hl{prior}{\mathbf z\sim\mathcal N(\mathbf 0,I)}\)，\(q_1=\tfrac{\pi}{2}+0.7z_1\)，\(q_2=z_2\)。我们看不到关节角，只看到手的位置，而且带噪声：\(\hl{lik}{\mathbf x\sim\mathcal N(f(\mathbf z),\sigma_x^2I)}\)，\(f\) 是正运动学。这就是一个隐变量模型：\(\mathbf z\) 是原因，\(\mathbf x\) 是观测，“解码器” \(f\) 是已知的。

从这个模型可以问两个问题。生成：手会出现在哪里？把很多 \(\mathbf z\) 推过去，得到下图左边的青色点云，它的密度就是边缘似然 \(\hl{px}{p(\mathbf x)=\int p(\mathbf z)\,p(\mathbf x|\mathbf z)\,d\mathbf z}\)。推断：看到手在某处，关节角可能是什么？这是贝叶斯公式（4.2 节）给出的后验 \(\hl{post}{p(\mathbf z|\mathbf x)}=p(\mathbf x|\mathbf z)p(\mathbf z)/p(\mathbf x)\)。因为 \(\mathbf z\) 只有两维，我们可以在一张 181×181 的网格上把它直接算出来，画在右边。拖动琥珀色的目标点，看后验怎样变。

Start with a generative model whose answers we know exactly. A planar two-link arm (links of 0.62 m and 0.50 m, the forward kinematics of Lesson 7.3) has its joint angles set by a two-dimensional latent variable: \(\hl{prior}{\mathbf z\sim\mathcal N(\mathbf 0,I)}\), \(q_1=\tfrac{\pi}{2}+0.7z_1\), \(q_2=z_2\). We never see the joint angles, only the hand position, and with noise: \(\hl{lik}{\mathbf x\sim\mathcal N(f(\mathbf z),\sigma_x^2I)}\), where \(f\) is the forward kinematics. This is a latent-variable model: \(\mathbf z\) is the cause, \(\mathbf x\) the observation, and the "decoder" \(f\) is known.

Two questions can be asked of it. Generation: where does the hand appear? Push many \(\mathbf z\) through and you get the cyan cloud on the left below; its density is the marginal likelihood \(\hl{px}{p(\mathbf x)=\int p(\mathbf z)\,p(\mathbf x|\mathbf z)\,d\mathbf z}\). Inference: the hand is seen somewhere; what might the joint angles be? That is the posterior of Bayes' rule (Lesson 4.2), \(\hl{post}{p(\mathbf z|\mathbf x)}=p(\mathbf x|\mathbf z)p(\mathbf z)/p(\mathbf x)\). Because \(\mathbf z\) has only two dimensions we can compute it directly on a 181×181 grid and draw it on the right. Drag the amber target and watch the posterior change.

现在换成真正的 VAE：数据是机械臂的图像，没有人告诉网络关节角。我们按 FIG 1 的机械臂在 \(q_1\in[\tfrac\pi2-1.1,\tfrac\pi2+1.1]\)、\(q_2\in[-1.7,1.7]\) 里均匀随机取 500 个姿态，各画一张 14×10 像素的灰度图（像素值在 0 到 1 之间，140 维），另取 100 张作测试集。模型有两个网络：

• 编码器 \(q_\phi(\mathbf z|\mathbf x)=\mathcal N(\boldsymbol\mu_\phi(\mathbf x),\operatorname{diag}\boldsymbol\sigma^2_\phi(\mathbf x))\)：MLP 140 → 48（ReLU）→ 4，输出 2 个均值和 2 个 \(\log\sigma^2\)。
• 解码器 \(p_\theta(\mathbf x|\mathbf z)\)：MLP 2 → 48（ReLU）→ 140，每个像素输出一个 logit，按伯努利分布计算似然，于是重建项就是 140 个像素的二元交叉熵之和（4.11 节）。

每一步取 16 张图，对每张：编码得到 \(\boldsymbol\mu,\boldsymbol\sigma\)，采一个 \(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\)，令 \(\mathbf z=\boldsymbol\mu+\boldsymbol\sigma\odot\boldsymbol\epsilon\)（重参数化，第 4 节），解码，损失为

Now a real VAE: the data are images of the arm, and nobody tells the network the joint angles. Using the arm of FIG 1 we draw 500 poses uniformly from \(q_1\in[\tfrac\pi2-1.1,\tfrac\pi2+1.1]\), \(q_2\in[-1.7,1.7]\) and render each as a 14×10-pixel grey image (pixel values between 0 and 1, 140 dimensions), plus 100 more as a test set. The model has two networks:

• Encoder \(q_\phi(\mathbf z|\mathbf x)=\mathcal N(\boldsymbol\mu_\phi(\mathbf x),\operatorname{diag}\boldsymbol\sigma^2_\phi(\mathbf x))\): an MLP 140 → 48 (ReLU) → 4 that outputs 2 means and 2 values of \(\log\sigma^2\).
• Decoder \(p_\theta(\mathbf x|\mathbf z)\): an MLP 2 → 48 (ReLU) → 140 with one logit per pixel and a Bernoulli likelihood, so the reconstruction term is the sum of binary cross-entropies over the 140 pixels (Lesson 4.11).

Each step takes 16 images, and for each: encode to get \(\boldsymbol\mu,\boldsymbol\sigma\), draw \(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\), set \(\mathbf z=\boldsymbol\mu+\boldsymbol\sigma\odot\boldsymbol\epsilon\) (the reparameterization, Section 4), decode, and the loss is

KL 项有闭式（4.11 节高斯 KL 公式的逐维求和），不用采样。\(\beta=1\) 时 \(\mathcal L\) 恰好是负的 ELBO（第 5 节推导），这就是标准 VAE；\(\beta\) 是你可以拉的旋钮。优化器是 Adam（学习率 0.002，6.3 节），全部由 nn.js 在你的浏览器里一步步算出，每步的随机数都有固定种子，所以下面引用的数字可以复现。

The KL term has a closed form (the Gaussian KL of Lesson 4.11, summed over dimensions), so it needs no sampling. With \(\beta=1\), \(\mathcal L\) is exactly the negative ELBO (derived in Section 5): the standard VAE. \(\beta\) is a knob you can turn. The optimizer is Adam (learning rate 0.002, Lesson 6.3), all computed step by step in your browser by nn.js; every random draw has a fixed seed, so the numbers quoted below are reproducible.

回到 FIG 1 的机械臂：在那里 \(\log p(\mathbf x)\) 和真实后验都能算，所以可以亲眼看到 VAE 的核心等式（第 5 节推导）：对任意分布 \(q(\mathbf z)\)，

Back to the arm of FIG 1: there \(\log p(\mathbf x)\) and the true posterior are computable, so we can watch the central identity of the VAE (derived in Section 5) with our own eyes: for any distribution \(q(\mathbf z)\),

左边与 \(q\) 无关，所以提高 ELBO 就等于缩小缺口，即让 \(q\) 贴近真实后验；缺口为零当且仅当 \(q\) 就是后验。下图中 \(q=\mathcal N(\mathbf m,\operatorname{diag}(s_1^2,s_2^2))\) 是对角高斯，和 VAE 编码器的输出形式一样。你可以手动拖它，也可以按“拟合 q”，用重参数化的随机梯度上升（每步 8 个样本，Adam，学习率 0.03，共 600 步）去最大化 ELBO。ELBO 用 24×24 点的高斯-埃尔米特求积精确计算，不含采样噪声。

The left side does not depend on \(q\), so raising the ELBO is the same as shrinking the gap, i.e. pulling \(q\) towards the true posterior; the gap is zero exactly when \(q\) is the posterior. Below, \(q=\mathcal N(\mathbf m,\operatorname{diag}(s_1^2,s_2^2))\) is a diagonal Gaussian, the same form a VAE encoder outputs. Drag it by hand, or press "Fit q" to maximize the ELBO by reparameterized stochastic gradient ascent (8 samples per step, Adam with learning rate 0.03, 600 steps). The ELBO itself is computed exactly by 24×24-point Gauss–Hermite quadrature, with no sampling noise.

这正是 4.11 节讲的反向 KL 的“抓峰”行为：缺口是 \(D_{\mathrm{KL}}(q\|p)\)，\(q\) 放在 \(p\) 很小的地方会被重罚，所以它宁可只蹲在一个峰上。对 VAE 而言，这意味着 ELBO 总是低估 \(\log p(\mathbf x)\)，低估多少取决于后验离“对角高斯”有多远。缺口有两个来源：近似族不够灵活（本图：对角高斯盖不住双峰和弯曲），以及摊销：VAE 不为每个 \(\mathbf x\) 单独优化 \(q\)（像本图这样），而是用一个编码网络一次性输出所有 \(\mathbf x\) 的 \(q\)，网络不完美又会多出一块“摊销缺口”。

This is the "mode-seeking" behaviour of the reverse KL from Lesson 4.11: the gap is \(D_{\mathrm{KL}}(q\|p)\), which punishes \(q\) heavily for putting mass where \(p\) is small, so \(q\) prefers to sit on one peak. For a VAE this means the ELBO always underestimates \(\log p(\mathbf x)\), by an amount that depends on how far the posterior is from a diagonal Gaussian. The gap has two sources: an inflexible approximating family (this figure: a diagonal Gaussian cannot cover two peaks or a bend), and amortization: a VAE does not optimize \(q\) separately for each \(\mathbf x\) as this figure does, but has one encoder network output \(q\) for every \(\mathbf x\) at once, and an imperfect network adds an "amortization gap" on top.

训练要对 \(\mathbb E_{q_\phi}[f(\mathbf z)]\) 关于编码器参数 \(\phi\) 求梯度，可“采样”这个操作本身没有导数。有两条路。得分函数（score function，也叫 REINFORCE、似然比）估计：\(\nabla_\mu\mathbb E[f]=\mathbb E\big[f(z)\,\nabla_\mu\log q(z)\big]\)，只需要能算 \(f\) 的值，不需要 \(f\) 可导，20.3 节的策略梯度用的就是它。重参数化估计：把随机性挪到一个与参数无关的噪声上，\(z=\mu+\sigma\hl{eps}{\epsilon}\)，\(\epsilon\sim\mathcal N(0,1)\)，于是 \(\nabla_\mu\mathbb E[f]=\mathbb E\big[f'(\mu+\sigma\epsilon)\big]\)，梯度直接沿网络反传（16.3 节）。两者的期望相同，方差却天差地别。下面在一维里比较：\(z\) 是抓取的横向偏移，\(f(z)=(z-1.5)^2\) 是离目标偏移 1.5 的平方误差，我们要 \(\nabla_\mu\mathbb E_{z\sim\mathcal N(\mu,\sigma^2)}[f(z)]\)，真值是 \(2(\mu-1.5)\)。

Training needs the gradient of \(\mathbb E_{q_\phi}[f(\mathbf z)]\) with respect to the encoder parameters \(\phi\), but "drawing a sample" has no derivative. There are two routes. The score-function estimator (also called REINFORCE or the likelihood ratio): \(\nabla_\mu\mathbb E[f]=\mathbb E\big[f(z)\,\nabla_\mu\log q(z)\big]\), which only needs values of \(f\), not its derivative; the policy gradient of Lesson 20.3 is exactly this. The reparameterization estimator: move the randomness into a noise that does not depend on the parameters, \(z=\mu+\sigma\hl{eps}{\epsilon}\), \(\epsilon\sim\mathcal N(0,1)\), so that \(\nabla_\mu\mathbb E[f]=\mathbb E\big[f'(\mu+\sigma\epsilon)\big]\) and the gradient flows straight back through the network (Lesson 16.3). Both have the same expectation, but their variances are worlds apart. Below we compare them in 1D: \(z\) is a lateral grasp offset, \(f(z)=(z-1.5)^2\) the squared miss from a target offset of 1.5, and we want \(\nabla_\mu\mathbb E_{z\sim\mathcal N(\mu,\sigma^2)}[f(z)]\), whose true value is \(2(\mu-1.5)\).

重参数化之所以好，是因为它用到了 \(f\) 的导数：它知道“往哪边挪 \(z\) 代价会降”，而得分函数只知道“这个样本代价高还是低”，要靠大量样本的相关性才能平均出方向。代价是要求 \(\mathbf z\) 连续、\(f\) 可微、分布能写成“参数的可微函数 + 固定噪声”（高斯、拉普拉斯、logistic 等可以；离散的类别变量不行，要用 Gumbel-softmax 之类的松弛，或者退回得分函数）。VAE 的解码器是可微网络，正好满足条件，这是 VAE 能在 2014 年训练起来的关键。

Reparameterization wins because it uses the derivative of \(f\): it knows "which way to move \(z\) for the cost to drop", while the score function only knows "this sample was expensive or cheap" and has to average the direction out of correlations over many samples. The price is that \(\mathbf z\) must be continuous, \(f\) differentiable, and the distribution writable as "a differentiable function of the parameters plus fixed noise" (Gaussian, Laplace and logistic distributions qualify; a discrete categorical variable does not and needs a relaxation such as Gumbel-softmax, or a return to the score function). A VAE's decoder is a differentiable network, so it qualifies, and that was the key to training VAEs in 2014.

训练好的 VAE 给了我们一张“隐空间地图”：因为 KL 项把每个编码都拉向 \(\mathcal N(\mathbf 0,I)\)、又让它们带着噪声，相邻的 \(\mathbf z\) 必须解码成相似的图像，先验覆盖的整个区域都被“填满”。下图用 FIG 2 的模型（两图共享，在这里也可以继续训练）把 \([-2.4,2.4]^2\) 上 9×9 个 \(\mathbf z\) 解码成图像，并比较两种插值：在隐空间里走直线 \((1-t)\mathbf z_A+t\mathbf z_B\) 再解码，和在像素空间里直接混合 \((1-t)\mathbf x_A+t\mathbf x_B\)。

A trained VAE gives us a "map" of its latent space: since the KL term pulls every code towards \(\mathcal N(\mathbf 0,I)\) and keeps it noisy, neighbouring \(\mathbf z\) must decode into similar images, and the whole region the prior covers gets "filled in". The figure below takes the model of FIG 2 (shared by both figures; you can keep training it here) and decodes a 9×9 grid of \(\mathbf z\) over \([-2.4,2.4]^2\), then compares two interpolations: a straight line in latent space, \((1-t)\mathbf z_A+t\mathbf z_B\), decoded; and a direct blend in pixel space, \((1-t)\mathbf x_A+t\mathbf x_B\).

生成模型是 \(p_\theta(\mathbf x,\mathbf z)=\hl{prior}{p(\mathbf z)}\,\hl{lik}{p_\theta(\mathbf x|\mathbf z)}\)，先验通常固定为 \(\mathcal N(\mathbf 0,I)\)，解码器是参数为 \(\theta\) 的网络。给定数据集 \(\{\mathbf x^{(n)}\}\)，最大似然（4.6 节）要最大化 \(\sum_n\log p_\theta(\mathbf x^{(n)})\)，其中

The generative model is \(p_\theta(\mathbf x,\mathbf z)=\hl{prior}{p(\mathbf z)}\,\hl{lik}{p_\theta(\mathbf x|\mathbf z)}\), with the prior usually fixed to \(\mathcal N(\mathbf 0,I)\) and the decoder a network with parameters \(\theta\). Given a data set \(\{\mathbf x^{(n)}\}\), maximum likelihood (Lesson 4.6) maximizes \(\sum_n\log p_\theta(\mathbf x^{(n)})\), where

这个积分在高维里没法算（FIG 1 的核心洞见）。用蒙特卡洛从先验采样 \(\frac1K\sum_k p_\theta(\mathbf x|\mathbf z_k)\) 理论上无偏，实际上几乎所有先验样本都解码成与 \(\mathbf x\) 无关的图像，估计的方差大到没法用：有用的 \(\mathbf z\) 集中在后验附近，而后验正是我们不知道的。

This integral cannot be computed in high dimensions (the key insight of FIG 1). A Monte Carlo estimate from the prior, \(\frac1K\sum_k p_\theta(\mathbf x|\mathbf z_k)\), is unbiased in principle, but in practice almost every prior sample decodes into an image unrelated to \(\mathbf x\), and the variance is far too large to be useful: the useful \(\mathbf z\) are concentrated near the posterior, which is exactly what we do not know.

引入任意一个分布 \(\hl{q}{q_\phi(\mathbf z|\mathbf x)}\)（在 \(p\) 非零处非零），对 \(\log p_\theta(\mathbf x)\) 做恒等变形。因为 \(\log p_\theta(\mathbf x)\) 不依赖 \(\mathbf z\)，可以对 \(q\) 取期望，再把 \(p_\theta(\mathbf x)=p_\theta(\mathbf x,\mathbf z)/p_\theta(\mathbf z|\mathbf x)\) 代入、上下同乘 \(q\)：

Bring in any distribution \(\hl{q}{q_\phi(\mathbf z|\mathbf x)}\) (non-zero wherever \(p\) is) and rewrite \(\log p_\theta(\mathbf x)\) exactly. Since \(\log p_\theta(\mathbf x)\) does not depend on \(\mathbf z\) we may take its expectation under \(q\), substitute \(p_\theta(\mathbf x)=p_\theta(\mathbf x,\mathbf z)/p_\theta(\mathbf z|\mathbf x)\), and multiply and divide by \(q\):

因为 KL 非负（4.11 节），\(\mathrm{ELBO}\le\log p_\theta(\mathbf x)\)，这就是“证据下界”名字的由来（证据 evidence 是 \(p(\mathbf x)\) 的别名）。也可以用 Jensen 不等式一步得到：\(\log\mathbb E_q[p(\mathbf x,\mathbf z)/q]\ge\mathbb E_q[\log(p(\mathbf x,\mathbf z)/q)]\)，但上面的恒等式多告诉了我们缺口是什么。于是同时对 \(\theta\) 和 \(\phi\) 最大化 ELBO 一举两得：对 \(\theta\) 是在（近似地）做最大似然，对 \(\phi\) 是在缩小缺口，让 \(q\) 逼近真实后验（FIG 3）。

最后一行的两项正是 FIG 2 的两条曲线。重建项要求“从 \(q\) 采出的编码能把 \(\mathbf x\) 解码回来”；KL 项要求“每个 \(\mathbf x\) 的编码分布别离先验太远”。对数据集平均后，KL 项还可以写成 \(I_q(\mathbf x;\mathbf z)+D_{\mathrm{KL}}(q(\mathbf z)\|p(\mathbf z))\)，其中 \(q(\mathbf z)=\mathbb E_{\mathbf x}[q_\phi(\mathbf z|\mathbf x)]\) 是“聚合后验”：它既限制了编码携带的信息量，又要求所有编码合起来像先验，后者正是“从先验采样能得到像样的图像”的保证。

Because the KL is non-negative (Lesson 4.11), \(\mathrm{ELBO}\le\log p_\theta(\mathbf x)\), hence the name "evidence lower bound" (the evidence is another name for \(p(\mathbf x)\)). Jensen's inequality gives the bound in one line, \(\log\mathbb E_q[p(\mathbf x,\mathbf z)/q]\ge\mathbb E_q[\log(p(\mathbf x,\mathbf z)/q)]\), but the identity above also tells us what the gap is. Maximizing the ELBO over \(\theta\) and \(\phi\) together therefore does two jobs at once: over \(\theta\) it does (approximate) maximum likelihood; over \(\phi\) it shrinks the gap and pulls \(q\) towards the true posterior (FIG 3).

The two terms of the last line are the two curves of FIG 2. The reconstruction term asks that "codes drawn from \(q\) decode back to \(\mathbf x\)"; the KL term asks that "the code distribution of each \(\mathbf x\) does not stray far from the prior". Averaged over the data set, the KL term can also be written \(I_q(\mathbf x;\mathbf z)+D_{\mathrm{KL}}(q(\mathbf z)\|p(\mathbf z))\), where \(q(\mathbf z)=\mathbb E_{\mathbf x}[q_\phi(\mathbf z|\mathbf x)]\) is the "aggregate posterior": it both limits how much information the codes carry and asks that all codes together look like the prior, which is what guarantees that sampling from the prior gives sensible images.

KL 项：\(q=\mathcal N(\boldsymbol\mu,\operatorname{diag}\boldsymbol\sigma^2)\)、\(p=\mathcal N(\mathbf 0,I)\) 时，把 4.11 节的一维公式逐维相加：

The KL term: for \(q=\mathcal N(\boldsymbol\mu,\operatorname{diag}\boldsymbol\sigma^2)\) and \(p=\mathcal N(\mathbf 0,I)\), add up the one-dimensional formula of Lesson 4.11 dimension by dimension:

每一维在 \(\mu_i=0,\sigma_i=1\) 处为零：一维“坍塌”就是指它停在了这里。编码器输出 \(\log\sigma^2\) 而不是 \(\sigma\)，这样不需要约束正负，数值也稳定（本页把它截在 \([-8,4]\)）。

重建项：对每张图只采一个 \(\mathbf z\) 做蒙特卡洛估计，它的形式取决于你给解码器选的似然。伯努利解码器（本页）：\(-\log p_\theta(\mathbf x|\mathbf z)=\sum_k \mathrm{BCE}(\ell_k,x_k)\)，\(\ell_k\) 是 logit。高斯解码器 \(\mathcal N(\hat{\mathbf x}_\theta(\mathbf z),s^2I)\)：\(-\log p=\frac{1}{2s^2}\|\mathbf x-\hat{\mathbf x}\|^2+\frac D2\log(2\pi s^2)\)。于是“均方误差 + KL”这个常见写法其实默认了 \(s^2=\tfrac12\)；用 \(\|\mathbf x-\hat{\mathbf x}\|^2+\beta D_{\mathrm{KL}}\) 等价于 \(s^2=\beta/2\)：β 和解码器噪声方差是同一个旋钮。

Each dimension is zero at \(\mu_i=0,\sigma_i=1\): a "collapsed" dimension is one that stopped there. The encoder outputs \(\log\sigma^2\) rather than \(\sigma\) so that no sign constraint is needed and the numbers stay stable (this page clips it to \([-8,4]\)).

The reconstruction term is estimated by Monte Carlo with one \(\mathbf z\) per image, and its form depends on the likelihood you choose for the decoder. A Bernoulli decoder (this page): \(-\log p_\theta(\mathbf x|\mathbf z)=\sum_k \mathrm{BCE}(\ell_k,x_k)\) with logits \(\ell_k\). A Gaussian decoder \(\mathcal N(\hat{\mathbf x}_\theta(\mathbf z),s^2I)\): \(-\log p=\frac{1}{2s^2}\|\mathbf x-\hat{\mathbf x}\|^2+\frac D2\log(2\pi s^2)\). So the common recipe "mean squared error + KL" silently assumes \(s^2=\tfrac12\), and \(\|\mathbf x-\hat{\mathbf x}\|^2+\beta D_{\mathrm{KL}}\) is equivalent to \(s^2=\beta/2\): β and the decoder noise variance are the same knob.

写 \(\mathbf z=g_\phi(\boldsymbol\epsilon,\mathbf x)=\boldsymbol\mu_\phi(\mathbf x)+\boldsymbol\sigma_\phi(\mathbf x)\odot\hl{eps}{\boldsymbol\epsilon}\)，\(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\) 与 \(\phi\) 无关，于是期望可以和梯度交换：

Write \(\mathbf z=g_\phi(\boldsymbol\epsilon,\mathbf x)=\boldsymbol\mu_\phi(\mathbf x)+\boldsymbol\sigma_\phi(\mathbf x)\odot\hl{eps}{\boldsymbol\epsilon}\) with \(\boldsymbol\epsilon\sim\mathcal N(\mathbf 0,I)\) independent of \(\phi\); then expectation and gradient commute:

这三行就是 FIG 2 每一步的反传：解码器的 backward 返回 \(\partial\mathcal L/\partial\mathbf z\)，乘上 \(1\) 和 \(\tfrac12\epsilon_i\sigma_i\) 送进编码器输出，再加上 KL 的闭式梯度。得分函数估计则是 \(\nabla_\mu\mathbb E[f]=\mathbb E[f(z)\,(z-\mu)/\sigma^2]\)。对 FIG 4 的 \(f(z)=(z-a)^2\)，记 \(d=\mu-a\)，按高斯矩 \(\mathbb E\epsilon^2=1,\ \mathbb E\epsilon^4=3,\ \mathbb E\epsilon^6=15\) 展开，可得三个单样本估计的方差：重参数化 \(\hat g_{\mathrm{rp}}\)、得分函数 \(\hat g_{\mathrm{sf}}\)，以及减去基线 \(b=\mathbb E_q f\) 的得分函数 \(\hat g_{\mathrm{sf}}^{\,b}\)：

These three expressions are the backward pass of every step in FIG 2: the decoder's backward returns \(\partial\mathcal L/\partial\mathbf z\), which is multiplied by \(1\) and by \(\tfrac12\epsilon_i\sigma_i\) and fed into the encoder outputs, plus the closed-form KL gradient. The score-function estimator instead is \(\nabla_\mu\mathbb E[f]=\mathbb E[f(z)\,(z-\mu)/\sigma^2]\). For the \(f(z)=(z-a)^2\) of FIG 4, with \(d=\mu-a\), expanding with the Gaussian moments \(\mathbb E\epsilon^2=1,\ \mathbb E\epsilon^4=3,\ \mathbb E\epsilon^6=15\) gives the variances of three single-sample estimators: the reparameterized \(\hat g_{\mathrm{rp}}\), the score function \(\hat g_{\mathrm{sf}}\), and the score function with the baseline \(b=\mathbb E_q f\) subtracted, \(\hat g_{\mathrm{sf}}^{\,b}\):

\(\sigma\to0\) 时重参数化的方差趋于 0，得分函数的却像 \(d^4/\sigma^2\) 那样发散；基线去掉了发散项，但仍留下与 \(\sigma\) 无关的 \(8d^2\)。例如 \(\mu=0,\sigma=0.2\)（\(d=-1.5\)）：两者标准差之比约 31，加基线后约 10.7，和 FIG 4 实测一致。

As \(\sigma\to0\) the reparameterized variance goes to 0 while the score-function one blows up like \(d^4/\sigma^2\); the baseline removes the divergent term but leaves an \(8d^2\) that does not shrink with \(\sigma\). For example at \(\mu=0,\sigma=0.2\) (\(d=-1.5\)) the ratio of standard deviations is about 31, and about 10.7 with the baseline, matching what FIG 4 measures.

Higgins 等（2017）的 β-VAE 最小化 \(-\mathbb E_q[\log p_\theta(\mathbf x|\mathbf z)]+\beta D_{\mathrm{KL}}(q\|p)\)。β > 1 逼编码器用更少的信息，常被用来换取更“解耦”的隐维度；β < 1 换取更好的重建。只有 β = 1 时目标才是 ELBO，此时比较不同模型的 ELBO 才有意义。后验坍塌是指 KL 项赢了：\(q_\phi(\mathbf z|\mathbf x)\approx p(\mathbf z)\) 对所有 \(\mathbf x\) 成立，\(\mathbf z\) 不再携带任何信息，解码器只好对所有输入输出数据的平均（FIG 2 里 β = 10 时重建项升到约 27 nats，β = 1 时还不到 15）。它在两种情况下最常见：KL 权重相对太大（大 β，或者重建项用了“对像素取平均”而不是求和，见第 7 节误区 3）；以及解码器强到不需要 \(\mathbf z\) 也能把 \(\mathbf x\) 建模好，例如自回归解码器（19.4 节，Bowman 等 2016 在文本 VAE 上首先报告）。常见对策：训练初期把 β 从 0 慢慢升到 1（KL 退火），“free bits”（每维 KL 低于某个阈值时不罚），或者削弱解码器。

The β-VAE of Higgins et al. (2017) minimizes \(-\mathbb E_q[\log p_\theta(\mathbf x|\mathbf z)]+\beta D_{\mathrm{KL}}(q\|p)\). β > 1 forces the encoder to use less information and is often used to buy more "disentangled" latent dimensions; β < 1 buys better reconstructions. Only at β = 1 is the objective the ELBO, and only then does comparing the ELBOs of different models make sense. Posterior collapse is the KL term winning: \(q_\phi(\mathbf z|\mathbf x)\approx p(\mathbf z)\) for every \(\mathbf x\), \(\mathbf z\) carries no information, and the decoder can only output the average of the data for every input (in FIG 2 at β = 10 the reconstruction term rises to about 27 nats, against under 15 at β = 1). It is most common in two situations: when the KL weight is relatively too large (a big β, or a reconstruction term averaged over pixels instead of summed, Section 7 pitfall 3); and when the decoder is strong enough to model \(\mathbf x\) without \(\mathbf z\), such as an autoregressive decoder (Lesson 19.4; first reported for text VAEs by Bowman et al., 2016). The usual remedies: ramp β from 0 up to 1 early in training (KL annealing), "free bits" (no penalty while the KL of a dimension stays below a threshold), or a weaker decoder.

机器人模仿学习最常见的难题之一：示教者面对同一个场景会做出不同但都正确的动作。下面的任务里，夹爪要从 (0, −1) 走到 (0, 1)，中间 \(y=0\) 处有一个半径 0.28 m 的障碍物，横向位置 \(o\) 每次不同。示教者从障碍物左边或右边 0.5 m 处绕过（加 0.06 m 的噪声），障碍物越靠右越倾向于走左边（走左边的概率 \(1/(1+e^{-4o})\)）。动作 \(a\) 就是路径在 \(y=0\) 处的横向位置，路径是过起点、\((a,0)\) 和终点的二次贝塞尔曲线。共 400 条示教。

用均方误差回归 \(a=\pi(o)\)（22.1 节的行为克隆）会学到条件均值 \(\mathbb E[a|o]\)，两个模式的平均，在 \(o\approx0\) 时正好穿过障碍物。条件 VAE（CVAE，Sohn 等 2015）把所有分布都以条件 \(\mathbf c=o\) 为条件：编码器 \(q_\phi(z|a,o)\)，解码器 \(p_\theta(a|z,o)=\mathcal N(\hat a_\theta(z,o),0.08^2)\)，先验仍用固定的 \(\mathcal N(0,1)\)，ELBO 变成

One of the most common difficulties of robot imitation learning: faced with the same scene, demonstrators make different but equally correct moves. In the task below a gripper travels from (0, −1) to (0, 1), and at \(y=0\) there is an obstacle of radius 0.28 m whose lateral position \(o\) changes every time. The demonstrator passes 0.5 m to the left or to the right of it (plus 0.06 m of noise), preferring the left the further right the obstacle is (left with probability \(1/(1+e^{-4o})\)). The action \(a\) is the lateral position of the path at \(y=0\); the path is the quadratic Bézier curve through the start, \((a,0)\) and the goal. There are 400 demonstrations.

Regressing \(a=\pi(o)\) with a squared error (the behaviour cloning of Lesson 22.1) learns the conditional mean \(\mathbb E[a|o]\), the average of the two modes, which at \(o\approx0\) goes straight through the obstacle. A conditional VAE (CVAE, Sohn et al., 2015) conditions every distribution on \(\mathbf c=o\): an encoder \(q_\phi(z|a,o)\), a decoder \(p_\theta(a|z,o)=\mathcal N(\hat a_\theta(z,o),0.08^2)\), and the prior kept as the fixed \(\mathcal N(0,1)\); the ELBO becomes

直觉上，一维的 \(\hl{cz}{z}\) 学会了编码“这次走哪边”，这正是光看 \(o\) 无法确定、只有看到动作才知道的那部分信息。测试时没有动作可编码，就从先验采 \(z\) 再解码，每个样本都是一个完整的、在某一边的动作。两个网络都是小型 tanh MLP（编码器 2 → 32 → 2，解码器 2 → 32 → 32 → 1；MSE 策略 1 → 32 → 32 → 1），每步 32 条示教，Adam，学习率 0.003。

Intuitively the one-dimensional \(\hl{cz}{z}\) learns to encode "which side this time", precisely the information that \(o\) alone cannot determine and only the action reveals. At test time there is no action to encode, so \(z\) is drawn from the prior and decoded, and every sample is one complete action on one side. The networks are small tanh MLPs (encoder 2 → 32 → 2, decoder 2 → 32 → 32 → 1; the MSE policy 1 → 32 → 32 → 1), 32 demonstrations per step, Adam with learning rate 0.003.

注意那 10% 到 15% 的碰撞样本：解码器是连续函数，\(z\) 从一边的模式连续变到另一边时必然经过中间，而先验在 \(z\approx0\) 附近恰好密度最大。所以“取先验均值 \(z=0\)”这个看似保险的做法，在两个模式差不多一样可能时（\(o\approx0\)），反而正落在缝里（读数“z = 0 的路径”）。这是连续隐变量模型对付多模态的一个结构性弱点，也是扩散模型和 flow matching（19.6、19.7 节）在动作生成上受欢迎的原因之一；22.2 节会把 ACT 与 Diffusion Policy 放在一起比较。

Note the 10% to 15% of colliding samples: the decoder is a continuous function, so as \(z\) moves from one mode to the other the output must pass through the middle, and the prior is densest right around \(z\approx0\). So the seemingly safe choice "take the prior mean \(z=0\)" lands right in the gap when the two modes are about equally likely (\(o\approx0\); see the readout "path at z = 0"). This is a structural weakness of continuous latent-variable models facing multimodality, and one reason diffusion and flow matching (Lessons 19.6 and 19.7) became popular for action generation; Lesson 22.2 puts ACT and Diffusion Policy side by side.