熵变分推断自编码 EVIA

用熵正则最优传输匹配聚合后验,得到玻尔兹曼后验与软重心编码器 —— WAE 的随机化推广

EVIA(Entropic Variational Inference Auto-encoding,熵变分推断自编码)是一类自编码算法:它用**熵正则最优传输**来匹配聚合后验与先验。可以把它看作 WAE 的推广 —— WAE 以最优传输代价匹配分布,EVIA 在该代价上再加一项熵,从而把”逐点确定性的编码器”换成”对每个输入给出一族潜变量的玻尔兹曼后验”。

软重心 T*(x)z ~ q*(z|x)观测 x编码器玻尔兹曼后验q*(z|x) ∝ κ·e^{G/γ}解码器 A重构
γ→0:后验坍缩为单点 —— 即 WAE / CryoGEN-II 的确定性映射;γ>0:一族编码,刻画缺失楔形带来的不确定性。

一句话概括三者的关系:VAE 逐样本地把每个 q(zx)q(z\mid x) 拉向先验;WAE 只约束聚合后验 qZ=q(zx)pX(x)dxq_Z=\int q(z\mid x)\,p_X(x)\,dx 而放任编码器收缩成一张确定性映射;EVIA 同样只约束聚合后验,却用一个温度 γ\gamma 把编码器”撑开”,让它在确定映射(γ0\gamma\to0)与弥散分布(γ\gamma 大)之间连续可调。这条温度旋钮正是后文所有结构的来源。

直觉

以最优传输匹配分布时,若不加正则,最优解往往退化成一张确定性映射(Monge 映射)—— 编码器对每个 xx 只吐一个 zz。但推断本应是分布性的:同一个 xx 可以对应多个合理的潜表示。EVIA 在传输代价里加一项熵,像一道势垒,逼迫解保持”摊开”,于是后验重新变成处处为正、可采样的分布。

把这件事放到缺失楔形修复上看就更具体:一张倾斜系列重构出的层析图,因为没有采到高角度投影,沿缺失方向的信息是真的丢了。一个确定性编码器只能挑出”一个”最可能的填充,把推断的不确定性藏起来;而 EVIA 给出的是”一族”与观测一致的填充,温度 γ\gamma 直接对应”我们对没看到的方向有多不确定”。

熵正则与玻尔兹曼后验

设数据 xx、潜变量 zz、解码器 A:zx\mathcal{A}:z\mapsto x,参考耦合 κ\kappa、温度 γ>0\gamma>0。EVIA 的原始目标是带熵罚的最优传输:

minπΠ(px,qz)  E(x,z)π[xA(z)2]  +  γKL(πκ).\min_{\pi\in\Pi(p_x,q_z)}\;\mathbb{E}_{(x,z)\sim\pi}\big[\|x-\mathcal{A}(z)\|^2\big]\;+\;\gamma\,\mathrm{KL}(\pi\,\|\,\kappa).

逐项读这个式子:πΠ(px,qz)\pi\in\Pi(p_x,q_z) 是以数据边缘 pxp_x 和潜分布边缘 qzq_z 为约束的所有联合耦合;第一项 ExA(z)2\mathbb{E}\,\|x-\mathcal{A}(z)\|^2 是传输代价,要求被配到一起的 (x,z)(x,z) 经解码后能重构出 xx;第二项 γKL(πκ)\gamma\,\mathrm{KL}(\pi\,\|\,\kappa) 是相对参考耦合 κ\kappa 的熵罚,温度 γ\gamma 是它的权重。γ\gamma 越大,越偏好接近 κ\kappa、越”摊开”的耦合;γ0\gamma\to0 时熵罚消失,目标退回纯传输,最优耦合塌成确定映射。把它和上一节 WAE 的目标比较:WAE 是 γ=0\gamma=0 的特例,EVIA 只是给同一个传输问题加了温度。

深入

