Variational Inference

Gaoustcer 4月 28, 2023

Variational Inference

近似推断

Reference

Bayes Inference and prediction

Inference意在求出参数在给定样本下的后验分布

$p(\theta|x) = \frac{p(x|\theta)p(\theta)}{p(x)}$

prediction的目标是预测新样本$\hat x$在给定样本x下的后验分布

$\begin{aligned} p(\hat x|x) &= \int_\theta p(\hat x,\theta|x) d\theta\\ &= \int_\theta p(\hat x|\theta)p(\theta|x) d\theta\\ &= E_{\theta|p(\cdot|x)} E[p(\hat x|\theta)] \end{aligned}$

计算后验可以通过

精确推断，得到$p(\theta|x)$的解析解
近似推断

ELBO和变分推断

Variational Inference and ELBO

给定观测数据$x$和隐变量z

$x = (x_1,x_2,\cdots,x_n)\\ z = (z_1,z_2,\cdots,z_m)$

推断出隐变量的后验概率$p(z|x)$，写成

$p(z|x) = \frac{p(z,x)}{p(x)}$

计算$p(x)$需要对全空间概率分布积分

$p(x )= \int p(x,z) dz$

一个思路是从分布簇Q中选择一个分布$q\in Q$，使得它和后验分布$p(z|x)$差异最小

$q^* =\arg\min_{q\in Q} KL(q(z)\parallel p(z|x))$

用分布参数$\Theta$表示为

$\theta^* = \arg\min_{\theta\in \Theta} KL(q_\theta(z)\parallel p(z|x))$

KL散度写成

$\begin{aligned} KL(q(z)\parallel p(z|x)) &= E_{z|q(\cdot)}[\log q(z)] - E_{z|q(\cdot)}[\log p(z|x)]\\ &=E_{z|p(\cdot)}[\log q(z)] -E_{z|q(\cdot)}[\log p(z,x)]+E_{z|q(\cdot)}[\log p(x)]\\ &=E_{z|p(\cdot)}[\log q(z)] -E_{z|q(\cdot)}[\log p(z,x)]+\log p(x)\\ &=\log p(x)- ELBO(q) \end{aligned}\\ ELBO (q) =\log p(x) - KL(q(z)\parallel p(z|x)) = E_{z|q(\cdot)}[\log p(z,x)]-E_{z|q(\cdot)}[\log q(z)]= \int_z q(z) \log p(z,x) dz-\int_z q(z)\log q(z) dz$

目标变成最大化ELBO

从最大化$\log p(x)$的角度上看，由于
$\log p(x) = ELBO(q) +KL(q(z)\parallel p(z|x))\geq ELBO$
ELBO是观测数据的极大似然估计的下确界

进一步分解ELBO

$EBLO(q) = E_{z|q(\cdot)}[\log p(x|z)] -KL(q(z)\parallel p(z))$

观测数据观测数据观测数据，$p(z)$是latent variable的先验分布，最大化ELBO实际上分解为两个目标

解释解释解释解释解释解释释*)
最小化隐变量分布和变分先验分布距离

选择分布族Q

分布族Q的复杂程度会影响优化算法的复杂度，这里假设分布族是平均场变分分布族，不需要了解详细的定义，只需要直到这种分布下满足

每一个隐变量$z_j$都相互独立，由自身的参数$q_j$影响，这样多个隐变量的联合分布写成
$q(z) = \prod_j q_j(z_j)$

求解最佳分布$q_{\theta_j}$(选择最佳参数$\theta^*$)

梯度上升最大化ELBO

$ELBO(q_\theta) = E_{z|q_\theta(\cdot)}[\log p(z,x)] - E_{z|q_\theta(\cdot)}[\log q_\theta(z)]$

我们采取Block Coordinate Descent的方法求解这个优化问题

Block Coordinate Descent，对于包含T个变量的损失函数
$J(\alpha_1,\alpha_2,\cdots,\alpha_T)$
希望找到$\alpha_i^*$使得J最小，迭代优化每个变量
$\alpha_t = \arg\min_{\alpha_t}J(\alpha_1,\cdots,\alpha_T)$
这里$z_i$相互独立，因此可以迭代优化各个变量的分布参数

利用分布的独立性拆分ELBO