引入先验势 w(z)w(z)(对应聚合后验的边缘约束;在 Cryo-ET 中由对抗学得的能量评判器 Dψ=EϕD_\psi=-E_\phi 实现),再以数据拟合权重 λ\lambda 定义效用 G(z;x)=w(z)λ2xA(z)2\mathcal{G}(z;x)=w(z)-\tfrac{\lambda}{2}\|x-\mathcal{A}(z)\|^2。由吉布斯变分原理(Donsker–Varadhan),固定 xx 时最优条件分布取吉布斯(玻尔兹曼)形式:

q(zx)    κ(zx)exp ⁣(G(z;x)γ),q^\star(z\mid x)\;\propto\;\kappa(z\mid x)\,\exp\!\Big(\frac{\mathcal{G}(z;x)}{\gamma}\Big),

处处为正。这与熵正则最优传输的玻尔兹曼耦合 πκec/γ\pi^\star\propto\kappa\,e^{-c/\gamma} 同源,区别仅在多了对抗项 wwγ0\gamma\to0 时退回 WAE 的确定性硬传输。

逐项的物理意义:效用 G\mathcal{G} 把两股力写在一起 —— w(z)w(z) 奖励”看起来像真分布”的潜码(高能量评判分),λ2xA(z)2-\tfrac{\lambda}{2}\|x-\mathcal{A}(z)\|^2 惩罚重构不上 xx 的潜码;λ\lambda 调两者的相对话语权。指数里的 G/γ\mathcal{G}/\gamma 是统计力学里熟悉的”负能量除以温度”:γ\gamma 小,分布尖锐地压在效用最高的 zz 上,逼近一个 argmax\arg\maxγ\gamma 大,分布被抹平,接近参考 κ\kappaww 这一项把单纯的传输匹配换成了能量基模型式的密度匹配,正是它把 EVIA 接到了能量基模型与对抗训练上。

温度 γ\gamma 如何控制后验的宽度与重构的不确定性 —— 拖动温度看:

能量最低 = MAP样本重构
能量 E(x)后验 q(x|y) ∝ e^(−E/γ)

后验宽 —— 一族重构 → 缺失楔形不确定性(CryoWGEN)

温度 γ 直接决定后验的宽度。把数据一致性写成能量 E(x)(琥珀井),后验就是井里的玻尔兹曼分布 q(x|y) ∝ e^(−E(x)/γ)(紫)。γ→0 时它塌成井底的一根尖峰 —— 单一确定重构,正是 WAE / CryoGEN-II;γ 变大,它在井里摊开成一族重构,这个宽度就是缺失楔形留下的不确定性,也就是 CryoWGEN 报告的东西。底部的紫色刻度是从后验抽出的样本重构,随 γ 增大而散开。

往两个极端拨这个旋钮:γ0\gamma\to0 时玻尔兹曼分布坍缩到效用最高的单点,EVIA 变成 WAE 的确定性编码器,给出”最佳猜测”的单一重构;γ\gamma 很大时后验摊回参考 κ\kappa,编码器几乎不看数据。Cryo-ET 里有意义的工作区在两者之间:γ\gamma 取得让后验恰好覆盖缺失楔形所允许的那一族解,既不假装确定,也不退化成噪声。

软重心编码器

最优编码器不是某个采样,而是后验下潜变量的条件期望——一个软重心投影:

T(x)=Eq(x)[z]    mωm(x)z(m),ωm(x)=exp(G(z(m);x)/γ)jexp(G(z(j);x)/γ),T^\star(x)=\mathbb{E}_{q^\star(\cdot\mid x)}[z]\;\approx\;\sum_{m}\omega_m(x)\,z^{(m)},\qquad \omega_m(x)=\frac{\exp(\mathcal{G}(z^{(m)};x)/\gamma)}{\sum_j\exp(\mathcal{G}(z^{(j)};x)/\gamma)},

其中 {z(m)}\{z^{(m)}\} 取自参考分布。权重是一个 softmax,因此该映射平滑可微,并随 γ0\gamma\to0 收敛到经典硬最优传输映射。

具体怎么算:从参考分布抽一批候选潜码 z(1),,z(M)z^{(1)},\dots,z^{(M)},给每个算效用 G(z(m);x)\mathcal{G}(z^{(m)};x),做一个温度为 γ\gamma 的 softmax 得到权重 ωm\omega_m,再把候选按权重加权平均。这就是”软”——不是硬选效用最高的那个 zz,而是按它们的吉布斯权重取重心。两个极端再次自洽:γ0\gamma\to0 时 softmax 趋于 one-hot,重心塌到单个 argmax\arg\max,恢复硬最优传输映射;γ\gamma 大时权重趋于均匀,T(x)T^\star(x) 退回候选的简单平均、几乎不依赖 xx。因为整条链路(效用、softmax、加权和)都可微,编码器能直接端到端反传训练。

两种实现:SGLD 与 amortized

把条件先验设为各向同性高斯 κ(zx)=N(zˉ(x),β1I)\kappa(z\mid x)=\mathcal{N}(\bar z(x),\beta^{-1}I),待最小化的有效势是一个可解的 Log-Sum-Exp:

Ψ(x)=log ⁣exp ⁣{w(z)λ2xA(z)2β2zzˉ2}dz.\Psi(x)=-\log\!\int\exp\!\Big\{w(z)-\tfrac{\lambda}{2}\|x-\mathcal{A}(z)\|^2-\tfrac{\beta}{2}\|z-\bar z\|^2\Big\}\,dz.

读这个有效势:高斯先验把 κ(zx)exp(β2zzˉ2)\kappa(z\mid x)\propto\exp(-\tfrac{\beta}{2}\|z-\bar z\|^2) 直接并进指数,于是积分里集齐了三项——能量评判 w(z)w(z)、重构罚 λ2xA(z)2-\tfrac{\lambda}{2}\|x-\mathcal{A}(z)\|^2、把 zz 拉回先验均值 zˉ(x)\bar z(x) 的高斯罚(强度 β\beta)。前面取负对数,Ψ(x)\Psi(x) 就是这族解的”自由能”;它对参数的梯度正是训练所需的信号。困难只在那个对 zz 的积分,两种实现就是对它的两种处理。

权衡很直接:SGLD 每个样本都要跑一段朗之万链,慢,但保留了完整的后验不确定性;amortized 用一次前向换掉内层采样,快得多,但交付的是后验的点估计(重心)。这条 SGLD 与 amortized 的分界,正对应到 Cryo-ET 落地时 CryoWGEN-II(朗之万、给出一族后验解)与单答案变体之间的取舍。

在自编码器谱系中的位置

算法后验匹配方式后验形态
VAE逐样本 KL高斯 qϕ(zx)q_\phi(z\mid x)
WAE聚合后验的最优传输常为确定性编码器
AAE对抗(密度比)由判别器隐式定义
EVIA熵正则最优传输玻尔兹曼后验(处处为正)

读这张表的一条主线是”后验被约束得多紧”:VAE 把每个样本都钉在先验附近,最严;WAE/AAE 只约束聚合后验、放开单样本,靠不同手段(闭式 MMD vs. 对抗判别器)估同一个散度;EVIA 同样只约束聚合,但用温度把”硬”传输连续地软化成处处为正的玻尔兹曼后验,恰好把前三者收进同一族中——γ0\gamma\to0 即 WAE,加上对抗能量项 ww 即接上 AAE 那条对抗匹配的思路。

EVIA 把缺失楔形修复奠基为一个分布化的反问题:观测(缺角的层析图)约束不出唯一解,正确的输出本就是一族与数据一致的体积,而非一张图。在 Cryo-ET 上,能量评判 w=Eϕw=-E_\phi 由对抗训练学得,SGLD/amortized 两条实现分别对应”采一族后验”与”取后验重心”。具体落地见 CryoWGEN;其朗之万版本给出一族后验解,见 CryoWGEN-II

← 生成与分布匹配