$A = E_{z|q(\cdot)} [\log p(z,x)],B = E_{z|q_\theta(\cdot)}[\log q_\theta(z)] \\ \begin{aligned} A&=\int _z \prod_{i=1}^m q_{\theta_i} (z_i) \log p(x,z) dz_1 dz_2\cdots d z_m\\ &=\int_{z_j} q_{j}(z_j)\prod_{i\neq j} q_i(z_i)\log p(x,z) dz_i dz_j \end{aligned}$

中间除去$z_j$的积分写成

$\int_{z_i,i\neq j} \log p(x,z) q_i(z_i) d_i = E_{\prod_{i\neq j} q_i(z_i)}[\log p(x,z)]$

$z_j$视为固定参数，$\log p(x,z)$视为m-1个随机变量的函数，因此A写成

$A= \int_{z_j} q_j(z_j)E_{\prod_{i\neq j} q_i(z_i)} [\log p(x,z)] dz_j= \int_{z_j} q_j(z_j) \log P(x,z_j) d z_j$

再拆分B

$\begin{aligned} B &= \int_{z} \prod_{i=1}^m q_i(z_i) \sum_{i=1}^m \log q_i(z_i) dz_i\\ &= \sum_{k=1}^m\int_z \prod_{i} q_i(z_i)\log q_{k}(z_k)\\ &=\sum_{k=1}^m \int_{z_k}q_k(z_k) \log q_k(z_k) \\ \end{aligned}$

与特定latent variable有关的仅包含$\int_{z_j} q_j(z_j)\log q_j(z_j)$，这样ELBO中和$z_j$有关的项写成

$\begin{aligned} ELBO(\theta_j)& = \int_{z_j} q_{\theta_j}(z_j) \log P(x,z_j) dz_j - \int _{z_j} q_{\theta_j}(z_j) \log q_{\theta_j}(z_j) dz_j\\&=\int_{z_j}q_{\theta_j}(z_j)\log \frac{P(x,z_j)}{q_{\theta_j}(z_j)} d z_j + \sum_{i\neq j}\int_{z_i} q_{\theta_i}(z_i)\log q_{\theta_i}(z_i) d z_i\\ &=-KL (q_{\theta_j}\parallel P(x,z_j))+C \end{aligned}$

KL散度最小，必然有

$q_{\theta_j}(z_j) = P(x,z_j) =\int_{z_i,i\neq j}P(x,z) q_{\theta_i}(z_i) dz_i$

坐标上升变分推断

变分推断目标是找到一个近似概率q，使得它和后验分布$p(z|x)$尽量相似

$\min KL(q\parallel p(z|x)) \to \arg\max_{q\in Q}ELBO(q)= E_{z|q_\theta(\cdot)}[\log p(z,x)] - E_{z|q_\theta(\cdot)}[\log q_\theta(z)]$

样本记作

$x = x^1,x^2,\cdots,x^n\\ z = z^1,z^2,\cdots,z^n$

$z^i$是$x^i$对应的latent variable，假设$x^i$ iid

$\log P_\theta(x)=\sum_{i=1}^n \log P_\theta(x^i)\Rightarrow \max \log P_\theta(x^i)$

严格意义上ELBO写成(instance wise)

$\log P_\theta(x^i) =ELBO+KL(q\parallel p) \\ ELBO(q(z) )=E_{q(z)}[\log P_\theta(x^i,z)]+H(q(z))$

注意区分$z_i$和$z^i$，前者是latent variable的某个维度，后者代表一个instance对应的随机变量，基本假设$z_i$相互独立

最优解满足(迭代求解)

$q_{\theta_j} (z_j) = E_{\prod_{i\neq j} p_{\theta_i}(z_i)}[\log p(x,z)]$

求解更新概率
带入最新概率求解下一个variable分布

$\hat q_1 (z_1) =\int_{q_2,\cdots,q_m} [\log P_\theta(x^i,z)] dq_1\cdots dq_m\\ \hat q_2 (z_2) =\int_{\hat q_1,q_3,\cdots,q_m} [\log P_\theta(x^i,z)] d\hat {q_1}d q_3 \cdots d q_m$

这就是坐标上升，每一个latent variable的每一维度都需要通过如下推断

Problems of Classical Variational Inference

latent variable每一维独立性很难保证

例子：深度神经网络latent embedding不能保证独立

递推本身需要求m-1维期望(intractable)

随机梯度变分推断(解决latent variable不独立的情况)

$z\to x$(generative model,Decoder) $x\to z$(inference model,Encoder)

Intuition：使用随机梯度上升解决ELBO最大

假设

待求解分布$q(z)$参数化$q_\phi(z)$，目标是找到最好的分布参数

$\max L(\phi) =E_{z|q_\phi(\cdot)}[\log p_\theta(z,x)] - E_{z|q_\phi(\cdot)}[\log q_\phi(z)]$

$\theta$是联合分布模型的参数

求梯度

$\begin{aligned} \nabla _\phi L(\phi) &= \nabla_\phi E_{z|q_\phi}[\log p_\theta(z,x)-\log q_\phi(z)]\\ &=\int_z \nabla_\phi q_\phi(z)[\log p_\theta(z,x) -\log q_\phi(z)] dz-\int_z q_\phi(z) \nabla \log q_\phi(z) dz \\ \int_z q_\phi(z)\nabla \log q_\phi(z) dz &= \int_z \nabla_\phi q_\phi(z) dz \\ &=\nabla _\phi \int_z q_\phi(z) dz = 0(constant)\\ \int _z \nabla q_\phi(z)[\log p_\theta(z,x) - \log q_\phi(z)] dz &=\int_z q_\phi(z)\nabla _\phi\log q_\phi(z) [\log_\theta(z,x) - \log q_\phi(z)] dz\\ (只对q_\phi(z)求偏导，对-\log q_\phi(z)不变)&=E_{z|q_\phi(\cdot)}[\nabla _\phi \log q_\phi(z)(\log p_\theta(z,x) - \log q_\phi(z)) \end{aligned}$

这个期望可以通过蒙特卡洛获得，对于给定样本$x^1,x^2,\cdots,x^n$，采样对应的n个latent variable $z^1,z^2,\cdots,z^n$，得到期望估计是

$\frac 1 n (\sum_{i=1} \nabla_\phi \log q_\phi(z^i)[\log p_\theta(z^i,x^i)-\log q_\phi(z^i)])$

容易导致梯度爆炸，这是因为容易采样到$z^i$使得$q_\phi(z^i)$非常小，导致期望估计不准(方差较大)

进行Variance Reduction(重参数化)

Reference

重参数化技巧

最初重参数化技巧用于解决蒙特卡洛从参数分布中采样无法进行梯度传播的问题，求解给定分布下某个随机变量函数的期望等于求解如下积分

$J_\theta =\int p_\theta(z) f(z) dz = E_{z|p_\theta(\cdot)}[f(z )]$

解决方法是将从$p_\theta(z)$中采样改成从无参数分布$q(\epsilon)$中采样，随后通过变换$z = g_\theta(\epsilon)$生成z，这样生成的z携带参数

$L_\theta = E_{\epsilon|q(\epsilon)} [f(g_\theta(\epsilon))]$

计算正态分布$N(\mu_\theta,\sigma_\theta^2)$下积分值，可以从$\mathcal N(0,1)$中采样$\epsilon_i$，计算
$z_i = \epsilon \times \sigma_\theta + \mu_\theta$
随后用
$\frac{\sum_i g(z_i)}{n}$
作为可以梯度传播的方差

离散情况直接用计算出的离散概率乘上函数值
$\sum_y p_\theta(y) f(y)$
但是这种做法在类别很多的情况不适用，Gumble Max在离散分布下采样实现梯度传播

Gumble Max

对于类别数为k的分布$p_1,p_2,\cdots,p_k$，从$U(0,1)$采样k个随机变量

$\epsilon_i|U(0,1)$

采样

$\arg\max_i (\log p_i-\log(-\log (\epsilon_i)))$

精确等价于按照$p_1,p_2,\cdots,p_k$p_k$**离散采样，将argmax输出的结果视为0-1向量，用softmax光滑之

为了解决argmax不可导的问题，用连续函数近似one-hot将其转化为平滑分布Gumble_softmax

$softmax(\frac{(\log p_i-\log(-\log \epsilon_i))}{\tau})$

这个想法实际上可以写成

$\arg\max(x) \approx\sum_i i softmax(x)$

本站访客数人次